Update app.py

app.py

@@ -3,37 +3,41 @@ import torch
 from diffusers import StableDiffusion3Pipeline
 
 def image_generation(prompt):
-    try:
-        device = "cpu"
-        pipeline = StableDiffusion3Pipeline.from_pretrained(
-            "stabilityai/stable-diffusion-3-medium-diffusers",
-            torch_dtype=torch.float32,  # ✅ Use float32 on CPU
-            text_encoder_3=None,
-            tokenizer_3=None
-        )
-        pipeline.to(device)
-
-        image = pipeline(
-            prompt=prompt,
-            negative_prompt="blurred, ugly, watermark, low resolution, blurry",
-            num_inference_steps=20,       # ✅ Reduce steps for performance
-            height=1024,                   # ✅ Smaller size to avoid memory error
-            width=1024,
-            guidance_scale=7.5,
-        ).images[0]
-
-        return image
-
-    except Exception as e:
-        print(f"[ERROR] {e}")
-        return f"Error: {e}"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
 
+    # Load the pipeline (with resume_download if interrupted previously)
+    pipeline = StableDiffusion3Pipeline.from_pretrained(
+        "stabilityai/stable-diffusion-3-medium-diffusers",
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+        text_encoder_3=None,
+        tokenizer_3=None,
+        resume_download=True
+    )
+
+    # ✅ Only use this line if you have GPU + Accelerate
+    # pipeline.enable_model_cpu_offload()
+
+    # ✅ Instead, move pipeline to CPU or CUDA manually
+    pipeline.to(device)
+
+    image = pipeline(
+        prompt=prompt,
+        negative_prompt="blurred,ugly,watermark, low resolution, blurry",
+        num_inference_steps=50,
+        height=1024,
+        width=1024,
+        guidance_scale=9.0,
+    ).images[0]
+
+    return image  # ✅ Return the image for Gradio
+
+# Gradio UI
 interface = gr.Interface(
     fn=image_generation,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your prompt..."),
-    outputs="image",
+    inputs=gr.Textbox(lines=2, placeholder="Enter your Prompt..."),
+    outputs=gr.Image(type="pil"),
     title="AI Image Generator By Arnav Anand",
-    description="Generate images using Stable Diffusion 3 (CPU-Only Version)"
+    description="This application will be used to generate awesome images using SD3 model"
 )
 
 interface.launch()