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