init

app.py

@@ -1,13 +1,51 @@
+from diffusers import StableDiffusionXLPipeline
+from fastapi.responses import FileResponse
+from pydantic import BaseModel
+from PIL import Image
 import gradio as gr
-import spaces
 import torch
+import uuid
+import io
+import os
+import spaces  # Assuming spaces is a valid module with GPU decorator
 
-zero = torch.Tensor([0]).cuda()
-print(zero.device) # <-- 'cpu' 🤔
+# Load your model
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    "segmind/SSD-1B", 
+    torch_dtype=torch.float16, 
+    use_safetensors=True, 
+    variant="fp16"
+)
+pipe.to("cuda:0")
 
-@spaces.GPU
-def greet(n):
-    print(zero.device) # <-- 'cuda:0' 🤗
-    return f"Hello {zero + n} Tensor"
+@spaces.GPU  # Apply the GPU decorator
+def generate_and_save_image(prompt, negative_prompt=''):
+    # Generate image using the provided prompts
+    image = pipe(prompt=prompt, negative_prompt=negative_prompt).images[0]
 
-gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text()).launch()
+    # Generate a unique UUID for the filename
+    unique_id = str(uuid.uuid4())
+    image_path = f"generated_images/{unique_id}.png"
+
+    # Save generated image locally
+    os.makedirs('generated_images', exist_ok=True)
+    image.save(image_path, format='PNG')
+
+    # Return the path of the saved image to display in Gradio interface
+    return image_path
+
+# Define the Gradio interface
+iface = gr.Interface(
+    fn=generate_and_save_image,
+    inputs=[
+        gr.Textbox(label="Enter prompt", placeholder="Type your prompt here..."),
+        gr.Textbox(label="Enter negative prompt (optional)", placeholder="Type your negative prompt here...")
+    ],
+    outputs=gr.Image(type="file", label="Generated Image"),
+    title="Image Generation with Stable Diffusion XL",
+    description="Enter a prompt and (optionally) a negative prompt to generate an image.",
+    layout="vertical"
+)
+
+# Launch the Gradio app
+iface.launch()