from diffusers import StableDiffusionPipeline, DiffusionPipeline
import torch
import gradio as gr
import spaces

css = """
#img-display-container {
    max-height: 50vh;
    }
#img-display-input {
    max-height: 40vh;
    }
#img-display-output {
    max-height: 40vh;
    }
"""

import os
os.makedirs("cache", exist_ok=True)

DEVICE = 'cuda'
model_id = "Onodofthenorth/SD_PixelArt_SpriteSheet_Generator"
pipe = StableDiffusionPipeline.from_pretrained(model_id, cache_dir="cache", torch_dtype=torch.float16)
pipe.to("cuda")


@spaces.GPU(enable_queue=True)
def generate_sprite(prompt):
    # pipe = pipe.to(DEVICE)
    image = pipe(prompt).images[0]
    return image

title = "# SD_PixelArt_SpriteSheet_Generator"
description = """Pixel Art Sprite Sheet Generator with Stable Diffusion Checkpoint."""

with gr.Blocks(css=css) as API:
    gr.Markdown(title)
    gr.Markdown(description)
        
    with gr.Column():
        inputs=gr.TextArea(label="Prompt", placeholder="Prompt")
        outputs=gr.Image(label="Ouput Image", type='pil', height=500)
    generate_btn = gr.Button(value="Generate")
    generate_btn.click(generate_sprite, inputs=inputs, outputs=outputs, api_name="generate_mesh")

API.launch()