Update app.py

app.py

@@ -1,17 +1,30 @@
 import torch
-from diffusers import FluxPipeline
+from diffusers import StableDiffusion3Pipeline
+import gradio as gr
 
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
-pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
+# Load model (auto CPU/GPU)
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-prompt = "A cat holding a sign that says hello world"
-image = pipe(
-    prompt,
-    height=1024,
-    width=1024,
-    guidance_scale=3.5,
-    num_inference_steps=50,
-    max_sequence_length=512,
-    generator=torch.Generator("cpu").manual_seed(0)
-).images[0]
-image.save("flux-dev.png")
+pipe = StableDiffusion3Pipeline.from_pretrained(
+    "stabilityai/stable-diffusion-3-medium-diffusers",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+)
+pipe.to(device)
+
+def generate(prompt):
+    image = pipe(
+        prompt,
+        negative_prompt="",
+        num_inference_steps=28,
+        guidance_scale=7.0,
+    ).images[0]
+    return image
+
+ui = gr.Interface(
+    fn=generate,
+    inputs=gr.Textbox(label="Enter your prompt"),
+    outputs=gr.Image(label="Generated Image"),
+    title="Stable Diffusion 3 Medium - HF Space",
+)
+
+ui.launch()