Update app.py

app.py

@@ -4,55 +4,64 @@ from PIL import Image, ImageDraw, ImageFont
 import gradio as gr
 from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 
-# --------------------
-# Config
-# --------------------
-MODEL_ID = os.getenv("MODEL_ID", "runwayml/stable-diffusion-v1-5")
+# =========================
+# SPEED PRESET
+# =========================
+# Use SD Turbo (1.5) – optimized for very few steps on CPU
+DEFAULT_MODEL_ID = "stabilityai/sd-turbo"
+MODEL_ID = os.getenv("MODEL_ID", DEFAULT_MODEL_ID)
+
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# Simple prompt guardrail (blocks obvious NSFW attempts)
+# Short NSFW guardrail (block, no blur)
 NSFW_TERMS = [
     r"\bnsfw\b", r"\bnude\b", r"\bnudity\b", r"\bsex\b", r"\bexplicit\b", r"\bporn\b",
-    r"\bboobs\b", r"\bbutt\b", r"\bass\b", r"\bnsfw\b", r"\bnaked\b", r"\btits\b",
+    r"\bboobs\b", r"\bbutt\b", r"\bass\b", r"\bnaked\b", r"\btits\b",
     r"\b18\+\b", r"\berotic\b", r"\bfetish\b"
 ]
 NSFW_REGEX = re.compile("|".join(NSFW_TERMS), flags=re.IGNORECASE)
 
-# --------------------
-# Load pipeline
-# --------------------
+def _blocked_tile(reason: str, w=384, h=384) -> Image.Image:
+    img = Image.new("RGB", (w, h), (18, 20, 26))
+    d = ImageDraw.Draw(img)
+    text = f"BLOCKED\n{reason}"
+    try:
+        font = ImageFont.truetype("DejaVuSans-Bold.ttf", 26)
+    except:
+        font = ImageFont.load_default()
+    box = d.multiline_textbbox((0,0), text, font=font, align="center")
+    tw, th = box[2]-box[0], box[3]-box[1]
+    d.multiline_text(((w-tw)//2, (h-th)//2), text, font=font, fill=(255,255,255), align="center")
+    return img
+
+def _is_nsfw(s: str) -> bool:
+    return bool(NSFW_REGEX.search(s or ""))
+
+# -------------------------
+# Load pipeline (fast path)
+# -------------------------
+torch.set_grad_enabled(False)
+
 pipe = StableDiffusionPipeline.from_pretrained(
     MODEL_ID,
-    torch_dtype=DTYPE
+    torch_dtype=DTYPE,
+    safety_checker=None  # let model config handle; we block explicitly on prompts
 )
+
+# Turbo still benefits from DPMSolver for CPU
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
 
-if torch.cuda.is_available():
+if DEVICE == "cuda":
     pipe = pipe.to("cuda")
     pipe.enable_attention_slicing()
     pipe.enable_vae_slicing()
 else:
     pipe = pipe.to("cpu")
 
-# --------------------
-# Helpers
-# --------------------
-def blocked_tile(reason: str, width=512, height=512) -> Image.Image:
-    img = Image.new("RGB", (width, height), (20, 20, 24))
-    draw = ImageDraw.Draw(img)
-    text = f"BLOCKED\n{reason}"
-    try:
-        font = ImageFont.truetype("DejaVuSans-Bold.ttf", 28)
-    except:
-        font = ImageFont.load_default()
-    tw, th = draw.multiline_textbbox((0,0), text, font=font)[2:]
-    draw.multiline_text(((width - tw)//2, (height - th)//2), text, fill=(255,255,255), font=font, align="center")
-    return img
-
-def is_prompt_nsfw(prompt: str) -> bool:
-    return bool(NSFW_REGEX.search(prompt or ""))
-
+# -------------------------
+# Generate fn (kept lean)
+# -------------------------
 def generate(
     prompt: str,
     negative_prompt: str,
@@ -64,12 +73,16 @@ def generate(
     batch_size: int
 ) -> Tuple[List[Image.Image], str]:
     if not prompt.strip():
-        return [], "Add a prompt to get rolling."
+        return [], "Add a prompt first."
+
+    # block obvious NSFW prompts
+    if _is_nsfw(prompt) or _is_nsfw(negative_prompt or ""):
+        return [_blocked_tile("NSFW prompt detected", width, height)], "Blocked: NSFW prompt."
 
-    # Hard block obvious NSFW prompts before hitting the model
-    if is_prompt_nsfw(prompt) or is_prompt_nsfw(negative_prompt or ""):
-        img = blocked_tile("NSFW prompt detected")
-        return [img], "Blocked: NSFW prompt."
+    # SD-Turbo is designed for tiny step counts + low/zero CFG
+    # guard rails on parameters
+    steps = max(1, min(int(steps), 12))
+    guidance = max(0.0, min(float(guidance), 2.0))
 
     # Seed
     if seed < 0:
@@ -78,7 +91,7 @@ def generate(
 
     out = pipe(
         prompt=prompt,
-        negative_prompt=negative_prompt or None,
+        negative_prompt=(negative_prompt or None),
         num_inference_steps=steps,
         guidance_scale=guidance,
         width=width,
@@ -87,29 +100,20 @@ def generate(
         generator=generator
     )
 
-    images = out.images
-    flags = getattr(out, "nsfw_content_detected", None)
-
-    # If the underlying safety checker flags NSFW, block it (no blur)
-    if flags:
-        for i, flagged in enumerate(flags):
-            if flagged:
-                images[i] = blocked_tile("NSFW content flagged")
-
-    msg = f"Seed: {seed} • Images: {len(images)}"
-    if flags is not None:
-        msg += f" • Flagged: {sum(1 for f in flags if f)}"
-    return images, msg
+    imgs = out.images
+    # Some sd-turbo configs may not return nsfw flags; we already block on prompt
+    msg = f"Model: {MODEL_ID} • Seed: {seed} • Steps: {steps} • CFG: {guidance} • {width}x{height} • Batch: {batch_size}"
+    return imgs, msg
 
-# --------------------
-# UI
-# --------------------
-with gr.Blocks(title="VibeForge — Clean Image Generator") as demo:
+# -------------------------
+# UI (defaults tuned for CPU)
+# -------------------------
+with gr.Blocks(title="VibeForge — Fast (CPU-friendly) Image Gen") as demo:
     gr.Markdown(
         """
-# VibeForge ⚒️  
-**Clean, creative image generation.**  
-NSFW inputs are blocked. Keep it classy and go wild on style, lighting, composition, mood.
+# VibeForge ⚒️
+**Fast, clean image generation (CPU-friendly).**  
+Uses **SD-Turbo** tuned for low steps. NSFW inputs are blocked.
         """
     )
 
@@ -117,23 +121,23 @@ NSFW inputs are blocked. Keep it classy and go wild on style, lighting, composit
         with gr.Column(scale=3):
             prompt = gr.Textbox(
                 label="Prompt",
-                placeholder="a cinematic photo of a vintage motorcycle by the ocean at sunset, golden hour, soft rim light, 50mm"
+                placeholder="a neon-lit lighthouse on a stormy cliff at night, cinematic, volumetric fog, high contrast"
             )
-            negative = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, watermark, jpeg artifacts")
+            negative = gr.Textbox(label="Negative Prompt", placeholder="low quality, watermark, overexposed")
             with gr.Row():
-                steps = gr.Slider(10, 50, value=28, step=1, label="Steps")
-                guidance = gr.Slider(1.0, 12.0, value=7.5, step=0.1, label="CFG")
+                steps = gr.Slider(1, 12, value=4, step=1, label="Steps (SD-Turbo sweet spot: 2-6)")
+                guidance = gr.Slider(0.0, 2.0, value=0.5, step=0.1, label="CFG (SD-Turbo likes low)")
             with gr.Row():
-                width = gr.Dropdown(choices=[384, 448, 512, 640, 768], value=512, label="Width")
-                height = gr.Dropdown(choices=[384, 448, 512, 640, 768], value=512, label="Height")
+                width = gr.Dropdown(choices=[384, 448, 512], value=384, label="Width")
+                height = gr.Dropdown(choices=[384, 448, 512], value=384, label="Height")
             with gr.Row():
                 seed = gr.Number(value=-1, label="Seed (-1 = random)", precision=0)
-                batch = gr.Slider(1, 4, value=1, step=1, label="Batch")
+                batch = gr.Slider(1, 2, value=1, step=1, label="Batch (keep small on CPU)")
 
             go = gr.Button("Generate", variant="primary")
 
         with gr.Column(scale=5):
-            gallery = gr.Gallery(label="Output", columns=2, height=512)
+            gallery = gr.Gallery(label="Output", columns=2, height=448)
             info = gr.Markdown()
 
     go.click(
@@ -143,4 +147,4 @@ NSFW inputs are blocked. Keep it classy and go wild on style, lighting, composit
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()