Update app.py

app.py

@@ -1,34 +1,37 @@
-import gradio as gr
-from diffusers import StableDiffusionPipeline
-import torch
-
-# Load base model
-pipe = StableDiffusionPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch.float16
-)
-
-# Load LoRA weights from the 'lora' folder
-pipe.load_attn_procs("./lora")
-
-# Move to GPU if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = pipe.to(device)
-
-# Define image generation function
-def generate(prompt):
-    image = pipe(prompt).images[0]
-    return image
-
-# Build Gradio UI
-demo = gr.Interface(
-    fn=generate,
-    inputs=gr.Textbox(label="Enter your prompt"),
-    outputs=gr.Image(label="Generated Image"),
-    title="Fine-tuned Stable Diffusion with LoRA",
-    description="Enter a prompt and get an AI-generated image based on your fine-tuned model."
-)
-
-# Launch app
-if __name__ == "__main__":
-    demo.launch()
+!pip install --upgrade diffusers accelerate transformers safetensors
+
+
+import gradio as gr
+from diffusers import StableDiffusionPipeline
+import torch
+
+# Load base model
+pipe = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16
+).to("cuda")
+
+# Load LoRA weights from the 'lora' folder
+pipe.load_attn_procs("./lora")
+
+# Move to GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = pipe.to(device)
+
+# Define image generation function
+def generate(prompt):
+    image = pipe(prompt).images[0]
+    return image
+
+# Build Gradio UI
+demo = gr.Interface(
+    fn=generate,
+    inputs=gr.Textbox(label="Enter your prompt"),
+    outputs=gr.Image(label="Generated Image"),
+    title="Fine-tuned Stable Diffusion with LoRA",
+    description="Enter a prompt and get an AI-generated image based on your fine-tuned model."
+)
+
+# Launch app
+if __name__ == "__main__":
+    demo.launch()