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



model_name ="cagliostrolab/animagine-xl-3.1"

base_dir = os.path.join('C:' + os.sep,'dev','ai','stable-diffusion-webui-forge')
lora_dir = os.path.join(base_dir, 'models', 'lora')
model_dir = os.path.join(base_dir, 'models', 'diffusers')
model = os.path.join(model_dir , model_name)




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


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


current_time = time.strftime("%Y%m%d%H%M%S")
prompt = "masterpiece, best quality, very aesthetic, 1girl, souryuu asuka langley, neon genesis evangelion, plugsuit, pilot suit, red bodysuit, sitting, crossing legs, black eye patch, cat hat, throne, symmetrical, looking down, from bottom, looking at viewer, outdoors"
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]"
output = "content/"+str(current_time)+".png"

generator = torch.Generator()
# generator.manual_seed(100)

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


image.save(output)
image = Image.open(output)
plt.imshow(image)
plt.axis('off') # to hide the axis