From c5be320a5ba8feaa4b76e9d5b04137f69acde209 Mon Sep 17 00:00:00 2001 From: Nikolay Semenov Date: Tue, 13 Oct 2026 02:03:47 +0300 Subject: [PATCH 01/12] add kandinsky-s2v --- docs/source/en/_toctree.yml | 6 + .../api/models/kandinsky_s2v_transformer.md | 22 + docs/source/en/api/pipelines/kandinsky_s2v.md | 106 +++ .../en/api/schedulers/kandinsky_s2v_piflow.md | 21 + src/diffusers/__init__.py | 12 + src/diffusers/models/__init__.py | 2 + src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_kandinsky_s2v.py | 838 ++++++++++++++++++ src/diffusers/modular_pipelines/__init__.py | 12 + .../kandinsky_s2v/__init__.py | 49 + .../kandinsky_s2v/before_denoise.py | 183 ++++ .../kandinsky_s2v/decoders.py | 97 ++ .../kandinsky_s2v/denoise.py | 364 ++++++++ .../kandinsky_s2v/encoders.py | 450 ++++++++++ .../modular_blocks_kandinsky_s2v.py | 99 +++ .../modular_blocks_kandinsky_s2v_distilled.py | 99 +++ .../kandinsky_s2v/modular_pipeline.py | 146 +++ .../modular_pipelines/modular_pipeline.py | 2 + src/diffusers/schedulers/__init__.py | 2 + .../scheduling_kandinsky_s2v_piflow.py | 305 +++++++ src/diffusers/utils/dummy_pt_objects.py | 30 + .../dummy_torch_and_transformers_objects.py | 60 ++ .../test_models_transformer_kandinsky_s2v.py | 143 +++ .../kandinsky_s2v/__init__.py | 0 .../test_modular_pipeline_kandinsky_s2v.py | 203 +++++ 25 files changed, 3252 insertions(+) create mode 100644 docs/source/en/api/models/kandinsky_s2v_transformer.md create mode 100644 docs/source/en/api/pipelines/kandinsky_s2v.md create mode 100644 docs/source/en/api/schedulers/kandinsky_s2v_piflow.md create mode 100644 src/diffusers/models/transformers/transformer_kandinsky_s2v.py create mode 100644 src/diffusers/modular_pipelines/kandinsky_s2v/__init__.py create mode 100644 src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py create mode 100644 src/diffusers/modular_pipelines/kandinsky_s2v/decoders.py create mode 100644 src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py create mode 100644 src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py create mode 100644 src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v.py create mode 100644 src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v_distilled.py create mode 100644 src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py create mode 100644 src/diffusers/schedulers/scheduling_kandinsky_s2v_piflow.py create mode 100644 tests/models/transformers/test_models_transformer_kandinsky_s2v.py create mode 100644 tests/modular_pipelines/kandinsky_s2v/__init__.py create mode 100644 tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index c51c6f3fe2c8..6e264325f2e7 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 @@ -838,6 +842,8 @@ title: HeunDiscreteScheduler - local: api/schedulers/ipndm title: IPNDMScheduler + - local: api/schedulers/kandinsky_s2v_piflow + title: KandinskyS2VPiflowScheduler - local: api/schedulers/dpm_discrete_ancestral title: KDPM2AncestralDiscreteScheduler - local: api/schedulers/dpm_discrete 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..2732214a58cf --- /dev/null +++ b/docs/source/en/api/pipelines/kandinsky_s2v.md @@ -0,0 +1,106 @@ + + +# 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`] | [`KandinskyS2VPiflowScheduler`] | 16 | off | +| [Pro distilled][pro-distill] | 21B | [`KandinskyS2VDistilledAutoBlocks`] | [`KandinskyS2VPiflowScheduler`] | 16 | off | + +[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, and the distilled ones +at 1.0, because guidance is already part of the student — they reject any other. + +## Usage + +```py +import torch +from diffusers import ModularPipeline +from diffusers.utils import export_to_video + +pipe = ModularPipeline.from_pretrained("kandinskylab/Kandinsky-S2V-1.0-Lite-distill-Diffusers") +pipe.load_components(dtype=torch.bfloat16) +pipe.to("cuda") + +video = pipe(image="reference.png", audio="speech.wav", output="videos")[0] +export_to_video(video, "avatar.mp4", fps=25) +``` + +`export_to_video` writes the frames only; add the speech track with any muxer, for +example `ffmpeg -i avatar.mp4 -i speech.wav -c:v copy -c:a aac -shortest out.mp4`. + +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. Keep one model on the GPU at a time, which +brings the peak to about 46 GiB: + +```py +from diffusers import ComponentsManager + +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", offload_strategy=lambda hooks, **_: hooks) +``` + +The default strategy is not enough here: it keeps the transformer beside the speech +encoder because both fit, and the VAE decode that follows needs the room they hold. +Enable the offload after `load_components`, so the strategy applies to every component. + +## Notes + +- 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/kandinsky_s2v_piflow.md b/docs/source/en/api/schedulers/kandinsky_s2v_piflow.md new file mode 100644 index 000000000000..a310bdfcde45 --- /dev/null +++ b/docs/source/en/api/schedulers/kandinsky_s2v_piflow.md @@ -0,0 +1,21 @@ + + +# KandinskyS2VPiflowScheduler + +`KandinskyS2VPiflowScheduler` is the few-step pi-Flow scheduler of the distilled Kandinsky S2V 1.0 +checkpoints. The student predicts a grid of x0 values per segment in one forward +pass, and the scheduler integrates that policy across the segment, so each +segment costs a single model evaluation. + +## KandinskyS2VPiflowScheduler +[[autodoc]] KandinskyS2VPiflowScheduler diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index f2147d993dea..00bcaac2f8a9 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -329,6 +329,7 @@ "Kandinsky6SRTransformer3DModel", "Kandinsky6SRVAE", "Kandinsky6Transformer3DModel", + "KandinskyS2VTransformer3DModel", "Krea2Transformer2DModel", "LatteTransformer3DModel", "LongCatAudioDiTTransformer", @@ -490,6 +491,7 @@ "HeliosScheduler", "HeunDiscreteScheduler", "IPNDMScheduler", + "KandinskyS2VPiflowScheduler", "KarrasVeScheduler", "KDPM2AncestralDiscreteScheduler", "KDPM2DiscreteScheduler", @@ -586,6 +588,10 @@ "HunyuanVideo15ModularPipeline", "Ideogram4AutoBlocks", "Ideogram4ModularPipeline", + "KandinskyS2VAutoBlocks", + "KandinskyS2VDistilledAutoBlocks", + "KandinskyS2VDistilledModularPipeline", + "KandinskyS2VModularPipeline", "Krea2AutoBlocks", "Krea2ModularPipeline", "Krea2TurboAutoBlocks", @@ -1256,6 +1262,7 @@ Kandinsky6SRTransformer3DModel, Kandinsky6SRVAE, Kandinsky6Transformer3DModel, + KandinskyS2VTransformer3DModel, Krea2Transformer2DModel, LatteTransformer3DModel, LongCatAudioDiTTransformer, @@ -1412,6 +1419,7 @@ HeliosScheduler, HeunDiscreteScheduler, IPNDMScheduler, + KandinskyS2VPiflowScheduler, KarrasVeScheduler, KDPM2AncestralDiscreteScheduler, KDPM2DiscreteScheduler, @@ -1491,6 +1499,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..4cf0d3e70b9d --- /dev/null +++ b/src/diffusers/models/transformers/transformer_kandinsky_s2v.py @@ -0,0 +1,838 @@ +# 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 # noqa: N812 +from torch import Tensor, nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin, 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 + + +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 to_rope(args: Tensor) -> Tensor: + 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) + + +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 scale_shift_norm(norm: nn.Module, x: Tensor, scale: Tensor, shift: Tensor, idx: Tensor) -> Tensor: + """AdaLN with the modulation row selected per token.""" + return norm(x.float()) * (scale.float().index_select(0, idx) + 1.0) + shift.float().index_select(0, idx) + + +def gate_sum(x: Tensor, out: Tensor, gate: Tensor, idx: Tensor) -> Tensor: + return x.float() + gate.float().index_select(0, idx) * out.float() + + +def split_lengths(total: int, groups: int, device: torch.device) -> Tensor: + """Split ``total`` items into ``groups`` with boundaries at ``total * i // groups``.""" + bounds = torch.arange(groups + 1, device=device) * total // groups + return torch.diff(bounds) + + +def audio_group_lengths(num_audio_tokens: int, duration: int, device: torch.device) -> Tensor: + """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" + ) + rest = split_lengths(num_audio_tokens - 1, duration - 1, device) + first = torch.ones(1, dtype=rest.dtype, device=device) + return torch.cat([first, rest]) + + +def pad_groups(x: Tensor, lengths: Tensor, max_len: int) -> tuple[Tensor, Tensor]: + """Scatter a packed ``(S, ...)`` tensor into ``(G, max_len, ...)`` plus a validity mask.""" + positions = torch.arange(max_len, device=x.device) + valid = positions[None, :] < lengths[:, None] + out = x.new_zeros(lengths.shape[0], max_len, *x.shape[1:]) + out[valid] = x + return out, valid + + +def grouped_attention( # noqa: PLR0913 + q: Tensor, + k: Tensor, + v: Tensor, + q_lengths: Tensor, + kv_lengths: Tensor, + backend=None, + parallel_config=None, +) -> Tensor: + """Block-diagonal attention: query group ``i`` attends exactly to key group ``i``. + + Groups are uneven, so they are padded into a batch axis and masked rather than reshaped. Inputs and output are + packed ``(S, heads, head_dim)``. + """ + q_pad, q_valid = pad_groups(q, q_lengths, int(q_lengths.max())) + kv_max = int(kv_lengths.max()) + k_pad, kv_valid = pad_groups(k, kv_lengths, kv_max) + v_pad, _ = pad_groups(v, kv_lengths, kv_max) + + out = dispatch_attention_fn( + q_pad, + k_pad, + v_pad, + attn_mask=kv_valid[:, None, None, :], + backend=backend, + parallel_config=parallel_config, + ) + return out[q_valid] + + +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}") + self.model_dim = model_dim + self.max_period = max_period + # 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]`` or ``[B,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: + batched = bool(x.ndim == 5) + offset = 1 if batched else 0 + if not batched and x.ndim != 4: + raise ValueError("visual input must have shape (T,H,W,C) or (B,T,H,W,C)") + + shape = x.shape + duration, height, width, channels = shape[offset:] + p_t, p_h, p_w = self.patch_size + if batched: + x = ( + x.view( + shape[0], + duration // p_t, + p_t, + height // p_h, + p_h, + width // p_w, + p_w, + channels, + ) + .permute(0, 1, 3, 5, 2, 4, 6, 7) + .flatten(4, 7) + ) + else: + 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) + + +# Not copied from transformer_kandinsky6: its Modulation casts the result back to +# the input dtype, while the AdaLN parameters here stay fp32 all the way through +# the per-token selection, as native S2V does. +class KandinskyS2VModulation(nn.Module): + 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) + self.out_layer.weight.data.zero_() + self.out_layer.bias.data.zero_() + + def forward(self, x: Tensor) -> Tensor: + # AdaLN runs in fp32 whatever the weights are stored as, matching native + # generation. Upcasting a bf16 weight is exact, so the result is identical. + return F.linear( + self.activation(x.float()), + self.out_layer.weight.float(), + self.out_layer.bias.float(), + ) + + +# 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__( # noqa: PLR0913 + self, + attn: KandinskyS2VAttention, + hidden_states: Tensor, + encoder_hidden_states: Tensor | None = None, + rotary_emb: Tensor | None = None, + q_lengths: Tensor | None = None, + kv_lengths: 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 q_lengths is not None: + if kv_lengths is None: + raise ValueError("kv_lengths is required when q_lengths is given") + hidden_states = grouped_attention( + query, + key, + value, + q_lengths, + kv_lengths, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + 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 = torch.chunk(self.modulation(time_embed), 2, dim=-1) + visual_embed = scale_shift_norm( + self.norm, visual_embed, scale[:, None, None], shift[:, None, None], time_embed_idx + ).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 Kandinsky6RoPE1D subclasses a base +# RoPE1D defined in that module, so adopting it would mean copying 45 lines to +# replace these 12. +class KandinskyS2VRoPE1D(nn.Module): + def __init__(self, dim: int, max_pos: int = 1024, max_period: float = 10000.0): + super().__init__() + self.dim = dim + self.max_pos = max_pos + self.max_period = max_period + freq = get_freqs(dim // 2, max_period) + pos = torch.arange(max_pos, dtype=freq.dtype) + self.register_buffer("args", torch.outer(pos, freq), persistent=False) + + def forward(self, pos: Tensor) -> Tensor: + return to_rope(self.args[pos]) + + +# Not copied from transformer_kandinsky6: 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.axes_dims = axes_dims + self.max_pos = max_pos + self.max_period = max_period + self.offsets_t = max_pos[0] + + 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.register_buffer(f"args_{i}", torch.outer(pos, freq), persistent=False) + + 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_t = self.args_0[pos[0] + self.offsets_t] / scale_factor[0] + args_h = self.args_1[pos[1]] / scale_factor[1] + args_w = self.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, + ) + return to_rope(args) + + +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 = scale_shift_norm(self.self_attention_norm, x, scale, shift, time_embed_idx).type_as(x) + out = self.self_attention(out, rotary_emb=rope) + x = gate_sum(x, out, gate, time_embed_idx).type_as(x) + + shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) + out = scale_shift_norm(self.feed_forward_norm, x, scale, shift, time_embed_idx).type_as(x) + out = self.feed_forward(out) + return gate_sum(x, out, gate, time_embed_idx).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__( # noqa: PLR0913 + 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( # noqa: PLR0913 + self, + visual_embed: Tensor, + text_embed: Tensor, + audio_embed: Tensor, + time_embed: Tensor, + rope: Tensor, + time_embed_idx: Tensor, + num_visual_tokens: int, + visual_group_lengths: Tensor, + audio_group_lengths: 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 = scale_shift_norm(self.self_attention_norm, visual_embed, scale, shift, time_embed_idx).type_as( + visual_embed + ) + out = self.self_attention(out, rotary_emb=rope) + visual_embed = gate_sum(visual_embed, out, gate, time_embed_idx).type_as(visual_embed) + + shift, scale, gate = torch.chunk(cross_attn_params, 3, dim=-1) + out = scale_shift_norm(self.cross_attention_norm, visual_embed, scale, shift, time_embed_idx).type_as( + visual_embed + ) + out = self.cross_attention(out, encoder_hidden_states=text_embed) + visual_embed = gate_sum(visual_embed, out, gate, time_embed_idx).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, + q_lengths=visual_group_lengths, + kv_lengths=audio_group_lengths, + ) + visual_embed = torch.cat([head, noisy + out], dim=0) + + shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) + out = scale_shift_norm(self.feed_forward_norm, visual_embed, scale, shift, time_embed_idx).type_as( + visual_embed + ) + out = self.feed_forward(out) + return gate_sum(visual_embed, out, gate, time_embed_idx).type_as(visual_embed) + + +class KandinskyS2VTransformer3DModel( + ModelMixin, + ConfigMixin, + PeftAdapterMixin, + FromOriginalModelMixin, + 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 + # the position tables are built 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", + "norm", + "self_attention_norm", + "cross_attention_norm", + "cross_attention_norm_v2a", + "feed_forward_norm", + "query_norm", + "key_norm", + "visual_rope_embeddings", + "text_rope_embeddings", + "audio_rope_embeddings", + ] + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( # noqa: PLR0913 + 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.in_visual_dim = in_visual_dim + self.model_dim = model_dim + self.patch_size = patch_size + self.frame_buckets = frame_buckets + self.enable_framepack = enable_framepack + self.index_rope_ref_image = index_rope_ref_image + self.audio_inject_layers = tuple(audio_inject_layers) + + 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) + for name, coef in ( + ("proj1x", patch_proj1x_coef), + ("proj2x", patch_proj2x_coef), + ("proj4x", patch_proj4x_coef), + ): + scaled = (patch_size[0] * coef[0], patch_size[1] * coef[1], patch_size[2] * coef[2]) + setattr( + self, + f"visual_embeddings_{name}", + KandinskyS2VVisualEmbeddings(in_visual_dim, model_dim, scaled), + ) + # 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 self.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 _embed_memory( + self, + memory_hidden_states: tuple[Tensor, Tensor, Tensor], + memory_rope_pos: tuple[list[Tensor], list[Tensor], list[Tensor]], + scale_factor: tuple[float, float, float], + ) -> tuple[Tensor, Tensor]: + """Embed the FramePack memory at its three temporal scales.""" + memory_4x, memory_2x, memory_1x = memory_hidden_states + rope_4x, rope_2x, rope_1x = memory_rope_pos + + parts = [(self.visual_embeddings_proj1x(memory_1x), rope_1x)] + if self.enable_framepack: + parts = [ + (self.visual_embeddings_proj4x(memory_4x), rope_4x), + (self.visual_embeddings_proj2x(memory_2x), rope_2x), + *parts, + ] + + embeds, ropes = [], [] + for embed, rope_pos in parts: + rope = self.visual_rope_embeddings(embed.shape[:-1], rope_pos, scale_factor) + embeds.append(embed.flatten(0, 2)) + ropes.append(rope.flatten(0, 2)) + return torch.cat(embeds, dim=0), torch.cat(ropes, dim=0) + + def forward( # noqa: PLR0913, PLR0915 + 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), + duration: int | None = None, + return_dict: bool = True, + ) -> Transformer2DModelOutput | 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. + duration: number of latent frames used to align speech to frames; + defaults to the temporal size of ``hidden_states``. + 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] if duration is None else duration + + 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) + + memory_embed, memory_rope = self._embed_memory(memory, memory_rope_pos, scale_factor) + + 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), + ] + ) + + visual_groups = split_lengths(num_visual_tokens, duration, device) + audio_groups = audio_group_lengths(audio_embed.shape[0], duration, device) + + for block in self.visual_transformer_blocks: + args = ( + visual_embed_full, + text_embed, + audio_embed, + time_embed_full, + visual_rope_full, + time_embed_idx_full, + num_visual_tokens, + visual_groups, + audio_groups, + ) + if torch.is_grad_enabled() and self.gradient_checkpointing: + visual_embed_full = self._gradient_checkpointing_func(block, *args) + else: + visual_embed_full = block(*args) + + 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..ab246093af28 --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py @@ -0,0 +1,183 @@ +# 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), + InputParam("audio_embeds", required=True, type_hint=torch.Tensor), + InputParam("num_chunks", required=True, type_hint=int), + InputParam("num_latent_frames", default=32, type_hint=int), + 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), + ] + + @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.template("negative_prompt_embeds"), + OutputParam("pooled_prompt_embeds", type_hint=torch.Tensor), + OutputParam("negative_pooled_prompt_embeds", type_hint=torch.Tensor), + OutputParam("text_rope_pos", type_hint=torch.Tensor), + OutputParam("negative_text_rope_pos", type_hint=torch.Tensor), + OutputParam("visual_rope_pos", type_hint=list), + OutputParam("reference_rope_pos", type_hint=list), + OutputParam("memory", type_hint=list, description="FramePack memory in [1x, 2x, 4x] order."), + ] + + @torch.no_grad() + def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + config = components.transformer.config + 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.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=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 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..991d8198a5f3 --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/decoders.py @@ -0,0 +1,97 @@ +# 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), + InputParam("track_seconds", required=True, type_hint=float), + InputParam("fps", default=25, type_hint=int), + InputParam.template( + "output_type", + description="Output format: 'pil', 'np', 'pt' or 'latent'. Latents are returned whole, untrimmed.", + ), + ] + + @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) + + if block_state.output_type == "latent": + block_state.videos = torch.cat(block_state.latent_chunks, dim=0).unsqueeze(0) + else: + 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..b3a77170a0ad --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py @@ -0,0 +1,364 @@ +# 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 +from ...schedulers.scheduling_kandinsky_s2v_piflow import KandinskyS2VPiflowScheduler +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]: + device = components._execution_device + if getattr(components.scheduler, "is_piflow", False): + components.scheduler.set_timesteps(block_state.num_inference_steps, device=device) + else: + # 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=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", KandinskyS2VPiflowScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [InputParam.template("num_inference_steps", default=16)] + + +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]: + if getattr(components.scheduler, "is_piflow", False) and components.requires_unconditional_embeds: + # Guidance on top of a distilled student over-guides it and doubles the cost. + raise ValueError( + f"guidance must be off with {type(components.scheduler).__name__}: CFG is distilled into the student" + ) + + config = components.transformer.config + transformer_dtype = components.transformer.dtype + device = components._execution_device + latents = block_state.latents + chunk_audio = block_state.chunk_audio + num_latent_frames = block_state.num_latent_frames + memory = block_state.memory + + # 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 = { + "reference_hidden_states": block_state.reference_hidden_states.unsqueeze(0), + "memory_hidden_states": tuple(m.unsqueeze(0).to(transformer_dtype) for m in reversed(memory)), + "visual_rope_pos": block_state.visual_rope_pos, + "reference_rope_pos": block_state.reference_rope_pos, + "memory_rope_pos": tuple(reversed(memory_rope_pos(config, memory, device))), + "audio_rope_pos": audio_rope_pos(chunk_audio.shape[0], num_latent_frames, device), + "scale_factor": block_state.scale_factor, + "duration": num_latent_frames, + "return_dict": False, + } + + 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 expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", KandinskyS2VTransformer3DModel), + ComponentSpec("scheduler", KandinskyS2VPiflowScheduler), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 1.0}), + default_creation_method="from_config", + ), + ] + + @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 + ] + + +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.")] + + @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..ef17b9c560f1 --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py @@ -0,0 +1,450 @@ +# 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 pathlib import Path +from typing import Any + +import PIL.Image +import torch +from transformers import ( + CLIPTextModel, + CLIPTokenizer, + Qwen2_5OmniProcessor, + Qwen2_5OmniThinkerForConditionalGeneration, +) + +from ...configuration_utils import FrozenDict +from ...guiders import ClassifierFreeGuidance +from ...models import AutoencoderKLHunyuanVideo +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 ( + AUDIO_SAMPLE_RATE, + AUDIO_TOKEN_TOLERANCE, + AUDIO_TOKENS_PER_SECOND, + DEFAULT_NEGATIVE_PROMPT, + DEFAULT_PROMPT, + SYSTEM_PROMPT, + TEMPORAL_COMPRESSION_RATIO, + KandinskyS2VModularPipeline, + audio_duration, + pad_with_silence, + select_resolution, +) + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def encode_qwen_omni( + components: KandinskyS2VModularPipeline, + prompt: str, + audio_path: str | Path | 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.""" + user_content: list[dict[str, Any]] = [{"type": "text", "text": prompt}] + if audio_path is not None: + item: dict[str, Any] = {"type": "audio", "audio": str(audio_path)} + if audio_end is not None: + item |= {"audio_start": 0.0, "audio_end": audio_end} + user_content.append(item) + conversation = [ + {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}, + {"role": "user", "content": user_content}, + ] + + chat_text = components.tokenizer.apply_chat_template([conversation], add_generation_prompt=False, tokenize=False) + audios = None + if audio_path is not None: + import librosa # noqa: PLC0415 + + audios = [librosa.load(str(audio_path), sr=AUDIO_SAMPLE_RATE, offset=0.0, duration=audio_end)[0]] + + inputs = components.tokenizer( + text=chat_text, + audio=audios, + return_tensors="pt", + padding=True, + use_audio_in_video=False, + ) + inputs = inputs.to(device).to(components.text_encoder.dtype) + + # The speech features are the audio tower's output, which the Thinker consumes internally; a hook on the tower + # reads them out of the same forward that produces the text hidden states. + captured: dict[str, torch.Tensor] = {} + + def hook(_module, _inputs, output): + captured["features"] = getattr(output, "last_hidden_state", output) + + handle = components.text_encoder.audio_tower.register_forward_hook(hook) + try: + output = components.text_encoder(**inputs, output_hidden_states=True, return_dict=True) + finally: + handle.remove() + + prompt_embeds = slice_user_text(components, output.hidden_states[-1], inputs) + return prompt_embeds, captured.get("features") + + +def slice_user_text(components: KandinskyS2VModularPipeline, 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 = components.tokenizer.tokenizer.convert_tokens_to_ids("<|im_start|>") + audio_start_token_id = components.text_encoder.config.audio_start_token_id + user_role_len = len(components.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( + components: KandinskyS2VModularPipeline, prompt: str, device: torch.device, max_length: int = 77 +) -> torch.Tensor: + inputs = components.tokenizer_2( + [prompt], + max_length=max_length, + truncation=True, + add_special_tokens=True, + padding="max_length", + return_tensors="pt", + ).to(device) + return components.text_encoder_2(**inputs)["pooler_output"] + + +class KandinskyS2VInputStep(ModularPipelineBlocks): + model_name = "kandinsky-s2v" + + @property + def description(self) -> str: + return ( + "Input step that checks the reference frame and the speech track, picks the resolution bucket and plans " + "the chunks: how many are needed to cover the track, and how many seconds of speech they take." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("tokenizer", Qwen2_5OmniProcessor), + ComponentSpec("transformer", KandinskyS2VTransformer3DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template( + "image", + type_hint=PIL.Image.Image | str | Path, + description="The reference frame to animate: a PIL image or a path to one.", + ), + InputParam( + "audio", + required=True, + type_hint=str | Path, + description="The speech track that drives the generation: a path to an audio file.", + ), + 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( + "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("reference_image", type_hint=PIL.Image.Image, description="The reference frame, in RGB."), + OutputParam("height", type_hint=int, description="The height in pixels of the generated video."), + OutputParam("width", type_hint=int, description="The width in pixels of the generated video."), + 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."), + OutputParam( + "target_seconds", type_hint=float, description="Seconds of speech the chunks cover, padded if needed." + ), + ] + + @staticmethod + def check_inputs(components, block_state): + if block_state.image is None: + raise ValueError("`image` is required: it is the reference frame to animate") + if block_state.audio is None: + raise ValueError("`audio` is required: it is the speech track that drives the generation") + 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.num_chunks is not None and block_state.num_chunks < 1: + raise ValueError(f"`num_chunks` must be >= 1, got {block_state.num_chunks}") + needed = sum(components.transformer.config.frame_buckets) + if block_state.num_latent_frames < needed: + raise ValueError( + f"`num_latent_frames` must be >= {needed}, the FramePack memory window " + f"sum(frame_buckets), got {block_state.num_latent_frames}" + ) + + @torch.no_grad() + def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(components, block_state) + + image = block_state.image + if not isinstance(image, PIL.Image.Image): + image = PIL.Image.open(image) + block_state.reference_image = image.convert("RGB") + if block_state.height is None: + block_state.height, block_state.width = select_resolution(block_state.reference_image) + + chunk_seconds = ((block_state.num_latent_frames - 1) * TEMPORAL_COMPRESSION_RATIO + 1) / block_state.fps + block_state.track_seconds = audio_duration(block_state.audio) + needed = max(1, math.ceil(block_state.track_seconds / chunk_seconds)) + if block_state.num_chunks is not None and block_state.num_chunks > needed: + # The video is trimmed to the track, so these are computed and then dropped. + logger.warning( + f"{block_state.num_chunks - needed} chunk(s) past the end of the speech track are generated and " + f"then trimmed; pass num_chunks={needed} to save the time" + ) + if block_state.num_chunks is None: + block_state.num_chunks = needed + block_state.target_seconds = block_state.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 = components.tokenizer.feature_extractor + max_seconds = extractor.n_samples / extractor.sampling_rate + if block_state.target_seconds > max_seconds + 0.01: + raise ValueError( + f"{block_state.num_chunks} chunks need {block_state.target_seconds:g} s of speech, but the speech " + f"encoder truncates at {max_seconds:g} s; trim the track or pass fewer `num_chunks`" + ) + + self.set_block_state(state, block_state) + return components, state + + +class KandinskyS2VTextAudioEncoderStep(ModularPipelineBlocks): + model_name = "kandinsky-s2v" + + @property + def description(self) -> str: + return ( + "Encoder step that turns the prompt and the speech track 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 only when the 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), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 3.0}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template( + "prompt", + required=False, + description="The prompt to guide the generation. Defaults to the reference prompt, `DEFAULT_PROMPT`.", + ), + InputParam.template( + "negative_prompt", + description="The prompt for the unconditional pass of the guidance. Defaults to " + "`DEFAULT_NEGATIVE_PROMPT`.", + ), + InputParam("audio", required=True, type_hint=str | Path, description="The speech track."), + InputParam("target_seconds", required=True, type_hint=float, description="Seconds of speech to encode."), + ] + + @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." + ), + ] + + @staticmethod + def check_inputs(components, block_state): + if block_state.prompt is None: + block_state.prompt = DEFAULT_PROMPT + if block_state.negative_prompt is None: + block_state.negative_prompt = DEFAULT_NEGATIVE_PROMPT + if not isinstance(block_state.prompt, str) or not isinstance(block_state.negative_prompt, str): + raise ValueError("`prompt` and `negative_prompt` must be strings; 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") + if components.requires_unconditional_embeds and not block_state.negative_prompt.strip(): + raise ValueError("`negative_prompt` must not be empty when classifier-free guidance is enabled") + + @torch.no_grad() + def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(components, block_state) + device = components._execution_device + target_seconds = block_state.target_seconds + + padded_audio = None + try: + padded_audio = pad_with_silence(block_state.audio, target_seconds) + audio_for_encode = padded_audio if padded_audio is not None else Path(block_state.audio) + prompt_embeds, audio_embeds = encode_qwen_omni( + components, block_state.prompt, audio_for_encode, target_seconds, device + ) + finally: + if padded_audio is not None: + padded_audio.unlink(missing_ok=True) + if audio_embeds is None: + raise ValueError("speech-to-video requires an audio track") + + # 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" + ) + + block_state.prompt_embeds = prompt_embeds + block_state.audio_embeds = audio_embeds + block_state.pooled_prompt_embeds = encode_clip(components, block_state.prompt, device) + block_state.negative_prompt_embeds = None + block_state.negative_pooled_prompt_embeds = None + if components.requires_unconditional_embeds: + block_state.negative_prompt_embeds, _ = encode_qwen_omni( + components, block_state.negative_prompt, None, None, device + ) + block_state.negative_pooled_prompt_embeds = encode_clip(components, block_state.negative_prompt, device) + + self.set_block_state(state, block_state) + return components, state + + +class KandinskyS2VDistilledTextAudioEncoderStep(KandinskyS2VTextAudioEncoderStep): + model_name = "kandinsky-s2v-distilled" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + *super().expected_components[:-1], + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 1.0}), + default_creation_method="from_config", + ), + ] + + +class KandinskyS2VVaeEncoderStep(ModularPipelineBlocks): + model_name = "kandinsky-s2v" + + @property + def description(self) -> str: + return "VAE encoder step that encodes the reference frame into a scaled latent, resized and center-cropped." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("vae", AutoencoderKLHunyuanVideo)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "reference_image", required=True, type_hint=PIL.Image.Image, description="The reference frame." + ), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + 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).", + ) + ] + + @torch.no_grad() + def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: + from torchvision import transforms # noqa: PLC0415 + + block_state = self.get_block_state(state) + device = components._execution_device + image, height, width = block_state.reference_image, block_state.height, block_state.width + + scale = max(width / image.width, height / image.height) + resized = transforms.functional.resize( + image, + (round(image.height * scale), round(image.width * scale)), + interpolation=transforms.InterpolationMode.BILINEAR, + ) + cropped = transforms.functional.center_crop(resized, (height, width)) + + tensor = transforms.ToTensor()(cropped).unsqueeze(0) + tensor = tensor.to(device=device, dtype=components.vae.dtype) * 2.0 - 1.0 + tensor = tensor.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..e32a8f86d8b8 --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v.py @@ -0,0 +1,99 @@ +# 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 KandinskyS2VChunkDenoiseStep +from .encoders import KandinskyS2VInputStep, KandinskyS2VTextAudioEncoderStep, KandinskyS2VVaeEncoderStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +AUTO_BLOCKS = InsertableDict( + [ + ("input", KandinskyS2VInputStep()), + ("text_encoder", KandinskyS2VTextAudioEncoderStep()), + ("vae_encoder", KandinskyS2VVaeEncoderStep()), + ("prepare", KandinskyS2VPrepareStep()), + ("denoise", KandinskyS2VChunkDenoiseStep()), + ("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: + tokenizer (`Qwen2_5OmniProcessor`) transformer (`KandinskyS2VTransformer3DModel`) text_encoder + (`Qwen2_5OmniThinkerForConditionalGeneration`) text_encoder_2 (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) + guider (`ClassifierFreeGuidance`) vae (`AutoencoderKLHunyuanVideo`) scheduler + (`FlowMatchEulerDiscreteScheduler`) video_processor (`VideoProcessor`) + + Inputs: + image (`Image | str | Path`): + The reference frame to animate: a PIL image or a path to one. + audio (`str | Path`): + The speech track that drives the generation: a path to an audio file. + 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`. + 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. + prompt (`str`, *optional*): + The prompt to guide the generation. Defaults to the reference prompt, `DEFAULT_PROMPT`. + negative_prompt (`str`, *optional*): + The prompt for the unconditional pass of the guidance. Defaults to `DEFAULT_NEGATIVE_PROMPT`. + 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' or 'latent'. Latents are returned whole, untrimmed. + + 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..9e6ac2e92b41 --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v_distilled.py @@ -0,0 +1,99 @@ +# 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, KandinskyS2VInputStep, KandinskyS2VVaeEncoderStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +DISTILLED_AUTO_BLOCKS = InsertableDict( + [ + ("input", KandinskyS2VInputStep()), + ("text_encoder", KandinskyS2VDistilledTextAudioEncoderStep()), + ("vae_encoder", KandinskyS2VVaeEncoderStep()), + ("prepare", KandinskyS2VPrepareStep()), + ("denoise", KandinskyS2VDistilledChunkDenoiseStep()), + ("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: + tokenizer (`Qwen2_5OmniProcessor`) transformer (`KandinskyS2VTransformer3DModel`) text_encoder + (`Qwen2_5OmniThinkerForConditionalGeneration`) text_encoder_2 (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) + guider (`ClassifierFreeGuidance`) vae (`AutoencoderKLHunyuanVideo`) scheduler (`KandinskyS2VPiflowScheduler`) + video_processor (`VideoProcessor`) + + Inputs: + image (`Image | str | Path`): + The reference frame to animate: a PIL image or a path to one. + audio (`str | Path`): + The speech track that drives the generation: a path to an audio file. + 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`. + 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. + prompt (`str`, *optional*): + The prompt to guide the generation. Defaults to the reference prompt, `DEFAULT_PROMPT`. + negative_prompt (`str`, *optional*): + The prompt for the unconditional pass of the guidance. Defaults to `DEFAULT_NEGATIVE_PROMPT`. + 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' or 'latent'. Latents are returned whole, untrimmed. + + 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..32bdf04df4eb --- /dev/null +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py @@ -0,0 +1,146 @@ +# 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 os +import tempfile +from pathlib import Path + +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, + allowed: tuple[tuple[int, int], ...] = ((512, 512), (768, 512), (512, 768)), +) -> tuple[int, int]: + """Pick the ``(height, width)`` bucket closest in aspect ratio to the image.""" + aspect = image.width / image.height + return min(allowed, key=lambda res: abs(aspect - res[1] / res[0])) + + +def audio_duration(path: str | Path) -> float: + import soundfile as sf # noqa: PLC0415 + + info = sf.info(str(path)) + return info.frames / info.samplerate + + +def pad_with_silence(path: str | Path, target_duration: float) -> Path | None: + """Append silence so the track covers ``target_duration``; ``None`` if it already does.""" + import soundfile as sf # noqa: PLC0415 + + original = audio_duration(path) + if target_duration <= original + 0.01: + return None + + data, sample_rate = sf.read(str(path)) + pad_samples = int((target_duration - original) * sample_rate) + shape = (pad_samples,) if data.ndim == 1 else (pad_samples, data.shape[1]) + padded = np.concatenate([data, np.zeros(shape, dtype=data.dtype)]) + + handle, temp_path = tempfile.mkstemp(suffix=".wav") + os.close(handle) + try: + sf.write(temp_path, padded, sample_rate) + except BaseException: + os.unlink(temp_path) + raise + return Path(temp_path) + + +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 = audio_group_lengths(num_audio_tokens, duration, 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 the guider defaults to `guidance_scale=1.0` 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/__init__.py b/src/diffusers/schedulers/__init__.py index 3909159729d6..a48ecd802be2 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -68,6 +68,7 @@ _import_structure["scheduling_helios_dmd"] = ["HeliosDMDScheduler"] _import_structure["scheduling_heun_discrete"] = ["HeunDiscreteScheduler"] _import_structure["scheduling_ipndm"] = ["IPNDMScheduler"] + _import_structure["scheduling_kandinsky_s2v_piflow"] = ["KandinskyS2VPiflowScheduler"] _import_structure["scheduling_k_dpm_2_ancestral_discrete"] = ["KDPM2AncestralDiscreteScheduler"] _import_structure["scheduling_k_dpm_2_discrete"] = ["KDPM2DiscreteScheduler"] _import_structure["scheduling_lcm"] = ["LCMScheduler"] @@ -156,6 +157,7 @@ from .scheduling_ipndm import IPNDMScheduler from .scheduling_k_dpm_2_ancestral_discrete import KDPM2AncestralDiscreteScheduler from .scheduling_k_dpm_2_discrete import KDPM2DiscreteScheduler + from .scheduling_kandinsky_s2v_piflow import KandinskyS2VPiflowScheduler from .scheduling_lcm import LCMScheduler from .scheduling_ltx_euler_ancestral_rf import LTXEulerAncestralRFScheduler from .scheduling_minimax_h3 import MiniMaxH3Scheduler diff --git a/src/diffusers/schedulers/scheduling_kandinsky_s2v_piflow.py b/src/diffusers/schedulers/scheduling_kandinsky_s2v_piflow.py new file mode 100644 index 000000000000..b45f2e11d295 --- /dev/null +++ b/src/diffusers/schedulers/scheduling_kandinsky_s2v_piflow.py @@ -0,0 +1,305 @@ +# 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. +"""pi-Flow scheduler for distilled Kandinsky S2V 1.0 checkpoints. + +pi-Flow (arXiv 2510.14974) replaces the per-step velocity with a *policy*: the student predicts a grid of x0 values +spanning the segment ``[t_dst, t_src]`` in a single forward, and the ODE is then integrated across the segment without +further network evaluations. One segment therefore costs exactly one model evaluation, and 16 segments replace the 50 +Euler steps of the base model. + +The distilled checkpoints widen the output head ``n_grid`` times; this scheduler splits that axis, so the transformer +itself is unchanged and its config only differs in ``out_visual_dim``. +""" + +from __future__ import annotations + +import torch +from torch import Tensor + +from ..configuration_utils import register_to_config +from .scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, + FlowMatchEulerDiscreteSchedulerOutput, +) + + +_MIN_GRID_POINTS = 2 + + +def shift_timesteps(t: Tensor, shift: float) -> Tensor: + """Flow-matching timestep warp.""" + return shift * t / (1 + (shift - 1) * t) + + +class DXPolicy: + """Dynamic-x0 policy: a grid of x0 predictions, linearly interpolated in raw time.""" + + def __init__( # noqa: PLR0913 + self, + denoising_output: Tensor, + x_t_src: Tensor, + sigma_t_src: Tensor, + segment_size: Tensor, + shift: float = 1.0, + eps: float = 1e-4, + ): + self.x_t_src = x_t_src + self.ndim = x_t_src.dim() + self.shift = shift + self.eps = eps + + self.sigma_t_src = sigma_t_src.reshape(*sigma_t_src.size(), *((self.ndim - sigma_t_src.dim()) * [1])) + self.raw_t_src = self._unwarp_t(self.sigma_t_src) + seg = segment_size + if isinstance(seg, Tensor) and seg.dim() < self.raw_t_src.dim(): + seg = seg.reshape(*seg.size(), *((self.ndim - seg.dim()) * [1])) + # clamp(min=0) here while the sampler clamps to eps — that asymmetry comes + # from the distillation fork (piflow_s2v.py vs utils_s2v_piflow.py) and the + # students were trained against it. Do not align the two. + self.raw_t_dst = (self.raw_t_src - seg).clamp(min=0) + self.segment_size = (self.raw_t_src - self.raw_t_dst).clamp(min=eps) + + self.denoising_output_x_0 = self._u_to_x_0(denoising_output, self.x_t_src, self.sigma_t_src) + + @staticmethod + def _interpolate(x: Tensor, t: Tensor) -> Tensor: + n = x.size(1) + if n < _MIN_GRID_POINTS: + return x.squeeze(1) + t = t.clamp(min=0, max=1) * (n - 1) + t0 = t.floor().to(torch.long).clamp(min=0, max=n - 2) + t1 = t0 + 1 + t0t1 = torch.stack([t0, t1], dim=1) + x0x1 = torch.gather(x, dim=1, index=t0t1.expand(-1, -1, *x.shape[2:])) + return (t1 - t) * x0x1[:, 0] + (t - t0) * x0x1[:, 1] + + def _unwarp_t(self, sigma_t: Tensor) -> Tensor: + return sigma_t / (self.shift + (1 - self.shift) * sigma_t) + + @staticmethod + def _u_to_x_0(denoising_output: Tensor, x_t: Tensor, sigma_t: Tensor) -> Tensor: + return x_t.unsqueeze(1) - sigma_t.unsqueeze(1) * denoising_output + + def pi(self, x_t: Tensor, sigma_t: Tensor) -> Tensor: + sigma_t = sigma_t.reshape(*sigma_t.size(), *((self.ndim - sigma_t.dim()) * [1])) + raw_t = self._unwarp_t(sigma_t) + x_0 = self._interpolate(self.denoising_output_x_0, (raw_t - self.raw_t_dst) / self.segment_size) + return (x_t - x_0) / sigma_t.clamp(min=self.eps) + + +def policy_rollout_fm( # noqa: PLR0913 + x_t_start: Tensor, + sigma_t_start: Tensor, + raw_t_start: Tensor, + raw_t_end: Tensor, + total_substeps: int, + policy: DXPolicy, +) -> Tensor: + """Integrate the policy ODE across a segment with a fixed substep size.""" + num_tokens = x_t_start.size(0) + ndim = x_t_start.dim() + raw_t_start = raw_t_start.reshape(num_tokens, *((ndim - 1) * [1])) + raw_t_end = raw_t_end.reshape(num_tokens, *((ndim - 1) * [1])) + + delta_raw_t = raw_t_start - raw_t_end + num_substeps = (delta_raw_t * total_substeps).round().to(torch.long).clamp(min=1) + substep_size = delta_raw_t / num_substeps + + raw_t, sigma_t, x_t = raw_t_start, sigma_t_start, x_t_start + for substep_id in range(int(num_substeps.max().item())): + u = policy.pi(x_t, sigma_t) + + raw_t_minus = (raw_t - substep_size).clamp(min=0) + sigma_t_minus = shift_timesteps(raw_t_minus, policy.shift) + x_t_minus = x_t + u * (sigma_t_minus - sigma_t) + + active = num_substeps > substep_id + x_t = torch.where(active, x_t_minus, x_t) + sigma_t = torch.where(active, sigma_t_minus, sigma_t) + raw_t = torch.where(active, raw_t_minus, raw_t) + return x_t + + +class KandinskyS2VPiflowScheduler(FlowMatchEulerDiscreteScheduler): + """Few-step pi-Flow scheduler for distilled Kandinsky S2V 1.0 checkpoints. + + Args: + num_train_timesteps (`int`, *optional*, defaults to 1000): + Number of training diffusion steps. + shift (`float`, *optional*, defaults to 10.0): + Flow-matching timestep shift. + n_grid (`int`, *optional*, defaults to 5): + Number of x0 predictions packed into the model's output channels. + eps (`float`, *optional*, defaults to 1e-6): + Minimum raw timestep and policy denominator. + final_step_size_scale (`float`, *optional*, defaults to 0.5): + Size of the last segment relative to the others. + num_policy_substeps (`int`, *optional*, defaults to 128): + Substeps per unit of raw time used to integrate the policy. + """ + + is_piflow = True + + @register_to_config + def __init__( # noqa: PLR0913 + self, + num_train_timesteps: int = 1000, + shift: float = 10.0, + n_grid: int = 5, + eps: float = 1e-6, + final_step_size_scale: float = 0.5, + num_policy_substeps: int = 128, + ): + if n_grid < _MIN_GRID_POINTS: + raise ValueError(f"n_grid must be >= {_MIN_GRID_POINTS}, got {n_grid}") + if eps <= 0: + raise ValueError(f"eps must be > 0, got {eps}") + if not 0 < final_step_size_scale <= 1: + raise ValueError(f"final_step_size_scale must be in (0, 1], got {final_step_size_scale}") + if num_policy_substeps < 1: + raise ValueError(f"num_policy_substeps must be >= 1, got {num_policy_substeps}") + super().__init__(num_train_timesteps=num_train_timesteps, shift=shift) + self._raw_timesteps = torch.empty(0) + self._segment_sizes = torch.empty(0) + + def set_timesteps( # noqa: PLR0913 + self, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + sigmas: list[float] | None = None, + mu: float | None = None, + timesteps: list[float] | None = None, + ) -> None: + """Build the distilled segment schedule. + + Every segment has the same size except the last one, which is scaled by ``final_step_size_scale``; together + they cover raw time from 1 down to 0. + """ + if sigmas is not None or mu is not None or timesteps is not None: + raise ValueError("this scheduler only supports its own distilled schedule") + if num_inference_steps is None or num_inference_steps < 1: + raise ValueError(f"num_inference_steps must be positive, got {num_inference_steps}") + + config = self.config + base_segment = 1.0 / (num_inference_steps - (1.0 - config.final_step_size_scale)) + + # Walked segment by segment rather than as `1 - cumsum(segments)`: the + # native sampler subtracts one segment at a time, and in float32 the two + # orders differ in the last bits, which the policy rollout then amplifies. + raws, segments = [], [] + raw_t = torch.ones(1, dtype=torch.float32, device=device) + for step in range(num_inference_steps): + is_final = step == num_inference_steps - 1 + segment = base_segment * (config.final_step_size_scale if is_final else 1.0) + raws.append(raw_t) + segments.append(segment) + raw_t = (raw_t - segment).clamp(min=config.eps) + + raw = torch.cat(raws) + segments = torch.tensor(segments, dtype=torch.float32, device=device) + raw_sigmas = shift_timesteps(raw, float(config.shift)) + + self.num_inference_steps = int(num_inference_steps) + self._raw_timesteps = raw + self._segment_sizes = segments + self.timesteps = raw_sigmas * config.num_train_timesteps + self.sigmas = torch.cat([raw_sigmas, raw_sigmas.new_zeros(1)]) + self._step_index = None + self._begin_index = None + + def _to_grid(self, model_output: Tensor, sample: Tensor) -> Tensor: + """``(T, H, W, n_grid * C)`` -> ``(T, n_grid, H, W, C)``. + + Grid-major, exactly as the distilled head emits it. Do not reorder: a channel-major split would silently + shuffle the policy grid while every shape still checks out. + """ + n_grid = self.config.n_grid + if model_output.shape[:-1] != sample.shape[:-1]: + raise ValueError(f"model output {tuple(model_output.shape)} does not match sample {tuple(sample.shape)}") + if model_output.shape[-1] != sample.shape[-1] * n_grid: + raise ValueError( + f"expected {sample.shape[-1] * n_grid} output channels for n_grid={n_grid}, " + f"got {model_output.shape[-1]}" + ) + return model_output.reshape(*model_output.shape[:-1], n_grid, sample.shape[-1]).movedim(-2, 1) + + def _resolve_step_index(self, timestep: Tensor | float) -> int: + if self.step_index is None: + if self.begin_index is not None: + self._step_index = self.begin_index + else: + schedule = self.timesteps + value = torch.as_tensor(timestep, device=schedule.device, dtype=schedule.dtype) + indices = torch.nonzero(schedule == value).flatten() + if not indices.numel(): + raise ValueError(f"timestep {float(value)} is not in the pi-Flow schedule") + self._step_index = int(indices[-1].item()) + if self.step_index >= self.num_inference_steps: + raise ValueError("step() called after the schedule was exhausted") + return int(self.step_index) + + def step( + self, + model_output: Tensor, + timestep: float | Tensor, + sample: Tensor, + return_dict: bool = True, + ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple: + """Advance one segment by integrating the policy. + + Args: + model_output: the widened prediction, ``n_grid`` x0 values per channel. + timestep: a value from ``scheduler.timesteps``. + sample: the current noisy latent. + return_dict: return a [`FlowMatchEulerDiscreteSchedulerOutput`]. + """ + if isinstance(timestep, int) or (isinstance(timestep, Tensor) and not timestep.is_floating_point()): + raise ValueError("pass a value from scheduler.timesteps, not an index") + + step_index = self._resolve_step_index(timestep) + config = self.config + # The policy rollout runs in float32, as the native sampler does. + sample = sample.to(torch.float32) + grid = self._to_grid(model_output.to(torch.float32), sample) + + device = sample.device + num_tokens = sample.shape[0] + extra_dims = (sample.dim() - 1) * [1] + + raw_src = self._raw_timesteps[step_index].to(device) + segment = self._segment_sizes[step_index].to(device) + raw_dst = (raw_src - segment).clamp(min=config.eps) + + sigma_src = self.sigmas[step_index].to(device).repeat_interleave(num_tokens).reshape(num_tokens, *extra_dims) + policy = DXPolicy( + denoising_output=grid, + x_t_src=sample, + sigma_t_src=sigma_src, + segment_size=torch.full_like(sigma_src, float(segment)), + shift=float(config.shift), + eps=config.eps, + ) + updated = policy_rollout_fm( + x_t_start=sample, + sigma_t_start=sigma_src, + raw_t_start=raw_src.repeat_interleave(num_tokens), + raw_t_end=raw_dst.repeat_interleave(num_tokens), + total_substeps=config.num_policy_substeps, + policy=policy, + ) + + self._step_index += 1 + if return_dict: + return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=updated) + return (updated,) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index b13d46f68ff5..35ad9d8844b2 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"] @@ -3921,6 +3936,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class KandinskyS2VPiflowScheduler(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 KarrasVeScheduler(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..caf31091799a --- /dev/null +++ b/tests/models/transformers/test_models_transformer_kandinsky_s2v.py @@ -0,0 +1,143 @@ +# 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 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 BaseModelTesterConfig, ModelTesterMixin + + +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) -> 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) + ), + "duration": DURATION, + } + + +class TestKandinskyS2VTransformer(KandinskyS2VTransformerTesterConfig, ModelTesterMixin): + pass 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..ad3ed0b97e7d --- /dev/null +++ b/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py @@ -0,0 +1,203 @@ +# 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 functools +import math +import struct +import tempfile +import wave +from pathlib import Path + +import pytest +import torch +from PIL import Image + +from diffusers import ClassifierFreeGuidance, ModularPipeline +from diffusers.modular_pipelines import ( + KandinskyS2VAutoBlocks, + KandinskyS2VDistilledAutoBlocks, + KandinskyS2VDistilledModularPipeline, + KandinskyS2VModularPipeline, +) + +from ..testing_utils import ( + BaseModularPipelineTesterConfig, + ModularLoadingTesterMixin, + ModularMemoryTesterMixin, + ModularPipelineTesterMixin, +) + + +# One tiny repository per variant, laid out like the real checkpoints. +TINY_REPO = "kandinskylab/tiny-kandinsky-s2v-modular-pipe" +TINY_DISTILLED_REPO = "kandinskylab/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 + + +@functools.lru_cache(maxsize=1) +def dummy_audio() -> str: + """A second of a plain tone; the speech is read from a file, not a tensor.""" + path = Path(tempfile.mkdtemp()) / "speech.wav" + with wave.open(str(path), "w") as handle: + handle.setnchannels(1) + handle.setsampwidth(2) + handle.setframerate(AUDIO_SAMPLE_RATE) + handle.writeframes(b"".join(struct.pack(" Date: Tue, 27 Oct 2026 02:58:07 +0300 Subject: [PATCH 02/12] Address review: reuse PiflowScheduler, keep RoPE tables fp32, validate inputs - Drop KandinskyS2VPiflowScheduler; the distilled blocks use PiflowScheduler. - PiflowScheduler: on the final step, pass the nominal segment to the policy and clamp only the rollout end to eps, as the native K6 and S2V samplers do. - Keep the S2V RoPE tables as plain attributes so `.to(torch.bfloat16)` leaves them in fp32. - Raise when only one of negative_prompt_embeds / negative_pooled_prompt_embeds is passed, and when the Euler path gets num_inference_steps < 2. --- docs/source/en/_toctree.yml | 2 - docs/source/en/api/pipelines/kandinsky_s2v.md | 6 +- .../en/api/schedulers/kandinsky_s2v_piflow.md | 21 -- docs/source/en/api/schedulers/piflow.md | 3 +- src/diffusers/__init__.py | 2 - .../transformers/transformer_kandinsky_s2v.py | 37 ++- .../kandinsky_s2v/before_denoise.py | 2 + .../kandinsky_s2v/denoise.py | 47 +-- .../modular_blocks_kandinsky_s2v_distilled.py | 2 +- src/diffusers/schedulers/__init__.py | 2 - .../scheduling_kandinsky_s2v_piflow.py | 305 ------------------ src/diffusers/schedulers/scheduling_piflow.py | 16 +- src/diffusers/utils/dummy_pt_objects.py | 15 - .../test_models_transformer_kandinsky_s2v.py | 9 +- .../test_modular_pipeline_kandinsky_s2v.py | 22 +- tests/schedulers/test_scheduler_piflow.py | 37 ++- 16 files changed, 131 insertions(+), 397 deletions(-) delete mode 100644 docs/source/en/api/schedulers/kandinsky_s2v_piflow.md delete mode 100644 src/diffusers/schedulers/scheduling_kandinsky_s2v_piflow.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 6e264325f2e7..824626bfc5aa 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -842,8 +842,6 @@ title: HeunDiscreteScheduler - local: api/schedulers/ipndm title: IPNDMScheduler - - local: api/schedulers/kandinsky_s2v_piflow - title: KandinskyS2VPiflowScheduler - local: api/schedulers/dpm_discrete_ancestral title: KDPM2AncestralDiscreteScheduler - local: api/schedulers/dpm_discrete diff --git a/docs/source/en/api/pipelines/kandinsky_s2v.md b/docs/source/en/api/pipelines/kandinsky_s2v.md index 2732214a58cf..0891ebc20a04 100644 --- a/docs/source/en/api/pipelines/kandinsky_s2v.md +++ b/docs/source/en/api/pipelines/kandinsky_s2v.md @@ -33,8 +33,10 @@ only; there is no standard `DiffusionPipeline` class for it. |---|---|---|---|---|---| | [Lite][lite] | 2B | [`KandinskyS2VAutoBlocks`] | [`FlowMatchEulerDiscreteScheduler`] | 50 | 3.0 | | [Pro][pro] | 21B | [`KandinskyS2VAutoBlocks`] | [`FlowMatchEulerDiscreteScheduler`] | 50 | 3.0 | -| [Lite distilled][lite-distill] | 2B | [`KandinskyS2VDistilledAutoBlocks`] | [`KandinskyS2VPiflowScheduler`] | 16 | off | -| [Pro distilled][pro-distill] | 21B | [`KandinskyS2VDistilledAutoBlocks`] | [`KandinskyS2VPiflowScheduler`] | 16 | off | +| [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 diff --git a/docs/source/en/api/schedulers/kandinsky_s2v_piflow.md b/docs/source/en/api/schedulers/kandinsky_s2v_piflow.md deleted file mode 100644 index a310bdfcde45..000000000000 --- a/docs/source/en/api/schedulers/kandinsky_s2v_piflow.md +++ /dev/null @@ -1,21 +0,0 @@ - - -# KandinskyS2VPiflowScheduler - -`KandinskyS2VPiflowScheduler` is the few-step pi-Flow scheduler of the distilled Kandinsky S2V 1.0 -checkpoints. The student predicts a grid of x0 values per segment in one forward -pass, and the scheduler integrates that policy across the segment, so each -segment costs a single model evaluation. - -## KandinskyS2VPiflowScheduler -[[autodoc]] KandinskyS2VPiflowScheduler 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 00bcaac2f8a9..940138e985ed 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -491,7 +491,6 @@ "HeliosScheduler", "HeunDiscreteScheduler", "IPNDMScheduler", - "KandinskyS2VPiflowScheduler", "KarrasVeScheduler", "KDPM2AncestralDiscreteScheduler", "KDPM2DiscreteScheduler", @@ -1419,7 +1418,6 @@ HeliosScheduler, HeunDiscreteScheduler, IPNDMScheduler, - KandinskyS2VPiflowScheduler, KarrasVeScheduler, KDPM2AncestralDiscreteScheduler, KDPM2DiscreteScheduler, diff --git a/src/diffusers/models/transformers/transformer_kandinsky_s2v.py b/src/diffusers/models/transformers/transformer_kandinsky_s2v.py index 4cf0d3e70b9d..c423ca77a5aa 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky_s2v.py +++ b/src/diffusers/models/transformers/transformer_kandinsky_s2v.py @@ -356,9 +356,8 @@ def forward(self, visual_embed: Tensor, time_embed: Tensor, time_embed_idx: Tens ) -# Not copied from transformer_kandinsky6: its Kandinsky6RoPE1D subclasses a base -# RoPE1D defined in that module, so adopting it would mean copying 45 lines to -# replace these 12. +# 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__() @@ -367,15 +366,16 @@ def __init__(self, dim: int, max_pos: int = 1024, max_period: float = 10000.0): self.max_period = max_period freq = get_freqs(dim // 2, max_period) pos = torch.arange(max_pos, dtype=freq.dtype) - self.register_buffer("args", torch.outer(pos, freq), persistent=False) + # 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: - return to_rope(self.args[pos]) + return to_rope(self.args.to(device=pos.device)[pos]) -# Not copied from transformer_kandinsky6: 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. +# 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.""" @@ -391,6 +391,8 @@ def __init__( self.max_period = max_period 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 = ( @@ -398,7 +400,7 @@ def __init__( if i == 0 else torch.arange(ax_max_pos, dtype=freq.dtype) ) - self.register_buffer(f"args_{i}", torch.outer(pos, freq), persistent=False) + self.args.append(torch.outer(pos, freq)) def forward( self, @@ -407,9 +409,10 @@ def forward( scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), ) -> Tensor: duration, height, width = shape - args_t = self.args_0[pos[0] + self.offsets_t] / scale_factor[0] - args_h = self.args_1[pos[1]] / scale_factor[1] - args_w = self.args_2[pos[2]] / scale_factor[2] + args_0, args_1, args_2 = (args.to(device=axis_pos.device) for args, axis_pos in zip(self.args, pos)) + 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( [ @@ -570,10 +573,9 @@ class KandinskyS2VTransformer3DModel( "KandinskyS2VTransformerDecoderBlock", ] _no_split_modules = _repeated_blocks - # The checkpoint keeps the time embedding and every normalisation in fp32, and - # the position tables are built 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. + # The checkpoint keeps the time embedding and every normalisation 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", "norm", @@ -583,9 +585,6 @@ class KandinskyS2VTransformer3DModel( "feed_forward_norm", "query_norm", "key_norm", - "visual_rope_embeddings", - "text_rope_embeddings", - "audio_rope_embeddings", ] _supports_gradient_checkpointing = True diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py b/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py index ab246093af28..ce5fd329e53a 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py @@ -156,6 +156,8 @@ def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState 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.negative_text_rope_pos = None + if (block_state.negative_prompt_embeds is None) != (block_state.negative_pooled_prompt_embeds is None): + raise ValueError("`negative_prompt_embeds` and `negative_pooled_prompt_embeds` must be passed together.") 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) diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py b/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py index b3a77170a0ad..7dd61e79a395 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py @@ -18,8 +18,7 @@ from ...configuration_utils import FrozenDict from ...guiders import ClassifierFreeGuidance from ...models.transformers.transformer_kandinsky_s2v import KandinskyS2VTransformer3DModel -from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...schedulers.scheduling_kandinsky_s2v_piflow import KandinskyS2VPiflowScheduler +from ...schedulers import FlowMatchEulerDiscreteScheduler, PiflowScheduler from ...utils import logging from ...utils.torch_utils import randn_tensor from ..modular_pipeline import ( @@ -97,14 +96,15 @@ def inputs(self) -> list[InputParam]: def __call__( self, components: KandinskyS2VModularPipeline, block_state: BlockState, k: int ) -> tuple[KandinskyS2VModularPipeline, BlockState]: - device = components._execution_device - if getattr(components.scheduler, "is_piflow", False): - components.scheduler.set_timesteps(block_state.num_inference_steps, device=device) - else: - # 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=device) + 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 @@ -118,12 +118,20 @@ def description(self) -> str: @property def expected_components(self) -> list[ComponentSpec]: - return [ComponentSpec("scheduler", KandinskyS2VPiflowScheduler)] + 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 + class KandinskyS2VChunkDenoiseInner(ModularPipelineBlocks): model_name = "kandinsky-s2v" @@ -176,12 +184,6 @@ def inputs(self) -> list[InputParam]: def __call__( self, components: KandinskyS2VModularPipeline, block_state: BlockState, k: int ) -> tuple[KandinskyS2VModularPipeline, BlockState]: - if getattr(components.scheduler, "is_piflow", False) and components.requires_unconditional_embeds: - # Guidance on top of a distilled student over-guides it and doubles the cost. - raise ValueError( - f"guidance must be off with {type(components.scheduler).__name__}: CFG is distilled into the student" - ) - config = components.transformer.config transformer_dtype = components.transformer.dtype device = components._execution_device @@ -249,7 +251,7 @@ class KandinskyS2VDistilledChunkDenoiseInner(KandinskyS2VChunkDenoiseInner): def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("transformer", KandinskyS2VTransformer3DModel), - ComponentSpec("scheduler", KandinskyS2VPiflowScheduler), + ComponentSpec("scheduler", PiflowScheduler), ComponentSpec( "guider", ClassifierFreeGuidance, @@ -265,6 +267,15 @@ def inputs(self) -> list[InputParam]: for param in super().inputs ] + @torch.no_grad() + def __call__( + self, components: KandinskyS2VModularPipeline, block_state: BlockState, k: int + ) -> tuple[KandinskyS2VModularPipeline, BlockState]: + if components.requires_unconditional_embeds: + # Guidance on top of a distilled student over-guides it and doubles the cost. + raise ValueError("guidance must be off with the distilled checkpoints: CFG is distilled into the student") + return super().__call__(components, block_state, k) + class KandinskyS2VChunkUpdateStep(ModularPipelineBlocks): model_name = "kandinsky-s2v" 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 index 9e6ac2e92b41..c3d567277e7a 100644 --- 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 @@ -45,7 +45,7 @@ class KandinskyS2VDistilledAutoBlocks(SequentialPipelineBlocks): Components: tokenizer (`Qwen2_5OmniProcessor`) transformer (`KandinskyS2VTransformer3DModel`) text_encoder (`Qwen2_5OmniThinkerForConditionalGeneration`) text_encoder_2 (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) - guider (`ClassifierFreeGuidance`) vae (`AutoencoderKLHunyuanVideo`) scheduler (`KandinskyS2VPiflowScheduler`) + guider (`ClassifierFreeGuidance`) vae (`AutoencoderKLHunyuanVideo`) scheduler (`PiflowScheduler`) video_processor (`VideoProcessor`) Inputs: diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index a48ecd802be2..3909159729d6 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -68,7 +68,6 @@ _import_structure["scheduling_helios_dmd"] = ["HeliosDMDScheduler"] _import_structure["scheduling_heun_discrete"] = ["HeunDiscreteScheduler"] _import_structure["scheduling_ipndm"] = ["IPNDMScheduler"] - _import_structure["scheduling_kandinsky_s2v_piflow"] = ["KandinskyS2VPiflowScheduler"] _import_structure["scheduling_k_dpm_2_ancestral_discrete"] = ["KDPM2AncestralDiscreteScheduler"] _import_structure["scheduling_k_dpm_2_discrete"] = ["KDPM2DiscreteScheduler"] _import_structure["scheduling_lcm"] = ["LCMScheduler"] @@ -157,7 +156,6 @@ from .scheduling_ipndm import IPNDMScheduler from .scheduling_k_dpm_2_ancestral_discrete import KDPM2AncestralDiscreteScheduler from .scheduling_k_dpm_2_discrete import KDPM2DiscreteScheduler - from .scheduling_kandinsky_s2v_piflow import KandinskyS2VPiflowScheduler from .scheduling_lcm import LCMScheduler from .scheduling_ltx_euler_ancestral_rf import LTXEulerAncestralRFScheduler from .scheduling_minimax_h3 import MiniMaxH3Scheduler diff --git a/src/diffusers/schedulers/scheduling_kandinsky_s2v_piflow.py b/src/diffusers/schedulers/scheduling_kandinsky_s2v_piflow.py deleted file mode 100644 index b45f2e11d295..000000000000 --- a/src/diffusers/schedulers/scheduling_kandinsky_s2v_piflow.py +++ /dev/null @@ -1,305 +0,0 @@ -# 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. -"""pi-Flow scheduler for distilled Kandinsky S2V 1.0 checkpoints. - -pi-Flow (arXiv 2510.14974) replaces the per-step velocity with a *policy*: the student predicts a grid of x0 values -spanning the segment ``[t_dst, t_src]`` in a single forward, and the ODE is then integrated across the segment without -further network evaluations. One segment therefore costs exactly one model evaluation, and 16 segments replace the 50 -Euler steps of the base model. - -The distilled checkpoints widen the output head ``n_grid`` times; this scheduler splits that axis, so the transformer -itself is unchanged and its config only differs in ``out_visual_dim``. -""" - -from __future__ import annotations - -import torch -from torch import Tensor - -from ..configuration_utils import register_to_config -from .scheduling_flow_match_euler_discrete import ( - FlowMatchEulerDiscreteScheduler, - FlowMatchEulerDiscreteSchedulerOutput, -) - - -_MIN_GRID_POINTS = 2 - - -def shift_timesteps(t: Tensor, shift: float) -> Tensor: - """Flow-matching timestep warp.""" - return shift * t / (1 + (shift - 1) * t) - - -class DXPolicy: - """Dynamic-x0 policy: a grid of x0 predictions, linearly interpolated in raw time.""" - - def __init__( # noqa: PLR0913 - self, - denoising_output: Tensor, - x_t_src: Tensor, - sigma_t_src: Tensor, - segment_size: Tensor, - shift: float = 1.0, - eps: float = 1e-4, - ): - self.x_t_src = x_t_src - self.ndim = x_t_src.dim() - self.shift = shift - self.eps = eps - - self.sigma_t_src = sigma_t_src.reshape(*sigma_t_src.size(), *((self.ndim - sigma_t_src.dim()) * [1])) - self.raw_t_src = self._unwarp_t(self.sigma_t_src) - seg = segment_size - if isinstance(seg, Tensor) and seg.dim() < self.raw_t_src.dim(): - seg = seg.reshape(*seg.size(), *((self.ndim - seg.dim()) * [1])) - # clamp(min=0) here while the sampler clamps to eps — that asymmetry comes - # from the distillation fork (piflow_s2v.py vs utils_s2v_piflow.py) and the - # students were trained against it. Do not align the two. - self.raw_t_dst = (self.raw_t_src - seg).clamp(min=0) - self.segment_size = (self.raw_t_src - self.raw_t_dst).clamp(min=eps) - - self.denoising_output_x_0 = self._u_to_x_0(denoising_output, self.x_t_src, self.sigma_t_src) - - @staticmethod - def _interpolate(x: Tensor, t: Tensor) -> Tensor: - n = x.size(1) - if n < _MIN_GRID_POINTS: - return x.squeeze(1) - t = t.clamp(min=0, max=1) * (n - 1) - t0 = t.floor().to(torch.long).clamp(min=0, max=n - 2) - t1 = t0 + 1 - t0t1 = torch.stack([t0, t1], dim=1) - x0x1 = torch.gather(x, dim=1, index=t0t1.expand(-1, -1, *x.shape[2:])) - return (t1 - t) * x0x1[:, 0] + (t - t0) * x0x1[:, 1] - - def _unwarp_t(self, sigma_t: Tensor) -> Tensor: - return sigma_t / (self.shift + (1 - self.shift) * sigma_t) - - @staticmethod - def _u_to_x_0(denoising_output: Tensor, x_t: Tensor, sigma_t: Tensor) -> Tensor: - return x_t.unsqueeze(1) - sigma_t.unsqueeze(1) * denoising_output - - def pi(self, x_t: Tensor, sigma_t: Tensor) -> Tensor: - sigma_t = sigma_t.reshape(*sigma_t.size(), *((self.ndim - sigma_t.dim()) * [1])) - raw_t = self._unwarp_t(sigma_t) - x_0 = self._interpolate(self.denoising_output_x_0, (raw_t - self.raw_t_dst) / self.segment_size) - return (x_t - x_0) / sigma_t.clamp(min=self.eps) - - -def policy_rollout_fm( # noqa: PLR0913 - x_t_start: Tensor, - sigma_t_start: Tensor, - raw_t_start: Tensor, - raw_t_end: Tensor, - total_substeps: int, - policy: DXPolicy, -) -> Tensor: - """Integrate the policy ODE across a segment with a fixed substep size.""" - num_tokens = x_t_start.size(0) - ndim = x_t_start.dim() - raw_t_start = raw_t_start.reshape(num_tokens, *((ndim - 1) * [1])) - raw_t_end = raw_t_end.reshape(num_tokens, *((ndim - 1) * [1])) - - delta_raw_t = raw_t_start - raw_t_end - num_substeps = (delta_raw_t * total_substeps).round().to(torch.long).clamp(min=1) - substep_size = delta_raw_t / num_substeps - - raw_t, sigma_t, x_t = raw_t_start, sigma_t_start, x_t_start - for substep_id in range(int(num_substeps.max().item())): - u = policy.pi(x_t, sigma_t) - - raw_t_minus = (raw_t - substep_size).clamp(min=0) - sigma_t_minus = shift_timesteps(raw_t_minus, policy.shift) - x_t_minus = x_t + u * (sigma_t_minus - sigma_t) - - active = num_substeps > substep_id - x_t = torch.where(active, x_t_minus, x_t) - sigma_t = torch.where(active, sigma_t_minus, sigma_t) - raw_t = torch.where(active, raw_t_minus, raw_t) - return x_t - - -class KandinskyS2VPiflowScheduler(FlowMatchEulerDiscreteScheduler): - """Few-step pi-Flow scheduler for distilled Kandinsky S2V 1.0 checkpoints. - - Args: - num_train_timesteps (`int`, *optional*, defaults to 1000): - Number of training diffusion steps. - shift (`float`, *optional*, defaults to 10.0): - Flow-matching timestep shift. - n_grid (`int`, *optional*, defaults to 5): - Number of x0 predictions packed into the model's output channels. - eps (`float`, *optional*, defaults to 1e-6): - Minimum raw timestep and policy denominator. - final_step_size_scale (`float`, *optional*, defaults to 0.5): - Size of the last segment relative to the others. - num_policy_substeps (`int`, *optional*, defaults to 128): - Substeps per unit of raw time used to integrate the policy. - """ - - is_piflow = True - - @register_to_config - def __init__( # noqa: PLR0913 - self, - num_train_timesteps: int = 1000, - shift: float = 10.0, - n_grid: int = 5, - eps: float = 1e-6, - final_step_size_scale: float = 0.5, - num_policy_substeps: int = 128, - ): - if n_grid < _MIN_GRID_POINTS: - raise ValueError(f"n_grid must be >= {_MIN_GRID_POINTS}, got {n_grid}") - if eps <= 0: - raise ValueError(f"eps must be > 0, got {eps}") - if not 0 < final_step_size_scale <= 1: - raise ValueError(f"final_step_size_scale must be in (0, 1], got {final_step_size_scale}") - if num_policy_substeps < 1: - raise ValueError(f"num_policy_substeps must be >= 1, got {num_policy_substeps}") - super().__init__(num_train_timesteps=num_train_timesteps, shift=shift) - self._raw_timesteps = torch.empty(0) - self._segment_sizes = torch.empty(0) - - def set_timesteps( # noqa: PLR0913 - self, - num_inference_steps: int | None = None, - device: str | torch.device | None = None, - sigmas: list[float] | None = None, - mu: float | None = None, - timesteps: list[float] | None = None, - ) -> None: - """Build the distilled segment schedule. - - Every segment has the same size except the last one, which is scaled by ``final_step_size_scale``; together - they cover raw time from 1 down to 0. - """ - if sigmas is not None or mu is not None or timesteps is not None: - raise ValueError("this scheduler only supports its own distilled schedule") - if num_inference_steps is None or num_inference_steps < 1: - raise ValueError(f"num_inference_steps must be positive, got {num_inference_steps}") - - config = self.config - base_segment = 1.0 / (num_inference_steps - (1.0 - config.final_step_size_scale)) - - # Walked segment by segment rather than as `1 - cumsum(segments)`: the - # native sampler subtracts one segment at a time, and in float32 the two - # orders differ in the last bits, which the policy rollout then amplifies. - raws, segments = [], [] - raw_t = torch.ones(1, dtype=torch.float32, device=device) - for step in range(num_inference_steps): - is_final = step == num_inference_steps - 1 - segment = base_segment * (config.final_step_size_scale if is_final else 1.0) - raws.append(raw_t) - segments.append(segment) - raw_t = (raw_t - segment).clamp(min=config.eps) - - raw = torch.cat(raws) - segments = torch.tensor(segments, dtype=torch.float32, device=device) - raw_sigmas = shift_timesteps(raw, float(config.shift)) - - self.num_inference_steps = int(num_inference_steps) - self._raw_timesteps = raw - self._segment_sizes = segments - self.timesteps = raw_sigmas * config.num_train_timesteps - self.sigmas = torch.cat([raw_sigmas, raw_sigmas.new_zeros(1)]) - self._step_index = None - self._begin_index = None - - def _to_grid(self, model_output: Tensor, sample: Tensor) -> Tensor: - """``(T, H, W, n_grid * C)`` -> ``(T, n_grid, H, W, C)``. - - Grid-major, exactly as the distilled head emits it. Do not reorder: a channel-major split would silently - shuffle the policy grid while every shape still checks out. - """ - n_grid = self.config.n_grid - if model_output.shape[:-1] != sample.shape[:-1]: - raise ValueError(f"model output {tuple(model_output.shape)} does not match sample {tuple(sample.shape)}") - if model_output.shape[-1] != sample.shape[-1] * n_grid: - raise ValueError( - f"expected {sample.shape[-1] * n_grid} output channels for n_grid={n_grid}, " - f"got {model_output.shape[-1]}" - ) - return model_output.reshape(*model_output.shape[:-1], n_grid, sample.shape[-1]).movedim(-2, 1) - - def _resolve_step_index(self, timestep: Tensor | float) -> int: - if self.step_index is None: - if self.begin_index is not None: - self._step_index = self.begin_index - else: - schedule = self.timesteps - value = torch.as_tensor(timestep, device=schedule.device, dtype=schedule.dtype) - indices = torch.nonzero(schedule == value).flatten() - if not indices.numel(): - raise ValueError(f"timestep {float(value)} is not in the pi-Flow schedule") - self._step_index = int(indices[-1].item()) - if self.step_index >= self.num_inference_steps: - raise ValueError("step() called after the schedule was exhausted") - return int(self.step_index) - - def step( - self, - model_output: Tensor, - timestep: float | Tensor, - sample: Tensor, - return_dict: bool = True, - ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple: - """Advance one segment by integrating the policy. - - Args: - model_output: the widened prediction, ``n_grid`` x0 values per channel. - timestep: a value from ``scheduler.timesteps``. - sample: the current noisy latent. - return_dict: return a [`FlowMatchEulerDiscreteSchedulerOutput`]. - """ - if isinstance(timestep, int) or (isinstance(timestep, Tensor) and not timestep.is_floating_point()): - raise ValueError("pass a value from scheduler.timesteps, not an index") - - step_index = self._resolve_step_index(timestep) - config = self.config - # The policy rollout runs in float32, as the native sampler does. - sample = sample.to(torch.float32) - grid = self._to_grid(model_output.to(torch.float32), sample) - - device = sample.device - num_tokens = sample.shape[0] - extra_dims = (sample.dim() - 1) * [1] - - raw_src = self._raw_timesteps[step_index].to(device) - segment = self._segment_sizes[step_index].to(device) - raw_dst = (raw_src - segment).clamp(min=config.eps) - - sigma_src = self.sigmas[step_index].to(device).repeat_interleave(num_tokens).reshape(num_tokens, *extra_dims) - policy = DXPolicy( - denoising_output=grid, - x_t_src=sample, - sigma_t_src=sigma_src, - segment_size=torch.full_like(sigma_src, float(segment)), - shift=float(config.shift), - eps=config.eps, - ) - updated = policy_rollout_fm( - x_t_start=sample, - sigma_t_start=sigma_src, - raw_t_start=raw_src.repeat_interleave(num_tokens), - raw_t_end=raw_dst.repeat_interleave(num_tokens), - total_substeps=config.num_policy_substeps, - policy=policy, - ) - - self._step_index += 1 - if return_dict: - return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=updated) - return (updated,) diff --git a/src/diffusers/schedulers/scheduling_piflow.py b/src/diffusers/schedulers/scheduling_piflow.py index d858b425a9fd..6df059f3babc 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_final_segment = None @property def step_index(self): @@ -247,6 +248,7 @@ def set_timesteps( sigmas = shift_timesteps(raw, float(self.config.shift)) self.num_inference_steps = int(num_inference_steps) self._piflow_raw_timesteps = raw + self._piflow_final_segment = segment * self.final_step_size_scale self.timesteps = sigmas * self.config.num_train_timesteps self.sigmas = torch.cat([sigmas, sigmas.new_zeros(1)]) self._step_index = None @@ -273,15 +275,17 @@ 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) + if step_index + 1 < self._piflow_raw_timesteps.numel(): + raw_dst = self._piflow_raw_timesteps[step_index + 1].to(device=sample.device) + segment = raw_src - raw_dst + else: + # The policy sees the nominal final segment, which ends at 0; only the rollout stops at eps. + segment = raw_src.new_full((), self._piflow_final_segment) + 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 = segment.expand(sample.shape[0]) policy = DXPolicy( model_output, sample, diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 35ad9d8844b2..9a734961af1f 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -3936,21 +3936,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class KandinskyS2VPiflowScheduler(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 KarrasVeScheduler(metaclass=DummyObject): _backends = ["torch"] diff --git a/tests/models/transformers/test_models_transformer_kandinsky_s2v.py b/tests/models/transformers/test_models_transformer_kandinsky_s2v.py index caf31091799a..9e3c7a661fcb 100644 --- a/tests/models/transformers/test_models_transformer_kandinsky_s2v.py +++ b/tests/models/transformers/test_models_transformer_kandinsky_s2v.py @@ -140,4 +140,11 @@ def randn(*shape: int) -> torch.Tensor: class TestKandinskyS2VTransformer(KandinskyS2VTransformerTesterConfig, ModelTesterMixin): - pass + 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) 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 index ad3ed0b97e7d..06bf41b3a484 100644 --- a/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py +++ b/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py @@ -31,6 +31,7 @@ KandinskyS2VDistilledModularPipeline, KandinskyS2VModularPipeline, ) +from diffusers.modular_pipelines.kandinsky_s2v.before_denoise import KandinskyS2VPrepareStep from ..testing_utils import ( BaseModularPipelineTesterConfig, @@ -122,6 +123,25 @@ def test_rejects_speech_past_the_encoder_limit(self): with pytest.raises(ValueError, match="truncates"): pipe(**inputs, output="videos") + def test_rejects_a_single_step(self): + """One step leaves an empty Euler schedule, which would decode the initial noise.""" + pipe = self.get_pipeline() + with pytest.raises(ValueError, match="at least 2"): + pipe(**{**self.get_dummy_inputs(), "num_inference_steps": 1}, output="videos") + + def test_rejects_negative_embeds_without_pooled(self): + prepare_pipe = KandinskyS2VPrepareStep().init_pipeline(self.pretrained_model_name_or_path) + prepare_pipe.load_components() + with pytest.raises(ValueError, match="must be passed together"): + prepare_pipe( + reference_latents=torch.randn(1, 8, 8, 4), + audio_embeds=torch.randn(16, 16), + num_chunks=1, + prompt_embeds=torch.randn(3, 16), + pooled_prompt_embeds=torch.randn(16), + negative_prompt_embeds=torch.randn(3, 16), + ) + class TestKandinskyS2VModularPipelineLoading(KandinskyS2VModularPipelineTesterConfig, ModularLoadingTesterMixin): pass @@ -140,7 +160,7 @@ class TestKandinskyS2VModularPipelineMemory(KandinskyS2VModularPipelineTesterCon class KandinskyS2VDistilledModularPipelineTesterConfig(KandinskyS2VModularPipelineTesterConfig): - """The pi-Flow checkpoints: a head widened to its grid, `KandinskyS2VPiflowScheduler`, guidance off.""" + """The pi-Flow checkpoints: a head widened to its grid, `PiflowScheduler`, guidance off.""" pipeline_class = KandinskyS2VDistilledModularPipeline pipeline_blocks_class = KandinskyS2VDistilledAutoBlocks 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) From 916f2faf7aaa49e763e0ec9e37eb7e9abfc31800 Mon Sep 17 00:00:00 2001 From: Nikolay Semenov Date: Tue, 27 Oct 2026 05:04:02 +0300 Subject: [PATCH 03/12] Address review: quantization-safe modulation, drop duration, guard optional deps - KandinskyS2VModulation calls its layer instead of reading the weight, and keeps the modulation layers in fp32 via _keep_in_fp32_modules (as in K6). - Remove the transformer's `duration` argument; it can only be the frame count. - Import librosa, soundfile and torchvision behind is_*_available(). - PiflowScheduler: walk raw time one segment at a time and give the policy the nominal segment on every step, as the native samplers do. - Update the Pro offload peak in the docs (about 50 GiB). --- docs/source/en/api/pipelines/kandinsky_s2v.md | 2 +- .../transformers/transformer_kandinsky_s2v.py | 37 ++++++++----------- .../kandinsky_s2v/denoise.py | 1 - .../kandinsky_s2v/encoders.py | 12 +++--- .../kandinsky_s2v/modular_pipeline.py | 10 ++--- src/diffusers/schedulers/scheduling_piflow.py | 25 +++++++------ .../test_models_transformer_kandinsky_s2v.py | 1 - 7 files changed, 43 insertions(+), 45 deletions(-) diff --git a/docs/source/en/api/pipelines/kandinsky_s2v.md b/docs/source/en/api/pipelines/kandinsky_s2v.md index 0891ebc20a04..d11b15b00d74 100644 --- a/docs/source/en/api/pipelines/kandinsky_s2v.md +++ b/docs/source/en/api/pipelines/kandinsky_s2v.md @@ -68,7 +68,7 @@ example `ffmpeg -i avatar.mp4 -i speech.wav -c:v copy -c:a aac -shortest out.mp4 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. Keep one model on the GPU at a time, which -brings the peak to about 46 GiB: +brings the peak to about 50 GiB: ```py from diffusers import ComponentsManager diff --git a/src/diffusers/models/transformers/transformer_kandinsky_s2v.py b/src/diffusers/models/transformers/transformer_kandinsky_s2v.py index c423ca77a5aa..043d1535ad80 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky_s2v.py +++ b/src/diffusers/models/transformers/transformer_kandinsky_s2v.py @@ -30,7 +30,7 @@ from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin +from ..modeling_utils import ModelMixin, get_parameter_dtype logger = logging.get_logger(__name__) @@ -204,25 +204,17 @@ def forward(self, x: Tensor) -> Tensor: return self.in_layer(x) -# Not copied from transformer_kandinsky6: its Modulation casts the result back to -# the input dtype, while the AdaLN parameters here stay fp32 all the way through -# the per-token selection, as native S2V does. +# 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) - self.out_layer.weight.data.zero_() - self.out_layer.bias.data.zero_() def forward(self, x: Tensor) -> Tensor: - # AdaLN runs in fp32 whatever the weights are stored as, matching native - # generation. Upcasting a bf16 weight is exact, so the result is identical. - return F.linear( - self.activation(x.float()), - self.out_layer.weight.float(), - self.out_layer.bias.float(), - ) + 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 @@ -409,7 +401,9 @@ def forward( 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)) + 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] @@ -573,11 +567,15 @@ class KandinskyS2VTransformer3DModel( "KandinskyS2VTransformerDecoderBlock", ] _no_split_modules = _repeated_blocks - # The checkpoint keeps the time embedding and every normalisation 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. + # 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", @@ -713,7 +711,6 @@ def forward( # noqa: PLR0913, PLR0915 text_rope_pos: Tensor, audio_rope_pos: Tensor, scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), - duration: int | None = None, return_dict: bool = True, ) -> Transformer2DModelOutput | Tensor: """Denoise one chunk. @@ -732,8 +729,6 @@ def forward( # noqa: PLR0913, PLR0915 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. - duration: number of latent frames used to align speech to frames; - defaults to the temporal size of ``hidden_states``. 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. @@ -752,7 +747,7 @@ def forward( # noqa: PLR0913, PLR0915 audio_embed = audio_hidden_states.squeeze(0) device = x.device - duration = x.shape[0] if duration is None else duration + duration = x.shape[0] text_embed = self.text_embeddings(text_embed) audio_embed = self.audio_embeddings(audio_embed) diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py b/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py index 7dd61e79a395..46b283ef9369 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py @@ -212,7 +212,6 @@ def __call__( "memory_rope_pos": tuple(reversed(memory_rope_pos(config, memory, device))), "audio_rope_pos": audio_rope_pos(chunk_audio.shape[0], num_latent_frames, device), "scale_factor": block_state.scale_factor, - "duration": num_latent_frames, "return_dict": False, } diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py index ef17b9c560f1..0a74afa2bfd8 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py @@ -29,7 +29,7 @@ from ...guiders import ClassifierFreeGuidance from ...models import AutoencoderKLHunyuanVideo from ...models.transformers.transformer_kandinsky_s2v import KandinskyS2VTransformer3DModel -from ...utils import logging +from ...utils import is_librosa_available, is_torchvision_available, logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import ( @@ -47,6 +47,12 @@ ) +if is_librosa_available(): + import librosa + +if is_torchvision_available(): + from torchvision import transforms + logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -72,8 +78,6 @@ def encode_qwen_omni( chat_text = components.tokenizer.apply_chat_template([conversation], add_generation_prompt=False, tokenize=False) audios = None if audio_path is not None: - import librosa # noqa: PLC0415 - audios = [librosa.load(str(audio_path), sr=AUDIO_SAMPLE_RATE, offset=0.0, duration=audio_end)[0]] inputs = components.tokenizer( @@ -425,8 +429,6 @@ def intermediate_outputs(self) -> list[OutputParam]: @torch.no_grad() def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: - from torchvision import transforms # noqa: PLC0415 - block_state = self.get_block_state(state) device = components._execution_device image, height, width = block_state.reference_image, block_state.height, block_state.width diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py index 32bdf04df4eb..b6eb301baefb 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py @@ -21,10 +21,14 @@ import torch from ...models.transformers.transformer_kandinsky_s2v import audio_group_lengths -from ...utils import logging +from ...utils import is_librosa_available, logging from ..modular_pipeline import ModularPipeline +# soundfile is a dependency of librosa. +if is_librosa_available(): + import soundfile as sf + logger = logging.get_logger(__name__) # pylint: disable=invalid-name SYSTEM_PROMPT = ( @@ -73,16 +77,12 @@ def select_resolution( def audio_duration(path: str | Path) -> float: - import soundfile as sf # noqa: PLC0415 - info = sf.info(str(path)) return info.frames / info.samplerate def pad_with_silence(path: str | Path, target_duration: float) -> Path | None: """Append silence so the track covers ``target_duration``; ``None`` if it already does.""" - import soundfile as sf # noqa: PLC0415 - original = audio_duration(path) if target_duration <= original + 0.01: return None diff --git a/src/diffusers/schedulers/scheduling_piflow.py b/src/diffusers/schedulers/scheduling_piflow.py index 6df059f3babc..7eaf7ccdbab8 100644 --- a/src/diffusers/schedulers/scheduling_piflow.py +++ b/src/diffusers/schedulers/scheduling_piflow.py @@ -195,7 +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_final_segment = None + self._piflow_segments = [] @property def step_index(self): @@ -244,11 +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_final_segment = segment * self.final_step_size_scale + 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 @@ -275,17 +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) - if step_index + 1 < self._piflow_raw_timesteps.numel(): - raw_dst = self._piflow_raw_timesteps[step_index + 1].to(device=sample.device) - segment = raw_src - raw_dst - else: - # The policy sees the nominal final segment, which ends at 0; only the rollout stops at eps. - segment = raw_src.new_full((), self._piflow_final_segment) - raw_dst = (raw_src - segment).clamp(min=self.eps) + # 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 = segment.expand(sample.shape[0]) + segment = raw_src.new_full((sample.shape[0],), segment) policy = DXPolicy( model_output, sample, diff --git a/tests/models/transformers/test_models_transformer_kandinsky_s2v.py b/tests/models/transformers/test_models_transformer_kandinsky_s2v.py index 9e3c7a661fcb..af5f24dfb7e5 100644 --- a/tests/models/transformers/test_models_transformer_kandinsky_s2v.py +++ b/tests/models/transformers/test_models_transformer_kandinsky_s2v.py @@ -135,7 +135,6 @@ def randn(*shape: int) -> torch.Tensor: "audio_rope_pos": torch.arange(DURATION, device=torch_device).repeat_interleave( torch.tensor([1] + [4] * (DURATION - 1), device=torch_device) ), - "duration": DURATION, } From ee542b7b713c9e3661c1b9bf49eef3ed125b1714 Mon Sep 17 00:00:00 2001 From: Nikolay Semenov Date: Wed, 28 Oct 2026 03:29:36 +0300 Subject: [PATCH 04/12] Address review: fold helpers, accept waveform audio, offload by margin - Fold the one-line helpers into their callers, inline the memory embedding, name the projection layers explicitly and spell out the block arguments. - Accept `audio` as a mono 16 kHz float waveform (tensor or array) besides a path. - Run the fp16 test with the pipeline loaded in fp16; drop the variant test. - Docs: offload Pro with memory_reserve_margin="40GiB" and show a full run. --- docs/source/en/api/pipelines/kandinsky_s2v.md | 22 +- .../transformers/transformer_kandinsky_s2v.py | 191 ++++++++++-------- .../kandinsky_s2v/encoders.py | 57 ++++-- .../modular_blocks_kandinsky_s2v.py | 5 +- .../modular_blocks_kandinsky_s2v_distilled.py | 5 +- .../kandinsky_s2v/modular_pipeline.py | 30 ++- .../test_modular_pipeline_kandinsky_s2v.py | 81 +++++--- 7 files changed, 253 insertions(+), 138 deletions(-) diff --git a/docs/source/en/api/pipelines/kandinsky_s2v.md b/docs/source/en/api/pipelines/kandinsky_s2v.md index d11b15b00d74..15c9fd7d9de4 100644 --- a/docs/source/en/api/pipelines/kandinsky_s2v.md +++ b/docs/source/en/api/pipelines/kandinsky_s2v.md @@ -67,24 +67,32 @@ export_to_video(video, "avatar.mp4", fps=25) example `ffmpeg -i avatar.mp4 -i speech.wav -c:v copy -c:a aac -shortest out.mp4`. 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. Keep one model on the GPU at a time, which -brings the peak to about 50 GiB: +when every component is on the device. Offload them through a `ComponentsManager`: ```py -from diffusers import ComponentsManager +import torch +from diffusers import ComponentsManager, ModularPipeline +from diffusers.utils import export_to_video 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", offload_strategy=lambda hooks, **_: hooks) +manager.enable_auto_cpu_offload(device="cuda", memory_reserve_margin="40GiB") + +video = pipe(image="reference.png", audio="speech.wav", output="videos")[0] +export_to_video(video, "avatar.mp4", fps=25) ``` -The default strategy is not enough here: it keeps the transformer beside the speech -encoder because both fit, and the VAE decode that follows needs the room they hold. -Enable the offload after `load_components`, so the strategy applies to every component. +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 +- `audio` is a path to an audio file, or a mono float waveform in [-1, 1] at 16 kHz, as a + tensor or an array shaped `(samples,)`. - 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. diff --git a/src/diffusers/models/transformers/transformer_kandinsky_s2v.py b/src/diffusers/models/transformers/transformer_kandinsky_s2v.py index 043d1535ad80..7f7796263092 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky_s2v.py +++ b/src/diffusers/models/transformers/transformer_kandinsky_s2v.py @@ -45,26 +45,12 @@ def get_freqs(dim, max_period=10000.0): return freqs -def to_rope(args: Tensor) -> Tensor: - 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) - - 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 scale_shift_norm(norm: nn.Module, x: Tensor, scale: Tensor, shift: Tensor, idx: Tensor) -> Tensor: - """AdaLN with the modulation row selected per token.""" - return norm(x.float()) * (scale.float().index_select(0, idx) + 1.0) + shift.float().index_select(0, idx) - - -def gate_sum(x: Tensor, out: Tensor, gate: Tensor, idx: Tensor) -> Tensor: - return x.float() + gate.float().index_select(0, idx) * out.float() - - def split_lengths(total: int, groups: int, device: torch.device) -> Tensor: """Split ``total`` items into ``groups`` with boundaries at ``total * i // groups``.""" bounds = torch.arange(groups + 1, device=device) * total // groups @@ -332,10 +318,11 @@ def __init__(self, model_dim: int, time_dim: int, visual_dim: int, patch_size: t 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 = torch.chunk(self.modulation(time_embed), 2, dim=-1) - visual_embed = scale_shift_norm( - self.norm, visual_embed, scale[:, None, None], shift[:, None, None], time_embed_idx - ).type_as(visual_embed) + 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 @@ -362,7 +349,9 @@ def __init__(self, dim: int, max_pos: int = 1024, max_period: float = 10000.0): self.args = torch.outer(pos, freq) def forward(self, pos: Tensor) -> Tensor: - return to_rope(self.args.to(device=pos.device)[pos]) + 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 @@ -416,7 +405,8 @@ def forward( ], dim=-1, ) - return to_rope(args) + 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): @@ -434,14 +424,20 @@ def forward(self, x: Tensor, time_embed: Tensor, rope: Tensor, time_embed_idx: T 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 = scale_shift_norm(self.self_attention_norm, x, scale, shift, time_embed_idx).type_as(x) + 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 = gate_sum(x, out, gate, time_embed_idx).type_as(x) + 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 = scale_shift_norm(self.feed_forward_norm, x, scale, shift, time_embed_idx).type_as(x) + 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 gate_sum(x, out, gate, time_embed_idx).type_as(x) + return (x.float() + gate.float().index_select(0, time_embed_idx) * out.float()).type_as(x) class KandinskyS2VTransformerAudioEncoderBlock(nn.Module): @@ -503,18 +499,24 @@ def forward( # noqa: PLR0913 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 = scale_shift_norm(self.self_attention_norm, visual_embed, scale, shift, time_embed_idx).type_as( + 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 ) - out = self.self_attention(out, rotary_emb=rope) - visual_embed = gate_sum(visual_embed, out, gate, time_embed_idx).type_as(visual_embed) shift, scale, gate = torch.chunk(cross_attn_params, 3, dim=-1) - out = scale_shift_norm(self.cross_attention_norm, visual_embed, scale, shift, time_embed_idx).type_as( + 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 ) - out = self.cross_attention(out, encoder_hidden_states=text_embed) - visual_embed = gate_sum(visual_embed, out, gate, time_embed_idx).type_as(visual_embed) if self.enable_audio: # Speech conditions only the noisy tail, frame by frame. The residual is @@ -529,11 +531,14 @@ def forward( # noqa: PLR0913 visual_embed = torch.cat([head, noisy + out], dim=0) shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) - out = scale_shift_norm(self.feed_forward_norm, visual_embed, scale, shift, time_embed_idx).type_as( + 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 ) - out = self.feed_forward(out) - return gate_sum(visual_embed, out, gate, time_embed_idx).type_as(visual_embed) class KandinskyS2VTransformer3DModel( @@ -627,17 +632,33 @@ def __init__( # noqa: PLR0913 self.audio_embeddings = KandinskyS2VGeneralEmbeddings(in_audio_dim, model_dim) self.visual_embeddings = KandinskyS2VVisualEmbeddings(in_visual_dim, model_dim, patch_size) - for name, coef in ( - ("proj1x", patch_proj1x_coef), - ("proj2x", patch_proj2x_coef), - ("proj4x", patch_proj4x_coef), - ): - scaled = (patch_size[0] * coef[0], patch_size[1] * coef[1], patch_size[2] * coef[2]) - setattr( - self, - f"visual_embeddings_{name}", - KandinskyS2VVisualEmbeddings(in_visual_dim, model_dim, scaled), - ) + 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) @@ -671,31 +692,6 @@ def __init__( # noqa: PLR0913 self.out_layer = KandinskyS2VOutLayer(model_dim, time_dim, out_visual_dim, patch_size) self.gradient_checkpointing = False - def _embed_memory( - self, - memory_hidden_states: tuple[Tensor, Tensor, Tensor], - memory_rope_pos: tuple[list[Tensor], list[Tensor], list[Tensor]], - scale_factor: tuple[float, float, float], - ) -> tuple[Tensor, Tensor]: - """Embed the FramePack memory at its three temporal scales.""" - memory_4x, memory_2x, memory_1x = memory_hidden_states - rope_4x, rope_2x, rope_1x = memory_rope_pos - - parts = [(self.visual_embeddings_proj1x(memory_1x), rope_1x)] - if self.enable_framepack: - parts = [ - (self.visual_embeddings_proj4x(memory_4x), rope_4x), - (self.visual_embeddings_proj2x(memory_2x), rope_2x), - *parts, - ] - - embeds, ropes = [], [] - for embed, rope_pos in parts: - rope = self.visual_rope_embeddings(embed.shape[:-1], rope_pos, scale_factor) - embeds.append(embed.flatten(0, 2)) - ropes.append(rope.flatten(0, 2)) - return torch.cat(embeds, dim=0), torch.cat(ropes, dim=0) - def forward( # noqa: PLR0913, PLR0915 self, hidden_states: Tensor, @@ -712,7 +708,7 @@ def forward( # noqa: PLR0913, PLR0915 audio_rope_pos: Tensor, scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), return_dict: bool = True, - ) -> Transformer2DModelOutput | Tensor: + ) -> Transformer2DModelOutput | tuple[Tensor]: """Denoise one chunk. Args: @@ -786,7 +782,24 @@ def forward( # noqa: PLR0913, PLR0915 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) - memory_embed, memory_rope = self._embed_memory(memory, memory_rope_pos, scale_factor) + # 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)) @@ -807,21 +820,31 @@ def forward( # noqa: PLR0913, PLR0915 audio_groups = audio_group_lengths(audio_embed.shape[0], duration, device) for block in self.visual_transformer_blocks: - args = ( - visual_embed_full, - text_embed, - audio_embed, - time_embed_full, - visual_rope_full, - time_embed_idx_full, - num_visual_tokens, - visual_groups, - audio_groups, - ) if torch.is_grad_enabled() and self.gradient_checkpointing: - visual_embed_full = self._gradient_checkpointing_func(block, *args) + 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, + visual_groups, + audio_groups, + ) else: - visual_embed_full = block(*args) + visual_embed_full = block( + visual_embed_full, + text_embed, + audio_embed, + time_embed_full, + visual_rope_full, + time_embed_idx_full, + num_visual_tokens, + visual_groups, + audio_groups, + ) 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) diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py index 0a74afa2bfd8..da215da947e1 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py @@ -16,6 +16,7 @@ from pathlib import Path from typing import Any +import numpy as np import PIL.Image import torch from transformers import ( @@ -42,6 +43,8 @@ TEMPORAL_COMPRESSION_RATIO, KandinskyS2VModularPipeline, audio_duration, + audio_waveform, + pad_waveform_with_silence, pad_with_silence, select_resolution, ) @@ -59,14 +62,17 @@ def encode_qwen_omni( components: KandinskyS2VModularPipeline, prompt: str, - audio_path: str | Path | None, + audio: str | Path | 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.""" + """Thinker hidden states for the prompt and audio-tower features for the speech. + + ``audio`` is a path to an audio file or a mono waveform at `AUDIO_SAMPLE_RATE`. + """ user_content: list[dict[str, Any]] = [{"type": "text", "text": prompt}] - if audio_path is not None: - item: dict[str, Any] = {"type": "audio", "audio": str(audio_path)} + if audio is not None: + item: dict[str, Any] = {"type": "audio", "audio": str(audio) if isinstance(audio, (str, Path)) else audio} if audio_end is not None: item |= {"audio_start": 0.0, "audio_end": audio_end} user_content.append(item) @@ -77,8 +83,11 @@ def encode_qwen_omni( chat_text = components.tokenizer.apply_chat_template([conversation], add_generation_prompt=False, tokenize=False) audios = None - if audio_path is not None: - audios = [librosa.load(str(audio_path), sr=AUDIO_SAMPLE_RATE, offset=0.0, duration=audio_end)[0]] + if isinstance(audio, (str, Path)): + audios = [librosa.load(str(audio), sr=AUDIO_SAMPLE_RATE, offset=0.0, duration=audio_end)[0]] + elif audio is not None: + # The same cut as `librosa.load(duration=...)` makes. + audios = [audio if audio_end is None else audio[: int(AUDIO_SAMPLE_RATE * audio_end)]] inputs = components.tokenizer( text=chat_text, @@ -167,8 +176,9 @@ def inputs(self) -> list[InputParam]: InputParam( "audio", required=True, - type_hint=str | Path, - description="The speech track that drives the generation: a path to an audio file.", + type_hint=str | Path | torch.Tensor | np.ndarray, + description="The speech track that drives the generation: a path to an audio file, or a mono " + f"float waveform in [-1, 1] at {AUDIO_SAMPLE_RATE} Hz shaped (samples,).", ), InputParam.template( "height", @@ -212,6 +222,8 @@ def check_inputs(components, block_state): raise ValueError("`image` is required: it is the reference frame to animate") if block_state.audio is None: raise ValueError("`audio` is required: it is the speech track that drives the generation") + if not isinstance(block_state.audio, (str, Path)): + audio_waveform(block_state.audio) 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.num_chunks is not None and block_state.num_chunks < 1: @@ -301,7 +313,12 @@ def inputs(self) -> list[InputParam]: description="The prompt for the unconditional pass of the guidance. Defaults to " "`DEFAULT_NEGATIVE_PROMPT`.", ), - InputParam("audio", required=True, type_hint=str | Path, description="The speech track."), + InputParam( + "audio", + required=True, + type_hint=str | Path | torch.Tensor | np.ndarray, + description=f"The speech track: a path to an audio file, or a mono float waveform at {AUDIO_SAMPLE_RATE} Hz.", + ), InputParam("target_seconds", required=True, type_hint=float, description="Seconds of speech to encode."), ] @@ -342,16 +359,22 @@ def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState device = components._execution_device target_seconds = block_state.target_seconds - padded_audio = None - try: - padded_audio = pad_with_silence(block_state.audio, target_seconds) - audio_for_encode = padded_audio if padded_audio is not None else Path(block_state.audio) + if isinstance(block_state.audio, (str, Path)): + padded_audio = None + try: + padded_audio = pad_with_silence(block_state.audio, target_seconds) + audio_for_encode = padded_audio if padded_audio is not None else Path(block_state.audio) + prompt_embeds, audio_embeds = encode_qwen_omni( + components, block_state.prompt, audio_for_encode, target_seconds, device + ) + finally: + if padded_audio is not None: + padded_audio.unlink(missing_ok=True) + else: + waveform = pad_waveform_with_silence(audio_waveform(block_state.audio), target_seconds) prompt_embeds, audio_embeds = encode_qwen_omni( - components, block_state.prompt, audio_for_encode, target_seconds, device + components, block_state.prompt, waveform, target_seconds, device ) - finally: - if padded_audio is not None: - padded_audio.unlink(missing_ok=True) if audio_embeds is None: raise ValueError("speech-to-video requires an audio track") 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 index e32a8f86d8b8..d3bfbae50b6d 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v.py @@ -51,8 +51,9 @@ class KandinskyS2VAutoBlocks(SequentialPipelineBlocks): Inputs: image (`Image | str | Path`): The reference frame to animate: a PIL image or a path to one. - audio (`str | Path`): - The speech track that drives the generation: a path to an audio file. + audio (`str | Path | Tensor | ndarray`): + The speech track that drives the generation: a path to an audio file, or a mono float waveform in [-1, 1] + at 16000 Hz shaped (samples,). height (`int`, *optional*): The height in pixels of the generated video. Defaults to the 512 bucket closest to the reference frame. width (`int`, *optional*): 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 index c3d567277e7a..ccadb7ca3ea0 100644 --- 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 @@ -51,8 +51,9 @@ class KandinskyS2VDistilledAutoBlocks(SequentialPipelineBlocks): Inputs: image (`Image | str | Path`): The reference frame to animate: a PIL image or a path to one. - audio (`str | Path`): - The speech track that drives the generation: a path to an audio file. + audio (`str | Path | Tensor | ndarray`): + The speech track that drives the generation: a path to an audio file, or a mono float waveform in [-1, 1] + at 16000 Hz shaped (samples,). height (`int`, *optional*): The height in pixels of the generated video. Defaults to the 512 bucket closest to the reference frame. width (`int`, *optional*): diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py index b6eb301baefb..03c35b8b6d6b 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py @@ -76,11 +76,37 @@ def select_resolution( return min(allowed, key=lambda res: abs(aspect - res[1] / res[0])) -def audio_duration(path: str | Path) -> float: - info = sf.info(str(path)) +def audio_waveform(audio: torch.Tensor | np.ndarray) -> np.ndarray: + """A speech track passed as an array, as the speech encoder takes it: mono float32 at `AUDIO_SAMPLE_RATE`.""" + 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` as an array must be a non-empty mono float waveform in [-1, 1] at {AUDIO_SAMPLE_RATE} Hz, " + f"shaped (samples,); got shape {tuple(audio.shape)} and dtype {audio.dtype}" + ) + if isinstance(audio, torch.Tensor): + return audio.detach().cpu().float().numpy() + return audio.astype(np.float32) + + +def audio_duration(audio: str | Path | torch.Tensor | np.ndarray) -> float: + if not isinstance(audio, (str, Path)): + return audio_waveform(audio).shape[0] / AUDIO_SAMPLE_RATE + info = sf.info(str(audio)) return info.frames / info.samplerate +def pad_waveform_with_silence(waveform: np.ndarray, target_duration: float) -> np.ndarray: + """`pad_with_silence` for a waveform at `AUDIO_SAMPLE_RATE`, in memory.""" + 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 pad_with_silence(path: str | Path, target_duration: float) -> Path | None: """Append silence so the track covers ``target_duration``; ``None`` if it already does.""" original = audio_duration(path) 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 index 06bf41b3a484..143abad27b4e 100644 --- a/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py +++ b/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py @@ -24,7 +24,7 @@ import torch from PIL import Image -from diffusers import ClassifierFreeGuidance, ModularPipeline +from diffusers import ClassifierFreeGuidance from diffusers.modular_pipelines import ( KandinskyS2VAutoBlocks, KandinskyS2VDistilledAutoBlocks, @@ -33,6 +33,7 @@ ) from diffusers.modular_pipelines.kandinsky_s2v.before_denoise import KandinskyS2VPrepareStep +from ...testing_utils import numpy_cosine_similarity_distance, require_accelerator, torch_device from ..testing_utils import ( BaseModularPipelineTesterConfig, ModularLoadingTesterMixin, @@ -50,18 +51,41 @@ NUM_LATENT_FRAMES = 8 -@functools.lru_cache(maxsize=1) -def dummy_audio() -> str: - """A second of a plain tone; the speech is read from a file, not a tensor.""" +def tone(seconds: int) -> list[int]: + """A plain tone as 16-bit samples.""" + return [int(8000 * math.sin(i / 24)) for i in range(seconds * AUDIO_SAMPLE_RATE)] + + +@functools.lru_cache(maxsize=2) +def dummy_audio(seconds: int = 1) -> str: + """The tone as a wav file.""" path = Path(tempfile.mkdtemp()) / "speech.wav" with wave.open(str(path), "w") as handle: handle.setnchannels(1) handle.setsampwidth(2) handle.setframerate(AUDIO_SAMPLE_RATE) - handle.writeframes(b"".join(struct.pack(" torch.Tensor: + """The same tone as a mono waveform, scaled as an audio loader reads 16-bit samples.""" + return torch.tensor(tone(seconds), dtype=torch.float32) / 32768 + + +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 @@ -100,9 +124,9 @@ def test_inference_batch_single_identical(self): def test_num_images_per_prompt(self): pass - @pytest.mark.skip(reason="The normalisations do their arithmetic in fp32, which casting the whole model bypasses") + @require_accelerator def test_float16_inference(self): - pass + 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.""" @@ -123,6 +147,30 @@ def test_rejects_speech_past_the_encoder_limit(self): with pytest.raises(ValueError, match="truncates"): pipe(**inputs, output="videos") + @pytest.mark.parametrize("seconds", [1, 2]) + def test_audio_as_a_waveform(self, seconds): + """A mono waveform gives the same video as a file holding the same samples, padded or cut to the chunks.""" + pipe = self.get_pipeline() + inputs = {**self.get_dummy_inputs(), "num_chunks": 1} + from_file = pipe(**{**inputs, "audio": dummy_audio(seconds)}, output="videos") + inputs = {**self.get_dummy_inputs(), "num_chunks": 1} + from_waveform = pipe(**{**inputs, "audio": dummy_waveform(seconds)}, output="videos") + assert torch.equal(from_file, from_waveform) + + @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") + def test_rejects_a_single_step(self): """One step leaves an empty Euler schedule, which would decode the initial noise.""" pipe = self.get_pipeline() @@ -155,7 +203,6 @@ class TestKandinskyS2VModularPipelineMemory(KandinskyS2VModularPipelineTesterCon "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", - "test_float16_inference": "The normalisations do their arithmetic in fp32, which casting the whole model bypasses", } @@ -182,9 +229,9 @@ def test_inference_batch_single_identical(self): def test_num_images_per_prompt(self): pass - @pytest.mark.skip(reason=FAST_SKIPS["test_float16_inference"]) + @require_accelerator def test_float16_inference(self): - pass + check_float16_inference(self) def test_guidance_is_off(self): assert self.get_pipeline().guider.guidance_scale == 1.0 @@ -207,17 +254,3 @@ class TestKandinskyS2VDistilledModularPipelineMemory( KandinskyS2VDistilledModularPipelineTesterConfig, ModularMemoryTesterMixin ): pass - - -@pytest.mark.parametrize( - ("repo", "pipeline_class", "blocks_class"), - [ - (TINY_REPO, KandinskyS2VModularPipeline, KandinskyS2VAutoBlocks), - (TINY_DISTILLED_REPO, KandinskyS2VDistilledModularPipeline, KandinskyS2VDistilledAutoBlocks), - ], -) -def test_from_pretrained_picks_the_variant(repo, pipeline_class, blocks_class): - """What a user does: the checkpoint alone decides the pipeline and blocks, base or distilled.""" - pipe = ModularPipeline.from_pretrained(repo) - assert type(pipe) is pipeline_class - assert type(pipe.blocks) is blocks_class From 07a55586e1825497f7cbddcf0aa15f220a63851c Mon Sep 17 00:00:00 2001 From: Nikolay Semenov Date: Thu, 29 Oct 2026 15:10:31 +0300 Subject: [PATCH 05/12] Address review: no file inputs, standalone top-level blocks - Take `image` as a PIL image and `audio` as a 16 kHz waveform; the docs load them. - Pass components to the encode helpers explicitly. - Drop the guider from the text encoder: standalone it encodes a negative prompt only when passed; in the pipeline it follows the guider. - Remove the input step: the text encoder plans the chunks, the VAE encoder picks the resolution. - Group prepare and the chunk loop into one `denoise` step. - Distilled blocks: no guider and no negative prompt. --- docs/source/en/api/pipelines/kandinsky_s2v.md | 26 +- .../kandinsky_s2v/before_denoise.py | 107 ++++- .../kandinsky_s2v/denoise.py | 70 ++-- .../kandinsky_s2v/encoders.py | 376 ++++++++---------- .../modular_blocks_kandinsky_s2v.py | 103 ++++- .../modular_blocks_kandinsky_s2v_distilled.py | 97 ++++- .../kandinsky_s2v/modular_pipeline.py | 49 +-- .../test_modular_pipeline_kandinsky_s2v.py | 76 ++-- 8 files changed, 513 insertions(+), 391 deletions(-) diff --git a/docs/source/en/api/pipelines/kandinsky_s2v.md b/docs/source/en/api/pipelines/kandinsky_s2v.md index 15c9fd7d9de4..6fdb8b8a190e 100644 --- a/docs/source/en/api/pipelines/kandinsky_s2v.md +++ b/docs/source/en/api/pipelines/kandinsky_s2v.md @@ -45,21 +45,25 @@ The distilled checkpoints configure `PiflowScheduler` with `shift=10.0` and `n_g `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, and the distilled ones -at 1.0, because guidance is already part of the student — they reject any other. +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 export_to_video +from diffusers.utils import export_to_video, load_image pipe = ModularPipeline.from_pretrained("kandinskylab/Kandinsky-S2V-1.0-Lite-distill-Diffusers") pipe.load_components(dtype=torch.bfloat16) pipe.to("cuda") -video = pipe(image="reference.png", audio="speech.wav", output="videos")[0] +image = load_image("reference.png") +audio, _ = librosa.load("speech.wav", sr=16000) # a mono waveform at 16 kHz + +video = pipe(image=image, audio=audio, output="videos")[0] export_to_video(video, "avatar.mp4", fps=25) ``` @@ -70,16 +74,20 @@ The Pro checkpoints do not fit an 80 GB GPU next to the 7B speech and text encod 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 export_to_video +from diffusers.utils import export_to_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") -video = pipe(image="reference.png", audio="speech.wav", output="videos")[0] +image = load_image("reference.png") +audio, _ = librosa.load("speech.wav", sr=16000) + +video = pipe(image=image, audio=audio, output="videos")[0] export_to_video(video, "avatar.mp4", fps=25) ``` @@ -91,8 +99,10 @@ so it applies to every component. ## Notes -- `audio` is a path to an audio file, or a mono float waveform in [-1, 1] at 16 kHz, as a - tensor or an array shaped `(samples,)`. +- `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. diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py b/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py index ce5fd329e53a..331097319e96 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py @@ -107,14 +107,38 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("reference_latents", required=True, type_hint=torch.Tensor), - InputParam("audio_embeds", required=True, type_hint=torch.Tensor), - InputParam("num_chunks", required=True, type_hint=int), - InputParam("num_latent_frames", default=32, type_hint=int), + 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.template("negative_prompt_embeds"), - InputParam("pooled_prompt_embeds", required=True, type_hint=torch.Tensor), - InputParam("negative_pooled_prompt_embeds", type_hint=torch.Tensor), + InputParam( + "pooled_prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="CLIP pooled prompt embedding.", + ), + InputParam( + "negative_pooled_prompt_embeds", + type_hint=torch.Tensor, + description="CLIP pooled negative prompt embedding.", + ), ] @property @@ -128,18 +152,50 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam("audio_chunks", type_hint=list, description="Speech features of each chunk."), OutputParam.template("prompt_embeds"), OutputParam.template("negative_prompt_embeds"), - OutputParam("pooled_prompt_embeds", type_hint=torch.Tensor), - OutputParam("negative_pooled_prompt_embeds", type_hint=torch.Tensor), - OutputParam("text_rope_pos", type_hint=torch.Tensor), - OutputParam("negative_text_rope_pos", type_hint=torch.Tensor), - OutputParam("visual_rope_pos", type_hint=list), - OutputParam("reference_rope_pos", type_hint=list), + 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("text_rope_pos", type_hint=torch.Tensor, description="RoPE positions of the prompt tokens."), + OutputParam( + "negative_text_rope_pos", + type_hint=torch.Tensor, + description="RoPE positions of the negative 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(components, block_state): + needed = sum(components.transformer.config.frame_buckets) + if block_state.num_latent_frames < needed: + raise ValueError( + f"`num_latent_frames` must be >= {needed}, the FramePack memory window " + f"sum(frame_buckets), got {block_state.num_latent_frames}" + ) + + @staticmethod + def prepare_negative_embeds(block_state, dtype: torch.dtype, device: torch.device) -> None: + block_state.negative_text_rope_pos = None + if (block_state.negative_prompt_embeds is None) != (block_state.negative_pooled_prompt_embeds is None): + raise ValueError("`negative_prompt_embeds` and `negative_pooled_prompt_embeds` must be passed together.") + 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=device + ) + @torch.no_grad() def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) + self.check_inputs(components, block_state) device = components._execution_device config = components.transformer.config dtype = components.transformer.dtype @@ -155,15 +211,7 @@ def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState 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.negative_text_rope_pos = None - if (block_state.negative_prompt_embeds is None) != (block_state.negative_pooled_prompt_embeds is None): - raise ValueError("`negative_prompt_embeds` and `negative_pooled_prompt_embeds` must be passed together.") - 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=device - ) + self.prepare_negative_embeds(block_state, dtype, device) block_state.audio_chunks = split_audio_embeds(block_state.audio_embeds.to(dtype), block_state.num_chunks) @@ -183,3 +231,20 @@ def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState self.set_block_state(state, block_state) return components, state + + +class KandinskyS2VDistilledPrepareStep(KandinskyS2VPrepareStep): + model_name = "kandinsky-s2v-distilled" + + @property + def inputs(self) -> list[InputParam]: + return [param for param in super().inputs if not param.name.startswith("negative_")] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [param for param in super().intermediate_outputs if not param.name.startswith("negative_")] + + @staticmethod + def prepare_negative_embeds(block_state, dtype: torch.dtype, device: torch.device) -> None: + # No guidance on the distilled checkpoints, so there are no negative embeddings to prepare. + pass diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py b/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py index 46b283ef9369..ec5e36a90375 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py @@ -180,17 +180,30 @@ def inputs(self) -> list[InputParam]: InputParam.template("num_inference_steps"), ] + @staticmethod + def shared_transformer_kwargs(components: KandinskyS2VModularPipeline, block_state: BlockState) -> dict: + """The transformer inputs of a chunk that are the same in every pass.""" + config = components.transformer.config + device = components._execution_device + memory = block_state.memory + return { + "reference_hidden_states": block_state.reference_hidden_states.unsqueeze(0), + "memory_hidden_states": tuple(m.unsqueeze(0).to(components.transformer.dtype) for m in reversed(memory)), + "visual_rope_pos": block_state.visual_rope_pos, + "reference_rope_pos": block_state.reference_rope_pos, + "memory_rope_pos": tuple(reversed(memory_rope_pos(config, memory, device))), + "audio_rope_pos": audio_rope_pos(block_state.chunk_audio.shape[0], block_state.num_latent_frames, device), + "scale_factor": block_state.scale_factor, + "return_dict": False, + } + @torch.no_grad() def __call__( self, components: KandinskyS2VModularPipeline, block_state: BlockState, k: int ) -> tuple[KandinskyS2VModularPipeline, BlockState]: - config = components.transformer.config transformer_dtype = components.transformer.dtype - device = components._execution_device latents = block_state.latents chunk_audio = block_state.chunk_audio - num_latent_frames = block_state.num_latent_frames - memory = block_state.memory # Only these differ between the conditional and the unconditional pass. negative_embeds = block_state.negative_prompt_embeds @@ -204,16 +217,7 @@ def __call__( "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 = { - "reference_hidden_states": block_state.reference_hidden_states.unsqueeze(0), - "memory_hidden_states": tuple(m.unsqueeze(0).to(transformer_dtype) for m in reversed(memory)), - "visual_rope_pos": block_state.visual_rope_pos, - "reference_rope_pos": block_state.reference_rope_pos, - "memory_rope_pos": tuple(reversed(memory_rope_pos(config, memory, device))), - "audio_rope_pos": audio_rope_pos(chunk_audio.shape[0], num_latent_frames, device), - "scale_factor": block_state.scale_factor, - "return_dict": False, - } + shared_kwargs = self.shared_transformer_kwargs(components, block_state) timesteps = block_state.timesteps for i, t in enumerate(tqdm(timesteps, desc=f"chunk {k + 1}", leave=False)): @@ -246,17 +250,18 @@ def __call__( 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), - ComponentSpec( - "guider", - ClassifierFreeGuidance, - config=FrozenDict({"guidance_scale": 1.0}), - default_creation_method="from_config", - ), ] @property @@ -264,16 +269,33 @@ 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]: - if components.requires_unconditional_embeds: - # Guidance on top of a distilled student over-guides it and doubles the cost. - raise ValueError("guidance must be off with the distilled checkpoints: CFG is distilled into the student") - return super().__call__(components, block_state, k) + transformer_dtype = components.transformer.dtype + latents = block_state.latents + shared_kwargs = self.shared_transformer_kwargs(components, block_state) + + 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): diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py index da215da947e1..70717ffc0166 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py @@ -13,7 +13,6 @@ # limitations under the License. import math -from pathlib import Path from typing import Any import numpy as np @@ -26,11 +25,8 @@ Qwen2_5OmniThinkerForConditionalGeneration, ) -from ...configuration_utils import FrozenDict -from ...guiders import ClassifierFreeGuidance from ...models import AutoencoderKLHunyuanVideo -from ...models.transformers.transformer_kandinsky_s2v import KandinskyS2VTransformer3DModel -from ...utils import is_librosa_available, is_torchvision_available, logging +from ...utils import is_torchvision_available, logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import ( @@ -42,37 +38,30 @@ SYSTEM_PROMPT, TEMPORAL_COMPRESSION_RATIO, KandinskyS2VModularPipeline, - audio_duration, audio_waveform, pad_waveform_with_silence, - pad_with_silence, select_resolution, ) -if is_librosa_available(): - import librosa - if is_torchvision_available(): from torchvision import transforms logger = logging.get_logger(__name__) # pylint: disable=invalid-name -def encode_qwen_omni( - components: KandinskyS2VModularPipeline, +def encode_qwen_omni( # noqa: PLR0913 + tokenizer: Qwen2_5OmniProcessor, + text_encoder: Qwen2_5OmniThinkerForConditionalGeneration, prompt: str, - audio: str | Path | np.ndarray | None, + 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. - - ``audio`` is a path to an audio file or a mono waveform at `AUDIO_SAMPLE_RATE`. - """ + """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: - item: dict[str, Any] = {"type": "audio", "audio": str(audio) if isinstance(audio, (str, Path)) else audio} + item: dict[str, Any] = {"type": "audio", "audio": audio} if audio_end is not None: item |= {"audio_start": 0.0, "audio_end": audio_end} user_content.append(item) @@ -81,22 +70,20 @@ def encode_qwen_omni( {"role": "user", "content": user_content}, ] - chat_text = components.tokenizer.apply_chat_template([conversation], add_generation_prompt=False, tokenize=False) + chat_text = tokenizer.apply_chat_template([conversation], add_generation_prompt=False, tokenize=False) audios = None - if isinstance(audio, (str, Path)): - audios = [librosa.load(str(audio), sr=AUDIO_SAMPLE_RATE, offset=0.0, duration=audio_end)[0]] - elif audio is not None: - # The same cut as `librosa.load(duration=...)` makes. + if audio is not None: + # Cut to the seconds the chunks cover. audios = [audio if audio_end is None else audio[: int(AUDIO_SAMPLE_RATE * audio_end)]] - inputs = components.tokenizer( + inputs = tokenizer( text=chat_text, audio=audios, return_tensors="pt", padding=True, use_audio_in_video=False, ) - inputs = inputs.to(device).to(components.text_encoder.dtype) + inputs = inputs.to(device).to(text_encoder.dtype) # The speech features are the audio tower's output, which the Thinker consumes internally; a hook on the tower # reads them out of the same forward that produces the text hidden states. @@ -105,26 +92,31 @@ def encode_qwen_omni( def hook(_module, _inputs, output): captured["features"] = getattr(output, "last_hidden_state", output) - handle = components.text_encoder.audio_tower.register_forward_hook(hook) + handle = text_encoder.audio_tower.register_forward_hook(hook) try: - output = components.text_encoder(**inputs, output_hidden_states=True, return_dict=True) + output = text_encoder(**inputs, output_hidden_states=True, return_dict=True) finally: handle.remove() - prompt_embeds = slice_user_text(components, output.hidden_states[-1], inputs) + prompt_embeds = slice_user_text(tokenizer, text_encoder, output.hidden_states[-1], inputs) return prompt_embeds, captured.get("features") -def slice_user_text(components: KandinskyS2VModularPipeline, hidden_states: torch.Tensor, inputs: Any) -> torch.Tensor: +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 = components.tokenizer.tokenizer.convert_tokens_to_ids("<|im_start|>") - audio_start_token_id = components.text_encoder.config.audio_start_token_id - user_role_len = len(components.tokenizer.tokenizer("user\n", add_special_tokens=False).input_ids) + 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 @@ -135,9 +127,9 @@ def slice_user_text(components: KandinskyS2VModularPipeline, hidden_states: torc def encode_clip( - components: KandinskyS2VModularPipeline, prompt: str, device: torch.device, max_length: int = 77 + tokenizer_2: CLIPTokenizer, text_encoder_2: CLIPTextModel, prompt: str, device: torch.device, max_length: int = 77 ) -> torch.Tensor: - inputs = components.tokenizer_2( + inputs = tokenizer_2( [prompt], max_length=max_length, truncation=True, @@ -145,48 +137,49 @@ def encode_clip( padding="max_length", return_tensors="pt", ).to(device) - return components.text_encoder_2(**inputs)["pooler_output"] + return text_encoder_2(**inputs)["pooler_output"] -class KandinskyS2VInputStep(ModularPipelineBlocks): +class KandinskyS2VTextAudioEncoderStep(ModularPipelineBlocks): model_name = "kandinsky-s2v" @property def description(self) -> str: return ( - "Input step that checks the reference frame and the speech track, picks the resolution bucket and plans " - "the chunks: how many are needed to cover the track, and how many seconds of speech they take." + "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("transformer", KandinskyS2VTransformer3DModel), + ComponentSpec("text_encoder_2", CLIPTextModel), + ComponentSpec("tokenizer_2", CLIPTokenizer), ] @property def inputs(self) -> list[InputParam]: return [ InputParam.template( - "image", - type_hint=PIL.Image.Image | str | Path, - description="The reference frame to animate: a PIL image or a path to one.", + "prompt", + required=False, + description="The prompt to guide the generation. Defaults to the reference prompt, `DEFAULT_PROMPT`.", + ), + 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=str | Path | torch.Tensor | np.ndarray, - description="The speech track that drives the generation: a path to an audio file, or a mono " - f"float waveform in [-1, 1] at {AUDIO_SAMPLE_RATE} Hz shaped (samples,).", - ), - 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`." + 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", @@ -206,49 +199,48 @@ def inputs(self) -> list[InputParam]: @property def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam("reference_image", type_hint=PIL.Image.Image, description="The reference frame, in RGB."), - OutputParam("height", type_hint=int, description="The height in pixels of the generated video."), - OutputParam("width", type_hint=int, description="The width in pixels of the generated video."), - 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."), + 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( - "target_seconds", type_hint=float, description="Seconds of speech the chunks cover, padded if needed." + "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(components, block_state): - if block_state.image is None: - raise ValueError("`image` is required: it is the reference frame to animate") + def check_inputs(block_state): + if block_state.prompt is None: + block_state.prompt = DEFAULT_PROMPT + 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") if block_state.audio is None: raise ValueError("`audio` is required: it is the speech track that drives the generation") - if not isinstance(block_state.audio, (str, Path)): - audio_waveform(block_state.audio) - 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") + audio_waveform(block_state.audio) 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}") - needed = sum(components.transformer.config.frame_buckets) - if block_state.num_latent_frames < needed: + if block_state.num_latent_frames < 1 or block_state.fps <= 0: raise ValueError( - f"`num_latent_frames` must be >= {needed}, the FramePack memory window " - f"sum(frame_buckets), got {block_state.num_latent_frames}" + 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(components, block_state) - - image = block_state.image - if not isinstance(image, PIL.Image.Image): - image = PIL.Image.open(image) - block_state.reference_image = image.convert("RGB") - if block_state.height is None: - block_state.height, block_state.width = select_resolution(block_state.reference_image) - + @staticmethod + def encode_prompt_and_speech(components: KandinskyS2VModularPipeline, block_state) -> None: + """Plan the chunks the speech track needs, then encode the prompt and the speech for them.""" + device = components._execution_device + waveform = audio_waveform(block_state.audio) chunk_seconds = ((block_state.num_latent_frames - 1) * TEMPORAL_COMPRESSION_RATIO + 1) / block_state.fps - block_state.track_seconds = audio_duration(block_state.audio) + block_state.track_seconds = waveform.shape[0] / AUDIO_SAMPLE_RATE needed = max(1, math.ceil(block_state.track_seconds / chunk_seconds)) if block_state.num_chunks is not None and block_state.num_chunks > needed: # The video is trimmed to the track, so these are computed and then dropped. @@ -258,126 +250,28 @@ def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState ) if block_state.num_chunks is None: block_state.num_chunks = needed - block_state.target_seconds = block_state.num_chunks * chunk_seconds + target_seconds = block_state.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 = components.tokenizer.feature_extractor max_seconds = extractor.n_samples / extractor.sampling_rate - if block_state.target_seconds > max_seconds + 0.01: + if target_seconds > max_seconds + 0.01: raise ValueError( - f"{block_state.num_chunks} chunks need {block_state.target_seconds:g} s of speech, but the speech " - f"encoder truncates at {max_seconds:g} s; trim the track or pass fewer `num_chunks`" + f"{block_state.num_chunks} chunks need {target_seconds:g} s of speech, but the speech encoder " + f"truncates at {max_seconds:g} s; trim the track or pass fewer `num_chunks`" ) - self.set_block_state(state, block_state) - return components, state - - -class KandinskyS2VTextAudioEncoderStep(ModularPipelineBlocks): - model_name = "kandinsky-s2v" - - @property - def description(self) -> str: - return ( - "Encoder step that turns the prompt and the speech track 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 only when the guider needs an unconditional pass." + prompt_embeds, audio_embeds = encode_qwen_omni( + components.tokenizer, + components.text_encoder, + block_state.prompt, + pad_waveform_with_silence(waveform, target_seconds), + target_seconds, + device, ) - - @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), - ComponentSpec( - "guider", - ClassifierFreeGuidance, - config=FrozenDict({"guidance_scale": 3.0}), - default_creation_method="from_config", - ), - ] - - @property - def inputs(self) -> list[InputParam]: - return [ - InputParam.template( - "prompt", - required=False, - description="The prompt to guide the generation. Defaults to the reference prompt, `DEFAULT_PROMPT`.", - ), - InputParam.template( - "negative_prompt", - description="The prompt for the unconditional pass of the guidance. Defaults to " - "`DEFAULT_NEGATIVE_PROMPT`.", - ), - InputParam( - "audio", - required=True, - type_hint=str | Path | torch.Tensor | np.ndarray, - description=f"The speech track: a path to an audio file, or a mono float waveform at {AUDIO_SAMPLE_RATE} Hz.", - ), - InputParam("target_seconds", required=True, type_hint=float, description="Seconds of speech to encode."), - ] - - @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." - ), - ] - - @staticmethod - def check_inputs(components, block_state): - if block_state.prompt is None: - block_state.prompt = DEFAULT_PROMPT - if block_state.negative_prompt is None: - block_state.negative_prompt = DEFAULT_NEGATIVE_PROMPT - if not isinstance(block_state.prompt, str) or not isinstance(block_state.negative_prompt, str): - raise ValueError("`prompt` and `negative_prompt` must be strings; 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") - if components.requires_unconditional_embeds and not block_state.negative_prompt.strip(): - raise ValueError("`negative_prompt` must not be empty when classifier-free guidance is enabled") - - @torch.no_grad() - def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: - block_state = self.get_block_state(state) - self.check_inputs(components, block_state) - device = components._execution_device - target_seconds = block_state.target_seconds - - if isinstance(block_state.audio, (str, Path)): - padded_audio = None - try: - padded_audio = pad_with_silence(block_state.audio, target_seconds) - audio_for_encode = padded_audio if padded_audio is not None else Path(block_state.audio) - prompt_embeds, audio_embeds = encode_qwen_omni( - components, block_state.prompt, audio_for_encode, target_seconds, device - ) - finally: - if padded_audio is not None: - padded_audio.unlink(missing_ok=True) - else: - waveform = pad_waveform_with_silence(audio_waveform(block_state.audio), target_seconds) - prompt_embeds, audio_embeds = encode_qwen_omni( - components, block_state.prompt, waveform, target_seconds, device - ) if audio_embeds is None: - raise ValueError("speech-to-video requires an audio track") - + raise ValueError("the speech encoder returned no features for the speech track") # 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) @@ -389,14 +283,39 @@ def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState block_state.prompt_embeds = prompt_embeds block_state.audio_embeds = audio_embeds - block_state.pooled_prompt_embeds = encode_clip(components, block_state.prompt, device) + block_state.pooled_prompt_embeds = encode_clip( + components.tokenizer_2, components.text_encoder_2, block_state.prompt, device + ) + + @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") + + self.encode_prompt_and_speech(components, block_state) + device = components._execution_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 - if components.requires_unconditional_embeds: + 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, block_state.negative_prompt, None, None, device + 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 ) - block_state.negative_pooled_prompt_embeds = encode_clip(components, block_state.negative_prompt, device) self.set_block_state(state, block_state) return components, state @@ -406,16 +325,27 @@ class KandinskyS2VDistilledTextAudioEncoderStep(KandinskyS2VTextAudioEncoderStep model_name = "kandinsky-s2v-distilled" @property - def expected_components(self) -> list[ComponentSpec]: - return [ - *super().expected_components[:-1], - ComponentSpec( - "guider", - ClassifierFreeGuidance, - config=FrozenDict({"guidance_scale": 1.0}), - default_creation_method="from_config", - ), - ] + 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) + self.encode_prompt_and_speech(components, block_state) + self.set_block_state(state, block_state) + return components, state class KandinskyS2VVaeEncoderStep(ModularPipelineBlocks): @@ -423,7 +353,10 @@ class KandinskyS2VVaeEncoderStep(ModularPipelineBlocks): @property def description(self) -> str: - return "VAE encoder step that encodes the reference frame into a scaled latent, resized and center-cropped." + 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]: @@ -432,11 +365,15 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam( - "reference_image", required=True, type_hint=PIL.Image.Image, description="The reference frame." + 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("height", required=True), - InputParam.template("width", required=True), InputParam.template("generator"), ] @@ -450,11 +387,24 @@ def intermediate_outputs(self) -> list[OutputParam]: ) ] + @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") + @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, height, width = block_state.reference_image, block_state.height, block_state.width + image = block_state.image.convert("RGB") + height, width = block_state.height, block_state.width + if height is None: + height, width = select_resolution(image) scale = max(width / image.width, height / image.height) resized = transforms.functional.resize( 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 index d3bfbae50b6d..496f14363ddd 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v.py @@ -18,19 +18,79 @@ from .before_denoise import KandinskyS2VPrepareStep from .decoders import KandinskyS2VDecodeStep from .denoise import KandinskyS2VChunkDenoiseStep -from .encoders import KandinskyS2VInputStep, KandinskyS2VTextAudioEncoderStep, KandinskyS2VVaeEncoderStep +from .encoders import KandinskyS2VTextAudioEncoderStep, KandinskyS2VVaeEncoderStep logger = logging.get_logger(__name__) # pylint: disable=invalid-name +KandinskyS2VCoreDenoiseBlocks = InsertableDict( + [ + ("prepare", KandinskyS2VPrepareStep()), + ("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. + negative_prompt_embeds (`Tensor`, *optional*): + negative text embeddings used to guide the image generation. Can be generated from text_encoder step. + pooled_prompt_embeds (`Tensor`): + CLIP pooled prompt embedding. + 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. + """ + + 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.")] + + @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( [ - ("input", KandinskyS2VInputStep()), ("text_encoder", KandinskyS2VTextAudioEncoderStep()), ("vae_encoder", KandinskyS2VVaeEncoderStep()), - ("prepare", KandinskyS2VPrepareStep()), - ("denoise", KandinskyS2VChunkDenoiseStep()), + ("denoise", KandinskyS2VCoreDenoiseStep()), ("decode", KandinskyS2VDecodeStep()), ] ) @@ -43,21 +103,20 @@ class KandinskyS2VAutoBlocks(SequentialPipelineBlocks): a speech track, generated chunk by chunk. Components: - tokenizer (`Qwen2_5OmniProcessor`) transformer (`KandinskyS2VTransformer3DModel`) text_encoder - (`Qwen2_5OmniThinkerForConditionalGeneration`) text_encoder_2 (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) - guider (`ClassifierFreeGuidance`) vae (`AutoencoderKLHunyuanVideo`) scheduler - (`FlowMatchEulerDiscreteScheduler`) video_processor (`VideoProcessor`) + text_encoder (`Qwen2_5OmniThinkerForConditionalGeneration`) tokenizer (`Qwen2_5OmniProcessor`) text_encoder_2 + (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) vae (`AutoencoderKLHunyuanVideo`) transformer + (`KandinskyS2VTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) video_processor (`VideoProcessor`) Inputs: - image (`Image | str | Path`): - The reference frame to animate: a PIL image or a path to one. - audio (`str | Path | Tensor | ndarray`): - The speech track that drives the generation: a path to an audio file, or a mono float waveform in [-1, 1] - at 16000 Hz shaped (samples,). - 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`. + prompt (`str`, *optional*): + The prompt to guide the generation. Defaults to the reference prompt, `DEFAULT_PROMPT`. + 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): @@ -65,10 +124,12 @@ class KandinskyS2VAutoBlocks(SequentialPipelineBlocks): 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. - prompt (`str`, *optional*): - The prompt to guide the generation. Defaults to the reference prompt, `DEFAULT_PROMPT`. - negative_prompt (`str`, *optional*): - The prompt for the unconditional pass of the guidance. Defaults to `DEFAULT_NEGATIVE_PROMPT`. + 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): 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 index ccadb7ca3ea0..1383d6171514 100644 --- 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 @@ -15,22 +15,78 @@ from ...utils import logging from ..modular_pipeline import SequentialPipelineBlocks from ..modular_pipeline_utils import InsertableDict, OutputParam -from .before_denoise import KandinskyS2VPrepareStep +from .before_denoise import KandinskyS2VDistilledPrepareStep from .decoders import KandinskyS2VDecodeStep from .denoise import KandinskyS2VDistilledChunkDenoiseStep -from .encoders import KandinskyS2VDistilledTextAudioEncoderStep, KandinskyS2VInputStep, KandinskyS2VVaeEncoderStep +from .encoders import KandinskyS2VDistilledTextAudioEncoderStep, KandinskyS2VVaeEncoderStep logger = logging.get_logger(__name__) # pylint: disable=invalid-name +KandinskyS2VDistilledCoreDenoiseBlocks = InsertableDict( + [ + ("prepare", KandinskyS2VDistilledPrepareStep()), + ("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. + """ + + 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.")] + + @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( [ - ("input", KandinskyS2VInputStep()), ("text_encoder", KandinskyS2VDistilledTextAudioEncoderStep()), ("vae_encoder", KandinskyS2VVaeEncoderStep()), - ("prepare", KandinskyS2VPrepareStep()), - ("denoise", KandinskyS2VDistilledChunkDenoiseStep()), + ("denoise", KandinskyS2VDistilledCoreDenoiseStep()), ("decode", KandinskyS2VDecodeStep()), ] ) @@ -43,21 +99,16 @@ class KandinskyS2VDistilledAutoBlocks(SequentialPipelineBlocks): segments per chunk and no guidance, which is distilled into the student. Components: - tokenizer (`Qwen2_5OmniProcessor`) transformer (`KandinskyS2VTransformer3DModel`) text_encoder - (`Qwen2_5OmniThinkerForConditionalGeneration`) text_encoder_2 (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) - guider (`ClassifierFreeGuidance`) vae (`AutoencoderKLHunyuanVideo`) scheduler (`PiflowScheduler`) - video_processor (`VideoProcessor`) + text_encoder (`Qwen2_5OmniThinkerForConditionalGeneration`) tokenizer (`Qwen2_5OmniProcessor`) text_encoder_2 + (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) vae (`AutoencoderKLHunyuanVideo`) transformer + (`KandinskyS2VTransformer3DModel`) scheduler (`PiflowScheduler`) video_processor (`VideoProcessor`) Inputs: - image (`Image | str | Path`): - The reference frame to animate: a PIL image or a path to one. - audio (`str | Path | Tensor | ndarray`): - The speech track that drives the generation: a path to an audio file, or a mono float waveform in [-1, 1] - at 16000 Hz shaped (samples,). - 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`. + prompt (`str`, *optional*): + The prompt to guide the generation. Defaults to the reference prompt, `DEFAULT_PROMPT`. + 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): @@ -65,10 +116,12 @@ class KandinskyS2VDistilledAutoBlocks(SequentialPipelineBlocks): 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. - prompt (`str`, *optional*): - The prompt to guide the generation. Defaults to the reference prompt, `DEFAULT_PROMPT`. - negative_prompt (`str`, *optional*): - The prompt for the unconditional pass of the guidance. Defaults to `DEFAULT_NEGATIVE_PROMPT`. + 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): diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py index 03c35b8b6d6b..dc766cccb549 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py @@ -12,23 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os -import tempfile -from pathlib import Path - import numpy as np import PIL.Image import torch from ...models.transformers.transformer_kandinsky_s2v import audio_group_lengths -from ...utils import is_librosa_available, logging +from ...utils import logging from ..modular_pipeline import ModularPipeline -# soundfile is a dependency of librosa. -if is_librosa_available(): - import soundfile as sf - logger = logging.get_logger(__name__) # pylint: disable=invalid-name SYSTEM_PROMPT = ( @@ -78,6 +70,11 @@ def select_resolution( def audio_waveform(audio: torch.Tensor | np.ndarray) -> np.ndarray: """A speech track passed as an array, as the speech encoder takes it: mono float32 at `AUDIO_SAMPLE_RATE`.""" + 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) ) @@ -91,15 +88,8 @@ def audio_waveform(audio: torch.Tensor | np.ndarray) -> np.ndarray: return audio.astype(np.float32) -def audio_duration(audio: str | Path | torch.Tensor | np.ndarray) -> float: - if not isinstance(audio, (str, Path)): - return audio_waveform(audio).shape[0] / AUDIO_SAMPLE_RATE - info = sf.info(str(audio)) - return info.frames / info.samplerate - - def pad_waveform_with_silence(waveform: np.ndarray, target_duration: float) -> np.ndarray: - """`pad_with_silence` for a waveform at `AUDIO_SAMPLE_RATE`, in memory.""" + """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 @@ -107,27 +97,6 @@ def pad_waveform_with_silence(waveform: np.ndarray, target_duration: float) -> n return np.concatenate([waveform, np.zeros(pad_samples, dtype=waveform.dtype)]) -def pad_with_silence(path: str | Path, target_duration: float) -> Path | None: - """Append silence so the track covers ``target_duration``; ``None`` if it already does.""" - original = audio_duration(path) - if target_duration <= original + 0.01: - return None - - data, sample_rate = sf.read(str(path)) - pad_samples = int((target_duration - original) * sample_rate) - shape = (pad_samples,) if data.ndim == 1 else (pad_samples, data.shape[1]) - padded = np.concatenate([data, np.zeros(shape, dtype=data.dtype)]) - - handle, temp_path = tempfile.mkstemp(suffix=".wav") - os.close(handle) - try: - sf.write(temp_path, padded, sample_rate) - except BaseException: - os.unlink(temp_path) - raise - return Path(temp_path) - - 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 @@ -164,8 +133,8 @@ class KandinskyS2VDistilledModularPipeline(KandinskyS2VModularPipeline): """ A ModularPipeline for the pi-Flow distilled Kandinsky S2V 1.0 checkpoints. - Guidance is distilled into the student, so the guider defaults to `guidance_scale=1.0` and a single forward pass - runs per segment. + Guidance is distilled into the student, so its blocks have no guider and no negative prompt, and a single forward + pass runs per segment. """ 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 index 143abad27b4e..22914f2e93a1 100644 --- a/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py +++ b/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py @@ -13,18 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -import functools import math -import struct -import tempfile -import wave -from pathlib import Path import pytest import torch from PIL import Image -from diffusers import ClassifierFreeGuidance from diffusers.modular_pipelines import ( KandinskyS2VAutoBlocks, KandinskyS2VDistilledAutoBlocks, @@ -32,6 +26,7 @@ KandinskyS2VModularPipeline, ) from diffusers.modular_pipelines.kandinsky_s2v.before_denoise import KandinskyS2VPrepareStep +from diffusers.modular_pipelines.kandinsky_s2v.encoders import KandinskyS2VTextAudioEncoderStep from ...testing_utils import numpy_cosine_similarity_distance, require_accelerator, torch_device from ..testing_utils import ( @@ -51,26 +46,9 @@ NUM_LATENT_FRAMES = 8 -def tone(seconds: int) -> list[int]: - """A plain tone as 16-bit samples.""" - return [int(8000 * math.sin(i / 24)) for i in range(seconds * AUDIO_SAMPLE_RATE)] - - -@functools.lru_cache(maxsize=2) -def dummy_audio(seconds: int = 1) -> str: - """The tone as a wav file.""" - path = Path(tempfile.mkdtemp()) / "speech.wav" - with wave.open(str(path), "w") as handle: - handle.setnchannels(1) - handle.setsampwidth(2) - handle.setframerate(AUDIO_SAMPLE_RATE) - handle.writeframes(b"".join(struct.pack(" torch.Tensor: - """The same tone as a mono waveform, scaled as an audio loader reads 16-bit samples.""" - return torch.tensor(tone(seconds), dtype=torch.float32) / 32768 + """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): @@ -98,7 +76,7 @@ class KandinskyS2VModularPipelineTesterConfig(BaseModularPipelineTesterConfig): def get_dummy_inputs(self, seed=0): return { "image": Image.new("RGB", (64, 64), color="red"), - "audio": dummy_audio(), + "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", @@ -147,15 +125,28 @@ def test_rejects_speech_past_the_encoder_limit(self): with pytest.raises(ValueError, match="truncates"): pipe(**inputs, output="videos") - @pytest.mark.parametrize("seconds", [1, 2]) - def test_audio_as_a_waveform(self, seconds): - """A mono waveform gives the same video as a file holding the same samples, padded or cut to the chunks.""" + 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} + latents = pipe(**{**inputs, "output_type": "latent"}, output="videos") + assert latents.shape[1] == NUM_LATENT_FRAMES + + def test_rejects_file_paths(self): + """The pipeline does not read files; the errors say how to load them.""" pipe = self.get_pipeline() - inputs = {**self.get_dummy_inputs(), "num_chunks": 1} - from_file = pipe(**{**inputs, "audio": dummy_audio(seconds)}, output="videos") - inputs = {**self.get_dummy_inputs(), "num_chunks": 1} - from_waveform = pipe(**{**inputs, "audio": dummy_waveform(seconds)}, output="videos") - assert torch.equal(from_file, from_waveform) + with pytest.raises(ValueError, match="load_image"): + pipe(**{**self.get_dummy_inputs(), "image": "reference.png"}, output="videos") + with pytest.raises(ValueError, match="librosa.load"): + pipe(**{**self.get_dummy_inputs(), "audio": "speech.wav"}, output="videos") + + 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", @@ -213,6 +204,11 @@ class KandinskyS2VDistilledModularPipelineTesterConfig(KandinskyS2VModularPipeli pipeline_blocks_class = KandinskyS2VDistilledAutoBlocks pretrained_model_name_or_path = TINY_DISTILLED_REPO + def get_dummy_inputs(self, seed=0): + inputs = super().get_dummy_inputs(seed) + inputs.pop("negative_prompt") + return inputs + class TestKandinskyS2VDistilledModularPipelineFast( KandinskyS2VDistilledModularPipelineTesterConfig, ModularPipelineTesterMixin @@ -233,15 +229,11 @@ def test_num_images_per_prompt(self): def test_float16_inference(self): check_float16_inference(self) - def test_guidance_is_off(self): - assert self.get_pipeline().guider.guidance_scale == 1.0 - - def test_rejects_guidance(self): - """Guidance on top of the student would over-guide it, quietly and at twice the cost.""" + def test_has_no_guidance(self): + """Guidance is distilled into the student: no guider, and no negative prompt to give it.""" pipe = self.get_pipeline() - pipe.update_components(guider=ClassifierFreeGuidance(guidance_scale=4.0)) - with pytest.raises(ValueError, match="distilled into the student"): - pipe(**self.get_dummy_inputs(), output="videos") + assert "guider" not in pipe.components + assert "negative_prompt" not in {param.name for param in pipe.blocks.inputs} class TestKandinskyS2VDistilledModularPipelineLoading( From 18cce4d0c562dae5978b37d8ceb78c41ef253684 Mon Sep 17 00:00:00 2001 From: Nikolay Semenov Date: Thu, 29 Oct 2026 17:20:15 +0300 Subject: [PATCH 06/12] Docs: save the video with the original speech track --- docs/source/en/api/pipelines/kandinsky_s2v.md | 27 +++++++++++++------ 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/docs/source/en/api/pipelines/kandinsky_s2v.md b/docs/source/en/api/pipelines/kandinsky_s2v.md index 6fdb8b8a190e..349e619be13e 100644 --- a/docs/source/en/api/pipelines/kandinsky_s2v.md +++ b/docs/source/en/api/pipelines/kandinsky_s2v.md @@ -54,22 +54,26 @@ no guider and no `negative_prompt`, because guidance is already part of the stud import librosa import torch from diffusers import ModularPipeline -from diffusers.utils import export_to_video, load_image +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 +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] -export_to_video(video, "avatar.mp4", fps=25) +encode_video( + video, + fps=25, + output_path="avatar.mp4", + audio=torch.from_numpy(speech)[None], + audio_sample_rate=sample_rate, +) ``` -`export_to_video` writes the frames only; add the speech track with any muxer, for -example `ffmpeg -i avatar.mp4 -i speech.wav -c:v copy -c:a aac -shortest out.mp4`. - 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`: @@ -77,7 +81,7 @@ when every component is on the device. Offload them through a `ComponentsManager import librosa import torch from diffusers import ComponentsManager, ModularPipeline -from diffusers.utils import export_to_video, load_image +from diffusers.utils import encode_video, load_image manager = ComponentsManager() pipe = ModularPipeline.from_pretrained("kandinskylab/Kandinsky-S2V-1.0-Pro-Diffusers", components_manager=manager) @@ -86,9 +90,16 @@ 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] -export_to_video(video, "avatar.mp4", fps=25) +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. From 40db4b8bd3b05b41a05b81c3614b9a3e1042d894 Mon Sep 17 00:00:00 2001 From: Nikolay Semenov Date: Fri, 30 Oct 2026 13:48:34 +0300 Subject: [PATCH 07/12] Address review: validation-only checks, no latent output, video_processor - Inline the grouped attention into the processor. - `prompt` defaults to DEFAULT_PROMPT in its signature; check_inputs only validates. - Add auto docstrings to the distilled text encoder and prepare steps. - Drop `output_type="latent"`; the latents are in the state as `latent_chunks`. - Preprocess the reference frame with the video processor. - Remove three tests the review asked to drop. --- .../transformers/transformer_kandinsky_s2v.py | 43 ++----- .../kandinsky_s2v/before_denoise.py | 46 ++++++++ .../kandinsky_s2v/decoders.py | 26 ++-- .../kandinsky_s2v/denoise.py | 9 +- .../kandinsky_s2v/encoders.py | 111 ++++++++++++++---- .../modular_blocks_kandinsky_s2v.py | 32 +++-- .../modular_blocks_kandinsky_s2v_distilled.py | 30 +++-- .../kandinsky_s2v/modular_pipeline.py | 20 ---- .../test_modular_pipeline_kandinsky_s2v.py | 38 +----- 9 files changed, 215 insertions(+), 140 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_kandinsky_s2v.py b/src/diffusers/models/transformers/transformer_kandinsky_s2v.py index 7f7796263092..b79bcacfc2d1 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky_s2v.py +++ b/src/diffusers/models/transformers/transformer_kandinsky_s2v.py @@ -81,36 +81,6 @@ def pad_groups(x: Tensor, lengths: Tensor, max_len: int) -> tuple[Tensor, Tensor return out, valid -def grouped_attention( # noqa: PLR0913 - q: Tensor, - k: Tensor, - v: Tensor, - q_lengths: Tensor, - kv_lengths: Tensor, - backend=None, - parallel_config=None, -) -> Tensor: - """Block-diagonal attention: query group ``i`` attends exactly to key group ``i``. - - Groups are uneven, so they are padded into a batch axis and masked rather than reshaped. Inputs and output are - packed ``(S, heads, head_dim)``. - """ - q_pad, q_valid = pad_groups(q, q_lengths, int(q_lengths.max())) - kv_max = int(kv_lengths.max()) - k_pad, kv_valid = pad_groups(k, kv_lengths, kv_max) - v_pad, _ = pad_groups(v, kv_lengths, kv_max) - - out = dispatch_attention_fn( - q_pad, - k_pad, - v_pad, - attn_mask=kv_valid[:, None, None, :], - backend=backend, - parallel_config=parallel_config, - ) - return out[q_valid] - - class KandinskyS2VTimeEmbeddings(nn.Module): def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): super().__init__() @@ -249,15 +219,20 @@ def __call__( # noqa: PLR0913 if q_lengths is not None: if kv_lengths is None: raise ValueError("kv_lengths is required when q_lengths is given") - hidden_states = grouped_attention( + # Block-diagonal attention: query group i attends exactly to key group i. The groups are uneven, so they + # are padded into a batch axis and masked rather than reshaped. + query, q_valid = pad_groups(query, q_lengths, int(q_lengths.max())) + kv_max = int(kv_lengths.max()) + key, kv_valid = pad_groups(key, kv_lengths, kv_max) + value, _ = pad_groups(value, kv_lengths, kv_max) + hidden_states = dispatch_attention_fn( query, key, value, - q_lengths, - kv_lengths, + attn_mask=kv_valid[:, None, None, :], backend=self._attention_backend, parallel_config=self._parallel_config, - ) + )[q_valid] else: hidden_states = dispatch_attention_fn( query.unsqueeze(0), diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py b/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py index 331097319e96..89f3dfdfa0cd 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py @@ -233,7 +233,53 @@ def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState return components, state +# auto_docstring class KandinskyS2VDistilledPrepareStep(KandinskyS2VPrepareStep): + """ + 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. + + Components: + transformer (`KandinskyS2VTransformer3DModel`) + + 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. + + Outputs: + reference_hidden_states (`Tensor`): + Reference latent, cropped and cast. + latent_height (`int`): + Height of the latents. + latent_width (`int`): + Width of the latents. + audio_chunks (`list`): + Speech features of each chunk. + prompt_embeds (`Tensor`): + The prompt embeddings. + pooled_prompt_embeds (`Tensor`): + CLIP pooled prompt embedding. + text_rope_pos (`Tensor`): + RoPE positions of the prompt tokens. + visual_rope_pos (`list`): + Per-axis RoPE positions of a chunk. + reference_rope_pos (`list`): + Per-axis RoPE positions of the reference frame. + memory (`list`): + FramePack memory in [1x, 2x, 4x] order. + """ + model_name = "kandinsky-s2v-distilled" @property diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/decoders.py b/src/diffusers/modular_pipelines/kandinsky_s2v/decoders.py index 991d8198a5f3..dfeee4db3151 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/decoders.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/decoders.py @@ -66,13 +66,16 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam("latent_chunks", required=True, type_hint=list), + 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", - description="Output format: 'pil', 'np', 'pt' or 'latent'. Latents are returned whole, untrimmed.", - ), + InputParam.template("output_type"), ] @property @@ -83,15 +86,10 @@ def intermediate_outputs(self) -> list[OutputParam]: def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - if block_state.output_type == "latent": - block_state.videos = torch.cat(block_state.latent_chunks, dim=0).unsqueeze(0) - else: - 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 - ) + 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 index ec5e36a90375..d03911657f6d 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py @@ -348,7 +348,14 @@ def loop_inputs(self) -> list[InputParam]: @property def loop_intermediate_outputs(self) -> list[OutputParam]: - return [OutputParam("latent_chunks", type_hint=list, description="Denoised latents of each chunk.")] + 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: diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py index 70717ffc0166..17592fb48d41 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py @@ -25,8 +25,10 @@ Qwen2_5OmniThinkerForConditionalGeneration, ) +from ...configuration_utils import FrozenDict from ...models import AutoencoderKLHunyuanVideo -from ...utils import is_torchvision_available, logging +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 ( @@ -38,15 +40,11 @@ SYSTEM_PROMPT, TEMPORAL_COMPRESSION_RATIO, KandinskyS2VModularPipeline, - audio_waveform, pad_waveform_with_silence, select_resolution, ) -if is_torchvision_available(): - from torchvision import transforms - logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -165,9 +163,7 @@ def expected_components(self) -> list[ComponentSpec]: def inputs(self) -> list[InputParam]: return [ InputParam.template( - "prompt", - required=False, - description="The prompt to guide the generation. Defaults to the reference prompt, `DEFAULT_PROMPT`.", + "prompt", required=False, default=DEFAULT_PROMPT, description="The prompt to guide the generation." ), InputParam.template( "negative_prompt", @@ -216,16 +212,25 @@ def intermediate_outputs(self) -> list[OutputParam]: @staticmethod def check_inputs(block_state): - if block_state.prompt is None: - block_state.prompt = DEFAULT_PROMPT 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") - if block_state.audio is None: - raise ValueError("`audio` is required: it is the speech track that drives the generation") - audio_waveform(block_state.audio) + 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: @@ -238,7 +243,10 @@ def check_inputs(block_state): def encode_prompt_and_speech(components: KandinskyS2VModularPipeline, block_state) -> None: """Plan the chunks the speech track needs, then encode the prompt and the speech for them.""" device = components._execution_device - waveform = audio_waveform(block_state.audio) + audio = block_state.audio + waveform = ( + audio.detach().cpu().float().numpy() if isinstance(audio, torch.Tensor) else audio.astype(np.float32) + ) chunk_seconds = ((block_state.num_latent_frames - 1) * TEMPORAL_COMPRESSION_RATIO + 1) / block_state.fps block_state.track_seconds = waveform.shape[0] / AUDIO_SAMPLE_RATE needed = max(1, math.ceil(block_state.track_seconds / chunk_seconds)) @@ -321,7 +329,52 @@ def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState 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 @@ -360,7 +413,15 @@ def description(self) -> str: @property def expected_components(self) -> list[ComponentSpec]: - return [ComponentSpec("vae", AutoencoderKLHunyuanVideo)] + 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]: @@ -395,6 +456,10 @@ def check_inputs(block_state): ) 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: @@ -406,17 +471,17 @@ def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState 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 = transforms.functional.resize( - image, - (round(image.height * scale), round(image.width * scale)), - interpolation=transforms.InterpolationMode.BILINEAR, + resized = image.resize( + (round(image.width * scale), round(image.height * scale)), resample=PIL.Image.Resampling.BILINEAR ) - cropped = transforms.functional.center_crop(resized, (height, width)) + left, top = round((resized.width - width) / 2), round((resized.height - height) / 2) + cropped = resized.crop((left, top, left + width, top + height)) - tensor = transforms.ToTensor()(cropped).unsqueeze(0) - tensor = tensor.to(device=device, dtype=components.vae.dtype) * 2.0 - 1.0 - tensor = tensor.unsqueeze(2) # (1, 3, 1, H, W) + 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) 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 index 496f14363ddd..5d21abfdcb49 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v.py @@ -67,7 +67,7 @@ class KandinskyS2VCoreDenoiseStep(SequentialPipelineBlocks): Outputs: latent_chunks (`list`): - Denoised latents of each chunk. + Denoised latents of each chunk, each shaped (num_latent_frames, latent_height, latent_width, channels). """ model_name = "kandinsky-s2v" @@ -76,7 +76,14 @@ class KandinskyS2VCoreDenoiseStep(SequentialPipelineBlocks): @property def outputs(self): - return [OutputParam("latent_chunks", type_hint=list, description="Denoised latents of each chunk.")] + 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): @@ -104,13 +111,22 @@ class KandinskyS2VAutoBlocks(SequentialPipelineBlocks): Components: text_encoder (`Qwen2_5OmniThinkerForConditionalGeneration`) tokenizer (`Qwen2_5OmniProcessor`) text_encoder_2 - (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) vae (`AutoencoderKLHunyuanVideo`) transformer - (`KandinskyS2VTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider - (`ClassifierFreeGuidance`) video_processor (`VideoProcessor`) + (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) vae (`AutoencoderKLHunyuanVideo`) video_processor + (`VideoProcessor`) transformer (`KandinskyS2VTransformer3DModel`) scheduler + (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) Inputs: - prompt (`str`, *optional*): - The prompt to guide the generation. Defaults to the reference prompt, `DEFAULT_PROMPT`. + 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. @@ -137,7 +153,7 @@ class KandinskyS2VAutoBlocks(SequentialPipelineBlocks): 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' or 'latent'. Latents are returned whole, untrimmed. + Output format: 'pil', 'np', 'pt'. Outputs: videos (`list`): 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 index 1383d6171514..3abbf8814828 100644 --- 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 @@ -63,7 +63,7 @@ class KandinskyS2VDistilledCoreDenoiseStep(SequentialPipelineBlocks): Outputs: latent_chunks (`list`): - Denoised latents of each chunk. + Denoised latents of each chunk, each shaped (num_latent_frames, latent_height, latent_width, channels). """ model_name = "kandinsky-s2v-distilled" @@ -72,7 +72,14 @@ class KandinskyS2VDistilledCoreDenoiseStep(SequentialPipelineBlocks): @property def outputs(self): - return [OutputParam("latent_chunks", type_hint=list, description="Denoised latents of each chunk.")] + 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): @@ -100,12 +107,21 @@ class KandinskyS2VDistilledAutoBlocks(SequentialPipelineBlocks): Components: text_encoder (`Qwen2_5OmniThinkerForConditionalGeneration`) tokenizer (`Qwen2_5OmniProcessor`) text_encoder_2 - (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) vae (`AutoencoderKLHunyuanVideo`) transformer - (`KandinskyS2VTransformer3DModel`) scheduler (`PiflowScheduler`) video_processor (`VideoProcessor`) + (`CLIPTextModel`) tokenizer_2 (`CLIPTokenizer`) vae (`AutoencoderKLHunyuanVideo`) video_processor + (`VideoProcessor`) transformer (`KandinskyS2VTransformer3DModel`) scheduler (`PiflowScheduler`) Inputs: - prompt (`str`, *optional*): - The prompt to guide the generation. Defaults to the reference prompt, `DEFAULT_PROMPT`. + 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,). @@ -129,7 +145,7 @@ class KandinskyS2VDistilledAutoBlocks(SequentialPipelineBlocks): 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' or 'latent'. Latents are returned whole, untrimmed. + Output format: 'pil', 'np', 'pt'. Outputs: videos (`list`): diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py index dc766cccb549..b9de1558f409 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py @@ -68,26 +68,6 @@ def select_resolution( return min(allowed, key=lambda res: abs(aspect - res[1] / res[0])) -def audio_waveform(audio: torch.Tensor | np.ndarray) -> np.ndarray: - """A speech track passed as an array, as the speech encoder takes it: mono float32 at `AUDIO_SAMPLE_RATE`.""" - 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` as an array must be a non-empty mono float waveform in [-1, 1] at {AUDIO_SAMPLE_RATE} Hz, " - f"shaped (samples,); got shape {tuple(audio.shape)} and dtype {audio.dtype}" - ) - if isinstance(audio, torch.Tensor): - return audio.detach().cpu().float().numpy() - return audio.astype(np.float32) - - 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 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 index 22914f2e93a1..52ef5082bbcf 100644 --- a/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py +++ b/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py @@ -25,7 +25,6 @@ KandinskyS2VDistilledModularPipeline, KandinskyS2VModularPipeline, ) -from diffusers.modular_pipelines.kandinsky_s2v.before_denoise import KandinskyS2VPrepareStep from diffusers.modular_pipelines.kandinsky_s2v.encoders import KandinskyS2VTextAudioEncoderStep from ...testing_utils import numpy_cosine_similarity_distance, require_accelerator, torch_device @@ -111,9 +110,9 @@ def test_two_chunks_carry_memory(self): pipe = self.get_pipeline() inputs = self.get_dummy_inputs() inputs["num_chunks"] = 2 - latents = pipe(**{**inputs, "output_type": "latent"}, output="videos") - assert latents.shape[1] == 2 * NUM_LATENT_FRAMES - assert torch.isfinite(latents).all() + 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() @@ -129,16 +128,8 @@ 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} - latents = pipe(**{**inputs, "output_type": "latent"}, output="videos") - assert latents.shape[1] == NUM_LATENT_FRAMES - - def test_rejects_file_paths(self): - """The pipeline does not read files; the errors say how to load them.""" - pipe = self.get_pipeline() - with pytest.raises(ValueError, match="load_image"): - pipe(**{**self.get_dummy_inputs(), "image": "reference.png"}, output="videos") - with pytest.raises(ValueError, match="librosa.load"): - pipe(**{**self.get_dummy_inputs(), "audio": "speech.wav"}, output="videos") + 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.""" @@ -162,25 +153,6 @@ def test_rejects_a_bad_waveform(self, waveform): with pytest.raises(ValueError, match="mono float waveform"): pipe(**{**self.get_dummy_inputs(), "audio": waveform}, output="videos") - def test_rejects_a_single_step(self): - """One step leaves an empty Euler schedule, which would decode the initial noise.""" - pipe = self.get_pipeline() - with pytest.raises(ValueError, match="at least 2"): - pipe(**{**self.get_dummy_inputs(), "num_inference_steps": 1}, output="videos") - - def test_rejects_negative_embeds_without_pooled(self): - prepare_pipe = KandinskyS2VPrepareStep().init_pipeline(self.pretrained_model_name_or_path) - prepare_pipe.load_components() - with pytest.raises(ValueError, match="must be passed together"): - prepare_pipe( - reference_latents=torch.randn(1, 8, 8, 4), - audio_embeds=torch.randn(16, 16), - num_chunks=1, - prompt_embeds=torch.randn(3, 16), - pooled_prompt_embeds=torch.randn(16), - negative_prompt_embeds=torch.randn(3, 16), - ) - class TestKandinskyS2VModularPipelineLoading(KandinskyS2VModularPipelineTesterConfig, ModularLoadingTesterMixin): pass From 05a66e082ca4b9fa29048c7947e15e8ccdb59ba7 Mon Sep 17 00:00:00 2001 From: Nikolay Semenov Date: Sat, 31 Oct 2026 05:26:17 +0300 Subject: [PATCH 08/12] Address review: explicit arguments for the helpers - encode_prompt_and_speech and shared_transformer_kwargs take their components and values explicitly and return their results. - The negative prompt gets its own prepare block in the base pipeline; the distilled pipeline uses the shared prepare step. - check_inputs of the prepare steps take explicit values. --- .../kandinsky_s2v/before_denoise.py | 150 +++++++---------- .../kandinsky_s2v/denoise.py | 65 +++++--- .../kandinsky_s2v/encoders.py | 154 +++++++++++------- .../modular_blocks_kandinsky_s2v.py | 7 +- .../modular_blocks_kandinsky_s2v_distilled.py | 4 +- 5 files changed, 207 insertions(+), 173 deletions(-) diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py b/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py index 89f3dfdfa0cd..7c988704eee4 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/before_denoise.py @@ -127,18 +127,12 @@ def inputs(self) -> list[InputParam]: description="Latent frames per chunk; 32 gives 125 frames (5 s at 25 fps).", ), InputParam.template("prompt_embeds", required=True), - InputParam.template("negative_prompt_embeds"), InputParam( "pooled_prompt_embeds", required=True, type_hint=torch.Tensor, description="CLIP pooled prompt embedding.", ), - InputParam( - "negative_pooled_prompt_embeds", - type_hint=torch.Tensor, - description="CLIP pooled negative prompt embedding.", - ), ] @property @@ -151,19 +145,8 @@ def intermediate_outputs(self) -> list[OutputParam]: 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.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("text_rope_pos", type_hint=torch.Tensor, description="RoPE positions of the prompt tokens."), - OutputParam( - "negative_text_rope_pos", - type_hint=torch.Tensor, - description="RoPE positions of the negative 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." @@ -172,32 +155,19 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @staticmethod - def check_inputs(components, block_state): - needed = sum(components.transformer.config.frame_buckets) - if block_state.num_latent_frames < needed: + 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 >= {needed}, the FramePack memory window " - f"sum(frame_buckets), got {block_state.num_latent_frames}" - ) - - @staticmethod - def prepare_negative_embeds(block_state, dtype: torch.dtype, device: torch.device) -> None: - block_state.negative_text_rope_pos = None - if (block_state.negative_prompt_embeds is None) != (block_state.negative_pooled_prompt_embeds is None): - raise ValueError("`negative_prompt_embeds` and `negative_pooled_prompt_embeds` must be passed together.") - 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=device + 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) - self.check_inputs(components, block_state) - device = components._execution_device 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 @@ -211,7 +181,6 @@ def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState 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) - self.prepare_negative_embeds(block_state, dtype, device) block_state.audio_chunks = split_audio_embeds(block_state.audio_embeds.to(dtype), block_state.num_chunks) @@ -233,64 +202,65 @@ def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState return components, state -# auto_docstring -class KandinskyS2VDistilledPrepareStep(KandinskyS2VPrepareStep): - """ - 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. - - Components: - transformer (`KandinskyS2VTransformer3DModel`) - - 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. - - Outputs: - reference_hidden_states (`Tensor`): - Reference latent, cropped and cast. - latent_height (`int`): - Height of the latents. - latent_width (`int`): - Width of the latents. - audio_chunks (`list`): - Speech features of each chunk. - prompt_embeds (`Tensor`): - The prompt embeddings. - pooled_prompt_embeds (`Tensor`): - CLIP pooled prompt embedding. - text_rope_pos (`Tensor`): - RoPE positions of the prompt tokens. - visual_rope_pos (`list`): - Per-axis RoPE positions of a chunk. - reference_rope_pos (`list`): - Per-axis RoPE positions of the reference frame. - memory (`list`): - FramePack memory in [1x, 2x, 4x] order. - """ - - model_name = "kandinsky-s2v-distilled" +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 [param for param in super().inputs if not param.name.startswith("negative_")] + 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 [param for param in super().intermediate_outputs if not param.name.startswith("negative_")] + 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 prepare_negative_embeds(block_state, dtype: torch.dtype, device: torch.device) -> None: - # No guidance on the distilled checkpoints, so there are no negative embeddings to prepare. - pass + 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/denoise.py b/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py index d03911657f6d..3747be705379 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/denoise.py @@ -133,6 +133,30 @@ def __call__( 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" @@ -180,23 +204,6 @@ def inputs(self) -> list[InputParam]: InputParam.template("num_inference_steps"), ] - @staticmethod - def shared_transformer_kwargs(components: KandinskyS2VModularPipeline, block_state: BlockState) -> dict: - """The transformer inputs of a chunk that are the same in every pass.""" - config = components.transformer.config - device = components._execution_device - memory = block_state.memory - return { - "reference_hidden_states": block_state.reference_hidden_states.unsqueeze(0), - "memory_hidden_states": tuple(m.unsqueeze(0).to(components.transformer.dtype) for m in reversed(memory)), - "visual_rope_pos": block_state.visual_rope_pos, - "reference_rope_pos": block_state.reference_rope_pos, - "memory_rope_pos": tuple(reversed(memory_rope_pos(config, memory, device))), - "audio_rope_pos": audio_rope_pos(block_state.chunk_audio.shape[0], block_state.num_latent_frames, device), - "scale_factor": block_state.scale_factor, - "return_dict": False, - } - @torch.no_grad() def __call__( self, components: KandinskyS2VModularPipeline, block_state: BlockState, k: int @@ -217,7 +224,17 @@ def __call__( "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 = self.shared_transformer_kwargs(components, block_state) + 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)): @@ -278,7 +295,17 @@ def __call__( ) -> tuple[KandinskyS2VModularPipeline, BlockState]: transformer_dtype = components.transformer.dtype latents = block_state.latents - shared_kwargs = self.shared_transformer_kwargs(components, block_state) + 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): diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py index 17592fb48d41..eaef5e9c2347 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py @@ -48,7 +48,7 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -def encode_qwen_omni( # noqa: PLR0913 +def encode_qwen_omni( tokenizer: Qwen2_5OmniProcessor, text_encoder: Qwen2_5OmniThinkerForConditionalGeneration, prompt: str, @@ -138,6 +138,64 @@ def encode_clip( 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 + ) + if audio_embeds is None: + raise ValueError("the speech encoder returned no features for the speech track") + # 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" @@ -239,62 +297,6 @@ def check_inputs(block_state): f"{block_state.fps}" ) - @staticmethod - def encode_prompt_and_speech(components: KandinskyS2VModularPipeline, block_state) -> None: - """Plan the chunks the speech track needs, then encode the prompt and the speech for them.""" - device = components._execution_device - audio = block_state.audio - waveform = ( - audio.detach().cpu().float().numpy() if isinstance(audio, torch.Tensor) else audio.astype(np.float32) - ) - chunk_seconds = ((block_state.num_latent_frames - 1) * TEMPORAL_COMPRESSION_RATIO + 1) / block_state.fps - block_state.track_seconds = waveform.shape[0] / AUDIO_SAMPLE_RATE - needed = max(1, math.ceil(block_state.track_seconds / chunk_seconds)) - if block_state.num_chunks is not None and block_state.num_chunks > needed: - # The video is trimmed to the track, so these are computed and then dropped. - logger.warning( - f"{block_state.num_chunks - needed} chunk(s) past the end of the speech track are generated and " - f"then trimmed; pass num_chunks={needed} to save the time" - ) - if block_state.num_chunks is None: - block_state.num_chunks = needed - target_seconds = block_state.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 = components.tokenizer.feature_extractor - max_seconds = extractor.n_samples / extractor.sampling_rate - if target_seconds > max_seconds + 0.01: - raise ValueError( - f"{block_state.num_chunks} chunks need {target_seconds:g} s of speech, but the speech encoder " - f"truncates at {max_seconds:g} s; trim the track or pass fewer `num_chunks`" - ) - - prompt_embeds, audio_embeds = encode_qwen_omni( - components.tokenizer, - components.text_encoder, - block_state.prompt, - pad_waveform_with_silence(waveform, target_seconds), - target_seconds, - device, - ) - if audio_embeds is None: - raise ValueError("the speech encoder returned no features for the speech track") - # 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" - ) - - block_state.prompt_embeds = prompt_embeds - block_state.audio_embeds = audio_embeds - block_state.pooled_prompt_embeds = encode_clip( - components.tokenizer_2, components.text_encoder_2, block_state.prompt, device - ) - @torch.no_grad() def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) @@ -307,8 +309,25 @@ def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState 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") - self.encode_prompt_and_speech(components, block_state) 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. @@ -396,7 +415,24 @@ def intermediate_outputs(self) -> list[OutputParam]: def __call__(self, components: KandinskyS2VModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) self.check_inputs(block_state) - self.encode_prompt_and_speech(components, 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 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 index 5d21abfdcb49..099b4353a25f 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_blocks_kandinsky_s2v.py @@ -15,7 +15,7 @@ from ...utils import logging from ..modular_pipeline import SequentialPipelineBlocks from ..modular_pipeline_utils import InsertableDict, OutputParam -from .before_denoise import KandinskyS2VPrepareStep +from .before_denoise import KandinskyS2VPrepareNegativeStep, KandinskyS2VPrepareStep from .decoders import KandinskyS2VDecodeStep from .denoise import KandinskyS2VChunkDenoiseStep from .encoders import KandinskyS2VTextAudioEncoderStep, KandinskyS2VVaeEncoderStep @@ -27,6 +27,7 @@ KandinskyS2VCoreDenoiseBlocks = InsertableDict( [ ("prepare", KandinskyS2VPrepareStep()), + ("prepare_negative", KandinskyS2VPrepareNegativeStep()), ("chunk_loop", KandinskyS2VChunkDenoiseStep()), ] ) @@ -52,10 +53,10 @@ class KandinskyS2VCoreDenoiseStep(SequentialPipelineBlocks): 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. - negative_prompt_embeds (`Tensor`, *optional*): - negative 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*): 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 index 3abbf8814828..f8fbb72c15e6 100644 --- 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 @@ -15,7 +15,7 @@ from ...utils import logging from ..modular_pipeline import SequentialPipelineBlocks from ..modular_pipeline_utils import InsertableDict, OutputParam -from .before_denoise import KandinskyS2VDistilledPrepareStep +from .before_denoise import KandinskyS2VPrepareStep from .decoders import KandinskyS2VDecodeStep from .denoise import KandinskyS2VDistilledChunkDenoiseStep from .encoders import KandinskyS2VDistilledTextAudioEncoderStep, KandinskyS2VVaeEncoderStep @@ -26,7 +26,7 @@ KandinskyS2VDistilledCoreDenoiseBlocks = InsertableDict( [ - ("prepare", KandinskyS2VDistilledPrepareStep()), + ("prepare", KandinskyS2VPrepareStep()), ("chunk_loop", KandinskyS2VDistilledChunkDenoiseStep()), ] ) From 133a6de20a613e390445127a06a1ad78029abf29 Mon Sep 17 00:00:00 2001 From: Nikolay Semenov Date: Sat, 31 Oct 2026 16:28:03 +0300 Subject: [PATCH 09/12] Tests: declare the canonical latent form main checks that the denoise group leaves latents in one canonical form (test_latents_output_in_canonical_form). Point it at latent_chunks: one chunk per entry, channels last. --- .../kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py | 3 +++ 1 file changed, 3 insertions(+) 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 index 52ef5082bbcf..ca8515fa47ea 100644 --- a/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py +++ b/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py @@ -71,6 +71,9 @@ class KandinskyS2VModularPipelineTesterConfig(BaseModularPipelineTesterConfig): 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) def get_dummy_inputs(self, seed=0): return { From 9ba233beabbbbb2f0589c27c685c2044ae3bc056 Mon Sep 17 00:00:00 2001 From: Nikolay Semenov Date: Sun, 1 Nov 2026 00:36:02 +0300 Subject: [PATCH 10/12] Address self-review: compile-friendly speech attention, full model testers - Speech cross-attention keeps static shapes: the per-frame audio groups are computed from shapes as Python ints and gathered into a padded, masked batch, so the transformer compiles with fullgraph=True. One forward is bit-identical to the original model on all four checkpoints. - Drop FromOriginalModelMixin: single-file loading is not wired up. - Model tests compose the testers generate_model_tests.py emits (memory, compile incl. dynamic shapes, training, attention); skips name their reason. - Modular tests pin expected_workflow_defaults for base and distilled and add the guider tester for base. - Remove dead code: unread attributes, the batched visual-embedding branch, unused helper parameters. --- .../transformers/transformer_kandinsky_s2v.py | 143 +++++++----------- .../kandinsky_s2v/encoders.py | 15 +- .../kandinsky_s2v/modular_pipeline.py | 9 +- .../test_models_transformer_kandinsky_s2v.py | 73 +++++++-- .../test_modular_pipeline_kandinsky_s2v.py | 82 +++++++++- 5 files changed, 197 insertions(+), 125 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_kandinsky_s2v.py b/src/diffusers/models/transformers/transformer_kandinsky_s2v.py index b79bcacfc2d1..4a1a539b6b11 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky_s2v.py +++ b/src/diffusers/models/transformers/transformer_kandinsky_s2v.py @@ -20,11 +20,11 @@ from typing import Any import torch -import torch.nn.functional as F # noqa: N812 +import torch.nn.functional as F from torch import Tensor, nn from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ...loaders import PeftAdapterMixin from ...utils import logging from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn @@ -51,13 +51,7 @@ def apply_rotary(x: Tensor, rope: Tensor) -> Tensor: return (rope.float() * x_).sum(dim=-1).reshape(*x.shape).to(dtype=x.dtype) -def split_lengths(total: int, groups: int, device: torch.device) -> Tensor: - """Split ``total`` items into ``groups`` with boundaries at ``total * i // groups``.""" - bounds = torch.arange(groups + 1, device=device) * total // groups - return torch.diff(bounds) - - -def audio_group_lengths(num_audio_tokens: int, duration: int, device: torch.device) -> Tensor: +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}") @@ -67,18 +61,24 @@ def audio_group_lengths(num_audio_tokens: int, duration: int, device: torch.devi raise ValueError( f"{num_audio_tokens} audio tokens cannot cover {duration} latent frames: every frame needs at least one" ) - rest = split_lengths(num_audio_tokens - 1, duration - 1, device) - first = torch.ones(1, dtype=rest.dtype, device=device) - return torch.cat([first, rest]) + # 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 pad_groups(x: Tensor, lengths: Tensor, max_len: int) -> tuple[Tensor, Tensor]: - """Scatter a packed ``(S, ...)`` tensor into ``(G, max_len, ...)`` plus a validity mask.""" - positions = torch.arange(max_len, device=x.device) - valid = positions[None, :] < lengths[:, None] - out = x.new_zeros(lengths.shape[0], max_len, *x.shape[1:]) - out[valid] = x - return out, valid +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): @@ -86,8 +86,6 @@ 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}") - self.model_dim = model_dim - self.max_period = max_period # A plain attribute rather than a buffer: it must survive a meta-device build. self.freqs = get_freqs(model_dim // 2, max_period) @@ -120,7 +118,7 @@ def forward(self, embed: Tensor) -> Tensor: class KandinskyS2VVisualEmbeddings(nn.Module): - """Patch projection for ``[T,H,W,C]`` or ``[B,T,H,W,C]`` visual tokens.""" + """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__() @@ -128,35 +126,13 @@ def __init__(self, visual_dim: int, model_dim: int, patch_size: tuple[int, int, self.in_layer = nn.Linear(math.prod(patch_size) * visual_dim, model_dim) def forward(self, x: Tensor) -> Tensor: - batched = bool(x.ndim == 5) - offset = 1 if batched else 0 - if not batched and x.ndim != 4: - raise ValueError("visual input must have shape (T,H,W,C) or (B,T,H,W,C)") - - shape = x.shape - duration, height, width, channels = shape[offset:] + duration, height, width, channels = x.shape p_t, p_h, p_w = self.patch_size - if batched: - x = ( - x.view( - shape[0], - duration // p_t, - p_t, - height // p_h, - p_h, - width // p_w, - p_w, - channels, - ) - .permute(0, 1, 3, 5, 2, 4, 6, 7) - .flatten(4, 7) - ) - else: - 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) - ) + 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) @@ -193,14 +169,14 @@ 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__( # noqa: PLR0913 + def __call__( self, attn: KandinskyS2VAttention, hidden_states: Tensor, encoder_hidden_states: Tensor | None = None, rotary_emb: Tensor | None = None, - q_lengths: Tensor | None = None, - kv_lengths: 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] @@ -216,23 +192,21 @@ def __call__( # noqa: PLR0913 query = apply_rotary(query, rotary_emb).type_as(query) key = apply_rotary(key, rotary_emb).type_as(key) - if q_lengths is not None: - if kv_lengths is None: - raise ValueError("kv_lengths is required when q_lengths is given") - # Block-diagonal attention: query group i attends exactly to key group i. The groups are uneven, so they - # are padded into a batch axis and masked rather than reshaped. - query, q_valid = pad_groups(query, q_lengths, int(q_lengths.max())) - kv_max = int(kv_lengths.max()) - key, kv_valid = pad_groups(key, kv_lengths, kv_max) - value, _ = pad_groups(value, kv_lengths, kv_max) + 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_valid[:, None, None, :], + attn_mask=kv_mask[:, None, None, :], backend=self._attention_backend, parallel_config=self._parallel_config, - )[q_valid] + ).flatten(0, 1) else: hidden_states = dispatch_attention_fn( query.unsqueeze(0), @@ -315,9 +289,6 @@ def forward(self, visual_embed: Tensor, time_embed: Tensor, time_embed_idx: Tens class KandinskyS2VRoPE1D(nn.Module): def __init__(self, dim: int, max_pos: int = 1024, max_period: float = 10000.0): super().__init__() - self.dim = dim - self.max_pos = max_pos - self.max_period = max_period 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. @@ -342,9 +313,6 @@ def __init__( max_period: float = 10000.0, ): super().__init__() - self.axes_dims = axes_dims - self.max_pos = max_pos - self.max_period = max_period self.offsets_t = max_pos[0] # Plain attributes rather than buffers, so `.to(torch.bfloat16)` leaves the tables in fp32. @@ -434,7 +402,7 @@ def forward(self, x: Tensor, rope: Tensor) -> Tensor: class KandinskyS2VTransformerDecoderBlock(nn.Module): """Visual block; ``enable_audio`` adds the frame-aligned speech adapter.""" - def __init__( # noqa: PLR0913 + def __init__( self, model_dim: int, time_dim: int, @@ -459,7 +427,7 @@ def __init__( # noqa: PLR0913 self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.feed_forward = KandinskyS2VFeedForward(model_dim, ff_dim) - def forward( # noqa: PLR0913 + def forward( self, visual_embed: Tensor, text_embed: Tensor, @@ -468,8 +436,8 @@ def forward( # noqa: PLR0913 rope: Tensor, time_embed_idx: Tensor, num_visual_tokens: int, - visual_group_lengths: Tensor, - audio_group_lengths: Tensor, + 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) @@ -500,8 +468,8 @@ def forward( # noqa: PLR0913 out = self.cross_attention_v2a( self.cross_attention_norm_v2a(noisy.float()).type_as(noisy), encoder_hidden_states=audio_embed, - q_lengths=visual_group_lengths, - kv_lengths=audio_group_lengths, + kv_index=audio_index, + kv_mask=audio_mask, ) visual_embed = torch.cat([head, noisy + out], dim=0) @@ -520,7 +488,6 @@ class KandinskyS2VTransformer3DModel( ModelMixin, ConfigMixin, PeftAdapterMixin, - FromOriginalModelMixin, CacheMixin, AttentionMixin, ): @@ -567,7 +534,7 @@ class KandinskyS2VTransformer3DModel( _supports_gradient_checkpointing = True @register_to_config - def __init__( # noqa: PLR0913 + def __init__( self, in_visual_dim: int = 16, out_visual_dim: int = 16, @@ -593,13 +560,8 @@ def __init__( # noqa: PLR0913 ): super().__init__() head_dim = sum(axes_dims) - self.in_visual_dim = in_visual_dim self.model_dim = model_dim - self.patch_size = patch_size - self.frame_buckets = frame_buckets self.enable_framepack = enable_framepack - self.index_rope_ref_image = index_rope_ref_image - self.audio_inject_layers = tuple(audio_inject_layers) self.time_embeddings = KandinskyS2VTimeEmbeddings(model_dim, time_dim) self.text_embeddings = KandinskyS2VGeneralEmbeddings(in_text_dim, model_dim) @@ -658,7 +620,7 @@ def __init__( # noqa: PLR0913 time_dim, ff_dim, head_dim, - enable_audio=(i in self.audio_inject_layers), + enable_audio=(i in audio_inject_layers), ) for i in range(num_visual_blocks) ] @@ -667,7 +629,7 @@ def __init__( # noqa: PLR0913 self.out_layer = KandinskyS2VOutLayer(model_dim, time_dim, out_visual_dim, patch_size) self.gradient_checkpointing = False - def forward( # noqa: PLR0913, PLR0915 + def forward( self, hidden_states: Tensor, encoder_hidden_states: Tensor, @@ -791,8 +753,7 @@ def forward( # noqa: PLR0913, PLR0915 ] ) - visual_groups = split_lengths(num_visual_tokens, duration, device) - audio_groups = audio_group_lengths(audio_embed.shape[0], duration, 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: @@ -805,8 +766,8 @@ def forward( # noqa: PLR0913, PLR0915 visual_rope_full, time_embed_idx_full, num_visual_tokens, - visual_groups, - audio_groups, + audio_index, + audio_mask, ) else: visual_embed_full = block( @@ -817,8 +778,8 @@ def forward( # noqa: PLR0913, PLR0915 visual_rope_full, time_embed_idx_full, num_visual_tokens, - visual_groups, - audio_groups, + audio_index, + audio_mask, ) visual_embed = visual_embed_full[-num_visual_tokens:].reshape(*visual_shape, self.model_dim) diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py index eaef5e9c2347..4128fd0baa11 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py @@ -59,10 +59,7 @@ def encode_qwen_omni( """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: - item: dict[str, Any] = {"type": "audio", "audio": audio} - if audio_end is not None: - item |= {"audio_start": 0.0, "audio_end": audio_end} - user_content.append(item) + 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}, @@ -72,7 +69,7 @@ def encode_qwen_omni( audios = None if audio is not None: # Cut to the seconds the chunks cover. - audios = [audio if audio_end is None else audio[: int(AUDIO_SAMPLE_RATE * audio_end)]] + audios = [audio[: int(AUDIO_SAMPLE_RATE * audio_end)]] inputs = tokenizer( text=chat_text, @@ -88,7 +85,7 @@ def encode_qwen_omni( captured: dict[str, torch.Tensor] = {} def hook(_module, _inputs, output): - captured["features"] = getattr(output, "last_hidden_state", output) + captured["features"] = output.last_hidden_state handle = text_encoder.audio_tower.register_forward_hook(hook) try: @@ -125,11 +122,11 @@ def slice_user_text( def encode_clip( - tokenizer_2: CLIPTokenizer, text_encoder_2: CLIPTextModel, prompt: str, device: torch.device, max_length: int = 77 + tokenizer_2: CLIPTokenizer, text_encoder_2: CLIPTextModel, prompt: str, device: torch.device ) -> torch.Tensor: inputs = tokenizer_2( [prompt], - max_length=max_length, + max_length=77, truncation=True, add_special_tokens=True, padding="max_length", @@ -182,8 +179,6 @@ def encode_prompt_and_speech( prompt_embeds, audio_embeds = encode_qwen_omni( tokenizer, text_encoder, prompt, pad_waveform_with_silence(waveform, target_seconds), target_seconds, device ) - if audio_embeds is None: - raise ValueError("the speech encoder returned no features for the speech track") # 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) diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py index b9de1558f409..33c29c89cec4 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/modular_pipeline.py @@ -59,13 +59,10 @@ def align(size: int, multiple: int) -> int: return (size // multiple) * multiple -def select_resolution( - image: PIL.Image.Image, - allowed: tuple[tuple[int, int], ...] = ((512, 512), (768, 512), (512, 768)), -) -> tuple[int, int]: +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(allowed, key=lambda res: abs(aspect - res[1] / res[0])) + 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: @@ -87,7 +84,7 @@ def split_audio_embeds(audio_embeds: torch.Tensor, chunks: int) -> list[torch.Te 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 = audio_group_lengths(num_audio_tokens, duration, device) + lengths = torch.tensor(audio_group_lengths(num_audio_tokens, duration), device=device) return torch.arange(duration, device=device).repeat_interleave(lengths).to(torch.int32) diff --git a/tests/models/transformers/test_models_transformer_kandinsky_s2v.py b/tests/models/transformers/test_models_transformer_kandinsky_s2v.py index af5f24dfb7e5..2a9bb44f0152 100644 --- a/tests/models/transformers/test_models_transformer_kandinsky_s2v.py +++ b/tests/models/transformers/test_models_transformer_kandinsky_s2v.py @@ -12,13 +12,21 @@ # 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 BaseModelTesterConfig, ModelTesterMixin +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) enable_full_determinism() @@ -98,7 +106,7 @@ def get_init_dict(self) -> dict: "text_rope_max_pos": 128, } - def get_dummy_inputs(self, batch_size: int = 1) -> dict: + 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: @@ -108,28 +116,28 @@ def randn(*shape: int) -> torch.Tensor: # 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), + 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]), + _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), + "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), + "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), + "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( @@ -147,3 +155,44 @@ def test_rope_tables_stay_fp32_after_cast(self): *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/test_modular_pipeline_kandinsky_s2v.py b/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py index ca8515fa47ea..dbb1d183d81c 100644 --- a/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py +++ b/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py @@ -26,13 +26,16 @@ 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, ) @@ -74,6 +77,37 @@ class KandinskyS2VModularPipelineTesterConfig(BaseModularPipelineTesterConfig): 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 { @@ -161,6 +195,14 @@ class TestKandinskyS2VModularPipelineLoading(KandinskyS2VModularPipelineTesterCo pass +class TestKandinskyS2VModularPipelineGuider(KandinskyS2VModularPipelineTesterConfig, ModularGuiderTesterMixin): + pass + + +class TestKandinskyS2VModularPipelineWorkflow(KandinskyS2VModularPipelineTesterConfig, ModularWorkflowTesterMixin): + pass + + class TestKandinskyS2VModularPipelineMemory(KandinskyS2VModularPipelineTesterConfig, ModularMemoryTesterMixin): pass @@ -178,6 +220,34 @@ class KandinskyS2VDistilledModularPipelineTesterConfig(KandinskyS2VModularPipeli 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) @@ -204,12 +274,6 @@ def test_num_images_per_prompt(self): def test_float16_inference(self): check_float16_inference(self) - def test_has_no_guidance(self): - """Guidance is distilled into the student: no guider, and no negative prompt to give it.""" - pipe = self.get_pipeline() - assert "guider" not in pipe.components - assert "negative_prompt" not in {param.name for param in pipe.blocks.inputs} - class TestKandinskyS2VDistilledModularPipelineLoading( KandinskyS2VDistilledModularPipelineTesterConfig, ModularLoadingTesterMixin @@ -217,6 +281,12 @@ class TestKandinskyS2VDistilledModularPipelineLoading( pass +class TestKandinskyS2VDistilledModularPipelineWorkflow( + KandinskyS2VDistilledModularPipelineTesterConfig, ModularWorkflowTesterMixin +): + pass + + class TestKandinskyS2VDistilledModularPipelineMemory( KandinskyS2VDistilledModularPipelineTesterConfig, ModularMemoryTesterMixin ): From df2dca1907618a1f4b648225a513d8e8ccd8195f Mon Sep 17 00:00:00 2001 From: Nikolay Semenov Date: Tue, 3 Nov 2026 09:59:25 +0300 Subject: [PATCH 11/12] Encode speech with get_audio_features instead of a hook The speech features now come from the Thinker's public get_audio_features rather than a forward hook on its audio tower. The tower runs twice per call, once inside the Thinker forward for the text hidden states and once for the features: +20 ms on a 12 s track (bf16, H100). The features and the text embeddings are bit-identical to before. --- .../kandinsky_s2v/encoders.py | 22 +++++++------------ 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py index 4128fd0baa11..a248582ae732 100644 --- a/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py +++ b/src/diffusers/modular_pipelines/kandinsky_s2v/encoders.py @@ -80,21 +80,15 @@ def encode_qwen_omni( ) inputs = inputs.to(device).to(text_encoder.dtype) - # The speech features are the audio tower's output, which the Thinker consumes internally; a hook on the tower - # reads them out of the same forward that produces the text hidden states. - captured: dict[str, torch.Tensor] = {} - - def hook(_module, _inputs, output): - captured["features"] = output.last_hidden_state - - handle = text_encoder.audio_tower.register_forward_hook(hook) - try: - output = text_encoder(**inputs, output_hidden_states=True, return_dict=True) - finally: - handle.remove() - + output = text_encoder(**inputs, output_hidden_states=True, return_dict=True) prompt_embeds = slice_user_text(tokenizer, text_encoder, output.hidden_states[-1], inputs) - return prompt_embeds, captured.get("features") + + 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( From ffdd8d8fa112d932d3dcd1aeba2d014e2e602605 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Mon, 2 Nov 2026 22:19:24 -1000 Subject: [PATCH 12/12] Apply batched suggestions from code review Co-authored-by: YiYi Xu --- .../kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 index dbb1d183d81c..8e665f82f21f 100644 --- a/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py +++ b/tests/modular_pipelines/kandinsky_s2v/test_modular_pipeline_kandinsky_s2v.py @@ -40,8 +40,8 @@ # One tiny repository per variant, laid out like the real checkpoints. -TINY_REPO = "kandinskylab/tiny-kandinsky-s2v-modular-pipe" -TINY_DISTILLED_REPO = "kandinskylab/tiny-kandinsky-s2v-distilled-modular-pipe" +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.