Update app.py

app.py

@@ -2,34 +2,53 @@ import gradio as gr
 from diffusers import StableDiffusionPipeline
 import torch
 
+# Detect device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if device == "cuda" else torch.float32
+
 # Load base model
 pipe = StableDiffusionPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch.float16
-).to("cuda")
+    torch_dtype=dtype
+)
 
-# Load LoRA weights from the 'lora' folder
-#pipe.load_attn_procs("./lora")
+# Load LoRA weights
 pipe.load_lora_weights("./lora")
 
-# Move to GPU if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Move to device
 pipe = pipe.to(device)
 
-# Define image generation function
-def generate(prompt):
-    image = pipe(prompt).images[0]
+# Extra CPU optimizations
+if device == "cpu":
+    pipe.enable_attention_slicing()
+    pipe.enable_sequential_cpu_offload()
+    pipe.enable_vae_tiling()
+
+# Define image generation function with quality toggle
+def generate(prompt, quality):
+    if quality == "Fast":
+        steps = 20
+        guidance_scale = 7.0
+    else:  # High Quality
+        steps = 40
+        guidance_scale = 8.5
+    
+    with torch.no_grad():
+        image = pipe(prompt, num_inference_steps=steps, guidance_scale=guidance_scale).images[0]
     return image
 
 # Build Gradio UI
 demo = gr.Interface(
     fn=generate,
-    inputs=gr.Textbox(label="Enter your prompt"),
+    inputs=[
+        gr.Textbox(label="Enter your prompt"),
+        gr.Dropdown(["Fast", "High Quality"], value="Fast", label="Generation Mode")
+    ],
     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."
+    title="Fine-tuned Stable Diffusion with LoRA (CPU-Optimized)",
+    description="Choose 'Fast' for quicker generation or 'High Quality' for better details."
 )
 
 # Launch app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()