import torch
from torch import autocast
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler

import matplotlib.pyplot as plt
from PIL import Image

# model = "linaqruf/animagine-xl"
# model = "WarriorMama777/OrangeMixs"
# # model = "cagliostrolab/animagine-xl-3.1"

# pipe = StableDiffusionXLPipeline.from_pretrained(
#     model,
#     torch_dtype=torch.float16,
#     # use_safetensors=True,
#     # variant="fp16",
#     # variant="None",
#     )
# pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
# pipe.to('cuda')


# from diffusers import DiffusionPipeline

# pipeline = DiffusionPipeline.from_pretrained("Vsukiyaki/ShiratakiMix")



model = "jianghushinian/animefull-final-pruned"



pipe = StableDiffusionXLPipeline.from_pretrained(
    model,
    torch_dtype=torch.float16,
    # use_safetensors=True,
    # variant="fp16",
    )

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


prompt = "face focus, cute, masterpiece, best quality, 1girl, green hair, sweater, looking at viewer, upper body, beanie, outdoors, night, turtleneck"
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"
output = "content/anime_girl.png"


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

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