Update app.py

app.py

@@ -1,301 +1,211 @@
-import os
-import io
-import json
-import uuid
-import time
-import shutil
-import mimetypes
-import requests
-import subprocess
+import os, uuid, time, shutil, mimetypes, subprocess, requests, gradio as gr
 from datetime import datetime
-from typing import List, Optional
+from typing import Optional, List
 
-import gradio as gr
-from PIL import Image
+# --- config ---
+ENDPOINT = "https://moonmath-ai-dev--moonmath-i2v-backend-moonmathinference-run.modal.run"
+FFMPEG = "ffmpeg"
+OUT, TMP = "outputs", "tmp"
+os.makedirs(OUT, exist_ok=True); os.makedirs(TMP, exist_ok=True)
 
-# ================================
-# Config
-# ================================
-ENDPOINT_URL = "https://moonmath-ai-dev--moonmath-i2v-backend-moonmathinference-run.modal.run"  # your backend
-OUT_DIR = "outputs"
-TMP_DIR = "tmp"
-os.makedirs(OUT_DIR, exist_ok=True)
-os.makedirs(TMP_DIR, exist_ok=True)
+ts     = lambda: datetime.utcnow().strftime("%Y%m%d_%H%M%S")
+fname  = lambda p,e: f"{p}_{ts()}_{uuid.uuid4().hex[:6]}.{e}"
+abspath= lambda p: os.path.abspath(p)
 
-# If ffmpeg is not on PATH for any reason, set absolute path here
-FFMPEG_BIN = "ffmpeg"
-
-# ================================
-# Helpers
-# ================================
-
-def _ts() -> str:
-    return datetime.utcnow().strftime("%Y%m%d_%H%M%S")
-
-def _safe_filename(prefix: str, ext: str) -> str:
-    return f"{prefix}_{_ts()}_{uuid.uuid4().hex[:8]}.{ext}"
-
-def _run_ffmpeg(args: List[str]) -> None:
-    """Run ffmpeg with given args; raise on failure with readable message."""
-    cmd = [FFMPEG_BIN] + args
-    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    if proc.returncode != 0:
-        raise RuntimeError(
-            f"ffmpeg failed ({proc.returncode}).\nSTDOUT:\n{proc.stdout.decode(errors='ignore')}\n\nSTDERR:\n{proc.stderr.decode(errors='ignore')}"
-        )
+def run_ffmpeg(args: List[str]):
+    p = subprocess.run([FFMPEG]+args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    if p.returncode:
+        raise RuntimeError(p.stderr.decode(errors="ignore"))
 
 def extract_last_frame(video_path: str) -> str:
-    """
-    Extract the very last frame as PNG using ffmpeg.
-    Uses -sseof -1 to seek to the last second.
-    """
-    out_path = os.path.join(TMP_DIR, _safe_filename("lastframe", "png"))
-    # Seek to last second and output a single frame
-    _run_ffmpeg([
-        "-sseof", "-1",
-        "-i", video_path,
-        "-update", "1",
-        "-frames:v", "1",
-        "-q:v", "2",
-        out_path
-    ])
-    return out_path
-
-def concat_videos_ffmpeg(video_paths: List[str]) -> str:
-    """
-    Concatenate MP4s using ffmpeg concat demuxer.
-    """
-    if len(video_paths) == 1:
-        # copy the single file to a new name to keep UX consistent
-        dst = os.path.join(OUT_DIR, _safe_filename("continuous", "mp4"))
-        shutil.copy(video_paths[0], dst)
-        return dst
-
-    list_file = os.path.join(TMP_DIR, f"concat_{uuid.uuid4().hex}.txt")
-    with open(list_file, "w") as f:
-        for p in video_paths:
-            # paths must be quoted if they may contain spaces
-            f.write(f"file '{os.path.abspath(p)}'\n")
-
-    out_path = os.path.join(OUT_DIR, _safe_filename("continuous", "mp4"))
-    _run_ffmpeg([
-        "-f", "concat",
-        "-safe", "0",
-        "-i", list_file,
-        "-c", "copy",
-        out_path
-    ])
-    try:
-        os.remove(list_file)
-    except Exception:
-        pass
-    return out_path
-
-def zip_files(file_paths: List[str]) -> str:
-    base = os.path.join(OUT_DIR, f"used_{_ts()}")
-    shutil.make_archive(base, "zip", root_dir=os.path.dirname(file_paths[0]) if file_paths else ".", base_dir=".")
-    # The above zips the entire directory — better: copy into temp pack dir first
-    # For precision, we’ll do that instead:
-    stamp = _ts()
-    pack_dir = os.path.join(TMP_DIR, f"pack_{stamp}_{uuid.uuid4().hex[:6]}")
-    os.makedirs(pack_dir, exist_ok=True)
-    for p in file_paths:
-        shutil.copy(p, pack_dir)
-    zip_base = os.path.join(OUT_DIR, f"used_{stamp}")
-    shutil.make_archive(zip_base, "zip", pack_dir)
-    shutil.rmtree(pack_dir, ignore_errors=True)
-    return f"{zip_base}.zip"
-
-def _guess_filename_from_response(resp: requests.Response, default_ext: str = "mp4") -> str:
-    # Try Content-Disposition
-    cd = resp.headers.get("Content-Disposition", "")
-    if "filename=" in cd:
-        name = cd.split("filename=")[-1].strip('"; ')
-        if name:
-            return name
-    # Otherwise, attempt from content-type
-    ctype = resp.headers.get("Content-Type", "")
-    ext = mimetypes.guess_extension(ctype) or f".{default_ext}"
-    return _safe_filename("gen", ext.lstrip("."))
-
-def save_bytes_to_file(content: bytes, ext: str = "mp4") -> str:
-    path = os.path.join(OUT_DIR, _safe_filename("gen", ext))
-    with open(path, "wb") as f:
-        f.write(content)
+    out = os.path.join(TMP, fname("last", "png"))
+    run_ffmpeg(["-sseof","-1","-i",video_path,"-frames:v","1","-q:v","2",out])
+    return out
+
+def concat_videos(paths: List[str]) -> str:
+    if not paths: raise ValueError("No videos selected.")
+    if len(paths)==1:
+        dst = os.path.join(OUT, fname("continuous","mp4")); shutil.copy(paths[0], dst); return dst
+    listfile = os.path.join(TMP, f"concat_{uuid.uuid4().hex}.txt")
+    with open(listfile,"w") as f:
+        for p in paths: f.write(f"file '{abspath(p)}'\n")
+    out = os.path.join(OUT, fname("continuous","mp4"))
+    run_ffmpeg(["-f","concat","-safe","0","-i",listfile,"-c","copy",out])
+    os.remove(listfile)
+    return out
+
+def zip_used(paths: List[str]) -> str:
+    pack = os.path.join(TMP, f"pack_{uuid.uuid4().hex[:6]}")
+    os.makedirs(pack, exist_ok=True)
+    for p in paths: shutil.copy(p, pack)
+    base = os.path.join(OUT, f"used_{ts()}")
+    shutil.make_archive(base, "zip", pack)
+    shutil.rmtree(pack, ignore_errors=True)
+    return base + ".zip"
+
+def save_video_bytes(content: bytes, content_type: str) -> str:
+    ext = (mimetypes.guess_extension(content_type) or ".mp4").lstrip(".")
+    path = os.path.join(OUT, fname("gen", ext))
+    with open(path,"wb") as f: f.write(content)
     return path
 
-# ================================
-# Backend call
-# ================================
-def call_generator(prompt: str, seed_frame_path: Optional[str]) -> str:
-    """
-    Calls your Modal endpoint. Supports two response modes:
-      1) Direct video bytes (Content-Type like video/mp4)
-      2) JSON with {"video_url": "..."} which we then download.
-
-    We send:
-      - form field: prompt
-      - optional file: seed_frame (PNG/JPG)
-    """
-    files = {}
-    data = {"prompt": prompt}
-
-    if seed_frame_path and os.path.exists(seed_frame_path):
-        files["seed_frame"] = (os.path.basename(seed_frame_path), open(seed_frame_path, "rb"), "image/png")
-
-    # Make the request
-    resp = requests.post(ENDPOINT_URL, data=data, files=files, timeout=600)
-    if files:
-        # close file handles
-        for _, f in files.items():
-            # f is a tuple (name, fileobj, mimetype)
-            try:
-                f[1].close()
-            except Exception:
-                pass
+def call_backend(
+    prompt: str,
+    image_bytes_path: str,
+    negative_prompt: Optional[str],
+    fps: int,
+    vlen: int,
+    steps: Optional[int],
+    seed: Optional[int]
+) -> str:
+    params = {
+        "prompt": prompt,
+        "frames_per_second": str(fps),
+        "video_length": str(vlen),
+    }
+    if negative_prompt:           params["negative_prompt"]      = negative_prompt
+    if steps is not None:         params["num_inference_steps"]  = str(steps)
+    if seed is None:              seed = int(time.time())
+    params["seed"] = str(seed)
+
+    files = {"image_bytes": (os.path.basename(image_bytes_path), open(image_bytes_path,"rb"),
+                             "application/octet-stream")}
+    try:
+        r = requests.post(ENDPOINT, params=params, files=files, headers={"accept":"application/json"}, timeout=600)
+    finally:
+        try: files["image_bytes"][1].close()
+        except: pass
 
-    if resp.status_code != 200:
-        raise RuntimeError(f"Generator returned {resp.status_code}: {resp.text[:500]}")
+    if r.status_code != 200:
+        raise RuntimeError(f"Backend {r.status_code}: {r.text[:500]}")
 
-    ctype = resp.headers.get("Content-Type", "")
-    if ctype.startswith("video/"):
-        # Direct bytes
-        ext = mimetypes.guess_extension(ctype) or ".mp4"
-        out_path = os.path.join(OUT_DIR, _safe_filename("gen", ext.lstrip(".")))
-        with open(out_path, "wb") as f:
-            f.write(resp.content)
-        return out_path
+    ctype = r.headers.get("Content-Type","")
+    if ctype.startswith("video/"):   # raw video bytes
+        return save_video_bytes(r.content, ctype)
 
-    # Assume JSON with URL
+    # expect JSON with { "video_url": ... }
     try:
-        payload = resp.json()
+        payload = r.json()
+        url = payload.get("video_url")
+        if not url: raise ValueError("no video_url in response")
+        r2 = requests.get(url, stream=True, timeout=600)
+        if r2.status_code != 200: raise RuntimeError(f"fetch video {r2.status_code}")
+        path = os.path.join(OUT, fname("gen","mp4"))
+        with open(path,"wb") as f:
+            for chunk in r2.iter_content(1<<20):
+                if chunk: f.write(chunk)
+        return path
     except Exception:
-        # fallback: treat body as raw mp4
-        return save_bytes_to_file(resp.content, "mp4")
-
-    video_url = payload.get("video_url")
-    if not video_url:
-        # Maybe it returned base64? If so, add handling here if needed.
-        raise RuntimeError(f"Unexpected response: {json.dumps(payload)[:500]}")
+        # fallback if backend mislabels content
+        return save_video_bytes(r.content, ctype or "video/mp4")
 
-    # Download the video
-    r2 = requests.get(video_url, stream=True, timeout=600)
-    if r2.status_code != 200:
-        raise RuntimeError(f"Failed to fetch video from URL: {r2.status_code}")
-    ext = mimetypes.guess_extension(r2.headers.get("Content-Type", "")) or ".mp4"
-    out_path = os.path.join(OUT_DIR, _safe_filename("gen", ext.lstrip(".")))
-    with open(out_path, "wb") as f:
-        for chunk in r2.iter_content(chunk_size=1024 * 1024):
-            if chunk:
-                f.write(chunk)
-    return out_path
-
-# ================================
-# Gradio UI
-# ================================
+# ================= UI =================
 with gr.Blocks() as demo:
-    gr.Markdown("## Continuous Video (Prompt → Use → Chain → Download)")
+    gr.Markdown("## Continuous Video — chain last frame → next first frame")
 
-    state_used = gr.State([])         # list[str] of used clip paths
-    state_seed = gr.State(None)       # path to last-frame image
-    state_current = gr.State(None)    # most recent generated video path
+    # state
+    used_videos   = gr.State([])      # list[str]
+    last_seed_img = gr.State(None)    # PNG path (becomes image_bytes)
+    current_video = gr.State(None)    # latest generated video
 
     with gr.Row():
-        prompt = gr.Textbox(
-            label="Prompt",
-            placeholder="Describe your next shot…",
-            lines=2,
-            autofocus=True
+        prompt = gr.Textbox(label="Prompt", placeholder="Describe your shot…", lines=2)
+    with gr.Row():
+        start_file = gr.File(
+            label="Initial start image or video (only for the very first clip)",
+            file_types=["image","video"]
         )
     with gr.Row():
-        video_out = gr.Video(label="Video Output", interactive=False)
-
+        negative = gr.Textbox(label="Negative prompt (optional)", placeholder="What to avoid…")
     with gr.Row():
-        btn_generate = gr.Button("Generate", variant="primary")
-        btn_use = gr.Button("Chain")
-        btn_download = gr.Button("Download All")
-        btn_reset = gr.Button("Reset Session", variant="stop")
-
-    files_out = gr.Files(label="Downloads", height=100)
-
-    # ---------- Handlers ----------
+        fps   = gr.Slider(1,60, value=24, step=1, label="Frames per second")
+        vlen  = gr.Slider(1,12, value=4,  step=1, label="Video length (seconds)")
+        steps = gr.Slider(1,100, value=30, step=1, label="Num inference steps (optional)")
+        seed  = gr.Number(label="Seed (optional)", precision=0)
 
-    def on_generate(prompt_text, seed_frame):
-        if not prompt_text or not prompt_text.strip():
-            return gr.update(value=None), None
+    video_out = gr.Video(label="Output")
+    status_md = gr.Markdown("")  # shows ✅ Saved… after Chain
+    with gr.Row():
+        gen_btn  = gr.Button("Generate", variant="primary")
+        chain_btn= gr.Button("Chain (save & clear)")
+        dl_btn   = gr.Button("Download (concatenate & ZIP)")
+
+    files_out = gr.Files(label="Downloads")
+
+    # ---- handlers ----
+    def on_generate(prompt_txt, start_file_obj, neg, fps_val, vlen_val, steps_val, seed_val, seed_img):
+        if not prompt_txt or not prompt_txt.strip():
+            raise gr.Error("Prompt is required.")
+        # choose image_bytes source: last stolen frame > uploaded start file
+        image_path = seed_img
+        if not image_path:
+            if not start_file_obj: raise gr.Error("First generation requires an initial image OR video.")
+            image_path = start_file_obj.name
         try:
-            vid_path = call_generator(prompt_text.strip(), seed_frame)
+            out = call_backend(
+                prompt=prompt_txt.strip(),
+                image_bytes_path=image_path,
+                negative_prompt=(neg.strip() if (neg and neg.strip()) else None),
+                fps=int(fps_val),
+                vlen=int(vlen_val),
+                steps=int(steps_val) if steps_val else None,
+                seed=int(seed_val) if (seed_val not in [None,""]) else None
+            )
         except Exception as e:
-            # Surface the error in the UI
-            err_mp4 = os.path.join(OUT_DIR, _safe_filename("error", "txt"))
-            with open(err_mp4, "w") as f:
-                f.write(str(e))
-            return gr.update(value=None), None
-        return vid_path, vid_path
+            raise gr.Error(str(e))
+        return out, out, gr.update(value="")  # show video, set state, clear status
 
-    btn_generate.click(
+    gen_btn.click(
         on_generate,
-        inputs=[prompt, state_seed],
-        outputs=[video_out, state_current],
+        inputs=[prompt, start_file, negative, fps, vlen, steps, seed, last_seed_img],
+        outputs=[video_out, current_video, status_md]
     )
 
-    def on_use(current_path, used_paths):
-        if not current_path or not os.path.exists(current_path):
-            return used_paths, None
-        new_used = list(used_paths)
-        if current_path not in new_used:
-            new_used.append(current_path)
-        # Extract last frame -> seed for next gen
-        try:
-            seed_img = extract_last_frame(current_path)
-        except Exception as e:
-            # If extraction fails, keep chaining logic but without seed
-            seed_img = None
-        return new_used, seed_img
+    def on_chain(curr, used):
+        """
+        Save current video, steal last frame for next seed,
+        and CLEAR: preview, prompt, negative, start_file.
+        Also show a visible 'saved' banner.
+        """
+        if not curr or not os.path.exists(curr):
+            raise gr.Error("Generate a video first.")
+        new_used = used[:] if used else []
+        if curr not in new_used: new_used.append(curr)
+        # last frame for next seed
+        seed_img = extract_last_frame(curr)
+        saved_idx = len(new_used)
+        saved_msg = f"✅ **Saved clip #{saved_idx}** — last frame captured. Ready for your next prompt."
+
+        return (
+            new_used,            # used_videos
+            seed_img,            # last_seed_img (image_bytes for next gen)
+            gr.update(value=None), # CLEAR prompt
+            gr.update(value=None), # CLEAR negative prompt
+            gr.update(value=None), # CLEAR start_file
+            gr.update(value=None), # CLEAR video preview
+            gr.update(value=saved_msg)  # status banner
+        )
 
-    btn_use.click(
-        on_use,
-        inputs=[state_current, state_used],
-        outputs=[state_used, state_seed],
+    chain_btn.click(
+        on_chain,
+        inputs=[current_video, used_videos],
+        outputs=[used_videos, last_seed_img, prompt, negative, start_file, video_out, status_md]
     )
 
-    def on_download(used_paths):
-        if not used_paths:
-            return []
+    def on_download(used):
+        if not used: raise gr.Error("Nothing to download. Click Chain after generating clips.")
+        files = []
         try:
-            concat_path = concat_videos_ffmpeg(used_paths)
+            merged = concat_videos(used); files.append(merged)
         except Exception as e:
-            # If concat fails, skip it but still offer the ZIP
-            concat_path = None
-        zip_path = zip_files(used_paths)
-        files = [zip_path]
-        if concat_path:
-            files.insert(0, concat_path)
-        return files
-
-    btn_download.click(
-        on_download,
-        inputs=[state_used],
-        outputs=[files_out],
-    )
-
-    def on_reset():
-        # Clear temp files and state
+            print("concat error:", e)
         try:
-            for f in os.listdir(TMP_DIR):
-                fp = os.path.join(TMP_DIR, f)
-                if os.path.isfile(fp):
-                    os.remove(fp)
-        except Exception:
-            pass
-        return None, [], None, gr.update(value=None), gr.update(value=None), gr.update(value=[])
+            zipped = zip_used(used); files.append(zipped)
+        except Exception as e:
+            print("zip error:", e)
+        return files
 
-    btn_reset.click(
-        on_reset,
-        inputs=None,
-        outputs=[state_seed, state_used, state_current, prompt, video_out, files_out],
-    )
+    dl_btn.click(on_download, inputs=[used_videos], outputs=[files_out])
 
 if __name__ == "__main__":
     demo.launch()
-