import torch
from torch import autocast
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler,AutoencoderKL,StableDiffusionPipeline,DiffusionPipeline
import matplotlib.pyplot as plt
from PIL import Image
import time
import os
import sys


# model = "./model/Counterfeit-V3.0_fp16.safetensors"
model = "cagliostrolab/animagine-xl-3.1"
# model = "gsdf/Counterfeit-V2.5"

output_dir = "./output"
if not os.path.exists(output_dir):
    # ディレクトリが存在しない場合、ディレクトリを作成する
    os.makedirs(output_dir)




# ダウンロードしたVAEファイルを指定
# vae = AutoencoderKL.from_single_file("./vae/Counterfeit-V2.5.vae.pt")
# vae = AutoencoderKL.from_single_file("./vae/vae-ft-mse-840000-ema-pruned.safetensors")
# vae = AutoencoderKL.from_single_file("./vae/clearvae_v23.safetensors")


# VAEを指定してPipeline作成
# pipe = StableDiffusionPipeline.from_single_file(
#     "./model/Counterfeit-V3.0_fp16.safetensors",
#     torch_dtype=torch.float16,
#     vae=vae
# )
pipe = StableDiffusionXLPipeline.from_pretrained(
    model,
    torch_dtype=torch.float16,
    # vae=vae,
    # extract_ema=True,
    # use_safetensors=True,
    # safety_checker = None,
    # requires_safety_checker = False
)

pipe.safety_checker = None
pipe.requires_safety_checker = False
# pipe.to("cuda")


prompt = "masterpiece, best quality, 1girl, solo, flower, long hair, outdoors, letterboxed, school uniform, day, sky, looking up, short sleeves, parted lips, shirt, cloud, black hair, sunlight, white shirt, serafuku, upper body, from side, pink flower, blurry, brown hair, blue sky, depth of field"
negative_prompt = "worst quality, low quality, nsfw"
# model ="cagliostrolab/animagine-xl-3.1"
# prompt = "1girl, souryuu asuka langley, neon genesis evangelion, solo, upper body, v, smile, looking at viewer, outdoors, night"
# negative_prompt = "nsfw, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]"



# pipe = DiffusionPipeline.from_pretrained(
#     model,
#     torch_dtype=torch.float16,
#     use_safetensors=True
# )


# pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.to('cuda')



image = pipe(
    prompt,
    negative_prompt=negative_prompt,
    width=832,
    height=1216,
    guidance_scale=7,
    num_inference_steps=28
).images[0]

current_time = time.strftime("%Y%m%d%H%M%S")



image.save(output_dir+"/"+current_time+".png")