From 198ce00a6ad3ede6a2a5c76ef063ea0d4ce5e953 Mon Sep 17 00:00:00 2001 From: huangfeice Date: Wed, 12 Aug 2026 10:58:26 +0800 Subject: [PATCH] Add JoyAI-Video-Edit pipeline Integrate JoyAI-Video-Edit with a causal streaming VAE, dual-stream 3D transformer, MiMo-VL prompt and image conditioning, chunk-wise KV caching, optional reference-image conditioning, and memory-efficient chunked decoding. Add checkpoint conversion with the model mixed-precision policy, lazy imports, API documentation, and comprehensive model and pipeline tests covering serialization, compilation, batching, callbacks, cache lifecycle, and CPU and group offloading. --- docs/source/en/_toctree.yml | 6 + .../api/models/autoencoder_kl_joyvideoedit.md | 37 + .../en/api/models/transformer_joyvideoedit.md | 31 + docs/source/en/api/pipelines/joyvideoedit.md | 84 ++ scripts/convert_joyvideoedit_to_diffusers.py | 177 ++++ src/diffusers/__init__.py | 12 + src/diffusers/hooks/__init__.py | 1 + src/diffusers/hooks/joyvideoedit_kv_cache.py | 74 ++ src/diffusers/models/__init__.py | 4 + src/diffusers/models/autoencoders/__init__.py | 1 + .../autoencoder_kl_joyvideoedit.py | 872 ++++++++++++++++ src/diffusers/models/cache_utils.py | 11 +- src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_joyvideoedit.py | 967 ++++++++++++++++++ src/diffusers/pipelines/__init__.py | 5 + .../pipelines/joyvideoedit/__init__.py | 48 + .../joyvideoedit/pipeline_joyvideoedit.py | 824 +++++++++++++++ .../pipelines/joyvideoedit/pipeline_output.py | 20 + src/diffusers/utils/dummy_pt_objects.py | 49 + .../dummy_torch_and_transformers_objects.py | 30 + ...test_models_autoencoder_kl_joyvideoedit.py | 137 +++ .../test_models_transformer_joyvideoedit.py | 143 +++ tests/pipelines/joyvideoedit/__init__.py | 0 .../test_pipeline_joyvideoedit.py | 429 ++++++++ 24 files changed, 3962 insertions(+), 1 deletion(-) create mode 100644 docs/source/en/api/models/autoencoder_kl_joyvideoedit.md create mode 100644 docs/source/en/api/models/transformer_joyvideoedit.md create mode 100644 docs/source/en/api/pipelines/joyvideoedit.md create mode 100644 scripts/convert_joyvideoedit_to_diffusers.py create mode 100644 src/diffusers/hooks/joyvideoedit_kv_cache.py create mode 100644 src/diffusers/models/autoencoders/autoencoder_kl_joyvideoedit.py create mode 100644 src/diffusers/models/transformers/transformer_joyvideoedit.py create mode 100644 src/diffusers/pipelines/joyvideoedit/__init__.py create mode 100644 src/diffusers/pipelines/joyvideoedit/pipeline_joyvideoedit.py create mode 100644 src/diffusers/pipelines/joyvideoedit/pipeline_output.py create mode 100644 tests/models/autoencoders/test_models_autoencoder_kl_joyvideoedit.py create mode 100644 tests/models/transformers/test_models_transformer_joyvideoedit.py create mode 100644 tests/pipelines/joyvideoedit/__init__.py create mode 100644 tests/pipelines/joyvideoedit/test_pipeline_joyvideoedit.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index f929a9840afb..e370ba25d427 100644 --- a/docs/source/en/_toctree.yml +++ b/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 diff --git a/docs/source/en/api/models/autoencoder_kl_joyvideoedit.md b/docs/source/en/api/models/autoencoder_kl_joyvideoedit.md new file mode 100644 index 000000000000..06c4807a79a9 --- /dev/null +++ b/docs/source/en/api/models/autoencoder_kl_joyvideoedit.md @@ -0,0 +1,37 @@ + + +# 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 diff --git a/docs/source/en/api/models/transformer_joyvideoedit.md b/docs/source/en/api/models/transformer_joyvideoedit.md new file mode 100644 index 000000000000..3f6e80f1f445 --- /dev/null +++ b/docs/source/en/api/models/transformer_joyvideoedit.md @@ -0,0 +1,31 @@ + + +# 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 diff --git a/docs/source/en/api/pipelines/joyvideoedit.md b/docs/source/en/api/pipelines/joyvideoedit.md new file mode 100644 index 000000000000..5d4e0c23e0ad --- /dev/null +++ b/docs/source/en/api/pipelines/joyvideoedit.md @@ -0,0 +1,84 @@ + + +# 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 diff --git a/scripts/convert_joyvideoedit_to_diffusers.py b/scripts/convert_joyvideoedit_to_diffusers.py new file mode 100644 index 000000000000..36b79a9ee1c1 --- /dev/null +++ b/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`.") diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 4bec0f5bd7ff..5f32744af561 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -200,6 +200,7 @@ "FasterCacheConfig", "FirstBlockCacheConfig", "HookRegistry", + "JoyVideoEditKVCacheConfig", "LayerSkipConfig", "MagCacheConfig", "PyramidAttentionBroadcastConfig", @@ -208,6 +209,7 @@ "TextKVCacheConfig", "apply_faster_cache", "apply_first_block_cache", + "apply_joyvideoedit_kv_cache", "apply_layer_skip", "apply_mag_cache", "apply_pyramid_attention_broadcast", @@ -242,6 +244,7 @@ "AutoencoderKLHunyuanImageRefiner", "AutoencoderKLHunyuanVideo", "AutoencoderKLHunyuanVideo15", + "AutoencoderKLJoyVideoEdit", "AutoencoderKLKVAE", "AutoencoderKLKVAEVideo", "AutoencoderKLLTX2Audio", @@ -300,6 +303,7 @@ "Ideogram4Transformer2DModel", "JoyImageEditPlusTransformer3DModel", "JoyImageEditTransformer3DModel", + "JoyVideoEditTransformer3DModel", "Kandinsky3UNet", "Kandinsky5Transformer3DModel", "Krea2Transformer2DModel", @@ -685,6 +689,8 @@ "JoyImageEditPipelineOutput", "JoyImageEditPlusPipeline", "JoyImageEditPlusPipelineOutput", + "JoyVideoEditPipeline", + "JoyVideoEditPipelineOutput", "Kandinsky3Img2ImgPipeline", "Kandinsky3Pipeline", "Kandinsky5I2IPipeline", @@ -1066,6 +1072,7 @@ FasterCacheConfig, FirstBlockCacheConfig, HookRegistry, + JoyVideoEditKVCacheConfig, LayerSkipConfig, MagCacheConfig, PyramidAttentionBroadcastConfig, @@ -1074,6 +1081,7 @@ TextKVCacheConfig, apply_faster_cache, apply_first_block_cache, + apply_joyvideoedit_kv_cache, apply_layer_skip, apply_mag_cache, apply_pyramid_attention_broadcast, @@ -1106,6 +1114,7 @@ AutoencoderKLHunyuanImageRefiner, AutoencoderKLHunyuanVideo, AutoencoderKLHunyuanVideo15, + AutoencoderKLJoyVideoEdit, AutoencoderKLKVAE, AutoencoderKLKVAEVideo, AutoencoderKLLTX2Audio, @@ -1164,6 +1173,7 @@ Ideogram4Transformer2DModel, JoyImageEditPlusTransformer3DModel, JoyImageEditTransformer3DModel, + JoyVideoEditTransformer3DModel, Kandinsky3UNet, Kandinsky5Transformer3DModel, Krea2Transformer2DModel, @@ -1524,6 +1534,8 @@ JoyImageEditPipelineOutput, JoyImageEditPlusPipeline, JoyImageEditPlusPipelineOutput, + JoyVideoEditPipeline, + JoyVideoEditPipelineOutput, Kandinsky3Img2ImgPipeline, Kandinsky3Pipeline, Kandinsky5I2IPipeline, diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py index 2a9aa81608e7..269be1f4a30e 100644 --- a/src/diffusers/hooks/__init__.py +++ b/src/diffusers/hooks/__init__.py @@ -21,6 +21,7 @@ from .first_block_cache import FirstBlockCacheConfig, apply_first_block_cache from .group_offloading import apply_group_offloading from .hooks import HookRegistry, ModelHook + from .joyvideoedit_kv_cache import JoyVideoEditKVCacheConfig, apply_joyvideoedit_kv_cache from .layer_skip import LayerSkipConfig, apply_layer_skip from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook from .mag_cache import MagCacheConfig, apply_mag_cache diff --git a/src/diffusers/hooks/joyvideoedit_kv_cache.py b/src/diffusers/hooks/joyvideoedit_kv_cache.py new file mode 100644 index 000000000000..26d3e13a4ccc --- /dev/null +++ b/src/diffusers/hooks/joyvideoedit_kv_cache.py @@ -0,0 +1,74 @@ +# 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. + +from dataclasses import dataclass +from typing import Dict + +import torch + +from .hooks import BaseState, HookRegistry, ModelHook, StateManager + + +_JOYVIDEOEDIT_KV_CACHE_HOOK = "joyvideoedit_kv_cache" + + +@dataclass +class JoyVideoEditKVCacheConfig: + """Enable the chunk-wise streaming KV cache used by `JoyVideoEditTransformer3DModel`. + + Chunks of clean video are fed through once (`kv_cache_mode="store"`) to populate a per-layer cache, then reused as + attention context for later chunks (`kv_cache_mode="reuse"`) without recomputing their key/value projections. + """ + + pass + + +class JoyVideoEditKVCacheState(BaseState): + """Holds the per-chunk-per-layer KV cache.""" + + def __init__(self): + self.chunk_cache: Dict[int, Dict[int, Dict[str, torch.Tensor]]] = {} + + def reset(self): + self.chunk_cache.clear() + + +class JoyVideoEditKVCacheHook(ModelHook): + """Routes `JoyVideoEditTransformer3DModel`'s KV-cache reads/writes through a `StateManager`-managed state. + + The hook owns the cache state and ensures a context is active before `forward` runs. Cache selection, assembly, and + eviction are handled by the model. + """ + + _is_stateful = True + + def __init__(self, state_manager: StateManager): + super().__init__() + self.state_manager = state_manager + + def new_forward(self, module: torch.nn.Module, *args, **kwargs): + if self.state_manager._current_context is None: + self.state_manager.set_context("inference") + return self.fn_ref.original_forward(*args, **kwargs) + + def reset_state(self, module: torch.nn.Module): + self.state_manager.reset() + return module + + +def apply_joyvideoedit_kv_cache(module: torch.nn.Module, config: JoyVideoEditKVCacheConfig) -> None: + registry = HookRegistry.check_if_exists_or_initialize(module) + state_manager = StateManager(JoyVideoEditKVCacheState) + hook = JoyVideoEditKVCacheHook(state_manager) + registry.register_hook(hook, _JOYVIDEOEDIT_KV_CACHE_HOOK) diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index d7f2c75c68e7..861a4dcf0e9e 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -40,6 +40,7 @@ _import_structure["autoencoders.autoencoder_kl_hunyuanimage"] = ["AutoencoderKLHunyuanImage"] _import_structure["autoencoders.autoencoder_kl_hunyuanimage_refiner"] = ["AutoencoderKLHunyuanImageRefiner"] _import_structure["autoencoders.autoencoder_kl_hunyuanvideo15"] = ["AutoencoderKLHunyuanVideo15"] + _import_structure["autoencoders.autoencoder_kl_joyvideoedit"] = ["AutoencoderKLJoyVideoEdit"] _import_structure["autoencoders.autoencoder_kl_kvae"] = ["AutoencoderKLKVAE"] _import_structure["autoencoders.autoencoder_kl_kvae_video"] = ["AutoencoderKLKVAEVideo"] _import_structure["autoencoders.autoencoder_kl_ltx"] = ["AutoencoderKLLTXVideo"] @@ -127,6 +128,7 @@ _import_structure["transformers.transformer_ideogram4"] = ["Ideogram4Transformer2DModel"] _import_structure["transformers.transformer_joyimage"] = ["JoyImageEditTransformer3DModel"] _import_structure["transformers.transformer_joyimage_edit_plus"] = ["JoyImageEditPlusTransformer3DModel"] + _import_structure["transformers.transformer_joyvideoedit"] = ["JoyVideoEditTransformer3DModel"] _import_structure["transformers.transformer_kandinsky"] = ["Kandinsky5Transformer3DModel"] _import_structure["transformers.transformer_krea2"] = ["Krea2Transformer2DModel"] _import_structure["transformers.transformer_longcat_audio_dit"] = ["LongCatAudioDiTTransformer"] @@ -183,6 +185,7 @@ AutoencoderKLHunyuanImageRefiner, AutoencoderKLHunyuanVideo, AutoencoderKLHunyuanVideo15, + AutoencoderKLJoyVideoEdit, AutoencoderKLKVAE, AutoencoderKLKVAEVideo, AutoencoderKLLTX2Audio, @@ -264,6 +267,7 @@ Ideogram4Transformer2DModel, JoyImageEditPlusTransformer3DModel, JoyImageEditTransformer3DModel, + JoyVideoEditTransformer3DModel, Kandinsky5Transformer3DModel, Krea2Transformer2DModel, LatteTransformer3DModel, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index 7d24611c825e..d5f441adc4d7 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -10,6 +10,7 @@ from .autoencoder_kl_hunyuanimage import AutoencoderKLHunyuanImage from .autoencoder_kl_hunyuanimage_refiner import AutoencoderKLHunyuanImageRefiner from .autoencoder_kl_hunyuanvideo15 import AutoencoderKLHunyuanVideo15 +from .autoencoder_kl_joyvideoedit import AutoencoderKLJoyVideoEdit from .autoencoder_kl_kvae import AutoencoderKLKVAE from .autoencoder_kl_kvae_video import AutoencoderKLKVAEVideo from .autoencoder_kl_ltx import AutoencoderKLLTXVideo diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_joyvideoedit.py b/src/diffusers/models/autoencoders/autoencoder_kl_joyvideoedit.py new file mode 100644 index 000000000000..458338fe2751 --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_kl_joyvideoedit.py @@ -0,0 +1,872 @@ +# Copyright 2026 The JoyAI-Video-Edit Team and 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. +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils.accelerate_utils import apply_forward_hook +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..modeling_outputs import AutoencoderKLOutput +from ..modeling_utils import ModelMixin +from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution + + +CACHE_T = 1 + + +class JoyVideoEditRMSNorm(nn.Module): + r""" + Channel-first RMS normalization (no learnable bias) used throughout the JoyVideoEdit VAE. + """ + + def __init__(self, dim: int) -> None: + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim, 1, 1, 1)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return F.normalize(hidden_states, dim=1) * self.scale * self.gamma + + +class JoyVideoEditCausalConv3d(nn.Conv3d): + r""" + A 3D convolution that is causal in the temporal dimension by construction of its *input*, rather than by padding + zeros at the front like a standard causal conv. + + Every forward call appends a duplicated copy of the last input frame at the end of the temporal axis before + convolving (`torch.cat([front, hidden_states, hidden_states[:, :, -1:, :, :]], dim=2)`), where `front` is either + the last `CACHE_T` frame(s) of the previous chunk (`cache_x`) or, for the very first chunk, the tensor itself + (which must have a single temporal frame in that case). Padding is applied block-internally rather than as a single + zero-pad at the very start of the sequence. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int, int], + stride: int | tuple[int, int, int] = 1, + padding: int | tuple[int, int, int] = 0, + ) -> None: + super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding) + assert self.padding[0] == 1, "Causal padding only supports padding of 1 in the temporal dimension." + self._padding = (self.padding[2], self.padding[2], self.padding[1], self.padding[1], 0, 0) + self.padding = (0, 0, 0) + + def forward(self, hidden_states: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor: + if cache_x is not None: + front = cache_x.to(hidden_states.device) + else: + assert hidden_states.shape[2] == 1, ( + f"Input temporal dimension is expected to be 1 when cache_x is None, got {hidden_states.shape}." + ) + front = hidden_states + hidden_states = torch.cat([front, hidden_states, hidden_states[:, :, -1:, :, :]], dim=2) + hidden_states = F.pad(hidden_states, self._padding) + return super().forward(hidden_states) + + +class JoyVideoEditResidualBlock(nn.Module): + def __init__(self, channels: int) -> None: + super().__init__() + self.norm1 = JoyVideoEditRMSNorm(channels) + self.conv1 = JoyVideoEditCausalConv3d(channels, channels, kernel_size=3, stride=1, padding=1) + self.norm2 = JoyVideoEditRMSNorm(channels) + self.conv2 = JoyVideoEditCausalConv3d(channels, channels, kernel_size=3, stride=1, padding=1) + + def forward(self, hidden_states: torch.Tensor, feat_cache: list, feat_idx: list[int]) -> torch.Tensor: + residual = hidden_states + + hidden_states = self.norm1(hidden_states) + hidden_states = F.silu(hidden_states) + idx = feat_idx[0] + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + hidden_states = self.conv1(hidden_states, cache_x=feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + hidden_states = self.norm2(hidden_states) + hidden_states = F.silu(hidden_states) + idx = feat_idx[0] + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + hidden_states = self.conv2(hidden_states, cache_x=feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + return hidden_states + residual + + +class JoyVideoEditVAEAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__(self, attn: "JoyVideoEditAttentionBlock", hidden_states: torch.Tensor) -> torch.Tensor: + identity = hidden_states + batch_size, channels, num_frames, height, width = hidden_states.shape + + hidden_states = attn.norm(hidden_states) + query = attn.q(hidden_states) + key = attn.k(hidden_states) + value = attn.v(hidden_states) + + # "b c t h w -> (b t) (h w) 1 c" + query = query.permute(0, 2, 3, 4, 1).contiguous().reshape(batch_size * num_frames, height * width, 1, channels) + key = key.permute(0, 2, 3, 4, 1).contiguous().reshape(batch_size * num_frames, height * width, 1, channels) + value = value.permute(0, 2, 3, 4, 1).contiguous().reshape(batch_size * num_frames, height * width, 1, channels) + + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=None, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + + # "(b t) (h w) 1 c -> b c t h w" + hidden_states = ( + hidden_states.reshape(batch_size, num_frames, height, width, channels).permute(0, 4, 1, 2, 3).contiguous() + ) + + hidden_states = attn.proj_out(hidden_states) + return identity + hidden_states + + +class JoyVideoEditAttentionBlock(nn.Module, AttentionModuleMixin): + r""" + Single-head spatial self-attention applied independently to every frame, with separate query, key, value, and + output projections. + """ + + _default_processor_cls = JoyVideoEditVAEAttnProcessor + _available_processors = [JoyVideoEditVAEAttnProcessor] + _supports_qkv_fusion = False + + def __init__(self, in_channels: int) -> None: + super().__init__() + self.in_channels = in_channels + + self.norm = JoyVideoEditRMSNorm(in_channels) + self.q = nn.Conv3d(in_channels, in_channels, kernel_size=1) + self.k = nn.Conv3d(in_channels, in_channels, kernel_size=1) + self.v = nn.Conv3d(in_channels, in_channels, kernel_size=1) + self.proj_out = nn.Conv3d(in_channels, in_channels, kernel_size=1) + self.set_processor(self._default_processor_cls()) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.processor(self, hidden_states) + + +class JoyVideoEditDownsampleBlock(nn.Module): + r""" + Causal-conv downsample followed by a space-to-depth fold: the conv reduces the channel count to `out_channels // + factor`, then `factor = 8` (temporal) or `4` (spatial-only) neighboring positions are folded into the channel + dimension so the output has `out_channels` channels at `1/factor` the number of positions. A mean-pooled version of + the *input* (folded the same way) is added back as a residual shortcut. + """ + + def __init__(self, in_channels: int, out_channels: int, temporal_downsample: bool) -> None: + super().__init__() + factor = 8 if temporal_downsample else 4 + self.conv = JoyVideoEditCausalConv3d(in_channels, out_channels // factor, kernel_size=3, stride=1, padding=1) + + self.temporal_downsample = temporal_downsample + self.group_size = factor * in_channels // out_channels + + @staticmethod + def _space_to_depth(hidden_states: torch.Tensor, factor_t: int) -> torch.Tensor: + # einops: "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=factor_t, r2=2, r3=2 + batch_size, channels, num_frames, height, width = hidden_states.shape + num_frames, height, width = num_frames // factor_t, height // 2, width // 2 + hidden_states = hidden_states.reshape(batch_size, channels, num_frames, factor_t, height, 2, width, 2) + hidden_states = hidden_states.permute(0, 3, 5, 7, 1, 2, 4, 6).contiguous() + return hidden_states.reshape(batch_size, factor_t * 4 * channels, num_frames, height, width) + + def forward( + self, + hidden_states: torch.Tensor, + feat_cache: list, + feat_idx: list[int], + first_chunk: bool = False, + ) -> torch.Tensor: + factor_t = 2 if self.temporal_downsample else 1 + + if self.temporal_downsample and first_chunk: + shortcut = torch.cat([hidden_states[:, :, :1, :, :], hidden_states], dim=2) + else: + shortcut = hidden_states + shortcut = self._space_to_depth(shortcut, factor_t) + + idx = feat_idx[0] + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + hidden_states = self.conv(hidden_states, cache_x=feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + if self.temporal_downsample and first_chunk: + hidden_states = torch.cat([hidden_states[:, :, :1, :, :], hidden_states], dim=2) + hidden_states = self._space_to_depth(hidden_states, factor_t) + + batch_size, channels, num_frames, height, width = shortcut.shape + shortcut = shortcut.view(batch_size, hidden_states.shape[1], self.group_size, num_frames, height, width).mean( + dim=2 + ) + return hidden_states + shortcut + + +class JoyVideoEditUpsampleBlock(nn.Module): + r""" + Causal-conv upsample followed by a depth-to-space unfold: the conv raises the channel count to `out_channels * + factor`, then unfolded into `factor = 8` (temporal) or `4` (spatial-only) neighboring positions. A + `repeat_interleave`-based upsample of the input (unfolded the same way) is added back as a residual shortcut. + """ + + def __init__(self, in_channels: int, out_channels: int, temporal_upsample: bool) -> None: + super().__init__() + factor = 8 if temporal_upsample else 4 + self.conv = JoyVideoEditCausalConv3d(in_channels, out_channels * factor, kernel_size=3, stride=1, padding=1) + + self.temporal_upsample = temporal_upsample + self.repeats = factor * out_channels // in_channels + + @staticmethod + def _depth_to_space(hidden_states: torch.Tensor, factor_t: int) -> torch.Tensor: + # einops: "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=factor_t, r2=2, r3=2 + batch_size, folded_channels, num_frames, height, width = hidden_states.shape + channels = folded_channels // (factor_t * 4) + hidden_states = hidden_states.reshape(batch_size, factor_t, 2, 2, channels, num_frames, height, width) + hidden_states = hidden_states.permute(0, 4, 5, 1, 6, 2, 7, 3).contiguous() + return hidden_states.reshape(batch_size, channels, num_frames * factor_t, height * 2, width * 2) + + def forward( + self, + hidden_states: torch.Tensor, + feat_cache: list, + feat_idx: list[int], + first_chunk: bool = False, + ) -> torch.Tensor: + factor_t = 2 if self.temporal_upsample else 1 + + shortcut = hidden_states.repeat_interleave(repeats=self.repeats, dim=1) + shortcut = self._depth_to_space(shortcut, factor_t) + + idx = feat_idx[0] + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + hidden_states = self.conv(hidden_states, cache_x=feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + hidden_states = self._depth_to_space(hidden_states, factor_t) + + hidden_states = hidden_states + shortcut + if self.temporal_upsample and first_chunk: + hidden_states = hidden_states[:, :, 1:, :, :] + + return hidden_states + + +class JoyVideoEditEncoder(nn.Module): + def __init__( + self, + in_channels: int, + z_channels: int, + num_res_blocks: int, + block_in_channels: tuple[int, ...], + temporal_downsample: tuple[bool, ...], + ) -> None: + super().__init__() + + self.conv_in = JoyVideoEditCausalConv3d(in_channels, block_in_channels[0], kernel_size=3, stride=1, padding=1) + + self.down_blocks = nn.ModuleList([]) + for i_level, block_in in enumerate(block_in_channels): + for _ in range(num_res_blocks): + self.down_blocks.append(JoyVideoEditResidualBlock(channels=block_in)) + + if i_level != len(block_in_channels) - 1: + block_out = block_in_channels[i_level + 1] + self.down_blocks.append(JoyVideoEditDownsampleBlock(block_in, block_out, temporal_downsample[i_level])) + + self.mid_blocks = nn.ModuleList( + [ + JoyVideoEditResidualBlock(channels=block_in), + JoyVideoEditAttentionBlock(block_in), + JoyVideoEditResidualBlock(channels=block_in), + ] + ) + + self.norm_out = JoyVideoEditRMSNorm(block_in) + self.conv_out = JoyVideoEditCausalConv3d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) + + def forward( + self, + hidden_states: torch.Tensor, + feat_cache: list, + feat_idx: list[int], + first_chunk: bool = False, + ) -> torch.Tensor: + idx = feat_idx[0] + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + hidden_states = self.conv_in(hidden_states, cache_x=feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + for block in self.down_blocks: + if isinstance(block, JoyVideoEditDownsampleBlock): + hidden_states = block(hidden_states, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk) + else: + hidden_states = block(hidden_states, feat_cache=feat_cache, feat_idx=feat_idx) + for block in self.mid_blocks: + if isinstance(block, JoyVideoEditResidualBlock): + hidden_states = block(hidden_states, feat_cache=feat_cache, feat_idx=feat_idx) + else: + hidden_states = block(hidden_states) + + hidden_states = self.norm_out(hidden_states) + hidden_states = F.silu(hidden_states) + idx = feat_idx[0] + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + hidden_states = self.conv_out(hidden_states, cache_x=feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return hidden_states + + +class JoyVideoEditDecoder(nn.Module): + def __init__( + self, + z_channels: int, + out_channels: int, + num_res_blocks: int, + block_in_channels: tuple[int, ...], + temporal_upsample: tuple[bool, ...], + ) -> None: + super().__init__() + + block_in = block_in_channels[0] + self.conv_in = JoyVideoEditCausalConv3d(z_channels, block_in, kernel_size=3, stride=1, padding=1) + + self.mid_blocks = nn.ModuleList( + [ + JoyVideoEditResidualBlock(channels=block_in), + JoyVideoEditAttentionBlock(block_in), + JoyVideoEditResidualBlock(channels=block_in), + ] + ) + + self.up_blocks = nn.ModuleList([]) + for i_level, block_in in enumerate(block_in_channels): + for _ in range(num_res_blocks + 1): + self.up_blocks.append(JoyVideoEditResidualBlock(channels=block_in)) + + if i_level != len(block_in_channels) - 1: + block_out = block_in_channels[i_level + 1] + self.up_blocks.append(JoyVideoEditUpsampleBlock(block_in, block_out, temporal_upsample[i_level])) + + self.norm_out = JoyVideoEditRMSNorm(block_in) + self.conv_out = JoyVideoEditCausalConv3d(block_in, out_channels, kernel_size=3, stride=1, padding=1) + + def forward( + self, + hidden_states: torch.Tensor, + feat_cache: list, + feat_idx: list[int], + first_chunk: bool = False, + ) -> torch.Tensor: + idx = feat_idx[0] + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + hidden_states = self.conv_in(hidden_states, cache_x=feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + for block in self.mid_blocks: + if isinstance(block, JoyVideoEditResidualBlock): + hidden_states = block(hidden_states, feat_cache=feat_cache, feat_idx=feat_idx) + else: + hidden_states = block(hidden_states) + + for block in self.up_blocks: + if isinstance(block, JoyVideoEditUpsampleBlock): + hidden_states = block(hidden_states, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk) + else: + hidden_states = block(hidden_states, feat_cache=feat_cache, feat_idx=feat_idx) + + hidden_states = self.norm_out(hidden_states) + hidden_states = F.silu(hidden_states) + idx = feat_idx[0] + cache_x = hidden_states[:, :, -CACHE_T:, :, :].clone() + hidden_states = self.conv_out(hidden_states, cache_x=feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return hidden_states + + +class JoyVideoEditStem(nn.Module): + """Extra 3/2 spatial down-projection applied to pixels before the encoder backbone. + + `pixel_unshuffle(stride)` folds a `stride x stride` neighborhood into channels, a 1x1 conv reprojects, then + `pixel_shuffle(group)` unfolds a `group x group` block back to space, giving a net spatial factor of `group / + stride = 2 / 3`. Combined with the backbone's `patch_size * 2 ** (len(temporal_downsample) - 1) = 16`, the VAE + reaches an effective spatial compression of 24. Frames are folded into the batch dim so the 2D shuffles act + per-frame. + """ + + def __init__(self, channels: int, stride: int = 3, group: int = 2) -> None: + super().__init__() + self.stride = stride + self.group = group + self.proj = nn.Conv2d(channels * stride * stride, channels * group * group, kernel_size=1, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + batch_size, channels, num_frames, height, width = x.shape + out_height, out_width = height * self.group // self.stride, width * self.group // self.stride + z = x.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) + z = F.pixel_unshuffle(z, self.stride) + z = self.proj(z) + z = F.pixel_shuffle(z, self.group) + return z.reshape(batch_size, num_frames, channels, out_height, out_width).permute(0, 2, 1, 3, 4) + + +class JoyVideoEditHeadResBlock(nn.Module): + """Depthwise-separable residual block used inside `JoyVideoEditHead`.""" + + def __init__(self, channels: int) -> None: + super().__init__() + self.dw = nn.Conv2d(channels, channels, kernel_size=3, padding=1, groups=channels) + self.pw = nn.Conv2d(channels, channels, kernel_size=1) + self.act = nn.ReLU(inplace=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.pw(self.act(self.dw(x))) + + +class JoyVideoEditHead(nn.Module): + """Learned 3/2 spatial up-projection applied after the decoder backbone, inverting `JoyVideoEditStem`. + + Bilinearly upsamples the decoded pixels by `scale` and adds a learned residual refinement. Frames are folded into + the batch dim so the 2D convolutions act per-frame. + """ + + def __init__( + self, channels: int, scale: float = 1.5, hidden: int = 32, num_blocks: int = 4, mid_channels: int = 12 + ) -> None: + super().__init__() + self.scale = float(scale) + self.conv_in = nn.Conv2d(channels, hidden, kernel_size=3, padding=1) + self.act = nn.ReLU(inplace=False) + self.blocks = nn.Sequential(*[JoyVideoEditHeadResBlock(hidden) for _ in range(num_blocks)]) + self.reduce = nn.Conv2d(hidden, mid_channels, kernel_size=3, padding=1) + self.conv_out = nn.Conv2d(mid_channels, channels, kernel_size=3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + batch_size, channels, num_frames, height, width = x.shape + out_height, out_width = round(height * self.scale), round(width * self.scale) + z = x.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) + f = self.act(self.conv_in(z)) + f = self.blocks(f) + f = self.reduce(f) + f = F.interpolate(f, size=(out_height, out_width), mode="bilinear", align_corners=False) + residual = self.conv_out(f) + base = F.interpolate(z, size=(out_height, out_width), mode="bilinear", align_corners=False) + out = (base + residual).clamp(-1.0, 1.0) + return out.reshape(batch_size, num_frames, channels, out_height, out_width).permute(0, 2, 1, 3, 4) + + +def patchify(hidden_states: torch.Tensor, patch_size: int) -> torch.Tensor: + if patch_size == 1: + return hidden_states + # einops: "b c t (h r1) (w r2) -> b (c r1 r2) t h w", r1=patch_size, r2=patch_size + batch_size, channels, num_frames, height, width = hidden_states.shape + height, width = height // patch_size, width // patch_size + hidden_states = hidden_states.reshape(batch_size, channels, num_frames, height, patch_size, width, patch_size) + hidden_states = hidden_states.permute(0, 1, 4, 6, 2, 3, 5).contiguous() + return hidden_states.reshape(batch_size, channels * patch_size * patch_size, num_frames, height, width) + + +def unpatchify(hidden_states: torch.Tensor, patch_size: int) -> torch.Tensor: + if patch_size == 1: + return hidden_states + # einops: "b (r1 r2 c) t h w -> b c t (h r1) (w r2)", r1=patch_size, r2=patch_size + batch_size, folded_channels, num_frames, height, width = hidden_states.shape + channels = folded_channels // (patch_size * patch_size) + hidden_states = hidden_states.reshape(batch_size, patch_size, patch_size, channels, num_frames, height, width) + hidden_states = hidden_states.permute(0, 3, 4, 5, 1, 6, 2).contiguous() + return hidden_states.reshape(batch_size, channels, num_frames, height * patch_size, width * patch_size) + + +class AutoencoderKLJoyVideoEdit(ModelMixin, AttentionMixin, AutoencoderMixin, ConfigMixin): + r""" + A causal, chunk-streamable VAE with KL loss for encoding videos into latents and decoding latent representations + into videos, used by the JoyVideoEdit pipeline. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + """ + + _supports_gradient_checkpointing = False + _repeated_blocks = None + _group_offload_block_modules = ["stem", "encoder", "decoder", "head"] + # keys to ignore when AlignDeviceHook moves inputs/outputs between devices + # these are shared mutable state modified in-place + _skip_keys = ["feat_cache", "feat_idx"] + + @register_to_config + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + patch_size: int = 2, + latent_channels: int = 64, + layers_per_block: int = 2, + block_in_channels: tuple[int, ...] = (128, 256, 512, 1024), + temporal_downsample: tuple[bool, ...] = (True, True, True, False), + chunk_size: int = 48, + latents_mean: list[float] = [ + 0.003708, + 0.018799, + -0.049072, + -1.171875, + 0.064453, + 0.648438, + -0.507812, + 0.030273, + -0.090332, + 0.10498, + -0.18457, + 0.667969, + -0.863281, + -0.12793, + 0.000511, + 0.472656, + -0.636719, + 0.761719, + 0.170898, + -0.482422, + 0.267578, + 0.092285, + -0.066406, + -0.002029, + 0.201172, + 0.026489, + -0.073242, + 0.016479, + -0.449219, + 0.070312, + -0.423828, + 0.804688, + -1.773438, + -0.117676, + 0.010986, + -0.092285, + -0.003448, + -0.133789, + -0.230469, + -0.410156, + -0.292969, + 0.414062, + -0.150391, + -0.045654, + -0.213867, + -0.126953, + -0.062012, + -1.039062, + 0.058838, + -0.015442, + -0.054932, + 0.100098, + -0.112793, + 0.0177, + 0.213867, + -0.003906, + 0.172852, + 0.003281, + -0.257812, + 0.010071, + 0.008362, + -0.163086, + 0.126953, + -1.34375, + ], + latents_std: list[float] = [ + 0.5625, + 1.710938, + 0.695312, + 2.453125, + 0.769531, + 3.265625, + 3.140625, + 0.835938, + 0.570312, + 0.757812, + 0.925781, + 2.046875, + 2.171875, + 0.503906, + 1.53125, + 1.03125, + 1.90625, + 2.375, + 0.5625, + 0.964844, + 0.699219, + 0.648438, + 3.890625, + 0.707031, + 2.265625, + 0.878906, + 0.550781, + 0.451172, + 2.46875, + 0.53125, + 1.914062, + 3.234375, + 4.65625, + 1.1875, + 0.65625, + 0.738281, + 0.851562, + 0.71875, + 0.796875, + 2.78125, + 1.445312, + 0.589844, + 0.535156, + 0.628906, + 0.734375, + 0.597656, + 0.921875, + 3.09375, + 0.585938, + 0.527344, + 0.570312, + 1.84375, + 0.574219, + 0.617188, + 0.65625, + 0.75, + 0.601562, + 0.539062, + 1.664062, + 0.777344, + 0.507812, + 0.652344, + 0.699219, + 2.8125, + ], + ) -> None: + super().__init__() + + if len(temporal_downsample) != len(block_in_channels): + raise ValueError( + "`temporal_downsample` must have one value per block in `block_in_channels`, got " + f"{len(temporal_downsample)} and {len(block_in_channels)}." + ) + if temporal_downsample[-1]: + raise ValueError( + "The last value must be `False` because the final encoder/decoder block has no temporal " + "downsample/upsample layer." + ) + + # The encoder/decoder backbone compresses space by `patch_size * 2 ** (len(temporal_downsample) - 1)`; the + # extra `JoyVideoEditStem` / `JoyVideoEditHead` around it multiply that by 3/2, so pixels are compressed 24x. + self.backbone_spatial_ratio = patch_size * 2 ** (len(temporal_downsample) - 1) + self.spatial_compression_ratio = self.backbone_spatial_ratio * 3 // 2 + self.temporal_compression_ratio = 2 ** sum(temporal_downsample[:-1]) + if chunk_size <= 0 or chunk_size % self.temporal_compression_ratio != 0: + raise ValueError( + f"`chunk_size` must be a positive multiple of the temporal compression ratio " + f"({self.temporal_compression_ratio}), got {chunk_size}." + ) + + self.stem = JoyVideoEditStem(in_channels) + self.encoder = JoyVideoEditEncoder( + in_channels=in_channels * patch_size**2, + z_channels=latent_channels, + num_res_blocks=layers_per_block, + block_in_channels=block_in_channels, + temporal_downsample=temporal_downsample, + ) + self.decoder = JoyVideoEditDecoder( + z_channels=latent_channels, + out_channels=out_channels * patch_size**2, + num_res_blocks=layers_per_block, + block_in_channels=tuple(reversed(block_in_channels)), + temporal_upsample=temporal_downsample, + ) + self.head = JoyVideoEditHead(out_channels) + + # When decoding a batch of video latents at a time, one can save memory by slicing across the batch + # dimension to perform decoding of a single video latent at a time. + self.use_slicing = False + + # Precompute and cache conv counts for encoder and decoder for clear_cache speedup + self._cached_conv_counts = { + "encoder": sum(isinstance(m, JoyVideoEditCausalConv3d) for m in self.encoder.modules()), + "decoder": sum(isinstance(m, JoyVideoEditCausalConv3d) for m in self.decoder.modules()), + } + self.clear_cache() + + def clear_cache(self) -> None: + self._enc_conv_idx = [0] + self._dec_conv_idx = [0] + self._enc_feat_map = [None] * self._cached_conv_counts["encoder"] + self._dec_feat_map = [None] * self._cached_conv_counts["decoder"] + + def _encode(self, x: torch.Tensor) -> torch.Tensor: + x = self.stem(x) + x = patchify(x, self.config.patch_size) + + self.clear_cache() + out = [] + num_chunks = 1 + math.ceil((x.shape[2] - 1) / self.config.chunk_size) + for i in range(num_chunks): + self._enc_conv_idx = [0] + if i == 0: + chunk = self.encoder( + x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx, first_chunk=True + ) + else: + start = 1 + (i - 1) * self.config.chunk_size + end = 1 + i * self.config.chunk_size + chunk = self.encoder( + x[:, :, start:end, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + first_chunk=False, + ) + out.append(chunk) + out = torch.cat(out, dim=2) + self.clear_cache() + return out + + @apply_forward_hook + def encode( + self, x: torch.Tensor, return_dict: bool = True + ) -> AutoencoderKLOutput | tuple[DiagonalGaussianDistribution]: + r""" + Encode a batch of videos into latents. + + Args: + x (`torch.Tensor`): Input batch of videos, shape `(batch, channels, frames, height, width)`. `frames` + must equal `self.temporal_compression_ratio * n + 1` for some integer `n`, and `height` / `width` must + be divisible by `self.spatial_compression_ratio`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.modeling_outputs.AutoencoderKLOutput`] instead of a plain tuple. + + Returns: + The latent representations of the encoded videos. If `return_dict` is True, a + [`~models.modeling_outputs.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned. + """ + _, _, num_frames, height, width = x.shape + if (num_frames - 1) % self.temporal_compression_ratio != 0: + raise ValueError(f"Temporal dimension must be {self.temporal_compression_ratio}n + 1, got {x.shape}.") + if height % self.spatial_compression_ratio != 0 or width % self.spatial_compression_ratio != 0: + raise ValueError( + f"Spatial dimensions must be divisible by {self.spatial_compression_ratio}, got {x.shape}." + ) + + if self.use_slicing and x.shape[0] > 1: + encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)] + h = torch.cat(encoded_slices) + else: + h = self._encode(x) + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) + + def _decode(self, z: torch.Tensor) -> torch.Tensor: + _, _, num_latent_frames, _, _ = z.shape + latent_chunk_size = self.config.chunk_size // self.temporal_compression_ratio + + self.clear_cache() + decoded = [] + num_chunks = 1 + math.ceil((num_latent_frames - 1) / latent_chunk_size) + for i in range(num_chunks): + self._dec_conv_idx = [0] + if i == 0: + chunk = self.decoder( + z[:, :, :1, :, :], feat_cache=self._dec_feat_map, feat_idx=self._dec_conv_idx, first_chunk=True + ) + else: + start = 1 + (i - 1) * latent_chunk_size + end = 1 + i * latent_chunk_size + chunk = self.decoder( + z[:, :, start:end, :, :], + feat_cache=self._dec_feat_map, + feat_idx=self._dec_conv_idx, + first_chunk=False, + ) + decoded.append(chunk) + decoded = torch.cat(decoded, dim=2) + self.clear_cache() + + decoded = unpatchify(decoded, self.config.patch_size) + return self.head(decoded) + + @apply_forward_hook + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor: + r""" + Decode a batch of latents into videos. + + Args: + z (`torch.Tensor`): Input batch of latent vectors. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoders.vae.DecoderOutput`] instead of a plain tuple. + + Returns: + [`~models.autoencoders.vae.DecoderOutput`] or `tuple`: + If return_dict is True, a [`~models.autoencoders.vae.DecoderOutput`] is returned, otherwise a plain + `tuple` is returned. + """ + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = self._decode(z) + + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: torch.Generator | None = None, + ) -> DecoderOutput | torch.Tensor: + r""" + Args: + sample (`torch.Tensor`): Input sample. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the posterior. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + generator (`torch.Generator`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling + deterministic. + + Returns: + [`~models.autoencoders.vae.DecoderOutput`] or `tuple`: + If `return_dict` is True, a [`~models.autoencoders.vae.DecoderOutput`] is returned, otherwise a plain + `tuple` is returned. + """ + posterior = self.encode(sample).latent_dist + z = posterior.sample(generator=generator) if sample_posterior else posterior.mode() + return self.decode(z, return_dict=return_dict) diff --git a/src/diffusers/models/cache_utils.py b/src/diffusers/models/cache_utils.py index 5aa189987ba2..1d057e93797d 100644 --- a/src/diffusers/models/cache_utils.py +++ b/src/diffusers/models/cache_utils.py @@ -41,12 +41,13 @@ def enable_cache(self, config) -> None: Enable caching techniques on the model. Args: - config (`PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig | TextKVCacheConfig`): + config (`PyramidAttentionBroadcastConfig | FasterCacheConfig | FirstBlockCacheConfig | TextKVCacheConfig | JoyVideoEditKVCacheConfig`): The configuration for applying the caching technique. Currently supported caching techniques are: - [`~hooks.PyramidAttentionBroadcastConfig`] - [`~hooks.FasterCacheConfig`] - [`~hooks.FirstBlockCacheConfig`] - [`~hooks.TextKVCacheConfig`] + - [`~hooks.JoyVideoEditKVCacheConfig`] Example: @@ -69,12 +70,14 @@ def enable_cache(self, config) -> None: from ..hooks import ( FasterCacheConfig, FirstBlockCacheConfig, + JoyVideoEditKVCacheConfig, MagCacheConfig, PyramidAttentionBroadcastConfig, TaylorSeerCacheConfig, TextKVCacheConfig, apply_faster_cache, apply_first_block_cache, + apply_joyvideoedit_kv_cache, apply_mag_cache, apply_pyramid_attention_broadcast, apply_taylorseer_cache, @@ -94,6 +97,8 @@ def enable_cache(self, config) -> None: apply_mag_cache(self, config) elif isinstance(config, TextKVCacheConfig): apply_text_kv_cache(self, config) + elif isinstance(config, JoyVideoEditKVCacheConfig): + apply_joyvideoedit_kv_cache(self, config) elif isinstance(config, PyramidAttentionBroadcastConfig): apply_pyramid_attention_broadcast(self, config) elif isinstance(config, TaylorSeerCacheConfig): @@ -108,6 +113,7 @@ def disable_cache(self) -> None: FasterCacheConfig, FirstBlockCacheConfig, HookRegistry, + JoyVideoEditKVCacheConfig, MagCacheConfig, PyramidAttentionBroadcastConfig, TaylorSeerCacheConfig, @@ -115,6 +121,7 @@ def disable_cache(self) -> None: ) from ..hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK from ..hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK + from ..hooks.joyvideoedit_kv_cache import _JOYVIDEOEDIT_KV_CACHE_HOOK from ..hooks.mag_cache import _MAG_CACHE_BLOCK_HOOK, _MAG_CACHE_LEADER_BLOCK_HOOK from ..hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK from ..hooks.taylorseer_cache import _TAYLORSEER_CACHE_HOOK @@ -139,6 +146,8 @@ def disable_cache(self) -> None: elif isinstance(self._cache_config, TextKVCacheConfig): registry.remove_hook(_TEXT_KV_CACHE_TRANSFORMER_HOOK, recurse=True) registry.remove_hook(_TEXT_KV_CACHE_BLOCK_HOOK, recurse=True) + elif isinstance(self._cache_config, JoyVideoEditKVCacheConfig): + registry.remove_hook(_JOYVIDEOEDIT_KV_CACHE_HOOK, recurse=True) elif isinstance(self._cache_config, TaylorSeerCacheConfig): registry.remove_hook(_TAYLORSEER_CACHE_HOOK, recurse=True) else: diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index e9fcfcf320dc..49ce10188b8b 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -44,6 +44,7 @@ from .transformer_ideogram4 import Ideogram4Transformer2DModel from .transformer_joyimage import JoyImageEditTransformer3DModel from .transformer_joyimage_edit_plus import JoyImageEditPlusTransformer3DModel + from .transformer_joyvideoedit import JoyVideoEditTransformer3DModel from .transformer_kandinsky import Kandinsky5Transformer3DModel from .transformer_krea2 import Krea2Transformer2DModel from .transformer_longcat_audio_dit import LongCatAudioDiTTransformer diff --git a/src/diffusers/models/transformers/transformer_joyvideoedit.py b/src/diffusers/models/transformers/transformer_joyvideoedit.py new file mode 100644 index 000000000000..44a7bbe926dc --- /dev/null +++ b/src/diffusers/models/transformers/transformer_joyvideoedit.py @@ -0,0 +1,967 @@ +# Copyright 2026 The JoyAI-Video-Edit Team and 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. + +import math +from collections.abc import Callable, Iterable + +import torch +import torch.nn as nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...hooks.joyvideoedit_kv_cache import _JOYVIDEOEDIT_KV_CACHE_HOOK, JoyVideoEditKVCacheState +from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward +from ..attention_dispatch import dispatch_attention_fn +from ..cache_utils import CacheMixin +from ..embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps, get_1d_rotary_pos_embed +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin + + +# Visual-token roles encoded by the source-id RoPE. +SOURCE_ID_TARGET = 0.0 +SOURCE_ID_EDIT_CONDITION = 1.0 +SOURCE_ID_EXTRA_REF_IMAGE = 2.0 + +TIME_FREQ_DIM = 256 + +NORM_EPS = 1e-6 +NUM_MODULATION_CHUNKS = 6 + +SELF_ATTN_MODE_REF_IMAGE_CACHE = "ref_image_cache" + + +# --------------------------------------------------------------------------- +# Rotary position embedding utilities +# --------------------------------------------------------------------------- + + +def _apply_rotary_emb(x: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: + """Apply rotary embeddings to a `(B, L, H, D)` tensor.""" + cos = freqs_cis[0].unsqueeze(2).to(x.device) + sin = freqs_cis[1].unsqueeze(2).to(x.device) + + x_real, x_imag = x.float().reshape(*x.shape[:-1], -1, 2).unbind(-1) + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(-2) + + return (x.float() * cos + x_rotated * sin).type_as(x) + + +def _concat_kv_entries( + entries: Iterable[dict[str, torch.Tensor]], + *, + device: torch.device, + dtype: torch.dtype, + cached_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Concatenate a sequence of stored KV-cache entries into a single `(key, value)` pair. + + Entries stored with `pre_rope=True` were cached *before* RoPE was applied to their key (so that RoPE can be + re-derived from `current_temporal_ids` at read time instead of being frozen at write time). For those entries, RoPE + is applied here using consecutive slices of `cached_freqs_cis`, in the same order the entries are concatenated. + """ + keys = [] + values = [] + pre_rope_offset = 0 + + for entry in entries: + if entry is None: + continue + key = entry.get("key") + value = entry.get("value") + if key is None or value is None: + continue + + key = key.to(device=device, dtype=dtype) + value = value.to(device=device, dtype=dtype) + + if entry.get("pre_rope", False) and cached_freqs_cis is not None: + cos_all, sin_all = cached_freqs_cis + seg_len = key.shape[1] + cos_seg = cos_all[..., pre_rope_offset : pre_rope_offset + seg_len, :] + sin_seg = sin_all[..., pre_rope_offset : pre_rope_offset + seg_len, :] + key = _apply_rotary_emb(key, (cos_seg, sin_seg)) + pre_rope_offset += seg_len + + keys.append(key) + values.append(value) + + if not keys: + return None, None + + return torch.cat(keys, dim=1), torch.cat(values, dim=1) + + +def _clone_kv_tensor(tensor: torch.Tensor | None) -> torch.Tensor | None: + if tensor is None: + return None + return tensor.detach().clone() + + +# --------------------------------------------------------------------------- +# Modulation +# --------------------------------------------------------------------------- + + +class JoyVideoEditModulate(nn.Module): + """Wan-style learnable modulation table. + + Produces `factor` modulation vectors by adding the conditioning signal to a learnable parameter table. + """ + + def __init__(self, hidden_size: int, factor: int, dtype=None, device=None): + super().__init__() + self.factor = factor + self.modulate_table = nn.Parameter( + torch.randn(1, factor, hidden_size, dtype=dtype, device=device) / hidden_size**0.5, + requires_grad=True, + ) + + def forward(self, x: torch.Tensor) -> list[torch.Tensor]: + if x.ndim != 3: + x = x.unsqueeze(1) + return [o.squeeze(1) for o in (self.modulate_table + x).chunk(self.factor, dim=1)] + + +# --------------------------------------------------------------------------- +# Attention processor +# --------------------------------------------------------------------------- + + +class JoyVideoEditAttnProcessor: + """Joint self-attention processor for `JoyVideoEditAttention`. + + Computes fused QKV projections for the image and (optionally) text streams, applies per-head RMSNorm and 3D RoPE + (with the source-id RoPE already folded into `image_rotary_emb` / `text_rotary_emb` by the caller), then runs joint + attention over `[img, txt, *cached_kv]`. KV-cache read/write is handled here since it must happen right after the + image stream's key/value are produced (before the joint concat with text). + """ + + _attention_backend = None + _parallel_config = None + + def __init__(self): + pass + + def __call__( + self, + attn: "JoyVideoEditAttention", + hidden_states: torch.Tensor, # image stream (B, S_img, D) + encoder_hidden_states: torch.Tensor | None = None, # text stream (B, S_txt, D) + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + encoder_hidden_states_mask: torch.Tensor | None = None, # text padding mask (B, S_txt), True = keep + skip_text_stream: bool = False, + kv_cache_reader: Callable[[int | None], Iterable[dict[str, torch.Tensor]]] | None = None, + kv_cache_writer: Callable[[int | None, torch.Tensor, torch.Tensor], None] | None = None, + layer_idx: int | None = None, + kv_cache_pre_rope: bool = False, + cached_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + heads = attn.heads + + img_qkv = attn.img_attn_qkv(hidden_states) + img_query, img_key, img_value = img_qkv.chunk(3, dim=-1) + img_query = img_query.unflatten(-1, (heads, -1)) + img_key = img_key.unflatten(-1, (heads, -1)) + img_value = img_value.unflatten(-1, (heads, -1)) + + if kv_cache_pre_rope: + # Cache the key right after QK-norm but *before* RoPE, so cached entries can be re-rotated later from + # whatever `current_temporal_ids` apply at read time (needed for non-contiguous / relative temporal ids). + img_key_for_cache = attn.img_attn_k_norm(img_key) + img_query = attn.img_attn_q_norm(img_query) + img_key = attn.img_attn_k_norm(img_key) + if image_rotary_emb is not None: + img_query, img_key = ( + _apply_rotary_emb(img_query, image_rotary_emb), + _apply_rotary_emb(img_key, image_rotary_emb), + ) + if not kv_cache_pre_rope: + img_key_for_cache = img_key + + if not skip_text_stream: + txt_qkv = attn.txt_attn_qkv(encoder_hidden_states) + txt_query, txt_key, txt_value = txt_qkv.chunk(3, dim=-1) + txt_query = txt_query.unflatten(-1, (heads, -1)) + txt_key = txt_key.unflatten(-1, (heads, -1)) + txt_value = txt_value.unflatten(-1, (heads, -1)) + txt_query = attn.txt_attn_q_norm(txt_query) + txt_key = attn.txt_attn_k_norm(txt_key) + + if skip_text_stream: + query = img_query + key = img_key + value = img_value + else: + query = torch.cat((img_query, txt_query), dim=1) + key = torch.cat((img_key, txt_key), dim=1) + value = torch.cat((img_value, txt_value), dim=1) + + if kv_cache_writer is not None: + kv_cache_writer(layer_idx, img_key_for_cache, img_value) + + if kv_cache_reader is not None: + cached_key, cached_value = _concat_kv_entries( + kv_cache_reader(layer_idx), + device=query.device, + dtype=query.dtype, + cached_freqs_cis=cached_freqs_cis if kv_cache_pre_rope else None, + ) + else: + cached_key = cached_value = None + + if cached_key is not None: + key = torch.cat([cached_key, key], dim=1) + value = torch.cat([cached_value, value], dim=1) + + # Build the joint-attention mask so padded text tokens never contribute to any query's softmax. Only the text + # stream carries padding; the image stream and any cached image KV are always valid. The key order is + # `[cached_img_kv, img, txt]` (text last), so the mask is all-ones over the leading visual span and equals + # `encoder_hidden_states_mask` over the trailing text span. Masking the key positions (broadcast over queries + # via the `(B, 1, 1, S_key)` shape) is sufficient — the padded text queries produce garbage rows that are + # discarded, since only the image tokens are read out downstream. + attn_mask = None + if not skip_text_stream and encoder_hidden_states_mask is not None: + num_visual_keys = key.shape[1] - encoder_hidden_states.shape[1] + visual_mask = encoder_hidden_states_mask.new_ones((key.shape[0], num_visual_keys)) + attn_mask = torch.cat([visual_mask, encoder_hidden_states_mask], dim=1)[:, None, None, :] + + joint_hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attn_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + joint_hidden_states = joint_hidden_states.flatten(2, 3).to(query.dtype) + + if skip_text_stream: + img_attn_output = joint_hidden_states + txt_attn_output = None + else: + img_attn_output = joint_hidden_states[:, : hidden_states.shape[1]] + txt_attn_output = joint_hidden_states[:, hidden_states.shape[1] :] + + img_attn_output = attn.img_attn_proj(img_attn_output) + if txt_attn_output is not None: + txt_attn_output = attn.txt_attn_proj(txt_attn_output) + + return img_attn_output, txt_attn_output + + +# --------------------------------------------------------------------------- +# Attention module +# --------------------------------------------------------------------------- + + +class JoyVideoEditAttention(nn.Module, AttentionModuleMixin): + """Joint attention module for JoyVideoEdit double-stream blocks. + + Wraps the fused QKV projections, per-head RMSNorm, and output projections for both the image and text streams. + Delegates the attention computation (RoPE, joint attention, KV-cache read/write) to a pluggable + `JoyVideoEditAttnProcessor`. + """ + + _default_processor_cls = JoyVideoEditAttnProcessor + _available_processors = [JoyVideoEditAttnProcessor] + _supports_qkv_fusion = False + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + eps: float = NORM_EPS, + processor=None, + ): + super().__init__() + + self.heads = num_attention_heads + self.head_dim = attention_head_dim + inner_dim = num_attention_heads * attention_head_dim + + self.img_attn_qkv = nn.Linear(dim, inner_dim * 3, bias=True) + self.img_attn_q_norm = nn.RMSNorm(attention_head_dim, eps=eps) + self.img_attn_k_norm = nn.RMSNorm(attention_head_dim, eps=eps) + self.img_attn_proj = nn.Linear(inner_dim, dim, bias=True) + + self.txt_attn_qkv = nn.Linear(dim, inner_dim * 3, bias=True) + self.txt_attn_q_norm = nn.RMSNorm(attention_head_dim, eps=eps) + self.txt_attn_k_norm = nn.RMSNorm(attention_head_dim, eps=eps) + self.txt_attn_proj = nn.Linear(inner_dim, dim, bias=True) + + if processor is None: + processor = self._default_processor_cls() + self.set_processor(processor) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + return self.processor(self, hidden_states, encoder_hidden_states, image_rotary_emb, **kwargs) + + +# --------------------------------------------------------------------------- +# Transformer block +# --------------------------------------------------------------------------- + + +class JoyVideoEditTransformerBlock(nn.Module): + """Double-stream transformer block for JoyVideoEdit.""" + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_width_ratio: float = 4.0, + eps: float = NORM_EPS, + ): + super().__init__() + + mlp_hidden_dim = int(dim * mlp_width_ratio) + + self.img_mod = JoyVideoEditModulate(dim, factor=NUM_MODULATION_CHUNKS) + self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.img_mlp = FeedForward(dim, inner_dim=mlp_hidden_dim, activation_fn="gelu-approximate") + + self.txt_mod = JoyVideoEditModulate(dim, factor=NUM_MODULATION_CHUNKS) + self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.txt_mlp = FeedForward(dim, inner_dim=mlp_hidden_dim, activation_fn="gelu-approximate") + + self.attn = JoyVideoEditAttention(dim, num_attention_heads, attention_head_dim, eps=eps) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + encoder_hidden_states_mask: torch.Tensor | None = None, + kv_cache_reader: Callable[[int | None], Iterable[dict[str, torch.Tensor]]] | None = None, + kv_cache_writer: Callable[[int | None, torch.Tensor, torch.Tensor], None] | None = None, + layer_idx: int | None = None, + skip_text_stream: bool = False, + kv_cache_pre_rope: bool = False, + cached_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + ( + img_mod1_shift, + img_mod1_scale, + img_mod1_gate, + img_mod2_shift, + img_mod2_scale, + img_mod2_gate, + ) = self.img_mod(temb) + if not skip_text_stream: + ( + txt_mod1_shift, + txt_mod1_scale, + txt_mod1_gate, + txt_mod2_shift, + txt_mod2_scale, + txt_mod2_gate, + ) = self.txt_mod(temb) + + img_modulated = self.img_norm1(hidden_states) * (1 + img_mod1_scale.unsqueeze(1)) + img_mod1_shift.unsqueeze(1) + txt_modulated = None + if not skip_text_stream: + txt_modulated = self.txt_norm1(encoder_hidden_states) * ( + 1 + txt_mod1_scale.unsqueeze(1) + ) + txt_mod1_shift.unsqueeze(1) + + img_attn, txt_attn = self.attn( + hidden_states=img_modulated, + encoder_hidden_states=txt_modulated, + image_rotary_emb=image_rotary_emb, + encoder_hidden_states_mask=encoder_hidden_states_mask, + skip_text_stream=skip_text_stream, + kv_cache_reader=kv_cache_reader, + kv_cache_writer=kv_cache_writer, + layer_idx=layer_idx, + kv_cache_pre_rope=kv_cache_pre_rope, + cached_freqs_cis=cached_freqs_cis, + ) + + hidden_states = hidden_states + img_attn * img_mod1_gate.unsqueeze(1) + img_mod2_modulated = self.img_norm2(hidden_states) * ( + 1 + img_mod2_scale.unsqueeze(1) + ) + img_mod2_shift.unsqueeze(1) + hidden_states = hidden_states + self.img_mlp(img_mod2_modulated) * img_mod2_gate.unsqueeze(1) + + if not skip_text_stream: + encoder_hidden_states = encoder_hidden_states + txt_attn * txt_mod1_gate.unsqueeze(1) + txt_mod2_modulated = self.txt_norm2(encoder_hidden_states) * ( + 1 + txt_mod2_scale.unsqueeze(1) + ) + txt_mod2_shift.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + self.txt_mlp(txt_mod2_modulated) * txt_mod2_gate.unsqueeze( + 1 + ) + + return hidden_states, encoder_hidden_states + + +# Copied from diffusers.models.transformers.transformer_joyimage.JoyImageTimeTextImageEmbedding with JoyImage->JoyVideoEdit +class JoyVideoEditTimeTextImageEmbedding(nn.Module): + def __init__( + self, + dim: int, + time_freq_dim: int, + time_proj_dim: int, + text_embed_dim: int, + ): + super().__init__() + + self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, time_embed_dim=dim) + self.act_fn = nn.SiLU() + self.time_proj = nn.Linear(dim, time_proj_dim) + self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh") + + def forward( + self, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + ): + timestep = self.timesteps_proj(timestep) + + time_embedder_dtype = next(iter(self.time_embedder.parameters())).dtype + if timestep.dtype != time_embedder_dtype and time_embedder_dtype != torch.int8: + timestep = timestep.to(time_embedder_dtype) + temb = self.time_embedder(timestep).type_as(encoder_hidden_states) + timestep_proj = self.time_proj(self.act_fn(temb)) + + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + + return temb, timestep_proj, encoder_hidden_states + + +# --------------------------------------------------------------------------- +# Main model +# --------------------------------------------------------------------------- + + +class JoyVideoEditTransformer3DModel(ModelMixin, ConfigMixin, AttentionMixin, CacheMixin): + """JoyVideoEdit streaming video-editing transformer. + + A dual-stream MM-DiT with source-id RoPE, cross-chunk KV caching, and reference-video latent conditioning for + chunk-wise causal video editing. + """ + + _skip_layerwise_casting_patterns = ["img_in", "condition_embedder", "norm"] + _no_split_modules = ["JoyVideoEditTransformerBlock"] + _supports_gradient_checkpointing = True + _keep_in_fp32_modules = [ + "time_embedder", + "norm1", + "norm2", + "norm_out", + ] + _repeated_blocks = ["JoyVideoEditTransformerBlock"] + + @register_to_config + def __init__( + self, + patch_size: list[int] = [1, 1, 1], + in_channels: int = 64, + out_channels: int | None = None, + hidden_size: int = 4096, + num_attention_heads: int = 32, + text_dim: int = 4096, + mlp_width_ratio: float = 4.0, + num_layers: int = 40, + rope_dim_list: list[int] = [16, 56, 56], + theta: int = 256, + chunk_size: int = 1, + local_window_size: int = 3, + global_sink_chunk: bool = True, + source_id_rope_dim: int = 128, + source_id_rope_theta: float = 256.0, + ): + if chunk_size <= 0: + raise ValueError(f"`chunk_size` must be positive, got {chunk_size}.") + if local_window_size <= 0: + raise ValueError(f"`local_window_size` must be positive, got {local_window_size}.") + if source_id_rope_dim < 0 or source_id_rope_dim % 2 != 0: + raise ValueError(f"`source_id_rope_dim` must be a non-negative even integer, got {source_id_rope_dim}.") + + super().__init__() + + self.out_channels = out_channels or in_channels + self.patch_size = patch_size + self.hidden_size = hidden_size + self.num_attention_heads = num_attention_heads + self.rope_dim_list = rope_dim_list + self.theta = theta + self.source_id_rope_dim = int(source_id_rope_dim) + self.source_id_rope_theta = float(source_id_rope_theta) + + if hidden_size % num_attention_heads != 0: + raise ValueError( + f"hidden_size ({hidden_size}) must be divisible by num_attention_heads ({num_attention_heads})" + ) + attention_head_dim = hidden_size // num_attention_heads + + self.img_in = nn.Conv3d(in_channels, hidden_size, kernel_size=patch_size, stride=patch_size) + + self.condition_embedder = JoyVideoEditTimeTextImageEmbedding( + dim=hidden_size, + time_freq_dim=TIME_FREQ_DIM, + time_proj_dim=hidden_size * NUM_MODULATION_CHUNKS, + text_embed_dim=text_dim, + ) + + self.double_blocks = nn.ModuleList( + [ + JoyVideoEditTransformerBlock( + dim=hidden_size, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_width_ratio=mlp_width_ratio, + ) + for _ in range(num_layers) + ] + ) + + self.norm_out = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=NORM_EPS) + self.proj_out = nn.Linear(hidden_size, self.out_channels * math.prod(patch_size)) + + self.gradient_checkpointing = False + self._kv_cache_chunk_id = None + self._kv_cache_selected_chunk_ids = None + self._kv_cache_pre_rope = False + + # ------------------------------------------------------------------ + # KV-cache bookkeeping + # + # Clean chunks are stored per layer and reused as attention context by later chunks. Cache tensors live in the + # state managed by `JoyVideoEditKVCacheHook`; the model controls cache selection, assembly, and eviction. + # ------------------------------------------------------------------ + + def _kv_cache_state(self) -> JoyVideoEditKVCacheState: + registry = getattr(self, "_diffusers_hook", None) + hook = registry.get_hook(_JOYVIDEOEDIT_KV_CACHE_HOOK) if registry is not None else None + if hook is None: + raise RuntimeError( + "The JoyVideoEdit KV cache is not enabled. Call `enable_cache(JoyVideoEditKVCacheConfig())` before " + "using `kv_cache_mode`/`kv_cache_selected_chunk_ids`." + ) + return hook.state_manager.get_state() + + def configure_inference_kv_cache( + self, + *, + chunk_id: int | None = None, + selected_chunk_ids: list[int] | None = None, + pre_rope: bool = False, + ) -> None: + self._kv_cache_chunk_id = chunk_id + self._kv_cache_selected_chunk_ids = list(selected_chunk_ids) if selected_chunk_ids is not None else None + self._kv_cache_pre_rope = bool(pre_rope) + + def _read_layer_kv_cache(self, layer_idx: int | None) -> list[dict[str, torch.Tensor]]: + if layer_idx is None: + return [] + chunk_cache = self._kv_cache_state().chunk_cache + selected_chunk_ids = self._kv_cache_selected_chunk_ids or [] + layer_entries = [] + for selected_chunk_id in selected_chunk_ids: + chunk_store = chunk_cache.get(selected_chunk_id) + if chunk_store is None: + continue + entry = chunk_store.get(layer_idx) + if entry is not None: + layer_entries.append(entry) + return layer_entries + + def _write_layer_kv_cache( + self, + layer_idx: int | None, + key: torch.Tensor, + value: torch.Tensor, + ) -> None: + if layer_idx is None or self._kv_cache_chunk_id is None: + return + state = self._kv_cache_state() + chunk_store = state.chunk_cache.setdefault(self._kv_cache_chunk_id, {}) + chunk_store[layer_idx] = { + "key": _clone_kv_tensor(key), + "value": _clone_kv_tensor(value), + "pre_rope": bool(self._kv_cache_pre_rope), + } + + def evict_kv_cache_chunks(self, chunk_ids_to_keep: set) -> None: + """Drop every cached chunk whose id is not in `chunk_ids_to_keep`.""" + state = self._kv_cache_state() + evict_ids = [cid for cid in state.chunk_cache if cid not in chunk_ids_to_keep] + for cid in evict_ids: + del state.chunk_cache[cid] + + # ------------------------------------------------------------------ + # RoPE helpers + # ------------------------------------------------------------------ + + def get_rotary_pos_embed_from_ids( + self, + *, + frame_ids: torch.Tensor, + spatial_shape: tuple[int, int], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Build 3D `(temporal, height, width)` RoPE from explicit per-token temporal positions. + + `frame_ids` has shape `(batch_size, sequence_length)` and supports non-contiguous or relative ids. + """ + post_patch_height, post_patch_width = spatial_shape + device = frame_ids.device + temporal_positions = frame_ids.to(dtype=torch.float32) + spatial_tokens_per_frame = post_patch_height * post_patch_width + if temporal_positions.shape[1] % spatial_tokens_per_frame != 0: + raise ValueError( + f"`frame_ids` length {temporal_positions.shape[1]} is not divisible by spatial token count " + f"{spatial_tokens_per_frame}." + ) + + h_positions = torch.arange(post_patch_height, dtype=torch.float32, device=device) + w_positions = torch.arange(post_patch_width, dtype=torch.float32, device=device) + h_grid, w_grid = torch.meshgrid(h_positions, w_positions, indexing="ij") + num_frames_in_grid = temporal_positions.shape[1] // spatial_tokens_per_frame + batch_size = temporal_positions.shape[0] + h_positions = h_grid.reshape(-1).repeat(num_frames_in_grid).unsqueeze(0).expand(batch_size, -1) + w_positions = w_grid.reshape(-1).repeat(num_frames_in_grid).unsqueeze(0).expand(batch_size, -1) + + head_dim = self.hidden_size // self.num_attention_heads + rope_dim_list = self.rope_dim_list + if sum(rope_dim_list) != head_dim: + raise ValueError("sum(rope_dim_list) should equal to head_dim of attention layer") + + cos_list = [] + sin_list = [] + for dim, positions in zip(rope_dim_list, (temporal_positions, h_positions, w_positions)): + cos, sin = get_1d_rotary_pos_embed(dim, positions.reshape(-1), theta=self.theta, use_real=True) + cos_list.append(cos.unflatten(0, (batch_size, -1))) + sin_list.append(sin.unflatten(0, (batch_size, -1))) + vis_freqs = (torch.cat(cos_list, dim=2), torch.cat(sin_list, dim=2)) + + return vis_freqs + + def generate_source_id_rope( + self, + source_id: torch.Tensor, + head_dim: int, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Per-token rotary phase encoding a token's *role* (target / edit-condition / extra-ref-image), independent + of its spatiotemporal position. Only the first `source_id_rope_dim` head channels carry a role-dependent angle; + the remaining channels get an identity rotation (cos=1, sin=0). Composed with the 3D spatiotemporal RoPE in + `forward` via `new_cos = cos_3d * cos_role - sin_3d * sin_role` (i.e. adding the two rotation angles), so a + token's final rotary phase is "spatiotemporal position" + "role". + """ + role_dim = max(0, min(int(self.source_id_rope_dim), int(head_dim))) + + half_head = head_dim // 2 + cos_half = torch.ones(*source_id.shape, half_head, device=device, dtype=torch.float32) + sin_half = torch.zeros(*source_id.shape, half_head, device=device, dtype=torch.float32) + + inv_freq = 1.0 / ( + self.source_id_rope_theta ** (torch.arange(0, role_dim, 2, device=device, dtype=torch.float32) / role_dim) + ) + + angles = source_id.unsqueeze(-1) * inv_freq + cos_half[..., : role_dim // 2] = torch.cos(angles) + sin_half[..., : role_dim // 2] = torch.sin(angles) + return ( + cos_half.repeat_interleave(2, dim=-1).to(dtype=dtype), + sin_half.repeat_interleave(2, dim=-1).to(dtype=dtype), + ) + + @staticmethod + def _get_patch_shape(latent: torch.Tensor, patch_size: tuple[int, int, int]) -> tuple[int, int, int]: + _, _, num_frames, height, width = latent.shape + return ( + num_frames // patch_size[0], + height // patch_size[1], + width // patch_size[2], + ) + + @staticmethod + def _get_token_frame_ids( + post_patch_shape: tuple[int, int, int], + device: torch.device, + temporal_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + num_frames, post_patch_height, post_patch_width = post_patch_shape + spatial_tokens_per_frame = post_patch_height * post_patch_width + if temporal_ids is None: + frame_ids = torch.arange(num_frames, device=device, dtype=torch.long).unsqueeze(0) + else: + frame_ids = torch.as_tensor(temporal_ids, device=device, dtype=torch.long) + if frame_ids.ndim != 2 or frame_ids.shape[1] != num_frames: + raise ValueError( + f"`temporal_ids` must have shape `(batch_size, {num_frames})`, got {tuple(frame_ids.shape)}." + ) + return frame_ids.repeat_interleave(spatial_tokens_per_frame, dim=1) + + # ------------------------------------------------------------------ + # Unpatchify + # ------------------------------------------------------------------ + + def unpatchify(self, x: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor: + c = self.out_channels + pt, ph, pw = self.patch_size + if t * h * w != x.shape[1]: + raise ValueError(f"Expected t*h*w ({t * h * w}) to equal x.shape[1] ({x.shape[1]})") + + x = x.reshape(x.shape[0], t, h, w, c, pt, ph, pw) + # (B, T, H, W, C, Pt, Ph, Pw) -> (B, C, T, Pt, H, Ph, W, Pw) + x = x.permute(0, 4, 1, 5, 2, 6, 3, 7) + return x.reshape(x.shape[0], c, t * pt, h * ph, w * pw) + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_mask: torch.Tensor | None = None, + ref_video_latent: torch.Tensor | None = None, + current_temporal_ids: torch.Tensor | None = None, + cached_temporal_ids: torch.Tensor | None = None, + kv_cache_mode: str | None = None, + kv_cache_chunk_id: int | None = None, + kv_cache_selected_chunk_ids: list[int] | None = None, + kv_cache_pre_rope: bool = False, + self_attn_input_mode: str | None = None, + skip_text_stream: bool = False, + return_dict: bool = True, + ): + """ + The [`JoyVideoEditTransformer3DModel`] forward method. + + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, num_channels, num_frames, height, width)`): + The noisy input latent patchified with `patch_size`. + timestep (`torch.Tensor` of shape `(batch_size,)`): + Denoising timestep, one scalar per batch element. + encoder_hidden_states (`torch.Tensor`): + Text conditioning embeddings. + encoder_hidden_states_mask (`torch.Tensor`, *optional*): + Boolean/int padding mask over `encoder_hidden_states` (`True`/non-zero keeps the token). When given, it + is folded into the joint self-attention so padded text tokens contribute nothing to any query's softmax + — required for correct batched inference with variable-length prompts. + ref_video_latent (`torch.Tensor`, *optional*): + A reference/edit-condition video latent, patchified with the same `img_in` and concatenated into the + image stream (tagged with the `SOURCE_ID_EDIT_CONDITION` role) so it participates in joint + self-attention as extra context, without being part of the denoised output. + current_temporal_ids (`torch.Tensor`, *optional*, shape `(batch_size, num_frames)`): + Explicit per-frame temporal ids for `hidden_states` (and, if provided, `ref_video_latent`). Falls back + to `0..num_frames-1` when omitted. + cached_temporal_ids (`torch.Tensor`, *optional*): + Temporal ids of the tokens stored in the selected KV-cache chunks, used to rebuild RoPE for + `pre_rope`-cached entries at read time. + kv_cache_mode (`str`, *optional*): One of `"store"`, `"reuse"`, or `None`. + kv_cache_chunk_id (`int`, *optional*): Identifier under which this call's image-stream KV is stored. + kv_cache_selected_chunk_ids (`list[int]`, *optional*): Chunk ids to read cached KV from. + kv_cache_pre_rope (`bool`, *optional*, defaults to `False`): + If `True`, cached keys are stored before RoPE and re-rotated at read time using `cached_temporal_ids`. + self_attn_input_mode (`str`, *optional*): + Set to `SELF_ATTN_MODE_REF_IMAGE_CACHE` when prefilling the KV cache from a static reference image + (tags the current tokens with the `SOURCE_ID_EXTRA_REF_IMAGE` role instead of `SOURCE_ID_TARGET`). + skip_text_stream (`bool`, *optional*, defaults to `False`): + If `True`, the text stream is not projected/updated and attention only runs over the image stream (+ + any cached KV) -- used for KV-cache prefill calls where only the image-stream cache is wanted. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to wrap the image-stream sample into a [`~models.modeling_outputs.Transformer2DModelOutput`]. + Pass `return_dict=False` to return both image and text streams. + + Returns: + `(img, txt)` when `return_dict=False`; otherwise a [`~models.modeling_outputs.Transformer2DModelOutput`] + containing `img`. + """ + if kv_cache_mode not in (None, "store", "reuse"): + raise ValueError(f"Unsupported cache mode: {kv_cache_mode!r}.") + if kv_cache_mode == "store" and kv_cache_chunk_id is None: + raise ValueError("A cache chunk id is required in store mode.") + if self_attn_input_mode not in (None, SELF_ATTN_MODE_REF_IMAGE_CACHE): + raise ValueError(f"Unsupported self-attention input mode: {self_attn_input_mode!r}.") + + self.configure_inference_kv_cache( + chunk_id=kv_cache_chunk_id, + selected_chunk_ids=kv_cache_selected_chunk_ids, + pre_rope=kv_cache_pre_rope, + ) + + batch_size = hidden_states.shape[0] + patch_size = tuple(self.patch_size) + current_patch_shape = self._get_patch_shape(hidden_states, patch_size) + current_seq_len = math.prod(current_patch_shape) + device = hidden_states.device + + if encoder_hidden_states_mask is not None: + encoder_hidden_states_mask = encoder_hidden_states_mask.to( + device=encoder_hidden_states.device, dtype=torch.bool + ) + + hidden_tokens = self.img_in(hidden_states).flatten(2).transpose(1, 2).contiguous() + temporal_ids = None + if current_temporal_ids is not None: + current_temporal_ids = torch.as_tensor(current_temporal_ids, device=device, dtype=torch.long) + if current_temporal_ids.shape != (batch_size, current_patch_shape[0]): + raise ValueError( + f"`current_temporal_ids` must have shape {(batch_size, current_patch_shape[0])}, " + f"got {tuple(current_temporal_ids.shape)}." + ) + temporal_ids = current_temporal_ids + + if self_attn_input_mode == SELF_ATTN_MODE_REF_IMAGE_CACHE: + current_source_id = torch.full( + (current_seq_len,), SOURCE_ID_EXTRA_REF_IMAGE, device=device, dtype=torch.float32 + ) + else: + current_source_id = torch.full((current_seq_len,), SOURCE_ID_TARGET, device=device, dtype=torch.float32) + current_frame_ids = self._get_token_frame_ids(current_patch_shape, device, temporal_ids=temporal_ids) + current_rotary = self.get_rotary_pos_embed_from_ids( + frame_ids=current_frame_ids, + spatial_shape=(current_patch_shape[1], current_patch_shape[2]), + ) + + latent_segments = [hidden_tokens] + rotary_segments = [current_rotary] + source_id_segments = [current_source_id] + + if ref_video_latent is not None: + if ref_video_latent.shape[0] != batch_size: + raise ValueError( + f"Ref video latent batch size {ref_video_latent.shape[0]} does not match hidden states batch " + f"size {batch_size}." + ) + ref_video_patch_shape = self._get_patch_shape(ref_video_latent, patch_size) + if ref_video_patch_shape[1:] != current_patch_shape[1:]: + raise ValueError( + "Ref video latent spatial patch shape must match noisy latent spatial patch shape: " + f"{ref_video_patch_shape[1:]} != {current_patch_shape[1:]}." + ) + ref_video_tokens = self.img_in(ref_video_latent).flatten(2).transpose(1, 2).contiguous() + video_frame_ids = self._get_token_frame_ids(ref_video_patch_shape, device, temporal_ids=temporal_ids) + latent_segments.append(ref_video_tokens) + rotary_segments.append( + self.get_rotary_pos_embed_from_ids( + frame_ids=video_frame_ids, + spatial_shape=(ref_video_patch_shape[1], ref_video_patch_shape[2]), + ) + ) + source_id_segments.append( + torch.full((ref_video_tokens.shape[1],), SOURCE_ID_EDIT_CONDITION, device=device, dtype=torch.float32) + ) + + img = torch.cat(latent_segments, dim=1) + visual_source_id = torch.cat(source_id_segments, dim=0).unsqueeze(0) + # `torch.gather` requires an explicit index row for every batch element. + current_indices = torch.arange(current_seq_len, device=device).unsqueeze(0).expand(batch_size, -1) + vis_freqs_cis = ( + torch.cat([rotary[0] for rotary in rotary_segments], dim=1), + torch.cat([rotary[1] for rotary in rotary_segments], dim=1), + ) + + head_dim = self.hidden_size // self.num_attention_heads + cos_3d, sin_3d = vis_freqs_cis + cos_role, sin_role = self.generate_source_id_rope( + source_id=visual_source_id, + head_dim=head_dim, + device=cos_3d.device, + dtype=cos_3d.dtype, + ) + # Compose the 3D spatiotemporal rotation with the role rotation: rotating by (angle_3d + angle_role) is + # equivalent to cos(a+b) = cos(a)cos(b) - sin(a)sin(b), sin(a+b) = sin(a)cos(b) + cos(a)sin(b). + new_cos = cos_3d * cos_role - sin_3d * sin_role + new_sin = sin_3d * cos_role + cos_3d * sin_role + vis_freqs_cis = (new_cos, new_sin) + + _, vec, txt = self.condition_embedder(timestep, encoder_hidden_states) + vec = vec.unflatten(-1, (NUM_MODULATION_CHUNKS, -1)) + + cached_freqs_cis = None + if kv_cache_pre_rope and cached_temporal_ids is not None: + cached_ids_tensor = torch.as_tensor(cached_temporal_ids, device=device, dtype=torch.long) + if cached_ids_tensor.ndim != 2 or cached_ids_tensor.shape[0] != batch_size: + raise ValueError( + "Cached temporal ids must have shape (batch_size, num_cached_frames), got " + f"{tuple(cached_ids_tensor.shape)}." + ) + cached_frame_ids = self._get_token_frame_ids( + (cached_ids_tensor.shape[1], current_patch_shape[1], current_patch_shape[2]), + device, + temporal_ids=cached_ids_tensor, + ) + cached_freqs_cis = self.get_rotary_pos_embed_from_ids( + frame_ids=cached_frame_ids, + spatial_shape=(current_patch_shape[1], current_patch_shape[2]), + ) + + for layer_idx, block in enumerate(self.double_blocks): + kv_cache_reader = self._read_layer_kv_cache if kv_cache_mode == "reuse" else None + kv_cache_writer = self._write_layer_kv_cache if kv_cache_mode == "store" else None + if torch.is_grad_enabled() and self.gradient_checkpointing: + # `_gradient_checkpointing_func` only forwards positional args to `torch.utils.checkpoint.checkpoint`, + # so every block argument (including the kv-cache callables) must be passed positionally here in the + # same order as `JoyVideoEditTransformerBlock.forward`. + img, txt = self._gradient_checkpointing_func( + block, + img, + txt, + vec, + vis_freqs_cis, + encoder_hidden_states_mask, + kv_cache_reader, + kv_cache_writer, + layer_idx, + skip_text_stream, + kv_cache_pre_rope, + cached_freqs_cis, + ) + else: + img, txt = block( + hidden_states=img, + encoder_hidden_states=txt, + temb=vec, + image_rotary_emb=vis_freqs_cis, + encoder_hidden_states_mask=encoder_hidden_states_mask, + kv_cache_reader=kv_cache_reader, + kv_cache_writer=kv_cache_writer, + layer_idx=layer_idx, + skip_text_stream=skip_text_stream, + kv_cache_pre_rope=kv_cache_pre_rope, + cached_freqs_cis=cached_freqs_cis, + ) + + img = self.proj_out(self.norm_out(img)) + + gather_index = current_indices.unsqueeze(-1).expand(-1, -1, img.shape[-1]) + img = torch.gather(img, dim=1, index=gather_index) + img = self.unpatchify(img, current_patch_shape[0], current_patch_shape[1], current_patch_shape[2]) + + if not return_dict: + return (img, txt) + return Transformer2DModelOutput(sample=img) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 50052b0ca887..a318df659339 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -352,6 +352,10 @@ "JoyImageEditPlusPipeline", "JoyImageEditPlusPipelineOutput", ] + _import_structure["joyvideoedit"] = [ + "JoyVideoEditPipeline", + "JoyVideoEditPipelineOutput", + ] _import_structure["lumina"] = ["LuminaPipeline", "LuminaText2ImgPipeline"] _import_structure["lumina2"] = ["Lumina2Pipeline", "Lumina2Text2ImgPipeline"] _import_structure["lucy"] = ["LucyEditPipeline"] @@ -740,6 +744,7 @@ JoyImageEditPlusPipeline, JoyImageEditPlusPipelineOutput, ) + from .joyvideoedit import JoyVideoEditPipeline, JoyVideoEditPipelineOutput from .kandinsky import ( KandinskyCombinedPipeline, KandinskyImg2ImgCombinedPipeline, diff --git a/src/diffusers/pipelines/joyvideoedit/__init__.py b/src/diffusers/pipelines/joyvideoedit/__init__.py new file mode 100644 index 000000000000..7eb3864f8e6a --- /dev/null +++ b/src/diffusers/pipelines/joyvideoedit/__init__.py @@ -0,0 +1,48 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa: F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["pipeline_joyvideoedit"] = ["JoyVideoEditPipeline"] + _import_structure["pipeline_output"] = ["JoyVideoEditPipelineOutput"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * + else: + from .pipeline_joyvideoedit import JoyVideoEditPipeline + from .pipeline_output import JoyVideoEditPipelineOutput +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/pipelines/joyvideoedit/pipeline_joyvideoedit.py b/src/diffusers/pipelines/joyvideoedit/pipeline_joyvideoedit.py new file mode 100644 index 000000000000..c8f3ba0b638a --- /dev/null +++ b/src/diffusers/pipelines/joyvideoedit/pipeline_joyvideoedit.py @@ -0,0 +1,824 @@ +# Copyright 2026 The JoyAI-Video-Edit Team and 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. + +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import numpy as np +import PIL.Image +import torch +from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor, Qwen2Tokenizer + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...hooks import JoyVideoEditKVCacheConfig +from ...image_processor import PipelineImageInput +from ...models import AutoencoderKLJoyVideoEdit, JoyVideoEditTransformer3DModel +from ...models.transformers.transformer_joyvideoedit import SELF_ATTN_MODE_REF_IMAGE_CACHE +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .pipeline_output import JoyVideoEditPipelineOutput + + +logger = logging.get_logger(__name__) + + +EXAMPLE_DOC_STRING = """ + Examples: + ```python + >>> import torch + >>> from diffusers import JoyVideoEditPipeline + >>> from diffusers.utils import export_to_video, load_video + >>> from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + >>> model_id = "jdopensource/JoyAI-Video-Edit-Diffusers" + >>> 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) + >>> pipe = JoyVideoEditPipeline.from_pretrained( + ... model_id, + ... text_encoder=text_encoder, + ... processor=processor, + ... dtype=torch.bfloat16, + ... ) + >>> pipe.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 = pipe( + ... video=video, + ... prompt=prompt, + ... num_inference_steps=2, + ... generator=torch.Generator(device="cpu").manual_seed(0), + ... ) + >>> export_to_video(output.frames[0], "joyvideoedit.mp4", fps=24) + ``` +""" + + +class JoyVideoEditPipeline(DiffusionPipeline): + r""" + Pipeline for chunk-wise causal video editing using the JoyAI-Video-Edit architecture. + + The source video is VAE-encoded into a latent sequence that conditions a dual-stream MM-DiT transformer. The + transformer denoises the output latents 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, so later chunks + stay temporally consistent with earlier ones without recomputing their key/value projections. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + MiMo-VL is an external runtime dependency. Load `XiaomiMiMo/MiMo-VL-7B-RL-2508` separately and pass its model and + processor to this pipeline, or provide precomputed prompt embeddings. + + Args: + transformer ([`JoyVideoEditTransformer3DModel`]): + The streaming video-editing transformer that denoises the output latents. + vae ([`AutoencoderKLJoyVideoEdit`]): + Causal, chunk-streamable VAE to encode the source video and decode the edited latents. + text_encoder ([`Qwen2_5_VLForConditionalGeneration`], *optional*): + MiMo-VL model used to encode the prompt and first video frame. Load it from MiMo-VL's repository. Required + unless `prompt_embeds` are provided. + tokenizer ([`Qwen2Tokenizer`], *optional*): + Tokenizer paired with `text_encoder`. Defaults to the processor's tokenizer when omitted. + processor ([`Qwen2_5_VLProcessor`], *optional*): + MiMo-VL processor loaded from `XiaomiMiMo/MiMo-VL-7B-RL-2508`. Required unless `prompt_embeds` are + provided. + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + Flow-matching scheduler used to denoise each chunk. + """ + + model_cpu_offload_seq = "text_encoder->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + _optional_components = ["text_encoder", "tokenizer", "processor"] + + # Sentinel chunk id under which the static reference image's KV is prefilled (never a real chunk index). + _KV_CACHE_ID_REF_IMAGE = -1 + + def __init__( + self, + transformer: JoyVideoEditTransformer3DModel, + vae: AutoencoderKLJoyVideoEdit, + text_encoder: Optional[Qwen2_5_VLForConditionalGeneration], + tokenizer: Optional[Qwen2Tokenizer], + processor: Optional[Qwen2_5_VLProcessor], + scheduler: FlowMatchEulerDiscreteScheduler, + ): + super().__init__() + + if tokenizer is None and processor is not None: + tokenizer = processor.tokenizer + + self.register_modules( + transformer=transformer, + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + processor=processor, + scheduler=scheduler, + ) + + # The KV cache is required for chunk-wise denoising. + if getattr(self, "transformer", None) is not None and not self.transformer.is_cache_enabled: + self.transformer.enable_cache(JoyVideoEditKVCacheConfig()) + + self.vae_scale_factor_spatial = self.vae.spatial_compression_ratio if getattr(self, "vae", None) else 16 + self.vae_scale_factor_temporal = self.vae.temporal_compression_ratio if getattr(self, "vae", None) else 8 + transformer_patch_size = ( + self.transformer.config.patch_size if getattr(self, "transformer", None) else (1, 1, 1) + ) + self.height_multiple = self.vae_scale_factor_spatial * transformer_patch_size[1] + self.width_multiple = self.vae_scale_factor_spatial * transformer_patch_size[2] + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + # Encode the prompt and the video's first frame with the image-description template. + self.prompt_template_encode = ( + "<|im_start|>system\n \\nDescribe the image by detailing the color, shape, size, texture, quantity, " + "text, spatial relationships of the objects and background:<|im_end|>\n" + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n" + ) + self.prompt_template_encode_start_idx = None + if self.tokenizer is not None: + prefix_ids = self.tokenizer(self.prompt_template_encode.split("{}")[0]).input_ids + user_id = self.tokenizer.convert_tokens_to_ids("user") + self.prompt_template_encode_start_idx = prefix_ids.index(user_id) + + # The anchor frame is resized to a fixed ViT input area before being packed by the processor. + self.vit_input_size = 512 + + # ------------------------------------------------------------------ + # Prompt encoding (multimodal: text + video's first frame as an image anchor) + # ------------------------------------------------------------------ + + # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._extract_masked_hidden + def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): + bool_mask = mask.bool() + valid_lengths = bool_mask.sum(dim=1) + selected = hidden_states[bool_mask] + split_result = torch.split(selected, valid_lengths.tolist(), dim=0) + + return split_result + + def _get_qwen_prompt_embeds( + self, + prompt: Union[str, List[str]], + image: PIL.Image.Image, + device: torch.device, + max_sequence_length: int, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + missing_components = [ + name for name in ("text_encoder", "tokenizer", "processor") if getattr(self, name, None) is None + ] + if missing_components: + raise ValueError( + f"Missing MiMo-VL components: {', '.join(missing_components)}. Load them from " + "`XiaomiMiMo/MiMo-VL-7B-RL-2508` and pass them to `JoyVideoEditPipeline.from_pretrained`, or pass " + "precomputed `prompt_embeds` and `prompt_embeds_mask`." + ) + + prompt = [prompt] if isinstance(prompt, str) else prompt + + if self.prompt_template_encode_start_idx is None: + prefix_ids = self.tokenizer(self.prompt_template_encode.split("{}")[0]).input_ids + user_id = self.tokenizer.convert_tokens_to_ids("user") + self.prompt_template_encode_start_idx = prefix_ids.index(user_id) + drop_idx = self.prompt_template_encode_start_idx + txt = [self.prompt_template_encode.format(e) for e in prompt] + + # Resize the anchor image to the fixed ViT input area while preserving its aspect ratio. + target_area = self.vit_input_size * self.vit_input_size + scale = (target_area / max(image.height * image.width, 1)) ** 0.5 + new_h = max(1, round(image.height * scale)) + new_w = max(1, round(image.width * scale)) + anchor = image.convert("RGB").resize((new_w, new_h), PIL.Image.BILINEAR) + + model_inputs = self.processor(text=txt, images=[anchor] * len(txt), padding=True, return_tensors="pt").to( + device + ) + # Forward all multimodal fields required for position encoding. + outputs = self.text_encoder(**model_inputs, output_hidden_states=True) + hidden_states = outputs.hidden_states[-1] + + # Remove padding and the template prefix, keep the most recent tokens, then left-align and pad the batch. + split_hidden_states = self._extract_masked_hidden(hidden_states, model_inputs["attention_mask"]) + split_hidden_states = [e[drop_idx:][-max_sequence_length:] for e in split_hidden_states] + attn_mask_list = [e.new_ones(e.size(0), dtype=torch.long) for e in split_hidden_states] + max_seq_len = max(e.size(0) for e in split_hidden_states) + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states] + ) + prompt_embeds_mask = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list] + ) + return prompt_embeds, prompt_embeds_mask + + def encode_prompt( + self, + prompt: Union[str, List[str]], + image: Optional[PIL.Image.Image] = None, + device: Optional[torch.device] = None, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds_mask: Optional[torch.Tensor] = None, + max_sequence_length: int = 1024, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + r""" + Encode a text prompt together with the video's first frame into multimodal embeddings. + + Args: + prompt (`str` or `List[str]`): Prompt(s) to encode. + image (`PIL.Image.Image`, *optional*): Anchor image (the video's first frame) spliced into the prompt as + an `` token. Required unless `prompt_embeds` are provided. + device (`torch.device`, *optional*): Target device. + prompt_embeds (`torch.Tensor`, *optional*): Pre-computed embeddings that bypass encoding. + prompt_embeds_mask (`torch.Tensor`, *optional*): Attention mask for pre-computed embeddings. + max_sequence_length (`int`, *optional*, defaults to 1024): Maximum prompt length. + + Returns: + Tuple of `(prompt_embeds, prompt_embeds_mask)`. + """ + device = device or self._execution_device + if prompt_embeds is None: + if image is None: + raise ValueError("`image` (the video's first frame) is required to encode a `prompt`.") + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds( + prompt, image, device, max_sequence_length + ) + prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype) + prompt_embeds = prompt_embeds[:, :max_sequence_length].to(device=device) + prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length].to(device=device) + # The reference DiT uses unmasked attention for fully valid prompt sequences. Preserve that numerical path and + # keep an explicit mask only when padding is present. + if prompt_embeds_mask.all(): + prompt_embeds_mask = None + return prompt_embeds, prompt_embeds_mask + + # ------------------------------------------------------------------ + # Latent (de)normalization + # ------------------------------------------------------------------ + + def normalize_latents(self, latent: torch.Tensor) -> torch.Tensor: + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(device=latent.device, dtype=latent.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(device=latent.device, dtype=latent.dtype) + ) + return (latent - latents_mean) / latents_std + + def denormalize_latents(self, latent: torch.Tensor) -> torch.Tensor: + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(device=latent.device, dtype=latent.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(device=latent.device, dtype=latent.dtype) + ) + return latent * latents_std + latents_mean + + # ------------------------------------------------------------------ + # Chunk and temporal-id helpers + # ------------------------------------------------------------------ + + def _kv_cache_memory_id(self, kind: str, chunk_id: Optional[int] = None) -> int: + if kind == "clean": + if chunk_id is None: + raise ValueError("`chunk_id` is required for clean cache ids.") + return int(chunk_id) + if kind == "ref_image": + return self._KV_CACHE_ID_REF_IMAGE + raise ValueError(f"Unsupported cache kind: {kind!r}") + + @staticmethod + def _get_chunk_windows( + total_latent_frames: int, + chunk_size: int, + window_size: int, + global_sink_chunk: bool, + ) -> List[Dict[str, Any]]: + if window_size <= 0: + raise ValueError(f"`window_size` must be positive, got {window_size}.") + + windows = [] + num_chunks = (total_latent_frames + chunk_size - 1) // chunk_size + for chunk_idx in range(num_chunks): + chunk_start = chunk_idx * chunk_size + chunk_end = min(total_latent_frames, chunk_start + chunk_size) + if global_sink_chunk and chunk_idx > 0: + tail_window_size = max(window_size - 1, 1) + tail_chunk_start = max(1, chunk_idx - tail_window_size + 1) + selected_chunk_ids = [0] + list(range(tail_chunk_start, chunk_idx + 1)) + else: + window_chunk_start = max(0, chunk_idx - window_size + 1) + selected_chunk_ids = list(range(window_chunk_start, chunk_idx + 1)) + + windows.append( + { + "chunk_start": chunk_start, + "chunk_end": chunk_end, + "selected_chunk_ids": selected_chunk_ids, + } + ) + return windows + + @staticmethod + def _chunk_frame_bounds(chunk_id: int, chunk_size: int, total_latent_frames: int) -> Tuple[int, int]: + chunk_start = chunk_id * chunk_size + chunk_end = min(total_latent_frames, chunk_start + chunk_size) + return chunk_start, chunk_end + + @classmethod + def _gather_window_temporal_ids( + cls, + selected_chunk_ids: List[int], + chunk_size: int, + total_latent_frames: int, + device: torch.device, + ) -> torch.Tensor: + temporal_ids = [] + offset = 0 + for cid in selected_chunk_ids: + frame_start, frame_end = cls._chunk_frame_bounds(cid, chunk_size, total_latent_frames) + chunk_len = frame_end - frame_start + temporal_ids.append(torch.arange(offset, offset + chunk_len, device=device, dtype=torch.long)) + offset += chunk_len + return torch.cat(temporal_ids, dim=0) + + # ------------------------------------------------------------------ + # Input validation + # ------------------------------------------------------------------ + + def check_inputs( + self, + video, + prompt, + height, + width, + num_inference_steps, + prompt_embeds=None, + prompt_embeds_mask=None, + chunk_size=None, + callback_on_step_end_tensor_inputs=None, + ): + if not isinstance(video, list) or len(video) == 0: + raise ValueError("`video` must be a non-empty list of PIL images.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError("`callback_on_step_end_tensor_inputs` has invalid keys.") + + if height <= 0 or width <= 0: + raise ValueError(f"`height` and `width` must be positive but are {height} and {width}.") + if height % self.height_multiple != 0 or width % self.width_multiple != 0: + raise ValueError( + f"`height` and `width` have to be divisible by {self.height_multiple} and {self.width_multiple} but " + f"are {height} and {width}." + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError("Cannot forward both `prompt` and `prompt_embeds`.") + if prompt is None and prompt_embeds is None: + raise ValueError("Provide either `prompt` or `prompt_embeds`.") + if prompt is not None and not isinstance(prompt, (str, list)): + raise ValueError("`prompt` has to be of type `str` or `list`.") + if prompt_embeds is not None and prompt_embeds_mask is None: + raise ValueError("If `prompt_embeds` are provided, `prompt_embeds_mask` is required.") + + if chunk_size is not None and chunk_size <= 0: + raise ValueError(f"`chunk_size` must be positive when provided, got {chunk_size}.") + if not isinstance(num_inference_steps, int) or num_inference_steps <= 0: + raise ValueError(f"`num_inference_steps` must be a positive integer, got {num_inference_steps}.") + + # ------------------------------------------------------------------ + # Pipeline properties + # ------------------------------------------------------------------ + + @property + def num_timesteps(self) -> int: + return self._num_timesteps + + @property + def interrupt(self) -> bool: + return self._interrupt + + # ------------------------------------------------------------------ + # Forward pass + # ------------------------------------------------------------------ + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + video: List[PIL.Image.Image] = None, + prompt: Union[str, List[str]] = None, + ref_image: Optional[PipelineImageInput] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 2, + chunk_size: Optional[int] = None, + local_window_size: Optional[int] = None, + global_sink_chunk: Optional[bool] = None, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + prompt_embeds_mask: Optional[torch.Tensor] = None, + output_type: Optional[str] = "np", + return_dict: bool = True, + callback_on_step_end: Optional[ + Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] + ] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + max_sequence_length: int = 1024, + ): + r""" + The call method of the pipeline for chunk-wise causal video editing. + + Args: + video (`List[PIL.Image.Image]`): + The source video to edit, as a sequence of frames. If its length does not satisfy + `temporal_compression_ratio * n + 1` for some integer `n`, trailing frames are truncated with a + warning. + prompt (`str` or `List[str]`): + The prompt describing the desired edited video. + ref_image (`PipelineImageInput`, *optional*): + An optional static reference image whose KV is prefilled into the cache and attended to by every chunk + (used to inject appearance/identity conditioning). + height (`int`, *optional*): + The height in pixels of the generated video. Defaults to the source video height and is adjusted down + to the nearest valid spatial multiple when needed. + width (`int`, *optional*): + The width in pixels of the generated video. Defaults to the source video width and is adjusted down to + the nearest valid spatial multiple when needed. + num_inference_steps (`int`, *optional*, defaults to 2): + The number of denoising steps applied per chunk. + chunk_size (`int`, *optional*): + Number of latent frames denoised per chunk. Defaults to `transformer.config.chunk_size`. + local_window_size (`int`, *optional*): + Number of recent chunks each chunk attends to. Defaults to `transformer.config.local_window_size`. + global_sink_chunk (`bool`, *optional*): + Whether every chunk additionally attends to chunk 0 (a global "sink"). Defaults to + `transformer.config.global_sink_chunk`. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A generator to make generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents for the first chunk, sampled from a Gaussian distribution when not + provided. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-computed text embeddings. When provided, `prompt` can be omitted. + prompt_embeds_mask (`torch.Tensor`, *optional*): + Attention mask for `prompt_embeds`. + output_type (`str`, *optional*, defaults to `"np"`): + The output format of the generated video. Choose between `"np"`, `"pt"`, `"pil"`, or `"latent"`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~pipelines.joyvideoedit.JoyVideoEditPipelineOutput`] instead of a plain tuple. + callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A callback invoked at the end of each denoising step. + callback_on_step_end_tensor_inputs (`List[str]`, *optional*, defaults to `["latents"]`): + Tensor keys passed to `callback_on_step_end`. + max_sequence_length (`int`, *optional*, defaults to 1024): + Maximum sequence length for prompt encoding. + + Examples: + + Returns: + [`~pipelines.joyvideoedit.JoyVideoEditPipelineOutput`] or `tuple`: + If `return_dict` is `True`, a [`~pipelines.joyvideoedit.JoyVideoEditPipelineOutput`] is returned, + otherwise a `tuple` where the first element is the generated frames. + """ + if not isinstance(video, list) or len(video) == 0: + raise ValueError("`video` must be a non-empty list of PIL images.") + height = height if height is not None else video[0].height + width = width if width is not None else video[0].width + if height < self.height_multiple or width < self.width_multiple: + raise ValueError( + f"`height` and `width` must be at least {self.height_multiple} and {self.width_multiple} but are " + f"{height} and {width}." + ) + adjusted_height = height // self.height_multiple * self.height_multiple + adjusted_width = width // self.width_multiple * self.width_multiple + if height != adjusted_height or width != adjusted_width: + logger.warning( + f"`height` and `width` must be multiples of ({self.height_multiple}, {self.width_multiple}). " + f"Adjusting ({height}, {width}) to ({adjusted_height}, {adjusted_width})." + ) + height, width = adjusted_height, adjusted_width + + self.check_inputs( + video, + prompt, + height, + width, + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + chunk_size=chunk_size, + num_inference_steps=num_inference_steps, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + + self._interrupt = False + device = self._execution_device + transformer_dtype = self.transformer.dtype + + chunk_size = chunk_size if chunk_size is not None else self.transformer.config.chunk_size + local_window_size = ( + local_window_size if local_window_size is not None else self.transformer.config.local_window_size + ) + global_sink_chunk = ( + global_sink_chunk if global_sink_chunk is not None else self.transformer.config.global_sink_chunk + ) + if chunk_size is None or chunk_size <= 0: + raise ValueError(f"`chunk_size` must resolve to a positive value, got {chunk_size}.") + + # 1. Encode prompt together with the video's first frame as an image anchor (no CFG, so no negative prompt). + prompt_embeds, prompt_embeds_mask = self.encode_prompt( + prompt=prompt, + image=video[0] if video is not None else None, + device=device, + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + max_sequence_length=max_sequence_length, + ) + prompt_embeds = prompt_embeds.to(transformer_dtype) + + # 2. Encode each conditioning latent frame from a causal pixel window and keep the window's final latent. + video_tensor = self.video_processor.preprocess_video(video, height=height, width=width) + video_tensor = video_tensor.to(device=device, dtype=self.vae.dtype) + + ffactor_temporal = self.vae_scale_factor_temporal + total_pixel_frames = video_tensor.shape[2] + valid_num_frames = (total_pixel_frames - 1) // ffactor_temporal * ffactor_temporal + 1 + if total_pixel_frames != valid_num_frames: + logger.warning( + f"Video contains {total_pixel_frames} frames, but its length must be of the form " + f"`k * {ffactor_temporal} + 1`. Truncating to {valid_num_frames} frames." + ) + video_tensor = video_tensor[:, :, :valid_num_frames] + total_pixel_frames = valid_num_frames + # The window size is fixed by the transformer's configured `chunk_size`, independent of any per-call + # `chunk_size` override that only affects the denoising chunk layout. + vae_chunk_size = self.transformer.config.chunk_size + window_pixels = vae_chunk_size * ffactor_temporal + window_frames = 1 + window_pixels + stride = ffactor_temporal + num_latents = (total_pixel_frames - 1) // stride + 1 + + latent_frames = [] + for k in range(num_latents): + if k == 0: + window = video_tensor[:, :, :1] + else: + end_frame = k * stride + start_frame = max(0, end_frame - window_pixels) + window = video_tensor[:, :, start_frame : end_frame + 1] + pad_needed = window_frames - window.shape[2] + if pad_needed > 0: + pad = video_tensor[:, :, :1].expand(-1, -1, pad_needed, -1, -1) + window = torch.cat([pad, window], dim=2) + del pad + window_latents = self.vae.encode(window).latent_dist.sample(generator=generator) + latent_frames.append(window_latents[:, :, -1:]) + ref_video_latents = torch.cat(latent_frames, dim=2) + ref_video_latents = self.normalize_latents(ref_video_latents).to(device=device, dtype=transformer_dtype) + del video_tensor, latent_frames, window, window_latents + + # A single source video conditions every prompt in the batch, so broadcast its latents to the number of + # prompts (`prompt=["edit A", "edit B"]` edits the same video two ways). + num_prompts = prompt_embeds.shape[0] + if ref_video_latents.shape[0] == 1 and num_prompts > 1: + ref_video_latents = ref_video_latents.repeat(num_prompts, 1, 1, 1, 1) + + batch_size, latent_channels, total_latent_frames, latent_height, latent_width = ref_video_latents.shape + + # 3. Optionally prefill the KV cache with a static reference image. + self.transformer._reset_stateful_cache() + ref_image_kv_prefilled = False + try: + if ref_image is not None: + ref_pixels = self.video_processor.preprocess( + ref_image, + height=latent_height * self.vae_scale_factor_spatial, + width=latent_width * self.vae_scale_factor_spatial, + ) + ref_pixels = ref_pixels.unsqueeze(2).to(device=device, dtype=self.vae.dtype) # (B, C, 1, H, W) + reference_image_latents = self.vae.encode(ref_pixels).latent_dist.sample(generator=generator) + reference_image_latents = self.normalize_latents(reference_image_latents) + reference_image_latents = reference_image_latents[:, :, :1].to(device=device, dtype=transformer_dtype) + if reference_image_latents.shape[0] == 1 and batch_size > 1: + reference_image_latents = reference_image_latents.repeat(batch_size, 1, 1, 1, 1) + ref_frames = reference_image_latents.shape[2] + with self.transformer.cache_context("inference"): + self.transformer( + hidden_states=reference_image_latents, + timestep=torch.zeros( + (reference_image_latents.shape[0],), device=device, dtype=transformer_dtype + ), + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_mask=prompt_embeds_mask, + current_temporal_ids=torch.zeros( + (reference_image_latents.shape[0], ref_frames), device=device, dtype=torch.long + ), + kv_cache_mode="store", + kv_cache_chunk_id=self._kv_cache_memory_id("ref_image"), + kv_cache_selected_chunk_ids=[], + self_attn_input_mode=SELF_ATTN_MODE_REF_IMAGE_CACHE, + skip_text_stream=True, + return_dict=False, + ) + ref_image_kv_prefilled = True + + # 4. Set up chunk-wise causal denoising while cache cleanup is still guarded. + windows = self._get_chunk_windows(total_latent_frames, chunk_size, local_window_size, global_sink_chunk) + num_chunks = len(windows) + self._num_timesteps = num_inference_steps + + # The first chunk may be seeded with user-provided `latents`; later chunks always start from fresh noise. + initial_latents = latents + raw_sigmas = torch.linspace(1, 0, num_inference_steps + 1)[:-1].numpy() + chunk_outputs = [] + except Exception: + self.transformer._reset_stateful_cache() + raise + + try: + for chunk_idx, window in enumerate(self.progress_bar(windows)): + chunk_start = window["chunk_start"] + chunk_end = window["chunk_end"] + selected_chunk_ids = window["selected_chunk_ids"] + history_chunk_ids = selected_chunk_ids[:-1] + active_chunk_id = selected_chunk_ids[-1] + current_chunk_len = chunk_end - chunk_start + + window_ids = self._gather_window_temporal_ids( + selected_chunk_ids, chunk_size, total_latent_frames, device + ) + current_temporal_ids = window_ids[-current_chunk_len:] + cached_temporal_ids = window_ids[:-current_chunk_len] + if cached_temporal_ids.numel() == 0: + cached_temporal_ids = None + + ref_chunk_latent = ref_video_latents[:, :, chunk_start:chunk_end] + + noise_shape = (batch_size, latent_channels, current_chunk_len, latent_height, latent_width) + # Keep Euler updates in float32 and cast only the transformer input. + if chunk_idx == 0 and initial_latents is not None: + latents = initial_latents.to(device=device, dtype=torch.float32) + else: + latents = randn_tensor(noise_shape, generator=generator, device=device, dtype=torch.float32) + + cache_memory_ids = [self._kv_cache_memory_id("clean", cid) for cid in history_chunk_ids] + if ref_image_kv_prefilled: + cache_memory_ids.append(self._kv_cache_memory_id("ref_image")) + + current_ids_batched = current_temporal_ids.unsqueeze(0).expand(batch_size, -1) + cached_ids_batched = ( + cached_temporal_ids.unsqueeze(0).expand(batch_size, -1) + if cached_temporal_ids is not None + else None + ) + + # Reset the scheduler for each independently denoised chunk. + self.scheduler.set_timesteps(sigmas=raw_sigmas, device=device) + timesteps = self.scheduler.timesteps + for i, t in enumerate(timesteps): + if self.interrupt: + continue + t_expand = t.repeat(latents.shape[0]) + with self.transformer.cache_context("inference"): + noise_pred = self.transformer( + hidden_states=latents.to(transformer_dtype), + timestep=t_expand, + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_mask=prompt_embeds_mask, + ref_video_latent=ref_chunk_latent, + current_temporal_ids=current_ids_batched, + cached_temporal_ids=cached_ids_batched, + kv_cache_mode="reuse", + kv_cache_chunk_id=active_chunk_id, + kv_cache_selected_chunk_ids=cache_memory_ids, + kv_cache_pre_rope=True, + return_dict=False, + )[0] + # Upcast the model output to keep the scheduler accumulator in float32. + latents = self.scheduler.step(noise_pred.float(), t, latents, return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + # Evict every cached chunk that neither the current history nor the reference image needs, then store + # this chunk's clean (denoised) KV so later chunks can attend to it. + keep_before_store = {self._kv_cache_memory_id("clean", cid) for cid in history_chunk_ids} + if ref_image_kv_prefilled: + keep_before_store.add(self._kv_cache_memory_id("ref_image")) + + # Keep only the chunks the next window will attend to (plus the reference image). + next_selected = windows[chunk_idx + 1]["selected_chunk_ids"] if chunk_idx + 1 < num_chunks else [] + keep_after_store = {self._kv_cache_memory_id("clean", cid) for cid in next_selected} + if ref_image_kv_prefilled: + keep_after_store.add(self._kv_cache_memory_id("ref_image")) + + # `evict_kv_cache_chunks` and the store-mode forward both read/write cache state through the KV-cache + # hook's `StateManager`, which requires an active context. + with self.transformer.cache_context("inference"): + self.transformer.evict_kv_cache_chunks(keep_before_store) + self.transformer( + hidden_states=latents.to(transformer_dtype), + timestep=torch.zeros((latents.shape[0],), device=device, dtype=transformer_dtype), + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_mask=prompt_embeds_mask, + current_temporal_ids=current_ids_batched, + kv_cache_mode="store", + kv_cache_chunk_id=self._kv_cache_memory_id("clean", active_chunk_id), + kv_cache_selected_chunk_ids=[], + kv_cache_pre_rope=True, + skip_text_stream=True, + return_dict=False, + ) + self.transformer.evict_kv_cache_chunks(keep_after_store) + + chunk_outputs.append(latents) + finally: + # Never let KV-cache state leak into the next `__call__`. + self.transformer._reset_stateful_cache() + + latents = torch.cat(chunk_outputs, dim=2) + del chunk_outputs, ref_video_latents, ref_chunk_latent, prompt_embeds, prompt_embeds_mask, initial_latents + + if output_type == "latent": + video = latents + else: + # Decode each chunk causally. Subsequent chunks prepend a latent encoded from the previous output frame. + latents = self.denormalize_latents(latents.to(self.vae.dtype)) + total_decoded_frames = 1 + (latents.shape[2] - 1) * ffactor_temporal + previous_frame = None + video = None + output_frame_start = 0 + for frame_start in range(0, latents.shape[2], chunk_size): + chunk_latents = latents[:, :, frame_start : frame_start + chunk_size] + if frame_start == 0: + decoded = self.vae.decode(chunk_latents, return_dict=False)[0] + else: + pseudo_latent = self.vae.encode(previous_frame).latent_dist.sample(generator=generator) + decoded = self.vae.decode(torch.cat([pseudo_latent, chunk_latents], dim=2), return_dict=False)[0] + decoded = decoded[:, :, -chunk_latents.shape[2] * ffactor_temporal :] + + previous_frame = decoded[:, :, -1:].clone() + num_decoded_frames = decoded.shape[2] + decoded = self.video_processor.postprocess_video(decoded, output_type=output_type) + if output_type == "pt": + decoded = decoded.cpu() + + if output_type == "pil": + if video is None: + video = [[] for _ in range(len(decoded))] + for batch_idx, frames in enumerate(decoded): + video[batch_idx].extend(frames) + else: + if video is None: + output_shape = (decoded.shape[0], total_decoded_frames, *decoded.shape[2:]) + if output_type == "np": + video = np.empty(output_shape, dtype=decoded.dtype) + else: + video = torch.empty(output_shape, dtype=decoded.dtype, device="cpu") + video[:, output_frame_start : output_frame_start + num_decoded_frames] = decoded + output_frame_start += num_decoded_frames + + if output_type != "pil": + video = video[:, :output_frame_start] + + self.maybe_free_model_hooks() + + if not return_dict: + return (video,) + return JoyVideoEditPipelineOutput(frames=video) diff --git a/src/diffusers/pipelines/joyvideoedit/pipeline_output.py b/src/diffusers/pipelines/joyvideoedit/pipeline_output.py new file mode 100644 index 000000000000..ebf5997d2af1 --- /dev/null +++ b/src/diffusers/pipelines/joyvideoedit/pipeline_output.py @@ -0,0 +1,20 @@ +from dataclasses import dataclass + +import torch + +from ...utils import BaseOutput + + +@dataclass +class JoyVideoEditPipelineOutput(BaseOutput): + r""" + Output class for JoyVideoEdit video-editing pipelines. + + Args: + frames (`torch.Tensor`, `np.ndarray`, or `List[List[PIL.Image.Image]]`): + List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing + denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape + `(batch_size, num_frames, channels, height, width)`. + """ + + frames: torch.Tensor diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 04bb0ceb7143..b1a9c1c6355d 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -212,6 +212,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class JoyVideoEditKVCacheConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class LayerSkipConfig(metaclass=DummyObject): _backends = ["torch"] @@ -310,6 +325,10 @@ def apply_first_block_cache(*args, **kwargs): requires_backends(apply_first_block_cache, ["torch"]) +def apply_joyvideoedit_kv_cache(*args, **kwargs): + requires_backends(apply_joyvideoedit_kv_cache, ["torch"]) + + def apply_layer_skip(*args, **kwargs): requires_backends(apply_layer_skip, ["torch"]) @@ -675,6 +694,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class AutoencoderKLJoyVideoEdit(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class AutoencoderKLKVAE(metaclass=DummyObject): _backends = ["torch"] @@ -1545,6 +1579,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class JoyVideoEditTransformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class Kandinsky3UNet(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index ebe36d242253..b5f1c45ee895 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -2582,6 +2582,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class JoyVideoEditPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class JoyVideoEditPipelineOutput(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class Kandinsky3Img2ImgPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/autoencoders/test_models_autoencoder_kl_joyvideoedit.py b/tests/models/autoencoders/test_models_autoencoder_kl_joyvideoedit.py new file mode 100644 index 000000000000..0488a298a1ee --- /dev/null +++ b/tests/models/autoencoders/test_models_autoencoder_kl_joyvideoedit.py @@ -0,0 +1,137 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# 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. + +import pytest +import torch +import torch.nn.functional as F + +from diffusers import AutoencoderKLJoyVideoEdit +from diffusers.models.autoencoders.autoencoder_kl_joyvideoedit import JoyVideoEditAttentionBlock +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, +) + + +enable_full_determinism() + + +class AutoencoderKLJoyVideoEditTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return AutoencoderKLJoyVideoEdit + + @property + def main_input_name(self) -> str: + return "sample" + + @property + def output_shape(self) -> tuple[int, ...]: + return (3, 5, 24, 24) + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict[str, int | list[int]]: + return { + "in_channels": 3, + "out_channels": 3, + "patch_size": 1, + "latent_channels": 4, + "layers_per_block": 1, + "block_in_channels": (8, 8), + "temporal_downsample": (True, False), + "chunk_size": 4, + "latents_mean": [0.0, 0.0, 0.0, 0.0], + "latents_std": [1.0, 1.0, 1.0, 1.0], + } + + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + batch_size = 1 + num_channels = 3 + num_frames = 5 # temporal_compression_ratio * n + 1 + sizes = (24, 24) # divisible by the spatial compression ratio (3 in this dummy config) + image = randn_tensor( + (batch_size, num_channels, num_frames, *sizes), generator=self.generator, device=torch_device + ) + return {"sample": image} + + +class TestAutoencoderKLJoyVideoEditModel(AutoencoderKLJoyVideoEditTesterConfig, ModelTesterMixin): + base_precision = 1e-2 + + @pytest.mark.parametrize("chunk_size", [0, 1]) + def test_invalid_chunk_size_raises(self, chunk_size): + init_dict = self.get_init_dict() + init_dict["chunk_size"] = chunk_size + + with pytest.raises(ValueError, match="positive multiple of the temporal compression ratio"): + self.model_class(**init_dict) + + def test_last_temporal_downsample_must_be_false(self): + init_dict = self.get_init_dict() + init_dict["temporal_downsample"] = (True, True) + + with pytest.raises(ValueError, match="last value must be `False`"): + self.model_class(**init_dict) + + def test_temporal_compression_ratio_one(self): + init_dict = self.get_init_dict() + init_dict["temporal_downsample"] = (False, False) + model = self.model_class(**init_dict).to(torch_device) + sample = randn_tensor((1, 3, 1, 24, 24), generator=self.generator, device=torch_device) + + latents = model.encode(sample).latent_dist.mode() + + assert latents.shape[2] == 1 + + def test_attention_processor_matches_sdpa(self): + attention = JoyVideoEditAttentionBlock(8).to(torch_device) + hidden_states = randn_tensor((2, 8, 3, 4, 4), generator=self.generator, device=torch_device) + + actual = attention(hidden_states) + + normed_hidden_states = attention.norm(hidden_states) + batch_size, channels, num_frames, height, width = hidden_states.shape + query = attention.q(normed_hidden_states) + key = attention.k(normed_hidden_states) + value = attention.v(normed_hidden_states) + query = query.permute(0, 2, 3, 4, 1).reshape(batch_size * num_frames, 1, height * width, channels) + key = key.permute(0, 2, 3, 4, 1).reshape(batch_size * num_frames, 1, height * width, channels) + value = value.permute(0, 2, 3, 4, 1).reshape(batch_size * num_frames, 1, height * width, channels) + expected = F.scaled_dot_product_attention(query, key, value) + expected = expected.reshape(batch_size, num_frames, height, width, channels).permute(0, 4, 1, 2, 3) + expected = hidden_states + attention.proj_out(expected) + + torch.testing.assert_close(actual, expected, atol=2e-4, rtol=2e-4) + + +class TestAutoencoderKLJoyVideoEditMemory(AutoencoderKLJoyVideoEditTesterConfig, MemoryTesterMixin): + pass + + +class TestAutoencoderKLJoyVideoEditCompile(AutoencoderKLJoyVideoEditTesterConfig, TorchCompileTesterMixin): + pass + + +class TestAutoencoderKLJoyVideoEditAttention(AutoencoderKLJoyVideoEditTesterConfig, AttentionTesterMixin): + pass diff --git a/tests/models/transformers/test_models_transformer_joyvideoedit.py b/tests/models/transformers/test_models_transformer_joyvideoedit.py new file mode 100644 index 000000000000..2a3dc919fcc7 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_joyvideoedit.py @@ -0,0 +1,143 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# 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. + +import pytest +import torch + +from diffusers import JoyVideoEditTransformer3DModel +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class JoyVideoEditTransformerTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return JoyVideoEditTransformer3DModel + + @property + def main_input_name(self) -> str: + return "hidden_states" + + @property + def uses_custom_attn_processor(self) -> bool: + return True + + @property + def input_shape(self) -> tuple[int, ...]: + return (8, 2, 4, 4) + + @property + def output_shape(self) -> tuple[int, ...]: + return (8, 2, 4, 4) + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict[str, int | list[int]]: + return { + "patch_size": [1, 1, 1], + "in_channels": 8, + "out_channels": 8, + "hidden_size": 32, + "num_attention_heads": 2, + "text_dim": 16, + "num_layers": 2, + "rope_dim_list": [4, 6, 6], + "theta": 256, + "chunk_size": 1, + "local_window_size": 3, + "global_sink_chunk": True, + "source_id_rope_dim": 4, + "source_id_rope_theta": 256.0, + } + + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + batch_size = 1 + num_frames, height, width = 2, 4, 4 + hidden_states = randn_tensor( + (batch_size, 8, num_frames, height, width), generator=self.generator, device=torch_device + ) + encoder_hidden_states = randn_tensor((batch_size, 8, 16), generator=self.generator, device=torch_device) + encoder_hidden_states_mask = torch.ones(batch_size, 8, dtype=torch.bool, device=torch_device) + encoder_hidden_states_mask[:, -1] = False + timestep = torch.tensor([1.0]).to(torch_device).expand(batch_size) + return { + "hidden_states": hidden_states, + "timestep": timestep, + "encoder_hidden_states": encoder_hidden_states, + "encoder_hidden_states_mask": encoder_hidden_states_mask, + } + + +class TestJoyVideoEditTransformerModel(JoyVideoEditTransformerTesterConfig, ModelTesterMixin): + def test_invalid_self_attn_input_mode_raises(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + + with pytest.raises(ValueError, match="Unsupported self-attention input mode"): + model(**self.get_dummy_inputs(), self_attn_input_mode="invalid") + + def test_temporal_ids_are_applied_per_batch_element(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + inputs["hidden_states"] = inputs["hidden_states"].expand(2, -1, -1, -1, -1).clone() + inputs["encoder_hidden_states"] = inputs["encoder_hidden_states"].expand(2, -1, -1).clone() + inputs["encoder_hidden_states_mask"] = inputs["encoder_hidden_states_mask"].expand(2, -1).clone() + inputs["timestep"] = inputs["timestep"].expand(2).clone() + inputs["current_temporal_ids"] = torch.tensor([[0, 1], [0, 3]], device=torch_device) + + with torch.no_grad(): + batch_output = model(**inputs).sample + first_output = model( + **{key: value[:1] for key, value in inputs.items() if key != "current_temporal_ids"}, + current_temporal_ids=inputs["current_temporal_ids"][:1], + ).sample + second_output = model( + **{key: value[1:] for key, value in inputs.items() if key != "current_temporal_ids"}, + current_temporal_ids=inputs["current_temporal_ids"][1:], + ).sample + + torch.testing.assert_close(batch_output[:1], first_output) + torch.testing.assert_close(batch_output[1:], second_output) + + +class TestJoyVideoEditTransformerMemory(JoyVideoEditTransformerTesterConfig, MemoryTesterMixin): + pass + + +class TestJoyVideoEditTransformerTraining(JoyVideoEditTransformerTesterConfig, TrainingTesterMixin): + def test_gradient_checkpointing_is_applied(self): + expected_set = {"JoyVideoEditTransformer3DModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class TestJoyVideoEditTransformerAttention(JoyVideoEditTransformerTesterConfig, AttentionTesterMixin): + pass + + +class TestJoyVideoEditTransformerCompile(JoyVideoEditTransformerTesterConfig, TorchCompileTesterMixin): + pass diff --git a/tests/pipelines/joyvideoedit/__init__.py b/tests/pipelines/joyvideoedit/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/pipelines/joyvideoedit/test_pipeline_joyvideoedit.py b/tests/pipelines/joyvideoedit/test_pipeline_joyvideoedit.py new file mode 100644 index 000000000000..b7faeb3f8d65 --- /dev/null +++ b/tests/pipelines/joyvideoedit/test_pipeline_joyvideoedit.py @@ -0,0 +1,429 @@ +# Copyright 2026 The HuggingFace Team. +# +# 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. + +import weakref + +import PIL.Image +import pytest +import torch +from transformers import Qwen2_5_VLConfig, Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor, Qwen2Tokenizer + +from diffusers import ( + AutoencoderKLJoyVideoEdit, + FlowMatchEulerDiscreteScheduler, + JoyVideoEditPipeline, + JoyVideoEditTransformer3DModel, +) +from diffusers.pipelines.joyvideoedit.pipeline_joyvideoedit import logger as joyvideoedit_logger +from scripts.convert_joyvideoedit_to_diffusers import save_joyvideoedit_pipeline + +from ...testing_utils import ( + CaptureLogger, + enable_full_determinism, + require_accelerate_version_greater, + require_accelerator, + torch_device, +) +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin + + +enable_full_determinism() + + +class JoyVideoEditPipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = JoyVideoEditPipeline + required_input_params_in_call_signature = frozenset(["prompt", "video", "height", "width", "num_inference_steps"]) + batch_input_params = frozenset() + # JoyVideoEdit edits a single source video per prompt, so there is no per-prompt sample-count knob. + optional_input_params = frozenset(["num_inference_steps", "generator", "latents", "output_type", "return_dict"]) + # Per-sample video shape (num_frames, channels, height, width) for the standard dummy inputs (3 input frames at + # 12x12 with a temporal compression ratio of 2 decode to 3 output frames). Spatial dims are a multiple of the + # VAE's spatial compression ratio (3 in this dummy config). + output_shape = (3, 3, 12, 12) + + def get_dummy_components(self, include_mimo_components=False): + torch.manual_seed(0) + transformer = JoyVideoEditTransformer3DModel( + patch_size=[1, 1, 1], + in_channels=2, + out_channels=2, + hidden_size=16, + num_attention_heads=2, + text_dim=16, + num_layers=2, + rope_dim_list=[2, 2, 4], + theta=256, + chunk_size=1, + local_window_size=3, + global_sink_chunk=True, + source_id_rope_dim=4, + source_id_rope_theta=256.0, + ) + + torch.manual_seed(0) + vae = AutoencoderKLJoyVideoEdit( + in_channels=3, + out_channels=3, + patch_size=1, + latent_channels=2, + layers_per_block=1, + block_in_channels=(4, 8), + temporal_downsample=(True, False), + chunk_size=2, + latents_mean=[0.0, 0.0], + latents_std=[1.0, 1.0], + ) + + scheduler = FlowMatchEulerDiscreteScheduler() + + text_encoder = tokenizer = processor = None + if include_mimo_components: + qwen_config = Qwen2_5_VLConfig( + text_config={ + "hidden_size": 16, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "rope_scaling": { + "mrope_section": [1, 1, 2], + "rope_type": "default", + "type": "default", + }, + "rope_theta": 1000000.0, + }, + vision_config={ + "depth": 1, + "hidden_size": 16, + "intermediate_size": 16, + "num_heads": 2, + "out_hidden_size": 16, + }, + hidden_size=16, + vocab_size=152064, + vision_end_token_id=151653, + vision_start_token_id=151652, + vision_token_id=151654, + ) + torch.manual_seed(0) + text_encoder = Qwen2_5_VLForConditionalGeneration(qwen_config) + tokenizer = Qwen2Tokenizer.from_pretrained( + "hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration" + ) + processor = Qwen2_5_VLProcessor.from_pretrained( + "hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration" + ) + + return { + "transformer": transformer, + "vae": vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "processor": processor, + } + + def get_dummy_inputs(self): + # 3 frames satisfy the VAE's `temporal_compression_ratio * n + 1` constraint (ratio 2 here); spatial dims are + # divisible by the spatial compression ratio (3 in this dummy config). + video = [PIL.Image.new("RGB", (12, 12), color=(i * 8, 0, 0)) for i in range(3)] + inputs = { + "video": video, + "prompt": None, + "prompt_embeds": torch.ones(1, 4, 16), + "prompt_embeds_mask": torch.ones(1, 4, dtype=torch.long), + "generator": self.get_generator(0), + "num_inference_steps": 2, + "height": 12, + "width": 12, + "max_sequence_length": 8, + "output_type": "pt", + } + return inputs + + +class TestJoyVideoEditPipeline(JoyVideoEditPipelineTesterConfig, PipelineTesterMixin): + def test_conversion_saves_loadable_pipeline(self, tmp_path): + components = self.get_dummy_components() + save_joyvideoedit_pipeline( + transformer=components["transformer"], + vae=components["vae"], + output_path=tmp_path, + ) + + assert (tmp_path / "model_index.json").is_file() + assert (tmp_path / "scheduler" / "scheduler_config.json").is_file() + + pipeline = self.pipeline_class.from_pretrained(tmp_path) + assert isinstance(pipeline.scheduler, FlowMatchEulerDiscreteScheduler) + + def test_encode_prompt_works_in_isolation(self): + components = self.get_dummy_components(include_mimo_components=True) + components.update(transformer=None, vae=None, scheduler=None) + pipe = self.pipeline_class(**components).to(torch_device) + + prompt_embeds, prompt_embeds_mask = pipe.encode_prompt( + prompt="add a hat", + image=PIL.Image.new("RGB", (12, 12)), + max_sequence_length=8, + ) + + assert prompt_embeds.shape[0] == 1 + assert prompt_embeds.shape[-1] == 16 + assert prompt_embeds_mask is None + + def test_encode_prompt_mask_handling(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + prompt_embeds = torch.ones(1, 4, 16, device=torch_device) + + _, prompt_embeds_mask = pipe.encode_prompt( + prompt=None, + prompt_embeds=prompt_embeds, + prompt_embeds_mask=torch.ones(1, 4, dtype=torch.long, device=torch_device), + ) + assert prompt_embeds_mask is None + + padded_mask = torch.tensor([[1, 1, 0, 0]], dtype=torch.long, device=torch_device) + _, prompt_embeds_mask = pipe.encode_prompt( + prompt=None, + prompt_embeds=prompt_embeds, + prompt_embeds_mask=padded_mask, + ) + torch.testing.assert_close(prompt_embeds_mask, padded_mask) + + def test_inference_batch_consistent(self, batch_size=2): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs() + inputs["prompt_embeds"] = inputs["prompt_embeds"].repeat(batch_size, 1, 1) + inputs["prompt_embeds_mask"] = inputs["prompt_embeds_mask"].repeat(batch_size, 1) + inputs["generator"] = [self.get_generator(i) for i in range(batch_size)] + + output = pipe(**inputs).frames + assert output.shape[0] == batch_size + + def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=2e-3): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs() + output = pipe(**inputs).frames + + inputs["prompt_embeds"] = inputs["prompt_embeds"].repeat(batch_size, 1, 1) + inputs["prompt_embeds_mask"] = inputs["prompt_embeds_mask"].repeat(batch_size, 1) + inputs["generator"] = [self.get_generator(i) for i in range(batch_size)] + output_batch = pipe(**inputs).frames + + torch.testing.assert_close(output_batch[0], output[0], atol=expected_max_diff, rtol=0) + + def test_inference_with_reference_image(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs() + inputs["ref_image"] = PIL.Image.new("RGB", (12, 12), color=(0, 128, 0)) + video = pipe(**inputs).frames + assert video.shape[-3:] == (3, 12, 12) + + def test_kv_cache_cleared_after_call(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe.set_progress_bar_config(disable=None) + + pipe(**self.get_dummy_inputs()) + + state_manager = pipe.transformer._diffusers_hook.get_hook("joyvideoedit_kv_cache").state_manager + assert state_manager._state_cache == {} + + def test_kv_cache_cleared_after_reference_prefill_error(self, monkeypatch): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs() + inputs["ref_image"] = PIL.Image.new("RGB", (12, 12)) + + def fail_prefill(*args, **kwargs): + raise RuntimeError("prefill failure") + + monkeypatch.setattr(pipe.transformer.double_blocks[1], "forward", fail_prefill) + + with pytest.raises(RuntimeError, match="prefill failure"): + pipe(**inputs) + + state_manager = pipe.transformer._diffusers_hook.get_hook("joyvideoedit_kv_cache").state_manager + assert state_manager._current_context is None + assert state_manager._state_cache == {} + + def test_missing_external_mimo_components_raises(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs() + inputs.update(prompt="make the sky orange", prompt_embeds=None, prompt_embeds_mask=None) + + with pytest.raises(ValueError, match="XiaomiMiMo/MiMo-VL-7B-RL-2508"): + pipe(**inputs) + + def test_tokenizer_defaults_to_processor_tokenizer(self): + components = self.get_dummy_components(include_mimo_components=True) + processor = components["processor"] + components["tokenizer"] = None + pipe = self.pipeline_class(**components) + + assert pipe.tokenizer is processor.tokenizer + + def test_empty_video_raises(self): + pipe = self.pipeline_class(**self.get_dummy_components(include_mimo_components=False)) + + with pytest.raises(ValueError, match="non-empty list"): + pipe(video=[], prompt="edit") + + def test_video_frames_are_truncated_to_temporal_grid(self): + pipe = self.pipeline_class(**self.get_dummy_components(include_mimo_components=False)).to(torch_device) + inputs = self.get_dummy_inputs() + inputs["video"].append(PIL.Image.new("RGB", (12, 12), color=(24, 0, 0))) + + with CaptureLogger(joyvideoedit_logger) as cap_logger: + video = pipe(**inputs).frames + + assert video.shape[1] == 3 + assert "Video contains 4 frames" in cap_logger.out + assert "Truncating to 3 frames" in cap_logger.out + + def test_height_and_width_are_adjusted_to_spatial_grid(self): + pipe = self.pipeline_class(**self.get_dummy_components(include_mimo_components=False)).to(torch_device) + inputs = self.get_dummy_inputs() + inputs.update(height=13, width=14) + + with CaptureLogger(joyvideoedit_logger) as cap_logger: + video = pipe(**inputs).frames + + assert video.shape[-2:] == (12, 12) + assert "Adjusting (13, 14) to (12, 12)" in cap_logger.out + + @pytest.mark.parametrize("dimension", ["height", "width"]) + def test_height_and_width_smaller_than_spatial_grid_raise(self, dimension): + pipe = self.pipeline_class(**self.get_dummy_components(include_mimo_components=False)).to(torch_device) + inputs = self.get_dummy_inputs() + inputs[dimension] = getattr(pipe, f"{dimension}_multiple") - 1 + + with pytest.raises(ValueError, match="must be at least"): + pipe(**inputs) + + @pytest.mark.parametrize("num_inference_steps", [0, -1, 1.5]) + def test_invalid_num_inference_steps_raises(self, num_inference_steps): + pipe = self.pipeline_class(**self.get_dummy_components(include_mimo_components=False)) + inputs = self.get_dummy_inputs() + inputs["num_inference_steps"] = num_inference_steps + + with pytest.raises(ValueError, match="positive integer"): + pipe(**inputs) + + def test_callback_can_update_prompt_embeds(self): + pipe = self.pipeline_class(**self.get_dummy_components(include_mimo_components=False)).to(torch_device) + pipe.set_progress_bar_config(disable=None) + + callback_prompt_embeds = [] + + def callback_on_step_end(pipe, step, timestep, callback_kwargs): + callback_prompt_embeds.append(callback_kwargs["prompt_embeds"].clone()) + return {"prompt_embeds": torch.zeros_like(callback_kwargs["prompt_embeds"])} + + inputs = self.get_dummy_inputs() + inputs.update( + prompt=None, + prompt_embeds=torch.ones(1, 4, 16), + prompt_embeds_mask=torch.ones(1, 4, dtype=torch.long), + callback_on_step_end=callback_on_step_end, + callback_on_step_end_tensor_inputs=["prompt_embeds"], + ) + pipe(**inputs) + + assert torch.count_nonzero(callback_prompt_embeds[0]) > 0 + assert torch.count_nonzero(callback_prompt_embeds[1]) == 0 + + def test_decoded_video_is_offloaded_to_cpu(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe.set_progress_bar_config(disable=None) + + video = pipe(**self.get_dummy_inputs()).frames + + assert video.device.type == "cpu" + + def test_input_video_tensor_is_released_before_decode(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + original_preprocess_video = pipe.video_processor.preprocess_video + original_decode = pipe.vae.decode + video_tensor_ref = None + + def preprocess_video(*args, **kwargs): + nonlocal video_tensor_ref + video_tensor = original_preprocess_video(*args, **kwargs) + video_tensor_ref = weakref.ref(video_tensor) + return video_tensor + + def decode(*args, **kwargs): + assert video_tensor_ref() is None + return original_decode(*args, **kwargs) + + pipe.video_processor.preprocess_video = preprocess_video + pipe.vae.decode = decode + + inputs = self.get_dummy_inputs() + inputs["video"] = [torch.rand(3, 3, 12, 12, device=torch_device)] + pipe(**inputs) + + @require_accelerator + @require_accelerate_version_greater("0.17.0") + def test_model_cpu_offload_releases_transformer_before_decode(self): + pipe = self.pipeline_class(**self.get_dummy_components()) + original_postprocess_video = pipe.video_processor.postprocess_video + component_devices = [] + + def postprocess_video(*args, **kwargs): + component_devices.append((pipe.transformer.device.type, pipe.vae.device.type)) + return original_postprocess_video(*args, **kwargs) + + pipe.video_processor.postprocess_video = postprocess_video + pipe.enable_model_cpu_offload(device=torch_device) + pipe(**self.get_dummy_inputs()) + + assert component_devices + assert component_devices[0] == ("cpu", torch_device) + + @pytest.mark.parametrize("offload_mode", ["model", "sequential", "group"]) + @require_accelerator + @require_accelerate_version_greater("0.17.0") + def test_offload_with_text_encoder(self, offload_mode): + pipe = self.pipeline_class(**self.get_dummy_components(include_mimo_components=True)) + if offload_mode == "model": + pipe.enable_model_cpu_offload(device=torch_device) + elif offload_mode == "sequential": + pipe.enable_sequential_cpu_offload(device=torch_device) + else: + pipe.enable_group_offload( + onload_device=torch.device(torch_device), offload_device=torch.device("cpu"), offload_type="leaf_level" + ) + + inputs = self.get_dummy_inputs() + inputs.update( + prompt="add a hat", + prompt_embeds=None, + prompt_embeds_mask=None, + ) + output = pipe(**inputs) + + assert output.frames.device.type == "cpu" + + +class TestJoyVideoEditPipelineMemory(JoyVideoEditPipelineTesterConfig, MemoryTesterMixin): + def test_group_offloading_inference(self): + # The required KV-cache hook is incompatible with this mixin's per-module hook assertion. + pytest.skip("Covered by test_pipeline_level_group_offloading_inference")