Update app.py

app.py

@@ -1,3 +1,44 @@
+from diffusers import StableDiffusionPipeline, DiffusionPipeline
+import torch
 import gradio as gr
+import spaces
 
-gr.Interface.load("models/Onodofthenorth/SD_PixelArt_SpriteSheet_Generator").launch()
+css = """
+#img-display-container {
+    max-height: 50vh;
+    }
+#img-display-input {
+    max-height: 40vh;
+    }
+#img-display-output {
+    max-height: 40vh;
+    }
+"""
+
+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()