Qwen-Video-Edit ZeroGPU demo

.gitattributes

@@ -35,3 +35,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 example1.mp4 filter=lfs diff=lfs merge=lfs -text
 example2.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/city_traffic_night.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/man_dancing.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/ocean_waves.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/pottery_wheel.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -4,37 +4,61 @@ emoji: 🎬
 colorFrom: gray
 colorTo: green
 sdk: gradio
-sdk_version: 6.15.1
+sdk_version: 6.24.0
 app_file: app.py
-short_description: Instruction-based video editing via an image editing model
+short_description: Instruction-based video editing with Qwen-Image-Edit
 python_version: "3.12"
 startup_duration_timeout: 1h
+license: mit
 ---
 
 # Qwen-Video-Edit
 
-Instruction-based **video editing by repurposing an image editing model**. Upload a source video and an editing instruction, and the model produces an edited video by applying the instruction to the video content.
+Demo of **[Qwen-Video-Edit: Instruction-Based Video Editing by Repurposing an Image
+Editing Model](https://huggingface.co/papers/2608.14790)** (Bai, Gandelsman, Gharbi, Huang).
 
-Built on [Qwen-Image-Edit](https://huggingface.co/Qwen/Qwen-Image-Edit)'s DiT, which directly edits Wan 2.1 video-VAE latents, bridged by two tiny trainable projections warm-started from the DiT's own input/output layers.
+* Paper: https://huggingface.co/papers/2608.14790
+* Project page: https://yunpeng1998.github.io/Qwen-Video-Edit-Page
+* Code: https://github.com/yunpeng1998/Qwen-Video-Edit
+* Weights: https://huggingface.co/yunpeng1998/Qwen-Video-Edit (`360P/step-30000.safetensors`)
 
-**Paper**: [Qwen-Video-Edit: Instruction-Based Video Editing by Repurposing an Image Editing Model](https://arxiv.org/abs/2608.14790)
+## What it does
 
-## How it works
+An image-editing DiT (Qwen-Image-Edit, 20.4B) is repurposed to edit **video latents**
+directly. The source clip is encoded by the frozen Wan 2.1 video VAE; two small trainable
+projections bridge the 16-channel video latents into and out of the DiT's token space, and
+the 12 latent frames are laid out as a 3×4 tile grid of one virtual image so the image
+model's 2-D RoPE covers the whole clip. Decoding the edited tokens back through the Wan VAE
+gives the edited video.
 
-```
-source video → (frozen Wan2.1 VAE) → video latents
-                    │
-           trainable in-projection
-                    ▼
-prompt + frame-grid preview → Qwen-Image-Edit DiT (LoRA) → grid RoPE over latent frames
-                                        │
-                              trainable out-projection
-                                        ▼
-                          edited latents → (frozen Wan2.1 VAE) → edited video
-```
+## What this Space runs
 
-## Credits
+This is **stage 1** of the authors' `infer.py`, matching the released 360P checkpoint's
+training configuration exactly (`--num_frames 45 --video_max_pixels 245760
+--latent_mode wan_compressed --pe_mode grid`), including the flow-matching schedule with
+`dynamic_shift_len`, the Qwen2.5-VL prompt embedding over a 3×3 preview grid of the clip,
+and true-CFG with the norm-preserving rescale.
 
-- Model: [yunpeng1998/Qwen-Video-Edit](https://huggingface.co/yunpeng1998/Qwen-Video-Edit)
-- Code: [yunpeng1998/Qwen-Video-Edit](https://github.com/yunpeng1998/Qwen-Video-Edit)
-- Built on [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio), [Qwen-Image-Edit](https://huggingface.co/Qwen/Qwen-Image-Edit), and [Wan 2.1](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B)
+Deliberate deviations, for ZeroGPU:
+
+* **One 45-frame (≈2.8 s) chunk per run**, at 360p — the window the checkpoint was
+  trained on. The reference script can loop over multiple chunks.
+* **The optional Wan 2.2 denoising-enhancement pass is not run** (it needs ~80 GB of extra
+  weights on top of the editing stack). This is equivalent to the reference's
+  `--skip_enhance`.
+* **Quantization**: the DiT transformer blocks run in fp8 (dynamic activation) and the
+  Qwen2.5-VL text encoder in int8 weight-only, via `torchao`. Entry/exit layers and both
+  trainable projections stay in bf16.
+* Uploads are temporally resampled to 16 fps (the training frame rate) before the
+  45-frame window is taken, so higher-fps clips are not played back in slow motion.
+* The default is 20 denoising steps (the reference default is 40) to keep a run inside a
+  reasonable ZeroGPU slice; raise it in **Advanced settings** for higher fidelity.
+
+## Example clips
+
+The bundled example videos come from
+[`linoyts/repo-to-space-example-videos`](https://huggingface.co/datasets/linoyts/repo-to-space-example-videos)
+(CC0-1.0), trimmed to the model's 45-frame window. The editing instructions are taken
+verbatim from the authors' own `examples/prompts.txt`.
+
+Model weights are MIT-licensed; the vendored DiffSynth-Studio code is Apache-2.0.

app.py

@@ -1,364 +1,416 @@
-"""Qwen-Video-Edit: Instruction-based video editing by repurposing an image editing model.
+"""Qwen-Video-Edit — instruction-based video editing on ZeroGPU.
 
-Gradio Space demo for yunpeng1998/Qwen-Video-Edit.
+Stage 1 of the official pipeline (`infer.py`), run on a single 45-frame chunk:
+
+    source video -> (frozen Wan 2.1 video VAE) -> wan latents
+                 -> trainable in-projection
+                 -> Qwen-Image-Edit DiT (full fine-tune, grid RoPE)
+                 -> trainable out-projection
+                 -> (frozen Wan 2.1 video VAE) -> edited video
+
+The optional Wan2.2 denoising-enhancement stage (~80GB of extra weights) is
+not run here — it does not fit alongside the editing stack on ZeroGPU.
 """
 
 import os
 
-os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+# DiffSynth's loader defaults to ModelScope; keep everything on the HF Hub.
 os.environ.setdefault("DIFFSYNTH_DOWNLOAD_SOURCE", "huggingface")
+os.environ.setdefault(
+    "DIFFSYNTH_MODEL_BASE_PATH",
+    os.path.join(os.path.dirname(os.path.abspath(__file__)), "models"),
+)
 
-import spaces  # MUST come before torch / any CUDA-touching import
-import functools
-import io
-import time
-import tempfile
-import math
-import gc
+import spaces  # noqa: E402  -- must precede any torch / CUDA-touching import
 
-import torch
-import numpy as np
-from PIL import Image
-import gradio as gr
+import random  # noqa: E402
+import tempfile  # noqa: E402
+import time  # noqa: E402
 
-print = functools.partial(print, flush=True)
+import gradio as gr  # noqa: E402
+import imageio.v3 as iio  # noqa: E402
+import numpy as np  # noqa: E402
+import torch  # noqa: E402
+from huggingface_hub import hf_hub_download  # noqa: E402
+from safetensors.torch import load_file  # noqa: E402
+from tqdm import tqdm  # noqa: E402
 
-# Apply rope patch guard (verifies vendored diffsynth is the patched copy)
-import rope_patch
-rope_patch.apply()
+import rope_patch  # noqa: E402
 
-from diffsynth.core import ModelConfig
-from diffsynth.pipelines.qwen_image import QwenImagePipeline, QwenImageUnit_PromptEmbedder
-from safetensors.torch import load_file
+rope_patch.apply()
 
-from dataset import (
-    adaptive_dims, build_preview_grid, decode_video_frames,
-    first_n_frames_padded, frames_to_tensor, open_maybe_remote,
-)
-from model import factorize_latent_grid, model_fn_video_tokens, num_token_groups
-from projections import (
-    WanToQwenProjection, QwenToWanProjection,
+from diffsynth.core import ModelConfig  # noqa: E402
+from diffsynth.core.vram.initialization import skip_model_initialization  # noqa: E402
+from diffsynth.models.qwen_image_dit import QwenImageDiT  # noqa: E402
+from diffsynth.pipelines.qwen_image import (  # noqa: E402
+    QwenImagePipeline,
+    QwenImageUnit_PromptEmbedder,
 )
 
-# ---- Model configuration (must match checkpoint training config) ----
-CHECKPOINT_REPO = "yunpeng1998/Qwen-Video-Edit"
-CHECKPOINT_PATH = "360P/step-30000.safetensors"
-NUM_FRAMES = 13  # Minimal frames to fit in VRAM (still satisfies (n-1)%4==0)
-VIDEO_MAX_PIXELS = 86400  # ~270p for VRAM
+from dataset import adaptive_dims, build_preview_grid, frames_to_tensor  # noqa: E402
+from model import factorize_latent_grid, model_fn_video_tokens, num_token_groups  # noqa: E402
+from projections import QwenToWanProjection, WanToQwenProjection  # noqa: E402
+
+
+# --------------------------------------------------------------------------
+# Checkpoint-pinned configuration (must match how the 360P model was trained;
+# see the repo README: --num_frames 45 --video_max_pixels 245760
+# --latent_mode wan_compressed --pe_mode grid, no --zero_cond_t).
+# --------------------------------------------------------------------------
+CKPT_REPO = "yunpeng1998/Qwen-Video-Edit"
+CKPT_FILE = "360P/step-30000.safetensors"
+NUM_FRAMES = 45
 LATENT_MODE = "wan_compressed"
 PE_MODE = "grid"
-ZERO_COND_T = True
-NUM_INFERENCE_STEPS = 30
-CFG_SCALE = 4.0
-DEFAULT_SEED = 42
+ZERO_COND_T = False
 FPS = 16
-
-DEVICE = "cuda"
-CPU = "cpu"
 DTYPE = torch.bfloat16
+MAX_SEED = np.iinfo(np.int32).max
 
-latent_grid = factorize_latent_grid(num_token_groups(NUM_FRAMES, LATENT_MODE))
-
-
-def load_checkpoint_into(pipe, in_proj, out_proj, checkpoint_path):
-    """Load the fine-tuned DiT weights + projections from a safetensors checkpoint."""
-    state_dict = load_file(checkpoint_path)
-    dit_sd = {k[len("pipe.dit."):]: v for k, v in state_dict.items() if k.startswith("pipe.dit.")}
-    in_proj.load_state_dict({k[len("in_proj."):]: v for k, v in state_dict.items() if k.startswith("in_proj.")})
-    out_proj.load_state_dict({k[len("out_proj."):]: v for k, v in state_dict.items() if k.startswith("out_proj.")})
-    if any("lora" in k for k in dit_sd):
-        pipe.load_lora(pipe.dit, state_dict=dit_sd, hotload=True)
-        print(f"[app] Loaded LoRA DiT weights ({len(dit_sd)} tensors) + projections.")
-    else:
-        pipe.dit.load_state_dict(dit_sd, strict=False)
-        print(f"[app] Loaded full DiT weights ({len(dit_sd)} tensors) + projections.")
-
-
-def tensor_to_uint8_frames(video):
-    """(C,T,H,W) [-1,1] -> list of HWC uint8 arrays."""
-    v = ((video.float().clamp(-1, 1) + 1) * 127.5).to(torch.uint8)
-    return [v[:, t].permute(1, 2, 0).cpu().numpy() for t in range(v.shape[1])]
-
-
-def save_video(frames_uint8, path_base, fps=16):
-    """Save frames as mp4."""
-    import imageio.v3 as iio
-    path = path_base + ".mp4"
-    iio.imwrite(path, frames_uint8, fps=fps, codec="libx264")
-    return path
-
-
-def free_cuda():
-    """Empty cache and collect garbage."""
-    torch.cuda.empty_cache()
-    gc.collect()
+DEFAULT_STEPS = 20
+DEFAULT_CFG = 4.0
+DEFAULT_NEG = " "
+DEFAULT_SEED = 42
+NATIVE_PIXELS = 245760  # 360p — the resolution the checkpoint was trained at
+DEFAULT_PIXELS = NATIVE_PIXELS
 
+# Calibrated against the live Space (see README): seconds of GPU time per DiT
+# forward at the native 360p token count, plus fixed VAE / text-encoder cost.
+SEC_PER_FORWARD = 3.0
+FIXED_OVERHEAD = 45.0
 
-# ---- Load model at module scope ----
-print("[app] Loading Qwen-Video-Edit pipeline "
-      "(downloading on first run, can take several minutes) ...")
-t0 = time.time()
+LATENT_GRID = factorize_latent_grid(num_token_groups(NUM_FRAMES, LATENT_MODE))
 
-# Load DiT (Qwen-Image-Edit transformer) — goes to CUDA for ZeroGPU packing
-model_configs = [
-    ModelConfig(
-        model_id="Qwen/Qwen-Image-Edit",
-        origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors",
-    ),
-]
-
-# Load text encoder separately — keep on CPU to save VRAM
-text_encoder_config = ModelConfig(
-    model_id="Qwen/Qwen-Image",
-    origin_file_pattern="text_encoder/model*.safetensors",
-)
 
+# --------------------------------------------------------------------------
+# Model loading (module scope, eagerly onto CUDA — ZeroGPU streams it in)
+# --------------------------------------------------------------------------
+print("[boot] Loading Qwen text encoder / tokenizer / processor ...", flush=True)
+_t0 = time.time()
 pipe = QwenImagePipeline.from_pretrained(
-    torch_dtype=DTYPE, device=DEVICE, model_configs=model_configs,
+    torch_dtype=DTYPE,
+    device="cpu",
+    model_configs=[
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="text_encoder/model*.safetensors"),
+    ],
     tokenizer_config=ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="tokenizer/"),
     processor_config=ModelConfig(model_id="Qwen/Qwen-Image-Edit", origin_file_pattern="processor/"),
 )
 
-# Load text encoder on CPU (NOT packed by ZeroGPU — saves ~17.5GB VRAM)
-text_encoder_configs = [text_encoder_config]
-text_encoder_pool = pipe.download_and_load_models(
-    text_encoder_configs, None
+print("[boot] Loading the Wan 2.1 video VAE ...", flush=True)
+_pool = pipe.download_and_load_models(
+    [ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="Wan2.1_VAE.pth")], None
 )
-pipe.text_encoder = text_encoder_pool.fetch_model("qwen_image_text_encoder")
-pipe.text_encoder.to(device=CPU)  # Keep on CPU
-
-# Load Wan 2.1 video VAE on CPU (loaded to CUDA only during encode/decode)
-wan_id = "Wan-AI/Wan2.1-T2V-1.3B"
-wan_pattern = "Wan2.1_VAE.pth"
-vae_pool = pipe.download_and_load_models(
-    [ModelConfig(model_id=wan_id, origin_file_pattern=wan_pattern)], None
+wan_vae = _pool.fetch_model("wan_video_vae")
+
+print(f"[boot] Downloading the Qwen-Video-Edit DiT checkpoint (~41GB) ...", flush=True)
+_ckpt_path = hf_hub_download(CKPT_REPO, CKPT_FILE)
+print("[boot] Loading the fine-tuned DiT + projections ...", flush=True)
+_state_dict = load_file(_ckpt_path, device="cpu")
+
+in_proj = WanToQwenProjection(16, 3072)
+out_proj = QwenToWanProjection(16, 3072)
+in_proj.load_state_dict({k[len("in_proj.") :]: v for k, v in _state_dict.items() if k.startswith("in_proj.")})
+out_proj.load_state_dict({k[len("out_proj.") :]: v for k, v in _state_dict.items() if k.startswith("out_proj.")})
+in_proj = in_proj.to(dtype=DTYPE).eval()
+out_proj = out_proj.to(dtype=DTYPE).eval()
+
+_dit_sd = {k[len("pipe.dit.") :]: v for k, v in _state_dict.items() if k.startswith("pipe.dit.")}
+with skip_model_initialization():
+    dit = QwenImageDiT()
+_missing, _unexpected = dit.load_state_dict(_dit_sd, strict=False, assign=True)
+_still_meta = [n for n, p in dit.named_parameters() if p.is_meta]
+if _still_meta:
+    raise RuntimeError(f"DiT parameters missing from the checkpoint: {_still_meta[:8]} ({len(_still_meta)} total)")
+print(
+    f"[boot] DiT loaded: {len(_dit_sd)} tensors, {sum(p.numel() for p in dit.parameters()) / 1e9:.2f}B params "
+    f"(missing={len(_missing)}, unexpected={len(_unexpected)})",
+    flush=True,
+)
+dit = dit.eval()
+del _state_dict, _dit_sd
+
+pipe.dit = dit
+pipe.to("cuda")
+wan_vae = wan_vae.to("cuda").eval()
+in_proj = in_proj.to("cuda")
+out_proj = out_proj.to("cuda")
+
+# The bf16 stack is ~57GB (20.4B DiT + 8.3B Qwen2.5-VL text encoder), which does
+# not fit a ZeroGPU slice. fp8 dynamic-activation quantization of the DiT blocks
+# and int8 weight-only on the text encoder bring it to ~29GB and roughly halve
+# the DiT matmul cost on Blackwell.
+from torchao.quantization import (  # noqa: E402
+    Float8DynamicActivationFloat8WeightConfig,
+    Int8WeightOnlyConfig,
+    quantize_,
 )
-wan_vae = vae_pool.fetch_model("wan_video_vae")
-wan_vae.to(device=CPU)  # Keep on CPU
 
-# Create projections
-in_proj = WanToQwenProjection(16, pipe.dit.img_in.out_features)
-out_proj = QwenToWanProjection(16, pipe.dit.img_in.out_features)
+print("[boot] Quantizing (fp8 DiT blocks / int8 text encoder) ...", flush=True)
+quantize_(dit.transformer_blocks, Float8DynamicActivationFloat8WeightConfig())
+quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
+print(f"[boot] Ready in {time.time() - _t0:.0f}s.", flush=True)
 
-# Download and load checkpoint
-from huggingface_hub import hf_hub_download
-checkpoint_path = hf_hub_download(CHECKPOINT_REPO, CHECKPOINT_PATH)
-load_checkpoint_into(pipe, in_proj, out_proj, checkpoint_path)
 
-in_proj.to(device=DEVICE, dtype=DTYPE)
-out_proj.to(device=DEVICE, dtype=DTYPE)
+# --------------------------------------------------------------------------
+# Helpers (ported 1:1 from the repo's infer.py / dataset.py)
+# --------------------------------------------------------------------------
+def tensor_to_uint8_frames(video):
+    """(C, T, H, W) in [-1, 1] -> list of HWC uint8 frames."""
+    v = ((video.float().clamp(-1, 1) + 1) * 127.5).to(torch.uint8)
+    return [v[:, t].permute(1, 2, 0).cpu().numpy() for t in range(v.shape[1])]
 
-print(f"[app] Pipeline ready in {time.time() - t0:.0f}s.")
 
+def read_chunk(video_path):
+    """Decode the source video and take the first NUM_FRAMES at ~16 fps.
 
-@spaces.GPU(duration=300, size="xlarge")
+    The checkpoint was trained on 16 fps clips, so a higher-fps upload is
+    temporally subsampled (rather than blindly taking 45 consecutive frames,
+    which would play back in slow motion at the fixed 16 fps output rate).
+    """
+    try:
+        meta = iio.immeta(video_path, plugin="pyav")
+        src_fps = float(meta.get("fps") or FPS)
+    except Exception:
+        src_fps = float(FPS)
+    if not np.isfinite(src_fps) or src_fps <= 0:
+        src_fps = float(FPS)
+
+    frames = [np.asarray(f) for f in iio.imiter(video_path, plugin="pyav")]
+    if not frames:
+        raise gr.Error("Could not decode any frames from that video.")
+    frames = [f[..., :3] if f.ndim == 3 else np.stack([f] * 3, -1) for f in frames]
+
+    idx = np.round(np.arange(NUM_FRAMES) * src_fps / FPS).astype(int)
+    idx = np.minimum(idx, len(frames) - 1)
+    return [frames[i] for i in idx]
+
+
+def _estimate_duration(*args, **kwargs):
+    def pick(name, pos, default):
+        if name in kwargs:
+            return kwargs[name]
+        return args[pos] if len(args) > pos else default
+
+    steps = int(pick("num_inference_steps", 2, DEFAULT_STEPS))
+    cfg = float(pick("cfg_scale", 3, DEFAULT_CFG))
+    pixels = float(pick("max_pixels", 7, DEFAULT_PIXELS))
+    forwards = steps * (2 if cfg > 1.0 else 1)
+    return int(FIXED_OVERHEAD + forwards * SEC_PER_FORWARD * (pixels / NATIVE_PIXELS))
+
+
+# --------------------------------------------------------------------------
+# Inference
+# --------------------------------------------------------------------------
+@spaces.GPU(duration=_estimate_duration)
 def edit_video(
-    source_video_path: str,
-    edit_instruction: str,
+    video_path: str,
+    instruction: str,
+    num_inference_steps: int = DEFAULT_STEPS,
+    cfg_scale: float = DEFAULT_CFG,
+    negative_prompt: str = DEFAULT_NEG,
     seed: int = DEFAULT_SEED,
-    num_inference_steps: int = NUM_INFERENCE_STEPS,
-    cfg_scale: float = CFG_SCALE,
-    negative_prompt: str = " ",
+    randomize_seed: bool = False,
+    max_pixels: int = DEFAULT_PIXELS,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Edit a source video based on a text instruction.
+    """Edit the first ~2.8 seconds of a video according to a text instruction.
 
     Args:
-        source_video_path: Path to the source video file (mp4).
-        edit_instruction: Text instruction describing how to edit the video.
-        seed: Random seed for reproducibility.
-        num_inference_steps: Number of denoising steps (more = higher quality, slower).
-        cfg_scale: Classifier-free guidance scale (1 = no CFG, 4 = default).
-        negative_prompt: Negative prompt for CFG guidance.
+        video_path: Path to the source video. Only the first 45 frames
+            (resampled to 16 fps) are edited.
+        instruction: Editing instruction, e.g. "Make it Ghibli style".
+        num_inference_steps: Flow-matching denoising steps.
+        cfg_scale: True-CFG scale. Values above 1 double the compute per step.
+        negative_prompt: Negative prompt used when cfg_scale > 1.
+        seed: Random seed for the initial noise.
+        randomize_seed: Draw a fresh random seed instead of using `seed`.
+        max_pixels: Pixel budget per frame (245760 = the native 360p training
+            resolution; lower is faster).
 
     Returns:
-        Path to the edited video file.
+        The edited video (mp4, 16 fps) and the seed that produced it.
     """
-    t_start = time.time()
-
-    # VRAM debug
-    if torch.cuda.is_available():
-        vram_total = torch.cuda.get_device_properties(0).total_memory / 1e9
-        vram_alloc = torch.cuda.memory_allocated(0) / 1e9
-        print(f"[app] VRAM: {vram_alloc:.1f} / {vram_total:.1f} GB allocated")
-
-    # Read and decode the source video
-    with open(source_video_path, "rb") as f:
-        all_frames = decode_video_frames(f.read())
-
-    # Use only the first NUM_FRAMES frames
-    chunk = first_n_frames_padded(all_frames[:NUM_FRAMES], NUM_FRAMES)
-    prompt = edit_instruction
+    if not video_path:
+        raise gr.Error("Please upload a source video.")
+    if not instruction or not instruction.strip():
+        raise gr.Error("Please write an editing instruction.")
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
+    num_inference_steps = int(num_inference_steps)
+    cfg_scale = float(cfg_scale)
+    max_pixels = int(max_pixels)
 
-    # Determine dimensions from max pixels
+    t_start = time.time()
+    chunk = read_chunk(video_path)
     h0, w0 = chunk[0].shape[:2]
-    width, height = adaptive_dims(w0, h0, VIDEO_MAX_PIXELS)
-
-    # Prepare source video tensor
+    width, height = adaptive_dims(w0, h0, max_pixels)
     source = frames_to_tensor(chunk, width, height)
     preview = build_preview_grid(chunk)
+    print(f"[run] {width}x{height}x{NUM_FRAMES} | steps={num_inference_steps} cfg={cfg_scale} seed={seed}", flush=True)
 
-    # --- Phase 1: VAE encode (load VAE to CUDA, then free) ---
-    wan_vae.to(device=DEVICE, dtype=DTYPE)
-    tiled = height * width >= 700_000
-    ref_latents = wan_vae.encode(
-        [source.to(dtype=DTYPE)], device=DEVICE, tiled=tiled
-    ).to(device=DEVICE, dtype=DTYPE)
-    wan_vae.to(device=CPU)
-    free_cuda()
-
-    # --- Phase 2: Prompt embedding (load text encoder to CUDA, then free) ---
-    pipe.text_encoder.to(device=DEVICE, dtype=DTYPE)
-    emb = QwenImageUnit_PromptEmbedder().process(pipe, prompt=prompt, edit_image=preview)
-    use_cfg = cfg_scale > 1.0
-    if use_cfg:
-        neg_emb = QwenImageUnit_PromptEmbedder().process(
-            pipe, prompt=negative_prompt, edit_image=preview
+    with torch.no_grad():
+        tiled = height * width >= 700_000
+        ref_latents = wan_vae.encode([source.to(dtype=DTYPE)], device="cuda", tiled=tiled).to(
+            device="cuda", dtype=DTYPE
         )
-    pipe.text_encoder.to(device=CPU)
-    free_cuda()
-
-    # --- Phase 3: Denoising loop (only DiT + projections in VRAM) ---
-    if torch.cuda.is_available():
-        vram_alloc = torch.cuda.memory_allocated(0) / 1e9
-        vram_total = torch.cuda.get_device_properties(0).total_memory / 1e9
-        print(f"[app] Before denoising: VRAM {vram_alloc:.1f} / {vram_total:.1f} GB")
-    group = getattr(in_proj, "group", 1)
-    noise_seq_len = (
-        (ref_latents.shape[2] // group)
-        * (ref_latents.shape[3] // 2)
-        * (ref_latents.shape[4] // 2)
-    )
-    pipe.scheduler.set_timesteps(
-        num_inference_steps, dynamic_shift_len=noise_seq_len
-    )
-    gen = torch.Generator(device="cpu").manual_seed(seed)
-    latents = torch.randn(ref_latents.shape, generator=gen).to(
-        device=DEVICE, dtype=DTYPE
-    )
 
-    def forward(e, timestep):
-        return model_fn_video_tokens(
-            pipe.dit, in_proj, out_proj,
-            latents=latents, ref_latents=ref_latents,
-            prompt_emb=e["prompt_emb"], prompt_emb_mask=e["prompt_emb_mask"],
-            timestep=timestep, latent_grid=latent_grid,
-            zero_cond_t=ZERO_COND_T, pe_mode=PE_MODE,
-        )
+        embedder = QwenImageUnit_PromptEmbedder()
+        emb = embedder.process(pipe, prompt=instruction, edit_image=preview)
+        use_cfg = cfg_scale > 1.0
+        neg_emb = embedder.process(pipe, prompt=negative_prompt, edit_image=preview) if use_cfg else None
+
+        group = getattr(in_proj, "group", 1)
+        noise_seq_len = (ref_latents.shape[2] // group) * (ref_latents.shape[3] // 2) * (ref_latents.shape[4] // 2)
+        pipe.scheduler.set_timesteps(num_inference_steps, dynamic_shift_len=noise_seq_len)
+        gen = torch.Generator(device="cpu").manual_seed(seed)
+        latents = torch.randn(ref_latents.shape, generator=gen).to(device="cuda", dtype=DTYPE)
+
+        def forward(e, timestep):
+            return model_fn_video_tokens(
+                dit,
+                in_proj,
+                out_proj,
+                latents=latents,
+                ref_latents=ref_latents,
+                prompt_emb=e["prompt_emb"],
+                prompt_emb_mask=e["prompt_emb_mask"],
+                timestep=timestep,
+                latent_grid=LATENT_GRID,
+                zero_cond_t=ZERO_COND_T,
+                pe_mode=PE_MODE,
+            )
 
-    from tqdm import tqdm
-    for pid, timestep in enumerate(tqdm(pipe.scheduler.timesteps, desc="denoising")):
-        timestep = timestep.unsqueeze(0).to(dtype=DTYPE, device=DEVICE)
-        try:
+        t_denoise = time.time()
+        for pid, timestep in enumerate(tqdm(pipe.scheduler.timesteps, desc="denoise")):
+            timestep = timestep.unsqueeze(0).to(dtype=DTYPE, device="cuda")
             noise_pred = forward(emb, timestep)
             if use_cfg:
+                # True CFG with a norm-preserving rescale over the latent
+                # channel dim (exactly as in the repo's infer.py).
                 neg_pred = forward(neg_emb, timestep)
                 comb = neg_pred + cfg_scale * (noise_pred - neg_pred)
-                comb = comb * (
-                    torch.norm(noise_pred, dim=1, keepdim=True)
-                    / torch.norm(comb, dim=1, keepdim=True)
-                )
+                comb = comb * (torch.norm(noise_pred, dim=1, keepdim=True) / torch.norm(comb, dim=1, keepdim=True))
                 noise_pred = comb
-            latents = pipe.scheduler.step(
-                noise_pred, pipe.scheduler.timesteps[pid], latents
-            )
-        except RuntimeError as e:
-            if "NVML_SUCCESS" in str(e) or "out of memory" in str(e).lower():
-                return None, f"GPU out of memory. The Qwen-Video-Edit DiT (~42GB) is very large. Try reducing inference steps or using a smaller video."
-            raise
-
-    # --- Phase 4: VAE decode (load VAE back to CUDA) ---
-    wan_vae.to(device=DEVICE, dtype=DTYPE)
-    tiled_dec = (latents.shape[3] * 8) * (latents.shape[4] * 8) >= 700_000
-    video = wan_vae.decode(latents, device=DEVICE, tiled=tiled_dec)[0].cpu()
-    wan_vae.to(device=CPU)
-    free_cuda()
-
-    # Save to temporary file
-    tmp_dir = tempfile.mkdtemp()
-    out_path = os.path.join(tmp_dir, "edited_video")
-    save_video(tensor_to_uint8_frames(video), out_path, fps=FPS)
-
-    elapsed = time.time() - t_start
-    print(f"[app] Video edit completed in {elapsed:.1f}s")
-
-    return out_path + ".mp4", f"Edit completed in {elapsed:.1f}s"
-
+            latents = pipe.scheduler.step(noise_pred, pipe.scheduler.timesteps[pid], latents)
+        denoise_s = time.time() - t_denoise
+
+        tiled = (latents.shape[3] * 8) * (latents.shape[4] * 8) >= 700_000
+        video = wan_vae.decode(latents, device="cuda", tiled=tiled)[0].cpu()
+
+    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    iio.imwrite(out_path, np.stack(tensor_to_uint8_frames(video)), fps=FPS, codec="libx264")
+    total = time.time() - t_start
+    forwards = num_inference_steps * (2 if cfg_scale > 1.0 else 1)
+    print(
+        f"[run] done in {total:.1f}s (denoise {denoise_s:.1f}s over {forwards} forwards "
+        f"= {denoise_s / max(forwards, 1):.2f}s/forward)",
+        flush=True,
+    )
+    return out_path, seed
+
+
+def run_example(video_path: str, instruction: str):
+    """Run an example row: edit a bundled clip with its instruction, all other
+    settings left at their defaults."""
+    return edit_video(video_path, instruction)
+
+
+# --------------------------------------------------------------------------
+# UI
+# --------------------------------------------------------------------------
+EXAMPLES = [
+    ["examples/man_dancing.mp4", "Turn the dancer into a glowing neon hologram"],
+    [
+        "examples/ocean_waves.mp4",
+        "Convert the video into a soft-focus, impressionist painting with brushstrokes "
+        "mimicking the movement of the waves and the warmth of the sunset.",
+    ],
+    ["examples/pottery_wheel.mp4", "Imitate the look of the Ghibli style."],
+    [
+        "examples/city_traffic_night.mp4",
+        "Convert the urban landscape into a retro-futuristic cityscape rendered in "
+        "neon-drenched 1980s synthwave style, with glowing rooftops, floating vehicles, "
+        "and chromatic halos around buildings.",
+    ],
+]
 
-# ---- Gradio UI ----
 CSS = """
 #col-container { max-width: 1100px; margin: 0 auto; }
-.dark .gradio-container { color: var(--body-text-color); }
 """
 
 with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("# 🎬 Qwen-Video-Edit")
         gr.Markdown(
-            "Instruction-based video editing by repurposing an image editing model. "
-            "Upload a video, describe how to edit it, and get the edited result.\n\n"
-            "[Paper](https://arxiv.org/abs/2608.14790) | "
-            "[Model](https://huggingface.co/yunpeng1998/Qwen-Video-Edit) | "
-            "[Code](https://github.com/yunpeng1998/Qwen-Video-Edit)"
-        )
+            """
+# 🎬 Qwen-Video-Edit
+
+**Instruction-based video editing by repurposing an image editing model.**
+Qwen-Image-Edit's DiT edits Wan 2.1 video-VAE latents directly, bridged by two tiny
+projections warm-started from the DiT's own input/output layers — no video-pretrained
+transformer required.
+
+[Paper](https://huggingface.co/papers/2608.14790) ·
+[Project page](https://yunpeng1998.github.io/Qwen-Video-Edit-Page) ·
+[Code](https://github.com/yunpeng1998/Qwen-Video-Edit) ·
+[Weights](https://huggingface.co/yunpeng1998/Qwen-Video-Edit)
 
+The model edits one **45-frame (≈2.8 s) chunk at 360p** per run — that is the window the
+released checkpoint was trained on.
+"""
+        )
         with gr.Row():
-            with gr.Column(scale=1):
-                source_video = gr.Video(
-                    label="Source Video",
-                    sources=["upload"],
-                )
-                edit_instruction = gr.Textbox(
-                    label="Edit Instruction",
-                    placeholder="e.g., Replace the woman with robot",
+            with gr.Column():
+                video_in = gr.Video(label="Source video", height=340)
+                instruction = gr.Textbox(
+                    label="Editing instruction",
+                    placeholder="e.g. Make it Ghibli style.",
                     lines=2,
                 )
-                run_btn = gr.Button("Edit Video", variant="primary")
-            with gr.Column(scale=1):
-                output_video = gr.Video(label="Edited Video")
-                status_text = gr.Textbox(label="Status", interactive=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            seed = gr.Number(label="Seed", value=DEFAULT_SEED, precision=0)
-            num_steps = gr.Slider(
-                label="Inference Steps",
-                minimum=10,
-                maximum=80,
-                value=NUM_INFERENCE_STEPS,
-                step=1,
-            )
-            cfg = gr.Slider(
-                label="CFG Scale",
-                minimum=1.0,
-                maximum=10.0,
-                value=CFG_SCALE,
-                step=0.5,
-            )
-            negative = gr.Textbox(
-                label="Negative Prompt",
-                value=" ",
-                lines=1,
+                run_btn = gr.Button("Edit video", variant="primary")
+            with gr.Column():
+                video_out = gr.Video(label="Edited video", height=340, autoplay=True, loop=True)
+                seed_out = gr.Number(label="Seed used", interactive=False)
+
+        with gr.Accordion("Advanced settings", open=False):
+            with gr.Row():
+                steps = gr.Slider(
+                    label="Denoising steps", minimum=8, maximum=40, step=1, value=DEFAULT_STEPS
+                )
+                cfg = gr.Slider(
+                    label="True-CFG scale (>1 doubles the compute)",
+                    minimum=1.0,
+                    maximum=7.0,
+                    step=0.1,
+                    value=DEFAULT_CFG,
+                )
+            negative = gr.Textbox(label="Negative prompt", value=DEFAULT_NEG)
+            with gr.Row():
+                seed_in = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=DEFAULT_SEED)
+                randomize = gr.Checkbox(label="Randomize seed", value=False)
+            pixels = gr.Radio(
+                label="Resolution",
+                choices=[("360p — native training resolution", NATIVE_PIXELS), ("256p — faster", 122880)],
+                value=DEFAULT_PIXELS,
             )
 
         gr.Examples(
-            examples=[
-                ["example1.mp4", "Replace the woman with robot"],
-                ["example2.mp4", "Replace the woman with an animated-style Captain Marvel"],
-                ["example1.mp4", "Make it watercolor drawing style."],
-                ["example2.mp4", "Convert to van Gogh style"],
-            ],
-            inputs=[source_video, edit_instruction],
-            outputs=[output_video, status_text],
-            fn=edit_video,
-            cache_examples=False,
-            run_on_click=True,
+            examples=EXAMPLES,
+            inputs=[video_in, instruction],
+            outputs=[video_out, seed_out],
+            fn=run_example,
+            cache_examples=True,
+            cache_mode="lazy",
+            label="Examples (instructions from the authors' prompt list)",
         )
 
-    run_btn.click(
+    gr.on(
+        triggers=[run_btn.click, instruction.submit],
         fn=edit_video,
-        inputs=[source_video, edit_instruction, seed, num_steps, cfg, negative],
-        outputs=[output_video, status_text],
-        api_name="edit_video",
+        inputs=[video_in, instruction, steps, cfg, negative, seed_in, randomize, pixels],
+        outputs=[video_out, seed_out],
     )
 
-demo.launch(mcp_server=True)
+demo.launch(mcp_server=True)

diffsynth/core/loader/model.py

@@ -20,7 +20,7 @@ def load_model(model_class, path, config=None, torch_dtype=torch.bfloat16, devic
         dtypes = [vram_config["offload_dtype"], vram_config["onload_dtype"], vram_config["preparing_dtype"], vram_config["computation_dtype"]]
         dtype = [d for d in dtypes if d != "disk"][0]
         if vram_config["offload_device"] != "disk":
-            if state_dict is None: state_dict = DiskMap(path, "cpu", torch_dtype=dtype)
+            if state_dict is None: state_dict = DiskMap(path, device, torch_dtype=dtype)
             if state_dict_converter is not None:
                 state_dict = state_dict_converter(state_dict)
             else:
@@ -28,7 +28,7 @@ def load_model(model_class, path, config=None, torch_dtype=torch.bfloat16, devic
             model.load_state_dict(state_dict, assign=True)
             model = enable_vram_management(model, module_map, vram_config=vram_config, disk_map=None, vram_limit=vram_limit)
         else:
-            disk_map = DiskMap(path, "cpu", state_dict_converter=state_dict_converter)
+            disk_map = DiskMap(path, device, state_dict_converter=state_dict_converter)
             model = enable_vram_management(model, module_map, vram_config=vram_config, disk_map=disk_map, vram_limit=vram_limit)
     else:
         # Why do we use `DiskMap`?
@@ -38,11 +38,9 @@ def load_model(model_class, path, config=None, torch_dtype=torch.bfloat16, devic
         if state_dict is not None:
             pass
         elif use_disk_map:
-            # On ZeroGPU, always load to CPU first — the .to(device) below
-            # is intercepted by the spaces hijack to pack weights to disk.
-            state_dict = DiskMap(path, "cpu", torch_dtype=torch_dtype)
+            state_dict = DiskMap(path, device, torch_dtype=torch_dtype)
         else:
-            state_dict = load_state_dict(path, torch_dtype, "cpu")
+            state_dict = load_state_dict(path, torch_dtype, device)
         # Why do we use `state_dict_converter`?
         # Some models are saved in complex formats,
         # and we need to convert the state dict into the appropriate format.
@@ -79,7 +77,7 @@ def load_model_with_disk_offload(model_class, path, config=None, torch_dtype=tor
         model = model_class(**config)
     if hasattr(model, "eval"):
         model = model.eval()
-    disk_map = DiskMap(path, "cpu", state_dict_converter=state_dict_converter)
+    disk_map = DiskMap(path, device, state_dict_converter=state_dict_converter)
     vram_config = {
         "offload_dtype": "disk",
         "offload_device": "disk",

diffsynth/core/vram/disk_map.py

@@ -47,13 +47,13 @@ class DiskMap:
         if len(self.files) == 0:
             for path in self.path:
                 if path.endswith(".safetensors"):
-                    self.files.append(safe_open(path, framework="pt", device="cpu"))
+                    self.files.append(safe_open(path, framework="pt", device=str(self.device)))
                 else:
-                    self.files.append(SafetensorsCompatibleBinaryLoader(path, device="cpu"))
+                    self.files.append(SafetensorsCompatibleBinaryLoader(path, device=self.device))
         else:
             for i, path in enumerate(self.path):
                 if path.endswith(".safetensors"):
-                    self.files[i] = safe_open(path, framework="pt", device="cpu")
+                    self.files[i] = safe_open(path, framework="pt", device=str(self.device))
         self.num_params = 0
 
     def __getitem__(self, name):

diffsynth/models/qwen_image_dit.py

@@ -22,12 +22,6 @@ try:
 except ModuleNotFoundError:
     FLASH_ATTN_3_AVAILABLE = False
 
-try:
-    from flash_attn import flash_attn_func as flash_attn_2_func
-    FLASH_ATTN_2_AVAILABLE = True
-except ModuleNotFoundError:
-    FLASH_ATTN_2_AVAILABLE = False
-
 
 def qwen_image_flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, attention_mask = None, enable_fp8_attention: bool = False):
     if FLASH_ATTN_3_AVAILABLE and attention_mask is None:
@@ -51,14 +45,6 @@ def qwen_image_flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
                 x = x[0]
             x = x.to(origin_dtype) * v_std
             x = rearrange(x, "b s n d -> b s (n d)", n=num_heads)
-    elif FLASH_ATTN_2_AVAILABLE and attention_mask is None:
-        q = rearrange(q, "b n s d -> b s n d", n=num_heads)
-        k = rearrange(k, "b n s d -> b s n d", n=num_heads)
-        v = rearrange(v, "b n s d -> b s n d", n=num_heads)
-        x = flash_attn_2_func(q, k, v)
-        if isinstance(x, tuple):
-            x = x[0]
-        x = rearrange(x, "b s n d -> b s (n d)", n=num_heads)
     else:
         x = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask)
         x = rearrange(x, "b n s d -> b s (n d)", n=num_heads)

examples/city_traffic_night.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed5ce2a6c4c7b951ff34d5d9ae282080c2f19d3864b5efa4ccc79b1ce9092b2c
+size 335551

examples/man_dancing.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6061764e56dd8cb6220894c254a639f0547869ea2609a9c971ee7dc9940fcc10
+size 333527

examples/ocean_waves.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:477dfe9c68d1ba5525320c07e597963d657ed14471ece6343fc667f865c209a5
+size 1366054

examples/pottery_wheel.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4fbb21011083c2a5e7f417dcce6aa92a54d124b2d4e2d926476875df6d4850fa
+size 380123

requirements.txt

@@ -1,7 +1,8 @@
-torch
 torchvision
+torchao
 transformers
 accelerate
+peft
 safetensors
 einops
 imageio
@@ -9,9 +10,9 @@ imageio[ffmpeg]
 av
 numpy
 Pillow
-modelscope
+pandas
 tqdm
+modelscope
+ftfy
+regex
 sentencepiece
-diffusers
-peft
-https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt211-cu130-cp312/flash_attn-2.8.3-cp312-cp312-linux_x86_64.whl