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.
"""

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 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 = 4.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 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=2,
    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)

    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).

[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=False,
                    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=2, 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)