Add JoyAI-Video-Edit pipeline

docs/source/en/_toctree.yml

@@ -361,6 +361,8 @@
         title: JoyImageEditPlusTransformer3DModel
       - local: api/models/transformer_joyimage
         title: JoyImageEditTransformer3DModel
+      - local: api/models/transformer_joyvideoedit
+        title: JoyVideoEditTransformer3DModel
       - local: api/models/krea2_transformer2d
         title: Krea2Transformer2DModel
       - local: api/models/latte_transformer3d
@@ -453,6 +455,8 @@
         title: AutoencoderKLHunyuanVideo
       - local: api/models/autoencoder_kl_hunyuan_video15
         title: AutoencoderKLHunyuanVideo15
+      - local: api/models/autoencoder_kl_joyvideoedit
+        title: AutoencoderKLJoyVideoEdit
       - local: api/models/autoencoder_kl_kvae
         title: AutoencoderKLKVAE
       - local: api/models/autoencoder_kl_kvae_video
@@ -575,6 +579,8 @@
         title: JoyImage Edit
       - local: api/pipelines/joyimage_edit_plus
         title: JoyImage Edit Plus
+      - local: api/pipelines/joyvideoedit
+        title: JoyVideo Edit
       - local: api/pipelines/kandinsky
         title: Kandinsky 2.1
       - local: api/pipelines/kandinsky_v22

docs/source/en/api/models/autoencoder_kl_joyvideoedit.md

@@ -0,0 +1,37 @@
+<!-- Copyright 2026 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License. -->
+
+# AutoencoderKLJoyVideoEdit
+
+The causal, chunk-streamable 3D variational autoencoder (VAE) model with KL loss used in [`JoyVideoEditPipeline`]. It
+encodes and decodes video in temporal chunks so that arbitrarily long sequences can be processed with bounded memory.
+
+The model can be loaded with the following code snippet.
+
+```python
+import torch
+
+from diffusers import AutoencoderKLJoyVideoEdit
+
+vae = AutoencoderKLJoyVideoEdit.from_pretrained(
+    "jdopensource/JoyAI-Video-Edit-Diffusers", subfolder="vae", dtype=torch.float32
+)
+```
+
+## AutoencoderKLJoyVideoEdit
+
+[[autodoc]] AutoencoderKLJoyVideoEdit
+  - decode
+  - all
+
+## DecoderOutput
+
+[[autodoc]] models.autoencoders.vae.DecoderOutput

docs/source/en/api/models/transformer_joyvideoedit.md

@@ -0,0 +1,31 @@
+<!--Copyright 2026 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+# JoyVideoEditTransformer3DModel
+
+A dual-stream MM-DiT transformer that denoises video latents one causal chunk at a time, used in
+[`JoyVideoEditPipeline`]. The model can be loaded with the following code snippet.
+
+```python
+import torch
+from diffusers import JoyVideoEditTransformer3DModel
+
+transformer = JoyVideoEditTransformer3DModel.from_pretrained("jdopensource/JoyAI-Video-Edit-Diffusers", subfolder="transformer", dtype=torch.bfloat16)
+```
+
+## JoyVideoEditTransformer3DModel
+
+[[autodoc]] JoyVideoEditTransformer3DModel
+
+## Transformer2DModelOutput
+
+[[autodoc]] models.modeling_outputs.Transformer2DModelOutput

docs/source/en/api/pipelines/joyvideoedit.md

@@ -0,0 +1,84 @@
+<!--Copyright 2026 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+# JoyAI-Video-Edit
+
+[JoyAI-Video-Edit](https://github.com/jd-opensource/JoyAI-Video-Edit) is an instruction-guided video-editing model built on the JoyAI streaming architecture. The source video is VAE-encoded into a latent sequence that conditions a dual-stream MM-DiT transformer, which denoises the edited output one causal chunk at a time. Each chunk attends to a sliding window of previously-denoised chunks (and an optional static reference image) through a per-layer KV cache, keeping later chunks temporally consistent with earlier ones without recomputing their key/value projections.
+
+| Model | Description | Download |
+|:-----:|:-----------:|:--------:|
+| JoyAI-Video-Edit | Instruction-guided causal video editing | [Hugging Face](https://huggingface.co/jdopensource/JoyAI-Video-Edit-Diffusers) |
+
+```python
+import torch
+from diffusers import JoyVideoEditPipeline
+from diffusers.utils import export_to_video, load_video
+from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
+
+mimo_id = "XiaomiMiMo/MiMo-VL-7B-RL-2508"
+processor = AutoProcessor.from_pretrained(mimo_id)
+text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    mimo_id, dtype=torch.bfloat16
+)
+
+pipeline = JoyVideoEditPipeline.from_pretrained(
+    "jdopensource/JoyAI-Video-Edit-Diffusers",
+    text_encoder=text_encoder,
+    processor=processor,
+    dtype=torch.bfloat16,
+)
+pipeline.enable_model_cpu_offload()
+
+video = load_video("https://raw.githubusercontent.com/jd-opensource/JoyAI-Video-Edit/main/assets/input.mp4")
+prompt = (
+    "Transform the scene into a British castle royal aristocratic style. Modify the characters' clothing to "
+    "aristocratic attire: dress the man in a tailored velvet suit with a ruffled cravat, and the women in elegant "
+    "silk gowns with lace details and embroidered bodices. Change their hairstyles to classic aristocratic styles, "
+    "such as elaborate updos with subtle jewels for the women and a neatly styled classic cut for the man. Change "
+    "the environmental decoration to a British castle interior: replace the plain walls and abstract painting with "
+    "stone walls and antique oil paintings in gilded frames, and replace the white window curtains with heavy velvet "
+    "drapes. The characters' ages and facial features must remain completely unchanged. The dining table, white "
+    "tablecloth, plates of food, wine glasses, water glasses, and the characters' positions and actions must remain "
+    "unchanged."
+)
+
+output = pipeline(
+    video=video,
+    prompt=prompt,
+    num_inference_steps=2,
+    generator=torch.Generator(device="cpu").manual_seed(0),
+)
+export_to_video(output.frames[0], "joyvideoedit_output.mp4", fps=24)
+```
+
+The pipeline denoises with a flow-matching scheduler and does not use classifier-free guidance, so it takes neither a
+`negative_prompt` nor a `guidance_scale` argument. An optional static reference image can be supplied through
+`ref_image`; its KV is prefilled into the cache and attended to by every chunk to inject appearance conditioning.
+
+The Diffusers checkpoint does not include MiMo-VL. Load [`XiaomiMiMo/MiMo-VL-7B-RL-2508`](https://huggingface.co/XiaomiMiMo/MiMo-VL-7B-RL-2508) from MiMo-VL's own repository
+and pass its model and processor to [`JoyVideoEditPipeline.from_pretrained`]. The tokenizer from the processor is used
+when a separate `tokenizer` is not provided. You can omit all MiMo-VL components when passing precomputed
+`prompt_embeds` and `prompt_embeds_mask`.
+
+Model CPU offloading is recommended because the transformer, MiMo-VL, and VAE are otherwise resident on the GPU at
+the same time. The pipeline also supports sequential CPU offloading for lower memory use and pipeline-level group
+offloading for a balance between transfer overhead and memory use.
+
+## JoyVideoEditPipeline
+
+[[autodoc]] JoyVideoEditPipeline
+  - all
+  - __call__
+
+## JoyVideoEditPipelineOutput
+
+[[autodoc]] pipelines.joyvideoedit.pipeline_output.JoyVideoEditPipelineOutput

scripts/convert_joyvideoedit_to_diffusers.py

@@ -0,0 +1,177 @@
+"""Convert JoyVideoEdit (JoyAI-Video-Edit) checkpoints to diffusers format.
+
+Converts the transformer and/or the VAE. The transformer checkpoint is a raw `.pth` state dict whose double-block
+attention keys need remapping under `.attn.`. The VAE checkpoint uses diffusers-format safetensors and is re-saved with
+the `AutoencoderKLJoyVideoEdit` configuration.
+
+Usage:
+
+```bash
+python scripts/convert_joyvideoedit_to_diffusers.py \
+    --transformer_ckpt_path /path/to/joyai_video_edit_dit_0804.pth \
+    --vae_dir /path/to/JoyAI-Video-Edit/vae \
+    --output_path /path/to/output
+```
+"""
+
+import argparse
+import json
+import os
+
+import torch
+from accelerate import init_empty_weights
+from safetensors.torch import load_file
+
+from diffusers import (
+    AutoencoderKLJoyVideoEdit,
+    FlowMatchEulerDiscreteScheduler,
+    JoyVideoEditPipeline,
+    JoyVideoEditTransformer3DModel,
+)
+
+
+TRANSFORMER_CONFIG = {
+    "patch_size": [1, 1, 1],
+    "in_channels": 64,
+    "out_channels": 64,
+    "hidden_size": 4096,
+    "num_attention_heads": 32,
+    "text_dim": 4096,
+    "num_layers": 40,
+    "rope_dim_list": [16, 56, 56],
+    "theta": 256,
+    "chunk_size": 1,
+    "local_window_size": 3,
+    "global_sink_chunk": True,
+    "source_id_rope_dim": 128,
+    "source_id_rope_theta": 256.0,
+}
+
+
+def convert_transformer(ckpt_path: str) -> JoyVideoEditTransformer3DModel:
+    checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+    original_state_dict = checkpoint["model"] if "model" in checkpoint else checkpoint
+
+    attn_suffixes = (
+        "img_attn_qkv.",
+        "img_attn_q_norm.",
+        "img_attn_k_norm.",
+        "img_attn_proj.",
+        "txt_attn_qkv.",
+        "txt_attn_q_norm.",
+        "txt_attn_k_norm.",
+        "txt_attn_proj.",
+    )
+    remapped = {}
+    for key, value in original_state_dict.items():
+        new_key = key
+        if key.startswith("double_blocks."):
+            for suffix in attn_suffixes:
+                if "." + suffix in key and ".attn." + suffix not in key:
+                    new_key = key.replace("." + suffix, ".attn." + suffix)
+                    break
+        remapped[new_key] = value
+
+    with init_empty_weights():
+        transformer = JoyVideoEditTransformer3DModel(**TRANSFORMER_CONFIG)
+    transformer.load_state_dict(remapped, strict=True, assign=True)
+    return transformer
+
+
+def convert_vae(vae_dir: str) -> AutoencoderKLJoyVideoEdit:
+    with open(os.path.join(vae_dir, "config.json")) as f:
+        config = json.load(f)
+    config = {k: v for k, v in config.items() if not k.startswith("_")}
+
+    state_dict = load_file(os.path.join(vae_dir, "diffusion_pytorch_model.safetensors"))
+
+    with init_empty_weights():
+        vae = AutoencoderKLJoyVideoEdit(**config)
+    vae.load_state_dict(state_dict, strict=True, assign=True)
+    return vae
+
+
+DTYPE_MAPPING = {
+    "fp32": torch.float32,
+    "fp16": torch.float16,
+    "bf16": torch.bfloat16,
+}
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="Convert JoyVideoEdit checkpoints to diffusers format")
+    parser.add_argument(
+        "--transformer_ckpt_path",
+        type=str,
+        default=None,
+        help="Path to the transformer checkpoint (e.g. joyai_video_edit_dit_0804.pth)",
+    )
+    parser.add_argument(
+        "--vae_dir",
+        type=str,
+        default=None,
+        help="Path to the VAE directory (with config.json + diffusion_pytorch_model.safetensors)",
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        required=True,
+        help=(
+            "Output directory. Saves a complete pipeline when both checkpoints are provided, or an individual "
+            "transformer/ or vae/ subdirectory otherwise."
+        ),
+    )
+    parser.add_argument("--dtype", choices=tuple(DTYPE_MAPPING), default="bf16", help="Torch dtype")
+    return parser.parse_args()
+
+
+def set_model_dtype(model: torch.nn.Module, dtype: torch.dtype) -> torch.nn.Module:
+    torch.nn.Module.to(model, dtype=dtype)
+    keep_in_fp32_modules = getattr(model, "_keep_in_fp32_modules", None) or []
+    for module_name, module in model.named_modules():
+        if any(pattern in module_name.split(".") for pattern in keep_in_fp32_modules):
+            torch.nn.Module.to(module, dtype=torch.float32)
+    return model
+
+
+def save_joyvideoedit_pipeline(
+    transformer: JoyVideoEditTransformer3DModel,
+    vae: AutoencoderKLJoyVideoEdit,
+    output_path: str,
+) -> None:
+    pipeline = JoyVideoEditPipeline(
+        transformer=transformer,
+        vae=vae,
+        text_encoder=None,
+        tokenizer=None,
+        processor=None,
+        scheduler=FlowMatchEulerDiscreteScheduler(),
+    )
+    pipeline.save_pretrained(output_path, safe_serialization=True, max_shard_size="5GB")
+
+
+if __name__ == "__main__":
+    args = get_args()
+    dtype = DTYPE_MAPPING[args.dtype]
+
+    transformer = None
+    vae = None
+
+    if args.transformer_ckpt_path is not None:
+        transformer = convert_transformer(args.transformer_ckpt_path)
+        transformer = set_model_dtype(transformer, dtype)
+
+    if args.vae_dir is not None:
+        vae = convert_vae(args.vae_dir)
+        vae = set_model_dtype(vae, dtype)
+
+    if transformer is not None and vae is not None:
+        save_joyvideoedit_pipeline(transformer, vae, args.output_path)
+    elif transformer is not None:
+        transformer.save_pretrained(
+            os.path.join(args.output_path, "transformer"), safe_serialization=True, max_shard_size="5GB"
+        )
+    elif vae is not None:
+        vae.save_pretrained(os.path.join(args.output_path, "vae"), safe_serialization=True)
+    else:
+        raise ValueError("Provide at least one of `--transformer_ckpt_path` or `--vae_dir`.")