Update app.py

app.py

@@ -1,47 +1,76 @@
+# app.py
+import os
 import torch
+import gradio as gr
 from diffusers import StableDiffusionPipeline
-from PIL import Image
-
-# ----------------------------
-# Settings
-# ----------------------------
-MODEL_ID = "stabilityai/stable-diffusion-1.5-turbo"  # Turbo model for CPU speed
-DEVICE = "cpu"  # CPU only
-DTYPE = torch.float32  # CPU-friendly
-
-PROMPT = "A cute cat wearing sunglasses, digital art"
-WIDTH = 768
-HEIGHT = 768
-NUM_STEPS = 1       # Fastest possible
-GUIDANCE_SCALE = 0  # Disable guidance for speed
-
-# ----------------------------
-# Load pipeline
-# ----------------------------
-pipe = StableDiffusionPipeline.from_pretrained(
-    MODEL_ID,
-    torch_dtype=DTYPE,
-    safety_checker=None,    # Optional, disable for speed
-    feature_extractor=None  # Optional, disable for speed
-)
-pipe.to(DEVICE)
-
-# ----------------------------
-# Generate image
-# ----------------------------
-generator = torch.Generator(device=DEVICE).manual_seed(42)  # Optional seed
-
-image: Image.Image = pipe(
-    prompt=PROMPT,
-    height=HEIGHT,
-    width=WIDTH,
-    num_inference_steps=NUM_STEPS,
-    guidance_scale=GUIDANCE_SCALE,
-    generator=generator
-).images[0]
-
-# ----------------------------
-# Save or show
-# ----------------------------
-image.save("output.png")
-image.show()
+
+# -----------------------
+# Config
+# -----------------------
+# Use a valid public turbo model
+MODEL_ID = "stabilityai/sd-turbo"   # <- change to "stabilityai/sdxl-turbo" if you prefer SDXL-Turbo
+
+# Optional: supply HF token in Spaces / env: HF_TOKEN=hf_xxx
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
+DEVICE = "cpu"
+DTYPE = torch.float32
+
+# Limit threads to avoid oversubscription on Spaces
+os.environ["OMP_NUM_THREADS"] = "4"
+os.environ["MKL_NUM_THREADS"] = "4"
+torch.set_num_threads(4)
+
+# -----------------------
+# Load pipeline (robust)
+# -----------------------
+download_kwargs = {}
+if HF_TOKEN:
+    download_kwargs["token"] = HF_TOKEN
+
+try:
+    pipe = StableDiffusionPipeline.from_pretrained(
+        MODEL_ID,
+        torch_dtype=DTYPE,
+        safety_checker=None,   # optional: speed + avoid HF safety issues on CPU
+        **download_kwargs
+    )
+except Exception as e:
+    raise RuntimeError(
+        f"Failed to download/load {MODEL_ID}. "
+        "If this is a gated model, add a valid HF_TOKEN env var. "
+        f"Underlying error: {e}"
+    )
+
+pipe = pipe.to(DEVICE)
+pipe.enable_attention_slicing()
+
+# -----------------------
+# Inference (1-step turbo)
+# -----------------------
+def generate(prompt: str, seed: int = 42):
+    generator = torch.Generator(device=DEVICE).manual_seed(seed) if seed else None
+
+    image = pipe(
+        prompt=prompt,
+        height=768,
+        width=768,
+        num_inference_steps=1,     # one step
+        guidance_scale=0.0,        # disable classifier-free guidance for turbo models
+        generator=generator,
+    ).images[0]
+    return image
+
+# -----------------------
+# Gradio UI
+# -----------------------
+with gr.Blocks() as demo:
+    gr.Markdown("## Turbo (1-step) Image Generation — CPU")
+    prompt = gr.Textbox(placeholder="A cute cat wearing sunglasses", lines=2)
+    seed = gr.Number(value=42, label="Seed (0 = random)")
+    btn = gr.Button("Generate")
+    out = gr.Image()
+    btn.click(generate, [prompt, seed], out)
+
+if __name__ == "__main__":
+    demo.launch()