Update app.py

app.py

@@ -1,66 +1,67 @@
-# app.py - MangaMorph (Gradio) CPU-friendly template
+# app.py — MangaMorph (Gradio) — colorful & polished CPU-friendly UI
 import os
 import random
+import time
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageOps
 import gradio as gr
 import torch
-from diffusers import DiffusionPipeline
-from diffusers import EulerDiscreteScheduler  # scheduler choice
+from diffusers import DiffusionPipeline, EulerDiscreteScheduler
 
-# ---- CONFIG ----
-# Change this model id if you prefer another (see note about license/access above)
-MODEL_ID = os.getenv("MODEL_ID", "hakurei/waifu-diffusion")  
-
-# If your model requires a token, set HUGGINGFACE_HUB_TOKEN in Space secrets
+# ---------- CONFIG ----------
+MODEL_ID = os.getenv("MODEL_ID", "hakurei/waifu-diffusion")  # change if needed
 HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN", None)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float32 if device == "cpu" else torch.float16
+torch_dtype = torch.float16 if device == "cuda" else torch.float32
 
-# Limits / defaults for CPU-friendly runs
-DEFAULT_WIDTH = 512
-DEFAULT_HEIGHT = 512
-DEFAULT_STEPS = 20
-DEFAULT_GUIDANCE = 7.5
+# CPU-friendly defaults & limits
+DEFAULT_WIDTH = 384
+DEFAULT_HEIGHT = 384
+DEFAULT_STEPS = 10
+DEFAULT_GUIDANCE = 5.5
 MAX_SEED = np.iinfo(np.int32).max
 
-# Load pipeline (wrapped in try/except so error messages are shown in app log)
+# ---------- Load pipeline (lazy) ----------
+PIPE = None
 def load_pipeline():
+    global PIPE
+    if PIPE is not None:
+        return PIPE
     try:
-        scheduler = EulerDiscreteScheduler.from_pretrained(MODEL_ID, subfolder="scheduler") if os.getenv("USE_EULER", "1") == "1" else None
-        pipe = DiffusionPipeline.from_pretrained(
-            MODEL_ID,
-            torch_dtype=torch_dtype,
-            use_auth_token=HF_TOKEN,
-        )
-        # attach scheduler only if available and desired
-        if isinstance(pipe.scheduler, type(None)) and scheduler is not None:
+        # Try to load model; if scheduler not present we keep default
+        pipe = DiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch_dtype, use_auth_token=HF_TOKEN)
+        # Try to set a faster scheduler if available
+        try:
+            scheduler = EulerDiscreteScheduler.from_pretrained(MODEL_ID, subfolder="scheduler")
             pipe.scheduler = scheduler
+        except Exception:
+            pass
+        # Move to device
         pipe = pipe.to(device)
-        # For CPU: disable safety checker to avoid long CPU runs (optional)
+        # Disable safety checker on CPU for speed (optional)
         try:
             pipe.safety_checker = None
         except Exception:
             pass
-        return pipe
+        PIPE = pipe
+        return PIPE
     except Exception as e:
-        raise RuntimeError(f"Failed to load model '{MODEL_ID}': {e}")
+        raise RuntimeError(f"Model load failed: {e}")
 
-# lazy load
-PIPE = None
-def get_pipe():
-    global PIPE
-    if PIPE is None:
-        PIPE = load_pipeline()
-    return PIPE
-
-# Default negative prompt tuned to reduce common artifacts
-DEFAULT_NEGATIVE_PROMPT = (
-    "low quality, bad anatomy, blurry, deformed, extra limbs, mutated hands, "
-    "poorly drawn face, watermark, text, signature, lowres, oversaturated"
+# ---------- Helpers ----------
+DEFAULT_NEG = (
+    "low quality, bad anatomy, blurry, extra limbs, malformed, deformed, watermark, text, signature, lowres"
 )
 
+def tidy_image(img: Image.Image, max_side=1024):
+    # ensure RGB and a consistent max size (safety)
+    img = img.convert("RGB")
+    if max(img.size) > max_side:
+        img = ImageOps.contain(img, (max_side, max_side))
+    return img
+
+# ---------- Inference function ----------
 def infer(
     prompt: str,
     negative_prompt: str,
@@ -71,100 +72,130 @@ def infer(
     guidance_scale: float,
     num_inference_steps: int,
 ):
-    if not prompt:
-        return None, "Please enter a prompt."
+    start = time.time()
+    if not prompt or prompt.strip() == "":
+        return None, "Enter a prompt first."
 
-    if randomize_seed:
+    if randomize_seed or int(seed) == 0:
         seed = random.randint(0, MAX_SEED)
+    else:
+        seed = int(seed) % MAX_SEED
 
-    gen = torch.Generator(device=device)
-    gen = gen.manual_seed(seed)
+    # load model (may download on first run)
+    try:
+        pipe = load_pipeline()
+    except Exception as e:
+        return None, f"Model load error: {e}"
 
-    pipe = get_pipe()
+    # Cap sizes for CPU safety
+    width = int(min(max(256, width), 512))
+    height = int(min(max(256, height), 512))
+    steps = int(min(max(4, num_inference_steps), 20))
 
-    # Cap size to avoid OOM on CPU
-    width = min(width, 768)
-    height = min(height, 768)
+    gen = torch.Generator(device=device).manual_seed(seed)
 
     try:
-        output = pipe(
+        out = pipe(
             prompt=prompt,
-            negative_prompt=(negative_prompt or DEFAULT_NEGATIVE_PROMPT),
+            negative_prompt=(negative_prompt or DEFAULT_NEG),
             width=width,
             height=height,
             guidance_scale=float(guidance_scale),
-            num_inference_steps=int(num_inference_steps),
+            num_inference_steps=steps,
             generator=gen,
         )
-        image = output.images[0]
-        # simple postprocessing: convert to RGB and return
-        if isinstance(image, Image.Image):
-            image = image.convert("RGB")
-        return image, f"Seed: {seed}"
+        image = out.images[0]
+        image = tidy_image(image, max_side=1024)
+        elapsed = time.time() - start
+        return image, f"✅ Done — Seed: {seed} • {int(elapsed)}s"
     except Exception as e:
-        # retry logic: try again with smaller steps/guidance if CPU fails
+        # try a lighter retry
         try:
-            output = pipe(
+            out = pipe(
                 prompt=prompt,
-                negative_prompt=(negative_prompt or DEFAULT_NEGATIVE_PROMPT),
+                negative_prompt=(negative_prompt or DEFAULT_NEG),
                 width=width,
                 height=height,
                 guidance_scale=max(3.0, float(guidance_scale) - 1.0),
-                num_inference_steps=max(5, int(num_inference_steps) - 5),
+                num_inference_steps=max(4, steps - 4),
                 generator=gen,
             )
-            image = output.images[0]
-            if isinstance(image, Image.Image):
-                image = image.convert("RGB")
-            return image, f"Recovered (retry) — Seed: {seed}"
+            image = tidy_image(out.images[0], max_side=1024)
+            elapsed = time.time() - start
+            return image, f"⚠ Recovered (retry) — Seed: {seed} • {int(elapsed)}s"
         except Exception as e2:
             return None, f"Generation failed: {e2}"
 
-# ---- UI ----
+# ---------- UI (Gradio Blocks) ----------
 css = """
-#main { max-width: 880px; margin: auto; }
-.header { text-align: center; }
-.small { font-size: 0.9rem; color: #666; }
+/* Gradient page background */
+body { background: linear-gradient(120deg,#f6f0ff 0%, #fff9f0 100%); }
+/* Card styling */
+.header { text-align: left; padding: 12px 18px; border-radius: 12px; background: linear-gradient(90deg,#ffd6e0,#ffe8a1); box-shadow: 0 6px 20px rgba(0,0,0,0.06); }
+.brand { font-weight: 800; font-size: 20px; letter-spacing: 0.2px; color: #5b1e72; }
+.subtitle { color:#333333; font-size:13px; margin-top:4px; }
+.controls { background: white; padding: 12px; border-radius: 10px; box-shadow: 0 4px 18px rgba(0,0,0,0.04); }
+.small { font-size:12px; color:#666; }
+.btn-primary { background: linear-gradient(90deg,#ff7ab6,#ffb86b); color: white; font-weight:700; }
+.footer { font-size:12px; color:#666; text-align:center; margin-top:8px; }
 """
 
 examples = [
-    "A young anime girl standing in a rain-soaked neon street, detailed lighting, cinematic",
-    "A samurai in traditional armor on a cliff at sunset, dramatic lighting, anime style",
-    "Cozy room with anime character by window reading, soft warm light"
+    "anime girl standing on a cherry-blossom bridge at sunset, cinematic lighting, detailed eyes",
+    "young samurai on a misty mountain path, dramatic clouds, anime style",
+    "cozy studio apartment with anime character reading by window, warm lighting"
 ]
 
-with gr.Blocks(css=css, theme=gr.themes.Default()) as demo:
-    with gr.Column(elem_id="main"):
-        gr.Markdown("<div class='header'><h2>MangaMorph — Anime Scene Generator</h2>"
-                    "<div class='small'>Text → Anime image | CPU-optimized | Use Model ID or set HF token in Secrets</div></div>")
-
-        with gr.Row():
-            prompt = gr.Textbox(label="Prompt", placeholder="Describe the anime scene you want...", lines=2)
-            run_btn = gr.Button("Generate", variant="primary")
-
-        with gr.Row():
-            gallery = gr.Image(label="Result")
-
-        with gr.Accordion("Advanced settings", open=False):
-            negative = gr.Textbox(label="Negative prompt (optional)", placeholder=DEFAULT_NEGATIVE_PROMPT, lines=2, value=DEFAULT_NEGATIVE_PROMPT)
-            seed = gr.Number(label="Seed (0 = randomize)", value=0)
-            randomize = gr.Checkbox(label="Randomize seed", value=True)
-            with gr.Row():
-                width = gr.Slider(label="Width", minimum=256, maximum=768, step=64, value=DEFAULT_WIDTH)
-                height = gr.Slider(label="Height", minimum=256, maximum=768, step=64, value=DEFAULT_HEIGHT)
-            with gr.Row():
-                guidance = gr.Slider(label="Guidance scale", minimum=1.0, maximum=15.0, step=0.1, value=DEFAULT_GUIDANCE)
-                steps = gr.Slider(label="Steps", minimum=5, maximum=50, step=1, value=DEFAULT_STEPS)
-
-        gr.Examples(examples=examples, inputs=[prompt])
-
-        status = gr.Textbox(label="Status / Seed", interactive=False)
+with gr.Blocks(css=css, title="MangaMorph — Anime Scene Generator") as demo:
+    with gr.Row():
+        with gr.Column(scale=2):
+            gr.HTML("<div class='header'><div class='brand'>MangaMorph</div>"
+                    "<div class='subtitle'>Text → Anime image • CPU-optimized • Try 384×384 & 10 steps for speed</div></div>")
+            with gr.Box(elem_id="controls", visible=True):
+                prompt = gr.Textbox(label="Describe your anime scene", placeholder="e.g. A cyberpunk anime girl on a rainy street, neon lights...", lines=3)
+                with gr.Row():
+                    run_btn = gr.Button("Generate", elem_id="run", variant="primary")
+                    download_btn = gr.Button("Download", elem_id="dl", variant="secondary")
+                with gr.Accordion("Advanced settings", open=False):
+                    negative = gr.Textbox(label="Negative prompt (optional)", placeholder="e.g. blurry, deformed, watermark", value=DEFAULT_NEG, lines=2)
+                    with gr.Row():
+                        seed = gr.Number(label="Seed (0 = random)", value=0)
+                        randomize = gr.Checkbox(label="Randomize seed", value=True)
+                    with gr.Row():
+                        width = gr.Slider(label="Width", minimum=256, maximum=512, step=64, value=DEFAULT_WIDTH)
+                        height = gr.Slider(label="Height", minimum=256, maximum=512, step=64, value=DEFAULT_HEIGHT)
+                    with gr.Row():
+                        guidance = gr.Slider(label="Guidance scale", minimum=1.0, maximum=12.0, step=0.1, value=DEFAULT_GUIDANCE)
+                        steps = gr.Slider(label="Steps", minimum=4, maximum=20, step=1, value=DEFAULT_STEPS)
+                gr.Examples(examples=examples, inputs=[prompt], label="Try examples")
+                status = gr.Textbox(label="Status", value="Ready", interactive=False)
+        with gr.Column(scale=1):
+            gr.HTML("<div style='padding:8px;text-align:center;'><b>Preview</b></div>")
+            result = gr.Image(label="Generated image", shape=(384,384))
+            gallery = gr.Gallery(label="History (latest first)", columns=1).style(height="auto")
+            gr.HTML("<div class='footer'>Tip: Use lower resolution & fewer steps for much faster results on CPU</div>")
+
+    # Click behavior
+    def generate_and_update(*args):
+        img, msg = infer(*args)
+        # manage gallery: return [img] to gallery; status msg to status
+        return img, msg, [img] if img is not None else [], img
 
     run_btn.click(
-        fn=infer,
+        fn=generate_and_update,
         inputs=[prompt, negative, seed, randomize, width, height, guidance, steps],
-        outputs=[gallery, status],
+        outputs=[result, status, gallery, result],
+        show_progress=True,
     )
 
+    # Download button behaviour: downloads currently previewed image
+    def download_current(img):
+        if img is None:
+            return gr.update(value=None)
+        # return PIL image to trigger download
+        return img
+
+    download_btn.click(fn=download_current, inputs=[result], outputs=[result])
+
 if __name__ == "__main__":
     demo.launch()