import gradio as gr
import torch
from torch import autocast
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler

model = "linaqruf/animagine-xl"

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


def launch(prompt, negative_prompt):
    prompt += " ,awesome, pixel art"
    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"
    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)
    return image.images[0]  # Assuming this is how you get the resulting image

iface = gr.Interface(fn=launch,
                     inputs=[gr.Textbox(label="Prompt"), gr.Textbox(label="Negative Prompt")],
                     outputs=gr.Image(type='pil'),
                     title="Generate Images",
                     description="Enter a prompt and a negative prompt to generate an image.")
iface.launch()