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


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



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 = StableDiffusionXLPipeline.from_pretrained(
    model,
    torch_dtype=torch.float16,
    use_safetensors=True
)


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



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




model_id = "stabilityai/sd-x2-latent-upscaler"
upscaler = StableDiffusionLatentUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
upscaler.to("cuda")

upscaled_image = upscaler(
    prompt=prompt,
    image=base_image,
    num_inference_steps=20,
    guidance_scale=0,
    # generator=generator,
).images[0]



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



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