Update app.py

app.py

@@ -1,23 +1,30 @@
-import streamlit as st
+import gradio as gr
 import torch
 from diffusers import DiffusionPipeline
 
-@st.cache_resource
-def load_pipe():
-    pipe = DiffusionPipeline.from_pretrained(
-        "stabilityai/stable-diffusion-xl-base-1.0",
-        torch_dtype=torch.float16,
-        variant="fp16",
-        use_safetensors=True
-    ).to("cuda")
-    pipe.load_lora_weights("ZB-Tech/Text-to-Image")
-    return pipe
+# Load model once
+pipe = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.float16,
+    variant="fp16",
+    use_safetensors=True
+).to("cuda")
 
-st.title("🎨 AI Image Generator")
-prompt = st.text_input("Enter a prompt:", "Draw a picture of some coding on a screen with a horror background")
+# Load the LoRA weights
+pipe.load_lora_weights("ZB-Tech/Text-to-Image")
 
-if st.button("Generate"):
-    with st.spinner("Generating image..."):
-        pipe = load_pipe()
-        image = pipe(prompt).images[0]
-        st.image(image, caption="Generated Image", use_column_width=True)
+# Define generation function
+def generate_image(prompt):
+    image = pipe(prompt).images[0]
+    return image
+
+# Gradio UI
+demo = gr.Interface(
+    fn=generate_image,
+    inputs=gr.Textbox(label="Enter a prompt", placeholder="e.g. Draw a picture of some coding on a screen with a horror background"),
+    outputs=gr.Image(type="pil", label="Generated Image"),
+    title="🎨 AI Image Generator",
+    description="This demo uses Stable Diffusion XL with custom LoRA weights to generate images from your prompts."
+)
+
+demo.launch()