Update app.py

app.py

@@ -1,243 +1,135 @@
-import os
-import io
-import time
-import random
-import base64
-from urllib.parse import quote_plus
-from typing import Optional, Tuple
+import os, io, time, base64, random
+from typing import Optional
+from urllib.parse import quote
 
 import requests
+from PIL import Image
 import gradio as gr
 
-# -----------------------------
-# Config
-# -----------------------------
-# You can set your Modal endpoint via env var MM_I2V_URL
-DEFAULT_API_URL = os.getenv(
-    "MM_I2V_URL",
-    "https://moonmath-ai--moonmath-i2v-backend-moonmathinference-run.modal.run",
-)
-
-SAVE_DIR = "outputs"
-os.makedirs(SAVE_DIR, exist_ok=True)
-
-
-# -----------------------------
-# Helpers
-# -----------------------------
-def _save_bytes_to_mp4(buf: bytes, name_prefix: str) -> str:
-    ts = int(time.time() * 1000)
-    path = os.path.join(SAVE_DIR, f"{name_prefix}-{ts}.mp4")
+# -------- Modal inference endpoint (dev) --------
+INFERENCE_URL = "https://moonmath-ai-dev--moonmath-i2v-backend-moonmathinference-run.modal.run"
+
+# -------- small helpers --------
+def _save_video_bytes(data: bytes, tag: str) -> str:
+    os.makedirs("/mnt/data", exist_ok=True)
+    path = f"/mnt/data/{tag}_{int(time.time())}.mp4"
     with open(path, "wb") as f:
-        f.write(buf)
+        f.write(data)
     return path
 
+def _png_bytes(img: Image.Image) -> bytes:
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    return buf.getvalue()
 
-def _download(url: str) -> bytes:
-    r = requests.get(url, timeout=600)
+def _download_to_bytes(url: str) -> bytes:
+    r = requests.get(url, timeout=180)
     r.raise_for_status()
     return r.content
 
+def stitch_call(start_img: Image.Image, end_img: Image.Image, prompt: str, seed: Optional[int]) -> Optional[str]:
+    """
+    JS equivalent:
+      const fd = new FormData();
+      fd.append("image_bytes", start);
+      fd.append("image_bytes_end", end);
+      fetch(`${INFERENCE_URL}?prompt=${prompt}&seed=${seed}`, { method:"POST", body: fd })
+    """
+    if start_img is None or end_img is None:
+        return None
+
+    if seed in (None, 0, -1):
+        seed = random.randint(1, 2**31 - 1)
 
-def call_i2v(
-    image_path: str,
-    prompt: str,
-    seed: Optional[int],
-    api_url: Optional[str] = None,
-) -> Tuple[Optional[str], Optional[str]]:
-    \"\"\"
-    Call the image->video backend and return (video_path, error_message).
-
-    Tries to handle several common response types:
-      1) raw mp4 bytes
-      2) JSON with {\"video\": \"<base64>\"} (mp4 base64)
-      3) JSON with {\"video_url\": \"https://...\"} (or \"result_url\")
-    \"\"\"
-    api = (api_url or DEFAULT_API_URL).strip().rstrip(\"/\")
-    used_seed = seed if (seed is not None and str(seed).strip() != \"\") else random.randint(0, 2**31 - 1)
-    url = f\"{api}?prompt={quote_plus(prompt)}&seed={used_seed}\"
+    url = f"{INFERENCE_URL}?prompt={quote(prompt or '')}&seed={seed}"
 
     files = {
-        \"image_bytes\": (os.path.basename(image_path), open(image_path, \"rb\"), \"application/octet-stream\")
+        "image_bytes": ("start.png", _png_bytes(start_img), "image/png"),
+        "image_bytes_end": ("end.png", _png_bytes(end_img), "image/png"),
     }
-    headers = {\"accept\": \"application/json\"}
+    headers = {"accept": "application/json"}
 
     try:
-        resp = requests.post(url, headers=headers, files=files, timeout=1200)
-        # Try to accommodate various backends
-        ctype = resp.headers.get(\"Content-Type\", \"\")
-        if \"application/json\" in ctype:
-            data = resp.json()
-            # base64 payload
-            if \"video\" in data and isinstance(data[\"video\"], str) and len(data[\"video\"]) > 50:
-                try:
-                    raw = base64.b64decode(data[\"video\"], validate=True)
-                    return _save_bytes_to_mp4(raw, \"clip\"), None
-                except Exception as e:
-                    return None, f\"Could not decode base64 video: {e}\"
-            # url payload
-            for key in (\"video_url\", \"result_url\", \"url\"):
-                if key in data and isinstance(data[key], str) and data[key].startswith(\"http\"):
-                    raw = _download(data[key])
-                    return _save_bytes_to_mp4(raw, \"clip\"), None
-            return None, 'JSON response did not include \"video\" (base64) or a known url key.'
-        # Raw bytes (ideally mp4)
-        elif \"video\" in ctype or \"octet-stream\" in ctype:
-            return _save_bytes_to_mp4(resp.content, \"clip\"), None
-        else:
-            # Some backends still reply bytes with missing/odd content-type
-            if resp.content and len(resp.content) > 1024:
-                return _save_bytes_to_mp4(resp.content, \"clip\"), None
-            return None, f\"Unexpected content type: {ctype}\"
-    except requests.RequestException as e:
-        return None, f\"Request failed: {e}\"
-
-
-def stitch_pair(
-    image_a: str,
-    image_b: str,
-    prompt: str,
-    seed: Optional[int],
-    api_url: Optional[str],
-    crossfade: float,
-) -> Tuple[Optional[str], str]:
-    \"\"\"
-    Strategy:
-      - Generate a short clip from image A
-      - Generate a short clip from image B (same prompt/seed unless user changes)
-      - Concatenate with a short crossfade in Python (moviepy)
-
-    If you already have a backend endpoint that does stitching directly,
-    replace this function body with a single backend call.
-    \"\"\"
-    if not image_a or not image_b:
-        return None, \"Please upload both images.\"
-
-    # First generate both clips
-    clip1_path, err1 = call_i2v(image_a, prompt, seed, api_url)
-    if err1:
-        return None, f\"Clip 1 failed: {err1}\"
-    clip2_path, err2 = call_i2v(image_b, prompt, seed, api_url)
-    if err2:
-        return None, f\"Clip 2 failed: {err2}\"
-
-    # If crossfade is 0, just concatenate directly
-    try:
-        from moviepy.editor import VideoFileClip, concatenate_videoclips
-    except Exception as e:
-        return None, f\"MoviePy import failed. Add moviepy & imageio-ffmpeg to requirements.txt. Error: {e}\"
+        resp = requests.post(url, files=files, headers=headers, timeout=600)
+        ctype = (resp.headers.get("content-type") or "").lower()
+
+        # Raw video bytes
+        if "application/json" not in ctype:
+            resp.raise_for_status()
+            return _save_video_bytes(resp.content, "stitch")
+
+        # JSON with url or base64
+        data = resp.json()
+        video_url = data.get("video_url") or data.get("url") or data.get("result") or data.get("output")
+        if isinstance(video_url, str) and (video_url.startswith("http://") or video_url.startswith("https://")):
+            b = _download_to_bytes(video_url)
+            return _save_video_bytes(b, "stitch")
+
+        video_b64 = data.get("video_b64") or data.get("videoBase64")
+        if isinstance(video_b64, str):
+            pad = (-len(video_b64)) % 4
+            if pad: video_b64 += "=" * pad
+            b = base64.b64decode(video_b64)
+            return _save_video_bytes(b, "stitch")
+
+    except Exception:
+        pass
+
+    return None
+
+# -------- Gradio callbacks (exactly two stitches) --------
+def stitch_12(prompt12, seed, img1, img2):
+    if img1 is None or img2 is None:
+        gr.Warning("Please upload Image 1 and Image 2.")
+        return None
+    path = stitch_call(img1, img2, prompt12 or "", int(seed or 0))
+    if path is None:
+        gr.Warning("Stitch 1&2 failed. Try again or adjust the prompt.")
+    return path
 
-    try:
-        c1 = VideoFileClip(clip1_path)
-        c2 = VideoFileClip(clip2_path)
-
-        # Enforce same size/fps (compose handles mismatches)
-        if crossfade and crossfade > 0:
-            # Apply crossfade (second clip fades in)
-            c2 = c2.crossfadein(crossfade)
-            c1 = c1.crossfadeout(crossfade)
-            merged = concatenate_videoclips([c1, c2], method=\"compose\", padding=-crossfade)
-        else:
-            merged = concatenate_videoclips([c1, c2], method=\"compose\")
-
-        out_path = os.path.join(SAVE_DIR, f\"stitched-{int(time.time()*1000)}.mp4\")
-        merged.write_videofile(out_path, codec=\"libx264\", audio_codec=\"aac\", verbose=False, logger=None)
-        c1.close(); c2.close(); merged.close()
-        return out_path, \"\"
-    except Exception as e:
-        return None, f\"Stitching failed: {e}\"
-
-
-# -----------------------------
-# UI
-# -----------------------------
-with gr.Blocks(title=\"Image Stitch to Video\", css=\"\"\"
-/* Rounded tiles like the sketch */
-.rounded { border-radius: 24px; }
-.tile { background: #f7f7ff; padding: 12px; }
-.tile-blue { background: #e8f0ff; }
-.tile-yellow { background: #fff7d6; }
-.small-btn button { padding: 6px 10px; border-radius: 999px; }
-.label-center label { text-align:center; width: 100%; }
-\"\"\") as demo:
-    gr.Markdown(\"### Image → Video (Stitch Adjacent Pairs)\\nUpload 3 images, enter prompts for each stitch, then click the stitch buttons.\")
-
-    with gr.Row(equal_height=True):
-        # Left column: images + add image
-        with gr.Column(scale=1):
-            gr.Markdown(\"**Images**\")
-            img1 = gr.Image(type=\"filepath\", label=\"Image 1\", height=220, elem_classes=[\"rounded\", \"tile\", \"tile-blue\"])
-            img2 = gr.Image(type=\"filepath\", label=\"Image 2\", height=220, elem_classes=[\"rounded\", \"tile\", \"tile-blue\"])
-            img3 = gr.Image(type=\"filepath\", label=\"Image 3\", height=220, elem_classes=[\"rounded\", \"tile\", \"tile-blue\"])
-
-            # Optional extra slots (hidden until added)
-            extra_imgs = []
-            for i in range(4, 9):
-                comp = gr.Image(type=\"filepath\", label=f\"Image {i}\", height=220, visible=False, elem_classes=[\"rounded\", \"tile\", \"tile-blue\"])
-                extra_imgs.append(comp)
-
-            add_btn = gr.Button(\"Add Image\", variant=\"secondary\")
-
-        # Middle column: prompts + stitch buttons
-        with gr.Column(scale=1):
-            gr.Markdown(\"**Prompts**\")
-            prompt12 = gr.Textbox(label=\"Prompt for Stitch 1 & 2\", lines=3, placeholder=\"Describe motion/style/etc.\", elem_classes=[\"rounded\", \"tile\"])
-            seed12 = gr.Number(label=\"Seed (optional)\", value=None, precision=0)
-            stitch12 = gr.Button(\"Stitch 1 & 2\", elem_classes=[\"small-btn\"])
-
-            prompt23 = gr.Textbox(label=\"Prompt for Stitch 2 & 3\", lines=3, placeholder=\"Describe motion/style/etc.\", elem_classes=[\"rounded\", \"tile\"])
-            seed23 = gr.Number(label=\"Seed (optional)\", value=None, precision=0)
-            stitch23 = gr.Button(\"Stitch 2 & 3\", elem_classes=[\"small-btn\"])
-
-            with gr.Accordion(\"Advanced (API & Stitch)\", open=False):
-                api_url = gr.Textbox(label=\"Backend API URL\", value=DEFAULT_API_URL)
-                crossfade = gr.Slider(0.0, 1.5, value=0.4, step=0.1, label=\"Crossfade seconds\")
-                clear_btn = gr.Button(\"Clear All\")
-
-        # Right column: video outputs
-        with gr.Column(scale=1):
-            gr.Markdown(\"**Outputs**\")
-            vid12 = gr.Video(label=\"Video (image 1 + 2) output\", elem_classes=[\"rounded\", \"tile\", \"tile-yellow\"])
-            vid23 = gr.Video(label=\"Video (image 2 + 3) output\", elem_classes=[\"rounded\", \"tile\", \"tile-yellow\"])
-
-    # Wire up actions
-    def _on_add(*imgs):
-        # Reveal the next hidden uploader
-        for comp in extra_imgs:
-            if comp.visible is False:
-                comp.visible = True
-                break
-        return [gr.update(visible=comp.visible) for comp in extra_imgs]
-
-    add_btn.click(
-        _on_add,
-        inputs=extra_imgs,
-        outputs=extra_imgs,
-    )
-
-    stitch12.click(
-        stitch_pair,
-        inputs=[img1, img2, prompt12, seed12, api_url, crossfade],
-        outputs=[vid12, gr.Textbox(visible=False)],
-    )
-
-    stitch23.click(
-        stitch_pair,
-        inputs=[img2, img3, prompt23, seed23, api_url, crossfade],
-        outputs=[vid23, gr.Textbox(visible=False)],
-    )
-
-    def _on_clear():
-        updates = []
-        for comp in [img1, img2, img3, *extra_imgs]:
-            updates.append(gr.update(value=None, visible=True if comp in [img1, img2, img3] else False))
-        return updates + [None, None, \"\", \"\", gr.update(value=DEFAULT_API_URL), 0.4]
-
-    clear_btn.click(
-        _on_clear,
-        inputs=None,
-        outputs=[img1, img2, img3, *extra_imgs, vid12, vid23, prompt12, prompt23, api_url, crossfade],
-    )
-
-if __name__ == \"__main__\":
-    demo.launch()
+def stitch_23(prompt23, seed, img2, img3):
+    if img2 is None or img3 is None:
+        gr.Warning("Please upload Image 2 and Image 3.")
+        return None
+    path = stitch_call(img2, img3, prompt23 or "", int(seed or 0))
+    if path is None:
+        gr.Warning("Stitch 2&3 failed. Try again or adjust the prompt.")
+    return path
+
+# -------- UI --------
+CSS = """
+.gradio-container { padding: 24px; }
+.pill button { border-radius: 999px !important; padding: 10px 18px; }
+.rounded textarea { border-radius: 16px !important; }
+"""
+
+with gr.Blocks(css=CSS, title="Stitch — 3 uploads, 2 stitches") as demo:
+    gr.Markdown("## Stitch — Upload 3 images, then generate videos between 1→2 and 2→3")
+
+    with gr.Row():
+        # Left: exactly 3 image inputs
+        with gr.Column(scale=1, min_width=360):
+            img1 = gr.Image(label="Image 1 upload", type="pil")
+            img2 = gr.Image(label="Image 2 upload", type="pil")
+            img3 = gr.Image(label="Image 3 upload", type="pil")
+
+        # Middle: prompts + buttons
+        with gr.Column(scale=1, min_width=360):
+            seed = gr.Number(value=0, precision=0, label="Seed (0 = random)")
+            prompt12 = gr.Textbox(placeholder="Prompt for stitching 1→2", lines=2, label="Prompt (1→2)", elem_classes=["rounded"])
+            btn12 = gr.Button("Stitch 1&2", elem_classes=["pill"])
+            gr.Markdown("---")
+            prompt23 = gr.Textbox(placeholder="Prompt for stitching 2→3", lines=2, label="Prompt (2→3)", elem_classes=["rounded"])
+            btn23 = gr.Button("Stitch 2&3", elem_classes=["pill"])
+
+        # Right: exactly 2 video outputs
+        with gr.Column(scale=1, min_width=360):
+            vid12 = gr.Video(label="Video (image 1+2) output")
+            vid23 = gr.Video(label="Video (image 2+3) output")
+
+    # Wire buttons
+    btn12.click(stitch_12, inputs=[prompt12, seed, img1, img2], outputs=[vid12])
+    btn23.click(stitch_23, inputs=[prompt23, seed, img2, img3], outputs=[vid23])
+
+if __name__ == "__main__":
+    demo.queue().launch()