Blame for app.py

"""JoyAI-Video-Edit — streaming instruction-guided video editing.

Chunk-causal port of the reference deployment
(https://github.com/jd-opensource/JoyAI-Video-Edit, `deploy/`) onto ZeroGPU.
Frames are edited 8 at a time by a 16.3B causal MMDiT with a rolling KV cache,
so the edited clip streams back chunk by chunk instead of appearing at the end.

The saved file keeps the source soundtrack. The live frame preview stays silent.
Requires `ffmpeg` in packages.txt.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import spaces  # noqa: E402  (must precede torch/CUDA imports)

import subprocess  # noqa: E402
import tempfile  # noqa: E402
import threading  # noqa: E402
import time  # noqa: E402

import gradio as gr  # noqa: E402
import imageio.v3 as iio  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
from PIL import Image  # noqa: E402

from xvideo.models.loader import load_everything  # noqa: E402
from xvideo.stream import StreamingEditor, StreamingSettings  # noqa: E402

TARGET_FPS = 24                 # matches `--record-fps 24` in the reference server
LANDSCAPE = (720, 1248)         # (height, width) — reference server default
PORTRAIT = (1248, 720)
FRAMES_PER_CHUNK = 8            # VAE temporal factor
MAX_SECONDS = 15.0
MIN_SECONDS = 0.5
MAX_SEED = 2**31 - 1

_GPU_LOCK = threading.Lock()

CFG, PIPELINE = load_everything(device="cuda")


# --------------------------------------------------------------------------- io
def read_frames(path: str, max_frames: int) -> tuple[list[Image.Image], float]:
    """Decode `path`, resample to TARGET_FPS, return at most `max_frames` PIL frames."""
    try:
        meta = iio.immeta(path, plugin="FFMPEG")
        src_fps = float(meta.get("fps") or TARGET_FPS)
    except Exception:  # noqa: BLE001
        src_fps = float(TARGET_FPS)
    if not np.isfinite(src_fps) or src_fps <= 0:
        src_fps = float(TARGET_FPS)

    step = src_fps / float(TARGET_FPS)
    frames: list[Image.Image] = []
    next_wanted = 0.0
    for idx, frame in enumerate(iio.imiter(path, plugin="FFMPEG")):
        if idx + 1e-6 >= next_wanted:
            frames.append(Image.fromarray(np.asarray(frame)[..., :3]))
            next_wanted += step
            if len(frames) >= max_frames:
                break
    return frames, src_fps


def write_video(frames: np.ndarray, fps: int = TARGET_FPS) -> str:
    path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    iio.imwrite(path, frames, plugin="FFMPEG", fps=fps, codec="libx264",
                output_params=["-pix_fmt", "yuv420p", "-crf", "18"])
    return path


def attach_source_audio(edited_path: str, source_path: str) -> str:
    """Copy the upload's audio onto the silent edited mp4.

    Returns the original path if the source has no audio stream or ffmpeg fails,
    so a missing soundtrack never drops the edited video.
    """
    if not source_path or not os.path.isfile(source_path):
        return edited_path
    probe = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "a",
         "-show_entries", "stream=index", "-of", "csv=p=0", source_path],
        capture_output=True, text=True,
    )
    if probe.returncode != 0 or not probe.stdout.strip():
        return edited_path
    out_path = edited_path.replace(".mp4", "_audio.mp4")
    mux = subprocess.run(
        ["ffmpeg", "-y", "-i", edited_path, "-i", source_path,
         "-map", "0:v:0", "-map", "1:a:0",
         "-c:v", "copy", "-c:a", "aac", "-b:a", "192k", "-shortest",
         out_path],
        capture_output=True, text=True,
    )
    if mux.returncode != 0 or not os.path.isfile(out_path):
        return edited_path
    return out_path


def chunks_for_seconds(seconds: float) -> int:
    wanted = max(1, int(round(float(seconds) * TARGET_FPS)))
    return max(1, 1 + (wanted - 1) // FRAMES_PER_CHUNK)


def _estimate_duration(
    video_path=None,
    instruction="",
    seconds=2.0,
    seed=42,
    randomize_seed=False,
    num_inference_steps=2,
    *args,
    **kwargs,
):
    """Reserve GPU time scaled to the real per-chunk cost.

    Measured on the live sm_120 slice: ~1.0 s per (chunk × step) plus fixed
    overhead (condition encode + mux). Worst case (12 chunks × 4 steps) runs in
    ~48 s of GPU compute, so 15 + units × 1.1 (≈68 s worst case) covers it with
    a modest margin while keeping short runs cheap and high-priority.
    """
    try:
        chunks = chunks_for_seconds(float(seconds))
        steps = max(1, int(num_inference_steps))
    except (TypeError, ValueError):
        chunks, steps = 12, 4
    return int(min(120, 15 + chunks * steps * 1.1))


# ---------------------------------------------------------------------- handler
@spaces.GPU(duration=_estimate_duration, size="large")
def edit_video(
    video_path,
    instruction,
    seconds=2.0,
    seed=42,
    randomize_seed=False,
    num_inference_steps=4,
    reference_image=None,
    progress=gr.Progress(),
):
    if not video_path:
        raise gr.Error("Please upload or record a source video first.")
    instruction = (instruction or "").strip()
    if not instruction:
        raise gr.Error("Please describe the edit you want (e.g. 'turn it into a watercolor wash').")

    seconds = float(np.clip(float(seconds), MIN_SECONDS, MAX_SECONDS))
    num_chunks = chunks_for_seconds(seconds)
    needed_frames = 1 + FRAMES_PER_CHUNK * (num_chunks - 1)

    if randomize_seed:
        seed = int(np.random.randint(0, MAX_SEED))
    seed = int(seed) % (MAX_SEED + 1)

    progress(0.0, desc="Decoding source video…")
    frames, src_fps = read_frames(video_path, needed_frames)
    if len(frames) < 1:
        raise gr.Error("Could not read any frames from that video.")
    if len(frames) < needed_frames:
        num_chunks = 1 + (len(frames) - 1) // FRAMES_PER_CHUNK
        needed_frames = 1 + FRAMES_PER_CHUNK * (num_chunks - 1)
        frames = frames[:needed_frames]

    src_w, src_h = frames[0].size
    height, width = PORTRAIT if src_h > src_w else LANDSCAPE

    ref_image = None
    if reference_image is not None:
        ref_image = reference_image if isinstance(reference_image, Image.Image) else Image.fromarray(reference_image)

    settings = StreamingSettings(
        height=height,
        width=width,
        num_inference_steps=int(num_inference_steps),
        seed=seed,
    )

    collected: list[np.ndarray] = []
    editor = None
    with _GPU_LOCK:
        try:
            editor = StreamingEditor(CFG, PIPELINE, settings)
            progress(0.02, desc="Encoding instruction + first frame…")
            t_start = time.perf_counter()
            first_chunk_at = None

            for chunk_idx, chunk_frames in editor.run(instruction, frames, ref_image=ref_image):
                collected.append(chunk_frames)
                if first_chunk_at is None:
                    first_chunk_at = time.perf_counter()
                done = chunk_idx + 1
                elapsed = time.perf_counter() - t_start
                n_out = sum(c.shape[0] for c in collected)
                status = (
                    f"**streaming** — chunk {done}/{num_chunks} · {n_out} frames "
                    f"({n_out / TARGET_FPS:.2f}s of video) · {elapsed:.1f}s elapsed · "
                    f"{n_out / max(elapsed, 1e-6):.1f} generated fps"
                )
                progress(min(0.98, 0.02 + 0.96 * done / num_chunks), desc=f"Chunk {done}/{num_chunks}")
                yield chunk_frames[-1], None, status

            total = time.perf_counter() - t_start
        finally:
            if editor is not None:
                editor.close()

    if not collected:
        raise gr.Error("The video was too short to edit — try a clip of at least half a second.")

    all_frames = np.concatenate(collected, axis=0)
    progress(0.99, desc="Muxing result…")
    out_path = write_video(all_frames, fps=TARGET_FPS)
    out_path = attach_source_audio(out_path, video_path)

    n_out = int(all_frames.shape[0])
    status = (
        f"**done** — {n_out} frames ({n_out / TARGET_FPS:.2f}s) at "
        f"{all_frames.shape[2]}×{all_frames.shape[1]} in {total:.1f}s "
        f"→ **{n_out / max(total, 1e-6):.1f} generated fps** "
        f"({len(collected)} chunks, {settings.num_inference_steps} steps/chunk, seed {seed})"
    )
    yield all_frames[-1], out_path, status


# --------------------------------------------------------------------------- ui
CSS = """
#col-container { max-width: 1180px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES = [
    ["examples/case02_watercolor.mp4", "Turn the video into a watercolor wash style."],
    ["examples/case01_castle.mp4", "Transform the people, hairstyles, and interior into a British castle aristocratic style."],
    ["examples/case03_dogs.mp4", "Make all dogs white, add colorful hats, and turn the sunglasses hot pink."],
    ["examples/case04_street.mp4", "Dress the girl in a brown down jacket and blue baseball cap."],
    ["examples/case05_cats.mp4", "Remove the two white cats in pink clothes on both sides."],
]

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
# JoyAI-Video-Edit — streaming video editing
Describe an edit in plain language and it is applied to your clip **chunk by chunk**:
a 16.3B chunk-causal MMDiT edits 8 frames at a time behind a rolling KV cache, so the
result streams back while it is still being generated (2 flow-matching steps per chunk).
The downloaded file keeps the original audio. The live preview stays silent.

[model](https://huggingface.co/jdopensource/JoyAI-Video-Edit) ·
[reference code](https://github.com/jd-opensource/JoyAI-Video-Edit) ·
[technical report](https://arxiv.org/pdf/2608.03974)
"""
        )

        with gr.Row():
            with gr.Column():
                video_in = gr.Video(
                    label="Source video",
                    sources=["upload", "webcam"],
                    include_audio=True,
                    height=320,
                )
                instruction = gr.Textbox(
                    label="Edit instruction",
                    placeholder="Turn the video into a watercolor wash style.",
                    lines=2,
                )
                seconds = gr.Slider(
                    MIN_SECONDS, MAX_SECONDS, value=2.0, step=0.5,
                    label="Seconds to edit",
                    info="Frames are taken from the start of the clip at 24 fps.",
                )
                run_btn = gr.Button("Edit video", variant="primary")

                with gr.Accordion("Advanced settings", open=False):
                    steps = gr.Slider(
                        1, 4, value=4, step=1,
                        label="Flow-matching steps per chunk",
                        info="The released checkpoint is distilled for 2 steps.",
                    )
                    seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
                    randomize_seed = gr.Checkbox(value=False, label="Randomize seed")
                    reference_image = gr.Image(
                        label="Reference image (optional, RV2V)",
                        type="pil",
                        height=200,
                    )

            with gr.Column():
                preview = gr.Image(
                    label="Live preview (latest generated frame)",
                    height=320,
                )
                video_out = gr.Video(label="Edited video", height=320, autoplay=True)
                status = gr.Markdown("")

        gr.Examples(
            examples=EXAMPLES,
            inputs=[video_in, instruction],
            outputs=[preview, video_out, status],
            fn=edit_video,
            cache_examples=False,
            run_on_click=True,
            label="Showcase clips from the reference repo",
        )

    inputs = [video_in, instruction, seconds, seed, randomize_seed, steps, reference_image]
    outputs = [preview, video_out, status]
    # Canonical API endpoint — the click handler owns `/edit_video` with the full
    # 7-input signature. The submit handler gets a distinct name so it can't shadow
    # the primary binding or register with a truncated signature.
    run_btn.click(edit_video, inputs=inputs, outputs=outputs,
                  concurrency_limit=1, api_name="edit_video")
    instruction.submit(edit_video, inputs=inputs, outputs=outputs,
                       concurrency_limit=1, api_name="edit_video_submit")

if __name__ == "__main__":
    demo.queue(max_size=12).launch(theme=gr.themes.Citrus(), css=CSS)