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 = "WarriorMama777/OrangeMixs"
model = "cagliostrolab/animagine-xl-3.1"
# model = "Vsukiyaki/ShiratakiMix"

# pipe = StableDiffusionXLPipeline.from_pretrained(
#     model,
#     torch_dtype=torch.float16,
#     use_safetensors=True,
#     variant="fp16",
#     )
# path = r"C:\Users\office54\Desktop"
# base_dir = r"C:\dev\ai\stable-diffusion-webui-forge"
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_name = ""
# print(lora_dir)
# print (os.path.join('C:' + os.sep,'tmp','test_data.csv'))
# print(lora_dir)
# print(os.listdir(lora_dir))

# sys.exit()

# model = model_dir + '/AbyssOrangeMix',


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

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


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

# pipe.load_lora_weights(".", weight_name="Kizuki - To Love Ru - Kurosaki Mea")
# pipe.load_lora_weights("lora/Kizuki - To Love Ru - Kurosaki Mea.safetensors")
# pipe.load_lora_weights(os.path.join(base_dir + os.sep + "Yabuki Kentarou_XL.safetensors"))
# pipe.load_lora_weights(

#     weight_name= os.path.join(base_dir + os.sep + "Yabuki Kentarou_XL.safetensors")
# )
# pipe.load_lora_weights(
#     lora_dir,
#     weight_name= "Yabuki Kentarou_XL.safetensors"
#     # weight_name= "mea-loveru-01.safetensors"
# )



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


# prompt = "face focus, cute, masterpiece, best quality, 1girl, green hair, sweater, looking at viewer, upper body, beanie, outdoors, night, turtleneck"
# prompt = "face focus, cute, masterpiece, best quality, 1girl, pink hair, virtual youtuber, playing a game, upper body, beanie, indoor"
# prompt = "cute girl,outdoor,scenery"
# negative_prompt = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry"
# prompt = "masterpiece, best quality,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"
# prompt = "masterpiece, best quality,1girl, Makinami Mari Illustrious, neon genesis evangelion, plugsuit, pilot suit, pink bodysuit, sitting, crossing legs, glass, cat hat, throne, symmetrical, looking down, from bottom, looking at viewer, outdoors"
# prompt = "masterpiece, best quality,1girl, ToLove-ru darkness Mea Kurosaki, black suit ,throne, symmetrical, looking down, from bottom, looking at viewer, outdoors"
# prompt = "1girl, kurosaki mea, hair ornament, short hair, hair intakes, absurdly long braid, red hair, purple eyes, medium breasts, braided ponytail,swimsuit bet in hotel, night,smile"
prompt = "1girl, solo, masterpiece, best quality, <lora:mea-loveru:0.8>, sweat, to love-ru, to love-ru darkness, mea kurosaki, kurosaki mea, blue eyes, red hair, long hair, braid, single braid, ahoge, hair ornament, gym uniform, gym shirt, white shirt, sleeveless, sleeveless shirt, collarbone, bare arms, cleavage, taut clothes, medium breasts, buruma, bare legs, breasts, cowboy shot, leaning forward, v arms, from above, looking up, looking at viewer, blush, smile, half-closed eyes, outdoors"
# negative_prompt = "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]"
negative_prompt = "verybadimagenegative_v1.3, (cameltoe:1.5), (black background, empty eyes:1.5)"

output = "content/"+str(current_time)+".png"

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


# image = pipe(
#     prompt,
#     negative_prompt=negative_prompt,
#     width=1024,
#     height=1024,
#     guidance_scale=12,
#     target_size=(1024,1024),
#     original_size=(4096,4096),
#     num_inference_steps=50
#     ).images[0]

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