diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index c51c6f3fe2c8..824626bfc5aa 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -381,6 +381,8 @@ title: JoyImageEditTransformer3DModel - local: api/models/kandinsky6_transformer title: Kandinsky 6 Transformers + - local: api/models/kandinsky_s2v_transformer + title: KandinskyS2VTransformer3DModel - local: api/models/krea2_transformer2d title: Krea2Transformer2DModel - local: api/models/latte_transformer3d @@ -757,6 +759,8 @@ title: Kandinsky 5.0 Video - local: api/pipelines/kandinsky6 title: Kandinsky 6 + - local: api/pipelines/kandinsky_s2v + title: Kandinsky S2V 1.0 - local: api/pipelines/latte title: Latte - local: api/pipelines/ltx2 diff --git a/docs/source/en/api/models/kandinsky_s2v_transformer.md b/docs/source/en/api/models/kandinsky_s2v_transformer.md new file mode 100644 index 000000000000..3285ebbef36f --- /dev/null +++ b/docs/source/en/api/models/kandinsky_s2v_transformer.md @@ -0,0 +1,22 @@ + + +# KandinskyS2VTransformer3DModel + +The Diffusion Transformer (DiT) of Kandinsky S2V 1.0, conditioned on a reference +frame and a speech track. + +## KandinskyS2VTransformer3DModel + +[[autodoc]] KandinskyS2VTransformer3DModel + - all + - forward diff --git a/docs/source/en/api/pipelines/kandinsky_s2v.md b/docs/source/en/api/pipelines/kandinsky_s2v.md new file mode 100644 index 000000000000..349e619be13e --- /dev/null +++ b/docs/source/en/api/pipelines/kandinsky_s2v.md @@ -0,0 +1,137 @@ + + +# Kandinsky S2V 1.0 + +Kandinsky S2V 1.0 generates an avatar video from a reference frame and a speech +track. The person in the frame is animated to match the speech — the whole figure, +not only the head. The speech is a condition, not an output: the model does not +generate audio, and the soundtrack of the result is the track you pass in. + +Both conditions go through [Qwen2.5-Omni-7B](https://huggingface.co/Qwen/Qwen2.5-Omni-7B): +its hidden states encode the prompt and its audio tower encodes the speech, which +enters the transformer through cross-attention adapters aligned frame by frame, so +each latent frame attends only to the speech it covers. A clip is generated in +chunks of 5 s (125 frames at 25 fps), each carrying FramePack memory of the +previous ones, so longer tracks stay consistent across chunk boundaries. + +Kandinsky S2V 1.0 is integrated as a [modular pipeline](../../modular_diffusers/overview) +only; there is no standard `DiffusionPipeline` class for it. + +## Checkpoints + +| Checkpoint | Transformer | Blocks | Scheduler | Steps | Guidance | +|---|---|---|---|---|---| +| [Lite][lite] | 2B | [`KandinskyS2VAutoBlocks`] | [`FlowMatchEulerDiscreteScheduler`] | 50 | 3.0 | +| [Pro][pro] | 21B | [`KandinskyS2VAutoBlocks`] | [`FlowMatchEulerDiscreteScheduler`] | 50 | 3.0 | +| [Lite distilled][lite-distill] | 2B | [`KandinskyS2VDistilledAutoBlocks`] | [`PiflowScheduler`] | 16 | off | +| [Pro distilled][pro-distill] | 21B | [`KandinskyS2VDistilledAutoBlocks`] | [`PiflowScheduler`] | 16 | off | + +The distilled checkpoints configure `PiflowScheduler` with `shift=10.0` and `n_grid=5`, not its defaults. + +[lite]: https://huggingface.co/kandinskylab/Kandinsky-S2V-1.0-Lite-Diffusers +[pro]: https://huggingface.co/kandinskylab/Kandinsky-S2V-1.0-Pro-Diffusers +[lite-distill]: https://huggingface.co/kandinskylab/Kandinsky-S2V-1.0-Lite-distill-Diffusers +[pro-distill]: https://huggingface.co/kandinskylab/Kandinsky-S2V-1.0-Pro-distill-Diffusers + +`ModularPipeline.from_pretrained` picks the blocks from the checkpoint. The steps +default to the values above. Classifier-free guidance is the `guider` component, +not a call argument: the base checkpoints ship it at 3.0. The distilled ones have +no guider and no `negative_prompt`, because guidance is already part of the student. + +## Usage + +```py +import librosa +import torch +from diffusers import ModularPipeline +from diffusers.utils import encode_video, load_image + +pipe = ModularPipeline.from_pretrained("kandinskylab/Kandinsky-S2V-1.0-Lite-distill-Diffusers") +pipe.load_components(dtype=torch.bfloat16) +pipe.to("cuda") + +image = load_image("reference.png") +audio, _ = librosa.load("speech.wav", sr=16000) # a mono waveform at 16 kHz, for the model +speech, sample_rate = librosa.load("speech.wav", sr=None) # the original track, added to the output video + +video = pipe(image=image, audio=audio, output="videos")[0] +encode_video( + video, + fps=25, + output_path="avatar.mp4", + audio=torch.from_numpy(speech)[None], + audio_sample_rate=sample_rate, +) +``` + +The Pro checkpoints do not fit an 80 GB GPU next to the 7B speech and text encoder +when every component is on the device. Offload them through a `ComponentsManager`: + +```py +import librosa +import torch +from diffusers import ComponentsManager, ModularPipeline +from diffusers.utils import encode_video, load_image + +manager = ComponentsManager() +pipe = ModularPipeline.from_pretrained("kandinskylab/Kandinsky-S2V-1.0-Pro-Diffusers", components_manager=manager) +pipe.load_components(dtype=torch.bfloat16) +manager.enable_auto_cpu_offload(device="cuda", memory_reserve_margin="40GiB") + +image = load_image("reference.png") +audio, _ = librosa.load("speech.wav", sr=16000) +speech, sample_rate = librosa.load("speech.wav", sr=None) + +video = pipe(image=image, audio=audio, output="videos")[0] +encode_video( + video, + fps=25, + output_path="avatar.mp4", + audio=torch.from_numpy(speech)[None], + audio_sample_rate=sample_rate, +) +``` + +The margin is the memory kept free beyond the weights when a model is moved to the GPU. +It has to hold the VAE decode of a chunk, about 33 GiB at 512x768: with the default 3 GB +the transformer stays on the GPU through the decode, which then runs out of memory. With +40 GiB the peak is about 50 GiB on an 80 GB GPU. Enable the offload after `load_components`, +so it applies to every component. + +## Notes + +- `image` is a PIL image and `audio` a mono float waveform in [-1, 1] at 16 kHz, as a tensor + or an array shaped `(samples,)`; the pipeline does not read files. +- The top-level blocks (`text_encoder`, `vae_encoder`, `denoise`, `decode`) run on their own too; + pass the same `num_latent_frames` and `fps` to each when you run them separately. +- The number of chunks follows the length of the speech track unless `num_chunks` + is passed. The video is trimmed to the track either way, so chunks past its end + are computed and dropped, with a warning. +- With guidance, the unconditional pass swaps in the negative prompt *and* silences + the speech, so the guidance is joint over text and speech. + +## KandinskyS2VModularPipeline + +[[autodoc]] KandinskyS2VModularPipeline + +## KandinskyS2VAutoBlocks + +[[autodoc]] KandinskyS2VAutoBlocks + +## KandinskyS2VDistilledModularPipeline + +[[autodoc]] KandinskyS2VDistilledModularPipeline + +## KandinskyS2VDistilledAutoBlocks + +[[autodoc]] KandinskyS2VDistilledAutoBlocks diff --git a/docs/source/en/api/schedulers/piflow.md b/docs/source/en/api/schedulers/piflow.md index 1845c05aa5cd..71fe81602ae9 100644 --- a/docs/source/en/api/schedulers/piflow.md +++ b/docs/source/en/api/schedulers/piflow.md @@ -13,7 +13,8 @@ specific language governing permissions and limitations under the License. # PiflowScheduler `PiflowScheduler` is the few-step scheduler of the distilled Kandinsky 6 checkpoints, both the -text/image-to-video-and-audio model and the video super-resolution model. It implements +text/image-to-video-and-audio model and the video super-resolution model, and of the distilled Kandinsky S2V 1.0 +checkpoints. It implements [π-Flow](https://huggingface.co/papers/2510.14974): the transformer predicts `n_grid` denoised estimates per latent channel at a small number of grid points, and the scheduler integrates a network-free policy between them. diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index f2147d993dea..940138e985ed 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -329,6 +329,7 @@ "Kandinsky6SRTransformer3DModel", "Kandinsky6SRVAE", "Kandinsky6Transformer3DModel", + "KandinskyS2VTransformer3DModel", "Krea2Transformer2DModel", "LatteTransformer3DModel", "LongCatAudioDiTTransformer", @@ -586,6 +587,10 @@ "HunyuanVideo15ModularPipeline", "Ideogram4AutoBlocks", "Ideogram4ModularPipeline", + "KandinskyS2VAutoBlocks", + "KandinskyS2VDistilledAutoBlocks", + "KandinskyS2VDistilledModularPipeline", + "KandinskyS2VModularPipeline", "Krea2AutoBlocks", "Krea2ModularPipeline", "Krea2TurboAutoBlocks", @@ -1256,6 +1261,7 @@ Kandinsky6SRTransformer3DModel, Kandinsky6SRVAE, Kandinsky6Transformer3DModel, + KandinskyS2VTransformer3DModel, Krea2Transformer2DModel, LatteTransformer3DModel, LongCatAudioDiTTransformer, @@ -1491,6 +1497,10 @@ HunyuanVideo15ModularPipeline, Ideogram4AutoBlocks, Ideogram4ModularPipeline, + KandinskyS2VAutoBlocks, + KandinskyS2VDistilledAutoBlocks, + KandinskyS2VDistilledModularPipeline, + KandinskyS2VModularPipeline, Krea2AutoBlocks, Krea2ModularPipeline, Krea2TurboAutoBlocks, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 4fa43aeee413..6e1743e514a8 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -149,6 +149,7 @@ _import_structure["transformers.transformer_kandinsky"] = ["Kandinsky5Transformer3DModel"] _import_structure["transformers.transformer_kandinsky6"] = ["Kandinsky6Transformer3DModel"] _import_structure["transformers.transformer_kandinsky6_sr"] = ["Kandinsky6SRTransformer3DModel"] + _import_structure["transformers.transformer_kandinsky_s2v"] = ["KandinskyS2VTransformer3DModel"] _import_structure["transformers.transformer_krea2"] = ["Krea2Transformer2DModel"] _import_structure["transformers.transformer_longcat_audio_dit"] = ["LongCatAudioDiTTransformer"] _import_structure["transformers.transformer_longcat_image"] = ["LongCatImageTransformer2DModel"] @@ -313,6 +314,7 @@ Kandinsky5Transformer3DModel, Kandinsky6SRTransformer3DModel, Kandinsky6Transformer3DModel, + KandinskyS2VTransformer3DModel, Krea2Transformer2DModel, LatteTransformer3DModel, LongCatAudioDiTTransformer, diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index dd7014096e54..81c620d3cfa3 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -49,6 +49,7 @@ from .transformer_kandinsky import Kandinsky5Transformer3DModel from .transformer_kandinsky6 import Kandinsky6Transformer3DModel from .transformer_kandinsky6_sr import Kandinsky6SRTransformer3DModel + from .transformer_kandinsky_s2v import KandinskyS2VTransformer3DModel from .transformer_krea2 import Krea2Transformer2DModel from .transformer_longcat_audio_dit import LongCatAudioDiTTransformer from .transformer_longcat_image import LongCatImageTransformer2DModel diff --git a/src/diffusers/models/transformers/transformer_kandinsky_s2v.py b/src/diffusers/models/transformers/transformer_kandinsky_s2v.py new file mode 100644 index 000000000000..4a1a539b6b11 --- /dev/null +++ b/src/diffusers/models/transformers/transformer_kandinsky_s2v.py @@ -0,0 +1,791 @@ +# Copyright 2025 The Kandinsky 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. +"""Kandinsky S2V 1.0 Diffusers transformer.""" + +from __future__ import annotations + +import inspect +import math +from typing import Any + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin +from ...utils import logging +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..cache_utils import CacheMixin +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin, get_parameter_dtype + + +logger = logging.get_logger(__name__) + +# The first latent frame takes a single audio token, so the rest need at least one frame. +_MIN_DURATION = 2 + + +# Copied from diffusers.models.transformers.transformer_kandinsky.get_freqs +def get_freqs(dim, max_period=10000.0): + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) + return freqs + + +def apply_rotary(x: Tensor, rope: Tensor) -> Tensor: + """RoPE apply in fp32 (rope tables are fp32), cast back to ``x.dtype``.""" + x_ = x.reshape(*x.shape[:-1], -1, 1, 2).float() + return (rope.float() * x_).sum(dim=-1).reshape(*x.shape).to(dtype=x.dtype) + + +def audio_group_lengths(num_audio_tokens: int, duration: int) -> list[int]: + """First latent frame gets one audio token, the remaining frames split the rest evenly.""" + if duration < _MIN_DURATION: + raise ValueError(f"duration must be >= {_MIN_DURATION} to split audio tokens, got {duration}") + # Fewer tokens than frames leaves empty groups, whose attention rows are fully + # masked; SDPA returns NaN for those and it spreads through every later block. + if num_audio_tokens < duration: + raise ValueError( + f"{num_audio_tokens} audio tokens cannot cover {duration} latent frames: every frame needs at least one" + ) + # Boundaries at `rest * i // (duration - 1)`. + rest = num_audio_tokens - 1 + bounds = [rest * i // (duration - 1) for i in range(duration)] + return [1] + [end - start for start, end in zip(bounds[:-1], bounds[1:], strict=True)] + + +def audio_group_index(num_audio_tokens: int, duration: int, device: torch.device) -> tuple[Tensor, Tensor]: + """Gather index ``(duration, max_len)`` of each latent frame's audio tokens, and a mask of the real ones. + + The lengths follow from the shapes alone, so the padded size is a Python int and the attention keeps static shapes. + """ + lengths = audio_group_lengths(num_audio_tokens, duration) + starts = [sum(lengths[:i]) for i in range(duration)] + positions = torch.arange(max(lengths), device=device) + mask = positions < torch.tensor(lengths, device=device)[:, None] + # The last group is the longest, so the padded index of every other group still stays in range. + index = torch.tensor(starts, device=device)[:, None] + positions + return index, mask + + +class KandinskyS2VTimeEmbeddings(nn.Module): + def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): + super().__init__() + if model_dim % 2 != 0: + raise ValueError(f"model_dim must be even, got {model_dim}") + # A plain attribute rather than a buffer: it must survive a meta-device build. + self.freqs = get_freqs(model_dim // 2, max_period) + + self.in_layer = nn.Linear(model_dim, time_dim, bias=True) + self.activation = nn.SiLU() + self.out_layer = nn.Linear(time_dim, time_dim, bias=True) + + # Not copied from transformer_kandinsky6: that version casts the result to the + # weight dtype, while native S2V keeps the time embedding in fp32 all the way + # into the AdaLN path. The cast would shift rounding in bf16. + def forward(self, time: Tensor) -> Tensor: + args = torch.outer(time.to(torch.float32), self.freqs.to(device=time.device)) + time_embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + return self.out_layer(self.activation(self.in_layer(time_embed))) + + +class KandinskyS2VGeneralEmbeddings(nn.Module): + """Linear projection plus LayerNorm; used for text, pooled text and audio.""" + + def __init__(self, dim: int, model_dim: int): + super().__init__() + self.in_layer = nn.Linear(dim, model_dim, bias=True) + self.norm = nn.LayerNorm(model_dim, elementwise_affine=True) + + def forward(self, embed: Tensor) -> Tensor: + embed = self.in_layer(embed) + # Norm weights are fp32 and activations bf16; native generation reconciles + # that through autocast, which a Diffusers pipeline does not set up. + return self.norm(embed.float()).type_as(embed) + + +class KandinskyS2VVisualEmbeddings(nn.Module): + """Patch projection for ``[T,H,W,C]`` visual tokens.""" + + def __init__(self, visual_dim: int, model_dim: int, patch_size: tuple[int, int, int]): + super().__init__() + self.patch_size = patch_size + self.in_layer = nn.Linear(math.prod(patch_size) * visual_dim, model_dim) + + def forward(self, x: Tensor) -> Tensor: + duration, height, width, channels = x.shape + p_t, p_h, p_w = self.patch_size + x = ( + x.view(duration // p_t, p_t, height // p_h, p_h, width // p_w, p_w, channels) + .permute(0, 2, 4, 1, 3, 5, 6) + .flatten(3, 6) + ) + return self.in_layer(x) + + +# Copied from diffusers.models.transformers.transformer_kandinsky6.Kandinsky6Modulation with Kandinsky6->KandinskyS2V +class KandinskyS2VModulation(nn.Module): + """AdaLN modulation projection.""" + + def __init__(self, time_dim: int, model_dim: int, num_params: int): + super().__init__() + self.activation = nn.SiLU() + self.out_layer = nn.Linear(time_dim, num_params * model_dim) + + def forward(self, x: Tensor) -> Tensor: + return self.out_layer(self.activation(x.to(get_parameter_dtype(self.out_layer)))) + + +# Copied from diffusers.models.transformers.transformer_kandinsky.Kandinsky5FeedForward with Kandinsky5->KandinskyS2V +class KandinskyS2VFeedForward(nn.Module): + def __init__(self, dim, ff_dim): + super().__init__() + self.in_layer = nn.Linear(dim, ff_dim, bias=False) + self.activation = nn.GELU() + self.out_layer = nn.Linear(ff_dim, dim, bias=False) + + def forward(self, x) -> torch.Tensor: + return self.out_layer(self.activation(self.in_layer(x))) + + +class KandinskyS2VAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __init__(self): + if not hasattr(F, "scaled_dot_product_attention"): + raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version.") + + def __call__( + self, + attn: KandinskyS2VAttention, + hidden_states: Tensor, + encoder_hidden_states: Tensor | None = None, + rotary_emb: Tensor | None = None, + kv_index: Tensor | None = None, + kv_mask: Tensor | None = None, + ) -> Tensor: + cond = hidden_states if encoder_hidden_states is None else encoder_hidden_states + shape, cond_shape = hidden_states.shape[:-1], cond.shape[:-1] + + query = attn.to_query(hidden_states).reshape(*shape, attn.num_heads, -1) + key = attn.to_key(cond).reshape(*cond_shape, attn.num_heads, -1) + value = attn.to_value(cond).reshape(*cond_shape, attn.num_heads, -1) + + query = attn.query_norm(query.float()).type_as(query) + key = attn.key_norm(key.float()).type_as(key) + + if rotary_emb is not None: + query = apply_rotary(query, rotary_emb).type_as(query) + key = apply_rotary(key, rotary_emb).type_as(key) + + if kv_index is not None: + # Block-diagonal attention: latent frame i attends exactly to its own audio tokens. The frames are equal + # in size, so the query is reshaped; the audio groups are not, so they are gathered into a padded batch + # axis and masked. + query = query.reshape(kv_index.shape[0], -1, *query.shape[1:]) + key = key[kv_index] + value = value[kv_index] + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=kv_mask[:, None, None, :], + backend=self._attention_backend, + parallel_config=self._parallel_config, + ).flatten(0, 1) + else: + hidden_states = dispatch_attention_fn( + query.unsqueeze(0), + key.unsqueeze(0), + value.unsqueeze(0), + backend=self._attention_backend, + parallel_config=self._parallel_config, + ).squeeze(0) + + return attn.out_layer(hidden_states.flatten(-2, -1)) + + +class KandinskyS2VAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = KandinskyS2VAttnProcessor + _available_processors = [KandinskyS2VAttnProcessor] + + def __init__(self, num_channels: int, head_dim: int, processor: Any | None = None): + super().__init__() + if num_channels % head_dim != 0: + raise ValueError(f"num_channels={num_channels} not divisible by head_dim={head_dim}") + self.num_heads = num_channels // head_dim + + self.to_query = nn.Linear(num_channels, num_channels, bias=True) + self.to_key = nn.Linear(num_channels, num_channels, bias=True) + self.to_value = nn.Linear(num_channels, num_channels, bias=True) + self.query_norm = nn.RMSNorm(head_dim) + self.key_norm = nn.RMSNorm(head_dim) + + self.out_layer = nn.Linear(num_channels, num_channels, bias=True) + self.set_processor(processor if processor is not None else self._default_processor_cls()) + + def forward(self, hidden_states: Tensor, **kwargs: Any) -> Tensor: + accepted = set(inspect.signature(self.processor.__call__).parameters.keys()) + unused = [k for k in kwargs if k not in accepted] + if unused: + logger.warning( + f"attention_processor_kwargs {unused} are not expected by " + f"{self.processor.__class__.__name__} and will be ignored." + ) + kwargs = {k: v for k, v in kwargs.items() if k in accepted} + return self.processor(self, hidden_states, **kwargs) + + +# Not copied from transformer_kandinsky6: that OutLayer takes one time embedding, +# while this one also takes `time_embed_idx` to pick a modulation row per token — +# memory and the reference frame are modulated at t=0, the noisy tail at t. +class KandinskyS2VOutLayer(nn.Module): + def __init__(self, model_dim: int, time_dim: int, visual_dim: int, patch_size: tuple[int, int, int]): + super().__init__() + if patch_size[0] != 1: + # Re-expanding the temporal patch needs cu_seqlens; every shipped S2V + # config uses patch_size[0] == 1. + raise ValueError(f"temporal patch size must be 1, got {patch_size[0]}") + self.patch_size = patch_size + self.modulation = KandinskyS2VModulation(time_dim, model_dim, 2) + self.norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.out_layer = nn.Linear(model_dim, math.prod(patch_size) * visual_dim, bias=True) + + def forward(self, visual_embed: Tensor, time_embed: Tensor, time_embed_idx: Tensor) -> Tensor: + """``(T, H, W, model_dim)`` -> ``(T*p0, H*p1, W*p2, visual_dim)``.""" + shift, scale = ( + p.float()[:, None, None].index_select(0, time_embed_idx) + for p in torch.chunk(self.modulation(time_embed), 2, dim=-1) + ) + visual_embed = (self.norm(visual_embed.float()) * (scale + 1.0) + shift).type_as(visual_embed) + x = self.out_layer(visual_embed) + + duration, height, width, _ = x.shape + return ( + x.view(duration, height, width, -1, *self.patch_size) + .permute(0, 4, 1, 5, 2, 6, 3) + .flatten(0, 1) + .flatten(1, 2) + .flatten(2, 3) + ) + + +# Not copied from transformer_kandinsky6: its tables are buffers, which +# `.to(torch.bfloat16)` casts down. +class KandinskyS2VRoPE1D(nn.Module): + def __init__(self, dim: int, max_pos: int = 1024, max_period: float = 10000.0): + super().__init__() + freq = get_freqs(dim // 2, max_period) + pos = torch.arange(max_pos, dtype=freq.dtype) + # A plain attribute rather than a buffer, so `.to(torch.bfloat16)` leaves the table in fp32. + self.args = torch.outer(pos, freq) + + def forward(self, pos: Tensor) -> Tensor: + args = self.args.to(device=pos.device)[pos] + rope = torch.stack([torch.cos(args), -torch.sin(args), torch.sin(args), torch.cos(args)], dim=-1) + return rope.view(*rope.shape[:-1], 2, 2).unsqueeze(-4) + + +# Not copied from transformer_kandinsky6: its tables are buffers, and its temporal +# table is built over arange(max_pos), which would silently wrap a negative index +# around the end of the tensor. FramePack memory sits at negative temporal positions. +class KandinskyS2VRoPE3D(nn.Module): + """3-D RoPE; the temporal axis spans negative positions for memory frames.""" + + def __init__( + self, + axes_dims: tuple[int, int, int], + max_pos: tuple[int, int, int] = (128, 128, 128), + max_period: float = 10000.0, + ): + super().__init__() + self.offsets_t = max_pos[0] + + # Plain attributes rather than buffers, so `.to(torch.bfloat16)` leaves the tables in fp32. + self.args = [] + for i, (axes_dim, ax_max_pos) in enumerate(zip(axes_dims, max_pos, strict=True)): + freq = get_freqs(axes_dim // 2, max_period) + pos = ( + torch.arange(-ax_max_pos, ax_max_pos, dtype=freq.dtype) + if i == 0 + else torch.arange(ax_max_pos, dtype=freq.dtype) + ) + self.args.append(torch.outer(pos, freq)) + + def forward( + self, + shape: tuple[int, int, int], + pos: list[Tensor], + scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), + ) -> Tensor: + duration, height, width = shape + args_0, args_1, args_2 = ( + args.to(device=axis_pos.device) for args, axis_pos in zip(self.args, pos, strict=True) + ) + args_t = args_0[pos[0] + self.offsets_t] / scale_factor[0] + args_h = args_1[pos[1]] / scale_factor[1] + args_w = args_2[pos[2]] / scale_factor[2] + + args = torch.cat( + [ + args_t.view(duration, 1, 1, -1).repeat(1, height, width, 1), + args_h.view(1, height, 1, -1).repeat(duration, 1, width, 1), + args_w.view(1, 1, width, -1).repeat(duration, height, 1, 1), + ], + dim=-1, + ) + rope = torch.stack([torch.cos(args), -torch.sin(args), torch.sin(args), torch.cos(args)], dim=-1) + return rope.view(*rope.shape[:-1], 2, 2).unsqueeze(-4) + + +class KandinskyS2VTransformerTextEncoderBlock(nn.Module): + def __init__(self, model_dim: int, time_dim: int, ff_dim: int, head_dim: int): + super().__init__() + self.text_modulation = KandinskyS2VModulation(time_dim, model_dim, 6) + + self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.self_attention = KandinskyS2VAttention(model_dim, head_dim) + + self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.feed_forward = KandinskyS2VFeedForward(model_dim, ff_dim) + + def forward(self, x: Tensor, time_embed: Tensor, rope: Tensor, time_embed_idx: Tensor) -> Tensor: + self_attn_params, ff_params = torch.chunk(self.text_modulation(time_embed), 2, dim=-1) + + shift, scale, gate = torch.chunk(self_attn_params, 3, dim=-1) + out = ( + self.self_attention_norm(x.float()) * (scale.float().index_select(0, time_embed_idx) + 1.0) + + shift.float().index_select(0, time_embed_idx) + ).type_as(x) + out = self.self_attention(out, rotary_emb=rope) + x = (x.float() + gate.float().index_select(0, time_embed_idx) * out.float()).type_as(x) + + shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) + out = ( + self.feed_forward_norm(x.float()) * (scale.float().index_select(0, time_embed_idx) + 1.0) + + shift.float().index_select(0, time_embed_idx) + ).type_as(x) + out = self.feed_forward(out) + return (x.float() + gate.float().index_select(0, time_embed_idx) * out.float()).type_as(x) + + +class KandinskyS2VTransformerAudioEncoderBlock(nn.Module): + """Plain pre-norm transformer block — the audio tower is not time-modulated.""" + + def __init__(self, model_dim: int, ff_dim: int, head_dim: int): + super().__init__() + self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=True) + self.self_attention = KandinskyS2VAttention(model_dim, head_dim) + + self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=True) + self.feed_forward = KandinskyS2VFeedForward(model_dim, ff_dim) + + def forward(self, x: Tensor, rope: Tensor) -> Tensor: + x = x + self.self_attention(self.self_attention_norm(x.float()).type_as(x), rotary_emb=rope) + return x + self.feed_forward(self.feed_forward_norm(x.float()).type_as(x)) + + +class KandinskyS2VTransformerDecoderBlock(nn.Module): + """Visual block; ``enable_audio`` adds the frame-aligned speech adapter.""" + + def __init__( + self, + model_dim: int, + time_dim: int, + ff_dim: int, + head_dim: int, + enable_audio: bool = False, + ): + super().__init__() + self.visual_modulation = KandinskyS2VModulation(time_dim, model_dim, 9) + + self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.self_attention = KandinskyS2VAttention(model_dim, head_dim) + + self.cross_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.cross_attention = KandinskyS2VAttention(model_dim, head_dim) + + self.enable_audio = enable_audio + if enable_audio: + self.cross_attention_norm_v2a = nn.LayerNorm(model_dim, elementwise_affine=True) + self.cross_attention_v2a = KandinskyS2VAttention(model_dim, head_dim) + + self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.feed_forward = KandinskyS2VFeedForward(model_dim, ff_dim) + + def forward( + self, + visual_embed: Tensor, + text_embed: Tensor, + audio_embed: Tensor, + time_embed: Tensor, + rope: Tensor, + time_embed_idx: Tensor, + num_visual_tokens: int, + audio_index: Tensor, + audio_mask: Tensor, + ) -> Tensor: + self_attn_params, cross_attn_params, ff_params = torch.chunk(self.visual_modulation(time_embed), 3, dim=-1) + + shift, scale, gate = torch.chunk(self_attn_params, 3, dim=-1) + out = ( + self.self_attention_norm(visual_embed.float()) * (scale.float().index_select(0, time_embed_idx) + 1.0) + + shift.float().index_select(0, time_embed_idx) + ).type_as(visual_embed) + out = self.self_attention(out, rotary_emb=rope) + visual_embed = (visual_embed.float() + gate.float().index_select(0, time_embed_idx) * out.float()).type_as( + visual_embed + ) + + shift, scale, gate = torch.chunk(cross_attn_params, 3, dim=-1) + out = ( + self.cross_attention_norm(visual_embed.float()) * (scale.float().index_select(0, time_embed_idx) + 1.0) + + shift.float().index_select(0, time_embed_idx) + ).type_as(visual_embed) + out = self.cross_attention(out, encoder_hidden_states=text_embed) + visual_embed = (visual_embed.float() + gate.float().index_select(0, time_embed_idx) * out.float()).type_as( + visual_embed + ) + + if self.enable_audio: + # Speech conditions only the noisy tail, frame by frame. The residual is + # ungated, as in native generation with audio_adaln=False. + head, noisy = visual_embed[:-num_visual_tokens], visual_embed[-num_visual_tokens:] + out = self.cross_attention_v2a( + self.cross_attention_norm_v2a(noisy.float()).type_as(noisy), + encoder_hidden_states=audio_embed, + kv_index=audio_index, + kv_mask=audio_mask, + ) + visual_embed = torch.cat([head, noisy + out], dim=0) + + shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) + out = ( + self.feed_forward_norm(visual_embed.float()) * (scale.float().index_select(0, time_embed_idx) + 1.0) + + shift.float().index_select(0, time_embed_idx) + ).type_as(visual_embed) + out = self.feed_forward(out) + return (visual_embed.float() + gate.float().index_select(0, time_embed_idx) * out.float()).type_as( + visual_embed + ) + + +class KandinskyS2VTransformer3DModel( + ModelMixin, + ConfigMixin, + PeftAdapterMixin, + CacheMixin, + AttentionMixin, +): + """Speech-to-video diffusion transformer. + + Parameter names match the training checkpoints, so a converted `.safetensors` loads with no key remapping. + Inference is batch size 1: the sequence is packed and the speech cross-attention is frame-aligned, so `forward` + takes a leading batch axis of 1 and squeezes it. + + The visual stream is one packed sequence ``[memory | ref | visual]``: + + * ``memory`` — FramePack context from the previous chunk at 4x / 2x / 1x temporal downsampling, placed at negative + temporal RoPE positions; + * ``ref`` — the reference frame, pinned to temporal RoPE index ``index_rope_ref_image``; + * ``visual`` — the noisy latent being denoised. + + ``memory`` and ``ref`` are clean, so they are modulated with ``t = 0`` while ``visual`` uses the real timestep; the + modulation holds both rows and an index selects per token. + """ + + _repeated_blocks = [ + "KandinskyS2VTransformerTextEncoderBlock", + "KandinskyS2VTransformerAudioEncoderBlock", + "KandinskyS2VTransformerDecoderBlock", + ] + _no_split_modules = _repeated_blocks + # The checkpoint keeps the time embedding and every normalisation in fp32, and + # AdaLN modulation runs in fp32; this stops `torch_dtype=` from flattening that. + # Matched both as substrings and as exact path components, hence the compound + # names beside the bare ones. + _keep_in_fp32_modules = [ + "time_embeddings", + "modulation", + "text_modulation", + "visual_modulation", + "norm", + "self_attention_norm", + "cross_attention_norm", + "cross_attention_norm_v2a", + "feed_forward_norm", + "query_norm", + "key_norm", + ] + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_visual_dim: int = 16, + out_visual_dim: int = 16, + in_text_dim: int = 3584, + in_text_dim2: int = 768, + in_audio_dim: int = 3584, + time_dim: int = 512, + patch_size: tuple[int, int, int] = (1, 2, 2), + frame_buckets: tuple[int, int, int] = (2, 4, 24), + enable_framepack: bool = True, + index_rope_ref_image: int = 40, + patch_proj1x_coef: tuple[int, int, int] = (1, 1, 1), + patch_proj2x_coef: tuple[int, int, int] = (2, 2, 2), + patch_proj4x_coef: tuple[int, int, int] = (4, 4, 4), + model_dim: int = 1792, + ff_dim: int = 7168, + num_text_blocks: int = 2, + num_audio_blocks: int = 2, + num_visual_blocks: int = 32, + axes_dims: tuple[int, int, int] = (16, 24, 24), + audio_inject_layers: tuple[int, ...] = (), + text_rope_max_pos: int = 4096, + ): + super().__init__() + head_dim = sum(axes_dims) + self.model_dim = model_dim + self.enable_framepack = enable_framepack + + self.time_embeddings = KandinskyS2VTimeEmbeddings(model_dim, time_dim) + self.text_embeddings = KandinskyS2VGeneralEmbeddings(in_text_dim, model_dim) + self.pooled_text_embeddings = KandinskyS2VGeneralEmbeddings(in_text_dim2, time_dim) + self.audio_embeddings = KandinskyS2VGeneralEmbeddings(in_audio_dim, model_dim) + + self.visual_embeddings = KandinskyS2VVisualEmbeddings(in_visual_dim, model_dim, patch_size) + self.visual_embeddings_proj1x = KandinskyS2VVisualEmbeddings( + in_visual_dim, + model_dim, + ( + patch_size[0] * patch_proj1x_coef[0], + patch_size[1] * patch_proj1x_coef[1], + patch_size[2] * patch_proj1x_coef[2], + ), + ) + self.visual_embeddings_proj2x = KandinskyS2VVisualEmbeddings( + in_visual_dim, + model_dim, + ( + patch_size[0] * patch_proj2x_coef[0], + patch_size[1] * patch_proj2x_coef[1], + patch_size[2] * patch_proj2x_coef[2], + ), + ) + self.visual_embeddings_proj4x = KandinskyS2VVisualEmbeddings( + in_visual_dim, + model_dim, + ( + patch_size[0] * patch_proj4x_coef[0], + patch_size[1] * patch_proj4x_coef[1], + patch_size[2] * patch_proj4x_coef[2], + ), + ) + # 0 = visual, 1 = reference frame, 2 = memory + self.mask_embeddings = nn.Embedding(3, model_dim) + + self.text_rope_embeddings = KandinskyS2VRoPE1D(head_dim, max_pos=text_rope_max_pos) + self.text_transformer_blocks = nn.ModuleList( + [ + KandinskyS2VTransformerTextEncoderBlock(model_dim, time_dim, ff_dim, head_dim) + for _ in range(num_text_blocks) + ] + ) + + self.audio_rope_embeddings = KandinskyS2VRoPE1D(head_dim) + self.audio_transformer_blocks = nn.ModuleList( + [KandinskyS2VTransformerAudioEncoderBlock(model_dim, ff_dim, head_dim) for _ in range(num_audio_blocks)] + ) + + self.visual_rope_embeddings = KandinskyS2VRoPE3D(axes_dims) + self.visual_transformer_blocks = nn.ModuleList( + [ + KandinskyS2VTransformerDecoderBlock( + model_dim, + time_dim, + ff_dim, + head_dim, + enable_audio=(i in audio_inject_layers), + ) + for i in range(num_visual_blocks) + ] + ) + + self.out_layer = KandinskyS2VOutLayer(model_dim, time_dim, out_visual_dim, patch_size) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: Tensor, + encoder_hidden_states: Tensor, + audio_hidden_states: Tensor, + reference_hidden_states: Tensor, + memory_hidden_states: tuple[Tensor, Tensor, Tensor], + timestep: Tensor, + pooled_projections: Tensor, + visual_rope_pos: list[Tensor], + reference_rope_pos: list[Tensor], + memory_rope_pos: tuple[list[Tensor], list[Tensor], list[Tensor]], + text_rope_pos: Tensor, + audio_rope_pos: Tensor, + scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), + return_dict: bool = True, + ) -> Transformer2DModelOutput | tuple[Tensor]: + """Denoise one chunk. + + Args: + hidden_states: noisy latent, ``(1, T, H, W, C)``. + encoder_hidden_states: text embeddings, ``(1, S_text, in_text_dim)``. + audio_hidden_states: speech features, ``(1, S_audio, in_audio_dim)``. + reference_hidden_states: reference frame latent, ``(1, 1, H, W, C)``. + memory_hidden_states: FramePack memory at 4x / 2x / 1x, each ``(1, T_i, H_i, W_i, C)``. + timestep: current timestep, ``(1,)``. + pooled_projections: pooled CLIP text embedding, ``(1, in_text_dim2)``. + visual_rope_pos: per-axis RoPE positions of the noisy latent. + reference_rope_pos: per-axis RoPE positions of the reference frame. + memory_rope_pos: per-axis RoPE positions of each FramePack memory scale. + text_rope_pos: 1-D RoPE positions of the text tokens. + audio_rope_pos: latent frame each speech token belongs to. + scale_factor: RoPE scale per axis, set by the resolution bucket. + return_dict: return a ``Transformer2DModelOutput`` instead of a tuple. + + Batch size must be 1: the sequence is packed and the speech cross-attention is frame-aligned. + + Returns: + [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple`: the predicted velocity, shaped like + ``hidden_states``. + """ + if hidden_states.shape[0] != 1: + raise ValueError(f"batch size must be 1, got {hidden_states.shape[0]}") + + x = hidden_states.squeeze(0) + ref = reference_hidden_states.squeeze(0) + memory = tuple(m.squeeze(0) for m in memory_hidden_states) + text_embed = encoder_hidden_states.squeeze(0) + audio_embed = audio_hidden_states.squeeze(0) + + device = x.device + duration = x.shape[0] + + text_embed = self.text_embeddings(text_embed) + audio_embed = self.audio_embeddings(audio_embed) + + time_embed = self.time_embeddings(timestep) + time_embed_clean = self.time_embeddings(torch.zeros_like(timestep)) + pooled = self.pooled_text_embeddings(pooled_projections) + time_embed = time_embed + pooled + time_embed_clean = time_embed_clean + pooled + # row 0 -> memory / ref (t = 0), row 1 -> visual (t) + time_embed_full = torch.stack((time_embed_clean, time_embed), dim=1).flatten(0, 1) + + text_rope = self.text_rope_embeddings(text_rope_pos) + text_time_embed_idx = torch.zeros(text_embed.shape[0], dtype=torch.int32, device=device) + for block in self.text_transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + text_embed = self._gradient_checkpointing_func( + block, text_embed, time_embed, text_rope, text_time_embed_idx + ) + else: + text_embed = block(text_embed, time_embed, text_rope, text_time_embed_idx) + + audio_rope = self.audio_rope_embeddings(audio_rope_pos) + for block in self.audio_transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + audio_embed = self._gradient_checkpointing_func(block, audio_embed, audio_rope) + else: + audio_embed = block(audio_embed, audio_rope) + + visual_embed = self.visual_embeddings(x) + visual_shape = visual_embed.shape[:-1] + visual_rope = self.visual_rope_embeddings(visual_shape, visual_rope_pos, scale_factor).flatten(0, 2) + visual_embed = visual_embed.flatten(0, 2) + + ref_embed = self.visual_embeddings(ref) + ref_rope = self.visual_rope_embeddings(ref_embed.shape[:-1], reference_rope_pos, scale_factor).flatten(0, 2) + ref_embed = ref_embed.flatten(0, 2) + + # The FramePack memory at its three temporal scales. + memory_4x, memory_2x, memory_1x = memory + rope_pos_4x, rope_pos_2x, rope_pos_1x = memory_rope_pos + memory_parts = [(self.visual_embeddings_proj1x(memory_1x), rope_pos_1x)] + if self.enable_framepack: + memory_parts = [ + (self.visual_embeddings_proj4x(memory_4x), rope_pos_4x), + (self.visual_embeddings_proj2x(memory_2x), rope_pos_2x), + *memory_parts, + ] + memory_embed = torch.cat([embed.flatten(0, 2) for embed, _ in memory_parts], dim=0) + memory_rope = torch.cat( + [ + self.visual_rope_embeddings(embed.shape[:-1], rope_pos, scale_factor).flatten(0, 2) + for embed, rope_pos in memory_parts + ], + dim=0, + ) + + visual_embed = visual_embed + self.mask_embeddings(torch.tensor(0, device=device)) + ref_embed = ref_embed + self.mask_embeddings(torch.tensor(1, device=device)) + memory_embed = memory_embed + self.mask_embeddings(torch.tensor(2, device=device)) + + num_visual_tokens = visual_embed.shape[0] + num_clean_tokens = memory_embed.shape[0] + ref_embed.shape[0] + visual_embed_full = torch.cat([memory_embed, ref_embed, visual_embed], dim=0) + visual_rope_full = torch.cat([memory_rope, ref_rope, visual_rope], dim=0) + time_embed_idx_full = torch.cat( + [ + torch.zeros(num_clean_tokens, dtype=torch.int32, device=device), + torch.ones(num_visual_tokens, dtype=torch.int32, device=device), + ] + ) + + audio_index, audio_mask = audio_group_index(audio_embed.shape[0], duration, device) + + for block in self.visual_transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + visual_embed_full = self._gradient_checkpointing_func( + block, + visual_embed_full, + text_embed, + audio_embed, + time_embed_full, + visual_rope_full, + time_embed_idx_full, + num_visual_tokens, + audio_index, + audio_mask, + ) + else: + visual_embed_full = block( + visual_embed_full, + text_embed, + audio_embed, + time_embed_full, + visual_rope_full, + time_embed_idx_full, + num_visual_tokens, + audio_index, + audio_mask, + ) + + visual_embed = visual_embed_full[-num_visual_tokens:].reshape(*visual_shape, self.model_dim) + visual_time_embed_idx = torch.zeros(visual_shape[0], dtype=torch.int32, device=device) + out = self.out_layer(visual_embed, time_embed, visual_time_embed_idx).unsqueeze(0) + + if not return_dict: + return (out,) + return Transformer2DModelOutput(sample=out) diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index fa1a5f0dad7d..7beb47d47008 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -91,6 +91,12 @@ "Ideogram4AutoBlocks", "Ideogram4ModularPipeline", ] + _import_structure["kandinsky_s2v"] = [ + "KandinskyS2VAutoBlocks", + "KandinskyS2VDistilledAutoBlocks", + "KandinskyS2VDistilledModularPipeline", + "KandinskyS2VModularPipeline", + ] _import_structure["krea2"] = [ "Krea2AutoBlocks", "Krea2ModularPipeline", @@ -211,6 +217,12 @@ Ideogram4AutoBlocks, Ideogram4ModularPipeline, ) + from .kandinsky_s2v import ( + KandinskyS2VAutoBlocks, + KandinskyS2VDistilledAutoBlocks, + KandinskyS2VDistilledModularPipeline, + KandinskyS2VModularPipeline, + ) from .krea2 import ( Krea2AutoBlocks, Krea2ModularPipeline, diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/__init__.py b/src/diffusers/modular_pipelines/kandinsky_s2v/__init__.py new file mode 100644 index 000000000000..a74da5e01c18 --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/__init__.py @@ -0,0 +1,49 @@ +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["modular_blocks_kandinsky_s2v"] = ["KandinskyS2VAutoBlocks"] + _import_structure["modular_blocks_kandinsky_s2v_distilled"] = ["KandinskyS2VDistilledAutoBlocks"] + _import_structure["modular_pipeline"] = ["KandinskyS2VDistilledModularPipeline", "KandinskyS2VModularPipeline"] + +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 * # noqa F403 + else: + from .modular_blocks_kandinsky_s2v import KandinskyS2VAutoBlocks + from .modular_blocks_kandinsky_s2v_distilled import KandinskyS2VDistilledAutoBlocks + from .modular_pipeline import KandinskyS2VDistilledModularPipeline, KandinskyS2VModularPipeline +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/modular_pipelines/kandinsky_s2v/before_denoise.py b/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py new file mode 100644 index 000000000000..7c988704eee4 --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py @@ -0,0 +1,266 @@ +# Copyright 2026 The Kandinsky 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 torch + +from ...models.transformers.transformer_kandinsky_s2v import KandinskyS2VTransformer3DModel +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import KandinskyS2VModularPipeline, align, split_audio_embeds + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def memory_coefs(config) -> tuple[tuple[int, int, int], ...]: + return (config.patch_proj1x_coef, config.patch_proj2x_coef, config.patch_proj4x_coef) + + +def init_memory(config, height: int, width: int, dtype: torch.dtype, device: torch.device) -> list[torch.Tensor]: + """Zero FramePack memory buffers for the first chunk, in ``[1x, 2x, 4x]`` order.""" + patch_size = config.patch_size + buffers = [] + for bucket, coef in zip(config.frame_buckets, memory_coefs(config), strict=True): + h_mult, w_mult = patch_size[1] * coef[1], patch_size[2] * coef[2] + buffers.append( + torch.zeros( + bucket, align(height, h_mult), align(width, w_mult), config.in_visual_dim, dtype=dtype, device=device + ) + ) + return buffers + + +def roll_memory(config, latents: torch.Tensor) -> list[torch.Tensor]: + """The tail of this chunk becomes the next chunk's memory, in ``[1x, 2x, 4x]`` order.""" + patch_size = config.patch_size + frame_buckets = config.frame_buckets + + needed = sum(frame_buckets) + if latents.shape[0] < needed: + raise ValueError( + f"latents have {latents.shape[0]} frames but the FramePack memory window " + f"sum(frame_buckets)={needed} requires at least that many" + ) + + offsets = [0] + for bucket in frame_buckets: + offsets.append(offsets[-1] + bucket) + + height, width = latents.shape[1], latents.shape[2] + buffers = [] + for i, coef in enumerate(memory_coefs(config)): + window = latents[latents.shape[0] - offsets[i + 1] : latents.shape[0] - offsets[i]] + h_mult, w_mult = patch_size[1] * coef[1], patch_size[2] * coef[2] + buffers.append(window[:, : align(height, h_mult), : align(width, w_mult), :].clone()) + return buffers + + +def memory_rope_pos(config, memory: list[torch.Tensor], device: torch.device) -> list[list[torch.Tensor]]: + """RoPE positions for ``[1x, 2x, 4x]`` memory, which sits at negative temporal indices.""" + patch_size = config.patch_size + + bounds = [0] + for bucket in config.frame_buckets: + bounds.append(bounds[-1] + bucket) + + out = [] + for i, (buffer, coef) in enumerate(zip(memory, memory_coefs(config), strict=True)): + height = buffer.shape[1] // patch_size[1] + width = buffer.shape[2] // patch_size[2] + out.append( + [ + torch.arange(-bounds[i + 1], -bounds[i], coef[0], device=device), + torch.arange(0, height, coef[1], device=device), + torch.arange(0, width, coef[2], device=device), + ] + ) + return out + + +class KandinskyS2VPrepareStep(ModularPipelineBlocks): + model_name = "kandinsky-s2v" + + @property + def description(self) -> str: + return ( + "Prepare step that crops the reference latent to whole patches, casts the conditions to the " + "transformer's dtype, splits the speech features evenly over the chunks, builds the RoPE positions and " + "the empty FramePack memory the first chunk starts from." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", KandinskyS2VTransformer3DModel)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "reference_latents", + required=True, + type_hint=torch.Tensor, + description="Scaled VAE latent of the reference frame, from the VAE encoder step.", + ), + InputParam( + "audio_embeds", + required=True, + type_hint=torch.Tensor, + description="Audio-tower features of the speech track, from the text encoder step.", + ), + InputParam("num_chunks", required=True, type_hint=int, description="Number of chunks to generate."), + InputParam( + "num_latent_frames", + default=32, + type_hint=int, + description="Latent frames per chunk; 32 gives 125 frames (5 s at 25 fps).", + ), + InputParam.template("prompt_embeds", required=True), + InputParam( + "pooled_prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="CLIP pooled prompt embedding.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "reference_hidden_states", type_hint=torch.Tensor, description="Reference latent, cropped and cast." + ), + OutputParam("latent_height", type_hint=int, description="Height of the latents."), + OutputParam("latent_width", type_hint=int, description="Width of the latents."), + OutputParam("audio_chunks", type_hint=list, description="Speech features of each chunk."), + OutputParam.template("prompt_embeds"), + OutputParam("pooled_prompt_embeds", type_hint=torch.Tensor, description="CLIP pooled prompt embedding."), + OutputParam("text_rope_pos", type_hint=torch.Tensor, description="RoPE positions of the prompt tokens."), + OutputParam("visual_rope_pos", type_hint=list, description="Per-axis RoPE positions of a chunk."), + OutputParam( + "reference_rope_pos", type_hint=list, description="Per-axis RoPE positions of the reference frame." + ), + OutputParam("memory", type_hint=list, description="FramePack memory in [1x, 2x, 4x] order."), + ] + + @staticmethod + def check_inputs(num_latent_frames: int, frame_buckets: list[int]): + if num_latent_frames < sum(frame_buckets): + raise ValueError( + f"`num_latent_frames` must be >= {sum(frame_buckets)}, the FramePack memory window " + f"sum(frame_buckets), got {num_latent_frames}" + ) + + @torch.no_grad() + def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + config = components.transformer.config + self.check_inputs(block_state.num_latent_frames, config.frame_buckets) + device = components._execution_device + dtype = components.transformer.dtype + patch_size = config.patch_size + + reference = block_state.reference_latents + reference = reference[ + :, : align(reference.shape[1], patch_size[1]), : align(reference.shape[2], patch_size[2]) + ] + block_state.latent_height, block_state.latent_width = reference.shape[1], reference.shape[2] + block_state.reference_hidden_states = reference.to(dtype) + + block_state.prompt_embeds = block_state.prompt_embeds.to(dtype) + block_state.pooled_prompt_embeds = block_state.pooled_prompt_embeds.to(dtype) + block_state.text_rope_pos = torch.arange(block_state.prompt_embeds.shape[0], device=device) + + block_state.audio_chunks = split_audio_embeds(block_state.audio_embeds.to(dtype), block_state.num_chunks) + + height, width = block_state.latent_height, block_state.latent_width + block_state.visual_rope_pos = [ + torch.arange(block_state.num_latent_frames, device=device), + torch.arange(height // patch_size[1], device=device), + torch.arange(width // patch_size[2], device=device), + ] + block_state.reference_rope_pos = [ + torch.tensor([config.index_rope_ref_image], device=device), + torch.arange(height // patch_size[1], device=device), + torch.arange(width // patch_size[2], device=device), + ] + # Latents and memory stay fp32; only the forward runs in the transformer's dtype. + block_state.memory = init_memory(config, height, width, torch.float32, device) + + self.set_block_state(state, block_state) + return components, state + + +class KandinskyS2VPrepareNegativeStep(ModularPipelineBlocks): + model_name = "kandinsky-s2v" + + @property + def description(self) -> str: + return ( + "Prepare step for the unconditional pass of the guidance: casts the negative prompt embeddings to the " + "transformer's dtype and builds their RoPE positions. Without a negative prompt it passes nothing on." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", KandinskyS2VTransformer3DModel)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("negative_prompt_embeds"), + InputParam( + "negative_pooled_prompt_embeds", + type_hint=torch.Tensor, + description="CLIP pooled negative prompt embedding.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("negative_prompt_embeds"), + OutputParam( + "negative_pooled_prompt_embeds", + type_hint=torch.Tensor, + description="CLIP pooled negative prompt embedding.", + ), + OutputParam( + "negative_text_rope_pos", + type_hint=torch.Tensor, + description="RoPE positions of the negative prompt tokens.", + ), + ] + + @staticmethod + def check_inputs(negative_prompt_embeds: torch.Tensor | None, negative_pooled_prompt_embeds: torch.Tensor | None): + if (negative_prompt_embeds is None) != (negative_pooled_prompt_embeds is None): + raise ValueError("`negative_prompt_embeds` and `negative_pooled_prompt_embeds` must be passed together.") + + @torch.no_grad() + def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state.negative_prompt_embeds, block_state.negative_pooled_prompt_embeds) + dtype = components.transformer.dtype + + block_state.negative_text_rope_pos = None + if block_state.negative_prompt_embeds is not None: + block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.to(dtype) + block_state.negative_pooled_prompt_embeds = block_state.negative_pooled_prompt_embeds.to(dtype) + block_state.negative_text_rope_pos = torch.arange( + block_state.negative_prompt_embeds.shape[0], device=components._execution_device + ) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/decoders.py b/src/diffusers/modular_pipelines/kandinsky_s2v/decoders.py new file mode 100644 index 000000000000..dfeee4db3151 --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/decoders.py @@ -0,0 +1,95 @@ +# Copyright 2026 The Kandinsky 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 torch + +from ...configuration_utils import FrozenDict +from ...models import AutoencoderKLHunyuanVideo +from ...utils import logging +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import KandinskyS2VModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def decode_latents(vae: AutoencoderKLHunyuanVideo, latents: torch.Tensor) -> torch.Tensor: + """``(T, H, W, C)`` latent -> ``(1, 3, T*4-3, H*8, W*8)`` in ``[-1, 1]``. + + The decode runs under bf16 autocast rather than in the VAE's own fp16, as the native pipeline does: decoding the + same latent the two ways differs by 4.4e-3 relative. + """ + images = latents.reshape(1, *latents.shape) + images = (images / vae.config.scaling_factor).permute(0, 4, 1, 2, 3) + device_type = latents.device.type + with torch.autocast(device_type, dtype=torch.bfloat16, enabled=device_type == "cuda"): + decoded = vae.decode(images.to(vae.dtype)).sample + # Autocast is a choice about the precision of the decode, not about the dtype handed back; the widening is exact. + return decoded.float() + + +class KandinskyS2VDecodeStep(ModularPipelineBlocks): + model_name = "kandinsky-s2v" + + @property + def description(self) -> str: + return ( + "Decode step that decodes each chunk's latents, joins the chunks and trims the video back to the length " + "of the speech track." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLHunyuanVideo), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 8}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "latent_chunks", + required=True, + type_hint=list, + description="Denoised latents of each chunk, each shaped " + "(num_latent_frames, latent_height, latent_width, channels).", + ), + InputParam("track_seconds", required=True, type_hint=float), + InputParam("fps", default=25, type_hint=int), + InputParam.template("output_type"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("videos")] + + @torch.no_grad() + def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + decoded = [decode_latents(components.vae, latents).cpu() for latents in block_state.latent_chunks] + frames = torch.cat(decoded, dim=2) # (1, C, T, H, W) + frames = frames[:, :, : min(frames.shape[2], int(block_state.track_seconds * block_state.fps))] + block_state.videos = components.video_processor.postprocess_video(frames, output_type=block_state.output_type) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py b/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py new file mode 100644 index 000000000000..3747be705379 --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py @@ -0,0 +1,430 @@ +# Copyright 2026 The Kandinsky 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 torch +from tqdm.auto import tqdm + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models.transformers.transformer_kandinsky_s2v import KandinskyS2VTransformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler, PiflowScheduler +from ...utils import logging +from ...utils.torch_utils import randn_tensor +from ..modular_pipeline import ( + BlockState, + LoopSequentialPipelineBlocks, + ModularPipelineBlocks, + PipelineState, +) +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .before_denoise import memory_rope_pos, roll_memory +from .modular_pipeline import KandinskyS2VModularPipeline, align, audio_rope_pos + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class KandinskyS2VChunkNoiseStep(ModularPipelineBlocks): + model_name = "kandinsky-s2v" + + @property + def description(self) -> str: + return "Draws the initial noise of one chunk and selects the speech features that chunk is driven by." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", KandinskyS2VTransformer3DModel)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("latent_height", required=True, type_hint=int), + InputParam("latent_width", required=True, type_hint=int), + InputParam("num_latent_frames", default=32, type_hint=int), + InputParam("audio_chunks", required=True, type_hint=list), + InputParam.template("generator"), + ] + + @torch.no_grad() + def __call__( + self, components: KandinskyS2VModularPipeline, block_state: BlockState, k: int + ) -> tuple[KandinskyS2VModularPipeline, BlockState]: + config = components.transformer.config + shape = ( + block_state.num_latent_frames, + align(block_state.latent_height, config.patch_size[1]), + align(block_state.latent_width, config.patch_size[2]), + config.in_visual_dim, + ) + block_state.latents = randn_tensor( + shape, generator=block_state.generator, device=components._execution_device, dtype=torch.float32 + ) + block_state.chunk_audio = block_state.audio_chunks[k] + return components, block_state + + +class KandinskyS2VChunkSchedulerResetStep(ModularPipelineBlocks): + model_name = "kandinsky-s2v" + + @property + def description(self) -> str: + return ( + "Resets the scheduler for one chunk. The flow-match scheduler is handed the raw grid the native sampler " + "uses and applies the shift itself, which leaves `num_inference_steps - 1` Euler updates." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [InputParam.template("num_inference_steps")] + + @torch.no_grad() + def __call__( + self, components: KandinskyS2VModularPipeline, block_state: BlockState, k: int + ) -> tuple[KandinskyS2VModularPipeline, BlockState]: + if block_state.num_inference_steps < 2: + raise ValueError( + f"`num_inference_steps` must be at least 2, got {block_state.num_inference_steps}: " + "the schedule runs `num_inference_steps - 1` Euler updates." + ) + # torch.linspace in float32, not numpy in float64: a one-ULP difference in the step size moves the + # trajectory, which then diverges on its own. + grid = torch.linspace(1.0, 0.0, block_state.num_inference_steps, dtype=torch.float32)[:-1] + components.scheduler.set_timesteps(sigmas=grid.numpy(), device=components._execution_device) + block_state.timesteps = components.scheduler.timesteps + return components, block_state + + +class KandinskyS2VDistilledChunkSchedulerResetStep(KandinskyS2VChunkSchedulerResetStep): + model_name = "kandinsky-s2v-distilled" + + @property + def description(self) -> str: + return "Resets the pi-Flow scheduler for one chunk: one model evaluation per segment." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", PiflowScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [InputParam.template("num_inference_steps", default=16)] + + @torch.no_grad() + def __call__( + self, components: KandinskyS2VModularPipeline, block_state: BlockState, k: int + ) -> tuple[KandinskyS2VModularPipeline, BlockState]: + components.scheduler.set_timesteps(block_state.num_inference_steps, device=components._execution_device) + block_state.timesteps = components.scheduler.timesteps + return components, block_state + + +def shared_transformer_kwargs( + transformer: KandinskyS2VTransformer3DModel, + reference_hidden_states: torch.Tensor, + memory: list[torch.Tensor], + visual_rope_pos: list[torch.Tensor], + reference_rope_pos: list[torch.Tensor], + num_audio_tokens: int, + num_latent_frames: int, + scale_factor: tuple[float, float, float], + device: torch.device, +) -> dict: + """The transformer inputs of a chunk that are the same in every pass.""" + return { + "reference_hidden_states": reference_hidden_states.unsqueeze(0), + "memory_hidden_states": tuple(m.unsqueeze(0).to(transformer.dtype) for m in reversed(memory)), + "visual_rope_pos": visual_rope_pos, + "reference_rope_pos": reference_rope_pos, + "memory_rope_pos": tuple(reversed(memory_rope_pos(transformer.config, memory, device))), + "audio_rope_pos": audio_rope_pos(num_audio_tokens, num_latent_frames, device), + "scale_factor": scale_factor, + "return_dict": False, + } + + +class KandinskyS2VChunkDenoiseInner(ModularPipelineBlocks): + model_name = "kandinsky-s2v" + + @property + def description(self) -> str: + return ( + "Inner denoising loop over the timesteps of one chunk. The guider runs the conditional pass and, when " + "guidance is on, the unconditional one, which swaps in the negative prompt *and* silences the speech, so " + "guidance is joint over text and speech." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", KandinskyS2VTransformer3DModel), + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 3.0}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt_embeds", required=True), + InputParam.template("negative_prompt_embeds"), + InputParam("pooled_prompt_embeds", required=True, type_hint=torch.Tensor), + InputParam("negative_pooled_prompt_embeds", type_hint=torch.Tensor), + InputParam("text_rope_pos", required=True, type_hint=torch.Tensor), + InputParam("negative_text_rope_pos", type_hint=torch.Tensor), + InputParam("reference_hidden_states", required=True, type_hint=torch.Tensor), + InputParam("visual_rope_pos", required=True, type_hint=list), + InputParam("reference_rope_pos", required=True, type_hint=list), + InputParam("memory", required=True, type_hint=list), + InputParam("num_latent_frames", default=32, type_hint=int), + InputParam( + "scale_factor", + default=(1.0, 2.0, 2.0), + type_hint=tuple, + description="Per-axis RoPE scale for the resolution bucket.", + ), + InputParam.template("num_inference_steps"), + ] + + @torch.no_grad() + def __call__( + self, components: KandinskyS2VModularPipeline, block_state: BlockState, k: int + ) -> tuple[KandinskyS2VModularPipeline, BlockState]: + transformer_dtype = components.transformer.dtype + latents = block_state.latents + chunk_audio = block_state.chunk_audio + + # Only these differ between the conditional and the unconditional pass. + negative_embeds = block_state.negative_prompt_embeds + guider_inputs = { + "encoder_hidden_states": ( + block_state.prompt_embeds.unsqueeze(0), + None if negative_embeds is None else negative_embeds.unsqueeze(0), + ), + "audio_hidden_states": (chunk_audio.unsqueeze(0), torch.zeros_like(chunk_audio).unsqueeze(0)), + "pooled_projections": (block_state.pooled_prompt_embeds, block_state.negative_pooled_prompt_embeds), + "text_rope_pos": (block_state.text_rope_pos, block_state.negative_text_rope_pos), + } + # Tuples here are not (cond, uncond) pairs, so they stay out of the guider's inputs. + shared_kwargs = shared_transformer_kwargs( + transformer=components.transformer, + reference_hidden_states=block_state.reference_hidden_states, + memory=block_state.memory, + visual_rope_pos=block_state.visual_rope_pos, + reference_rope_pos=block_state.reference_rope_pos, + num_audio_tokens=block_state.chunk_audio.shape[0], + num_latent_frames=block_state.num_latent_frames, + scale_factor=block_state.scale_factor, + device=components._execution_device, + ) + + timesteps = block_state.timesteps + for i, t in enumerate(tqdm(timesteps, desc=f"chunk {k + 1}", leave=False)): + timestep = t.expand(1).to(latents.dtype) + latent_model_input = latents.unsqueeze(0).to(transformer_dtype) + + components.guider.set_state(step=i, num_inference_steps=len(timesteps), timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = {key: getattr(guider_state_batch, key) for key in guider_inputs} + context_name = getattr(guider_state_batch, components.guider._identifier_key) + with components.transformer.cache_context(context_name): + guider_state_batch.noise_pred = components.transformer( + hidden_states=latent_model_input, + timestep=timestep, + **cond_kwargs, + **shared_kwargs, + )[0] + components.guider.cleanup_models(components.transformer) + noise_pred = components.guider(guider_state)[0] + + # float() keeps the scheduler from handing the latents back in bf16 + latents = components.scheduler.step(noise_pred.squeeze(0).float(), t, latents, return_dict=False)[0] + + block_state.latents = latents + return components, block_state + + +class KandinskyS2VDistilledChunkDenoiseInner(KandinskyS2VChunkDenoiseInner): + model_name = "kandinsky-s2v-distilled" + + @property + def description(self) -> str: + return ( + "Inner denoising loop over the segments of one chunk for the distilled checkpoints: one conditional pass " + "per segment, since classifier-free guidance is distilled into the student." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", KandinskyS2VTransformer3DModel), + ComponentSpec("scheduler", PiflowScheduler), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + param if param.name != "num_inference_steps" else InputParam.template("num_inference_steps", default=16) + for param in super().inputs + if not param.name.startswith("negative_") + ] + + @torch.no_grad() + def __call__( + self, components: KandinskyS2VModularPipeline, block_state: BlockState, k: int + ) -> tuple[KandinskyS2VModularPipeline, BlockState]: + transformer_dtype = components.transformer.dtype + latents = block_state.latents + shared_kwargs = shared_transformer_kwargs( + transformer=components.transformer, + reference_hidden_states=block_state.reference_hidden_states, + memory=block_state.memory, + visual_rope_pos=block_state.visual_rope_pos, + reference_rope_pos=block_state.reference_rope_pos, + num_audio_tokens=block_state.chunk_audio.shape[0], + num_latent_frames=block_state.num_latent_frames, + scale_factor=block_state.scale_factor, + device=components._execution_device, + ) + + timesteps = block_state.timesteps + for t in tqdm(timesteps, desc=f"chunk {k + 1}", leave=False): + noise_pred = components.transformer( + hidden_states=latents.unsqueeze(0).to(transformer_dtype), + timestep=t.expand(1).to(latents.dtype), + encoder_hidden_states=block_state.prompt_embeds.unsqueeze(0), + audio_hidden_states=block_state.chunk_audio.unsqueeze(0), + pooled_projections=block_state.pooled_prompt_embeds, + text_rope_pos=block_state.text_rope_pos, + **shared_kwargs, + )[0] + # float() keeps the scheduler from handing the latents back in bf16 + latents = components.scheduler.step(noise_pred.squeeze(0).float(), t, latents, return_dict=False)[0] + + block_state.latents = latents + return components, block_state + + +class KandinskyS2VChunkUpdateStep(ModularPipelineBlocks): + model_name = "kandinsky-s2v" + + @property + def description(self) -> str: + return ( + "Collects the denoised latents of one chunk and rolls its tail into the FramePack memory the next chunk " + "starts from." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", KandinskyS2VTransformer3DModel)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("memory", required=True, type_hint=list), + InputParam("num_chunks", required=True, type_hint=int), + ] + + @torch.no_grad() + def __call__( + self, components: KandinskyS2VModularPipeline, block_state: BlockState, k: int + ) -> tuple[KandinskyS2VModularPipeline, BlockState]: + block_state.latent_chunks.append(block_state.latents) + # Not after the last chunk: nothing consumes that memory, and a chunk shorter than the FramePack window + # would fail here having already spent the whole denoising loop. + if k + 1 < block_state.num_chunks: + block_state.memory = roll_memory(components.transformer.config, block_state.latents) + return components, block_state + + +class KandinskyS2VChunkLoopWrapper(LoopSequentialPipelineBlocks): + model_name = "kandinsky-s2v" + + @property + def description(self) -> str: + return ( + "Pipeline block that iterates over the chunks of the video. Each chunk is denoised from fresh noise, " + "conditioned on the reference frame, its own slice of the speech and the FramePack memory of the chunks " + "before it." + ) + + @property + def loop_inputs(self) -> list[InputParam]: + return [InputParam("num_chunks", required=True, type_hint=int)] + + @property + def loop_intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latent_chunks", + type_hint=list, + description="Denoised latents of each chunk, each shaped " + "(num_latent_frames, latent_height, latent_width, channels).", + ) + ] + + @torch.no_grad() + def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.latent_chunks = [] + for k in range(block_state.num_chunks): + components, block_state = self.loop_step(components, block_state, k=k) + self.set_block_state(state, block_state) + return components, state + + +class KandinskyS2VChunkDenoiseStep(KandinskyS2VChunkLoopWrapper): + block_classes = [ + KandinskyS2VChunkNoiseStep, + KandinskyS2VChunkSchedulerResetStep, + KandinskyS2VChunkDenoiseInner, + KandinskyS2VChunkUpdateStep, + ] + block_names = ["noise", "scheduler_reset", "denoise_inner", "update_chunk"] + + @property + def description(self) -> str: + return ( + "Chunk denoise step that iterates over the chunks of the video.\n" + "At each chunk: noise -> scheduler_reset -> denoise_inner -> update_chunk." + ) + + +class KandinskyS2VDistilledChunkDenoiseStep(KandinskyS2VChunkLoopWrapper): + model_name = "kandinsky-s2v-distilled" + + block_classes = [ + KandinskyS2VChunkNoiseStep, + KandinskyS2VDistilledChunkSchedulerResetStep, + KandinskyS2VDistilledChunkDenoiseInner, + KandinskyS2VChunkUpdateStep, + ] + block_names = ["noise", "scheduler_reset", "denoise_inner", "update_chunk"] + + @property + def description(self) -> str: + return ( + "Chunk denoise step of the distilled checkpoints, with the pi-Flow scheduler and no guidance.\n" + "At each chunk: noise -> scheduler_reset -> denoise_inner -> update_chunk." + ) diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py new file mode 100644 index 000000000000..a248582ae732 --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py @@ -0,0 +1,515 @@ +# Copyright 2026 The Kandinsky 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 typing import Any + +import numpy as np +import PIL.Image +import torch +from transformers import ( + CLIPTextModel, + CLIPTokenizer, + Qwen2_5OmniProcessor, + Qwen2_5OmniThinkerForConditionalGeneration, +) + +from ...configuration_utils import FrozenDict +from ...models import AutoencoderKLHunyuanVideo +from ...utils import logging +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import ( + AUDIO_SAMPLE_RATE, + AUDIO_TOKEN_TOLERANCE, + AUDIO_TOKENS_PER_SECOND, + DEFAULT_NEGATIVE_PROMPT, + DEFAULT_PROMPT, + SYSTEM_PROMPT, + TEMPORAL_COMPRESSION_RATIO, + KandinskyS2VModularPipeline, + pad_waveform_with_silence, + select_resolution, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def encode_qwen_omni( + tokenizer: Qwen2_5OmniProcessor, + text_encoder: Qwen2_5OmniThinkerForConditionalGeneration, + prompt: str, + audio: np.ndarray | None, + audio_end: float | None, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Thinker hidden states for the prompt and audio-tower features for the speech, a waveform at 16 kHz.""" + user_content: list[dict[str, Any]] = [{"type": "text", "text": prompt}] + if audio is not None: + user_content.append({"type": "audio", "audio": audio, "audio_start": 0.0, "audio_end": audio_end}) + conversation = [ + {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}, + {"role": "user", "content": user_content}, + ] + + chat_text = tokenizer.apply_chat_template([conversation], add_generation_prompt=False, tokenize=False) + audios = None + if audio is not None: + # Cut to the seconds the chunks cover. + audios = [audio[: int(AUDIO_SAMPLE_RATE * audio_end)]] + + inputs = tokenizer( + text=chat_text, + audio=audios, + return_tensors="pt", + padding=True, + use_audio_in_video=False, + ) + inputs = inputs.to(device).to(text_encoder.dtype) + + output = text_encoder(**inputs, output_hidden_states=True, return_dict=True) + prompt_embeds = slice_user_text(tokenizer, text_encoder, output.hidden_states[-1], inputs) + + audio_embeds = None + if audio is not None: + audio_embeds = text_encoder.get_audio_features( + inputs["input_features"], inputs["feature_attention_mask"] + ).last_hidden_state + return prompt_embeds, audio_embeds + + +def slice_user_text( + tokenizer: Qwen2_5OmniProcessor, + text_encoder: Qwen2_5OmniThinkerForConditionalGeneration, + hidden_states: torch.Tensor, + inputs: Any, +) -> torch.Tensor: + """Keep only the hidden states of the user's own text span.""" + input_ids = inputs["input_ids"] + attention_mask = inputs.get("attention_mask", torch.ones_like(input_ids)) + seq_len = input_ids.shape[1] + device = input_ids.device + + im_start_id = tokenizer.tokenizer.convert_tokens_to_ids("<|im_start|>") + audio_start_token_id = text_encoder.config.audio_start_token_id + user_role_len = len(tokenizer.tokenizer("user\n", add_special_tokens=False).input_ids) + + positions = torch.arange(seq_len, device=device).unsqueeze(0) + last_im_start = torch.where(input_ids == im_start_id, positions, -1).max(dim=1).values + start = (last_im_start + 1 + user_role_len).unsqueeze(1) + end = torch.where(input_ids == audio_start_token_id, positions, seq_len - 2).min(dim=1).values.unsqueeze(1) + mask = (positions >= start) & (positions < end) & (attention_mask == 1) + return hidden_states[mask] + + +def encode_clip( + tokenizer_2: CLIPTokenizer, text_encoder_2: CLIPTextModel, prompt: str, device: torch.device +) -> torch.Tensor: + inputs = tokenizer_2( + [prompt], + max_length=77, + truncation=True, + add_special_tokens=True, + padding="max_length", + return_tensors="pt", + ).to(device) + return text_encoder_2(**inputs)["pooler_output"] + + +def encode_prompt_and_speech( + tokenizer: Qwen2_5OmniProcessor, + text_encoder: Qwen2_5OmniThinkerForConditionalGeneration, + tokenizer_2: CLIPTokenizer, + text_encoder_2: CLIPTextModel, + prompt: str, + audio: torch.Tensor | np.ndarray, + num_latent_frames: int, + fps: int, + num_chunks: int | None, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, float]: + """Plan the chunks the speech track needs, then encode the prompt and the speech for them. + + Returns the prompt embeddings, the pooled prompt embedding, the speech features, the number of chunks and the + length of the track in seconds. + """ + waveform = audio.detach().cpu().float().numpy() if isinstance(audio, torch.Tensor) else audio.astype(np.float32) + chunk_seconds = ((num_latent_frames - 1) * TEMPORAL_COMPRESSION_RATIO + 1) / fps + track_seconds = waveform.shape[0] / AUDIO_SAMPLE_RATE + needed = max(1, math.ceil(track_seconds / chunk_seconds)) + if num_chunks is not None and num_chunks > needed: + # The video is trimmed to the track, so these are computed and then dropped. + logger.warning( + f"{num_chunks - needed} chunk(s) past the end of the speech track are generated and then trimmed; " + f"pass num_chunks={needed} to save the time" + ) + if num_chunks is None: + num_chunks = needed + target_seconds = num_chunks * chunk_seconds + + # The speech encoder truncates its input; chunks past its limit would split too few speech tokens and drift + # out of sync. Refused here, before any encoding is spent on it. + extractor = tokenizer.feature_extractor + max_seconds = extractor.n_samples / extractor.sampling_rate + if target_seconds > max_seconds + 0.01: + raise ValueError( + f"{num_chunks} chunks need {target_seconds:g} s of speech, but the speech encoder truncates at " + f"{max_seconds:g} s; trim the track or pass fewer `num_chunks`" + ) + + prompt_embeds, audio_embeds = encode_qwen_omni( + tokenizer, text_encoder, prompt, pad_waveform_with_silence(waveform, target_seconds), target_seconds, device + ) + # The feature extractor truncates beyond its own chunk_length, and a short tensor would just be split + # thinner, desyncing the speech silently. Checked rather than trusted. + expected_tokens = round(target_seconds * AUDIO_TOKENS_PER_SECOND) + if audio_embeds.shape[0] < expected_tokens * AUDIO_TOKEN_TOLERANCE: + raise ValueError( + f"the speech encoder returned {audio_embeds.shape[0]} tokens for {target_seconds:.1f} s, " + f"expected about {expected_tokens}; the track is likely longer than the encoder accepts" + ) + pooled_prompt_embeds = encode_clip(tokenizer_2, text_encoder_2, prompt, device) + return prompt_embeds, pooled_prompt_embeds, audio_embeds, num_chunks, track_seconds + + +class KandinskyS2VTextAudioEncoderStep(ModularPipelineBlocks): + model_name = "kandinsky-s2v" + + @property + def description(self) -> str: + return ( + "Encoder step that plans the chunks the speech track needs and turns the prompt and the speech into the " + "transformer's conditions: Qwen2.5-Omni Thinker hidden states for the prompt, its audio-tower features " + "for the speech, and a CLIP pooled embedding. The negative prompt is encoded when it is passed, or when " + "the pipeline's guider needs an unconditional pass." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen2_5OmniThinkerForConditionalGeneration), + ComponentSpec("tokenizer", Qwen2_5OmniProcessor), + ComponentSpec("text_encoder_2", CLIPTextModel), + ComponentSpec("tokenizer_2", CLIPTokenizer), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template( + "prompt", required=False, default=DEFAULT_PROMPT, description="The prompt to guide the generation." + ), + InputParam.template( + "negative_prompt", + description="The prompt for the unconditional pass of the guidance. Defaults to " + "`DEFAULT_NEGATIVE_PROMPT` when the guidance needs one.", + ), + InputParam( + "audio", + required=True, + type_hint=torch.Tensor | np.ndarray, + description="The speech track that drives the generation: a mono float waveform in [-1, 1] at " + f"{AUDIO_SAMPLE_RATE} Hz, shaped (samples,).", + ), + InputParam( + "num_latent_frames", + default=32, + type_hint=int, + description="Latent frames per chunk; 32 gives 125 frames (5 s at 25 fps).", + ), + InputParam("fps", default=25, type_hint=int, description="Frame rate the chunk length is computed at."), + InputParam( + "num_chunks", + type_hint=int, + description="Number of chunks to generate. Defaults to however many cover the speech track; the " + "video is trimmed to the track either way.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("prompt_embeds"), + OutputParam.template("negative_prompt_embeds"), + OutputParam("pooled_prompt_embeds", type_hint=torch.Tensor, description="CLIP pooled prompt embedding."), + OutputParam( + "negative_pooled_prompt_embeds", + type_hint=torch.Tensor, + description="CLIP pooled negative prompt embedding.", + ), + OutputParam( + "audio_embeds", type_hint=torch.Tensor, description="Audio-tower features of the speech track." + ), + OutputParam("num_chunks", type_hint=int, description="Number of chunks to generate."), + OutputParam("track_seconds", type_hint=float, description="Length of the speech track in seconds."), + ] + + @staticmethod + def check_inputs(block_state): + if not isinstance(block_state.prompt, str): + raise ValueError("`prompt` must be a string; batching is not supported") + # Checked here because an empty prompt fails as a reshape error deep inside the transformer. + if not block_state.prompt.strip(): + raise ValueError("`prompt` must not be empty; pass DEFAULT_PROMPT to use the reference prompt") + audio = block_state.audio + if not isinstance(audio, (torch.Tensor, np.ndarray)): + raise ValueError( + f"`audio` must be a mono float waveform at {AUDIO_SAMPLE_RATE} Hz as a tensor or an array, got " + f"{type(audio).__name__}; load a file with e.g. `librosa.load(path, sr={AUDIO_SAMPLE_RATE})[0]`" + ) + is_float = ( + audio.is_floating_point() if isinstance(audio, torch.Tensor) else np.issubdtype(audio.dtype, np.floating) + ) + if audio.ndim != 1 or audio.shape[0] == 0 or not is_float: + raise ValueError( + f"`audio` must be a non-empty mono float waveform in [-1, 1] at {AUDIO_SAMPLE_RATE} Hz, shaped " + f"(samples,); got shape {tuple(audio.shape)} and dtype {audio.dtype}" + ) + if block_state.num_chunks is not None and block_state.num_chunks < 1: + raise ValueError(f"`num_chunks` must be >= 1, got {block_state.num_chunks}") + if block_state.num_latent_frames < 1 or block_state.fps <= 0: + raise ValueError( + f"`num_latent_frames` and `fps` must be positive, got {block_state.num_latent_frames} and " + f"{block_state.fps}" + ) + + @torch.no_grad() + def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + if block_state.negative_prompt is None and components.requires_unconditional_embeds: + block_state.negative_prompt = DEFAULT_NEGATIVE_PROMPT + if not isinstance(block_state.negative_prompt, (str, type(None))): + raise ValueError("`negative_prompt` must be a string; batching is not supported") + # The same reshape error as an empty prompt, once the negative embeddings reach a guided denoise. + if block_state.negative_prompt is not None and not block_state.negative_prompt.strip(): + raise ValueError("`negative_prompt` must not be empty; omit it to skip the unconditional pass") + + device = components._execution_device + ( + block_state.prompt_embeds, + block_state.pooled_prompt_embeds, + block_state.audio_embeds, + block_state.num_chunks, + block_state.track_seconds, + ) = encode_prompt_and_speech( + tokenizer=components.tokenizer, + text_encoder=components.text_encoder, + tokenizer_2=components.tokenizer_2, + text_encoder_2=components.text_encoder_2, + prompt=block_state.prompt, + audio=block_state.audio, + num_latent_frames=block_state.num_latent_frames, + fps=block_state.fps, + num_chunks=block_state.num_chunks, + device=device, + ) + # The guider is not a component of this block: inside the pipeline, `requires_unconditional_embeds` reads the + # denoise step's guider and the negative prompt defaults to `DEFAULT_NEGATIVE_PROMPT`; standalone, there is no + # guider and the negative prompt is encoded only when it is passed. + block_state.negative_prompt_embeds = None + block_state.negative_pooled_prompt_embeds = None + guidance_off = getattr(components, "guider", None) is not None and not components.requires_unconditional_embeds + if block_state.negative_prompt is not None and guidance_off: + logger.warning("`negative_prompt` is ignored: the guider runs no unconditional pass") + elif block_state.negative_prompt is not None: + block_state.negative_prompt_embeds, _ = encode_qwen_omni( + components.tokenizer, components.text_encoder, block_state.negative_prompt, None, None, device + ) + block_state.negative_pooled_prompt_embeds = encode_clip( + components.tokenizer_2, components.text_encoder_2, block_state.negative_prompt, device + ) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class KandinskyS2VDistilledTextAudioEncoderStep(KandinskyS2VTextAudioEncoderStep): + """ + Encoder step of the distilled checkpoints, which sample without classifier-free guidance and so take no negative + prompt: it plans the chunks the speech track needs and encodes the prompt and the speech. + + Components: + text_encoder (`Qwen2_5OmniThinkerForConditionalGeneration`) tokenizer (`Qwen2_5OmniProcessor`) text_encoder_2 + (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) + + Inputs: + prompt (`str`, *optional*, defaults to The video presents a stable medium close-up of a subject facing the + camera directly. The subject exhibits natural, lifelike body dynamics: the head and shoulders move fluidly + with the rhythm of speech. The lighting is even, maintaining a balanced natural exposure to clearly define + facial features without washout. Throughout the sequence, the subject engages in continuous speech marked by + precise, high-fidelity articulation. The lips move in a rhythmic, natural pattern: alternating between brief + firm closures and relaxed openings. During these openings, the tips of the upper teeth are momentarily and + realistically visible, exhibiting natural enamel texture. The upper lip lifts subtly over the teeth during + enunciation, while the jaw performs minimal vertical oscillations synchronized with the speech cadence. + Facial micro-expressions are active yet restrained; eyebrows rise and lower slightly with minute asymmetry. + The eyes maintain a steady, engaged gaze forward, punctuated by smooth, deliberate blinks. ): + The prompt to guide the generation. + audio (`Tensor | ndarray`): + The speech track that drives the generation: a mono float waveform in [-1, 1] at 16000 Hz, shaped + (samples,). + num_latent_frames (`int`, *optional*, defaults to 32): + Latent frames per chunk; 32 gives 125 frames (5 s at 25 fps). + fps (`int`, *optional*, defaults to 25): + Frame rate the chunk length is computed at. + num_chunks (`int`, *optional*): + Number of chunks to generate. Defaults to however many cover the speech track; the video is trimmed to + the track either way. + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + pooled_prompt_embeds (`Tensor`): + CLIP pooled prompt embedding. + audio_embeds (`Tensor`): + Audio-tower features of the speech track. + num_chunks (`int`): + Number of chunks to generate. + track_seconds (`float`): + Length of the speech track in seconds. + """ + + model_name = "kandinsky-s2v-distilled" + + @property + def description(self) -> str: + return ( + "Encoder step of the distilled checkpoints, which sample without classifier-free guidance and so take no " + "negative prompt: it plans the chunks the speech track needs and encodes the prompt and the speech." + ) + + @property + def inputs(self) -> list[InputParam]: + return [param for param in super().inputs if param.name != "negative_prompt"] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [param for param in super().intermediate_outputs if not param.name.startswith("negative_")] + + @torch.no_grad() + def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + ( + block_state.prompt_embeds, + block_state.pooled_prompt_embeds, + block_state.audio_embeds, + block_state.num_chunks, + block_state.track_seconds, + ) = encode_prompt_and_speech( + tokenizer=components.tokenizer, + text_encoder=components.text_encoder, + tokenizer_2=components.tokenizer_2, + text_encoder_2=components.text_encoder_2, + prompt=block_state.prompt, + audio=block_state.audio, + num_latent_frames=block_state.num_latent_frames, + fps=block_state.fps, + num_chunks=block_state.num_chunks, + device=components._execution_device, + ) + self.set_block_state(state, block_state) + return components, state + + +class KandinskyS2VVaeEncoderStep(ModularPipelineBlocks): + model_name = "kandinsky-s2v" + + @property + def description(self) -> str: + return ( + "VAE encoder step that picks the resolution bucket closest to the reference frame unless one is given, " + "and encodes the frame, resized and center-cropped, into a scaled latent." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLHunyuanVideo), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 8}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image", type_hint=PIL.Image.Image, description="The reference frame to animate."), + InputParam.template( + "height", + description="The height in pixels of the generated video. Defaults to the 512 bucket closest to the " + "reference frame.", + ), + InputParam.template( + "width", description="The width in pixels of the generated video, chosen together with `height`." + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "reference_latents", + type_hint=torch.Tensor, + description="Scaled VAE latent of the reference frame, (1, H/8, W/8, C).", + ) + ] + + @staticmethod + def check_inputs(block_state): + if not isinstance(block_state.image, PIL.Image.Image): + raise ValueError( + "`image` is required and must be a PIL image; load a file with `diffusers.utils.load_image`" + ) + if (block_state.height is None) != (block_state.width is None): + raise ValueError("pass both `height` and `width`, or neither to pick a bucket from the image") + if block_state.height is not None and (block_state.height % 16 or block_state.width % 16): + raise ValueError( + f"`height` and `width` must be multiples of 16, got {block_state.height} and {block_state.width}" + ) + + @torch.no_grad() + def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state) + device = components._execution_device + image = block_state.image.convert("RGB") + height, width = block_state.height, block_state.width + if height is None: + height, width = select_resolution(image) + + # Resized to cover the bucket and center-cropped with the rounding the model was trained with; the + # processor's own `resize_mode="crop"` rounds differently and shifts some aspect ratios by a pixel. + scale = max(width / image.width, height / image.height) + resized = image.resize( + (round(image.width * scale), round(image.height * scale)), resample=PIL.Image.Resampling.BILINEAR + ) + left, top = round((resized.width - width) / 2), round((resized.height - height) / 2) + cropped = resized.crop((left, top, left + width, top + height)) + + tensor = components.video_processor.preprocess(cropped, height=height, width=width) + tensor = tensor.to(device=device, dtype=components.vae.dtype).unsqueeze(2) # (1, 3, 1, H, W) + # The posterior is sampled, so it takes the generator like every other draw. + latent = components.vae.encode(tensor).latent_dist.sample(block_state.generator).permute(0, 2, 3, 4, 1) + block_state.reference_latents = (latent * components.vae.config.scaling_factor).squeeze(1) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v.py b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v.py new file mode 100644 index 000000000000..099b4353a25f --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v.py @@ -0,0 +1,178 @@ +# Copyright 2026 The Kandinsky 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 ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import KandinskyS2VPrepareNegativeStep, KandinskyS2VPrepareStep +from .decoders import KandinskyS2VDecodeStep +from .denoise import KandinskyS2VChunkDenoiseStep +from .encoders import KandinskyS2VTextAudioEncoderStep, KandinskyS2VVaeEncoderStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +KandinskyS2VCoreDenoiseBlocks = InsertableDict( + [ + ("prepare", KandinskyS2VPrepareStep()), + ("prepare_negative", KandinskyS2VPrepareNegativeStep()), + ("chunk_loop", KandinskyS2VChunkDenoiseStep()), + ] +) + + +# auto_docstring +class KandinskyS2VCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that prepares the conditions and the FramePack memory, then denoises the video chunk by chunk. + + Components: + transformer (`KandinskyS2VTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) + + Inputs: + reference_latents (`Tensor`): + Scaled VAE latent of the reference frame, from the VAE encoder step. + audio_embeds (`Tensor`): + Audio-tower features of the speech track, from the text encoder step. + num_chunks (`int`): + Number of chunks to generate. + num_latent_frames (`int`, *optional*, defaults to 32): + Latent frames per chunk; 32 gives 125 frames (5 s at 25 fps). + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + pooled_prompt_embeds (`Tensor`): + CLIP pooled prompt embedding. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + negative_pooled_prompt_embeds (`Tensor`, *optional*): + CLIP pooled negative prompt embedding. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + scale_factor (`tuple`, *optional*, defaults to (1.0, 2.0, 2.0)): + Per-axis RoPE scale for the resolution bucket. + + Outputs: + latent_chunks (`list`): + Denoised latents of each chunk, each shaped (num_latent_frames, latent_height, latent_width, channels). + """ + + model_name = "kandinsky-s2v" + block_classes = KandinskyS2VCoreDenoiseBlocks.values() + block_names = KandinskyS2VCoreDenoiseBlocks.keys() + + @property + def outputs(self): + return [ + OutputParam( + "latent_chunks", + type_hint=list, + description="Denoised latents of each chunk, each shaped " + "(num_latent_frames, latent_height, latent_width, channels).", + ) + ] + + @property + def description(self): + return ( + "Core denoise step that prepares the conditions and the FramePack memory, then denoises the video " + "chunk by chunk." + ) + + +AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", KandinskyS2VTextAudioEncoderStep()), + ("vae_encoder", KandinskyS2VVaeEncoderStep()), + ("denoise", KandinskyS2VCoreDenoiseStep()), + ("decode", KandinskyS2VDecodeStep()), + ] +) + + +# auto_docstring +class KandinskyS2VAutoBlocks(SequentialPipelineBlocks): + """ + Modular pipeline for speech-to-video generation with Kandinsky S2V 1.0: an avatar video from a reference frame and + a speech track, generated chunk by chunk. + + Components: + text_encoder (`Qwen2_5OmniThinkerForConditionalGeneration`) tokenizer (`Qwen2_5OmniProcessor`) text_encoder_2 + (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) vae (`AutoencoderKLHunyuanVideo`) video_processor + (`VideoProcessor`) transformer (`KandinskyS2VTransformer3DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) + + Inputs: + prompt (`str`, *optional*, defaults to The video presents a stable medium close-up of a subject facing the + camera directly. The subject exhibits natural, lifelike body dynamics: the head and shoulders move fluidly + with the rhythm of speech. The lighting is even, maintaining a balanced natural exposure to clearly define + facial features without washout. Throughout the sequence, the subject engages in continuous speech marked by + precise, high-fidelity articulation. The lips move in a rhythmic, natural pattern: alternating between brief + firm closures and relaxed openings. During these openings, the tips of the upper teeth are momentarily and + realistically visible, exhibiting natural enamel texture. The upper lip lifts subtly over the teeth during + enunciation, while the jaw performs minimal vertical oscillations synchronized with the speech cadence. + Facial micro-expressions are active yet restrained; eyebrows rise and lower slightly with minute asymmetry. + The eyes maintain a steady, engaged gaze forward, punctuated by smooth, deliberate blinks. ): + The prompt to guide the generation. + negative_prompt (`str`, *optional*): + The prompt for the unconditional pass of the guidance. Defaults to `DEFAULT_NEGATIVE_PROMPT` when the + guidance needs one. + audio (`Tensor | ndarray`): + The speech track that drives the generation: a mono float waveform in [-1, 1] at 16000 Hz, shaped + (samples,). + num_latent_frames (`int`, *optional*, defaults to 32): + Latent frames per chunk; 32 gives 125 frames (5 s at 25 fps). + fps (`int`, *optional*, defaults to 25): + Frame rate the chunk length is computed at. + num_chunks (`int`, *optional*): + Number of chunks to generate. Defaults to however many cover the speech track; the video is trimmed to + the track either way. + image (`Image`): + The reference frame to animate. + height (`int`, *optional*): + The height in pixels of the generated video. Defaults to the 512 bucket closest to the reference frame. + width (`int`, *optional*): + The width in pixels of the generated video, chosen together with `height`. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. + scale_factor (`tuple`, *optional*, defaults to (1.0, 2.0, 2.0)): + Per-axis RoPE scale for the resolution bucket. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "kandinsky-s2v" + + block_classes = AUTO_BLOCKS.values() + block_names = AUTO_BLOCKS.keys() + + @property + def description(self): + return ( + "Modular pipeline for speech-to-video generation with Kandinsky S2V 1.0: an avatar video from a reference " + "frame and a speech track, generated chunk by chunk." + ) + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v_distilled.py b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v_distilled.py new file mode 100644 index 000000000000..f8fbb72c15e6 --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v_distilled.py @@ -0,0 +1,169 @@ +# Copyright 2026 The Kandinsky 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 ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict, OutputParam +from .before_denoise import KandinskyS2VPrepareStep +from .decoders import KandinskyS2VDecodeStep +from .denoise import KandinskyS2VDistilledChunkDenoiseStep +from .encoders import KandinskyS2VDistilledTextAudioEncoderStep, KandinskyS2VVaeEncoderStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +KandinskyS2VDistilledCoreDenoiseBlocks = InsertableDict( + [ + ("prepare", KandinskyS2VPrepareStep()), + ("chunk_loop", KandinskyS2VDistilledChunkDenoiseStep()), + ] +) + + +# auto_docstring +class KandinskyS2VDistilledCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step of the distilled checkpoints that prepares the conditions and the FramePack memory, then denoises + the video chunk by chunk with the pi-Flow scheduler. + + Components: + transformer (`KandinskyS2VTransformer3DModel`) scheduler (`PiflowScheduler`) + + Inputs: + reference_latents (`Tensor`): + Scaled VAE latent of the reference frame, from the VAE encoder step. + audio_embeds (`Tensor`): + Audio-tower features of the speech track, from the text encoder step. + num_chunks (`int`): + Number of chunks to generate. + num_latent_frames (`int`, *optional*, defaults to 32): + Latent frames per chunk; 32 gives 125 frames (5 s at 25 fps). + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + pooled_prompt_embeds (`Tensor`): + CLIP pooled prompt embedding. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 16): + The number of denoising steps. + scale_factor (`tuple`, *optional*, defaults to (1.0, 2.0, 2.0)): + Per-axis RoPE scale for the resolution bucket. + + Outputs: + latent_chunks (`list`): + Denoised latents of each chunk, each shaped (num_latent_frames, latent_height, latent_width, channels). + """ + + model_name = "kandinsky-s2v-distilled" + block_classes = KandinskyS2VDistilledCoreDenoiseBlocks.values() + block_names = KandinskyS2VDistilledCoreDenoiseBlocks.keys() + + @property + def outputs(self): + return [ + OutputParam( + "latent_chunks", + type_hint=list, + description="Denoised latents of each chunk, each shaped " + "(num_latent_frames, latent_height, latent_width, channels).", + ) + ] + + @property + def description(self): + return ( + "Core denoise step of the distilled checkpoints that prepares the conditions and the FramePack memory, " + "then denoises the video chunk by chunk with the pi-Flow scheduler." + ) + + +DISTILLED_AUTO_BLOCKS = InsertableDict( + [ + ("text_encoder", KandinskyS2VDistilledTextAudioEncoderStep()), + ("vae_encoder", KandinskyS2VVaeEncoderStep()), + ("denoise", KandinskyS2VDistilledCoreDenoiseStep()), + ("decode", KandinskyS2VDecodeStep()), + ] +) + + +# auto_docstring +class KandinskyS2VDistilledAutoBlocks(SequentialPipelineBlocks): + """ + Modular pipeline for speech-to-video generation with the pi-Flow distilled Kandinsky S2V 1.0 checkpoints: 16 + segments per chunk and no guidance, which is distilled into the student. + + Components: + text_encoder (`Qwen2_5OmniThinkerForConditionalGeneration`) tokenizer (`Qwen2_5OmniProcessor`) text_encoder_2 + (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) vae (`AutoencoderKLHunyuanVideo`) video_processor + (`VideoProcessor`) transformer (`KandinskyS2VTransformer3DModel`) scheduler (`PiflowScheduler`) + + Inputs: + prompt (`str`, *optional*, defaults to The video presents a stable medium close-up of a subject facing the + camera directly. The subject exhibits natural, lifelike body dynamics: the head and shoulders move fluidly + with the rhythm of speech. The lighting is even, maintaining a balanced natural exposure to clearly define + facial features without washout. Throughout the sequence, the subject engages in continuous speech marked by + precise, high-fidelity articulation. The lips move in a rhythmic, natural pattern: alternating between brief + firm closures and relaxed openings. During these openings, the tips of the upper teeth are momentarily and + realistically visible, exhibiting natural enamel texture. The upper lip lifts subtly over the teeth during + enunciation, while the jaw performs minimal vertical oscillations synchronized with the speech cadence. + Facial micro-expressions are active yet restrained; eyebrows rise and lower slightly with minute asymmetry. + The eyes maintain a steady, engaged gaze forward, punctuated by smooth, deliberate blinks. ): + The prompt to guide the generation. + audio (`Tensor | ndarray`): + The speech track that drives the generation: a mono float waveform in [-1, 1] at 16000 Hz, shaped + (samples,). + num_latent_frames (`int`, *optional*, defaults to 32): + Latent frames per chunk; 32 gives 125 frames (5 s at 25 fps). + fps (`int`, *optional*, defaults to 25): + Frame rate the chunk length is computed at. + num_chunks (`int`, *optional*): + Number of chunks to generate. Defaults to however many cover the speech track; the video is trimmed to + the track either way. + image (`Image`): + The reference frame to animate. + height (`int`, *optional*): + The height in pixels of the generated video. Defaults to the 512 bucket closest to the reference frame. + width (`int`, *optional*): + The width in pixels of the generated video, chosen together with `height`. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 16): + The number of denoising steps. + scale_factor (`tuple`, *optional*, defaults to (1.0, 2.0, 2.0)): + Per-axis RoPE scale for the resolution bucket. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + + Outputs: + videos (`list`): + The generated videos. + """ + + model_name = "kandinsky-s2v-distilled" + + block_classes = DISTILLED_AUTO_BLOCKS.values() + block_names = DISTILLED_AUTO_BLOCKS.keys() + + @property + def description(self): + return ( + "Modular pipeline for speech-to-video generation with the pi-Flow distilled Kandinsky S2V 1.0 " + "checkpoints: 16 segments per chunk and no guidance, which is distilled into the student." + ) + + @property + def outputs(self): + return [OutputParam.template("videos")] diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py new file mode 100644 index 000000000000..33c29c89cec4 --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py @@ -0,0 +1,118 @@ +# Copyright 2026 The Kandinsky 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 numpy as np +import PIL.Image +import torch + +from ...models.transformers.transformer_kandinsky_s2v import audio_group_lengths +from ...utils import logging +from ..modular_pipeline import ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +SYSTEM_PROMPT = ( + "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, " + "capable of perceiving auditory and visual inputs, as well as generating " + "text and speech." +) +AUDIO_SAMPLE_RATE = 16000 +AUDIO_TOKENS_PER_SECOND = 25 +AUDIO_TOKEN_TOLERANCE = 0.95 +TEMPORAL_COMPRESSION_RATIO = 4 + +# The reference S2V prompts. Text conditioning is not optional: an empty prompt +# yields zero text tokens and the cross-attention has nothing to attend to. +DEFAULT_PROMPT = ( + "The video presents a stable medium close-up of a subject facing the camera directly. The subject " + "exhibits natural, lifelike body dynamics: the head and shoulders move fluidly with the rhythm of " + "speech. The lighting is even, maintaining a balanced natural exposure to clearly define facial " + "features without washout. Throughout the sequence, the subject engages in continuous speech marked by " + "precise, high-fidelity articulation. The lips move in a rhythmic, natural pattern: alternating " + "between brief firm closures and relaxed openings. During these openings, the tips of the upper teeth " + "are momentarily and realistically visible, exhibiting natural enamel texture. The upper lip lifts " + "subtly over the teeth during enunciation, while the jaw performs minimal vertical oscillations " + "synchronized with the speech cadence. Facial micro-expressions are active yet restrained; eyebrows " + "rise and lower slightly with minute asymmetry. The eyes maintain a steady, engaged gaze forward, " + "punctuated by smooth, deliberate blinks. " +) +DEFAULT_NEGATIVE_PROMPT = ( + "vivid color tones, static, motionless, frozen, subtitles, text, watermark, low quality, jpeg " + "artifacts, blurred, malformed anatomy, deformed hands, extra limbs, mutation, disfigured, teeth, " + "fangs, chaotic background, camera shake, painting, illustration " +) + + +def align(size: int, multiple: int) -> int: + return (size // multiple) * multiple + + +def select_resolution(image: PIL.Image.Image) -> tuple[int, int]: + """Pick the ``(height, width)`` bucket closest in aspect ratio to the image.""" + aspect = image.width / image.height + return min(((512, 512), (768, 512), (512, 768)), key=lambda res: abs(aspect - res[1] / res[0])) + + +def pad_waveform_with_silence(waveform: np.ndarray, target_duration: float) -> np.ndarray: + """Append silence so a waveform at `AUDIO_SAMPLE_RATE` covers ``target_duration``.""" + original = waveform.shape[0] / AUDIO_SAMPLE_RATE + if target_duration <= original + 0.01: + return waveform + pad_samples = int((target_duration - original) * AUDIO_SAMPLE_RATE) + return np.concatenate([waveform, np.zeros(pad_samples, dtype=waveform.dtype)]) + + +def split_audio_embeds(audio_embeds: torch.Tensor, chunks: int) -> list[torch.Tensor]: + """Split ``(S_audio, D)`` evenly over the chunks, dropping the ragged tail.""" + usable = (audio_embeds.shape[0] // chunks) * chunks + if usable == 0: + raise ValueError(f"{audio_embeds.shape[0]} audio tokens cannot be split into {chunks} chunks") + return list(audio_embeds[:usable].view(chunks, -1, audio_embeds.shape[1]).unbind(0)) + + +def audio_rope_pos(num_audio_tokens: int, duration: int, device: torch.device) -> torch.Tensor: + """Audio token -> index of the latent frame it belongs to.""" + lengths = torch.tensor(audio_group_lengths(num_audio_tokens, duration), device=device) + return torch.arange(duration, device=device).repeat_interleave(lengths).to(torch.int32) + + +class KandinskyS2VModularPipeline(ModularPipeline): + """ + A ModularPipeline for Kandinsky S2V 1.0 speech-to-video generation. + + """ + + default_blocks_name = "KandinskyS2VAutoBlocks" + + @property + def requires_unconditional_embeds(self): + requires_unconditional_embeds = False + + if hasattr(self, "guider") and self.guider is not None: + requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1 + + return requires_unconditional_embeds + + +class KandinskyS2VDistilledModularPipeline(KandinskyS2VModularPipeline): + """ + A ModularPipeline for the pi-Flow distilled Kandinsky S2V 1.0 checkpoints. + + Guidance is distilled into the student, so its blocks have no guider and no negative prompt, and a single forward + pass runs per segment. + + """ + + default_blocks_name = "KandinskyS2VDistilledAutoBlocks" diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index fc2c412b93df..979a80197936 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -210,6 +210,8 @@ def _helios_pyramid_map_fn(config_dict=None): ("cosmos3-omni", _create_default_map_fn("Cosmos3OmniModularPipeline")), ("helios", _create_default_map_fn("HeliosModularPipeline")), ("helios-pyramid", _helios_pyramid_map_fn), + ("kandinsky-s2v", _create_default_map_fn("KandinskyS2VModularPipeline")), + ("kandinsky-s2v-distilled", _create_default_map_fn("KandinskyS2VDistilledModularPipeline")), ("hunyuan-video-1.5", _create_default_map_fn("HunyuanVideo15ModularPipeline")), ("echo", _create_default_map_fn("EchoModularPipeline")), ("ltx", _create_default_map_fn("LTXModularPipeline")), diff --git a/src/diffusers/schedulers/scheduling_piflow.py b/src/diffusers/schedulers/scheduling_piflow.py index d858b425a9fd..7eaf7ccdbab8 100644 --- a/src/diffusers/schedulers/scheduling_piflow.py +++ b/src/diffusers/schedulers/scheduling_piflow.py @@ -195,6 +195,7 @@ def __init__( self.final_step_size_scale = float(final_step_size_scale) self.num_policy_substeps = int(num_policy_substeps) self._piflow_raw_timesteps = torch.empty(0) + self._piflow_segments = [] @property def step_index(self): @@ -243,10 +244,18 @@ def set_timesteps( raise ValueError(f"num_inference_steps must be positive, got {num_inference_steps}") one_minus_final = 1.0 - self.final_step_size_scale segment = 1.0 / (num_inference_steps - one_minus_final) - raw = 1.0 - torch.arange(num_inference_steps, dtype=torch.float32, device=device) * segment + segments = [segment] * (num_inference_steps - 1) + [segment * self.final_step_size_scale] + # Walked one segment at a time, as the native samplers do; `1 - i * segment` differs in the last bits. + raw_t = torch.ones((), dtype=torch.float32, device=device) + raws = [] + for segment in segments: + raws.append(raw_t) + raw_t = (raw_t - segment).clamp(min=self.eps) + raw = torch.stack(raws) sigmas = shift_timesteps(raw, float(self.config.shift)) self.num_inference_steps = int(num_inference_steps) self._piflow_raw_timesteps = raw + self._piflow_segments = segments self.timesteps = sigmas * self.config.num_train_timesteps self.sigmas = torch.cat([sigmas, sigmas.new_zeros(1)]) self._step_index = None @@ -273,15 +282,13 @@ def _to_grid(self, model_output: torch.Tensor, sample: torch.Tensor) -> torch.Te def _policy_step(self, model_output: torch.Tensor, sample: torch.Tensor, step_index: int) -> torch.Tensor: model_output = self._to_grid(model_output, sample) raw_src = self._piflow_raw_timesteps[step_index].to(device=sample.device) - raw_dst = ( - self._piflow_raw_timesteps[step_index + 1] - if step_index + 1 < self._piflow_raw_timesteps.numel() - else self._piflow_raw_timesteps.new_full((), self.eps) - ).to(device=sample.device) + # The policy sees the nominal segment, so the final one ends at 0; only the rollout stops at eps. + segment = self._piflow_segments[step_index] + raw_dst = (raw_src - segment).clamp(min=self.eps) sigma_src = self.sigmas[step_index].to(device=sample.device) token_shape = (sample.shape[0], *((sample.ndim - 1) * [1])) sigma = sigma_src.expand(sample.shape[0]).reshape(token_shape) - segment = (raw_src - raw_dst).expand(sample.shape[0]) + segment = raw_src.new_full((sample.shape[0],), segment) policy = DXPolicy( model_output, sample, diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index b13d46f68ff5..9a734961af1f 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1789,6 +1789,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class KandinskyS2VTransformer3DModel(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 Krea2Transformer2DModel(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 0f2b55484c0b..1b7b4d541a60 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -512,6 +512,66 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class KandinskyS2VAutoBlocks(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 KandinskyS2VDistilledAutoBlocks(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 KandinskyS2VDistilledModularPipeline(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 KandinskyS2VModularPipeline(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 Krea2AutoBlocks(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/transformers/test_models_transformer_kandinsky_s2v.py b/tests/models/transformers/test_models_transformer_kandinsky_s2v.py new file mode 100644 index 000000000000..2a9bb44f0152 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_kandinsky_s2v.py @@ -0,0 +1,198 @@ +# Copyright 2025 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 pytest +import torch + +from diffusers import KandinskyS2VTransformer3DModel +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() + + +# The visual stream is one packed sequence [memory | ref | visual], so the shapes +# below are tied together: memory is patched more coarsely than the noisy latent +# (1x, 2x and 4x), and the speech tokens have to cover every latent frame. +NUM_CHANNELS = 4 +DURATION = 4 +HEIGHT = 16 +WIDTH = 16 +PATCH_SIZE = (1, 2, 2) +FRAME_BUCKETS = (1, 2, 4) +TEXT_DIM = 8 +POOLED_DIM = 6 +AUDIO_DIM = 10 +NUM_TEXT_TOKENS = 7 +# One token for the first frame, four for each of the rest: >= DURATION, which the +# frame-aligned speech attention requires. +NUM_AUDIO_TOKENS = 1 + (DURATION - 1) * 4 +INDEX_ROPE_REF_IMAGE = 40 + + +def _rope_pos(duration: int, height: int, width: int, start: int = 0) -> list[torch.Tensor]: + return [ + torch.arange(start, start + duration, device=torch_device), + torch.arange(height, device=torch_device), + torch.arange(width, device=torch_device), + ] + + +class KandinskyS2VTransformerTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return KandinskyS2VTransformer3DModel + + @property + def main_input_name(self) -> str: + return "hidden_states" + + @property + def input_shape(self) -> tuple: + return (DURATION, HEIGHT, WIDTH, NUM_CHANNELS) + + @property + def output_shape(self) -> tuple: + return (DURATION, HEIGHT, WIDTH, NUM_CHANNELS) + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + @property + def torch_dtype(self): + return None + + def get_init_dict(self) -> dict: + return { + "in_visual_dim": NUM_CHANNELS, + "out_visual_dim": NUM_CHANNELS, + "in_text_dim": TEXT_DIM, + "in_text_dim2": POOLED_DIM, + "in_audio_dim": AUDIO_DIM, + "time_dim": 16, + "patch_size": PATCH_SIZE, + "frame_buckets": FRAME_BUCKETS, + "enable_framepack": True, + "index_rope_ref_image": INDEX_ROPE_REF_IMAGE, + "model_dim": 32, + "ff_dim": 64, + "num_text_blocks": 1, + "num_audio_blocks": 1, + "num_visual_blocks": 2, + "axes_dims": (4, 6, 6), + "audio_inject_layers": (0,), + "text_rope_max_pos": 128, + } + + def get_dummy_inputs(self, batch_size: int = 1, height: int = HEIGHT, width: int = WIDTH) -> dict: + dtype = self.torch_dtype + + def randn(*shape: int) -> torch.Tensor: + return randn_tensor(shape, generator=self.generator, device=torch_device, dtype=dtype) + + patch_h, patch_w = PATCH_SIZE[1], PATCH_SIZE[2] + # 1x keeps the visual patch, 2x doubles it and 4x quadruples it, so each + # memory buffer needs a frame count divisible by its temporal patch. + memory = ( + randn(batch_size, FRAME_BUCKETS[2], height, width, NUM_CHANNELS), + randn(batch_size, FRAME_BUCKETS[1], height, width, NUM_CHANNELS), + randn(batch_size, FRAME_BUCKETS[0], height, width, NUM_CHANNELS), + ) + memory_rope_pos = ( + _rope_pos(1, height // (patch_h * 4), width // (patch_w * 4), start=-FRAME_BUCKETS[2]), + _rope_pos(1, height // (patch_h * 2), width // (patch_w * 2), start=-FRAME_BUCKETS[1]), + _rope_pos(FRAME_BUCKETS[0], height // patch_h, width // patch_w, start=-FRAME_BUCKETS[0]), + ) + + return { + "hidden_states": randn(batch_size, DURATION, height, width, NUM_CHANNELS), + "encoder_hidden_states": randn(batch_size, NUM_TEXT_TOKENS, TEXT_DIM), + "audio_hidden_states": randn(batch_size, NUM_AUDIO_TOKENS, AUDIO_DIM), + "reference_hidden_states": randn(batch_size, 1, height, width, NUM_CHANNELS), + "memory_hidden_states": memory, + "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to( + torch_device, dtype=dtype or torch.float32 + ), + "pooled_projections": randn(batch_size, POOLED_DIM), + "visual_rope_pos": _rope_pos(DURATION, height // patch_h, width // patch_w), + "reference_rope_pos": _rope_pos(1, height // patch_h, width // patch_w, start=INDEX_ROPE_REF_IMAGE), + "memory_rope_pos": memory_rope_pos, + "text_rope_pos": torch.arange(NUM_TEXT_TOKENS, device=torch_device), + "audio_rope_pos": torch.arange(DURATION, device=torch_device).repeat_interleave( + torch.tensor([1] + [4] * (DURATION - 1), device=torch_device) + ), + } + + +class TestKandinskyS2VTransformer(KandinskyS2VTransformerTesterConfig, ModelTesterMixin): + def test_rope_tables_stay_fp32_after_cast(self): + model = self.model_class(**self.get_init_dict()).to(torch.bfloat16) + tables = [ + model.text_rope_embeddings.args, + model.audio_rope_embeddings.args, + *model.visual_rope_embeddings.args, + ] + assert all(table.dtype == torch.float32 for table in tables) + + +class TestKandinskyS2VTransformerMemory(KandinskyS2VTransformerTesterConfig, MemoryTesterMixin): + def test_layerwise_casting_memory(self): + pytest.skip( + "The affine normalisations take float32 input, as the checkpoints keep them in float32 " + "(`_keep_in_fp32_modules`); this test first casts the whole model to bfloat16 with `.to()`." + ) + + +class TestKandinskyS2VTransformerTorchCompile(KandinskyS2VTransformerTesterConfig, TorchCompileTesterMixin): + @property + def different_shapes_for_compilation(self): + # Divisible by the 4x memory patch. + return [(16, 16), (16, 32), (32, 32)] + + def test_torch_compile_repeated_blocks(self): + # The visual blocks come in two variants, with and without the speech adapter (`audio_inject_layers`), + # so we need recompile_limit=2 instead of the default 1. + super().test_torch_compile_repeated_blocks(recompile_limit=2) + + @pytest.mark.skip( + reason="AOTInductor lowers the lookup into the RoPE tables, which are plain float32 attributes rather than " + "buffers so that casting the model keeps them in float32, to a CPU kernel on a CUDA tensor." + ) + def test_compile_works_with_aot(self, tmp_path): + pass + + +class TestKandinskyS2VTransformerTraining(KandinskyS2VTransformerTesterConfig, TrainingTesterMixin): + def test_gradient_checkpointing_is_applied(self): + super().test_gradient_checkpointing_is_applied(expected_set={"KandinskyS2VTransformer3DModel"}) + + +class TestKandinskyS2VTransformerAttention(KandinskyS2VTransformerTesterConfig, AttentionTesterMixin): + def test_fuse_unfuse_qkv_projections(self, *args, **kwargs): + pytest.skip( + "KandinskyS2VAttention names its projections to_query/to_key/to_value/out_layer (matching the " + "training checkpoints) rather than to_q/to_k/to_v/to_out, which " + "AttentionModuleMixin.fuse_projections hardcodes." + ) diff --git a/tests/modular_pipelines/kandinsky_s2v/__init__.py b/tests/modular_pipelines/kandinsky_s2v/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py b/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py new file mode 100644 index 000000000000..8e665f82f21f --- /dev/null +++ b/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py @@ -0,0 +1,293 @@ +# coding=utf-8 +# Copyright 2026 The Kandinsky 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 pytest +import torch +from PIL import Image + +from diffusers.modular_pipelines import ( + KandinskyS2VAutoBlocks, + KandinskyS2VDistilledAutoBlocks, + KandinskyS2VDistilledModularPipeline, + KandinskyS2VModularPipeline, +) +from diffusers.modular_pipelines.kandinsky_s2v.encoders import KandinskyS2VTextAudioEncoderStep +from diffusers.modular_pipelines.kandinsky_s2v.modular_pipeline import DEFAULT_PROMPT + +from ...testing_utils import numpy_cosine_similarity_distance, require_accelerator, torch_device +from ..testing_utils import ( + BaseModularPipelineTesterConfig, + ModularGuiderTesterMixin, + ModularLoadingTesterMixin, + ModularMemoryTesterMixin, + ModularPipelineTesterMixin, + ModularWorkflowTesterMixin, +) + + +# One tiny repository per variant, laid out like the real checkpoints. +TINY_REPO = "hf-internal-testing/tiny-kandinsky-s2v-modular-pipe" +TINY_DISTILLED_REPO = "hf-internal-testing/tiny-kandinsky-s2v-distilled-modular-pipe" +AUDIO_SAMPLE_RATE = 16000 +# One chunk of NUM_LATENT_FRAMES covers ((8 - 1) * 4 + 1) / 25 = 1.16 s at 25 fps; the speech is shorter, so the +# video is trimmed back to it. NUM_LATENT_FRAMES has to reach the FramePack memory window, sum(frame_buckets) = 7. +NUM_LATENT_FRAMES = 8 + + +def dummy_waveform(seconds: int = 1) -> torch.Tensor: + """A plain tone, as a mono waveform at 16 kHz.""" + return 0.25 * torch.sin(torch.arange(seconds * AUDIO_SAMPLE_RATE) / 24) + + +def check_float16_inference(tester, expected_max_diff=5e-2): + """The mixin's fp16 check, with the pipeline loaded in fp16 rather than cast to it after loading. + + Loading in a dtype keeps `_keep_in_fp32_modules` in fp32; `.to(torch.float16)` would cast them too. + """ + outputs = [] + for dtype in (torch.float32, torch.float16): + pipe = tester.get_pipeline(dtype=dtype).to(torch_device) + outputs.append(pipe(**tester.get_dummy_inputs(), output=tester.output_name).float().cpu()) + max_diff = numpy_cosine_similarity_distance(outputs[0].flatten().numpy(), outputs[1].flatten().numpy()) + assert max_diff < expected_max_diff, f"FP16 inference is different from FP32 inference (max_diff: {max_diff})" + + +class KandinskyS2VModularPipelineTesterConfig(BaseModularPipelineTesterConfig): + pipeline_class = KandinskyS2VModularPipeline + pipeline_blocks_class = KandinskyS2VAutoBlocks + pretrained_model_name_or_path = TINY_REPO + params = frozenset(["image", "audio", "prompt", "height", "width", "num_latent_frames"]) + batch_params = frozenset(["prompt"]) + optional_params = frozenset(["num_inference_steps", "output_type"]) + output_name = "videos" + latents_output_name = "latent_chunks" + # One chunk, channels last: (num_latent_frames, latent_height, latent_width, channels). + expected_latents_shape = (NUM_LATENT_FRAMES, 8, 8, 4) + expected_workflow_defaults = { + None: { + "components": { + "text_encoder": "Qwen2_5OmniThinkerForConditionalGeneration", + "tokenizer": "Qwen2_5OmniProcessor", + "text_encoder_2": "CLIPTextModel", + "tokenizer_2": "CLIPTokenizer", + "vae": "AutoencoderKLHunyuanVideo", + "video_processor": "VideoProcessor", + "transformer": "KandinskyS2VTransformer3DModel", + "scheduler": "FlowMatchEulerDiscreteScheduler", + "guider": "ClassifierFreeGuidance", + }, + # The speech track drives the video; the prompt defaults to the one the checkpoints were trained with. + "required_inputs": ["image", "audio"], + "inputs": { + "prompt": DEFAULT_PROMPT, + "negative_prompt": None, + "num_latent_frames": 32, + "fps": 25, + "num_chunks": None, + "height": None, + "width": None, + "generator": None, + "num_inference_steps": 50, + "scale_factor": (1.0, 2.0, 2.0), + "output_type": "pil", + }, + "component_configs": {"video_processor": {"vae_scale_factor": 8}, "guider": {"guidance_scale": 3.0}}, + }, + } + + def get_dummy_inputs(self, seed=0): + return { + "image": Image.new("RGB", (64, 64), color="red"), + "audio": dummy_waveform(), + # Short prompts: the tiny transformer's text RoPE table is shorter than the default prompt. + "prompt": "a person speaking", + "negative_prompt": "static, motionless", + "height": 64, + "width": 64, + "num_latent_frames": NUM_LATENT_FRAMES, + "num_inference_steps": 2, + "generator": self.get_generator(seed), + "output_type": "pt", + } + + +class TestKandinskyS2VModularPipelineFast(KandinskyS2VModularPipelineTesterConfig, ModularPipelineTesterMixin): + @pytest.mark.skip(reason="The pipeline takes one prompt at a time; batching is not supported") + def test_inference_batch_consistent(self): + pass + + @pytest.mark.skip(reason="The pipeline takes one prompt at a time; batching is not supported") + def test_inference_batch_single_identical(self): + pass + + @pytest.mark.skip(reason="One video per call; there is no num_images_per_prompt") + def test_num_images_per_prompt(self): + pass + + @require_accelerator + def test_float16_inference(self): + check_float16_inference(self) + + def test_two_chunks_carry_memory(self): + """Two chunks end to end: the second is conditioned on the FramePack memory rolled from the first.""" + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs["num_chunks"] = 2 + latent_chunks = pipe(**inputs, output="latent_chunks") + assert len(latent_chunks) == 2 + assert all(chunk.shape[0] == NUM_LATENT_FRAMES and torch.isfinite(chunk).all() for chunk in latent_chunks) + + def test_rejects_speech_past_the_encoder_limit(self): + pipe = self.get_pipeline() + extractor = pipe.tokenizer.feature_extractor + limit = extractor.n_samples / extractor.sampling_rate + chunk_seconds = ((NUM_LATENT_FRAMES - 1) * 4 + 1) / 25 + inputs = self.get_dummy_inputs() + inputs["num_chunks"] = math.ceil(limit / chunk_seconds) + 1 + with pytest.raises(ValueError, match="truncates"): + pipe(**inputs, output="videos") + + def test_track_longer_than_the_chunks(self): + """The speech is cut to the chunks: a 2 s track in one 1.16 s chunk.""" + pipe = self.get_pipeline() + inputs = {**self.get_dummy_inputs(), "audio": dummy_waveform(seconds=2), "num_chunks": 1} + latent_chunks = pipe(**inputs, output="latent_chunks") + assert len(latent_chunks) == 1 and latent_chunks[0].shape[0] == NUM_LATENT_FRAMES + + def test_text_encoder_standalone_encodes_negative_prompt_only_when_passed(self): + """Standalone, the text encoder has no guider: a negative prompt is encoded only when it is passed.""" + text_pipe = KandinskyS2VTextAudioEncoderStep().init_pipeline(self.pretrained_model_name_or_path) + text_pipe.load_components() + inputs = {"prompt": "a person speaking", "audio": dummy_waveform(), "num_latent_frames": NUM_LATENT_FRAMES} + assert text_pipe(**inputs, output="negative_prompt_embeds") is None + assert text_pipe(**inputs, negative_prompt="static", output="negative_prompt_embeds") is not None + + @pytest.mark.parametrize( + "waveform", + [ + dummy_waveform().expand(2, -1), + torch.zeros(0), + (dummy_waveform() * 32768).to(torch.int16), + ], + ids=["multichannel", "empty", "int16"], + ) + def test_rejects_a_bad_waveform(self, waveform): + pipe = self.get_pipeline() + with pytest.raises(ValueError, match="mono float waveform"): + pipe(**{**self.get_dummy_inputs(), "audio": waveform}, output="videos") + + +class TestKandinskyS2VModularPipelineLoading(KandinskyS2VModularPipelineTesterConfig, ModularLoadingTesterMixin): + pass + + +class TestKandinskyS2VModularPipelineGuider(KandinskyS2VModularPipelineTesterConfig, ModularGuiderTesterMixin): + pass + + +class TestKandinskyS2VModularPipelineWorkflow(KandinskyS2VModularPipelineTesterConfig, ModularWorkflowTesterMixin): + pass + + +class TestKandinskyS2VModularPipelineMemory(KandinskyS2VModularPipelineTesterConfig, ModularMemoryTesterMixin): + pass + + +FAST_SKIPS = { + "test_inference_batch_consistent": "The pipeline takes one prompt at a time; batching is not supported", + "test_inference_batch_single_identical": "The pipeline takes one prompt at a time; batching is not supported", + "test_num_images_per_prompt": "One video per call; there is no num_images_per_prompt", +} + + +class KandinskyS2VDistilledModularPipelineTesterConfig(KandinskyS2VModularPipelineTesterConfig): + """The pi-Flow checkpoints: a head widened to its grid, `PiflowScheduler`, guidance off.""" + + pipeline_class = KandinskyS2VDistilledModularPipeline + pipeline_blocks_class = KandinskyS2VDistilledAutoBlocks + pretrained_model_name_or_path = TINY_DISTILLED_REPO + expected_workflow_defaults = { + None: { + "components": { + "text_encoder": "Qwen2_5OmniThinkerForConditionalGeneration", + "tokenizer": "Qwen2_5OmniProcessor", + "text_encoder_2": "CLIPTextModel", + "tokenizer_2": "CLIPTokenizer", + "vae": "AutoencoderKLHunyuanVideo", + "video_processor": "VideoProcessor", + "transformer": "KandinskyS2VTransformer3DModel", + "scheduler": "PiflowScheduler", + }, + "required_inputs": ["image", "audio"], + "inputs": { + "prompt": DEFAULT_PROMPT, + "num_latent_frames": 32, + "fps": 25, + "num_chunks": None, + "height": None, + "width": None, + "generator": None, + "num_inference_steps": 16, + "scale_factor": (1.0, 2.0, 2.0), + "output_type": "pil", + }, + "component_configs": {"video_processor": {"vae_scale_factor": 8}}, + }, + } + + def get_dummy_inputs(self, seed=0): + inputs = super().get_dummy_inputs(seed) + inputs.pop("negative_prompt") + return inputs + + +class TestKandinskyS2VDistilledModularPipelineFast( + KandinskyS2VDistilledModularPipelineTesterConfig, ModularPipelineTesterMixin +): + @pytest.mark.skip(reason=FAST_SKIPS["test_inference_batch_consistent"]) + def test_inference_batch_consistent(self): + pass + + @pytest.mark.skip(reason=FAST_SKIPS["test_inference_batch_single_identical"]) + def test_inference_batch_single_identical(self): + pass + + @pytest.mark.skip(reason=FAST_SKIPS["test_num_images_per_prompt"]) + def test_num_images_per_prompt(self): + pass + + @require_accelerator + def test_float16_inference(self): + check_float16_inference(self) + + +class TestKandinskyS2VDistilledModularPipelineLoading( + KandinskyS2VDistilledModularPipelineTesterConfig, ModularLoadingTesterMixin +): + pass + + +class TestKandinskyS2VDistilledModularPipelineWorkflow( + KandinskyS2VDistilledModularPipelineTesterConfig, ModularWorkflowTesterMixin +): + pass + + +class TestKandinskyS2VDistilledModularPipelineMemory( + KandinskyS2VDistilledModularPipelineTesterConfig, ModularMemoryTesterMixin +): + pass diff --git a/tests/schedulers/test_scheduler_piflow.py b/tests/schedulers/test_scheduler_piflow.py index 53b7d5d2c043..743ad0a668b9 100644 --- a/tests/schedulers/test_scheduler_piflow.py +++ b/tests/schedulers/test_scheduler_piflow.py @@ -17,7 +17,7 @@ import torch from diffusers import PiflowScheduler -from diffusers.schedulers.scheduling_piflow import PiflowSchedulerOutput +from diffusers.schedulers.scheduling_piflow import DXPolicy, PiflowSchedulerOutput, policy_rollout_fm class PiflowSchedulerTest(unittest.TestCase): @@ -196,3 +196,38 @@ def test_step_recovers_known_clean_target(self): sample = scheduler.step(model_output, t, sample).prev_sample torch.testing.assert_close(sample, target, atol=1e-3, rtol=1e-3) + + def test_final_step_policy_spans_the_nominal_segment(self): + """On the last step the policy's grid spans the nominal segment down to 0, while the rollout stops at + `eps`, as in the native samplers the distilled checkpoints were trained with.""" + torch.manual_seed(0) + scheduler = self.scheduler_class(**self.get_default_config()) + scheduler.set_timesteps(4) + sample = torch.randn(2, 3) + model_output = torch.randn(2, 3 * scheduler.n_grid) + + for t in scheduler.timesteps[:-1]: + sample = scheduler.step(model_output, t, sample).prev_sample + out = scheduler.step(model_output, scheduler.timesteps[-1], sample).prev_sample + + raw_src = scheduler._piflow_raw_timesteps[-1] + segment = 1.0 / (4 - 0.5) * 0.5 + sigma = scheduler.sigmas[-2].expand(2).reshape(2, 1) + policy = DXPolicy( + scheduler._to_grid(model_output, sample), + sample, + sigma, + torch.full((2,), segment), + shift=scheduler.config.shift, + eps=scheduler.eps, + ) + self.assertTrue(torch.all(policy.raw_t_dst == 0)) + expected, _, _ = policy_rollout_fm( + sample, + sigma, + raw_src.expand(2), + (raw_src - segment).clamp(min=scheduler.eps).expand(2), + scheduler.num_policy_substeps, + policy, + ) + torch.testing.assert_close(out, expected, atol=0, rtol=0)