From 81f3b12b3b716c355a93f84d366bda24e731cbea Mon Sep 17 00:00:00 2001 From: Denis Koposov Date: Wed, 16 Sep 2026 10:29:07 +0000 Subject: [PATCH 01/62] implement Kandinsky 6 - add TI2VA and SR pipelines - additonally implement PiflowScheduler and MMAudioVAE --- src/diffusers/__init__.py | 15 + src/diffusers/models/__init__.py | 8 + src/diffusers/models/autoencoders/__init__.py | 2 + .../autoencoders/autoencoder_kandinsky6_sr.py | 1178 +++++++ .../autoencoders/autoencoder_mmaudio.py | 1621 +++++++++ src/diffusers/models/transformers/__init__.py | 2 + .../transformers/transformer_kandinsky6.py | 1448 ++++++++ .../transformers/transformer_kandinsky6_sr.py | 2496 +++++++++++++ src/diffusers/pipelines/__init__.py | 14 + .../pipelines/kandinsky6/__init__.py | 39 + .../pipelines/kandinsky6/latent_upscaler.py | 2723 ++++++++++++++ .../kandinsky6/pipeline_kandinsky6_sr.py | 3129 +++++++++++++++++ .../kandinsky6/pipeline_kandinsky6_ti2va.py | 1599 +++++++++ .../pipelines/kandinsky6/pipeline_output.py | 36 + src/diffusers/schedulers/__init__.py | 2 + src/diffusers/schedulers/scheduling_piflow.py | 271 ++ 16 files changed, 14583 insertions(+) create mode 100644 src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py create mode 100644 src/diffusers/models/autoencoders/autoencoder_mmaudio.py create mode 100644 src/diffusers/models/transformers/transformer_kandinsky6.py create mode 100644 src/diffusers/models/transformers/transformer_kandinsky6_sr.py create mode 100644 src/diffusers/pipelines/kandinsky6/__init__.py create mode 100644 src/diffusers/pipelines/kandinsky6/latent_upscaler.py create mode 100644 src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py create mode 100644 src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py create mode 100644 src/diffusers/pipelines/kandinsky6/pipeline_output.py create mode 100644 src/diffusers/schedulers/scheduling_piflow.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index b7d79b8ee97d..a5e9286e7d42 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -302,6 +302,10 @@ "JoyImageEditPlusTransformer3DModel", "JoyImageEditTransformer3DModel", "Kandinsky3UNet", + "MMAudioVAE", + "Kandinsky6SRTransformer3DModel", + "Kandinsky6SRVAE", + "Kandinsky6Transformer3DModel", "Kandinsky5Transformer3DModel", "Krea2Transformer2DModel", "LatteTransformer3DModel", @@ -443,6 +447,7 @@ "EulerAncestralDiscreteScheduler", "EulerDiscreteScheduler", "FlowMapEulerDiscreteScheduler", + "PiflowScheduler", "FlowMatchEulerDiscreteScheduler", "FlowMatchHeunDiscreteScheduler", "FlowMatchLCMScheduler", @@ -693,6 +698,11 @@ "Kandinsky5I2IPipeline", "Kandinsky5I2VPipeline", "Kandinsky5T2IPipeline", + "Kandinsky6SRPipeline", + "Kandinsky6SRPipelineOutput", + "Kandinsky6SRLatentUpscalerBank", + "Kandinsky6TI2VAPipeline", + "Kandinsky6TI2VAPipelineOutput", "Kandinsky5T2VPipeline", "KandinskyCombinedPipeline", "KandinskyImg2ImgCombinedPipeline", @@ -1173,6 +1183,10 @@ JoyImageEditPlusTransformer3DModel, JoyImageEditTransformer3DModel, Kandinsky3UNet, + MMAudioVAE, + Kandinsky6SRTransformer3DModel, + Kandinsky6SRVAE, + Kandinsky6Transformer3DModel, Kandinsky5Transformer3DModel, Krea2Transformer2DModel, LatteTransformer3DModel, @@ -1310,6 +1324,7 @@ EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, FlowMapEulerDiscreteScheduler, + PiflowScheduler, FlowMatchEulerDiscreteScheduler, FlowMatchHeunDiscreteScheduler, FlowMatchLCMScheduler, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 8ba17d896434..64948c22ddcc 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -128,6 +128,10 @@ _import_structure["transformers.transformer_ideogram4"] = ["Ideogram4Transformer2DModel"] _import_structure["transformers.transformer_joyimage"] = ["JoyImageEditTransformer3DModel"] _import_structure["transformers.transformer_joyimage_edit_plus"] = ["JoyImageEditPlusTransformer3DModel"] + _import_structure["autoencoders.autoencoder_mmaudio"] = ["MMAudioVAE"] + _import_structure["autoencoders.autoencoder_kandinsky6_sr"] = ["Kandinsky6SRVAE"] + _import_structure["transformers.transformer_kandinsky6_sr"] = ["Kandinsky6SRTransformer3DModel"] + _import_structure["transformers.transformer_kandinsky6"] = ["Kandinsky6Transformer3DModel"] _import_structure["transformers.transformer_kandinsky"] = ["Kandinsky5Transformer3DModel"] _import_structure["transformers.transformer_krea2"] = ["Krea2Transformer2DModel"] _import_structure["transformers.transformer_longcat_audio_dit"] = ["LongCatAudioDiTTransformer"] @@ -183,6 +187,8 @@ AutoencoderKLFlux2, AutoencoderKLHunyuanImage, AutoencoderKLHunyuanImageRefiner, + MMAudioVAE, + Kandinsky6SRVAE, AutoencoderKLHunyuanVideo, AutoencoderKLHunyuanVideo15, AutoencoderKLKVAE, @@ -267,6 +273,8 @@ Ideogram4Transformer2DModel, JoyImageEditPlusTransformer3DModel, JoyImageEditTransformer3DModel, + Kandinsky6SRTransformer3DModel, + Kandinsky6Transformer3DModel, Kandinsky5Transformer3DModel, Krea2Transformer2DModel, LatteTransformer3DModel, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index 607704343743..08f422cbd4fd 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -6,6 +6,8 @@ from .autoencoder_kl_cogvideox import AutoencoderKLCogVideoX from .autoencoder_kl_cosmos import AutoencoderKLCosmos from .autoencoder_kl_flux2 import AutoencoderKLFlux2 +from .autoencoder_mmaudio import MMAudioVAE +from .autoencoder_kandinsky6_sr import Kandinsky6SRVAE from .autoencoder_kl_hunyuan_video import AutoencoderKLHunyuanVideo from .autoencoder_kl_hunyuanimage import AutoencoderKLHunyuanImage from .autoencoder_kl_hunyuanimage_refiner import AutoencoderKLHunyuanImageRefiner diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py new file mode 100644 index 000000000000..d87fd0df4a6d --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -0,0 +1,1178 @@ +"""Kandinsky 6 SR KVAE Diffusers component.""" + +from __future__ import annotations + +import functools +import math +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, Tuple, Union + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from ...configuration_utils import ConfigMixin, register_to_config +from ..modeling_utils import ModelMixin +from ...utils.accelerate_utils import apply_forward_hook +from einops import rearrange +from safetensors.torch import load_file as safetensors_load_file +from torch.utils.checkpoint import checkpoint + + +def cast_tuple(t, length=1): + return t if isinstance(t, tuple) else ((t,) * length) + + +class SafeConv3d(nn.Conv3d): + def forward(self, x, write_to=None, transform=None): + if transform is None: + transform = lambda x: x + + memory_count = x.numel() / (10**9) + if memory_count > 2: + kernel_size = self.kernel_size[0] + part_num = math.ceil(memory_count / 2) + input_chunks = torch.chunk(x, part_num, dim=2) # NCTHW + + if any(ch.size(2) < kernel_size for ch in input_chunks) and kernel_size > 1: + assert input_chunks[0].numel() * (kernel_size / input_chunks[0].size(2)) < (2 * 10**9), ( + "frames are too big for Conv3d" + ) + + t_stride, output = self.stride[0], [] + for i in range(0, x.size(2) - kernel_size + 1, t_stride): + chunk = transform(x[:, :, i : i + kernel_size]) + output.append(super(SafeConv3d, self).forward(chunk)) + output = torch.cat(output, dim=2) + return output + + if write_to is None: + output = [] + for i, chunk in enumerate(input_chunks): + if i == 0 or kernel_size == 1: + z = torch.clone(chunk) + else: + z = torch.cat([z[:, :, -kernel_size + 1 :], chunk], dim=2) + output.append(super(SafeConv3d, self).forward(transform(z))) + output = torch.cat(output, dim=2) + return output + else: + time_offset = 0 + for i, chunk in enumerate(input_chunks): + if i == 0 or kernel_size == 1: + z = torch.clone(chunk) + else: + z = torch.cat([z[:, :, -kernel_size + 1 :], chunk], dim=2) + z_time = z.size(2) - (kernel_size - 1) + write_to[:, :, time_offset : time_offset + z_time] = super(SafeConv3d, self).forward(transform(z)) + time_offset += z_time + return write_to + else: + if write_to is None: + return super(SafeConv3d, self).forward(transform(x)) + else: + write_to[...] = super(SafeConv3d, self).forward(transform(x)) + return write_to + + +def nonlinearity(x): + # swish + return x * torch.sigmoid(x) + + +def custom_forward(module): + def inside_fn(args): + if isinstance(args, tuple): + return module(*args) + else: + return module(args) + + return inside_fn + + +def ckpt_wrap(module, args): + fn = custom_forward(module) + return checkpoint(fn, args=args, use_reentrant=False) + + +def identity(module, args): + if not isinstance(args, tuple): + return module(args) + return module(*args) + + +Conv3d = SafeConv3d + + +class CausalConv3d(nn.Module): + def __init__( + self, + chan_in, + chan_out, + kernel_size: Union[int, Tuple[int, int, int]], + stride=(1, 1, 1), + dilation=(1, 1, 1), + padding_mode=None, + **kwargs, + ): + super().__init__() + kernel_size = cast_tuple(kernel_size, 3) + + time_kernel_size, height_kernel_size, width_kernel_size = kernel_size + + assert (height_kernel_size % 2) and (width_kernel_size % 2) + + self.height_pad = height_kernel_size // 2 + self.width_pad = width_kernel_size // 2 + self.time_pad = time_kernel_size - 1 + self.time_kernel_size = time_kernel_size + self.temporal_dim = 2 + + self.stride = stride + self.conv = Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) + self.cache_padding = None + self.padding_mode = padding_mode + + def forward(self, input_): + input_parallel = input_ + + padding_3d = (self.width_pad, self.width_pad, self.height_pad, self.height_pad, self.time_pad, 0) + input_parallel = F.pad(input_parallel, padding_3d, mode=self.padding_mode or "replicate") + + output = self.conv(input_parallel) + return output + + +def RMSNorm(in_channels, *args, **kwargs): + return WanRMS_norm(n_ch=in_channels, bias=False) + + +class WanRMS_norm(nn.Module): + r""" + A custom RMS normalization layer. + + Args: + dim (int): The number of dimensions to normalize over. + bias (bool, optional): Whether to include a learnable bias term. Default is False. + """ + + def __init__(self, n_ch: int, bias: bool = False) -> None: + super().__init__() + shape = (n_ch, 1, 1, 1) + + self.scale = n_ch**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x, *args, **kwargs): + return F.normalize(x, dim=1) * self.scale * self.gamma + self.bias + + +class CachedGroupNorm(nn.GroupNorm): + def group_forward(self, x_input, expectation=None, variance=None, return_stat=False): + input_dtype = x_input.dtype + x = x_input.to(torch.float32) + chunks = torch.chunk(x, self.num_groups, dim=1) + if expectation is None: + ch_mean = [torch.mean(chunk, dim=(1, 2, 3, 4), keepdim=True) for chunk in chunks] + else: + ch_mean = expectation + + if variance is None: + ch_var = [torch.var(chunk, dim=(1, 2, 3, 4), keepdim=True, unbiased=False) for chunk in chunks] + else: + ch_var = variance + + x_norm = [(chunk - mean) / torch.sqrt(var + self.eps) for chunk, mean, var in zip(chunks, ch_mean, ch_var)] + x_norm = torch.cat(x_norm, dim=1) + + x_norm.mul_(self.weight.data.view(1, -1, 1, 1, 1)) + x_norm.add_(self.bias.data.view(1, -1, 1, 1, 1)) + + x_out = x_norm.to(input_dtype) + if return_stat: + return x_out, ch_mean, ch_var + return x_out + + def forward(self, x, cache: dict): + out = super().forward(x) + if cache.get("mean") is None and cache.get("var") is None: + cache["mean"] = 1 + cache["var"] = 1 + return out + + +def Normalize(in_channels, gather=False, **kwargs): + return CachedGroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + + +class CachedCausalConv3d(CausalConv3d): + def forward(self, input_, cache: dict, fixed_stride=False): + t_stride = self.stride[0] + padding_3d = (self.height_pad, self.height_pad, self.width_pad, self.width_pad, 0, 0) + input_parallel = F.pad( + input_, + padding_3d, + mode="constant" if self.padding_mode == "zeros" else (self.padding_mode or "replicate"), + value=0 if self.padding_mode in ["constant", "zeros"] else None, + ) + + if cache["padding"] is None: + first_frame = input_parallel[:, :, :1] + time_pad_shape = [i for i in first_frame.shape] + time_pad_shape[2] = self.time_pad + padding = first_frame.expand(time_pad_shape) + else: + padding = cache["padding"] + + out_size = [i for i in input_.shape] + out_size[1] = self.conv.out_channels + if t_stride == 2: + out_size[2] = (input_.size(2) + 1) // 2 + output = torch.empty(tuple(out_size), dtype=input_.dtype, device=input_.device) + + offset_out = math.ceil( + padding.size(2) / t_stride + ) # forward on `padding_poisoned` should take exactly this range + offset_in = offset_out * t_stride - padding.size( + 2 + ) # to make forward on `input_parallel` take slice starting with this index + + if offset_out > 0: + padding_poisoned = torch.cat( + [padding, input_parallel[:, :, : offset_in + self.time_kernel_size - t_stride]], dim=2 + ) + output[:, :, :offset_out] = self.conv(padding_poisoned) + + if offset_out < output.size(2): + output[:, :, offset_out:] = self.conv(input_parallel[:, :, offset_in:]) + + if t_stride == 2 and not fixed_stride and cache["padding"] is not None: + expected_pad_size = padding.size(2) - 1 + offset_out = math.ceil(expected_pad_size / t_stride) + offset_in = offset_out * t_stride - expected_pad_size + + # exact formula, doesn't depend on size of segments + pad_offset = ( + offset_in + + t_stride * math.trunc((input_parallel.size(2) - offset_in - self.time_kernel_size) / t_stride) + + t_stride + ) + + # this condition will be executed ONLY for old models + if t_stride == 2 and not fixed_stride: + pad_offset -= 1 + + if pad_offset < 0: # specific to small chunks (for inference on high resolution videos) + cache["padding"] = torch.cat([padding[:, :, pad_offset:], input_parallel], dim=2) + else: + cache["padding"] = torch.clone(input_parallel[:, :, pad_offset:]) + + return output + + +class CachedCausalResnetBlock3D(nn.Module): + def __init__( + self, + *, + in_channels, + out_channels=None, + conv_shortcut=False, + dropout, + temb_channels=512, + zq_ch=None, + add_conv=False, + gather_norm=False, + normalization=Normalize, + padding_mode=None, + ): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + + self.norm1 = normalization(in_channels, zq_ch=zq_ch, add_conv=add_conv) + + self.conv1 = CachedCausalConv3d( + chan_in=in_channels, chan_out=out_channels, kernel_size=3, padding_mode=padding_mode + ) + if temb_channels > 0: + self.temb_proj = torch.nn.Linear(temb_channels, out_channels) + self.norm2 = normalization(out_channels, zq_ch=zq_ch, add_conv=add_conv) + # self.dropout = torch.nn.Dropout(dropout) + self.conv2 = CachedCausalConv3d( + chan_in=out_channels, chan_out=out_channels, kernel_size=3, padding_mode=padding_mode + ) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = CachedCausalConv3d( + chan_in=in_channels, chan_out=out_channels, kernel_size=3, padding_mode=padding_mode + ) + else: + self.nin_shortcut = Conv3d( + in_channels, + out_channels, + kernel_size=1, + stride=1, + padding=0, + ) + + def forward(self, x, temb, layer_cache, zq=None): + if x.size(2) == 17 and x.size(3) == 1080 and zq is not None: + torch.cuda.empty_cache() + h = x + + if zq is None: + h = self.norm1(h, cache=layer_cache["norm1"]) + else: + h = self.norm1(h, zq, cache=layer_cache["norm1"]) + + if x.size(2) == 17 and x.size(3) == 1080 and zq is not None: + torch.cuda.empty_cache() + + # h = nonlinearity(h) + h = F.silu(h, inplace=True) + if x.size(2) == 17 and x.size(3) == 1080 and zq is not None: + torch.cuda.empty_cache() + + h = self.conv1(h, cache=layer_cache["conv1"]) + if x.size(2) == 17 and x.size(3) == 1080 and zq is not None: + torch.cuda.empty_cache() + + if temb is not None: + h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None, None] + + if zq is None: + h = self.norm2(h, cache=layer_cache["norm2"]) + else: + h = self.norm2(h, zq, cache=layer_cache["norm2"]) + + # h = nonlinearity(h) + h = F.silu(h, inplace=True) + # h = self.dropout(h) + h = self.conv2(h, cache=layer_cache["conv2"]) + + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + x = self.conv_shortcut(x, cache=layer_cache["conv_shortcut"]) + else: + x = self.nin_shortcut(x) + + return x + h + + +class CachedPXSDownsample(nn.Module): + def __init__( + self, in_channels: int, compress_time: bool, factor: int = 2, version=1, fixed_stride=False, padding_mode=None + ): + super().__init__() + self.temporal_compress = compress_time + self.fixed_stride = fixed_stride + self.factor = factor + self.unshuffle = nn.PixelUnshuffle(self.factor) + self.s_pool = nn.AvgPool3d((1, 2, 2), (1, 2, 2)) + + self.version = version + out_channels = in_channels * 2 if version > 1 else in_channels + + self.spatial_conv = Conv3d( + in_channels, + out_channels, + kernel_size=(1, 3, 3), + stride=(1, 2, 2), + padding=(0, 1, 1), + padding_mode=padding_mode or "reflect", + ) + + if self.temporal_compress: + if version == 2: + self.temporal_conv = nn.Sequential( + CachedCausalConv3d( + out_channels, + out_channels, + kernel_size=(2, 1, 1), + stride=(1, 1, 1), + dilation=(1, 1, 1), + padding_mode=padding_mode, + ), + CachedCausalConv3d( + out_channels, + out_channels, + kernel_size=(2, 1, 1), + stride=(2, 1, 1), + dilation=(1, 1, 1), + padding_mode=padding_mode, + ), + ) + else: # 1 or 3 + self.temporal_conv = CachedCausalConv3d( + out_channels, + out_channels, + kernel_size=(3, 1, 1), + stride=(2, 1, 1), + dilation=(1, 1, 1), + padding_mode=padding_mode, + ) + + self.linear = nn.Conv3d(out_channels, out_channels, kernel_size=1, stride=1) + + def spatial_downsample(self, input_): + # PixelShuffle part + pxs_input = rearrange(input_, "b c t h w -> (b t) c h w") + pxs_interm = self.unshuffle(pxs_input) + b, c, h, w = pxs_interm.shape + if self.version > 1: + pxs_interm_view = pxs_interm.view(b, c // self.factor, self.factor, h, w) + else: # + pxs_interm_view = pxs_interm.view(b, c // self.factor**2, self.factor**2, h, w) + pxs_out = torch.mean(pxs_interm_view, dim=2) + pxs_out = rearrange(pxs_out, "(b t) c h w -> b c t h w", t=input_.size(2)) + + # Downsampling by 3D-convolution + conv_out = self.spatial_conv(input_) + + # adding it all together + return conv_out + pxs_out + + def temporal_downsample(self, input_, cache): + # Interpolation part + permuted = rearrange(input_, "b c t h w -> (b h w) c t") + if cache[0]["padding"] is None: + first, rest = permuted[..., :1], permuted[..., 1:] + + if rest.size(-1) > 0: + rest_interp = F.avg_pool1d(rest, kernel_size=2, stride=2) + full_interp = torch.cat([first, rest_interp], dim=-1) + else: + full_interp = first + else: + rest = permuted + if rest.size(-1) > 0: + full_interp = F.avg_pool1d(rest, kernel_size=2, stride=2) + + full_interp = rearrange(full_interp, "(b h w) c t -> b c t h w", h=input_.size(-2), w=input_.size(-1)) + + # Downsampling by convolution + if self.version == 1: + conv_out = self.temporal_conv(input_, cache[0], fixed_stride=self.fixed_stride) + elif self.version == 2: + conv_out = self.temporal_conv[0](input_, cache[0], fixed_stride=self.fixed_stride) + conv_out = self.temporal_conv[1](conv_out, cache[1], fixed_stride=self.fixed_stride) + + return conv_out + full_interp + + def forward(self, x, cache): + # SPATIAL DOWNSAMPLE + out = self.spatial_downsample(x) + + if self.temporal_compress: + # TEMPORAL DOWNSAMPLE + out = self.temporal_downsample(out, cache=cache) + + return self.linear(out) + + +class CachedSpatialNorm3D(nn.Module): + """ + Looking at `forward`, it seems this class should be renamed to `DecoderSpatialNorm` or something similar, + because it's not usual forward, but a conditional one. + """ + + def __init__( + self, + f_channels, + zq_channels, + freeze_norm_layer=False, + add_conv=False, + padding_mode=None, + normalization=Normalize, + **norm_layer_params, + ): + super().__init__() + self.norm_layer = normalization(in_channels=f_channels, **norm_layer_params) + + self.add_conv = add_conv + if add_conv: + self.conv = CachedCausalConv3d( + chan_in=zq_channels, chan_out=zq_channels, kernel_size=3, padding_mode=padding_mode + ) + + self.conv_y = Conv3d( + zq_channels, + f_channels, + kernel_size=1, + ) + self.conv_b = Conv3d( + zq_channels, + f_channels, + kernel_size=1, + ) + + def forward(self, f, zq, cache): + f_shape = [s for s in f.shape] + + if cache["norm"]["mean"] is None and cache["norm"]["var"] is None: + f_first, f_rest = f[:, :, :1], f[:, :, 1:] + f_first_size, f_rest_size = f_first.shape[-3:], f_rest.shape[-3:] + zq_first, zq_rest = zq[:, :, :1], zq[:, :, 1:] + + zq_first = F.interpolate(zq_first, size=f_first_size, mode="nearest") + + if zq.size(2) > 1: + zq_rest_splits = torch.split(zq_rest, 32, dim=1) + interpolated_splits = [ + F.interpolate(split, size=f_rest_size, mode="nearest") for split in zq_rest_splits + ] + + zq_rest = torch.cat(interpolated_splits, dim=1) + zq = torch.cat([zq_first, zq_rest], dim=2) + else: + zq = zq_first + else: + f_size = f.shape[-3:] + zq_splits = torch.split(zq, 32, dim=1) + interpolated_splits = [F.interpolate(split, size=f_size, mode="nearest") for split in zq_splits] + zq = torch.cat(interpolated_splits, dim=1) + + if self.add_conv: + zq = self.conv(zq, cache["add_conv"]) + + norm_f = self.norm_layer(f, cache["norm"]) + norm_f.mul_(self.conv_y(zq)) + norm_f.add_(self.conv_b(zq)) + + if cache["norm"]["mean"] is None and cache["norm"]["var"] is None: + cache["norm"]["mean"] = 1 + cache["norm"]["var"] = 1 + + return norm_f + + +def Normalize3D(in_channels, zq_ch, add_conv, normalization=Normalize): + return CachedSpatialNorm3D( + in_channels, + zq_ch, + freeze_norm_layer=False, + add_conv=add_conv, + num_groups=32, + eps=1e-6, + affine=True, + normalization=normalization, + ) + + +class CachedPXSUpsample(nn.Module): + def __init__(self, in_channels: int, compress_time: bool, factor: int = 2, padding_mode=None): + super().__init__() + self.temporal_compress = compress_time + self.factor = factor + self.shuffle = nn.PixelShuffle(self.factor) + self.spatial_conv = Conv3d( + in_channels, + in_channels, + kernel_size=(1, 3, 3), + stride=(1, 1, 1), + padding=(0, 1, 1), + padding_mode=padding_mode or "reflect", + ) + + if self.temporal_compress: + self.temporal_conv = CachedCausalConv3d( + in_channels, + in_channels, + kernel_size=(3, 1, 1), + stride=(1, 1, 1), + dilation=(1, 1, 1), + padding_mode=padding_mode, + ) + + self.linear = Conv3d(in_channels, in_channels, kernel_size=1, stride=1) + + def spatial_upsample_NEW(self, input_): + def conv_part(x): + to = torch.empty_like(x) + out = self.spatial_conv(x, write_to=to) + return out + + # 5D interpolate keeps channels_last_3d layout intact; merging dims + # via view() is impossible for channels_last strides (dim 1 is innermost) + input_interp = F.interpolate(input_, scale_factor=(1, 2, 2), mode="nearest") + input_interp.add_(conv_part(input_interp)) + return input_interp + + def temporal_upsample(self, input_, cache): + # input_ : (T + 1) x H x W + + repeated = input_.repeat_interleave(2, dim=2) + # repeated: (2T + 2) x H x W + + if cache["padding"] is None: + tail = repeated[..., 1:, :, :] # tail: (2T + 1) x H x W + else: + tail = repeated + + conv_out = self.temporal_conv(tail, cache) + return conv_out + tail + + def forward(self, x, cache): + if self.temporal_compress: + # TEMPORAL UPSAMPLE + x = self.temporal_upsample(x, cache) + + # SPATIAL UPSAMPLE + s_out = self.spatial_upsample_NEW(x) + to = torch.empty_like(s_out) + + lin_out = self.linear(s_out, write_to=to) + + return lin_out + + +class DiagonalGaussianDistribution(object): + def __init__(self, parameters, deterministic=False): + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like(self.mean) # .to(device=self.parameters.device) + + def sample(self): + x = self.mean + self.std * torch.randn_like(self.mean) + return x + + def kl(self, other=None): + if self.deterministic: + return torch.Tensor([0.0]) + else: + if other is None: + return 0.5 * torch.sum( + torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, + dim=[1, 2, 3, 4], + ) + else: + return 0.5 * torch.sum( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var + - 1.0 + - self.logvar + + other.logvar, + dim=[1, 2, 3, 4], + ) + + def mode(self): + return self.mean + + +class DiagonalGaussianRegularizer(nn.Module): + def __init__(self, sample: bool = True): + super().__init__() + self.sample = sample + + def get_trainable_parameters(self) -> Any: + yield from () + + def get_kl_loss(self, z: torch.Tensor) -> torch.Tensor: + posterior = DiagonalGaussianDistribution(z) + kl_loss = posterior.kl() + kl_loss = torch.mean(kl_loss, dim=0) + return kl_loss + + def get_sample(self, z: torch.Tensor) -> torch.Tensor: + posterior = DiagonalGaussianDistribution(z) + if self.sample: + z = posterior.sample() + else: + z = posterior.mode() + return z + + def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]: + log = dict() + posterior = DiagonalGaussianDistribution(z) + if self.sample: + z = posterior.sample() + else: + z = posterior.mode() + kl_loss = posterior.kl() + kl_loss = torch.sum(kl_loss) / kl_loss.shape[0] + log["kl_loss"] = kl_loss + return z, log + + +class CachedEncoder3D(nn.Module): + def __init__( + self, + *, + ch=128, + ch_mult=(1, 2, 4, 8), + num_res_blocks, + dropout=0.0, + resamp_with_conv=True, + in_channels, + resolution=0, + z_channels, + double_z=True, + padding_mode=None, + temporal_compress_times=4, + gather_norm=False, + fix_pxs=False, + norm_type="group_norm", + downsample_version=1, + temporal_compress_start_level=0, + skip_last_resolution=False, + **ignore_kwargs, + ): + super().__init__() + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.skip_last_resolution = skip_last_resolution + + # log2 of temporal_compress_times + temporal_compress_level = int(np.log2(temporal_compress_times)) + temporal_compress_start_level + + in_ch_mult = (ch_mult[0],) + tuple(ch_mult) + self.conv_in = CachedCausalConv3d( + chan_in=in_channels, chan_out=round(in_ch_mult[0] * self.ch), kernel_size=3, padding_mode=padding_mode + ) + + normalization = Normalize if norm_type == "group_norm" else RMSNorm + + curr_res = resolution + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + + block_in = round(ch * in_ch_mult[i_level]) + block_out = round(ch * ch_mult[i_level]) + + if downsample_version > 1 and i_level > 0: + block_in *= 2 + + for i_block in range(self.num_res_blocks): + block.append( + CachedCausalResnetBlock3D( + in_channels=block_in, + out_channels=block_out, + dropout=dropout, + temb_channels=self.temb_ch, + gather_norm=gather_norm, + normalization=normalization, + padding_mode=padding_mode, + ) + ) + block_in = block_out + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions - 1: + if temporal_compress_start_level <= i_level < temporal_compress_level: + down.downsample = CachedPXSDownsample( + block_in, + compress_time=True, + version=downsample_version, + fixed_stride=fix_pxs, + padding_mode=padding_mode, + ) + else: + down.downsample = CachedPXSDownsample( + block_in, + compress_time=False, + version=downsample_version, + fixed_stride=fix_pxs, + padding_mode=padding_mode, + ) + curr_res = curr_res // 2 + if not skip_last_resolution or i_level != self.num_resolutions - 1: + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = CachedCausalResnetBlock3D( + in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + gather_norm=gather_norm, + normalization=normalization, + padding_mode=padding_mode, + ) + + self.mid.block_2 = CachedCausalResnetBlock3D( + in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + gather_norm=gather_norm, + normalization=normalization, + padding_mode=padding_mode, + ) + + # end + self.norm_out = normalization(block_in, gather=gather_norm) + + self.conv_out = CachedCausalConv3d( + chan_in=block_in, + chan_out=2 * z_channels if double_z else z_channels, + kernel_size=3, + padding_mode=padding_mode, + ) + + def forward(self, x, cache_dict, use_cp=True): + # timestep embedding + temb = None + + # downsampling + h = self.conv_in(x, cache=cache_dict["conv_in"]) + for i_level in range(self.num_resolutions): + if not self.skip_last_resolution or i_level != self.num_resolutions - 1: + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](h, temb, layer_cache=cache_dict[i_level][i_block]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + if i_level != self.num_resolutions - 1: + h = self.down[i_level].downsample(h, cache=cache_dict[i_level]["down"]) + + # middle + h = self.mid.block_1(h, temb, layer_cache=cache_dict["mid_1"]) + h = self.mid.block_2(h, temb, layer_cache=cache_dict["mid_2"]) + + # end + h = self.norm_out(h, cache=cache_dict["norm_out"]) + h = nonlinearity(h) + h = self.conv_out(h, cache=cache_dict["conv_out"]) + + return h + + def get_last_layer(self): + return self.conv_out.conv.weight + + +class CachedDecoder3D(nn.Module): + def __init__( + self, + ch=128, + out_ch=None, + ch_mult=(1, 2, 4, 8), + num_res_blocks=2, + dropout=0.0, + resamp_with_conv=True, + resolution=0, + z_channels=16, + give_pre_end=False, + zq_ch=None, + add_conv=False, + padding_mode=None, + temporal_compress_times=4, + gather_norm=False, + norm_type="group_norm", + temporal_compress_start_level=0, + **kwargs, + ): + super().__init__() + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.give_pre_end = give_pre_end + + # log2 of temporal_compress_times + temporal_compress_level = int(np.log2(temporal_compress_times)) + temporal_compress_start_level + + if zq_ch is None: + zq_ch = z_channels + + # compute in_ch_mult, block_in and curr_res at lowest res + block_in = round(ch * ch_mult[self.num_resolutions - 1]) + curr_res = resolution // 2 ** (self.num_resolutions - 1) + self.z_shape = (1, z_channels, curr_res, curr_res) + + self.conv_in = CachedCausalConv3d( + chan_in=z_channels, chan_out=block_in, kernel_size=3, padding_mode=padding_mode + ) + + modulated_norm = functools.partial( + Normalize3D, normalization=Normalize if norm_type == "group_norm" else RMSNorm + ) + + # middle + self.mid = nn.Module() + self.mid.block_1 = CachedCausalResnetBlock3D( + in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + zq_ch=zq_ch, + add_conv=add_conv, + normalization=modulated_norm, + gather_norm=gather_norm, + padding_mode=padding_mode, + ) + + self.mid.block_2 = CachedCausalResnetBlock3D( + in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + zq_ch=zq_ch, + add_conv=add_conv, + normalization=modulated_norm, + gather_norm=gather_norm, + padding_mode=padding_mode, + ) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = round(ch * ch_mult[i_level]) + for i_block in range(self.num_res_blocks + 1): + block.append( + CachedCausalResnetBlock3D( + in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout, + zq_ch=zq_ch, + add_conv=add_conv, + normalization=modulated_norm, + gather_norm=gather_norm, + padding_mode=padding_mode, + ) + ) + block_in = block_out + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + if ( + self.num_resolutions - temporal_compress_start_level + > i_level + >= self.num_resolutions - temporal_compress_level + ): + up.upsample = CachedPXSUpsample(block_in, compress_time=True, padding_mode=padding_mode) + else: + up.upsample = CachedPXSUpsample(block_in, compress_time=False, padding_mode=padding_mode) + self.up.insert(0, up) + + self.norm_out = modulated_norm(block_in, zq_ch, add_conv=add_conv) # , gather=gather_norm) + + self.conv_out = CachedCausalConv3d(chan_in=block_in, chan_out=out_ch, kernel_size=3, padding_mode=padding_mode) + + def forward(self, z, cache_dict): + self.last_z_shape = z.shape + + # timestep embedding + temb = None + + t = z.shape[2] + + zq = z + h = self.conv_in(z, cache_dict["conv_in"]) + + # middle + h = self.mid.block_1(h, temb, layer_cache=cache_dict["mid_1"], zq=zq) + h = self.mid.block_2(h, temb, layer_cache=cache_dict["mid_2"], zq=zq) + + # upsampling + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h, temb, layer_cache=cache_dict[i_level][i_block], zq=zq) + + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h, zq) + if i_level != 0: + h = self.up[i_level].upsample(h, cache_dict[i_level]["up"]) + + # end + if self.give_pre_end: + return h + + h = self.norm_out(h, zq, cache_dict["norm_out"]) + h = nonlinearity(h) + h = self.conv_out(h, cache_dict["conv_out"]) + + return h + + def get_last_layer(self): + return self.conv_out.conv.weight + + +@dataclass +class DecoderOutput: + sample: torch.Tensor + + +class Kandinsky6SRVAE(ModelMixin, ConfigMixin): + @staticmethod + def normalize_data(data): + return data / 128 - 1.0 + + @staticmethod + def denormalize_data(data): + return (data + 1) * 128 + + @register_to_config + def __init__( + self, + vae_type: str, + encoder_config: Mapping[str, Any], + decoder_config: Mapping[str, Any], + scaling_factor: float = 1.0, + spatial_factor: int = 16, + temporal_factor: int = 4, + ) -> None: + super().__init__() + if vae_type != "video-kvae": + raise ValueError(f"Kandinsky6SRVAE supports only 'video-kvae', got: {vae_type!r}") + self.conf = {"enc": encoder_config, "dec": decoder_config} + self.encoder = CachedEncoder3D(**encoder_config) + self.decoder = CachedDecoder3D(**decoder_config) + self.regularizer = DiagonalGaussianRegularizer(sample=False) + self.spatial_factor = int(spatial_factor) + self.temporal_factor = int(temporal_factor) + + def init_from_ckpt(self, path): + if str(path).endswith(".safetensors"): + # Release checkpoints are flat safetensors state dicts already in + # the new-style key naming. + self.load_state_dict(safetensors_load_file(str(path)), strict=True) + return + sd = torch.load(path, map_location="cpu")["state_dict"] + + # Fix checkpoint if starting from new style + replace_keys = dict() + delete_keys = list() + for k in sd: + if k.startswith("loss."): + delete_keys.append(k) + continue + if "encoder.down" in k and "downsample.temporal_conv.conv" in k: + continue + + if k.startswith("decoder"): + if "upsample" in k: + continue + elif ".conv_b.conv" in k: + replace_keys[k] = k.replace(".conv_b.conv", ".conv_b") + continue + elif ".conv_y.conv" in k: + replace_keys[k] = k.replace(".conv_y.conv", ".conv_y") + continue + + if k.endswith("sample.temporal_conv.conv.weight") or k.endswith("sample.temporal_conv.conv.bias"): + replace_keys[k] = k.replace(".temporal_conv.conv.", ".temporal_conv.") + for old_k, new_k in replace_keys.items(): + sd[new_k] = sd[old_k] + del sd[old_k] + for k in delete_keys: + del sd[k] + + self.load_state_dict(sd, strict=True) + + def make_empty_cache(self, block: str): + def make_dict(name, p=None): + if name == "conv": + return {"padding": None} + + layer, module = name.split("_") + if layer == "norm": + if module == "enc": + return {"mean": None, "var": None} + else: + return {"norm": make_dict("norm_enc"), "add_conv": make_dict("conv")} + elif layer == "resblock": + return { + "norm1": make_dict(f"norm_{module}"), + "norm2": make_dict(f"norm_{module}"), + "conv1": make_dict("conv"), + "conv2": make_dict("conv"), + "conv_shortcut": make_dict("conv"), + } + elif layer.isdigit(): + out_dict = {"down": [make_dict("conv"), make_dict("conv")], "up": make_dict("conv")} + for i in range(p): + out_dict[i] = make_dict(f"resblock_{module}") + + return out_dict + + cache = { + "conv_in": make_dict("conv"), + "mid_1": make_dict(f"resblock_{block}"), + "mid_2": make_dict(f"resblock_{block}"), + "norm_out": make_dict(f"norm_{block}"), + "conv_out": make_dict("conv"), + } + for i in range(len(self.conf[block].get("ch_mult", [1, 2, 4, 8]))): + cache[i] = make_dict(f"{i}_block", p=self.conf[block]["num_res_blocks"] + 1) + return cache + + @apply_forward_hook + def encode(self, x, seg_len=16): + cache = self.make_empty_cache("enc") + + ## get segments size + split_list = [seg_len + 1] + n_frames = x.size(2) - (seg_len + 1) + while n_frames > 0: + split_list.append(seg_len) + n_frames -= seg_len + + split_list[-1] += n_frames + + ## encode by segments + latent = [] + for chunk in torch.split(x, split_list, dim=2): + l = self.encoder(chunk, cache) + sample = self.regularizer.get_sample(l) + latent.append(sample) + + latent = torch.cat(latent, dim=2) + return latent, split_list + + @apply_forward_hook + def decode(self, z, split_list=None): + cache = self.make_empty_cache("dec") + + ## get segments size + if split_list is None: + default_split_size = 16 // self.conf["enc"]["temporal_compress_times"] + time_dim = z.shape[2] + if time_dim == 1: + # image + split_list = [1] + else: + splits_num = (time_dim - 1) // default_split_size + split_list = [default_split_size] * splits_num + if (time_dim - 1) % default_split_size != 0: + split_list.append((time_dim - 1) % default_split_size) + split_list[0] += 1 + else: + split_list = [math.ceil(size / self.conf["enc"]["temporal_compress_times"]) for size in split_list] + + ## decode by segments + recs = [] + for chunk in torch.split(z, split_list, dim=2): + out = self.decoder(chunk, cache) + recs.append(out) + + recs = torch.cat(recs, dim=2) + return DecoderOutput(sample=recs) + + def forward(self, x, seg_len: int = 16): + latent, split_list = self.encode(x, seg_len) + recs = self.decode(latent, split_list) + return recs + + +__all__ = ["Kandinsky6SRVAE"] diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py new file mode 100644 index 000000000000..243bd3c91f5d --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -0,0 +1,1621 @@ +"""MMAudio VAE Diffusers component.""" + +from __future__ import annotations + +import json +import math +import os +import shutil +from os import PathLike +from pathlib import Path +from typing import Any, Dict, Literal, Optional, Union + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from ...configuration_utils import ConfigMixin, register_to_config +from ..modeling_utils import ModelMixin +from ...utils.accelerate_utils import apply_forward_hook +from einops import rearrange +from huggingface_hub import PyTorchModelHubMixin, hf_hub_download +from librosa.filters import mel as librosa_mel_fn +from torch import nn, pow, sin +from torch.nn import Conv1d, ConvTranspose1d, Parameter +from torch.nn import functional as F +from torch.nn.utils import weight_norm as legacy_weight_norm +from torch.nn.utils.parametrizations import weight_norm +from torch.nn.utils.parametrize import remove_parametrizations + +# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# This work is licensed under a Creative Commons +# Attribution-NonCommercial-ShareAlike 4.0 International License. +# You should have received a copy of the license along with this +# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/ +"""Improved diffusion model architecture proposed in the paper +"Analyzing and Improving the Training Dynamics of Diffusion Models".""" + + +# ---------------------------------------------------------------------------- +# Variant of constant() that inherits dtype and device from the given +# reference tensor by default. + +_constant_cache = dict() + + +def constant(value, shape=None, dtype=None, device=None, memory_format=None): + value = np.asarray(value) + if shape is not None: + shape = tuple(shape) + if dtype is None: + dtype = torch.get_default_dtype() + if device is None: + device = torch.device("cpu") + if memory_format is None: + memory_format = torch.contiguous_format + + key = (value.shape, value.dtype, value.tobytes(), shape, dtype, device, memory_format) + tensor = _constant_cache.get(key, None) + if tensor is None: + tensor = torch.as_tensor(value.copy(), dtype=dtype, device=device) + if shape is not None: + tensor, _ = torch.broadcast_tensors(tensor, torch.empty(shape)) + tensor = tensor.contiguous(memory_format=memory_format) + _constant_cache[key] = tensor + return tensor + + +def const_like(ref, value, shape=None, dtype=None, device=None, memory_format=None): + if dtype is None: + dtype = ref.dtype + if device is None: + device = ref.device + return constant(value, shape=shape, dtype=dtype, device=device, memory_format=memory_format) + + +# ---------------------------------------------------------------------------- +# Normalize given tensor to unit magnitude with respect to the given +# dimensions. Default = all dimensions except the first. + + +def normalize(x, dim=None, eps=1e-4): + if dim is None: + dim = list(range(1, x.ndim)) + norm = torch.linalg.vector_norm(x, dim=dim, keepdim=True, dtype=torch.float32) + norm = torch.add(eps, norm, alpha=np.sqrt(norm.numel() / x.numel())) + return x / norm.to(x.dtype) + + +class Normalize(torch.nn.Module): + def __init__(self, dim=None, eps=1e-4): + super().__init__() + self.dim = dim + self.eps = eps + + def forward(self, x): + return normalize(x, dim=self.dim, eps=self.eps) + + +# ---------------------------------------------------------------------------- +# Upsample or downsample the given tensor with the given filter, +# or keep it as is. + + +def resample(x, f=[1, 1], mode="keep"): + if mode == "keep": + return x + f = np.float32(f) + assert f.ndim == 1 and len(f) % 2 == 0 + pad = (len(f) - 1) // 2 + f = f / f.sum() + f = np.outer(f, f)[np.newaxis, np.newaxis, :, :] + f = const_like(x, f) + c = x.shape[1] + if mode == "down": + return torch.nn.functional.conv2d(x, f.tile([c, 1, 1, 1]), groups=c, stride=2, padding=(pad,)) + assert mode == "up" + return torch.nn.functional.conv_transpose2d(x, (f * 4).tile([c, 1, 1, 1]), groups=c, stride=2, padding=(pad,)) + + +# ---------------------------------------------------------------------------- +# Magnitude-preserving SiLU (Equation 81). + + +def mp_silu(x): + return torch.nn.functional.silu(x) / 0.596 + + +class MPSiLU(torch.nn.Module): + def forward(self, x): + return mp_silu(x) + + +# ---------------------------------------------------------------------------- +# Magnitude-preserving sum (Equation 88). + + +def mp_sum(a, b, t=0.5): + return a.lerp(b, t) / np.sqrt((1 - t) ** 2 + t**2) + + +# ---------------------------------------------------------------------------- +# Magnitude-preserving concatenation (Equation 103). + + +def mp_cat(a, b, dim=1, t=0.5): + Na = a.shape[dim] + Nb = b.shape[dim] + C = np.sqrt((Na + Nb) / ((1 - t) ** 2 + t**2)) + wa = C / np.sqrt(Na) * (1 - t) + wb = C / np.sqrt(Nb) * t + return torch.cat([wa * a, wb * b], dim=dim) + + +# ---------------------------------------------------------------------------- +# Magnitude-preserving convolution or fully-connected layer (Equation 47) +# with force weight normalization (Equation 66). + + +class MPConv1D(torch.nn.Module): + def __init__(self, in_channels, out_channels, kernel_size): + super().__init__() + self.out_channels = out_channels + self.weight = torch.nn.Parameter(torch.randn(out_channels, in_channels, kernel_size)) + + self.weight_norm_removed = False + + def forward(self, x, gain=1): + assert self.weight_norm_removed, "call remove_weight_norm() before inference" + + w = self.weight * gain + if w.ndim == 2: + return x @ w.t() + assert w.ndim == 3 + return torch.nn.functional.conv1d(x, w, padding=(w.shape[-1] // 2,)) + + def remove_weight_norm(self): + w = self.weight.to(torch.float32) + w = normalize(w) # traditional weight normalization + w = w / np.sqrt(w[0].numel()) + w = w.to(self.weight.dtype) + self.weight.data.copy_(w) + + self.weight_norm_removed = True + return self + + +class DiagonalGaussianDistribution: + def __init__(self, parameters, deterministic=False): + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device) + + def sample(self, rng: Optional[torch.Generator] = None): + # x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device) + + r = torch.empty_like(self.mean).normal_(generator=rng) + x = self.mean + self.std * r + + return x + + def kl(self, other=None): + if self.deterministic: + return torch.Tensor([0.0]) + else: + if other is None: + return 0.5 * torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar + else: + return 0.5 * ( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var + - 1.0 + - self.logvar + + other.logvar + ) + + def nll(self, sample, dims=[1, 2, 3]): + if self.deterministic: + return torch.Tensor([0.0]) + logtwopi = np.log(2.0 * np.pi) + return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, dim=dims) + + def mode(self): + return self.mean + + +def nonlinearity(x): + # swish + return mp_silu(x) + + +class ResnetBlock1D(nn.Module): + def __init__(self, *, in_dim, out_dim=None, conv_shortcut=False, kernel_size=3, use_norm=True): + super().__init__() + self.in_dim = in_dim + out_dim = in_dim if out_dim is None else out_dim + self.out_dim = out_dim + self.use_conv_shortcut = conv_shortcut + self.use_norm = use_norm + + self.conv1 = MPConv1D(in_dim, out_dim, kernel_size=kernel_size) + self.conv2 = MPConv1D(out_dim, out_dim, kernel_size=kernel_size) + if self.in_dim != self.out_dim: + if self.use_conv_shortcut: + self.conv_shortcut = MPConv1D(in_dim, out_dim, kernel_size=kernel_size) + else: + self.nin_shortcut = MPConv1D(in_dim, out_dim, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + + # pixel norm + if self.use_norm: + x = normalize(x, dim=1) + + h = x + h = nonlinearity(h) + h = self.conv1(h) + + h = nonlinearity(h) + h = self.conv2(h) + + if self.in_dim != self.out_dim: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + + return mp_sum(x, h, t=0.3) + + +class AttnBlock1D(nn.Module): + def __init__(self, in_channels, num_heads=1): + super().__init__() + self.in_channels = in_channels + + self.num_heads = num_heads + self.qkv = MPConv1D(in_channels, in_channels * 3, kernel_size=1) + self.proj_out = MPConv1D(in_channels, in_channels, kernel_size=1) + + def forward(self, x): + h = x + y = self.qkv(h) + y = y.reshape(y.shape[0], self.num_heads, -1, 3, y.shape[-1]) + q, k, v = normalize(y, dim=2).unbind(3) + + q = rearrange(q, "b h c l -> b h l c") + k = rearrange(k, "b h c l -> b h l c") + v = rearrange(v, "b h c l -> b h l c") + + h = F.scaled_dot_product_attention(q, k, v) + h = rearrange(h, "b h l c -> b (h c) l") + + h = self.proj_out(h) + + return mp_sum(x, h, t=0.3) + + +class Upsample1D(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + self.conv = MPConv1D(in_channels, in_channels, kernel_size=3) + + def forward(self, x): + x = F.interpolate(x, scale_factor=2.0, mode="nearest-exact") # support 3D tensor(B,C,T) + if self.with_conv: + x = self.conv(x) + return x + + +class Downsample1D(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + # no asymmetric padding in torch conv, must do it ourselves + self.conv1 = MPConv1D(in_channels, in_channels, kernel_size=1) + self.conv2 = MPConv1D(in_channels, in_channels, kernel_size=1) + + def forward(self, x): + + if self.with_conv: + x = self.conv1(x) + + x = F.avg_pool1d(x, kernel_size=2, stride=2) + + if self.with_conv: + x = self.conv2(x) + + return x + + +# fmt: off +DATA_MEAN_128D = [ + -3.3462, -2.6723, -2.4893, -2.3143, -2.2664, -2.3317, -2.1802, -2.4006, -2.2357, -2.4597, + -2.3717, -2.4690, -2.5142, -2.4919, -2.6610, -2.5047, -2.7483, -2.5926, -2.7462, -2.7033, + -2.7386, -2.8112, -2.7502, -2.9594, -2.7473, -3.0035, -2.8891, -2.9922, -2.9856, -3.0157, + -3.1191, -2.9893, -3.1718, -3.0745, -3.1879, -3.2310, -3.1424, -3.2296, -3.2791, -3.2782, + -3.2756, -3.3134, -3.3509, -3.3750, -3.3951, -3.3698, -3.4505, -3.4509, -3.5089, -3.4647, + -3.5536, -3.5788, -3.5867, -3.6036, -3.6400, -3.6747, -3.7072, -3.7279, -3.7283, -3.7795, + -3.8259, -3.8447, -3.8663, -3.9182, -3.9605, -3.9861, -4.0105, -4.0373, -4.0762, -4.1121, + -4.1488, -4.1874, -4.2461, -4.3170, -4.3639, -4.4452, -4.5282, -4.6297, -4.7019, -4.7960, + -4.8700, -4.9507, -5.0303, -5.0866, -5.1634, -5.2342, -5.3242, -5.4053, -5.4927, -5.5712, + -5.6464, -5.7052, -5.7619, -5.8410, -5.9188, -6.0103, -6.0955, -6.1673, -6.2362, -6.3120, + -6.3926, -6.4797, -6.5565, -6.6511, -6.8130, -6.9961, -7.1275, -7.2457, -7.3576, -7.4663, + -7.6136, -7.7469, -7.8815, -8.0132, -8.1515, -8.3071, -8.4722, -8.7418, -9.3975, -9.6628, + -9.7671, -9.8863, -9.9992, -10.0860, -10.1709, -10.5418, -11.2795, -11.3861 +] + +DATA_STD_128D = [ + 2.3804, 2.4368, 2.3772, 2.3145, 2.2803, 2.2510, 2.2316, 2.2083, 2.1996, 2.1835, 2.1769, 2.1659, + 2.1631, 2.1618, 2.1540, 2.1606, 2.1571, 2.1567, 2.1612, 2.1579, 2.1679, 2.1683, 2.1634, 2.1557, + 2.1668, 2.1518, 2.1415, 2.1449, 2.1406, 2.1350, 2.1313, 2.1415, 2.1281, 2.1352, 2.1219, 2.1182, + 2.1327, 2.1195, 2.1137, 2.1080, 2.1179, 2.1036, 2.1087, 2.1036, 2.1015, 2.1068, 2.0975, 2.0991, + 2.0902, 2.1015, 2.0857, 2.0920, 2.0893, 2.0897, 2.0910, 2.0881, 2.0925, 2.0873, 2.0960, 2.0900, + 2.0957, 2.0958, 2.0978, 2.0936, 2.0886, 2.0905, 2.0845, 2.0855, 2.0796, 2.0840, 2.0813, 2.0817, + 2.0838, 2.0840, 2.0917, 2.1061, 2.1431, 2.1976, 2.2482, 2.3055, 2.3700, 2.4088, 2.4372, 2.4609, + 2.4731, 2.4847, 2.5072, 2.5451, 2.5772, 2.6147, 2.6529, 2.6596, 2.6645, 2.6726, 2.6803, 2.6812, + 2.6899, 2.6916, 2.6931, 2.6998, 2.7062, 2.7262, 2.7222, 2.7158, 2.7041, 2.7485, 2.7491, 2.7451, + 2.7485, 2.7233, 2.7297, 2.7233, 2.7145, 2.6958, 2.6788, 2.6439, 2.6007, 2.4786, 2.2469, 2.1877, + 2.1392, 2.0717, 2.0107, 1.9676, 1.9140, 1.7102, 0.9101, 0.7164 +] +# fmt: on + + +class VAE(nn.Module): + def __init__( + self, + *, + data_dim: int, + embed_dim: int, + hidden_dim: int, + ): + super().__init__() + + if data_dim != 128: + raise ValueError(f"Only 44k audio VAE data is supported, got data_dim={data_dim}") + self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_128D, dtype=torch.float32)) + self.data_std = nn.Buffer(torch.tensor(DATA_STD_128D, dtype=torch.float32)) + + self.data_mean = self.data_mean.view(1, -1, 1) + self.data_std = self.data_std.view(1, -1, 1) + + self.encoder = Encoder1D( + dim=hidden_dim, + ch_mult=(1, 2, 4), + num_res_blocks=2, + attn_layers=[3], + down_layers=[0], + in_dim=data_dim, + embed_dim=embed_dim, + ) + self.decoder = Decoder1D( + dim=hidden_dim, + ch_mult=(1, 2, 4), + num_res_blocks=2, + attn_layers=[3], + down_layers=[0], + in_dim=data_dim, + out_dim=data_dim, + embed_dim=embed_dim, + ) + + self.embed_dim = embed_dim + # self.quant_conv = nn.Conv1d(2 * embed_dim, 2 * embed_dim, 1) + # self.post_quant_conv = nn.Conv1d(embed_dim, embed_dim, 1) + + def encode(self, x: torch.Tensor, normalize: bool = True) -> DiagonalGaussianDistribution: + if normalize: + x = self.normalize(x) + moments = self.encoder(x) + posterior = DiagonalGaussianDistribution(moments) + return posterior + + def decode(self, z: torch.Tensor, unnormalize: bool = True) -> torch.Tensor: + dec = self.decoder(z) + if unnormalize: + dec = self.unnormalize(dec) + return dec + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + return (x - self.data_mean) / self.data_std + + def unnormalize(self, x: torch.Tensor) -> torch.Tensor: + return x * self.data_std + self.data_mean + + def forward( + self, + x: torch.Tensor, + sample_posterior: bool = True, + rng: Optional[torch.Generator] = None, + normalize: bool = True, + unnormalize: bool = True, + ) -> tuple[torch.Tensor, DiagonalGaussianDistribution]: + + posterior = self.encode(x, normalize=normalize) + if sample_posterior: + z = posterior.sample(rng) + else: + z = posterior.mode() + dec = self.decode(z, unnormalize=unnormalize) + return dec, posterior + + def load_weights(self, src_dict) -> None: + self.load_state_dict(src_dict, strict=True) + + @property + def device(self) -> torch.device: + return next(self.parameters()).device + + def get_last_layer(self): + return self.decoder.conv_out.weight + + def remove_weight_norm(self): + for name, m in self.named_modules(): + if isinstance(m, MPConv1D): + m.remove_weight_norm() + + return self + + +class Encoder1D(nn.Module): + def __init__( + self, + *, + dim: int, + ch_mult: tuple[int] = (1, 2, 4, 8), + num_res_blocks: int, + attn_layers: list[int] = [], + down_layers: list[int] = [], + resamp_with_conv: bool = True, + in_dim: int, + embed_dim: int, + double_z: bool = True, + kernel_size: int = 3, + clip_act: float = 256.0, + ): + super().__init__() + self.dim = dim + self.num_layers = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.in_channels = in_dim + self.clip_act = clip_act + self.down_layers = down_layers + self.attn_layers = attn_layers + self.conv_in = MPConv1D(in_dim, self.dim, kernel_size=kernel_size) + + in_ch_mult = (1,) + tuple(ch_mult) + self.in_ch_mult = in_ch_mult + # downsampling + self.down = nn.ModuleList() + for i_level in range(self.num_layers): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = dim * in_ch_mult[i_level] + block_out = dim * ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append(ResnetBlock1D(in_dim=block_in, out_dim=block_out, kernel_size=kernel_size, use_norm=True)) + block_in = block_out + if i_level in attn_layers: + attn.append(AttnBlock1D(block_in)) + down = nn.Module() + down.block = block + down.attn = attn + if i_level in down_layers: + down.downsample = Downsample1D(block_in, resamp_with_conv) + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, kernel_size=kernel_size, use_norm=True) + self.mid.attn_1 = AttnBlock1D(block_in) + self.mid.block_2 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, kernel_size=kernel_size, use_norm=True) + + # end + self.conv_out = MPConv1D(block_in, 2 * embed_dim if double_z else embed_dim, kernel_size=kernel_size) + + self.learnable_gain = nn.Parameter(torch.zeros([])) + + def forward(self, x): + + # downsampling + hs = [self.conv_in(x)] + for i_level in range(self.num_layers): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + h = h.clamp(-self.clip_act, self.clip_act) + hs.append(h) + if i_level in self.down_layers: + hs.append(self.down[i_level].downsample(hs[-1])) + + # middle + h = hs[-1] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + h = h.clamp(-self.clip_act, self.clip_act) + + # end + h = nonlinearity(h) + h = self.conv_out(h, gain=(self.learnable_gain + 1)) + return h + + +class Decoder1D(nn.Module): + def __init__( + self, + *, + dim: int, + out_dim: int, + ch_mult: tuple[int] = (1, 2, 4, 8), + num_res_blocks: int, + attn_layers: list[int] = [], + down_layers: list[int] = [], + kernel_size: int = 3, + resamp_with_conv: bool = True, + in_dim: int, + embed_dim: int, + clip_act: float = 256.0, + ): + super().__init__() + self.ch = dim + self.num_layers = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.in_channels = in_dim + self.clip_act = clip_act + self.down_layers = [i + 1 for i in down_layers] # each downlayer add one + + # compute in_ch_mult, block_in and curr_res at lowest res + block_in = dim * ch_mult[self.num_layers - 1] + + # z to block_in + self.conv_in = MPConv1D(embed_dim, block_in, kernel_size=kernel_size) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True) + self.mid.attn_1 = AttnBlock1D(block_in) + self.mid.block_2 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_layers)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = dim * ch_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + block.append(ResnetBlock1D(in_dim=block_in, out_dim=block_out, use_norm=True)) + block_in = block_out + if i_level in attn_layers: + attn.append(AttnBlock1D(block_in)) + up = nn.Module() + up.block = block + up.attn = attn + if i_level in self.down_layers: + up.upsample = Upsample1D(block_in, resamp_with_conv) + self.up.insert(0, up) # prepend to get consistent order + + # end + self.conv_out = MPConv1D(block_in, out_dim, kernel_size=kernel_size) + self.learnable_gain = nn.Parameter(torch.zeros([])) + + def forward(self, z): + # z to block_in + h = self.conv_in(z) + + # middle + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + h = h.clamp(-self.clip_act, self.clip_act) + + # upsampling + for i_level in reversed(range(self.num_layers)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + h = h.clamp(-self.clip_act, self.clip_act) + if i_level in self.down_layers: + h = self.up[i_level].upsample(h) + + h = nonlinearity(h) + h = self.conv_out(h, gain=(self.learnable_gain + 1)) + return h + + +# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license. +# LICENSE is in incl_licenses directory. + + +class Snake(nn.Module): + """ + Implementation of a sine-based periodic activation function + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter + References: + - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snake(256) + >>> x = torch.randn(256) + >>> x = a1(x) + """ + + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + """ + Initialization. + INPUT: + - in_features: shape of the input + - alpha: trainable parameter + alpha is initialized to 1 by default, higher values = higher-frequency. + alpha will be trained along with the rest of your model. + """ + super(Snake, self).__init__() + self.in_features = in_features + + # Initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # Log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + else: # Linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + """ + Forward pass of the function. + Applies the function to the input elementwise. + Snake ∶= x + 1/a * sin^2 (xa) + """ + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # Line up with x to [B, C, T] + if self.alpha_logscale: + alpha = torch.exp(alpha) + x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + return x + + +class SnakeBeta(nn.Module): + """ + A modified Snake function which uses separate parameters for the magnitude of the periodic components + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + References: + - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snakebeta(256) + >>> x = torch.randn(256) + >>> x = a1(x) + """ + + def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): + """ + Initialization. + INPUT: + - in_features: shape of the input + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + alpha is initialized to 1 by default, higher values = higher-frequency. + beta is initialized to 1 by default, higher values = higher-magnitude. + alpha will be trained along with the rest of your model. + """ + super(SnakeBeta, self).__init__() + self.in_features = in_features + + # Initialize alpha + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: # Log scale alphas initialized to zeros + self.alpha = Parameter(torch.zeros(in_features) * alpha) + self.beta = Parameter(torch.zeros(in_features) * alpha) + else: # Linear scale alphas initialized to ones + self.alpha = Parameter(torch.ones(in_features) * alpha) + self.beta = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + """ + Forward pass of the function. + Applies the function to the input elementwise. + SnakeBeta ∶= x + 1/b * sin^2 (xa) + """ + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # Line up with x to [B, C, T] + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2) + + return x + + +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 +# LICENSE is in incl_licenses directory. + + +if "sinc" in dir(torch): + sinc = torch.sinc +else: + # This code is adopted from adefossez's julius.core.sinc under the MIT License + # https://adefossez.github.io/julius/julius/core.html + # LICENSE is in incl_licenses directory. + def sinc(x: torch.Tensor): + """ + Implementation of sinc, i.e. sin(pi * x) / (pi * x) + __Warning__: Different to julius.sinc, the input is multiplied by `pi`! + """ + return torch.where( + x == 0, + torch.tensor(1.0, device=x.device, dtype=x.dtype), + torch.sin(math.pi * x) / math.pi / x, + ) + + +# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License +# https://adefossez.github.io/julius/julius/lowpass.html +# LICENSE is in incl_licenses directory. +def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): # return filter [1,1,kernel_size] + even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + + # For kaiser window + delta_f = 4 * half_width + A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + if A > 50.0: + beta = 0.1102 * (A - 8.7) + elif A >= 21.0: + beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + + # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio + if even: + time = torch.arange(-half_size, half_size) + 0.5 + else: + time = torch.arange(kernel_size) - half_size + if cutoff == 0: + filter_ = torch.zeros_like(time) + else: + filter_ = 2 * cutoff * window * sinc(2 * cutoff * time) + """ + Normalize filter to have sum = 1, otherwise we will have a small leakage of the constant component in the input signal. + """ + filter_ /= filter_.sum() + filter = filter_.view(1, 1, kernel_size) + + return filter + + +class LowPassFilter1d(nn.Module): + def __init__( + self, + cutoff=0.5, + half_width=0.6, + stride: int = 1, + padding: bool = True, + padding_mode: str = "replicate", + kernel_size: int = 12, + ): + """ + kernel_size should be even number for stylegan3 setup, in this implementation, odd number is also possible. + """ + super().__init__() + if cutoff < -0.0: + raise ValueError("Minimum cutoff must be larger than zero.") + if cutoff > 0.5: + raise ValueError("A cutoff above 0.5 does not make sense.") + self.kernel_size = kernel_size + self.even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(self.even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.padding = padding + self.padding_mode = padding_mode + filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) + self.register_buffer("filter", filter) + + # Input [B, C, T] + def forward(self, x): + _, C, _ = x.shape + + if self.padding: + x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) + out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + + return out + + +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 +# LICENSE is in incl_licenses directory. + + +class UpSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.stride = ratio + self.pad = self.kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 + filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size) + self.register_buffer("filter", filter) + + # x: [B, C, T] + def forward(self, x): + _, C, _ = x.shape + + x = F.pad(x, (self.pad, self.pad), mode="replicate") + x = self.ratio * F.conv_transpose1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + x = x[..., self.pad_left : -self.pad_right] + + return x + + +class DownSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.lowpass = LowPassFilter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + stride=ratio, + kernel_size=self.kernel_size, + ) + + def forward(self, x): + xx = self.lowpass(x) + + return xx + + +# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 +# LICENSE is in incl_licenses directory. + + +class Activation1d(nn.Module): + def __init__( + self, + activation, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + ): + super().__init__() + self.up_ratio = up_ratio + self.down_ratio = down_ratio + self.act = activation + self.upsample = UpSample1d(up_ratio, up_kernel_size) + self.downsample = DownSample1d(down_ratio, down_kernel_size) + + # x: [B,C,T] + def forward(self, x): + x = self.upsample(x) + x = self.act(x) + x = self.downsample(x) + + return x + + +# Adapted from https://github.com/jik876/hifi-gan under the MIT license. +# LICENSE is in incl_licenses directory. + + +class AttrDict(dict): + def __init__(self, *args, **kwargs): + super(AttrDict, self).__init__(*args, **kwargs) + self.__dict__ = self + + +def build_env(config, config_name, path): + t_path = os.path.join(path, config_name) + if config != t_path: + os.makedirs(path, exist_ok=True) + shutil.copyfile(config, os.path.join(path, config_name)) + + +# Adapted from https://github.com/jik876/hifi-gan under the MIT license. +# LICENSE is in incl_licenses directory. + + +def init_weights(m, mean=0.0, std=0.01): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + m.weight.data.normal_(mean, std) + + +def apply_weight_norm(m): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + legacy_weight_norm(m) + + +def get_padding(kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + +def load_checkpoint(filepath, device): + assert os.path.isfile(filepath) + print(f"Loading '{filepath}'") + checkpoint_dict = torch.load(filepath, map_location=device) + print("Complete.") + return checkpoint_dict + + +# Copyright (c) 2024 NVIDIA CORPORATION. +# Licensed under the MIT license. + +# Adapted from https://github.com/jik876/hifi-gan under the MIT license. +# LICENSE is in incl_licenses directory. + + +def load_hparams_from_json(path) -> AttrDict: + with open(path) as f: + data = f.read() + return AttrDict(json.loads(data)) + + +class AMPBlock1(torch.nn.Module): + """ + AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer. + AMPBlock1 has additional self.convs2 that contains additional Conv1d layers with a fixed dilation=1 followed by each layer in self.convs1 + + Args: + h (AttrDict): Hyperparameters. + channels (int): Number of convolution channels. + kernel_size (int): Size of the convolution kernel. Default is 3. + dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5). + activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Default is None. + """ + + def __init__( + self, + h: AttrDict, + channels: int, + kernel_size: int = 3, + dilation: tuple = (1, 3, 5), + activation: str = None, + ): + super().__init__() + + self.h = h + + self.convs1 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + stride=1, + dilation=d, + padding=get_padding(kernel_size, d), + ) + ) + for d in dilation + ] + ) + self.convs1.apply(init_weights) + + self.convs2 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + stride=1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ) + for _ in range(len(dilation)) + ] + ) + self.convs2.apply(init_weights) + + self.num_layers = len(self.convs1) + len(self.convs2) # Total number of conv layers + + # Activation functions + if activation == "snake": + self.activations = nn.ModuleList( + [ + TorchActivation1d(activation=Snake(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ] + ) + elif activation == "snakebeta": + self.activations = nn.ModuleList( + [ + TorchActivation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ] + ) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + def forward(self, x): + acts1, acts2 = self.activations[::2], self.activations[1::2] + for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2): + xt = a1(x) + xt = c1(xt) + xt = a2(xt) + xt = c2(xt) + x = xt + x + + return x + + def remove_weight_norm(self): + for l in self.convs1: + remove_parametrizations(l, "weight") + for l in self.convs2: + remove_parametrizations(l, "weight") + + +class AMPBlock2(torch.nn.Module): + """ + AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer. + Unlike AMPBlock1, AMPBlock2 does not contain extra Conv1d layers with fixed dilation=1 + + Args: + h (AttrDict): Hyperparameters. + channels (int): Number of convolution channels. + kernel_size (int): Size of the convolution kernel. Default is 3. + dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5). + activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Default is None. + """ + + def __init__( + self, + h: AttrDict, + channels: int, + kernel_size: int = 3, + dilation: tuple = (1, 3, 5), + activation: str = None, + ): + super().__init__() + + self.h = h + + self.convs = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + stride=1, + dilation=d, + padding=get_padding(kernel_size, d), + ) + ) + for d in dilation + ] + ) + self.convs.apply(init_weights) + + self.num_layers = len(self.convs) # Total number of conv layers + + # Activation functions + if activation == "snake": + self.activations = nn.ModuleList( + [ + TorchActivation1d(activation=Snake(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ] + ) + elif activation == "snakebeta": + self.activations = nn.ModuleList( + [ + TorchActivation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ] + ) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + def forward(self, x): + for c, a in zip(self.convs, self.activations): + xt = a(x) + xt = c(xt) + x = xt + x + return x + + def remove_weight_norm(self): + for l in self.convs: + remove_weight_norm(l) + + +TorchActivation1d = Activation1d + + +class BigVGANv2( + torch.nn.Module, + PyTorchModelHubMixin, + library_name="bigvgan", + repo_url="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/NVIDIA/BigVGAN", + docs_url="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/NVIDIA/BigVGAN/blob/main/README.md", + pipeline_tag="audio-to-audio", + license="mit", + tags=["neural-vocoder", "audio-generation", "arxiv:2206.04658"], +): + """ + BigVGAN is a neural vocoder model that applies anti-aliased periodic activation for residual blocks (resblocks). + + Args: + h (AttrDict): Hyperparameters. + + Note: + - Ensure that the activation function is correctly specified in the hyperparameters (h.activation). + """ + + def __init__(self, h: AttrDict): + super().__init__() + self.h = h + + self.num_kernels = len(h.resblock_kernel_sizes) + self.num_upsamples = len(h.upsample_rates) + + # Pre-conv + self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)) + + # Define which AMPBlock to use. BigVGAN uses AMPBlock1 as default + if h.resblock == "1": + resblock_class = AMPBlock1 + elif h.resblock == "2": + resblock_class = AMPBlock2 + else: + raise ValueError(f"Incorrect resblock class specified in hyperparameters. Got {h.resblock}") + + # Transposed conv-based upsamplers. does not apply anti-aliasing + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)): + self.ups.append( + nn.ModuleList( + [ + weight_norm( + ConvTranspose1d( + h.upsample_initial_channel // (2**i), + h.upsample_initial_channel // (2 ** (i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ) + ] + ) + ) + + # Residual blocks using anti-aliased multi-periodicity composition modules (AMP) + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = h.upsample_initial_channel // (2 ** (i + 1)) + for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)): + self.resblocks.append(resblock_class(h, ch, k, d, activation=h.activation)) + + # Post-conv + activation_post = ( + Snake(ch, alpha_logscale=h.snake_logscale) + if h.activation == "snake" + else (SnakeBeta(ch, alpha_logscale=h.snake_logscale) if h.activation == "snakebeta" else None) + ) + if activation_post is None: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + self.activation_post = TorchActivation1d(activation=activation_post) + + # Whether to use bias for the final conv_post. Default to True for backward compatibility + self.use_bias_at_final = h.get("use_bias_at_final", True) + self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final)) + + # Weight initialization + for i in range(len(self.ups)): + self.ups[i].apply(init_weights) + self.conv_post.apply(init_weights) + + # Final tanh activation. Defaults to True for backward compatibility + self.use_tanh_at_final = h.get("use_tanh_at_final", True) + + def forward(self, x): + # Pre-conv + x = self.conv_pre(x) + + for i in range(self.num_upsamples): + # Upsampling + for i_up in range(len(self.ups[i])): + x = self.ups[i][i_up](x) + # AMP blocks + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + # Post-conv + x = self.activation_post(x) + x = self.conv_post(x) + # Final tanh activation + if self.use_tanh_at_final: + x = torch.tanh(x) + else: + x = torch.clamp(x, min=-1.0, max=1.0) # Bound the output to [-1, 1] + + return x + + def remove_weight_norm(self): + try: + # print("Removing weight norm...") + for l in self.ups: + for l_i in l: + remove_parametrizations(l_i, "weight") + for l in self.resblocks: + l.remove_weight_norm() + remove_parametrizations(self.conv_pre, "weight") + remove_parametrizations(self.conv_post, "weight") + except ValueError: + print("[INFO] Model already removed weight norm. Skipping!") + pass + + # Additional methods for huggingface_hub support + def _save_pretrained(self, save_directory: Path) -> None: + """Save weights and config.json from a Pytorch model to a local directory.""" + + model_path = save_directory / "bigvgan_generator.pt" + torch.save({"generator": self.state_dict()}, model_path) + + config_path = save_directory / "config.json" + with open(config_path, "w") as config_file: + json.dump(self.h, config_file, indent=4) + + @classmethod + def _from_pretrained( + cls, + *, + model_id: str, + revision: str, + cache_dir: str, + force_download: bool, + proxies: Optional[Dict], + resume_download: bool, + local_files_only: bool, + token: Union[str, bool, None], + map_location: str = "cpu", # Additional argument + strict: bool = False, # Additional argument + **model_kwargs, + ): + """Load Pytorch pretrained weights and return the loaded model.""" + + # Download and load hyperparameters (h) used by BigVGAN + if os.path.isdir(model_id): + # print("Loading config.json from local directory") + config_file = os.path.join(model_id, "config.json") + else: + config_file = hf_hub_download( + repo_id=model_id, + filename="config.json", + revision=revision, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + token=token, + local_files_only=local_files_only, + ) + h = load_hparams_from_json(config_file) + + model = cls(h) + + # Download and load pretrained generator weight + if os.path.isdir(model_id): + # print("Loading weights from local directory") + model_file = os.path.join(model_id, "bigvgan_generator.pt") + else: + print(f"Loading weights from {model_id}") + model_file = hf_hub_download( + repo_id=model_id, + filename="bigvgan_generator.pt", + revision=revision, + cache_dir=cache_dir, + force_download=force_download, + proxies=proxies, + resume_download=resume_download, + token=token, + local_files_only=local_files_only, + ) + + checkpoint_dict = torch.load(model_file, map_location=map_location, weights_only=True) + + try: + model.load_state_dict(checkpoint_dict["generator"]) + except RuntimeError: + print( + f"[INFO] the pretrained checkpoint does not contain weight norm. Loading the checkpoint after removing weight norm!" + ) + model.remove_weight_norm() + model.load_state_dict(checkpoint_dict["generator"]) + + return model + + +def _load_bigvgan_v2( + vocoder_ckpt_path: str | None = None, + *, + vocoder_config: dict | None = None, +) -> BigVGANv2: + """Load BigVGAN-v2 from a local directory (avoids HubMixin/huggingface_hub API drift).""" + model_file = None if vocoder_ckpt_path is None else os.path.join(vocoder_ckpt_path, "bigvgan_generator.pt") + if vocoder_config is None: + if vocoder_ckpt_path is None: + raise ValueError("Either vocoder_ckpt_path or vocoder_config is required") + config_file = os.path.join(vocoder_ckpt_path, "config.json") + if not os.path.isfile(config_file) or not os.path.isfile(model_file): + raise FileNotFoundError(f"Expected config.json and bigvgan_generator.pt under {vocoder_ckpt_path}") + h = load_hparams_from_json(config_file) + else: + if model_file is not None and not os.path.isfile(model_file): + raise FileNotFoundError(f"BigVGAN checkpoint does not exist: {model_file}") + h = AttrDict(vocoder_config) + + model = BigVGANv2(h) + if vocoder_ckpt_path is not None: + assert model_file is not None + checkpoint_dict = torch.load(model_file, map_location="cpu", weights_only=True) + try: + model.load_state_dict(checkpoint_dict["generator"]) + except RuntimeError: + model.remove_weight_norm() + model.load_state_dict(checkpoint_dict["generator"]) + model.remove_weight_norm() + else: + model.remove_weight_norm() + return model + + +class AutoEncoderModule(nn.Module): + def __init__( + self, + *, + vae_ckpt_path: str | None = None, + vocoder_ckpt_path: str | None = None, + vocoder_config: dict | None = None, + mode: Literal["44k"], + need_vae_encoder: bool = True, + need_vae_decoder: bool = True, + ): + super().__init__() + if mode != "44k": + raise ValueError(f"Unknown model: {mode}") + self.vae: VAE = VAE(data_dim=128, embed_dim=40, hidden_dim=512).eval() + if vae_ckpt_path is not None: + vae_state_dict = torch.load(vae_ckpt_path, weights_only=True, map_location="cpu") + self.vae.load_state_dict(vae_state_dict) + self.vae.remove_weight_norm() + + if need_vae_decoder: + self.vocoder = _load_bigvgan_v2( + vocoder_ckpt_path, + vocoder_config=vocoder_config, + ).eval() + else: + del self.vae.decoder + + if not need_vae_encoder: + del self.vae.encoder + + for param in self.parameters(): + param.requires_grad = False + + def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution: + return self.vae.encode(x) + + def decode(self, z: torch.Tensor) -> torch.Tensor: + return self.vae.decode(z) + + def vocode(self, spec: torch.Tensor) -> torch.Tensor: + return self.vocoder(spec) + + +# Reference: # https://github.com/bytedance/Make-An-Audio-2 + + +def dynamic_range_compression_torch(x, C=1, clip_val=1e-5, *, norm_fn): + return norm_fn(torch.clamp(x, min=clip_val) * C) + + +def spectral_normalize_torch(magnitudes, norm_fn): + output = dynamic_range_compression_torch(magnitudes, norm_fn=norm_fn) + return output + + +class MelConverter(nn.Module): + def __init__( + self, + *, + sampling_rate: float, + n_fft: int, + num_mels: int, + hop_size: int, + win_size: int, + fmin: float, + fmax: float, + norm_fn, + ): + super().__init__() + self.sampling_rate = sampling_rate + self.n_fft = n_fft + self.num_mels = num_mels + self.hop_size = hop_size + self.win_size = win_size + self.fmin = fmin + self.fmax = fmax + self.norm_fn = norm_fn + + mel = librosa_mel_fn( + sr=self.sampling_rate, n_fft=self.n_fft, n_mels=self.num_mels, fmin=self.fmin, fmax=self.fmax + ) + mel_basis = torch.from_numpy(mel).float() + hann_window = torch.hann_window(self.win_size) + + self.register_buffer("mel_basis", mel_basis) + self.register_buffer("hann_window", hann_window) + + @property + def device(self): + return self.mel_basis.device + + def forward(self, waveform: torch.Tensor, center: bool = False) -> torch.Tensor: + waveform = waveform.clamp(min=-1.0, max=1.0).to(self.device) + + waveform = torch.nn.functional.pad( + waveform.unsqueeze(1), + [int((self.n_fft - self.hop_size) / 2), int((self.n_fft - self.hop_size) / 2)], + mode="reflect", + ) + waveform = waveform.squeeze(1) + + spec = torch.stft( + waveform, + self.n_fft, + hop_length=self.hop_size, + win_length=self.win_size, + window=self.hann_window, + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + + spec = torch.view_as_real(spec) + spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)).float() + spec = torch.matmul(self.mel_basis, spec) + spec = spectral_normalize_torch(spec, self.norm_fn) + + return spec + + +def get_mel_converter(mode: Literal["44k"]) -> MelConverter: + return MelConverter( + sampling_rate=44_100, + n_fft=2048, + num_mels=128, + hop_size=512, + win_size=2048, + fmin=0, + fmax=44100 / 2, + norm_fn=torch.log, + ) + + +"""Audio feature / VAE facade.""" + + +class MMAudioVAE(ModelMixin, ConfigMixin): + ignore_for_config = ["tod_vae_ckpt", "bigvgan_vocoder_ckpt", "device"] + + @register_to_config + def __init__( + self, + *, + tod_vae_ckpt: str | PathLike[str] | None = None, + bigvgan_vocoder_ckpt: str | PathLike[str] | None = None, + vocoder_config: dict[str, Any] | None = None, + mode: str = "44k", + need_vae_encoder: bool = True, + need_vae_decoder: bool = True, + scaling_factor: float = 1.0, + sample_rate: int = 44_100, + downsample_factor: int = 1_024, + device: str | torch.device | None = None, + ) -> None: + super().__init__() + self.mel_converter = get_mel_converter(mode) + self.tod = AutoEncoderModule( + vae_ckpt_path=tod_vae_ckpt, + vocoder_ckpt_path=bigvgan_vocoder_ckpt, + vocoder_config=vocoder_config, + mode=mode, + need_vae_encoder=need_vae_encoder, + need_vae_decoder=need_vae_decoder, + ) + self.scaling_factor = float(scaling_factor) + self.sample_rate = int(sample_rate) + self.downsample_factor = int(downsample_factor) + if device is not None: + self.to(device) + + def compile(self): + self.decode = torch.compile(self.decode) + self.vocode = torch.compile(self.vocode) + + def train(self, mode: bool = True) -> MMAudioVAE: + return super().train(False) + + def encode_audio(self, x) -> DiagonalGaussianDistribution: + assert self.tod is not None, "VAE is not loaded" + mel = self.mel_converter(x) + return self.tod.encode(mel) + + def vocode(self, mel: torch.Tensor) -> torch.Tensor: + assert self.tod is not None, "VAE is not loaded" + return self.tod.vocode(mel) + + def decode(self, z: torch.Tensor) -> torch.Tensor: + assert self.tod is not None, "VAE is not loaded" + return self.tod.decode(z) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + @apply_forward_hook + def wrapped_decode(self, z): + mel_decoded = self.decode(z.to(dtype=self.dtype)) + return self.vocode(mel_decoded) + + @apply_forward_hook + def wrapped_encode(self, audio): + dist = self.encode_audio(audio.to(dtype=self.dtype)) + return dist.mean + + +__all__ = ["MMAudioVAE"] diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 0e167812ad88..acf8f8be3fa4 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -44,6 +44,8 @@ from .transformer_ideogram4 import Ideogram4Transformer2DModel from .transformer_joyimage import JoyImageEditTransformer3DModel from .transformer_joyimage_edit_plus import JoyImageEditPlusTransformer3DModel + from .transformer_kandinsky6_sr import Kandinsky6SRTransformer3DModel + from .transformer_kandinsky6 import Kandinsky6Transformer3DModel from .transformer_kandinsky import Kandinsky5Transformer3DModel from .transformer_krea2 import Krea2Transformer2DModel from .transformer_longcat_audio_dit import LongCatAudioDiTTransformer diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py new file mode 100644 index 000000000000..cafc554628fb --- /dev/null +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -0,0 +1,1448 @@ +"""Kandinsky 6 Diffusers transformer.""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +import torch.nn.functional as functional +from ...configuration_utils import ConfigMixin, register_to_config +from ...hooks import MagCacheConfig +from ...hooks.hooks import HookRegistry, ModelHook, StateManager +from ...hooks.mag_cache import MagCacheState +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import ( + _CAN_USE_FLEX_ATTN, + AttentionBackendName, + dispatch_attention_fn, +) +from ..cache_utils import CacheMixin +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from torch import Tensor, nn +from torch.nn.attention.flex_attention import BlockMask + + +def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: + return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) + + +def apply_scale_shift_norm(norm, x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: + """AdaLN-style affine in fp32, cast back to ``x.dtype``.""" + if x.ndim > 2 and scale.ndim == 2: + shape = (scale.shape[0],) + (1,) * (x.ndim - 2) + (scale.shape[-1],) + scale, shift = scale.reshape(shape), shift.reshape(shape) + return (norm(x.float()) * (scale.float() + 1.0) + shift.float()).to(dtype=x.dtype) + + +def apply_gate_sum(x: Tensor, out: Tensor, gate: Tensor) -> Tensor: + """Residual gate in fp32, cast back to ``x.dtype``.""" + if x.ndim > 2 and gate.ndim == 2: + gate = gate.reshape((gate.shape[0],) + (1,) * (x.ndim - 2) + (gate.shape[-1],)) + return (x.float() + gate.float() * out.float()).to(dtype=x.dtype) + + +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 _local_patch(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Tensor: + T, H, W = shape + g1, g2, g3 = group_size + x = x.reshape(*x.shape[:dim], T // g1, g1, H // g2, g2, W // g3, g3, *x.shape[dim + 3 :]) + d = len(x.shape[:dim]) + x = x.permute(*range(d), d, d + 2, d + 4, d + 1, d + 3, d + 5, *range(d + 6, len(x.shape))) + return x.flatten(dim, dim + 2).flatten(dim + 1, dim + 3) + + +def _local_merge(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Tensor: + T, H, W = shape + g1, g2, g3 = group_size + x = x.reshape(*x.shape[:dim], T // g1, H // g2, W // g3, g1, g2, g3, *x.shape[dim + 2 :]) + d = len(x.shape[:dim]) + x = x.permute(*range(d), d, d + 3, d + 1, d + 4, d + 2, d + 5, *range(d + 6, len(x.shape))) + return x.flatten(dim, dim + 1).flatten(dim + 1, dim + 2).flatten(dim + 2, dim + 3) + + +def fractal_flatten(x: Tensor, rope: Tensor, shape: tuple, block_mask: bool = False): + if block_mask: + ps = 8 + x = _local_patch(x, shape, (1, ps, ps), dim=0) + rope = _local_patch(rope, shape, (1, ps, ps), dim=0) + return x.flatten(0, 1), rope.flatten(0, 1) + return x.flatten(0, 2), rope.flatten(0, 2) + + +def fractal_unflatten(x: Tensor, shape: tuple, block_mask: bool = False) -> Tensor: + if block_mask: + ps = 8 + x = x.reshape(-1, ps * ps, x.shape[-1]) + return _local_merge(x, shape, (1, ps, ps), dim=0) + return x.reshape(*shape, x.shape[-1]) + + +def fast_sta_nabla( + T: int, + H: int, + W: int, + wT: int = 3, + wH: int = 3, + wW: int = 3, + device: str | torch.device = "cuda", +) -> Tensor: + """Precomputes the Sliding Tile Attention (STA) boolean mask for nabla attention.""" + l = max(T, H, W) + r = torch.arange(l, dtype=torch.int16, device=device) + mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs() + + sta_t = mat[:T, :T].flatten() <= wT // 2 + sta_h = mat[:H, :H].flatten() <= wH // 2 + sta_w = mat[:W, :W].flatten() <= wW // 2 + + sta_hw = (sta_h.unsqueeze(1) * sta_w.unsqueeze(0)).reshape(H, H, W, W).transpose(1, 2).flatten() + sta = (sta_t.unsqueeze(1) * sta_hw.unsqueeze(0)).reshape(T, T, H * W, H * W).transpose(1, 2) + return sta.reshape(T * H * W, T * H * W) + + +def nabla_block_mask( + q: Tensor, + k: Tensor, + sta: Tensor, + thr: float = 0.9, + block_size: int = 64, +) -> BlockMask: + """Builds a dynamic nabla BlockMask from query/key statistics + STA prior.""" + B, h, S, D = q.shape + s1 = S // block_size + qa = q.reshape(B, h, s1, block_size, D).mean(-2) + ka = k.reshape(B, h, s1, block_size, D).mean(-2).transpose(-2, -1) + attn_map = torch.softmax((qa @ ka) / math.sqrt(D), dim=-1) + + vals, inds = attn_map.sort(-1) + mask = (vals.cumsum_(-1) >= 1 - thr).int().gather(-1, inds.argsort(-1)) + mask = torch.logical_or(mask, sta) + + kv_nb = mask.sum(-1).to(torch.int32) + kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) + return BlockMask.from_kv_blocks( + torch.zeros_like(kv_nb), + kv_inds, + kv_nb, + kv_inds, + BLOCK_SIZE=block_size, + mask_mod=None, + ) + + +class RoPE1D(nn.Module): + """1-D Rotary Position Embedding — used for text and audio sequences.""" + + def __init__( + self, + dim: int, + max_pos: int = 2048, + max_period: float = 10000.0, + freqs_scaling: float = 1.0, + ): + super().__init__() + self.dim = dim + self.max_pos = max_pos + self.max_period = max_period + self.freqs_scaling = freqs_scaling + freq = get_freqs(dim // 2, max_period) * freqs_scaling + self.register_buffer("args", torch.outer(torch.arange(max_pos, dtype=freq.dtype), freq), persistent=False) + + def forward(self, pos: Tensor) -> Tensor: + # RoPE tables are fp32; keep trig in fp32. + args = self.args[pos] # (seq_len, dim//2) + 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 reset_parameters(self) -> None: + freq = get_freqs(self.dim // 2, self.max_period).to(self.args.device) * self.freqs_scaling + self.args = torch.outer(torch.arange(self.max_pos, dtype=freq.dtype, device=freq.device), freq) + + +class RoPE3D(nn.Module): + """3-D Rotary Position Embedding — used for video spatial-temporal tokens (T, H, W).""" + + 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 + for i, (d, mp) in enumerate(zip(axes_dims, max_pos)): + freq = get_freqs(d // 2, max_period) + self.register_buffer(f"args_{i}", torch.outer(torch.arange(mp, dtype=freq.dtype), freq), persistent=False) + + def forward( + self, + shape: tuple, + pos: list[Tensor], + scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), + ) -> Tensor: + T, H, W = shape + args_t = getattr(self, "args_0")[pos[0]] / scale_factor[0] # (T, d//2) + args_h = getattr(self, "args_1")[pos[1]] / scale_factor[1] # (H, d//2) + args_w = getattr(self, "args_2")[pos[2]] / scale_factor[2] # (W, d//2) + + args = torch.cat( + [ + args_t.view(T, 1, 1, -1).expand(T, H, W, -1), + args_h.view(1, H, 1, -1).expand(T, H, W, -1), + args_w.view(1, 1, W, -1).expand(T, H, W, -1), + ], + dim=-1, + ) + cos, sin = torch.cos(args), torch.sin(args) + rope = torch.stack([cos, -sin, sin, cos], dim=-1) # (T, H, W, total_dim, 4) + rope = rope.view(*rope.shape[:-1], 2, 2) # (T, H, W, total_dim, 2, 2) + return rope.unsqueeze(-4) # (T, H, W, 1, total_dim, 2, 2) + + def reset_parameters(self) -> None: + for i, (d, mp) in enumerate(zip(self.axes_dims, self.max_pos)): + freq = get_freqs(d // 2, self.max_period).to(getattr(self, f"args_{i}").device) + setattr(self, f"args_{i}", torch.outer(torch.arange(mp, dtype=freq.dtype, device=freq.device), freq)) + + +"""Diffusers MagCache adapter for the K6 multimodal transformer. + +Diffusers provides the public :class:`MagCacheConfig` and the stateful hook +infrastructure. K6's fused visual blocks carry video and audio streams +together, so the stock hook needs a small adapter to preserve both tensors +when a block is skipped. The adapter keeps the Diffusers ``enable_cache`` / +``disable_cache`` API and uses the Diffusers MagCache state and configuration. +""" + + +_HEAD_HOOK = "kandinsky6_mag_cache_head" +_BLOCK_HOOK = "kandinsky6_mag_cache_block" + + +def _streams_from_args(args: tuple[Any, ...], kwargs: dict[str, Any]) -> tuple[Tensor | None, Tensor | None]: + video = kwargs.get("vis", args[0] if args else None) + audio = kwargs.get("aud", args[1] if len(args) > 1 else None) + return video, audio + + +def _pack_streams(video: Tensor | None, audio: Tensor | None) -> Tensor | tuple[Tensor | None, Tensor | None]: + return video if audio is None else (video, audio) + + +def _add_residual( + video: Tensor | None, + audio: Tensor | None, + residual: Tensor | tuple[Tensor | None, Tensor | None], +) -> Tensor | tuple[Tensor | None, Tensor | None]: + if audio is None: + if not isinstance(residual, Tensor): + raise RuntimeError("K6 MagCache residual does not match a video-only block") + return video + residual if video is not None else video + if not isinstance(residual, tuple): + raise RuntimeError("K6 MagCache residual does not contain an audio stream") + video_residual, audio_residual = residual + return ( + video + video_residual if video is not None and video_residual is not None else video, + audio + audio_residual if audio_residual is not None else audio, + ) + + +def _residual( + output: Tensor | tuple[Tensor | None, Tensor | None], + input_video: Tensor | None, + input_audio: Tensor | None, +) -> Tensor | tuple[Tensor | None, Tensor | None]: + output_video, output_audio = _streams_from_args((output,), {}) if isinstance(output, Tensor) else output + if input_audio is None: + if output_video is None or input_video is None: + return output_video + return output_video - input_video + return ( + output_video - input_video if output_video is not None and input_video is not None else None, + output_audio - input_audio if output_audio is not None else None, + ) + + +def _should_compute(state: MagCacheState, config: MagCacheConfig, *, lane: int, num_steps: int) -> bool: + if config.calibrate: + return True + + ratio_index = state.step_index * 2 + lane + if config.mag_ratios is None or ratio_index >= len(config.mag_ratios): + current_scale = 1.0 + else: + current_scale = float(config.mag_ratios[ratio_index]) + + retention_step = int(config.retention_ratio * num_steps + 0.5) + if state.step_index < retention_step: + return True + + state.accumulated_ratio *= current_scale + state.accumulated_steps += 1 + state.accumulated_err += abs(1.0 - state.accumulated_ratio) + if ( + state.previous_residual is not None + and state.accumulated_err <= config.threshold + and state.accumulated_steps <= config.max_skip_steps + ): + return False + + state.accumulated_ratio = 1.0 + state.accumulated_steps = 0 + state.accumulated_err = 0.0 + return True + + +def _advance(state: MagCacheState, config: MagCacheConfig, num_steps: int) -> None: + state.step_index += 1 + if state.step_index < num_steps: + return + state.step_index = 0 + state.accumulated_ratio = 1.0 + state.accumulated_steps = 0 + state.accumulated_err = 0.0 + state.previous_residual = None + state.head_block_input = None + state.should_compute = True + state.calibration_ratios = [] + + +class _Kandinsky6MagCacheHeadHook(ModelHook): + _is_stateful = True + + def __init__(self, state_manager: StateManager, config: MagCacheConfig, num_steps: int): + super().__init__() + self.state_manager = state_manager + self.config = config + self.num_steps = num_steps + + @torch.compiler.disable + def new_forward(self, module: nn.Module, *args, **kwargs): + if self.state_manager._current_context is None: + self.state_manager.set_context("inference") + state: MagCacheState = self.state_manager.get_state() + video, audio = _streams_from_args(args, kwargs) + state.head_block_input = (video, audio) + lane = 1 if self.state_manager._current_context in {"uncond", "negative"} else 0 + state.should_compute = _should_compute(state, self.config, lane=lane, num_steps=self.num_steps) + + if not state.should_compute: + if state.previous_residual is None: + raise RuntimeError("K6 MagCache requested a skip before a residual was computed") + return _add_residual(video, audio, state.previous_residual) + return self.fn_ref.original_forward(*args, **kwargs) + + def reset_state(self, module: nn.Module): + self.state_manager.reset() + return module + + +class _Kandinsky6MagCacheBlockHook(ModelHook): + def __init__(self, state_manager: StateManager, config: MagCacheConfig, num_steps: int, is_tail: bool): + super().__init__() + self.state_manager = state_manager + self.config = config + self.num_steps = num_steps + self.is_tail = is_tail + + @torch.compiler.disable + def new_forward(self, module: nn.Module, *args, **kwargs): + if self.state_manager._current_context is None: + self.state_manager.set_context("inference") + state: MagCacheState = self.state_manager.get_state() + video, audio = _streams_from_args(args, kwargs) + + if not state.should_compute: + if self.is_tail: + _advance(state, self.config, self.num_steps) + return _pack_streams(video, audio) + + output = self.fn_ref.original_forward(*args, **kwargs) + if self.is_tail: + input_video, input_audio = state.head_block_input + state.previous_residual = _residual(output, input_video, input_audio) + _advance(state, self.config, self.num_steps) + return output + + def reset_state(self, module: nn.Module): + self.state_manager.reset() + return module + + +def _apply_kandinsky6_mag_cache(module: nn.Module, config: MagCacheConfig) -> None: + blocks = getattr(module, "visual_transformer_blocks", None) + if not isinstance(blocks, nn.ModuleList) or len(blocks) == 0: + raise ValueError("K6 MagCache requires a non-empty transformer.visual_transformer_blocks ModuleList") + + registry = HookRegistry.check_if_exists_or_initialize(module) + registry.remove_hook(_HEAD_HOOK, recurse=True) + registry.remove_hook(_BLOCK_HOOK, recurse=True) + state_manager = StateManager(MagCacheState, (), {}) + # K6 stores conditional and unconditional coefficients interleaved. The + # Diffusers pipeline sets separate cache contexts for the two CFG passes. + num_steps = config.num_inference_steps // 2 + + head_registry = HookRegistry.check_if_exists_or_initialize(blocks[0]) + head_registry.register_hook( + _Kandinsky6MagCacheHeadHook(state_manager, config, num_steps), + _HEAD_HOOK, + ) + for block in blocks[1:-1]: + block_registry = HookRegistry.check_if_exists_or_initialize(block) + block_registry.register_hook( + _Kandinsky6MagCacheBlockHook(state_manager, config, num_steps, is_tail=False), + _BLOCK_HOOK, + ) + if len(blocks) > 1: + tail_registry = HookRegistry.check_if_exists_or_initialize(blocks[-1]) + tail_registry.register_hook( + _Kandinsky6MagCacheBlockHook(state_manager, config, num_steps, is_tail=True), + _BLOCK_HOOK, + ) + + +class Kandinsky6MagCacheMixin(CacheMixin): + """Expose Diffusers' cache API for K6's dual-stream visual blocks.""" + + def enable_cache(self, config) -> None: + if not isinstance(config, MagCacheConfig): + return super().enable_cache(config) + if self.is_cache_enabled: + raise ValueError(f"Caching has already been enabled with {type(self._cache_config)}") + if config.num_inference_steps % 2 != 0: + raise ValueError( + "K6 MagCacheConfig.num_inference_steps must count interleaved CFG forwards (an even number)" + ) + _apply_kandinsky6_mag_cache(self, config) + self._cache_config = config + + def disable_cache(self) -> None: + if not isinstance(self._cache_config, MagCacheConfig): + return super().disable_cache() + registry = HookRegistry.check_if_exists_or_initialize(self) + registry.remove_hook(_HEAD_HOOK, recurse=True) + registry.remove_hook(_BLOCK_HOOK, recurse=True) + self._cache_config = None + + +"""Diffusers-style K6 transformer components for the source port. + +The assembler inlines this module after the portable attention, RoPE, and +tensor helpers. The classes use the same inner-module vocabulary as the +Diffusers Kandinsky5 transformer (for example ``in_layer``, ``modulation``, +``self_attention`` and ``feed_forward``). Checkpoint conversion maps the +native K6 names to this public Diffusers layout. +""" + + +_MASKED_ATTENTION_BACKENDS = { + "flash": "flash_varlen", + "_flash_3": "_flash_varlen_3", + "sage": "sage_varlen", + "native": "native", +} + + +class Kandinsky6TimeEmbeddings(nn.Module): + """Sinusoidal timestep embedding with a K6-compatible parameter layout.""" + + def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): + super().__init__() + if model_dim % 2: + raise ValueError("model_dim must be even") + self.register_buffer("freqs", get_freqs(model_dim // 2, max_period), persistent=False) + self.in_layer = nn.Linear(model_dim, time_dim) + self.activation = nn.SiLU() + self.out_layer = nn.Linear(time_dim, time_dim) + + def forward(self, time: Tensor) -> Tensor: + freqs = self.freqs.to(device=time.device, dtype=torch.float32) + args = torch.outer(time.float(), freqs) + embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + h = functional.linear(embed, self.in_layer.weight.float(), self.in_layer.bias.float()) + out = functional.linear(self.activation(h), self.out_layer.weight.float(), self.out_layer.bias.float()) + return out.to(dtype=self.out_layer.weight.dtype) + + +class Kandinsky6TextEmbeddings(nn.Module): + """Text projection and normalization used by K6 text branches.""" + + def __init__(self, text_dim: int, model_dim: int): + super().__init__() + self.in_layer = nn.Linear(text_dim, model_dim) + self.norm = nn.LayerNorm(model_dim, elementwise_affine=True) + + def forward(self, x: Tensor) -> Tensor: + x = self.in_layer(x.to(dtype=self.in_layer.weight.dtype)) + return self.norm(x).to(dtype=x.dtype) + + +class Kandinsky6VisualEmbeddings(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.to(dtype=self.in_layer.weight.dtype)) + + +class Kandinsky6Modulation(nn.Module): + """Zero-initialized 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) + nn.init.zeros_(self.out_layer.weight) + nn.init.zeros_(self.out_layer.bias) + + def forward(self, x: Tensor) -> Tensor: + out = functional.linear(self.activation(x.float()), self.out_layer.weight.float(), self.out_layer.bias.float()) + return out.to(dtype=x.dtype) + + +class Kandinsky6FeedForward(nn.Module): + """K6 bias-free GELU feed-forward network.""" + + def __init__(self, dim: int, ff_dim: int): + 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: Tensor) -> Tensor: + return self.out_layer(self.activation(self.in_layer(x))) + + +class Kandinsky6AttnProcessor: + """Diffusers attention-processor boundary for K6 QKV projections.""" + + _attention_backend = None + _parallel_config = None + + def __init__(self, attention_backend=None, parallel_config=None): + if not hasattr(functional, "scaled_dot_product_attention"): + raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0 or newer.") + self._masked = False + self._attention_backend = attention_backend + self._parallel_config = parallel_config + + @property + def _attention_backend(self): + return self.__attention_backend + + @_attention_backend.setter + def _attention_backend(self, backend): + if self._masked and backend is not None: + name = getattr(backend, "value", backend).lower() + backend = _MASKED_ATTENTION_BACKENDS.get(name, name) + backend = AttentionBackendName(backend) + self.__attention_backend = backend + + def __call__( + self, + attn: Kandinsky6Attention, + hidden_states: Tensor, + encoder_hidden_states: Tensor | None = None, + rotary_emb: Tensor | None = None, + rotary_emb_kv: Tensor | None = None, + sparse_params: dict[str, Any] | None = None, + attn_mask: Tensor | None = None, + ) -> Tensor: + query = attn.to_query(hidden_states) + if encoder_hidden_states is None: + key = attn.to_key(hidden_states) + value = attn.to_value(hidden_states) + else: + key = attn.to_key(encoder_hidden_states) + value = attn.to_value(encoder_hidden_states) + + query = query.reshape(*query.shape[:-1], attn.num_heads, -1) + key = key.reshape(*key.shape[:-1], attn.num_heads, -1) + value = value.reshape(*value.shape[:-1], attn.num_heads, -1) + query = attn.query_norm(query) + key = attn.key_norm(key) + + if rotary_emb is not None: + query = apply_rotary(query, rotary_emb).to(dtype=query.dtype) + if rotary_emb_kv is not None: + key = apply_rotary(key, rotary_emb_kv).to(dtype=key.dtype) + + # Diffusers attention kernels use [B,S,H,D]. Text and unbatched + # cross-attention conditions are promoted to a singleton batch while + # packed visual streams already carry their batch dimension. + query_was_batched = query.dim() == 4 + if not query_was_batched: + query, key, value = query.unsqueeze(0), key.unsqueeze(0), value.unsqueeze(0) + elif key.dim() == 3: + key, value = key.unsqueeze(0), value.unsqueeze(0) + + if sparse_params is not None: + q = query.transpose(1, 2).contiguous() + k = key.transpose(1, 2).contiguous() + v = value.transpose(1, 2).contiguous() + block_mask = nabla_block_mask( + q, + k, + sparse_params["sta_mask"], + thr=sparse_params["P"], + ) + if not _CAN_USE_FLEX_ATTN: + raise ValueError("Nabla attention requires PyTorch 2.5 or newer") + output = ( + torch.nn.attention.flex_attention.flex_attention( + q, + k, + v, + block_mask=block_mask, + ) + .transpose(1, 2) + .contiguous() + ) + else: + output = dispatch_attention_fn( + query, + key, + value, + attn_mask=attn_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + + if not query_was_batched: + output = output[0] + return attn.out_layer(output.flatten(-2, -1)) + + +class Kandinsky6Attention(nn.Module, AttentionModuleMixin): + """K6 attention with the Diffusers ``set_processor`` contract.""" + + _default_processor_cls = Kandinsky6AttnProcessor + _available_processors = [Kandinsky6AttnProcessor] + + def __init__( + self, + num_channels: int, + head_dim: int, + kv_dim: int | None = None, + text_token_padding: bool = False, + visual: bool = False, + processor: Kandinsky6AttnProcessor | None = None, + ): + super().__init__() + if num_channels % head_dim: + raise ValueError("num_channels must be divisible by head_dim") + kv_dim = kv_dim or num_channels + self.num_heads = num_channels // head_dim + self.to_query = nn.Linear(num_channels, num_channels) + self.to_key = nn.Linear(kv_dim, num_channels) + self.to_value = nn.Linear(kv_dim, num_channels) + self.query_norm = nn.RMSNorm(head_dim) + self.key_norm = nn.RMSNorm(head_dim) + self.out_layer = nn.Linear(num_channels, num_channels) + self.visual = visual + self.text_token_padding = text_token_padding + self.set_processor(processor or self._default_processor_cls()) + if self.text_token_padding: + self.processor._masked = True + self.processor._attention_backend = AttentionBackendName.NATIVE + + def forward( + self, + hidden_states: Tensor, + encoder_hidden_states: Tensor | None = None, + attn_mask: Tensor | None = None, + rotary_emb: Tensor | None = None, + sparse_params: dict[str, Any] | None = None, + rope_q: Tensor | None = None, + rope_kv: Tensor | None = None, + **kwargs: Any, + ) -> Tensor: + # Native K6 blocks pass ``(hidden, rope, mask_or_sparse)`` for + # self-attention. Keep that call shape while exposing Diffusers' + # encoder_hidden_states/rotary_emb keyword boundary. + if attn_mask is None: + attn_mask = kwargs.get("attention_mask") + rotary_emb = rope_q if rope_q is not None else rotary_emb + rotary_emb_kv = rope_kv if rope_kv is not None else rotary_emb + return self.processor( + self, + hidden_states, + encoder_hidden_states=encoder_hidden_states, + rotary_emb=rotary_emb, + rotary_emb_kv=rotary_emb_kv, + sparse_params=sparse_params, + attn_mask=attn_mask, + ) + + +class Kandinsky6OutLayer(nn.Module): + """Projects visual hidden states back to packed latent patches.""" + + def __init__(self, model_dim: int, time_dim: int, visual_dim: int, patch_size: tuple[int, int, int]): + super().__init__() + self.patch_size = patch_size + self.modulation = Kandinsky6Modulation(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) + + def forward(self, visual_embed: Tensor, time_embed: Tensor) -> Tensor: + shift, scale = torch.chunk(self.modulation(time_embed), 2, dim=-1) + condition_shape = (scale.shape[0],) + (1,) * (visual_embed.ndim - 2) + (scale.shape[-1],) + x = apply_scale_shift_norm( + self.norm, + visual_embed, + scale.reshape(condition_shape), + shift.reshape(condition_shape), + ) + x = self.out_layer(x) + + # batch_dim variant + if x.ndim == 5: + batch, duration, height, width = x.shape[:4] + p_t, p_h, p_w = self.patch_size + return ( + x.view(batch, duration, height, width, -1, p_t, p_h, p_w) + .permute(0, 1, 5, 2, 6, 3, 7, 4) + .flatten(1, 2) + .flatten(2, 3) + .flatten(3, 4) + ) + + duration, height, width, _ = x.shape + p_t, p_h, p_w = self.patch_size + return ( + x.view(duration, height, width, -1, p_t, p_h, p_w) + .permute(0, 4, 1, 5, 2, 6, 3) + .flatten(0, 1) + .flatten(1, 2) + .flatten(2, 3) + ) + + +class Kandinsky6OutLayerAudio(nn.Module): + """Projects audio hidden states back to audio latent channels.""" + + def __init__(self, model_dim: int, time_dim: int, audio_dim: int): + super().__init__() + self.modulation = Kandinsky6Modulation(time_dim, model_dim, 2) + self.norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.out_layer = nn.Linear(model_dim, audio_dim) + + def forward(self, audio_embed: Tensor, time_embed: Tensor) -> Tensor: + shift, scale = torch.chunk(self.modulation(time_embed), 2, dim=-1) + x = apply_scale_shift_norm(self.norm, audio_embed, scale, shift) + x = self.norm(x) + return self.out_layer(x) + + +class Kandinsky6TransformerEncoderBlock(nn.Module): + """Text self-attention + feed-forward block in Diffusers style.""" + + def __init__( + self, + model_dim: int, + time_dim: int, + ff_dim: int, + head_dim: int, + text_token_padding: bool = False, + ): + super().__init__() + self.text_modulation = Kandinsky6Modulation(time_dim, model_dim, 6) + self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.self_attention = Kandinsky6Attention(model_dim, head_dim, text_token_padding=text_token_padding) + self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.feed_forward = Kandinsky6FeedForward(model_dim, ff_dim) + + def forward(self, x: Tensor, time_embed: Tensor, rope: Tensor, attn_mask: Tensor | None = None) -> Tensor: + sa_params, ff_params = torch.chunk(self.text_modulation(time_embed), 2, dim=-1) + shift, scale, gate = torch.chunk(sa_params, 3, dim=-1) + x = apply_gate_sum( + x, + self.self_attention( + apply_scale_shift_norm(self.self_attention_norm, x, scale, shift), + rotary_emb=rope, + attn_mask=attn_mask, + ), + gate, + ) + shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) + return apply_gate_sum( + x, self.feed_forward(apply_scale_shift_norm(self.feed_forward_norm, x, scale, shift)), gate + ) + + +class Kandinsky6TransformerDecoderBlock(nn.Module): + """Visual self-attention, text cross-attention, and feed-forward block.""" + + def __init__( + self, + model_dim: int, + time_dim: int, + ff_dim: int, + head_dim: int, + text_token_padding: bool = False, + ): + super().__init__() + self.visual_modulation = Kandinsky6Modulation(time_dim, model_dim, 9) + self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.self_attention = Kandinsky6Attention(model_dim, head_dim, visual=True) + self.cross_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.cross_attention = Kandinsky6Attention( + model_dim, + head_dim, + kv_dim=model_dim, + text_token_padding=text_token_padding, + ) + self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.feed_forward = Kandinsky6FeedForward(model_dim, ff_dim) + + def forward( + self, vis: Tensor, text: Tensor, time_embed: Tensor, rope: Tensor, sparse_params: dict | None, attn_mask=None + ) -> Tensor: + sa_params, ca_params, ff_params = torch.chunk(self.visual_modulation(time_embed), 3, dim=-1) + shift, scale, gate = torch.chunk(sa_params, 3, dim=-1) + vis = apply_gate_sum( + vis, + self.self_attention( + apply_scale_shift_norm(self.self_attention_norm, vis, scale, shift), + rotary_emb=rope, + sparse_params=sparse_params, + ), + gate, + ) + shift, scale, gate = torch.chunk(ca_params, 3, dim=-1) + vis = apply_gate_sum( + vis, + self.cross_attention( + apply_scale_shift_norm(self.cross_attention_norm, vis, scale, shift), + encoder_hidden_states=text, + attn_mask=attn_mask, + ), + gate, + ) + shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) + return apply_gate_sum( + vis, self.feed_forward(apply_scale_shift_norm(self.feed_forward_norm, vis, scale, shift)), gate + ) + + +class Kandinsky6FusedTransformerDecoderBlock(nn.Module): + """Fused K6 video/audio block with cross-modal attention.""" + + def __init__( + self, + model_dim: int, + time_dim: int, + ff_dim: int, + head_dim: int, + model_dim_a: int, + time_dim_a: int, + ff_dim_a: int, + head_dim_a: int, + text_token_padding: bool = False, + ca_rope: bool = False, + cross_gates: bool = False, + fix_modulation: bool = False, + ): + super().__init__() + self.videoT = Kandinsky6TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim, text_token_padding) + self.audioT = Kandinsky6TransformerDecoderBlock( + model_dim_a, time_dim_a, ff_dim_a, head_dim_a, text_token_padding + ) + self.va_cross_attention = Kandinsky6Attention( + model_dim, + head_dim, + kv_dim=model_dim_a, + ) + self.av_cross_attention = Kandinsky6Attention( + model_dim_a, + head_dim_a, + kv_dim=model_dim, + ) + self.va_modulation = Kandinsky6Modulation( + time_dim, + model_dim if not cross_gates else model_dim * 2 + model_dim_a, + 1 if cross_gates else 3, + ) + self.av_modulation = Kandinsky6Modulation( + time_dim_a, + model_dim_a if not cross_gates else model_dim_a * 2 + model_dim, + 1 if cross_gates else 3, + ) + self.va_normalization = nn.LayerNorm(model_dim, elementwise_affine=False) + self.av_normalization = nn.LayerNorm(model_dim_a, elementwise_affine=False) + self.ca_rope = ca_rope + self.cross_gates = cross_gates + self.fix_modulation = fix_modulation + self.model_dim = model_dim + self.model_dim_a = model_dim_a + + def forward( + self, + vis: Tensor | None, + aud: Tensor | None, + text_v: Tensor, + text_a: Tensor, + time_embed: tuple[Tensor, Tensor], + vis_rope: Tensor | None, + aud_rope: Tensor | None, + sparse_params: dict | None, + attn_mask=None, + modality_mask=None, + av_gate_scale: float = 1.0, + va_gate_scale: float = 1.0, + ) -> tuple[Tensor | None, Tensor | None]: + fake_audio = modality_mask[0] if modality_mask is not None else 0 + fake_video = modality_mask[1] if modality_mask is not None else 0 + t_v, t_a = time_embed + if vis is not None: + sa_p, ca_p, ff_p = torch.chunk(self.videoT.visual_modulation(t_v), 3, dim=-1) + shift, scale, gate = torch.chunk(sa_p, 3, dim=-1) + vis = apply_gate_sum( + vis, + self.videoT.self_attention( + apply_scale_shift_norm(self.videoT.self_attention_norm, vis, scale, shift), + rotary_emb=vis_rope, + sparse_params=sparse_params, + ), + gate, + ).type_as(vis) + shift, scale, gate_v = torch.chunk(ca_p, 3, dim=-1) + vis_pre_ca = apply_scale_shift_norm(self.videoT.cross_attention_norm, vis, scale, shift) + vis_out_t = self.videoT.cross_attention( + vis_pre_ca, + encoder_hidden_states=text_v, + attn_mask=attn_mask, + ) + + if aud is not None: + sa_p, ca_p, ff_p_a = torch.chunk(self.audioT.visual_modulation(t_a), 3, dim=-1) + shift, scale, gate = torch.chunk(sa_p, 3, dim=-1) + aud = apply_gate_sum( + aud, + self.audioT.self_attention( + apply_scale_shift_norm(self.audioT.self_attention_norm, aud, scale, shift), + rotary_emb=aud_rope, + ), + gate, + ).type_as(aud) + shift, scale, gate_a = torch.chunk(ca_p, 3, dim=-1) + aud_pre_ca = apply_scale_shift_norm(self.audioT.cross_attention_norm, aud, scale, shift) + aud_out_t = self.audioT.cross_attention( + aud_pre_ca, + encoder_hidden_states=text_a, + attn_mask=attn_mask, + ) + aud = apply_gate_sum(aud, aud_out_t, gate_a).type_as(aud) + + if vis is not None: + t_va_mod = t_a if not self.fix_modulation else t_v + t_av_mod = t_v if not self.fix_modulation else t_a + va_params = self.va_modulation(t_va_mod) + av_params = self.av_modulation(t_av_mod) + if self.cross_gates: + va_shift, va_scale, va_gate = torch.split( + va_params, [self.model_dim, self.model_dim, self.model_dim_a], dim=-1 + ) + av_shift, av_scale, av_gate = torch.split( + av_params, [self.model_dim_a, self.model_dim_a, self.model_dim], dim=-1 + ) + else: + va_shift, va_scale, va_gate = torch.chunk(va_params, 3, dim=-1) + av_shift, av_scale, av_gate = torch.chunk(av_params, 3, dim=-1) + vis = apply_gate_sum(vis, vis_out_t, gate_v).type_as(vis) + vis_for_va = apply_scale_shift_norm(self.va_normalization, vis, va_scale, va_shift) + aud_for_av = apply_scale_shift_norm(self.av_normalization, aud, av_scale, av_shift) + rq_v = vis_rope if self.ca_rope else None + rk_a = aud_rope if self.ca_rope else None + vis_from_aud = ( + self.va_cross_attention( + vis_for_va, + encoder_hidden_states=aud_pre_ca, + rope_q=rq_v, + rope_kv=rk_a, + ) + * (1 - fake_audio) + * (1 - fake_video) + ) + aud_from_vis = ( + self.av_cross_attention( + aud_for_av, + encoder_hidden_states=vis_pre_ca, + rope_q=rk_a, + rope_kv=rq_v, + ) + * (1 - fake_audio) + * (1 - fake_video) + ) + vis = apply_gate_sum( + vis, vis_from_aud, (va_gate if not self.cross_gates else av_gate) * va_gate_scale + ).type_as(vis) + aud = apply_gate_sum( + aud, aud_from_vis, (av_gate if not self.cross_gates else va_gate) * av_gate_scale + ).type_as(aud) + elif vis is not None: + vis = apply_gate_sum(vis, vis_out_t, gate_v).type_as(vis) + + if vis is not None: + shift, scale, gate = torch.chunk(ff_p, 3, dim=-1) + vis = apply_gate_sum( + vis, + self.videoT.feed_forward(apply_scale_shift_norm(self.videoT.feed_forward_norm, vis, scale, shift)), + gate, + ).type_as(vis) + if aud is not None: + shift, scale, gate = torch.chunk(ff_p_a, 3, dim=-1) + aud = apply_gate_sum( + aud, + self.audioT.feed_forward(apply_scale_shift_norm(self.audioT.feed_forward_norm, aud, scale, shift)), + gate, + ).type_as(aud) + return vis, aud + + +class Kandinsky6RoPE1D(RoPE1D): + """Diffusers-exported name for the K6 one-dimensional RoPE.""" + + +class Kandinsky6RoPE3D(RoPE3D): + """Diffusers-exported name for the K6 three-dimensional RoPE.""" + + +class Kandinsky6Transformer3DModel( + Kandinsky6MagCacheMixin, + ModelMixin, + ConfigMixin, + PeftAdapterMixin, + FromOriginalModelMixin, + CacheMixin, + AttentionMixin, +): + """Standalone K6 T2V/T2VA transformer with Diffusers model mixins.""" + + _repeated_blocks = ["Kandinsky6TransformerEncoderBlock", "Kandinsky6TransformerDecoderBlock"] + _keep_in_fp32_modules = ["time_embeddings", "modulation", "visual_modulation", "text_modulation"] + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_visual_dim: int = 16, + out_visual_dim: int = 16, + in_text_dim: int = 3584, + in_text_dim2: int = 768, + time_dim: int = 1024, + patch_size: tuple = (1, 2, 2), + model_dim: int = 4096, + ff_dim: int = 16384, + num_text_blocks: int = 4, + num_visual_blocks: int = 60, + axes_dims: tuple = (32, 48, 48), + visual_cond: bool = True, + is_multimodal: bool = False, + in_audio_dim: int = 20, + out_audio_dim: int = 20, + model_dim_a: int | None = None, + time_dim_a: int | None = None, + ff_dim_a: int | None = None, + axes_dims_a: tuple | None = None, + audio_freqs_scaling: float = 1.0, + text_token_padding: bool = False, + scale_factor: tuple | list[float] = (1.0, 2.0, 2.0), + ca_rope: bool = False, + cross_gates: bool = False, + fix_modulation: bool = False, + visual_token_type_num_embeddings: int = 0, + magcache: dict[str, Any] | None = None, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.visual_cond = visual_cond + self.is_multimodal = is_multimodal + self.in_visual_dim = in_visual_dim + self.in_audio_dim = in_audio_dim + self.text_token_padding = text_token_padding + self.scale_factor = tuple(float(value) for value in scale_factor) + self.visual_token_type_num_embeddings = int(visual_token_type_num_embeddings or 0) + self._text_proj_cache: dict[tuple, object] = {} + head_dim = sum(axes_dims) + model_dim_a = model_dim_a or model_dim + time_dim_a = time_dim_a or time_dim + ff_dim_a = ff_dim_a or ff_dim + axes_dims_a = axes_dims_a or axes_dims + head_dim_a = sum(axes_dims_a) + + vis_in_dim = (2 * in_visual_dim + 1) if visual_cond else in_visual_dim + self.visual_embeddings = Kandinsky6VisualEmbeddings(vis_in_dim, model_dim, patch_size) + if self.visual_token_type_num_embeddings > 0: + self.visual_token_type_embeddings = nn.Embedding(self.visual_token_type_num_embeddings, model_dim) + self.visual_rope_embeddings = Kandinsky6RoPE3D(axes_dims) + self.out_layer = Kandinsky6OutLayer(model_dim, time_dim, out_visual_dim, patch_size) + + if not is_multimodal: + self.time_embeddings = Kandinsky6TimeEmbeddings(model_dim, time_dim) + self.text_embeddings = Kandinsky6TextEmbeddings(in_text_dim, model_dim) + self.pooled_text_embeddings = Kandinsky6TextEmbeddings(in_text_dim2, time_dim) + self.text_rope_embeddings = Kandinsky6RoPE1D(head_dim) + self.text_transformer_blocks = nn.ModuleList( + [ + Kandinsky6TransformerEncoderBlock(model_dim, time_dim, ff_dim, head_dim, text_token_padding) + for _ in range(num_text_blocks) + ] + ) + self.visual_transformer_blocks = nn.ModuleList( + [ + Kandinsky6TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim, text_token_padding) + for _ in range(num_visual_blocks) + ] + ) + else: + self.audio_embeddings = Kandinsky6TextEmbeddings(in_audio_dim, model_dim_a) + self.audio_rope_embeddings = Kandinsky6RoPE1D(head_dim_a, freqs_scaling=audio_freqs_scaling) + self.audio_out_layer = Kandinsky6OutLayerAudio(model_dim_a, time_dim_a, out_audio_dim) + for prefix, md, td, fd, hd in ( + ("video", model_dim, time_dim, ff_dim, head_dim), + ("audio", model_dim_a, time_dim_a, ff_dim_a, head_dim_a), + ): + setattr(self, f"{prefix}_time_embeddings", Kandinsky6TimeEmbeddings(md, td)) + setattr(self, f"{prefix}_text_embeddings", Kandinsky6TextEmbeddings(in_text_dim, md)) + setattr(self, f"{prefix}_pooled_text_embeddings", Kandinsky6TextEmbeddings(in_text_dim2, td)) + setattr(self, f"{prefix}_text_rope_embeddings", Kandinsky6RoPE1D(hd)) + setattr( + self, + f"{prefix}_text_transformer_blocks", + nn.ModuleList( + [ + Kandinsky6TransformerEncoderBlock(md, td, fd, hd, text_token_padding) + for _ in range(num_text_blocks) + ] + ), + ) + self.visual_transformer_blocks = nn.ModuleList( + [ + Kandinsky6FusedTransformerDecoderBlock( + model_dim, + time_dim, + ff_dim, + head_dim, + model_dim_a, + time_dim_a, + ff_dim_a, + head_dim_a, + text_token_padding, + ca_rope, + cross_gates, + fix_modulation, + ) + for _ in range(num_visual_blocks) + ] + ) + + def clear_text_proj_cache(self) -> None: + self._text_proj_cache.clear() + + def _project_pooled(self, prefix: str | None, pooled: Tensor) -> Tensor: + key = ("pe", prefix, pooled.data_ptr(), tuple(pooled.shape)) + hit = self._text_proj_cache.get(key) + if hit is not None: + return hit # type: ignore[return-value] + pe = ( + self.pooled_text_embeddings(pooled) + if prefix is None + else getattr(self, f"{prefix}_pooled_text_embeddings")(pooled) + ) + self._text_proj_cache[key] = pe + return pe + + def _project_text_tokens(self, prefix: str | None, text_embed: Tensor, pooled: Tensor) -> tuple[Tensor, Tensor]: + key = ("te", prefix, text_embed.data_ptr(), pooled.data_ptr(), tuple(text_embed.shape), tuple(pooled.shape)) + hit = self._text_proj_cache.get(key) + if hit is not None: + return hit # type: ignore[return-value] + te = ( + self.text_embeddings(text_embed) + if prefix is None + else getattr(self, f"{prefix}_text_embeddings")(text_embed) + ) + pe = self._project_pooled(prefix, pooled) + self._text_proj_cache[key] = (te, pe) + return te, pe + + def _time_embed(self, prefix: str | None, time: Tensor, pooled_proj: Tensor) -> Tensor: + return ( + self.time_embeddings(time) if prefix is None else getattr(self, f"{prefix}_time_embeddings")(time) + ) + pooled_proj + + @staticmethod + def _normalize_attn_mask(attn_mask: Tensor | None) -> Tensor | None: + if attn_mask is None: + return None + return attn_mask.unsqueeze(0) if attn_mask.dim() == 1 else attn_mask + + def _run_text_blocks( + self, prefix: str | None, te: Tensor, tm: Tensor, text_rope: Tensor, attn_mask: Tensor | None = None + ) -> Tensor: + blocks = self.text_transformer_blocks if prefix is None else getattr(self, f"{prefix}_text_transformer_blocks") + for block in blocks: + te = block(te, tm, text_rope, attn_mask) + return te + + def _encode_text( + self, + prefix: str, + text_embed: Tensor, + pooled: Tensor, + time: Tensor, + text_rope: Tensor, + attn_mask: Tensor | None = None, + ) -> tuple[Tensor, Tensor]: + te, pe = self._project_text_tokens(prefix, text_embed, pooled) + tm = self._time_embed(prefix, time, pe) + return self._run_text_blocks(prefix, te, tm, text_rope, attn_mask), tm + + def _encode_t2v( + self, text_embed: Tensor, pooled: Tensor, time: Tensor, text_rope: Tensor, attn_mask: Tensor | None = None + ) -> tuple[Tensor, Tensor]: + te, pe = self._project_text_tokens(None, text_embed, pooled) + tm = self._time_embed(None, time, pe) + return self._run_text_blocks(None, te, tm, text_rope, attn_mask), tm + + def _time_only(self, prefix: str | None, pooled: Tensor, time: Tensor) -> Tensor: + return self._time_embed(prefix, time, self._project_pooled(prefix, pooled)) + + def _embed_visual( + self, + x_video: Tensor, + visual_rope: Tensor, + sparse_params: dict | None, + *, + apply_fractal: bool = True, + visual_token_type_ids: Tensor | None = None, + ) -> tuple[Tensor, tuple, Tensor]: + if x_video.ndim == 4: + x_video = x_video.unsqueeze(0) + + visual_embed = self.visual_embeddings(x_video) + if hasattr(self, "visual_token_type_embeddings") and visual_token_type_ids is not None: + if visual_token_type_ids.ndim == 1: + visual_token_type_ids = visual_token_type_ids.unsqueeze(0) + token_types = self.visual_token_type_embeddings(visual_token_type_ids.to(device=visual_embed.device)) + token_types = token_types[:, :, None, None, :] + visual_embed = visual_embed + token_types + visual_shape = visual_embed.shape[-4:-1] + to_fractal = sparse_params["to_fractal"] if sparse_params and apply_fractal else False + if to_fractal: + visual_embed = _local_patch(visual_embed, visual_shape, (1, 8, 8), dim=1).flatten(1, 2) + visual_rope = _local_patch(visual_rope, visual_shape, (1, 8, 8), dim=0).flatten(0, 1) + else: + visual_embed = visual_embed.flatten(1, 3) + visual_rope = visual_rope.flatten(0, 2) + return visual_embed, visual_shape, visual_rope + + def _embed_audio(self, x_audio: Tensor, audio_rope: Tensor) -> tuple[Tensor, Tensor]: + if x_audio.ndim == 2: + x_audio = x_audio.unsqueeze(0) + embed = self.audio_embeddings(x_audio) + return embed, audio_rope + + def _run_visual_blocks_single( + self, + vis_embed: Tensor | None, + aud_embed: Tensor | None, + te: Tensor, + tm: Tensor, + vis_rope: Tensor | None, + aud_rope: Tensor | None, + sparse_params: dict | None, + attn_mask: Tensor | None = None, + ) -> tuple[Tensor | None, Tensor | None]: + for block in self.visual_transformer_blocks: + if self.is_multimodal: + if vis_embed is not None and aud_embed is None: + vis_embed, _ = block(vis_embed, None, te, te, (tm, tm), vis_rope, None, sparse_params, attn_mask) + elif aud_embed is not None and vis_embed is None: + _, aud_embed = block(None, aud_embed, te, te, (tm, tm), None, aud_rope, None, attn_mask) + else: + raise RuntimeError("single-modality fused path expects exactly one of video/audio") + elif vis_embed is not None: + vis_embed = block(vis_embed, te, tm, vis_rope, sparse_params, attn_mask) + else: + aud_embed = block(aud_embed, te, tm, aud_rope, None, attn_mask) + return vis_embed, aud_embed + + def _run_visual_blocks_fused( + self, + vis_embed: Tensor, + aud_embed: Tensor, + video_te: Tensor, + audio_te: Tensor, + video_tm: Tensor, + audio_tm: Tensor, + vis_rope: Tensor, + aud_rope: Tensor, + sparse_params: dict | None, + attn_mask: Tensor | None = None, + ) -> tuple[Tensor, Tensor]: + for block in self.visual_transformer_blocks: + vis_embed, aud_embed = block( + vis_embed, + aud_embed, + video_te, + audio_te, + (video_tm, audio_tm), + vis_rope, + aud_rope, + sparse_params, + attn_mask, + ) + return vis_embed, aud_embed + + def _project_video(self, vis_embed: Tensor, vis_shape: tuple, tm: Tensor, sparse_params: dict | None) -> Tensor: + if sparse_params and sparse_params.get("to_fractal"): + vis_embed = _local_merge( + vis_embed.reshape(vis_embed.shape[0], -1, 64, vis_embed.shape[-1]), + vis_shape, + (1, 8, 8), + dim=1, + ) + else: + vis_embed = vis_embed.reshape(-1, *vis_shape, vis_embed.shape[-1]) + return self.out_layer(vis_embed, tm) + + def _project_audio(self, aud_embed: Tensor, tm: Tensor) -> Tensor: + return self.audio_out_layer(aud_embed, tm) + + def _project_fused( + self, vis_embed: Tensor, aud_embed: Tensor, vis_shape: tuple, video_tm: Tensor, audio_tm: Tensor + ) -> tuple[Tensor, Tensor]: + video = self.out_layer(vis_embed.reshape(-1, *vis_shape, vis_embed.shape[-1]), video_tm) + return video, self.audio_out_layer(aud_embed, audio_tm) + + def forward( + self, + x_video: Tensor | None = None, + x_audio: Tensor | None = None, + text_embed: Tensor | list[Tensor] | None = None, + pooled_text_embed: Tensor | list[Tensor] | None = None, + time: Tensor | list[Tensor] | None = None, + visual_rope: Tensor | None = None, + audio_rope: Tensor | None = None, + text_rope: Tensor | list[Tensor] | None = None, + sparse_params: dict | None = None, + attention_mask: Tensor | None = None, + visual_token_type_ids: Tensor | None = None, + return_dict: bool = False, + **kwargs: Any, + ) -> Tensor | tuple[Tensor, Tensor] | Transformer2DModelOutput: + if text_embed is None or pooled_text_embed is None or time is None or text_rope is None: + raise ValueError("text_embed, pooled_text_embed, time, and text_rope are required") + both = x_video is not None and x_audio is not None and self.is_multimodal + attn_mask = self._normalize_attn_mask(attention_mask) + if not both: + te_in = text_embed[0] if isinstance(text_embed, list) else text_embed + pe_in = pooled_text_embed[0] if isinstance(pooled_text_embed, list) else pooled_text_embed + rope_in = text_rope[0] if isinstance(text_rope, list) else text_rope + t_in = time[0] if isinstance(time, list) else time + if self.is_multimodal: + prefix = "audio" if x_audio is not None else "video" + if isinstance(text_rope, list): + rope_in = text_rope[1] if prefix == "audio" else text_rope[0] + te, tm = self._encode_text(prefix, te_in, pe_in, t_in, rope_in, attn_mask) + else: + te, tm = self._encode_t2v(te_in, pe_in, t_in, rope_in, attn_mask) + if x_video is not None: + vis_embed, vis_shape, vis_rope = self._embed_visual( + x_video, visual_rope, sparse_params, visual_token_type_ids=visual_token_type_ids + ) + vis_embed, _ = self._run_visual_blocks_single( + vis_embed, None, te, tm, vis_rope, None, sparse_params, attn_mask + ) + result: Tensor | tuple[Tensor, Tensor] = self._project_video(vis_embed, vis_shape, tm, sparse_params) + else: + aud_embed, aud_rope = self._embed_audio(x_audio, audio_rope) + _, aud_embed = self._run_visual_blocks_single(None, aud_embed, te, tm, None, aud_rope, None, attn_mask) + result = self._project_audio(aud_embed, tm) + else: + te_v, pe_v = ( + (text_embed[0], pooled_text_embed[0]) + if isinstance(text_embed, list) + else (text_embed, pooled_text_embed) + ) + te_a, pe_a = ( + (text_embed[1], pooled_text_embed[1]) + if isinstance(text_embed, list) + else (text_embed, pooled_text_embed) + ) + rope_v, rope_a = (text_rope[0], text_rope[1]) if isinstance(text_rope, list) else (text_rope, text_rope) + t_v, t_a = (time[0], time[1]) if isinstance(time, list) else (time, time) + video_te, video_tm = self._encode_text("video", te_v, pe_v, t_v, rope_v, attn_mask) + audio_te, audio_tm = self._encode_text("audio", te_a, pe_a, t_a, rope_a, attn_mask) + vis_embed, vis_shape, vis_rope = self._embed_visual( + x_video, visual_rope, sparse_params, apply_fractal=False, visual_token_type_ids=visual_token_type_ids + ) + aud_embed, aud_rope = self._embed_audio(x_audio, audio_rope) + vis_embed, aud_embed = self._run_visual_blocks_fused( + vis_embed, + aud_embed, + video_te, + audio_te, + video_tm, + audio_tm, + vis_rope, + aud_rope, + sparse_params, + attn_mask, + ) + result = self._project_fused(vis_embed, aud_embed, vis_shape, video_tm, audio_tm) + + return Transformer2DModelOutput(sample=result) if return_dict else result + + +__all__ = [ + "Kandinsky6Transformer3DModel", + "Kandinsky6TransformerEncoderBlock", + "Kandinsky6TransformerDecoderBlock", + "Kandinsky6FusedTransformerDecoderBlock", +] diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py new file mode 100644 index 000000000000..a5f4fb35411a --- /dev/null +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -0,0 +1,2496 @@ +"""Kandinsky 6 SR transformer Diffusers component.""" + +from __future__ import annotations + +import dataclasses +import math +from collections.abc import Mapping +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any + +import torch +from ...configuration_utils import ConfigMixin, register_to_config +from ..modeling_utils import ModelMixin +from diffusers.utils import logging +from flash_attn import flash_attn_varlen_func, flash_attn_varlen_qkvpacked_func +from torch import Tensor, nn +from torch.nn.attention.flex_attention import BlockMask, flex_attention + +# Side of the local 8x8 token block used by fractal (NABLA) attention. +# Independent of the VAE spatial compression — do not swap for VAE_SPATIAL_FACTOR. +FRACTAL_BLOCK_SIZE = 8 + + +@torch.autocast(device_type="cuda", enabled=False) +def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: + """Compute sinusoidal frequency schedule for rotary embeddings. + + Args: + dim: Embedding dimension. + max_period: Maximum period of the sinusoidal frequencies. + + Returns: + Frequency tensor of shape ``(dim,)``. + """ + return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) + + +def fractal_flatten( + x: Tensor, + rope: Tensor, + cu_seqlens: Tensor, + shape: tuple[int, int, int], + *, + fractal: bool = False, +) -> tuple[Tensor, Tensor, Tensor]: + """Flatten spatial dimensions with optional fractal (local block) patching. + + Args: + x: Input tensor with spatial dimensions. + rope: Rotary position embedding tensor matching ``x`` layout. + cu_seqlens: Cumulative sequence lengths (per-frame counts). + shape: Spatial shape as ``(length, height, width)``. + fractal: If True, apply local 8x8 block patching before flattening. + + Returns: + Tuple of (flattened ``x``, flattened ``rope``, scaled ``cu_seqlens``). + """ + _, height, width = shape + if fractal: + block = FRACTAL_BLOCK_SIZE + x = local_patching(x, shape, (1, block, block), dim=0) + rope = local_patching(rope, shape, (1, block, block), dim=0) + x = x.flatten(0, 1) + rope = rope.flatten(0, 1) + else: + x = x.flatten(0, 2) + rope = rope.flatten(0, 2) + cu_seqlens = cu_seqlens * (height * width) + return x, rope, cu_seqlens + + +def fractal_unflatten( + x: Tensor, + cu_seqlens: Tensor, + shape: tuple[int, int, int], + *, + fractal: bool = False, +) -> tuple[Tensor, Tensor]: + """Unflatten spatial dimensions, inverse of ``fractal_flatten``. + + Args: + x: Flattened input tensor. + cu_seqlens: Scaled cumulative sequence lengths. + shape: Target spatial shape as ``(length, height, width)``. + fractal: If True, reverse local 8x8 block patching. + + Returns: + Tuple of (unflattened ``x``, rescaled ``cu_seqlens``). + """ + _, height, width = shape + if fractal: + block = FRACTAL_BLOCK_SIZE + x = x.reshape(-1, block**2, *x.shape[1:]) + x = local_merge(x, shape, (1, block, block), dim=0) + else: + x = x.reshape(*shape, *x.shape[1:]) + cu_seqlens = cu_seqlens // (height * width) + return x, cu_seqlens + + +def local_patching(x: Tensor, shape: tuple[int, int, int], group_size: tuple[int, int, int], dim: int = 0) -> Tensor: + """Rearrange tensor into local spatial patches. + + Groups neighboring elements along each spatial axis into patches, + producing a tensor with patch-count and intra-patch dimensions. + + Args: + x: Input tensor with spatial dimensions starting at ``dim``. + shape: Spatial shape as ``(duration, height, width)``. + group_size: Patch size per axis ``(g1, g2, g3)``. + dim: First spatial dimension index. + + Returns: + Patched tensor with shape ``(..., num_patches, patch_elems, ...)``. + """ + duration, height, width = shape + g1, g2, g3 = group_size + x = x.reshape(*x.shape[:dim], duration // g1, g1, height // g2, g2, width // g3, g3, *x.shape[dim + 3 :]) + x = x.permute( + *range(len(x.shape[:dim])), dim, dim + 2, dim + 4, dim + 1, dim + 3, dim + 5, *range(dim + 6, len(x.shape)) + ) + return x.flatten(dim, dim + 2).flatten(dim + 1, dim + 3) + + +def local_merge(x: Tensor, shape: tuple[int, int, int], group_size: tuple[int, int, int], dim: int = 0) -> Tensor: + """Reverse local patching, restoring the original spatial layout. + + Args: + x: Patched tensor with ``(num_patches, patch_elems)`` at ``dim``. + shape: Original spatial shape as ``(duration, height, width)``. + group_size: Patch size per axis ``(g1, g2, g3)``. + dim: First spatial dimension index. + + Returns: + Tensor with restored spatial dimensions. + """ + duration, height, width = shape + g1, g2, g3 = group_size + x = x.reshape(*x.shape[:dim], duration // g1, height // g2, width // g3, g1, g2, g3, *x.shape[dim + 2 :]) + x = x.permute( + *range(len(x.shape[:dim])), dim, dim + 3, dim + 1, dim + 4, dim + 2, dim + 5, *range(dim + 6, len(x.shape)) + ) + return x.flatten(dim, dim + 1).flatten(dim + 1, dim + 2).flatten(dim + 2, dim + 3) + + +@torch.compile() +@torch.no_grad() +def fast_sta_nabla(T: int, H: int, W: int, wT: int = 3, wH: int = 3, wW: int = 3, device: str = "cuda") -> Tensor: + """Build a static spatiotemporal neighborhood attention mask. + + Each token attends to spatial and temporal neighbors within + the given window sizes. + + Args: + T: Temporal dimension (number of frames). + H: Height dimension. + W: Width dimension. + wT: Temporal window size (odd). + wH: Height window size (odd). + wW: Width window size (odd). + device: Device for the output tensor. + + Returns: + Boolean mask of shape ``(T*H*W, T*H*W)``. + """ + max_dim = max(T, H, W) + r = torch.arange(0, max_dim, 1, dtype=torch.int16, device=device) + mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs() + sta_t, sta_h, sta_w = mat[:T, :T].flatten(), mat[:H, :H].flatten(), mat[:W, :W].flatten() + sta_t = sta_t <= wT // 2 + sta_h = sta_h <= wH // 2 + sta_w = sta_w <= wW // 2 + sta_hw = (sta_h.unsqueeze(1) * sta_w.unsqueeze(0)).reshape(H, H, W, W).transpose(1, 2).flatten() + sta = (sta_t.unsqueeze(1) * sta_hw.unsqueeze(0)).reshape(T, T, H * W, H * W).transpose(1, 2) + return sta.reshape(T * H * W, T * H * W) + + +@torch.compile(dynamic=True) +@torch.no_grad() +def nablaT_v2_doc( + q: Tensor, + k: Tensor, + seq: Tensor, + T: int, + H: int, + W: int, + *, + wT: int = 3, + wH: int = 3, + wW: int = 3, + thr: float = 0.9, + add_sta: bool = True, + method: str = "topcdf", +) -> BlockMask: + """Build a dynamic sparse attention BlockMask with document boundaries. + + Estimates an approximate attention map from block-averaged queries and + keys, then binarizes it to select the most relevant blocks per query. + + Args: + q: Query tensor of shape ``(B, heads, seq_len, dim)``. + k: Key tensor of shape ``(B, heads, seq_len, dim)``. + seq: Cumulative document lengths (e.g. [0, 31, 51, 66, 97] - video boundaries in seq_len). + T: Temporal dimension. + H: Height dimension. + W: Width dimension. + wT: Temporal window for static attention. + wH: Height window for static attention. + wW: Width window for static attention. + thr: Threshold for CDF cutoff (``topcdf``) or token fraction (``topk``). + add_sta: Whether to add static local attention. + method: Binarization method, ``"topcdf"`` or ``"topk"``. + + Returns: + Sparse ``BlockMask`` with block size 64. + """ + if method not in {"topcdf", "topk"}: + msg = f"nabla method should be topcdf or topk, got {method}" + raise ValueError(msg) + # Q/K are the authoritative execution tensors when this function is + # reached through attention. Keep every mask intermediate on their device. + device = q.device + seq = seq.to(device=device) + + # Map estimation + B, h, S, D = q.shape + qa = q.reshape(B, h, S // 64, 64, D).mean(-2) + ka = k.reshape(B, h, S // 64, 64, D).mean(-2).transpose(-2, -1) + attn_map = qa @ ka + + d = torch.diff(seq) + doc = ( + torch.eye(d.numel(), dtype=torch.bool, device=device) + .repeat_interleave(d * H * W, dim=0) + .repeat_interleave(d * H * W, dim=1) + ) + attn_map += doc.log() + attn_map = torch.softmax(attn_map / math.sqrt(D), dim=-1) + if method == "topcdf": + # Map binarization + vals, inds = attn_map.sort(-1) + cvals = vals.cumsum_(-1) + mask = (cvals >= 1 - thr).int() + mask = mask.gather(-1, inds.argsort(-1)) + else: + attn_map = attn_map.reshape(B * h * S // 64, S // 64) + dl = d.tolist() + start_row = 0 + mask = torch.zeros_like(attn_map) + for di in dl: + d_full = di * W * H * h * B + end_row = start_row + d_full + k = max(1, int(thr * di * W * H)) + group = attn_map[start_row:end_row, :] + _, topk_indices = torch.topk(group, k, dim=-1) + row_indices = torch.arange(start_row, end_row, device=mask.device).view(-1, 1) + mask[row_indices, topk_indices] = 1 + start_row = end_row + mask = mask.reshape(B, h, S // 64, S // 64) + + if add_sta: + sta = fast_sta_nabla(T, H, W, wT, wH, wW, device=device).unsqueeze_(0).unsqueeze_(0) + mask = torch.logical_or(mask, sta) + mask = torch.logical_and(mask, doc) + + # BlockMask creation + kv_nb = mask.sum(-1).to(torch.int32) + kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) + return BlockMask.from_kv_blocks(torch.zeros_like(kv_nb), kv_inds, kv_nb, kv_inds, BLOCK_SIZE=64, mask_mod=None) + + +def block_mask_from_bool(mask: Tensor) -> BlockMask: + """Convert a boolean block-attention matrix into a flex-attention ``BlockMask``. + + Args: + mask: Boolean tensor broadcastable to ``(B, heads, S // 64, S // 64)``; + ``True`` at ``(..., i, j)`` lets query block ``i`` attend key block ``j``. + + Returns: + Sparse ``BlockMask`` with block size 64 and no intra-block ``mask_mod``. + """ + kv_nb = mask.sum(-1).to(torch.int32) + kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) + return BlockMask.from_kv_blocks(torch.zeros_like(kv_nb), kv_inds, kv_nb, kv_inds, BLOCK_SIZE=64, mask_mod=None) + + +def build_framewise_causal_block_doc( + seq: Tensor, + H: int, + W: int, + *, + mf: int = 2, +) -> Tensor: + """Build the frame-wise causal block-attention structure (no content sparsity). + + Causal across frames (a frame attends only to current and past frames), fully + bidirectional within each ``mf``-frame group and within a frame, and isolated + across packed documents. Operates at 64-block granularity, so ``H * W`` is the + number of 64-token blocks per frame. Weight-independent: depends only on the + document layout, never on Q/K content. + + Args: + seq: Cumulative per-document frame counts, ``(num_docs + 1,)``. + H: Block-grid height per frame. + W: Block-grid width per frame (``H * W`` = 64-token blocks per frame). + mf: Multi-frame group size; frames in a group attend bidirectionally. + + Returns: + Boolean ``(S // 64, S // 64)`` block-attention matrix. + """ + device = seq.device + d = torch.diff(seq) + doc1 = ( + torch.eye(d.numel(), dtype=torch.bool, device=device) + .repeat_interleave(d, dim=0) + .repeat_interleave(d, dim=1) + .tril() + ) + group_sizes = [[c.sum().item() for c in torch.ones((dd,)).split(mf)] for dd in d] + cl = torch.tensor([x for xs in group_sizes for x in xs], dtype=torch.int32, device=device) + doc2 = ( + torch.eye(cl.numel(), dtype=torch.bool, device=device).repeat_interleave(cl, dim=0).repeat_interleave(cl, dim=1) + ) + return torch.logical_or(doc1, doc2).repeat_interleave(H * W, dim=0).repeat_interleave(H * W, dim=1) + + +@torch.compile(dynamic=True) +@torch.no_grad() +def nablaT_v2_doc_mfcausal( + q: Tensor, + k: Tensor, + seq: Tensor, + T: int, + H: int, + W: int, + *, + wT: int = 3, + wH: int = 3, + wW: int = 3, + thr: float = 0.9, + add_sta: bool = True, + mf: int = 2, +) -> BlockMask: + """Build a dynamic sparse attention BlockMask with multi-frame causal masking. + + Similar to ``nablaT_v2_doc`` but adds causal constraints across + multi-frame groups within each document. + + Args: + q: Query tensor of shape ``(B, heads, seq_len, dim)``. + k: Key tensor of shape ``(B, heads, seq_len, dim)``. + seq: Cumulative document lengths. + T: Temporal dimension. + H: Height dimension. + W: Width dimension. + wT: Temporal window for static attention. + wH: Height window for static attention. + wW: Width window for static attention. + thr: CDF threshold for binarization. + add_sta: Whether to add static local attention. + mf: Multi-frame group size for causal masking. + + Returns: + Sparse ``BlockMask`` with block size 64. + """ + # Q/K are the authoritative execution tensors. + device = q.device + seq = seq.to(device=device) + + # Map estimation + B, h, S, D = q.shape + qa = q.reshape(B, h, S // 64, 64, D).mean(-2) + ka = k.reshape(B, h, S // 64, 64, D).mean(-2).transpose(-2, -1) + attn_map = qa @ ka + + doc = build_framewise_causal_block_doc(seq, H, W, mf=mf) + attn_map += doc.log() + attn_map = torch.softmax(attn_map / math.sqrt(D), dim=-1) + + # Map binarization + vals, inds = attn_map.sort(-1) + cvals = vals.cumsum_(-1) + mask = (cvals >= 1 - thr).int() + mask = mask.gather(-1, inds.argsort(-1)) + + if add_sta: + sta = fast_sta_nabla(T, H, W, wT, wH, wW, device=device).unsqueeze_(0).unsqueeze_(0) + mask = torch.logical_or(mask, sta) + mask = torch.logical_and(mask, doc) + + return block_mask_from_bool(mask) + + +@torch.no_grad() +def framewise_causal_dense( + seq: Tensor, + H: int, + W: int, + *, + mf: int = 2, +) -> BlockMask: + """Build the dense (non-sparse) frame-wise causal ``BlockMask``. + + The exact, weight-independent causal reference: the frame-wise causal block + structure (causal across frames, bidirectional within an ``mf``-group and within + a frame) with no content-based sparsity. Runs through the same ``flex_attention`` + kernel as the NABLA path; only the ``BlockMask`` content differs. Used as the + streaming ground-truth mask and to validate that the NABLA causal selection stays + within this causal envelope. + + Args: + seq: Cumulative per-document frame counts, ``(num_docs + 1,)``. + H: Block-grid height per frame. + W: Block-grid width per frame (``H * W`` = 64-token blocks per frame). + mf: Multi-frame group size; frames in a group attend bidirectionally. + + Returns: + Dense causal ``BlockMask`` with block size 64. + """ + doc = build_framewise_causal_block_doc(seq, H, W, mf=mf) + return block_mask_from_bool(doc[None, None]) + + +"""Pure-torch helpers for the streaming KV-cache in self-attention. + +Kept free of ``flash_attn`` imports so the cache bookkeeping is importable and +testable on CPU. The actual attention kernel call lives in +``kandinsky_sr.model.nn``; this module only assembles its packed inputs. +""" + + +if TYPE_CHECKING: + from collections.abc import Iterator + + from torch import nn + + +def assemble_cached_attention_inputs( + new_key: Tensor, + new_value: Tensor, + cached_key: Tensor, + cached_value: Tensor, + cu_seqlens_new: Tensor, + cached_cu_seqlens: Tensor, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """Merge cached and new K/V into packed varlen attention inputs. + + For each packed sequence ``i`` the keys/values become ``[cache_i | new_i]``; + the query packing stays the new-token packing (queries are the freshly + arrived tokens that attend back over the cache plus themselves). + + Args: + new_key: New keys, packed ``(sum(new_lens), num_heads, head_dim)``. + new_value: New values, same packing as ``new_key``. + cached_key: Cached keys, packed ``(sum(cache_lens), num_heads, head_dim)``. + cached_value: Cached values, same packing as ``cached_key``. + cu_seqlens_new: Cumulative sequence lengths of the new tokens, ``(num_seqs + 1,)``. + cached_cu_seqlens: Cumulative sequence lengths of the cache, ``(num_seqs + 1,)``. + + Returns: + Tuple ``(key, value, cu_seqlens_q, cu_seqlens_k)`` where ``key``/``value`` + are the merged packed tensors, ``cu_seqlens_q`` is the query packing + (``= cu_seqlens_new``) and ``cu_seqlens_k`` the merged key packing. + """ + num_seqs = cu_seqlens_new.numel() - 1 + key_parts: list[Tensor] = [] + value_parts: list[Tensor] = [] + for i in range(num_seqs): + cache_start, cache_end = int(cached_cu_seqlens[i]), int(cached_cu_seqlens[i + 1]) + new_start, new_end = int(cu_seqlens_new[i]), int(cu_seqlens_new[i + 1]) + key_parts.append(cached_key[cache_start:cache_end]) + key_parts.append(new_key[new_start:new_end]) + value_parts.append(cached_value[cache_start:cache_end]) + value_parts.append(new_value[new_start:new_end]) + + key = torch.cat(key_parts, dim=0) + value = torch.cat(value_parts, dim=0) + + merged_lens = torch.diff(cached_cu_seqlens) + torch.diff(cu_seqlens_new) + cu_seqlens_k = torch.cat([cu_seqlens_new.new_zeros(1), torch.cumsum(merged_lens, dim=0)]).to(cu_seqlens_new.dtype) + cu_seqlens_q = cu_seqlens_new + return key, value, cu_seqlens_q, cu_seqlens_k + + +def evict_cached_kv( + cached_key: Tensor, + cached_value: Tensor, + cached_cu_seqlens: Tensor, + max_tokens_per_seq: int, +) -> tuple[Tensor, Tensor, Tensor]: + """Trim a packed KV-cache to its last ``max_tokens_per_seq`` tokens per sequence. + + Implements the fixed-length rolling window: for each packed sequence the + oldest tokens are dropped so at most ``max_tokens_per_seq`` most-recent ones + remain. Sequences already within the limit are kept whole. ``cu_seqlens`` is + rebuilt over the trimmed packing. Pure torch (no ``flash_attn``) so it stays + CPU-testable. + + Args: + cached_key: Cached keys, packed ``(sum(cache_lens), num_heads, head_dim)``. + cached_value: Cached values, same packing as ``cached_key``. + cached_cu_seqlens: Cumulative sequence lengths of the cache, ``(num_seqs + 1,)``. + max_tokens_per_seq: Maximum tokens to keep per packed sequence (the rolling + window length in tokens, i.e. ``cache_frames * H * W``). + + Returns: + Tuple ``(key, value, cu_seqlens)`` of the trimmed packed tensors and the + rebuilt cumulative sequence lengths. + """ + num_seqs = cached_cu_seqlens.numel() - 1 + key_parts: list[Tensor] = [] + value_parts: list[Tensor] = [] + kept_lens: list[int] = [] + for i in range(num_seqs): + start, end = int(cached_cu_seqlens[i]), int(cached_cu_seqlens[i + 1]) + keep_start = max(start, end - max_tokens_per_seq) + key_parts.append(cached_key[keep_start:end]) + value_parts.append(cached_value[keep_start:end]) + kept_lens.append(end - keep_start) + + key = torch.cat(key_parts, dim=0) + value = torch.cat(value_parts, dim=0) + lens = torch.tensor(kept_lens, device=cached_cu_seqlens.device) + cu_seqlens = torch.cat([cached_cu_seqlens.new_zeros(1), torch.cumsum(lens, dim=0)]).to(cached_cu_seqlens.dtype) + return key, value, cu_seqlens + + +def iter_self_attentions(model: nn.Module) -> Iterator[nn.Module]: + """Yield every submodule exposing the streaming KV-cache API. + + Duck-typed on the presence of ``reset_kv_cache`` so this module stays free of + ``flash_attn`` (which importing ``MultiheadSelfAttention`` would pull in) and + avoids a circular import with ``kandinsky_sr.model.nn``. + + Args: + model: Module to scan recursively. + + Yields: + Each self-attention submodule with a rolling KV-cache. + """ + for module in model.modules(): + if callable(getattr(module, "reset_kv_cache", None)): + yield module + + +def set_return_kv(model: nn.Module, *, flag: bool) -> None: + """Toggle ``return_kv`` on every self-attention submodule. + + Args: + model: Model whose self-attention layers to update. + flag: When ``True`` the next forward writes its merged K/V into the cache. + """ + for attn in iter_self_attentions(model): + attn.return_kv = flag + + +def reset_kv_caches(model: nn.Module) -> None: + """Clear the rolling KV-cache on every self-attention submodule. + + Args: + model: Model whose self-attention caches to reset. + """ + for attn in iter_self_attentions(model): + attn.reset_kv_cache() + + +@contextmanager +def kv_cache_session(model: nn.Module) -> Iterator[None]: + """Scope a streaming generation: caches are guaranteed clean before and after. + + Entry reset protects against state leaked by a previously crashed session; + the ``finally`` reset guarantees that even if generation dies mid-clip (e.g. + OOM) no layer keeps stale ``cached_k`` or a dangling ``return_kv=True``, so + subsequent non-streaming calls on the same model stay correct. + + Args: + model: Model whose self-attention caches to scope. + + Yields: + Nothing; run the chunked generation loop inside the block. + """ + reset_kv_caches(model) + try: + yield + finally: + reset_kv_caches(model) + + +@contextmanager +def kv_writer(model: nn.Module) -> Iterator[None]: + """Mark the enclosed forward as the cache writer (``return_kv=True``). + + ``return_kv`` is switched off in a ``finally`` so a crash inside the writer + forward cannot leave the model silently appending K/V on later calls. + + Args: + model: Model whose self-attention layers write K/V inside the block. + + Yields: + Nothing; run the writer forward (and its state capture) inside the block. + """ + set_return_kv(model, flag=True) + try: + yield + finally: + set_return_kv(model, flag=False) + + +# A captured per-step cache "slot": one ``(key, value, cu_seqlens)`` triple per +# self-attention layer (aligned with ``iter_self_attentions`` order), or ``None`` +# for an empty slot (no past frames cached yet at that denoising step). +LayerCache = tuple["Tensor | None", "Tensor | None", "Tensor | None"] +KVSlot = "list[LayerCache] | None" + + +def capture_kv_state(model: nn.Module) -> list[LayerCache]: + """Snapshot every self-attention's rolling KV-cache. + + Stores the current ``(cached_k, cached_v, cached_cu_seqlens)`` references for + each self-attention layer. The references are safe to hold: the attention + forward reassigns these attributes to freshly concatenated tensors rather + than mutating them in place, so a snapshot never aliases future writes. + + Args: + model: Model whose self-attention caches to snapshot. + + Returns: + Per-layer ``(key, value, cu_seqlens)`` triples in ``iter_self_attentions`` order. + """ + return [(attn.cached_k, attn.cached_v, attn.cached_cu_seqlens) for attn in iter_self_attentions(model)] + + +def restore_kv_state(model: nn.Module, slot: KVSlot) -> None: + """Load a captured slot (or clear) into every self-attention's cache. + + Args: + model: Model whose self-attention caches to overwrite. + slot: Per-layer ``(key, value, cu_seqlens)`` triples from + :func:`capture_kv_state`, or ``None`` to clear all caches (an empty + slot — no past frames yet). + """ + attns = list(iter_self_attentions(model)) + if slot is None: + for attn in attns: + attn.cached_k, attn.cached_v, attn.cached_cu_seqlens = None, None, None + return + for attn, (key, value, cu_seqlens) in zip(attns, slot, strict=True): + attn.cached_k, attn.cached_v, attn.cached_cu_seqlens = key, value, cu_seqlens + + +def evict_kv_slot(slot: KVSlot, max_tokens_per_seq: int) -> KVSlot: + """Apply fixed-length rolling eviction to a captured slot, per layer. + + Args: + slot: A captured per-step slot, or ``None`` (empty slot passes through). + max_tokens_per_seq: Rolling window length in tokens (``cache_frames * H * W``). + + Returns: + The evicted slot, or ``None`` if the input was ``None``. + """ + if slot is None: + return None + evicted: list[LayerCache] = [] + for key, value, cu_seqlens in slot: + if key is None: + evicted.append((None, None, None)) + else: + evicted.append(evict_cached_kv(key, value, cu_seqlens, max_tokens_per_seq)) + return evicted + + +"""Neural network building blocks for the diffusion transformer.""" + + +logger = logging.get_logger(__name__) + +try: + import flash_attn_interface # pyright: ignore + + FA3 = True +except ImportError: + FA3 = False + + +def _ensure_nabla_compatible_flex_bwd_configs() -> None: + """Guarantee flex_attention's backward keeps a 64-block-compatible config. + + ``nablaT_v2_doc`` builds a ``BLOCK_SIZE=64`` sparse mask. flex_attention's + backward lowering drops every autotune config whose block sizes don't divide + the mask block size (``SPARSE_BLOCK_SIZE % conf.block_* == 0``), checking each + config's OWN block sizes — ``kernel_options`` cannot inject one. On torch + builds whose head_dim=64 backward configs use ``BLOCK_N=128`` (e.g. the sm90 + default ``FlexBwDConfig(64, 128, 128, 64)``) every choice is filtered out and + backward compilation aborts with ``NoValidChoicesError``. This is hit only by + nabla TRAINING — forward-only validation never compiles the backward, which + is why nabla validation worked but training did not. + + Harmless for non-nabla (flash) runs: ``flex_attention`` is only reached from + ``attention_flex`` under ``sparse_params is not None``, so a flash run never + lowers a flex backward and the wrapper installed here is never invoked — it + only ever fires for a configured nabla run. + + Newer torch fixed this upstream by restricting the backward autotune blocks + to ``{32, 64}``. This shim back-ports the guarantee: it wraps the inductor + config provider and prepends an all-64 config whenever none of the returned + configs divide 64. Fully defensive — any structural mismatch (renamed method, + different config type) makes it a silent no-op rather than breaking startup, + and it never removes existing choices. + """ + try: + from torch._inductor.choices import InductorChoices # noqa: PLC0415 — optional, version-dependent internal + except Exception: + return + + original = getattr(InductorChoices, "get_flex_attention_bwd_configs", None) + if original is None or getattr(original, "nabla_patched", False): + return + + def divides_64(conf: Any) -> bool: + try: + return all(64 % b == 0 for b in (conf.block_m1, conf.block_n1, conf.block_m2, conf.block_n2)) + except Exception: + # Unknown config shape — assume compatible so we don't interfere. + return True + + def with_64_blocks(conf: Any) -> Any | None: + fields = {"block_m1": 64, "block_n1": 64, "block_m2": 64, "block_n2": 64} + try: + return dataclasses.replace(conf, **fields) + except Exception: + try: + return conf._replace(**fields) # NamedTuple fallback + except Exception: + return None + + def patched(self: Any, head_dim: int, dtype: Any, device_type: str = "cuda") -> Any: + configs = original(self, head_dim, dtype, device_type) + try: + if configs and not any(divides_64(c) for c in configs): + compat = with_64_blocks(configs[0]) + if compat is not None: + configs = [compat, *configs] + except Exception as exp: + logger.warning("nabla flex bwd-config shim no-op (config introspection failed): %s", exp) + return configs + + patched.nabla_patched = True # type: ignore[attr-defined] + try: + InductorChoices.get_flex_attention_bwd_configs = patched # type: ignore[method-assign] + except Exception as exp: + logger.warning("Could not install nabla flex bwd-config shim: %s", exp) + + +_ensure_nabla_compatible_flex_bwd_configs() + +flex = torch.compile(flex_attention, mode="max-autotune-no-cudagraphs", dynamic=True) + +# Without FA3 the only varlen entry points available are FA2's +# ``flash_attn_varlen_*_func``, whose C++ binding declares ``max_seqlen_q``/ +# ``max_seqlen_k`` as ``SymInt`` but does not support FakeTensors → dynamo +# can't trace through it. Skip dynamo on the affected methods so the +# surrounding code in ``TransformerDecoderBlock.forward`` (norms, QKV +# projections, modulation, FFN) still gets ``@torch.compile``-fused while +# only the FA call runs in eager. When FA3 is available the dedicated +# ``flash_attn_interface.flash_attn_varlen_func`` does register a custom op +# and traces cleanly, so this decorator becomes a no-op. +_disable_dynamo_if_fa2 = (lambda fn: fn) if FA3 else torch.compiler.disable + + +@torch.autocast(device_type="cuda", dtype=torch.float32) +def apply_scale_shift_norm(norm: nn.Module, x: Tensor, scale: Tensor, shift: Tensor, idx: Tensor) -> Tensor: + """Apply adaptive normalization with indexed scale and shift.""" + return norm(x) * (scale.index_select(0, idx) + 1.0) + shift.index_select(0, idx) + + +@torch.autocast(device_type="cuda", dtype=torch.float32) +def apply_gate_sum(x: Tensor, out: Tensor, gate: Tensor, idx: Tensor) -> Tensor: + """Add gated residual output to input using indexed gate values.""" + return x + gate.index_select(0, idx) * out + + +@torch.autocast(device_type="cuda", dtype=torch.float32) +def apply_scale_shift_norm_spatial(norm: nn.Module, x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: + """Apply adaptive normalization with per-token scale and shift. + + Unlike ``apply_scale_shift_norm`` which uses ``index_select`` for per-sample + params, this function takes spatially-varying (per-token) params directly. + + Args: + norm: Normalization layer (e.g. LayerNorm). + x: Input tensor of shape ``(tokens, dim)``. + scale: Per-token scale of shape ``(tokens, dim)``. + shift: Per-token shift of shape ``(tokens, dim)``. + """ + return norm(x) * (scale + 1.0) + shift + + +@torch.autocast(device_type="cuda", dtype=torch.float32) +def apply_gate_sum_spatial(x: Tensor, out: Tensor, gate: Tensor) -> Tensor: + """Add gated residual output to input using per-token gate values. + + Args: + x: Input tensor of shape ``(tokens, dim)``. + out: Residual output of shape ``(tokens, dim)``. + gate: Per-token gate of shape ``(tokens, dim)``. + """ + return x + gate * out + + +@torch.autocast(device_type="cuda", enabled=False) +def apply_rotary(x: Tensor, rope: Tensor) -> Tensor: + """Apply rotary position embeddings to input tensor.""" + x_ = x.reshape(*x.shape[:-1], -1, 1, 2).to(torch.float32) + x_out = rope[..., 0] * x_[..., 0] + rope[..., 1] * x_[..., 1] + return x_out.reshape(*x.shape) + + +class TimeEmbeddings(nn.Module): + """Sinusoidal time step embeddings projected through a two-layer MLP.""" + + def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0) -> None: + """Initialize time embeddings. + + Args: + model_dim: Dimension of sinusoidal encoding. Must be even. + time_dim: Output dimension after MLP projection. + max_period: Maximum period for frequency computation. + """ + super().__init__() + if model_dim % 2 != 0: + msg = "model_dim must be even" + raise ValueError(msg) + self.model_dim = model_dim + self.max_period = max_period + self.register_buffer("freqs", get_freqs(model_dim // 2, max_period), persistent=False) + + 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) + + @torch.autocast(device_type="cuda", dtype=torch.float32) + def forward(self, time: Tensor) -> tuple[Tensor, Tensor]: + """Encode time steps into embeddings. + + Args: + time: Diffusion time steps. + + Returns: + Tuple of (time embeddings, index tensor mapping tokens to embeddings). + """ + args = torch.outer(time, self.freqs.to(device=time.device)) + time_embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + time_embed = self.out_layer(self.activation(self.in_layer(time_embed))) + time_embed_idx = torch.arange(time_embed.shape[0], device=time_embed.device, dtype=torch.int32) + return time_embed, time_embed_idx + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + self.in_layer.reset_parameters() + self.out_layer.reset_parameters() + self.freqs = get_freqs(self.model_dim // 2, self.max_period) + + +class TextEmbeddings(nn.Module): + """Linear projection with layer normalization for text encoder outputs.""" + + def __init__(self, text_dim: int, model_dim: int) -> None: + """Initialize text embeddings. + + Args: + text_dim: Input dimension of text encoder hidden states. + model_dim: Output dimension after projection. + """ + super().__init__() + self.in_layer = nn.Linear(text_dim, model_dim, bias=True) + self.norm = nn.LayerNorm(model_dim, elementwise_affine=True) + + def forward(self, text_embed: Tensor) -> Tensor: + """Project and normalize text embeddings. + + Args: + text_embed: Raw text encoder hidden states. + + Returns: + Projected and normalized text embeddings. + """ + text_embed = self.in_layer(text_embed) + return self.norm(text_embed).type_as(text_embed) + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + self.in_layer.reset_parameters() + self.norm.reset_parameters() + + +class VisualEmbeddings(nn.Module): + """Patchify and project visual tokens into the transformer hidden space.""" + + def __init__( # noqa: PLR0913 + self, + visual_dim: int, + model_dim: int, + patch_size: tuple[int, int, int], + *, + use_lq_modulation: bool = False, + lq_channels: int = 16, + time_dim: int = 512, + zero_lq_in_main_path: bool = False, + ) -> None: + """Initialize visual embeddings. + + Args: + visual_dim: Number of input visual channels. + model_dim: Output dimension after projection. + patch_size: Patch size as (temporal, height, width). + use_lq_modulation: Whether to extract and encode LQ patches. + lq_channels: Number of LQ channels in the input (used only when + ``use_lq_modulation=True``). + time_dim: Output dimension for the LQ encoder (used only when + ``use_lq_modulation=True``). + zero_lq_in_main_path: Zero out LQ and mask channels in the main + ``in_layer`` input so LQ reaches the model only via per-block + ``ModulationLQ``. Only used when ``use_lq_modulation=True``. + """ + super().__init__() + self.patch_size = patch_size + self.use_lq_modulation = use_lq_modulation + self.zero_lq_in_main_path = zero_lq_in_main_path + self.in_layer = nn.Linear(math.prod(patch_size) * visual_dim, model_dim) + + if use_lq_modulation: + self.lq_channels = lq_channels + # input will be concat [hq(16) | lq(16) | mask(1)] + self.lq_start = visual_dim - lq_channels - 1 + self.lq_end = visual_dim - 1 + self.lq_encoder = LQEncoder( + input_dim=math.prod(patch_size) * lq_channels, + time_dim=time_dim, + ) + + def _patchify(self, x: Tensor, visual_cu_seqlens: Tensor) -> tuple[Tensor, Tensor]: + """Shared patchification logic for both full input and LQ-only slices. + + Args: + x: Input tensor of shape ``(duration, height, width, channels)``. + visual_cu_seqlens: Cumulative sequence lengths. + + Returns: + Tuple of (patchified tensor, updated cumulative sequence lengths). + """ + if self.patch_size[0] > 1: + idxs = torch.ones(x.shape[0], dtype=torch.int32, device=visual_cu_seqlens.device) + idxs[visual_cu_seqlens[:-1]] += self.patch_size[0] - 1 + x = torch.repeat_interleave(x, idxs, dim=0) + visual_cu_seqlens = visual_cu_seqlens + torch.arange( + visual_cu_seqlens.shape[0], device=visual_cu_seqlens.device, dtype=torch.int32 + ) + + duration, height, width, dim = x.shape + # [T, H, W, C] -> [T/pt, H/ph, W/pw, pt*ph*pw*C] + # Groups of (pt, ph, pw) neighboring pixels are concatenated into one token vector + x = ( + x.view( + duration // self.patch_size[0], + self.patch_size[0], + height // self.patch_size[1], + self.patch_size[1], + width // self.patch_size[2], + self.patch_size[2], + dim, + ) + .permute(0, 2, 4, 1, 3, 5, 6) + .flatten(3, 6) + ) + visual_cu_seqlens = visual_cu_seqlens // self.patch_size[0] + return x, visual_cu_seqlens + + def forward(self, x: Tensor, visual_cu_seqlens: Tensor) -> tuple[Tensor, Tensor] | tuple[Tensor, Tensor, Tensor]: + """Patchify input and project to model dimension. + + Args: + x: Visual input of shape ``(duration, height, width, channels)``. + visual_cu_seqlens: Cumulative sequence lengths for packed sequences. + + Returns: + When ``use_lq_modulation=False``: ``(projected_patches, cu_seqlens)``. + When ``use_lq_modulation=True``: ``(projected_patches, cu_seqlens, lq_encoded)``. + """ + if not self.use_lq_modulation: + # [T, H, W, 33] -> [T, H/2, W/2, 132] + # (T x H/2 x W/2) tokens, additive mix of patched HQ/LQ/mask, no cross-terms + patches, cu = self._patchify(x, visual_cu_seqlens) + # Linear(132 -> 1792) + return self.in_layer(patches), cu + + lq_slice = x[..., self.lq_start : self.lq_end] + if self.zero_lq_in_main_path: + # Zero out LQ+mask so in_layer sees [HQ | zeros | 0]; LQ only via ModulationLQ + x = x.clone() + x[..., self.lq_start :] = 0.0 + # Main path: Linear(132 -> 1792) on all 33 channels (LQ zeroed if zero_lq_in_main_path) + patches, cu = self._patchify(x, visual_cu_seqlens.clone()) + # LQ path: [T, H, W, 16] -> [T, H/2, W/2, 64] + # Per-token LQ encodings for bilinear modulation HQ * f(LQ) in decoder blocks + lq_patches, _ = self._patchify(lq_slice, visual_cu_seqlens.clone()) + # LQEncoder(64 -> 512) + lq_encoded = self.lq_encoder(lq_patches) + # Linear(132 -> 1792) + return self.in_layer(patches), cu, lq_encoded + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + self.in_layer.reset_parameters() + if self.use_lq_modulation: + self.lq_encoder.reset_parameters() + + +class RoPE1D(nn.Module): + """1D rotary position embeddings for text tokens.""" + + def __init__(self, dim: int, max_pos: int = 1024, max_period: float = 10000.0) -> None: + """Initialize 1D rotary position embeddings. + + Args: + dim: Embedding dimension. Must be even. + max_pos: Maximum number of positions to precompute. + max_period: Maximum period for frequency computation. + """ + super().__init__() + self.max_period = max_period + self.dim = dim + self.max_pos = max_pos + 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) + + @torch.autocast(device_type="cuda", enabled=False) + def forward(self, pos: Tensor) -> Tensor: + """Compute rotary embeddings for given positions. + + Args: + pos: Position indices. + + Returns: + Rotary embedding matrix with shape ``(*pos.shape, 1, dim//2, 2, 2)``. + """ + args = self.args[pos] + rope = torch.stack([torch.cos(args), -torch.sin(args), torch.sin(args), torch.cos(args)], dim=-1) + rope = rope.view(*rope.shape[:-1], 2, 2) + return rope.unsqueeze(-4) + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + freq = get_freqs(self.dim // 2, self.max_period).to(self.args.device) + pos = torch.arange(self.max_pos, dtype=freq.dtype, device=freq.device) + self.args = torch.outer(pos, freq) + + +class RoPE3D(nn.Module): + """3D rotary position embeddings for visual tokens along temporal, height, and width axes.""" + + def __init__( + self, + axes_dims: tuple[int, int, int], + max_pos: tuple[int, int, int] = (128, 128, 128), + max_period: float = 10000.0, + ) -> None: + """Initialize 3D rotary position embeddings. + + Args: + axes_dims: Per-axis embedding dimensions (temporal, height, width). + max_pos: Maximum positions per axis. + max_period: Maximum period for frequency computation. + """ + super().__init__() + self.axes_dims = axes_dims + self.max_pos = max_pos + self.max_period = max_period + + for i, (axes_dim, ax_max_pos) in enumerate(zip(axes_dims, max_pos, strict=False)): + freq = get_freqs(axes_dim // 2, max_period) + pos = torch.arange(ax_max_pos, dtype=freq.dtype) + self.register_buffer(f"args_{i}", torch.outer(pos, freq), persistent=False) + + @torch.autocast(device_type="cuda", enabled=False) + def forward( + self, + shape: tuple[int, int, int], + pos: tuple[Tensor, Tensor, Tensor], + scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), + ) -> Tensor: + """Compute 3D rotary embeddings for given positions. + + Args: + shape: Spatial dimensions as (duration, height, width). + pos: Position indices per axis (temporal, height, width). + scale_factor: Frequency scaling per axis. + + Returns: + Rotary embedding matrix broadcast to (duration, height, width, 1, dim//2, 2, 2). + """ + duration, height, width = shape + args_t = self.args_0[pos[0]] / 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, + ) + rope = torch.stack([torch.cos(args), -torch.sin(args), torch.sin(args), torch.cos(args)], dim=-1) + rope = rope.view(*rope.shape[:-1], 2, 2) + return rope.unsqueeze(-4) + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + for i, (axes_dim, ax_max_pos) in enumerate(zip(self.axes_dims, self.max_pos, strict=False)): + freq = get_freqs(axes_dim // 2, self.max_period).to(self.args_0.device) + pos = torch.arange(ax_max_pos, dtype=freq.dtype, device=freq.device) + setattr(self, f"args_{i}", torch.outer(pos, freq)) + + +class Modulation(nn.Module): + """Adaptive modulation layer producing scale, shift, and gate parameters from time embeddings.""" + + def __init__(self, time_dim: int, model_dim: int, num_params: int) -> None: + """Initialize modulation layer. + + Args: + time_dim: Input dimension of time embeddings. + model_dim: Per-parameter output dimension. + num_params: Number of modulation parameters to produce. + """ + 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_() + + @torch.autocast(device_type="cuda", dtype=torch.float32) + def forward(self, x: Tensor) -> Tensor: + """Compute modulation parameters from time embeddings. + + Args: + x: Time embedding tensor. + + Returns: + Concatenated modulation parameters. + """ + return self.out_layer(self.activation(x)) + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + self.out_layer.weight.data.zero_() + self.out_layer.bias.data.zero_() + + +class LQEncoder(nn.Module): + """Lightweight MLP that maps patchified LQ tokens to the time embedding space. + + Architecture: Linear -> SiLU -> Linear. Used to produce per-token conditioning + vectors from LQ patches for spatially-varying modulation in decoder blocks. + + Args: + input_dim: Dimension of patchified LQ tokens (patch_volume * lq_channels). + time_dim: Output dimension matching the time embedding space. + """ + + def __init__(self, input_dim: int, time_dim: int) -> None: + """Initialize LQ encoder layers. + + Args: + input_dim: Dimension of patchified LQ tokens (patch_volume * lq_channels). + time_dim: Output dimension matching the time embedding space. + """ + super().__init__() + self.in_layer = nn.Linear(input_dim, time_dim, bias=True) + self.activation = nn.SiLU() + self.out_layer = nn.Linear(time_dim, time_dim, bias=True) + + def forward(self, x: Tensor) -> Tensor: + """Encode patchified LQ tokens. + + Args: + x: Patchified LQ tokens of shape ``(total_tokens, input_dim)``. + + Returns: + Encoded LQ tokens of shape ``(total_tokens, time_dim)``. + """ + return self.out_layer(self.activation(self.in_layer(x))) + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + self.in_layer.reset_parameters() + self.out_layer.reset_parameters() + + +class ModulationLQ(nn.Module): + """Per-token modulation layer producing scale, shift, and gate from LQ encodings. + + Identical architecture to ``Modulation`` but designed for per-token (spatially-varying) + input from ``LQEncoder`` rather than per-sample time embeddings. + Zero-initialized so the model starts identically to a pretrained checkpoint + without LQ modulation. + + Args: + time_dim: Input dimension of LQ encodings (same as time embedding dim). + model_dim: Per-parameter output dimension. + num_params: Number of modulation parameters to produce. + """ + + def __init__(self, time_dim: int, model_dim: int, num_params: int) -> None: + """Initialize modulation layer with zero weights. + + Args: + time_dim: Input dimension of LQ encodings (same as time embedding dim). + model_dim: Per-parameter output dimension. + num_params: Number of modulation parameters to produce. + """ + 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: + """Compute per-token modulation parameters from LQ encodings. + + Args: + x: LQ encoding tensor of shape ``(total_tokens, time_dim)``. + + Returns: + Concatenated modulation parameters ``(total_tokens, num_params * model_dim)``. + """ + return self.out_layer(self.activation(x)) + + def reset_parameters(self) -> None: + """Re-initialize with zeros for identity behavior.""" + self.out_layer.weight.data.zero_() + self.out_layer.bias.data.zero_() + + +class MultiheadSelfAttention(nn.Module): + """Multi-head self-attention with flash attention and optional sparse flex attention.""" + + def __init__(self, num_channels: int, head_dim: int) -> None: + """Initialize multi-head self-attention. + + Args: + num_channels: Total number of channels. Must be divisible by head_dim. + head_dim: Dimension per attention head. + """ + super().__init__() + if num_channels % head_dim != 0: + msg = "num_channels must be divisible by head_dim" + raise ValueError(msg) + 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.cached_k: Tensor | None = None + self.cached_v: Tensor | None = None + self.cached_cu_seqlens: Tensor | None = None + self.return_kv = False + + def get_qkv(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: + """Project input into query, key, and value tensors. + + Args: + x: Input tensor. + + Returns: + Tuple of (query, key, value) reshaped to ``(seq_len, num_heads, head_dim)``. + """ + query = self.to_query(x) + key = self.to_key(x) + value = self.to_value(x) + + shape = query.shape[:-1] # for TP compatibility + query = query.reshape(*shape, self.num_heads, -1) + key = key.reshape(*shape, self.num_heads, -1) + value = value.reshape(*shape, self.num_heads, -1) + + return query, key, value + + def norm_qk(self, q: Tensor, k: Tensor) -> tuple[Tensor, Tensor]: + """Apply RMS normalization to query and key. + + Args: + q: Query tensor. + k: Key tensor. + + Returns: + Tuple of (normalized query, normalized key). + """ + q = self.query_norm(q.float()).type_as(q) + k = self.key_norm(k.float()).type_as(k) + return q, k + + @_disable_dynamo_if_fa2 + def scaled_dot_product_attention( + self, + query: Tensor, + key: Tensor, + value: Tensor, + cu_seqlens: Tensor, + *, + return_attn_probs: bool = False, + ) -> Tensor | tuple[Tensor, Tensor, None]: + """Compute self-attention using flash attention. + + When a KV-cache is present (``cached_k is not None``) the queries are the + new tokens and attention runs over ``[cache | new]`` keys/values via + :meth:`attend_with_kv_cache`. When ``return_kv`` is set, the (merged) K/V + are stored as the cache for the next streaming step. + + Args: + query: Query tensor. + key: Key tensor. + value: Value tensor. + cu_seqlens: Cumulative sequence lengths for packed sequences. + return_attn_probs: Whether to return attention log-sum-exp values. + + Returns: + Attention output, or tuple of (output, softmax_lse, None) if returning probs. + """ + if self.cached_k is not None: + return self.attend_with_kv_cache(query, key, value, cu_seqlens, return_attn_probs=return_attn_probs) + + if self.return_kv: + self.cached_k, self.cached_v, self.cached_cu_seqlens = key, value, cu_seqlens + + max_seqlen = torch.diff(cu_seqlens).max() + if FA3: + out, softmax_lse = flash_attn_interface.flash_attn_varlen_func( + q=query, + k=key, + v=value, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + ) + else: + query_key_value = torch.stack([query, key, value], dim=-3) + out, softmax_lse, _ = flash_attn_varlen_qkvpacked_func( + query_key_value, cu_seqlens, max_seqlen, return_attn_probs=True + ) + out = out.flatten(-2, -1) + + if return_attn_probs: + return out, softmax_lse, None + return out + + @torch.compiler.disable + def attend_with_kv_cache( + self, + query: Tensor, + key: Tensor, + value: Tensor, + cu_seqlens: Tensor, + *, + return_attn_probs: bool = False, + ) -> Tensor | tuple[Tensor, Tensor, None]: + """Dense flash attention consuming the rolling KV-cache. + + Queries are the new tokens; keys/values are the per-sequence concatenation + ``[cache | new]`` (with absolute RoPE already applied upstream). When + ``return_kv`` is set, the merged K/V become the cache for the next + streaming step. Kept out of ``torch.compile`` because the cache assembly + indexes ``cu_seqlens`` host-side; it only runs during streaming inference. + + Args: + query: New-token query tensor. + key: New-token key tensor (post-RoPE). + value: New-token value tensor. + cu_seqlens: Cumulative sequence lengths of the new tokens. + return_attn_probs: Whether to return attention log-sum-exp values. + + Returns: + Attention output for the new tokens, or tuple of (output, softmax_lse, + None) if returning probs. + """ + key, value, cu_seqlens_q, cu_seqlens_k = assemble_cached_attention_inputs( + key, value, self.cached_k, self.cached_v, cu_seqlens, self.cached_cu_seqlens + ) + if self.return_kv: + self.cached_k, self.cached_v, self.cached_cu_seqlens = key, value, cu_seqlens_k + + max_seqlen_q = torch.diff(cu_seqlens_q).max() + max_seqlen_k = torch.diff(cu_seqlens_k).max() + if FA3: + out, softmax_lse = flash_attn_interface.flash_attn_varlen_func( + q=query, + k=key, + v=value, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + ) + else: + out, softmax_lse, _ = flash_attn_varlen_func( + q=query, + k=key, + v=value, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + return_attn_probs=True, + ) + out = out.flatten(-2, -1) + + if return_attn_probs: + return out, softmax_lse, None + return out + + def reset_kv_cache(self) -> None: + """Clear the rolling KV-cache and stop emitting K/V.""" + self.cached_k = None + self.cached_v = None + self.cached_cu_seqlens = None + self.return_kv = False + + def attention_flex( + self, + query: Tensor, + key: Tensor, + value: Tensor, + sparse_params: dict[str, Any] | None = None, + *, + return_sparsity: bool = False, + ) -> Tensor | tuple[Tensor, float]: + """Compute sparse self-attention using flex attention with block masks. + + Args: + query: Query tensor. + key: Key tensor. + value: Value tensor. + sparse_params: Sparse attention configuration. + return_sparsity: Whether to return sparsity percentage. + + Returns: + Attention output, or tuple of (output, sparsity percentage) if requested. + """ + if self.cached_k is not None: + msg = "KV-cache on the flex/NABLA path is out of scope (dense flash only in M0.2; see LAY-429)." + raise NotImplementedError(msg) + + query = query.unsqueeze(0).transpose(1, 2).contiguous() + key = key.unsqueeze(0).transpose(1, 2).contiguous() + value = value.unsqueeze(0).transpose(1, 2).contiguous() + + t, h, w = sparse_params["visual_shape"] + # Token grid -> fractal 8x8 block grid; the block size matches + # fractal_flatten and is unrelated to the VAE spatial compression. + h, w = h // FRACTAL_BLOCK_SIZE, w // FRACTAL_BLOCK_SIZE + visual_seqlens = sparse_params["visual_seqlens"].to(device=query.device) + if sparse_params["attention_type"] == "dense_framewise_causal": + block_mask = framewise_causal_dense( + visual_seqlens, + h, + w, + mf=sparse_params["mf"], + ) + elif "mf" not in sparse_params: + block_mask = nablaT_v2_doc( + query, + key, + visual_seqlens, + t, + h, + w, + wT=sparse_params["wT"], + wH=sparse_params["wH"], + wW=sparse_params["wW"], + thr=sparse_params["P"], + add_sta=sparse_params["add_sta"], + method=sparse_params["method"], + ) + else: + block_mask = nablaT_v2_doc_mfcausal( + query, + key, + visual_seqlens, + t, + h, + w, + wT=sparse_params["wT"], + wH=sparse_params["wH"], + wW=sparse_params["wW"], + thr=sparse_params["P"], + add_sta=sparse_params["add_sta"], + mf=sparse_params["mf"], + ) + out = ( + flex( + query, + key, + value, + block_mask=block_mask, + kernel_options={"BLOCK_M": 64, "BLOCK_N": 64}, + ) + .transpose(1, 2) + .squeeze(0) + .contiguous() + ) + out = out.flatten(-2, -1) + + if return_sparsity: + sparsity = 100.0 * (1 - (1 - block_mask.sparsity() / 100) * (sparse_params["visual_seqlens"].shape[0] - 1)) + return out, sparsity + return out + + def forward( + self, + x: Tensor, + rope: Tensor, + cu_seqlens: Tensor, + sparse_params: dict[str, Any] | None = None, + ) -> Tensor: + """Run self-attention with rotary embeddings. + + Args: + x: Input tensor. + rope: Rotary position embeddings. + cu_seqlens: Cumulative sequence lengths for packed sequences. + sparse_params: Optional sparse attention parameters. + + Returns: + Self-attention output. + """ + query, key, value = self.get_qkv(x) + query, key = self.norm_qk(query, key) + query = apply_rotary(query, rope).type_as(query) + key = apply_rotary(key, rope).type_as(key) + + if sparse_params is not None: + out = self.attention_flex(query, key, value, sparse_params=sparse_params) + else: + out = self.scaled_dot_product_attention(query, key, value, cu_seqlens) + + return self.out_layer(out) + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + self.to_query.reset_parameters() + self.to_key.reset_parameters() + self.to_value.reset_parameters() + + self.out_layer.reset_parameters() + + self.query_norm.reset_parameters() + self.key_norm.reset_parameters() + + +class MultiheadCrossAttention(nn.Module): + """Multi-head cross-attention with flash attention.""" + + def __init__(self, num_channels: int, head_dim: int) -> None: + """Initialize multi-head cross-attention. + + Args: + num_channels: Total number of channels. Must be divisible by head_dim. + head_dim: Dimension per attention head. + """ + super().__init__() + if num_channels % head_dim != 0: + msg = "num_channels must be divisible by head_dim" + raise ValueError(msg) + 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) + + def get_qkv(self, x: Tensor, cond: Tensor) -> tuple[Tensor, Tensor, Tensor]: + """Project input and condition into query, key, and value tensors. + + Args: + x: Input tensor for query projection. + cond: Conditioning tensor for key and value projections. + + Returns: + Tuple of (query, key, value) reshaped to ``(seq_len, num_heads, head_dim)``. + """ + query = self.to_query(x) + key = self.to_key(cond) + value = self.to_value(cond) + + shape, cond_shape = query.shape[:-1], key.shape[:-1] # for TP compatibility + query = query.reshape(*shape, self.num_heads, -1) + key = key.reshape(*cond_shape, self.num_heads, -1) + value = value.reshape(*cond_shape, self.num_heads, -1) + + return query, key, value + + def norm_qk(self, q: Tensor, k: Tensor) -> tuple[Tensor, Tensor]: + """Apply RMS normalization to query and key. + + Args: + q: Query tensor. + k: Key tensor. + + Returns: + Tuple of (normalized query, normalized key). + """ + q = self.query_norm(q.float()).type_as(q) + k = self.key_norm(k.float()).type_as(k) + return q, k + + @_disable_dynamo_if_fa2 + def scaled_dot_product_attention( + self, + query: Tensor, + key: Tensor, + value: Tensor, + cu_seqlens: Tensor, + cond_cu_seqlens: Tensor, + *, + return_attn_probs: bool = False, + ) -> Tensor | tuple[Tensor, Tensor, None]: + """Compute cross-attention using flash attention. + + Args: + query: Query tensor. + key: Key tensor from conditioning input. + value: Value tensor from conditioning input. + cu_seqlens: Cumulative sequence lengths for query sequences. + cond_cu_seqlens: Cumulative sequence lengths for conditioning sequences. + return_attn_probs: Whether to return attention log-sum-exp values. + + Returns: + Attention output, or tuple of (output, softmax_lse, None) if returning probs. + """ + max_seqlen = torch.diff(cu_seqlens).max() + cond_max_seqlen = torch.diff(cond_cu_seqlens).max() + out, softmax_lse, _ = flash_attn_varlen_func( + query, key, value, cu_seqlens, cond_cu_seqlens, max_seqlen, cond_max_seqlen, return_attn_probs=True + ) + out = out.flatten(-2, -1) + + if return_attn_probs: + return out, softmax_lse, None + return out + + def forward(self, x: Tensor, cond: Tensor, cu_seqlens: Tensor, cond_cu_seqlens: Tensor) -> Tensor: + """Run cross-attention between input and conditioning. + + Args: + x: Input tensor for query computation. + cond: Conditioning tensor for key and value computation. + cu_seqlens: Cumulative sequence lengths for input sequences. + cond_cu_seqlens: Cumulative sequence lengths for conditioning sequences. + + Returns: + Cross-attention output. + """ + query, key, value = self.get_qkv(x, cond) + query, key = self.norm_qk(query, key) + + out = self.scaled_dot_product_attention(query, key, value, cu_seqlens, cond_cu_seqlens) + return self.out_layer(out) + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + self.to_query.reset_parameters() + self.to_key.reset_parameters() + self.to_value.reset_parameters() + + self.out_layer.reset_parameters() + + self.query_norm.reset_parameters() + self.key_norm.reset_parameters() + + +class FeedForward(nn.Module): + """Two-layer feed-forward network with GELU activation.""" + + def __init__(self, dim: int, ff_dim: int) -> None: + """Initialize feed-forward network. + + Args: + dim: Input and output dimension. + ff_dim: Hidden layer dimension. + """ + 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: Tensor) -> Tensor: + """Apply feed-forward transformation. + + Args: + x: Input tensor. + + Returns: + Transformed tensor. + """ + return self.out_layer(self.activation(self.in_layer(x))) + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + self.in_layer.reset_parameters() + self.out_layer.reset_parameters() + + +class OutLayer(nn.Module): + """Final output layer that unpatchifies and projects visual tokens to pixel space.""" + + def __init__(self, model_dim: int, time_dim: int, visual_dim: int, patch_size: tuple[int, int, int]) -> None: + """Initialize output layer. + + Args: + model_dim: Input dimension from the transformer. + time_dim: Dimension of time embeddings for modulation. + visual_dim: Number of output visual channels. + patch_size: Patch size as (temporal, height, width). + """ + super().__init__() + self.patch_size = patch_size + self.modulation = Modulation(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, + _text_embed: Tensor, + time_embed: Tensor, + visual_cu_seqlens: Tensor, + time_embed_idx: Tensor, + ) -> Tensor: + """Unpatchify and project visual embeddings to output channels. + + Args: + visual_embed: Visual token embeddings from the transformer. + _text_embed: Text embeddings (unused, reserved for interface compatibility). + time_embed: Time step embeddings for adaptive modulation. + visual_cu_seqlens: Cumulative sequence lengths for visual tokens. + time_embed_idx: Index mapping visual tokens to their time embeddings. + + Returns: + Denoised visual output tensor. + """ + shift, scale = torch.chunk(self.modulation(time_embed), 2, dim=-1) + visual_embed = apply_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, _dim = x.shape + x = ( + x.view( + duration, + height, + width, + -1, + self.patch_size[0], + self.patch_size[1], + self.patch_size[2], + ) + .permute(0, 4, 1, 5, 2, 6, 3) + .flatten(0, 1) + .flatten(1, 2) + .flatten(2, 3) + ) + visual_cu_seqlens = visual_cu_seqlens * self.patch_size[0] + + if self.patch_size[0] > 1: + idxs = torch.ones(duration * self.patch_size[0], dtype=torch.int32, device=visual_cu_seqlens.device) + idxs[visual_cu_seqlens[:-1]] -= self.patch_size[0] - 1 + x = torch.repeat_interleave(x, idxs, dim=0) + return x + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + self.norm.reset_parameters() + self.out_layer.reset_parameters() + self.modulation.reset_parameters() + + +"""Diffusion transformer architecture for text-conditioned visual generation.""" + + +class TransformerEncoderBlock(nn.Module): + """Transformer encoder block with self-attention, feed-forward, and adaptive modulation.""" + + def __init__(self, model_dim: int, time_dim: int, ff_dim: int, head_dim: int) -> None: + """Initialize encoder block layers. + + Args: + model_dim: Hidden dimension of the transformer. + time_dim: Dimension of time step embeddings. + ff_dim: Inner dimension of the feed-forward network. + head_dim: Per-head dimension for self-attention. + """ + super().__init__() + self.text_modulation = Modulation(time_dim, model_dim, 6) + + self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.self_attention = MultiheadSelfAttention(model_dim, head_dim) + + self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.feed_forward = FeedForward(model_dim, ff_dim) + + def forward( + self, + x: Tensor, + time_embed: Tensor, + rope: Tensor, + cu_seqlens: Tensor, + time_embed_idx: Tensor, + ) -> Tensor: + """Run self-attention and feed-forward with time-conditioned modulation. + + Args: + x: Input text embeddings. + time_embed: Time step embeddings for modulation. + rope: Rotary position embeddings. + cu_seqlens: Cumulative sequence lengths for packed sequences. + time_embed_idx: Index mapping tokens to their time embeddings. + + Returns: + Modulated text embeddings. + """ + 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 = apply_scale_shift_norm(self.self_attention_norm, x, scale, shift, time_embed_idx).type_as(x) + out = self.self_attention(out, rope, cu_seqlens) + x = apply_gate_sum(x, out, gate, time_embed_idx).type_as(x) + + shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) + out = apply_scale_shift_norm(self.feed_forward_norm, x, scale, shift, time_embed_idx).type_as(x) + out = self.feed_forward(out) + return apply_gate_sum(x, out, gate, time_embed_idx).type_as(x) + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + self.text_modulation.reset_parameters() + + self.self_attention_norm.reset_parameters() + self.self_attention.reset_parameters() + + self.feed_forward_norm.reset_parameters() + self.feed_forward.reset_parameters() + + +class TransformerDecoderBlock(nn.Module): + """Transformer decoder block with self-attention, cross-attention, and feed-forward. + + When ``use_lq_modulation=True``, a parallel ``ModulationLQ`` produces per-token + scale/shift/gate from LQ encodings, added to the per-sample time modulation. + Zero-initialized so that at init the block behaves identically to the baseline. + """ + + LQ_SUBLAYER_ORDER = ("self_attention", "cross_attention", "ffn") + + def __init__( + self, + model_dim: int, + time_dim: int, + ff_dim: int, + head_dim: int, + *, + use_lq_modulation: bool = False, + lq_modulation_sublayers: dict[str, bool] | None = None, + use_text: bool = True, + ) -> None: + """Initialize decoder block layers. + + Args: + model_dim: Hidden dimension of the transformer. + time_dim: Dimension of time step embeddings. + ff_dim: Inner dimension of the feed-forward network. + head_dim: Per-head dimension for attention layers. + use_lq_modulation: Whether to add per-token LQ modulation. + lq_modulation_sublayers: Which sub-layers to modulate. Keys: + ``self_attention``, ``cross_attention``, ``ffn``. Defaults to all True. + use_text: When False, the block has no text cross-attention: no + ``cross_attention``/``cross_attention_norm`` modules and the + modulation produces 6 params (self-attn + FFN) instead of 9. + Used by the text-free SR variant. + """ + super().__init__() + self.use_lq_modulation = use_lq_modulation + self.use_text = use_text + self.visual_modulation = Modulation(time_dim, model_dim, 9 if use_text else 6) + + if lq_modulation_sublayers is None: + lq_modulation_sublayers = dict.fromkeys(self.LQ_SUBLAYER_ORDER, True) + self.lq_sublayers = lq_modulation_sublayers + self.lq_num_active = sum(lq_modulation_sublayers[k] for k in self.LQ_SUBLAYER_ORDER) + + if use_lq_modulation and self.lq_num_active > 0: + self.lq_modulation = ModulationLQ(time_dim, model_dim, 3 * self.lq_num_active) + + self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.self_attention = MultiheadSelfAttention(model_dim, head_dim) + + if use_text: + self.cross_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.cross_attention = MultiheadCrossAttention(model_dim, head_dim) + + self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.feed_forward = FeedForward(model_dim, ff_dim) + + @torch.compile(dynamic=True) + def forward( + self, + visual_embed: Tensor, + text_embed: Tensor, + time_embed: Tensor, + rope: Tensor, + visual_cu_seqlens: Tensor, + text_cu_seqlens: Tensor, + time_embed_idx: Tensor, + sparse_params: dict[str, Any] | None, + lq_tokens: Tensor | None = None, + ) -> Tensor: + """Run self-attention, cross-attention, and feed-forward with modulation. + + Args: + visual_embed: Visual token embeddings. + text_embed: Encoded text embeddings used as cross-attention context. + time_embed: Time step embeddings for modulation. + rope: Rotary position embeddings for visual tokens. + visual_cu_seqlens: Cumulative sequence lengths for visual tokens. + text_cu_seqlens: Cumulative sequence lengths for text tokens. + time_embed_idx: Index mapping visual tokens to their time embeddings. + sparse_params: Optional sparse attention parameters. + lq_tokens: Per-token LQ encodings ``(total_tokens, time_dim)`` for + spatially-varying modulation. ``None`` when LQ modulation is off. + + Returns: + Updated visual embeddings. + """ + if self.use_text: + self_attn_params, cross_attn_params, ff_params = torch.chunk(self.visual_modulation(time_embed), 3, dim=-1) + else: + self_attn_params, ff_params = torch.chunk(self.visual_modulation(time_embed), 2, dim=-1) + cross_attn_params = None + + has_lq = self.use_lq_modulation and lq_tokens is not None and self.lq_num_active > 0 + lq_sa_shift = lq_sa_scale = lq_sa_gate = None + lq_ca_shift = lq_ca_scale = lq_ca_gate = None + lq_ff_shift = lq_ff_scale = lq_ff_gate = None + if has_lq: + # lq_modulation output: (total_tokens, 3 * num_active * model_dim) + lq_chunks = torch.chunk(self.lq_modulation(lq_tokens), self.lq_num_active, dim=-1) + # Each chunk: (total_tokens, 3 * model_dim) -> (shift, scale, gate) per-token + idx = 0 + if self.lq_sublayers["self_attention"]: + lq_sa_shift, lq_sa_scale, lq_sa_gate = torch.chunk(lq_chunks[idx], 3, dim=-1) + idx += 1 + if self.lq_sublayers["cross_attention"]: + lq_ca_shift, lq_ca_scale, lq_ca_gate = torch.chunk(lq_chunks[idx], 3, dim=-1) + idx += 1 + if self.lq_sublayers["ffn"]: + lq_ff_shift, lq_ff_scale, lq_ff_gate = torch.chunk(lq_chunks[idx], 3, dim=-1) + + # --- Self-attention --- + # Norm -> modulate by (scale_t + scale_lq + 1, shift_t + shift_lq) -> SelfAttn -> gate by (gate_t + gate_lq) + shift_t, scale_t, gate_t = torch.chunk(self_attn_params, 3, dim=-1) + with torch.autocast(device_type="cuda", dtype=torch.float32): + normed = self.self_attention_norm(visual_embed) + scale = scale_t.index_select(0, time_embed_idx) + shift = shift_t.index_select(0, time_embed_idx) + visual_out = normed * (scale + 1.0) + shift + if lq_sa_scale is not None: + visual_out = visual_out + lq_sa_scale * normed + lq_sa_shift + visual_out = visual_out.type_as(visual_embed) + visual_out = self.self_attention(visual_out, rope, visual_cu_seqlens, sparse_params) + with torch.autocast(device_type="cuda", dtype=torch.float32): + visual_embed = visual_embed + gate_t.index_select(0, time_embed_idx) * visual_out + if lq_sa_gate is not None: + visual_embed = visual_embed + lq_sa_gate * visual_out + + # --- Cross-attention --- (skipped when text-free) + # Norm -> modulate by (scale_t + scale_lq + 1, shift_t + shift_lq) -> CrossAttn -> gate by (gate_t + gate_lq) + if self.use_text: + shift_t, scale_t, gate_t = torch.chunk(cross_attn_params, 3, dim=-1) + with torch.autocast(device_type="cuda", dtype=torch.float32): + normed = self.cross_attention_norm(visual_embed) + scale = scale_t.index_select(0, time_embed_idx) + shift = shift_t.index_select(0, time_embed_idx) + visual_out = normed * (scale + 1.0) + shift + if lq_ca_scale is not None: + visual_out = visual_out + lq_ca_scale * normed + lq_ca_shift + visual_out = visual_out.type_as(visual_embed) + visual_out = self.cross_attention(visual_out, text_embed, visual_cu_seqlens, text_cu_seqlens) + with torch.autocast(device_type="cuda", dtype=torch.float32): + visual_embed = visual_embed + gate_t.index_select(0, time_embed_idx) * visual_out + if lq_ca_gate is not None: + visual_embed = visual_embed + lq_ca_gate * visual_out + + # --- Feed-forward --- + # Norm -> modulate by (scale_t + scale_lq + 1, shift_t + shift_lq) -> FF -> gate by (gate_t + gate_lq) + shift_t, scale_t, gate_t = torch.chunk(ff_params, 3, dim=-1) + with torch.autocast(device_type="cuda", dtype=torch.float32): + normed = self.feed_forward_norm(visual_embed) + scale = scale_t.index_select(0, time_embed_idx) + shift = shift_t.index_select(0, time_embed_idx) + visual_out = normed * (scale + 1.0) + shift + if lq_ff_scale is not None: + visual_out = visual_out + lq_ff_scale * normed + lq_ff_shift + visual_out = visual_out.type_as(visual_embed) + visual_out = self.feed_forward(visual_out) + with torch.autocast(device_type="cuda", dtype=torch.float32): + visual_embed = visual_embed + gate_t.index_select(0, time_embed_idx) * visual_out + if lq_ff_gate is not None: + visual_embed = visual_embed + lq_ff_gate * visual_out + return visual_embed + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + self.visual_modulation.reset_parameters() + if self.use_lq_modulation and self.lq_num_active > 0: + self.lq_modulation.reset_parameters() + + self.self_attention_norm.reset_parameters() + self.self_attention.reset_parameters() + + if self.use_text: + self.cross_attention_norm.reset_parameters() + self.cross_attention.reset_parameters() + + self.feed_forward_norm.reset_parameters() + self.feed_forward.reset_parameters() + + +class VideoAdapter(nn.Module): + """Lighter DiT that conditions a frozen backbone on degraded video input. + + The adapter has the same block design as the backbone but half as many blocks. + LQ latent tokens c = P(z̃) are processed through adapter blocks whose outputs + are injected into alternating backbone blocks via learnable γ_ℓ scalars. + + γ_ℓ are zero-initialized so the adapter has no effect at init, enabling + fine-tuning from a pretrained backbone checkpoint. + """ + + def __init__( + self, + in_visual_dim: int, + model_dim: int, + time_dim: int, + ff_dim: int, + head_dim: int, + patch_size: tuple[int, int, int], + num_blocks: int, + gamma: float = 0.0, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.gamma = gamma + patch_vol = math.prod(patch_size) + self.lq_proj = nn.Linear(patch_vol * in_visual_dim, model_dim) + self.blocks = nn.ModuleList( + [TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim) for _ in range(num_blocks)] + ) + + def _patchify(self, x: Tensor, visual_cu_seqlens: Tensor) -> tuple[Tensor, Tensor]: + """Patchify LQ latent with the same logic as VisualEmbeddings._patchify.""" + pt, ph, pw = self.patch_size + if pt > 1: + idxs = torch.ones(x.shape[0], dtype=torch.int32, device=visual_cu_seqlens.device) + idxs[visual_cu_seqlens[:-1]] += pt - 1 + x = torch.repeat_interleave(x, idxs, dim=0) + visual_cu_seqlens = visual_cu_seqlens + torch.arange( + visual_cu_seqlens.shape[0], device=visual_cu_seqlens.device, dtype=torch.int32 + ) + T, H, W, C = x.shape + x = x.view(T // pt, pt, H // ph, ph, W // pw, pw, C).permute(0, 2, 4, 1, 3, 5, 6).flatten(3, 6) + return x, visual_cu_seqlens // pt + + def forward( + self, + lq_visual: Tensor, + pre_patch_cu_seqlens: Tensor, + time_embed: Tensor, + text_embed: Tensor, + visual_rope: Tensor, + visual_cu_seqlens: Tensor, + text_cu_seqlens: Tensor, + visual_time_embed_idx: Tensor, + sparse_params: dict[str, Any] | None, + to_fractal: bool, + ) -> list[Tensor]: + """Process LQ tokens through adapter blocks, return γ-scaled features.""" + lq_patches, _ = self._patchify(lq_visual, pre_patch_cu_seqlens) + adapter_embed = self.lq_proj(lq_patches) + if to_fractal: + visual_shape = adapter_embed.shape[:-1] + adapter_embed = local_patching(adapter_embed, visual_shape, (1, 8, 8), dim=0) + adapter_embed = adapter_embed.flatten(0, 1) + else: + adapter_embed = adapter_embed.flatten(0, 2) + features = [] + for block in self.blocks: + adapter_embed = block( + adapter_embed, + text_embed, + time_embed, + visual_rope, + visual_cu_seqlens, + text_cu_seqlens, + visual_time_embed_idx, + sparse_params, + ) + features.append(self.gamma * adapter_embed) + return features + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters.""" + self.lq_proj.reset_parameters() + for block in self.blocks: + block.reset_parameters() + + +class Kandinsky6SRTransformer3DModel(ModelMixin, ConfigMixin): + """3D diffusion transformer with text-conditioned visual generation. + + Processes visual tokens through encoder (text) and decoder (visual) transformer blocks + with RoPE, adaptive modulation, and optional fractal attention sparsity. + """ + + @register_to_config + def __init__( + self, + in_visual_dim: int = 4, + in_text_dim: int = 3584, + in_text_dim2: int = 768, + time_dim: int = 512, + out_visual_dim: int = 4, + patch_size: tuple[int, int, int] = (1, 2, 2), + model_dim: int = 2048, + ff_dim: int = 5120, + num_text_blocks: int = 2, + num_visual_blocks: int = 32, + axes_dims: tuple[int, int, int] = (16, 24, 24), + *, + visual_cond: bool = False, + instruct_type: str | None = None, + attention_params: dict[str, Any] | None = None, + use_motion_score: bool = False, + use_lq_modulation: bool = False, + lq_modulation_sublayers: dict[str, bool] | None = None, + zero_lq_in_main_path: bool = False, + use_adapter: bool = False, + adapter_gamma: float = 0.0, + use_text: bool = True, + use_lq_noise_cond: bool = False, + attribute_overrides: Mapping[str, Any] | None = None, + sr_params: Mapping[str, Any] | None = None, + ) -> None: + """Initialize the 3D diffusion transformer. + + Args: + in_visual_dim: Number of input visual channels. + in_text_dim: Dimension of text encoder hidden states. + in_text_dim2: Dimension of pooled text embeddings. + time_dim: Dimension of time step embeddings. + out_visual_dim: Number of output visual channels. + patch_size: Patch size as (temporal, height, width). + model_dim: Hidden dimension of the transformer. + ff_dim: Inner dimension of feed-forward networks. + num_text_blocks: Number of text encoder blocks. + num_visual_blocks: Number of visual decoder blocks. + axes_dims: RoPE dimension split per axis (temporal, height, width). + visual_cond: Whether to use visual conditioning input. + instruct_type: Conditioning mode (e.g. ``"channel"``). + attention_params: Extra attention configuration. + use_motion_score: Whether to add motion score embeddings to time conditioning. + use_lq_modulation: Whether to enable per-block spatially-varying LQ modulation. + lq_modulation_sublayers: Which sub-layers to modulate with LQ. Keys: + ``self_attention``, ``cross_attention``, ``ffn``. Defaults to all True. + zero_lq_in_main_path: Zero out LQ+mask channels in ``in_layer`` input so + LQ reaches the model only via ``ModulationLQ``. Ablation experiment flag. + use_adapter: Whether to attach a VideoAdapter that conditions the backbone + on the degraded LQ input. The adapter has ``num_visual_blocks // 2`` + blocks and injects γ-scaled features into alternating backbone blocks. + adapter_gamma: Fixed scalar γ (not a trainable parameter) controlling + adapter influence on the backbone. Small values keep the backbone + dominant; larger values allow stronger structural correction. + Only used when ``use_adapter=True``. + use_text: When False, builds a text-free model with no text encoder, + text RoPE/blocks, pooled-text conditioning, or per-block cross- + attention. Default True keeps the standard architecture. + use_lq_noise_cond: When True, adds a second ``TimeEmbeddings`` that + encodes the LQ-endpoint noise fraction ``s`` (``lq_noise_scale``) + and sums it into the adaLN conditioning bus (like motion_score). + Its ``out_layer`` is zero-initialised so a checkpoint without the + module warm-starts as an exact no-op. + """ + super().__init__() + head_dim = sum(axes_dims) + self.in_visual_dim = in_visual_dim + self.instruct_type = instruct_type + self.model_dim = model_dim + self.patch_size = patch_size + self.visual_cond = visual_cond + self.attention_params = attention_params + self.use_lq_modulation = use_lq_modulation + if use_lq_modulation and instruct_type not in ("channel", "hybrid"): + msg = "use_lq_modulation requires instruct_type='channel' or 'hybrid'" + raise ValueError(msg) + + self.use_adapter = use_adapter + + visual_embed_dim = ( + 2 * in_visual_dim + 1 + if visual_cond or instruct_type in ("channel", "hybrid", "hybrid_anchor") + else in_visual_dim + ) + self.time_embeddings = TimeEmbeddings(model_dim, time_dim) + if use_motion_score: + self.motion_embeddings = TimeEmbeddings(model_dim, time_dim) + if use_lq_noise_cond: + self.lq_noise_embeddings = TimeEmbeddings(model_dim, time_dim) + self.zero_init_lq_noise_out() + + self.use_motion_score = use_motion_score + self.use_lq_noise_cond = use_lq_noise_cond + self.use_text = use_text + if use_adapter and not use_text: + msg = "use_adapter requires use_text=True: VideoAdapter blocks use text cross-attention" + raise ValueError(msg) + if use_text: + self.text_embeddings = TextEmbeddings(in_text_dim, model_dim) + self.pooled_text_embeddings = TextEmbeddings(in_text_dim2, time_dim) + else: + # Text-free: under the empty caption the pooled-text contribution + # `pooled_text_embeddings("")` is a constant added to the time embedding. + # build_text_free_dit.py bakes it into this bias; dropping it instead corrupts + # the time conditioning, so the model still needs it (cross-attention does not). + self.pooled_bias = nn.Parameter(torch.zeros(time_dim)) + self.visual_embeddings = VisualEmbeddings( + visual_embed_dim, + model_dim, + patch_size, + use_lq_modulation=use_lq_modulation, + lq_channels=in_visual_dim, + time_dim=time_dim, + zero_lq_in_main_path=zero_lq_in_main_path, + ) + + if use_text: + self.text_rope_embeddings = RoPE1D(head_dim) + self.text_transformer_blocks = nn.ModuleList( + [TransformerEncoderBlock(model_dim, time_dim, ff_dim, head_dim) for _ in range(num_text_blocks)] + ) + + self.visual_rope_embeddings = RoPE3D(axes_dims) + self.visual_transformer_blocks = nn.ModuleList( + [ + TransformerDecoderBlock( + model_dim, + time_dim, + ff_dim, + head_dim, + use_lq_modulation=use_lq_modulation, + lq_modulation_sublayers=lq_modulation_sublayers, + use_text=use_text, + ) + for _ in range(num_visual_blocks) + ] + ) + + if use_adapter: + self.adapter = VideoAdapter( + in_visual_dim=in_visual_dim, + model_dim=model_dim, + time_dim=time_dim, + ff_dim=ff_dim, + head_dim=head_dim, + patch_size=patch_size, + num_blocks=num_visual_blocks // 2, + gamma=adapter_gamma, + ) + + self.out_layer = OutLayer(model_dim, time_dim, out_visual_dim, patch_size) + for name, value in (attribute_overrides or {}).items(): + setattr(self, str(name), value) + + def forward( + self, + x: Tensor, + text_embed: Tensor, + pooled_text_embed: Tensor, + time: Tensor, + visual_cu_seqlens: Tensor, + text_cu_seqlens: Tensor, + visual_rope_pos: tuple[Tensor, Tensor, Tensor], + text_rope_pos: Tensor, + scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), + sparse_params: dict[str, Any] | None = None, + motion_score: Tensor | None = None, + lq_latent: Tensor | None = None, + lq_noise_level: Tensor | None = None, + ) -> Tensor: + """Run the full diffusion transformer forward pass. + + Args: + x: Noisy visual input tensor. + text_embed: Raw text encoder hidden states. + pooled_text_embed: Pooled text embeddings added to time conditioning. + time: Diffusion time steps. + visual_cu_seqlens: Cumulative sequence lengths for visual tokens. + text_cu_seqlens: Cumulative sequence lengths for text tokens. + visual_rope_pos: 3D positional indices (temporal, height, width) for visual RoPE. + text_rope_pos: 1D positional indices for text RoPE. + scale_factor: RoPE frequency scaling per axis. + sparse_params: Optional sparse/fractal attention parameters. + motion_score: Optional motion score for temporal conditioning. + lq_latent: Degraded video latent passed exclusively to the + adapter when ``use_adapter=True``. Not noised. ``None`` when adapter + is disabled. + lq_noise_level: Per-sequence LQ-endpoint noise fraction ``s`` in + ``[0, 1]`` (shape ``[N_seqs]`` or broadcastable ``[1]``). Only + used when ``use_lq_noise_cond=True``; embedded like the timestep + (scaled x1000) and summed into the adaLN conditioning bus. + + Returns: + Denoised visual output tensor. + """ + text_embed = self.text_embeddings(text_embed) if self.use_text else None + time_embed, time_embed_idx = self.time_embeddings(time) + if motion_score is not None and self.use_motion_score: + ms_embed, _ = self.motion_embeddings(motion_score) + time_embed = time_embed + ms_embed + if lq_noise_level is not None and self.use_lq_noise_cond: + # Same input scale as the timestep (t*1000 at every call site). + nl_embed, _ = self.lq_noise_embeddings(1000.0 * lq_noise_level) + time_embed = time_embed + nl_embed + + if self.use_text: + time_embed = time_embed + self.pooled_text_embeddings(pooled_text_embed) + else: + time_embed = time_embed + self.pooled_bias + + # Save cu_seqlens before patchification (adapter needs it for its own _patchify) + if self.use_adapter: + pre_patch_cu_seqlens = visual_cu_seqlens + + ve_result = self.visual_embeddings(x, visual_cu_seqlens) + if self.use_lq_modulation: + visual_embed, visual_cu_seqlens, lq_tokens = ve_result + else: + visual_embed, visual_cu_seqlens = ve_result + lq_tokens = None + + if self.use_text: + text_rope = self.text_rope_embeddings(text_rope_pos) + text_time_embed_idx = time_embed_idx.repeat_interleave(torch.diff(text_cu_seqlens), dim=0) + for text_transformer_block in self.text_transformer_blocks: + text_embed = text_transformer_block( + text_embed, time_embed, text_rope, text_cu_seqlens, text_time_embed_idx + ) + + visual_shape = visual_embed.shape[:-1] + visual_rope = self.visual_rope_embeddings(visual_shape, visual_rope_pos, scale_factor) + to_fractal = sparse_params["to_fractal"] if sparse_params is not None else False + # [T, 32, 32, 1792] -> [T*32*32, 1792] + visual_embed, visual_rope, visual_cu_seqlens = fractal_flatten( + visual_embed, visual_rope, visual_cu_seqlens, visual_shape, fractal=to_fractal + ) + if lq_tokens is not None: + # Flatten lq_tokens the same way as visual_embed (spatial dims only). + # fractal_flatten already modified visual_cu_seqlens, so we inline the + # flatten here rather than calling fractal_flatten again. + if to_fractal: + # [T, 32, 32, dim] -> [T * (32/8)*(32/8), 8*8, dim] -> [T*16, 64, dim] + lq_tokens = local_patching(lq_tokens, visual_shape, (1, 8, 8), dim=0) + lq_tokens = lq_tokens.flatten(0, 1) + else: + # [T, 32, 32, 512] -> [T*32*32, 512] + lq_tokens = lq_tokens.flatten(0, 2) + + visual_time_embed_idx = time_embed_idx.repeat_interleave(torch.diff(visual_cu_seqlens), dim=0) + + # Run adapter on LQ tokens and collect γ-scaled features for injection + adapter_features: list[Tensor] | None = None + if self.use_adapter: + adapter_features = self.adapter( + lq_latent, + pre_patch_cu_seqlens, + time_embed, + text_embed, + visual_rope, + visual_cu_seqlens, + text_cu_seqlens, + visual_time_embed_idx, + sparse_params, + to_fractal, + ) + + for i, visual_transformer_block in enumerate(self.visual_transformer_blocks): + # Inject adapter features before alternating backbone blocks (0, 2, 4, ...) + if adapter_features is not None and i % 2 == 0: + visual_embed = visual_embed + adapter_features[i // 2] + visual_embed = visual_transformer_block( + visual_embed, + text_embed, + time_embed, + visual_rope, + visual_cu_seqlens, + text_cu_seqlens, + visual_time_embed_idx, + sparse_params, + lq_tokens, + ) + visual_embed, visual_cu_seqlens = fractal_unflatten( + visual_embed, visual_cu_seqlens, visual_shape, fractal=to_fractal + ) + + visual_time_embed_idx = time_embed_idx.repeat_interleave(torch.diff(visual_cu_seqlens), dim=0) + return self.out_layer(visual_embed, text_embed, time_embed, visual_cu_seqlens, visual_time_embed_idx) + + def enable_explicit_forward_prefetch(self, factor: int = 1) -> None: + """Enable FSDP forward prefetching for visual transformer blocks. + + Args: + factor: Number of blocks to prefetch ahead. Skipped if <= 0 or already enabled. + + Raises: + RuntimeError: If the model is not wrapped with FSDP. + """ + if factor <= 0 or getattr(self, "forward_prefetch", False): + return + for i, layer in enumerate(self.visual_transformer_blocks): + if not hasattr(layer, "set_modules_to_forward_prefetch"): + msg = "wrap DiT with FSDP first to use forward_prefetch" + raise RuntimeError(msg) + + start, end = i, i + factor + prefetch_modules = self.visual_transformer_blocks[start:end] + layer.set_modules_to_forward_prefetch(prefetch_modules) + self.forward_prefetch = True + + def zero_init_lq_noise_out(self) -> None: + """Zero the s-embedder output layer so its contribution starts at exactly 0. + + Keeps warm-start from a checkpoint without ``lq_noise_embeddings`` an + exact no-op (adaLN-zero pattern). Must be re-applied after any + ``reset_parameters`` (FSDP meta-init materialisation re-randomises it). + """ + nn.init.zeros_(self.lq_noise_embeddings.out_layer.weight) + nn.init.zeros_(self.lq_noise_embeddings.out_layer.bias) + + def reset_parameters(self) -> None: + """Re-initialize all learnable parameters across every sub-module.""" + self.time_embeddings.reset_parameters() + if self.use_lq_noise_cond: + self.lq_noise_embeddings.reset_parameters() + self.zero_init_lq_noise_out() + if self.use_text: + self.text_embeddings.reset_parameters() + self.pooled_text_embeddings.reset_parameters() + else: + self.pooled_bias.data.zero_() + self.visual_embeddings.reset_parameters() + + if self.use_text: + self.text_rope_embeddings.reset_parameters() + for module in self.text_transformer_blocks: + module.reset_parameters() + + self.visual_rope_embeddings.reset_parameters() + for module in self.visual_transformer_blocks: + module.reset_parameters() + + self.out_layer.reset_parameters() + + if self.use_adapter: + self.adapter.reset_parameters() + + +__all__ = ["Kandinsky6SRTransformer3DModel"] diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index ef1814bbebcb..ff4162bdc7a8 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -448,6 +448,13 @@ "Kandinsky5T2IPipeline", "Kandinsky5I2IPipeline", ] + _import_structure["kandinsky6"] = [ + "Kandinsky6SRPipeline", + "Kandinsky6SRPipelineOutput", + "Kandinsky6TI2VAPipeline", + "Kandinsky6SRLatentUpscalerBank", + "Kandinsky6TI2VAPipelineOutput", + ] _import_structure["z_image"] = [ "ZImageControlNetInpaintPipeline", "ZImageControlNetPipeline", @@ -777,6 +784,13 @@ Kandinsky5T2IPipeline, Kandinsky5T2VPipeline, ) + from .kandinsky6 import ( + Kandinsky6SRPipeline, + Kandinsky6SRPipelineOutput, + Kandinsky6TI2VAPipeline, + Kandinsky6SRLatentUpscalerBank, + Kandinsky6TI2VAPipelineOutput, + ) from .krea2 import Krea2Pipeline from .latent_consistency_models import ( LatentConsistencyModelImg2ImgPipeline, diff --git a/src/diffusers/pipelines/kandinsky6/__init__.py b/src/diffusers/pipelines/kandinsky6/__init__.py new file mode 100644 index 000000000000..61d785b59ea8 --- /dev/null +++ b/src/diffusers/pipelines/kandinsky6/__init__.py @@ -0,0 +1,39 @@ +"""Kandinsky 6 pipelines.""" + +from typing import TYPE_CHECKING + +from ...utils import DIFFUSERS_SLOW_IMPORT, _LazyModule, is_torch_available + + +_import_structure = {} + +if is_torch_available(): + _import_structure["pipeline_kandinsky6_ti2va"] = ["Kandinsky6TI2VAPipeline"] + _import_structure["pipeline_kandinsky6_sr"] = ["Kandinsky6SRPipeline"] + _import_structure["latent_upscaler"] = ["Kandinsky6SRLatentUpscalerBank"] + _import_structure["pipeline_output"] = ["Kandinsky6SRPipelineOutput", "Kandinsky6TI2VAPipelineOutput"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + if is_torch_available(): + from .pipeline_kandinsky6_ti2va import Kandinsky6TI2VAPipeline + from .pipeline_kandinsky6_sr import Kandinsky6SRPipeline + from .latent_upscaler import Kandinsky6SRLatentUpscalerBank + from .pipeline_output import Kandinsky6SRPipelineOutput, Kandinsky6TI2VAPipelineOutput +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + +__all__ = [ + "Kandinsky6TI2VAPipeline", + "Kandinsky6TI2VAPipelineOutput", + "Kandinsky6SRPipeline", + "Kandinsky6SRLatentUpscalerBank", + "Kandinsky6SRPipelineOutput", +] diff --git a/src/diffusers/pipelines/kandinsky6/latent_upscaler.py b/src/diffusers/pipelines/kandinsky6/latent_upscaler.py new file mode 100644 index 000000000000..ec81765c4e0a --- /dev/null +++ b/src/diffusers/pipelines/kandinsky6/latent_upscaler.py @@ -0,0 +1,2723 @@ +"""Kandinsky 6 SR latent-upscaler bank Diffusers component.""" + +from __future__ import annotations + +import copy +import functools +import importlib +import math +from collections import defaultdict +from collections.abc import Mapping, Sequence +from functools import partial +from typing import TYPE_CHECKING, Annotated, Any, Literal, NamedTuple + +import pydantic +import torch +from ...configuration_utils import ConfigMixin, register_to_config +from ...models.modeling_utils import ModelMixin +from einops import rearrange +from pydantic import BaseModel, Field, TypeAdapter +from torch import Tensor, nn +from torch.autograd import Function +from torch.nn import functional +from torchvision.ops import StochasticDepth + + +@torch.autocast(device_type="cuda", enabled=False) +def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: + return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) + + +"""Convolution primitives shared by the latent upsampler architectures.""" + + +DIMS_2 = 2 +DIMS_3 = 3 + +TemporalPadding = Literal["zeros", "replicate", "causal"] +UpsamplePaddingMode = Literal["reflect", "zeros"] + + +def as_triple(value: int | tuple[int, int, int]) -> tuple[int, int, int]: + """Expand a scalar kernel/padding spec into an explicit ``(t, h, w)`` triple.""" + return (value, value, value) if isinstance(value, int) else value + + +class TemporalReplicateConv3d(nn.Conv3d): + """``Conv3d`` that repeats the edge frame along T and zero-pads H/W. + + K-VAE extends the temporal axis by repeating a boundary frame — the encoder + and decoder both seed their causal padding with a replica of frame 0 — while + padding H/W with zeros (``padding_mode: zeros`` in the shipped sidecar). + ``nn.Conv3d`` cannot express that split, because ``padding_mode`` applies to + every padded dim at once; here T is padded explicitly and H/W is left to the + convolution. + + Without this, a zero-padded temporal axis makes the first and last latent + frames see a hole where a neighbour should be — on K-VAE latents that hole + lands on frame 0, which is already the odd one out (it encodes a single + pixel frame, while every later latent aggregates four). + + Parameter names and shapes are those of ``nn.Conv3d``, so checkpoints + trained before this padding change still load. + + Args: + in_channels: Number of input channels. + out_channels: Number of output channels. + kernel_size: Kernel as ``(kt, kh, kw)`` or a single int applied to all dims. + padding: The "same" padding the caller would have passed to ``nn.Conv3d``. + Its H/W entries go to the convolution; its T entry becomes the width + of the replicate pad. + groups: Convolution groups (``in_channels`` for a depthwise conv). + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int, int], + padding: int | tuple[int, int, int] = 0, + *, + groups: int = 1, + ) -> None: + """Initialize the convolution with H/W padding only, keeping T for forward.""" + pad_t, pad_h, pad_w = as_triple(padding) + super().__init__( + in_channels, + out_channels, + kernel_size, + padding=(0, pad_h, pad_w), + groups=groups, + ) + self.temporal_pad = pad_t + + def temporal_pad_lr(self) -> tuple[int, int]: + """Split the temporal pad budget into ``(before, after)`` frame counts.""" + return (self.temporal_pad, self.temporal_pad) + + def forward(self, x: Tensor) -> Tensor: + """Replicate-pad T, then convolve with zero padding on H and W.""" + before, after = self.temporal_pad_lr() + if before or after: + x = functional.pad(x, (0, 0, 0, 0, before, after), mode="replicate") + return super().forward(x) + + +class TemporalCausalConv3d(TemporalReplicateConv3d): + """``TemporalReplicateConv3d`` with K-VAE's causal temporal window. + + K-VAE's ``CausalConv3d`` spends the whole temporal pad budget in front: + ``kT - 1`` replicas of frame 0 precede the clip and nothing follows it, so + each 3-tap kernel reads ``(t-2, t-1, t)``. Moving the symmetric budget + (``pad_t`` per side) to the front reproduces that window exactly, which is + what lets decoder kernels load without recentring — and keeps every frame's + output independent of its future, the property a chunked streaming + inference would rely on. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int, int], + padding: int | tuple[int, int, int] = 0, + *, + groups: int = 1, + ) -> None: + """Validate that the pad budget covers exactly the causal window.""" + super().__init__(in_channels, out_channels, kernel_size, padding, groups=groups) + kernel_t = self.kernel_size[0] + if 2 * self.temporal_pad != kernel_t - 1: + msg = ( + f"a causal window needs the full kT-1 = {kernel_t - 1} pad budget in front, " + f"but padding supplies 2 * {self.temporal_pad}" + ) + raise ValueError(msg) + + def temporal_pad_lr(self) -> tuple[int, int]: + """Put the whole budget before the clip: taps end on the current frame.""" + return (2 * self.temporal_pad, 0) + + +def make_conv(dims: int, temporal_padding: TemporalPadding) -> type[nn.Module]: + """Pick the convolution class for a stack of ``dims``-dimensional convs. + + Args: + dims: 2 for ``Conv2d``, 3 for ``Conv3d``. + temporal_padding: How ``dims == 3`` convolutions extend T. ``"zeros"`` + keeps the stock convolution; ``"replicate"`` repeats the edge frame; + ``"causal"`` pads the past only (K-VAE semantics). + Ignored for ``dims == 2``, which has no temporal axis. + + Returns: + The convolution class to instantiate. + """ + if dims == DIMS_2: + return nn.Conv2d + if temporal_padding == "replicate": + return TemporalReplicateConv3d + if temporal_padding == "causal": + return TemporalCausalConv3d + return nn.Conv3d + + +"""Reusable building blocks for latent upsampler architectures.""" + + +class RMSNorm(nn.Module): + """Channel-first Root Mean Square normalization with learnable gamma. + + Args: + dim: Number of channels. + dims: Convolution dimensionality — 2 for ``(C,1,1)``, 3 for ``(C,1,1,1)``. + """ + + def __init__(self, dim: int, dims: int = 2) -> None: + """Initialize RMSNorm. + + Args: + dim: Number of channels. + dims: Convolution dimensionality — 2 for ``(C,1,1)``, 3 for ``(C,1,1,1)``. + """ + super().__init__() + broadcastable_dims = (1, 1) if dims == 2 else (1, 1, 1) + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim, *broadcastable_dims)) + + def forward(self, x: Tensor) -> Tensor: + """Apply RMS normalization along the channel dimension.""" + return functional.normalize(x, dim=1) * self.scale * self.gamma + + +class ModulatedRMSNorm(nn.Module): + """RMSNorm followed by spatial FiLM modulation conditioned on a side tensor. + + Computes ``RMSNorm(x) * conv_y(zq) + conv_b(zq)`` where ``conv_y`` and + ``conv_b`` are 1x1 convolutions of ``zq`` (the conditioning tensor, e.g. + the LQ latent). ``zq`` is aligned to ``x`` by nearest-neighbor + interpolation along the spatial dims, so the same ``zq`` can be reused + across feature pyramids of different resolutions. + + Both convs keep the stock PyTorch initialization, so modulation is active + from the first step — the same regime as K-VAE's ``CachedSpatialNorm3D``, + whose ``conv_y`` / ``conv_b`` are plain 1x1 convs with no custom init. + + Args: + dim: Number of feature channels. + zq_dim: Number of channels in the conditioning tensor. + dims: Convolution dimensionality — 2 for Conv2d, 3 for Conv3d. + """ + + def __init__(self, dim: int, zq_dim: int, dims: int = 2) -> None: + """Initialize ModulatedRMSNorm. + + Args: + dim: Number of feature channels. + zq_dim: Number of channels in the conditioning tensor. + dims: Convolution dimensionality — 2 for Conv2d, 3 for Conv3d. + """ + super().__init__() + conv = nn.Conv2d if dims == 2 else nn.Conv3d + self.norm = RMSNorm(dim, dims=dims) + self.conv_y = conv(zq_dim, dim, kernel_size=1) + self.conv_b = conv(zq_dim, dim, kernel_size=1) + + def forward(self, x: Tensor, zq: Tensor) -> Tensor: + """Apply RMS norm + FiLM modulation by spatially-aligned ``zq``.""" + if zq.shape[2:] != x.shape[2:]: + zq = functional.interpolate(zq, size=x.shape[2:], mode="nearest") + return self.norm(x) * self.conv_y(zq) + self.conv_b(zq) + + +class LayerScale(nn.Module): + """Learnable per-channel scaling applied before residual addition. + + Args: + channels: Number of channels. + init_value: Initial scale value (small for training stability). + dims: Convolution dimensionality — 2 for ``(C,1,1)``, 3 for ``(C,1,1,1)``. + """ + + def __init__(self, channels: int, init_value: float = 1e-6, dims: int = 2) -> None: + """Initialize LayerScale. + + Args: + channels: Number of channels. + init_value: Initial scale value. + dims: Convolution dimensionality — 2 for ``(C,1,1)``, 3 for ``(C,1,1,1)``. + """ + super().__init__() + broadcastable_dims = (1, 1) if dims == 2 else (1, 1, 1) + self.gamma = nn.Parameter(init_value * torch.ones(channels, *broadcastable_dims)) + + def forward(self, x: Tensor) -> Tensor: + """Scale input by learnable per-channel gamma.""" + return x * self.gamma + + +class GRN(nn.Module): + """Global Response Normalization (ConvNeXtV2). + + Aggregates global spatial info per channel and normalizes across channels + to encourage feature diversity and prevent feature collapse. + + Args: + channels: Number of channels. + dims: Convolution dimensionality — 2 for spatial dims ``(2,3)``, 3 for ``(2,3,4)``. + """ + + def __init__(self, channels: int, dims: int = 2) -> None: + """Initialize GRN. + + Args: + channels: Number of channels. + dims: Convolution dimensionality — 2 for spatial dims ``(2,3)``, 3 for ``(2,3,4)``. + """ + super().__init__() + self.spatial_dims: tuple[int, ...] = (2, 3) if dims == 2 else (2, 3, 4) + broadcastable = (1, 1) if dims == 2 else (1, 1, 1) + self.gamma = nn.Parameter(torch.zeros(1, channels, *broadcastable)) + self.beta = nn.Parameter(torch.zeros(1, channels, *broadcastable)) + + def forward(self, x: Tensor) -> Tensor: + """Apply global response normalization.""" + gx = torch.norm(x, p=2, dim=self.spatial_dims, keepdim=True) + nx = gx / (gx.mean(dim=1, keepdim=True) + 1e-6) + return self.gamma * (x * nx) + self.beta + x + + +class ResidualBlock(nn.Module): + """Pre-activation residual block with optional LayerScale, GRN, and StochasticDepth. + + Standard path: ``RMSNorm → SiLU → Conv3x3 → RMSNorm → SiLU → [GRN →] Conv3x3``. + Depthwise path: ``DWConv_KxK → RMSNorm → PWConv1x1 → SiLU → [GRN →] PWConv1x1``. + Skip: identity when ``in_channels == out_channels``, else ``Conv 1x1``. + + When ``zq_dim`` is set, every ``RMSNorm`` is replaced by a + ``ModulatedRMSNorm`` conditioned on a side tensor ``zq``, and the inner + layers are stored as named submodules so ``zq`` can be threaded through + ``forward(x, zq)``. When ``zq_dim is None``, the legacy ``nn.Sequential`` + layout is preserved bit-identically — the parameter names match older + checkpoints, so loading them remains backward compatible. + """ + + def __init__( + self, + in_channels: int, + out_channels: int | None = None, + mid_channels: int | None = None, + dims: int = 2, + layer_scale_init: float | None = None, + *, + grn: bool = False, + stochastic_depth_prob: float = 0.0, + depthwise: bool = False, + kernel_size: int = 3, + zq_dim: int | None = None, + temporal_padding: TemporalPadding = "zeros", + ) -> None: + """Initialize residual block. + + Args: + in_channels: Number of input channels. + out_channels: Number of output channels (defaults to ``in_channels``). + mid_channels: Internal channel width (defaults to ``out_channels``). + dims: Convolution dimensionality — 2 for Conv2d, 3 for Conv3d. + layer_scale_init: Initial LayerScale gamma value. ``None`` disables LayerScale. + grn: Whether to insert Global Response Normalization on mid_channels. + stochastic_depth_prob: Drop probability for stochastic depth (0 = off). + depthwise: Use depthwise separable convolutions instead of standard convs. + kernel_size: Kernel size for the depthwise convolution (only used when ``depthwise=True``). + zq_dim: When set, replace every ``RMSNorm`` with ``ModulatedRMSNorm(zq_dim)``; + the block then accepts a conditioning tensor in ``forward(x, zq)``. + temporal_padding: How ``dims == 3`` convolutions extend T — ``"zeros"`` + (default) or ``"replicate"``, which repeats the edge frame as K-VAE does. + """ + super().__init__() + out_channels = out_channels or in_channels + mid_channels = mid_channels or out_channels + self.depthwise = depthwise + self.has_grn = grn + self.zq_dim = zq_dim + conv = make_conv(dims, temporal_padding) + + if zq_dim is None: + self._build_legacy(in_channels, out_channels, mid_channels, dims, conv, grn=grn, kernel_size=kernel_size) + else: + self._build_modulated( + in_channels, + out_channels, + mid_channels, + dims, + conv, + grn=grn, + kernel_size=kernel_size, + zq_dim=zq_dim, + ) + self.shortcut = nn.Identity() if in_channels == out_channels else conv(in_channels, out_channels, kernel_size=1) + self.layer_scale: nn.Module = ( + LayerScale(out_channels, init_value=layer_scale_init, dims=dims) + if layer_scale_init is not None + else nn.Identity() + ) + self.stochastic_depth: nn.Module = ( + StochasticDepth(stochastic_depth_prob, mode="row") if stochastic_depth_prob > 0 else nn.Identity() + ) + + def _build_legacy( + self, + in_channels: int, + out_channels: int, + mid_channels: int, + dims: int, + conv: type[nn.Module], + *, + grn: bool, + kernel_size: int, + ) -> None: + """Build the original ``nn.Sequential`` layout (parameter names unchanged).""" + if self.depthwise: + ks = kernel_size if dims == 2 else (min(kernel_size, 5), kernel_size, kernel_size) + pad = kernel_size // 2 if dims == 2 else (min(kernel_size, 5) // 2, kernel_size // 2, kernel_size // 2) + layers: list[nn.Module] = [ + conv(in_channels, in_channels, kernel_size=ks, padding=pad, groups=in_channels), + RMSNorm(in_channels, dims=dims), + conv(in_channels, mid_channels, kernel_size=1), + nn.SiLU(), + ] + if grn: + layers.append(GRN(mid_channels, dims=dims)) + layers.append(conv(mid_channels, out_channels, kernel_size=1)) + else: + layers = [ + RMSNorm(in_channels, dims=dims), + nn.SiLU(), + conv(in_channels, mid_channels, kernel_size=3, padding=1), + RMSNorm(mid_channels, dims=dims), + nn.SiLU(), + ] + if grn: + layers.append(GRN(mid_channels, dims=dims)) + layers.append(conv(mid_channels, out_channels, kernel_size=3, padding=1)) + self.block = nn.Sequential(*layers) + + def _build_modulated( + self, + in_channels: int, + out_channels: int, + mid_channels: int, + dims: int, + conv: type[nn.Module], + *, + grn: bool, + kernel_size: int, + zq_dim: int, + ) -> None: + """Build a path with FiLM-modulated RMSNorms accepting ``zq`` in forward.""" + if self.depthwise: + ks = kernel_size if dims == 2 else (min(kernel_size, 5), kernel_size, kernel_size) + pad = kernel_size // 2 if dims == 2 else (min(kernel_size, 5) // 2, kernel_size // 2, kernel_size // 2) + self.dwconv = conv(in_channels, in_channels, kernel_size=ks, padding=pad, groups=in_channels) + self.norm1 = ModulatedRMSNorm(in_channels, zq_dim, dims=dims) + self.pwconv1 = conv(in_channels, mid_channels, kernel_size=1) + self.grn: nn.Module | None = GRN(mid_channels, dims=dims) if grn else None + self.pwconv2 = conv(mid_channels, out_channels, kernel_size=1) + else: + self.norm1 = ModulatedRMSNorm(in_channels, zq_dim, dims=dims) + self.conv1 = conv(in_channels, mid_channels, kernel_size=3, padding=1) + self.norm2 = ModulatedRMSNorm(mid_channels, zq_dim, dims=dims) + self.grn = GRN(mid_channels, dims=dims) if grn else None + self.conv2 = conv(mid_channels, out_channels, kernel_size=3, padding=1) + + def forward(self, x: Tensor, zq: Tensor | None = None) -> Tensor: + """Apply residual block with skip connection. + + Args: + x: Feature tensor. + zq: Optional conditioning tensor for FiLM-modulated norms. Required + when the block was constructed with ``zq_dim`` set; ignored otherwise. + """ + out = self.block(x) if self.zq_dim is None else self.forward_modulated_path(x, zq) # type: ignore + out = self.layer_scale(out) + return self.shortcut(x) + self.stochastic_depth(out) + + def forward_modulated_path( + self, + x: Tensor, + zq: Tensor, + norm1: ModulatedRMSNorm | None = None, + norm2: ModulatedRMSNorm | None = None, + ) -> Tensor: + """Evaluate the residual path with optional scale-specific norms.""" + if self.zq_dim is None: + msg = "forward_modulated_path requires a block with zq_dim" + raise ValueError(msg) + active_norm1 = norm1 if norm1 is not None else self.norm1 + if self.depthwise: + h = self.dwconv(x) + h = active_norm1(h, zq) + h = self.pwconv1(h) + h = functional.silu(h) + if self.grn is not None: + h = self.grn(h) + return self.pwconv2(h) + active_norm2 = norm2 if norm2 is not None else self.norm2 + h = active_norm1(x, zq) + h = functional.silu(h) + h = self.conv1(h) + h = active_norm2(h, zq) + h = functional.silu(h) + if self.grn is not None: + h = self.grn(h) + return self.conv2(h) + + def forward_with_norms( + self, + x: Tensor, + zq: Tensor, + norm1: ModulatedRMSNorm, + norm2: ModulatedRMSNorm, + ) -> Tensor: + """Apply the block with x2-specific modulation and shared transforms.""" + if self.depthwise: + msg = "scale-specific normalization does not support depthwise residual blocks" + raise ValueError(msg) + out = self.forward_modulated_path(x, zq, norm1=norm1, norm2=norm2) + out = self.layer_scale(out) + return self.shortcut(x) + self.stochastic_depth(out) + + def zero_init_residual(self) -> None: + """Initialize the residual path to zero so the block starts as identity.""" + if self.zq_dim is None: + final_conv = self.block[-1] + elif self.depthwise: + final_conv = self.pwconv2 + else: + final_conv = self.conv2 + nn.init.zeros_(final_conv.weight) # type: ignore[arg-type] + if final_conv.bias is not None: + nn.init.zeros_(final_conv.bias) # type: ignore[arg-type] + + +"""Exact merging of flash-style attention branches with a closed-form null branch.""" + + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + +def merge_math( + outputs: Sequence[Tensor], + lses: Sequence[Tensor], + null_value: Tensor, + null_logits: Tensor, +) -> tuple[Tensor, Tensor]: + """Combine attention branches and the null token into one softmax, in fp32. + + Args: + outputs: Branch outputs shaped ``(batch, seq, heads, dim)``. + lses: Branch log-sum-exps shaped ``(batch, seq, heads)``. + null_value: Null-token value shaped ``(heads, dim)``. + null_logits: Null-token logits shaped ``(batch, seq, heads)``. + + Returns: + Merged output ``(batch, seq, heads, dim)`` and total lse ``(batch, seq, heads)``. + """ + lse_stack = torch.stack([*[lse.float() for lse in lses], null_logits.float()], dim=0) + total_lse = torch.logsumexp(lse_stack, dim=0) + weights = torch.exp(lse_stack - total_lse) + merged = weights[-1].unsqueeze(-1) * null_value.float() + for weight, output in zip(weights[:-1], outputs, strict=True): + merged = merged + weight.unsqueeze(-1) * output.float() + return merged, total_lse + + +@functools.cache +def compiled_merge_math() -> Callable[..., tuple[Tensor, Tensor]]: + """Compile the merge math once and reuse it across calls.""" + return torch.compile(merge_math, fullgraph=True) + + +class MergeWithNullBranch(Function): + """Merge kernel attention branches with the null branch, with exact gradients. + + Flash-style kernels (NATTEN ``na2d``) return ``(output, lse)`` but their + backward ignores incoming lse gradients, so plain autograd through a softmax + merge silently drops part of their q/k/v gradients. This op follows the + ring-attention convention instead: each kernel branch receives the full + upstream gradient while its saved output/lse tensors are overwritten in + place with the merged values, letting the kernel backward reconstruct exact + global gradients. The null branch is single-key attention, so its gradients + have a closed form and are computed here directly. + + The in-place overwrite means branch tensors must not be read after backward, + and the op is incompatible with gradient recomputation (activation + checkpointing) and double backward. + """ + + @staticmethod + def forward(ctx: Any, *args: Any) -> Tensor: + """Merge branches; see :func:`merge_attention_branches` for the layout of ``args``.""" + null_value, null_logits, num_branches, use_compile, *branch_tensors = args + outputs = branch_tensors[:num_branches] + lses = branch_tensors[num_branches:] + math_fn = compiled_merge_math() if use_compile else merge_math + merged, total_lse = math_fn(outputs, lses, null_value, null_logits) + ctx.num_branches = num_branches + ctx.save_for_backward(null_value, null_logits, merged, total_lse, *branch_tensors) + target = outputs[0] if outputs else null_value + return merged.to(target.dtype) + + @staticmethod + def backward(ctx: Any, grad_out: Tensor) -> tuple[Tensor | None, ...]: + """Route full gradients to kernel branches and closed-form ones to the null pair.""" + num_branches = ctx.num_branches + null_value, null_logits, merged, total_lse, *branch_tensors = ctx.saved_tensors + outputs = branch_tensors[:num_branches] + lses = branch_tensors[num_branches:] + grad = grad_out.float() + null_weight = torch.exp(null_logits.float() - total_lse) + d_null_value = torch.einsum("bsh,bshd->hd", null_weight, grad) + d_null_logits = null_weight * ( + torch.einsum("bshd,hd->bsh", grad, null_value.float()) - (grad * merged).sum(dim=-1) + ) + for output, lse in zip(outputs, lses, strict=True): + output.data.copy_(merged.to(output.dtype)) + lse.data.copy_(total_lse.to(lse.dtype)) + return ( + d_null_value.to(null_value.dtype), + d_null_logits.to(null_logits.dtype), + None, + None, + *(grad_out for _ in outputs), + *(None for _ in lses), + ) + + +def merge_attention_branches( + outputs: Sequence[Tensor], + lses: Sequence[Tensor], + null_value: Tensor, + null_logits: Tensor, + *, + torch_compile: bool = False, +) -> Tensor: + """Merge kernel attention branches with the closed-form null branch. + + Args: + outputs: Kernel branch outputs shaped ``(batch, seq, heads, dim)``; may be empty. + lses: Matching branch log-sum-exps shaped ``(batch, seq, heads)``. + null_value: Null-token value shaped ``(heads, dim)``. + null_logits: Null-token logits shaped ``(batch, seq, heads)``. + torch_compile: Compile the merge math with ``torch.compile``. + + Returns: + Merged attention output shaped ``(batch, seq, heads, dim)`` in the branch dtype + (or the null-value dtype when no branches are given). + """ + return MergeWithNullBranch.apply(null_value, null_logits, len(outputs), torch_compile, *outputs, *lses) + + +"""Runtime helpers shared by latent-upscaler model components.""" + + +if TYPE_CHECKING: + from collections.abc import Callable + + +def forward_with_checkpointing( + module: Callable[..., Tensor], + *inputs: Tensor, + use_checkpointing: bool = False, +) -> Tensor: + """Run a callable directly or through non-reentrant activation checkpointing.""" + if use_checkpointing: + return torch.utils.checkpoint.checkpoint(module, *inputs, use_reentrant=False) + return module(*inputs) + + +"""Spatial upsampling operations: pixel-shuffle, bilinear, and K-VAE PXS v2.""" + + +if TYPE_CHECKING: + from collections.abc import Callable + +DIMS_2 = 2 +DIMS_3 = 3 +VALID_DIMS = {DIMS_2, DIMS_3} + + +def icnr_( + weight: torch.Tensor, + factor: int, + init_fn: Callable[[torch.Tensor], torch.Tensor] = nn.init.kaiming_normal_, +) -> None: + """Initialize sub-pixel-conv weight so initial output equals nearest 2x. + + Implements the ICNR scheme of Aitken et al. (2017): the ``factor**2`` + sub-channel groups consumed by ``PixelShuffle(factor)`` are initialized + with identical kernels, so on the first forward pass the output is a + pure nearest-neighbor upsample. This eliminates checkerboard artifacts + at the start of training. + + Args: + weight: Conv weight tensor of shape ``(C_out * factor**2, C_in, *kernel)``. + Modified in-place. + factor: Upscale factor used by the downstream ``PixelShuffle``. + init_fn: Base initializer applied to the underlying ``C_out`` kernel + before replication. Defaults to Kaiming normal. + + Raises: + ValueError: If ``weight.shape[0]`` is not divisible by ``factor**2``. + """ + out_c = weight.shape[0] // factor**2 + if out_c * factor**2 != weight.shape[0]: + msg = f"weight.shape[0] ({weight.shape[0]}) must be divisible by factor**2 ({factor**2})" + raise ValueError(msg) + base = torch.empty(out_c, *weight.shape[1:], device=weight.device, dtype=weight.dtype) + init_fn(base) + weight.data.copy_(base.repeat_interleave(factor**2, dim=0)) + + +def spatial_nearest_2x(x: torch.Tensor, dims: int, factor: int) -> torch.Tensor: + """Apply nearest-neighbor upsampling to H, W only, preserving B (and T). + + For 5D ``(B, C, T, H, W)`` tensors, T is merged into batch, interpolated, + and unmerged. For 4D tensors, interpolation runs directly. + + Args: + x: Input tensor of shape ``(B, C, H, W)`` or ``(B, C, T, H, W)``. + dims: Tensor dimensionality — 2 for 4D, 3 for 5D. + factor: Spatial upscale factor. + + Returns: + Upsampled tensor with H, W scaled by ``factor``. + """ + if dims == DIMS_3: + b, _c, t, _h, _w = x.shape + x = rearrange(x, "b c t h w -> (b t) c h w") + x = functional.interpolate(x, scale_factor=factor, mode="nearest") + return rearrange(x, "(b t) c h w -> b c t h w", b=b, t=t) + return functional.interpolate(x, scale_factor=factor, mode="nearest") + + +class PixelShuffleND(nn.Module): + """Conv + pixel-shuffle that upscales only H and W. + + Internally applies a convolution to expand channels by ``upscale_factor ** 2``, + then rearranges channels into spatial dimensions. + + Supports both 4D ``(B, C, H, W)`` and 5D ``(B, C, T, H, W)`` tensors. + The temporal dimension is always preserved. + + Args: + channels: Number of input (and output) channels. + dims: Tensor dimensionality — 2 for 4D input, 3 for 5D (video) input. + factor: Spatial upscale factor applied to both H and W. + temporal_mix: Whether the 3D convolution mixes across the temporal dimension. + When ``False`` and ``dims == 3``, uses ``kernel_size=(1, 3, 3)`` so that + each frame is convolved independently in H/W. Ignored when ``dims == 2``. + Defaults to ``True`` (full 3x3x3 kernel). + icnr: Apply ICNR initialization (Aitken et al., 2017) to the conv weight + and zero its bias, so the module starts as a nearest-neighbor upsample. + Removes checkerboard at init. Defaults to ``False``. + temporal_padding: How the ``temporal_mix=True`` kernel extends T — + ``"zeros"`` (default) or ``"replicate"``, matching K-VAE's edge repeat. + """ + + def __init__( + self, + channels: int, + dims: int, + factor: int = 2, + *, + temporal_mix: bool = True, + icnr: bool = False, + temporal_padding: TemporalPadding = "zeros", + ) -> None: + """Initialize PixelShuffleND.""" + super().__init__() + if dims not in VALID_DIMS: + msg = f"dims must be one of {VALID_DIMS}, got {dims}" + raise ValueError(msg) + self.dims = dims + self.factor = factor + + if dims == DIMS_2: + self.conv = nn.Conv2d(channels, channels * factor**2, kernel_size=3, padding=1) + elif temporal_mix: + conv3d = make_conv(DIMS_3, temporal_padding) + self.conv = conv3d(channels, channels * factor**2, kernel_size=3, padding=1) + else: + self.conv = nn.Conv3d(channels, channels * factor**2, kernel_size=(1, 3, 3), padding=(0, 1, 1)) + + if icnr: + icnr_(self.conv.weight, factor) + if self.conv.bias is not None: + nn.init.zeros_(self.conv.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Expand channels via convolution, then rearrange into spatial dimensions.""" + x = self.conv(x) + if self.dims == DIMS_3: + return rearrange(x, "b (c p1 p2) t h w -> b c t (h p1) (w p2)", p1=self.factor, p2=self.factor) + return rearrange(x, "b (c p1 p2) h w -> b c (h p1) (w p2)", p1=self.factor, p2=self.factor) + + +class BilinearUpsampleND(nn.Module): + """Parameter-free bilinear spatial upsampling for 4D and 5D tensors. + + Upscales only H and W dimensions. For 5D ``(B, C, T, H, W)`` tensors, + the temporal dimension is merged into the batch, interpolated in 2D, + and unmerged. + + Args: + dims: Tensor dimensionality — 2 for 4D input, 3 for 5D (video) input. + factor: Spatial upscale factor applied to both H and W. + """ + + def __init__(self, dims: int, factor: int = 2) -> None: + """Initialize BilinearUpsampleND.""" + super().__init__() + if dims not in VALID_DIMS: + msg = f"dims must be one of {VALID_DIMS}, got {dims}" + raise ValueError(msg) + self.dims = dims + self.factor = factor + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply bilinear interpolation to spatial dimensions.""" + if self.dims == DIMS_3: + b, _c, t, _h, _w = x.shape + x = rearrange(x, "b c t h w -> (b t) c h w") + x = functional.interpolate(x, scale_factor=self.factor, mode="bilinear", align_corners=False) + return rearrange(x, "(b t) c h w -> b c t h w", b=b, t=t) + return functional.interpolate(x, scale_factor=self.factor, mode="bilinear", align_corners=False) + + +class PXSv2UpsampleND(nn.Module): + """K-VAE PXS v2 spatial upsample: ``nearest 2x + Conv_(1,3,3) residual + Conv_1x1``. + + Replicates ``CachedPXSUpsample.spatial_upsample_NEW`` followed by the + post-``linear`` 1x1x1 conv from K-VAE-3D-2.0. The base path is a + parameter-free nearest upsample; the residual conv adds learnable detail + on top, and a final pointwise conv mixes channels. + + Spatial conv is hard-coded to ``kernel=(1, 3, 3)`` for ``dims == 3``, + matching K-VAE: per-frame, no temporal mixing through the upsample itself. + Temporal correlation is restored by surrounding ResBlocks. + + Args: + channels: Number of input (and output) channels. + dims: Tensor dimensionality — 2 for 4D, 3 for 5D. + factor: Spatial upscale factor for H and W. + with_linear: Append a final 1x1 (or 1x1x1) pointwise conv. Disable + when this module is composed inside a parent that owns its own + post-mixer (e.g. ``PXSv2HybridUpsampleND``). Defaults to ``True``. + padding_mode: Edge handling for the residual conv. ``"reflect"`` is the + historical default; ``"zeros"`` matches the production K-VAE, whose + ``CachedPXSUpsample`` resolves ``padding_mode or 'reflect'`` against a + sidecar that passes ``'zeros'``. + """ + + def __init__( + self, + channels: int, + dims: int, + factor: int = 2, + *, + with_linear: bool = True, + padding_mode: UpsamplePaddingMode = "reflect", + ) -> None: + """Initialize PXSv2UpsampleND.""" + super().__init__() + if dims not in VALID_DIMS: + msg = f"dims must be one of {VALID_DIMS}, got {dims}" + raise ValueError(msg) + self.dims = dims + self.factor = factor + + if dims == DIMS_2: + self.spatial_conv: nn.Module = nn.Conv2d( + channels, + channels, + kernel_size=3, + padding=1, + padding_mode=padding_mode, + ) + self.linear: nn.Module = nn.Conv2d(channels, channels, kernel_size=1) if with_linear else nn.Identity() + else: + self.spatial_conv = nn.Conv3d( + channels, + channels, + kernel_size=(1, 3, 3), + padding=(0, 1, 1), + padding_mode=padding_mode, + ) + self.linear = nn.Conv3d(channels, channels, kernel_size=1) if with_linear else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Compute ``Linear( NN_2x(x) + Conv·NN_2x(x) )``.""" + up = spatial_nearest_2x(x, self.dims, self.factor) + return self.linear(up + self.spatial_conv(up)) + + +class PXSv2HybridUpsampleND(nn.Module): + """Hybrid upsample: ICNR-PXS branch parallel to PXS v2 branch, mixed by 1x1. + + Combines the learnable sub-pixel-conv branch (``PixelShuffleND``, optionally + ICNR-initialized so it starts as nearest) with the K-VAE PXS v2 branch + (``nearest + Conv_(1,3,3) residual``), then mixes both via a final 1x1 + pointwise conv. + + Args: + channels: Number of input (and output) channels. + dims: Tensor dimensionality — 2 for 4D, 3 for 5D. + factor: Spatial upscale factor for H and W. + temporal_mix: Forwarded to the inner ``PixelShuffleND`` (controls 3x3x3 + vs (1,3,3) kernel for ``dims == 3``). The PXS v2 branch always uses + (1,3,3) regardless. Defaults to ``True``. + icnr: Apply ICNR initialization to the inner ``PixelShuffleND`` so it + starts as nearest, matching the PXS v2 branch and removing + checkerboard at init. Defaults to ``True``. + temporal_padding: Forwarded to the inner ``PixelShuffleND``. + padding_mode: Forwarded to the inner ``PXSv2UpsampleND``. + """ + + def __init__( + self, + channels: int, + dims: int, + factor: int = 2, + *, + temporal_mix: bool = True, + icnr: bool = True, + temporal_padding: TemporalPadding = "zeros", + padding_mode: UpsamplePaddingMode = "reflect", + ) -> None: + """Initialize PXSv2HybridUpsampleND.""" + super().__init__() + if dims not in VALID_DIMS: + msg = f"dims must be one of {VALID_DIMS}, got {dims}" + raise ValueError(msg) + self.pxs = PixelShuffleND( + channels, + dims=dims, + factor=factor, + temporal_mix=temporal_mix, + icnr=icnr, + temporal_padding=temporal_padding, + ) + self.v3 = PXSv2UpsampleND( + channels, + dims=dims, + factor=factor, + with_linear=False, + padding_mode=padding_mode, + ) + if dims == DIMS_2: + self.linear: nn.Module = nn.Conv2d(channels, channels, kernel_size=1) + else: + self.linear = nn.Conv3d(channels, channels, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Compute ``Linear( PXS(x) + (NN_2x + Conv·NN_2x)(x) )``.""" + return self.linear(self.pxs(x) + self.v3(x)) + + +"""Motion-correspondence attention for latent video features.""" + + +if TYPE_CHECKING: + from types import ModuleType + +MotionAttentionBackend = Literal["sdpa", "natten"] +NattenBackend = Literal["cutlass-fna", "hopper-fna", "flex-fna"] + + +class MotionCorrespondenceSpec(BaseModel): + """Construction parameters for one motion-correspondence block.""" + + channels: int + num_heads: int + head_dim: int + spatial_kernel_size: int + temporal_offsets: tuple[int, ...] + backend: MotionAttentionBackend + natten_backend: NattenBackend = "hopper-fna" + merge_compile: bool = False + + +def temporal_neighbor_indices( + query_index: int, + num_frames: int, + temporal_offsets: tuple[int, ...], +) -> tuple[int, ...]: + """Return valid neighbor-frame indices without shifting or duplication. + + Args: + query_index: Query frame index. + num_frames: Number of frames in the clip. + temporal_offsets: Allowed non-zero offsets from the query frame. + + Returns: + Valid neighbor indices in the configured offset order. + """ + return tuple(query_index + offset for offset in temporal_offsets if 0 <= query_index + offset < num_frames) + + +def spatial_neighborhood_mask( + height: int, + width: int, + kernel_size: int, + device: torch.device, +) -> Tensor: + """Build the shifted sliding-window mask used by neighborhood attention. + + Args: + height: Feature-map height. + width: Feature-map width. + kernel_size: Odd spatial neighborhood size. + device: Device for the returned mask. + + Returns: + Boolean mask shaped ``(height * width, height * width)``. + """ + if kernel_size > min(height, width): + msg = f"spatial kernel {kernel_size} exceeds feature map {height}x{width}" + raise ValueError(msg) + + query_h, query_w = torch.meshgrid( + torch.arange(height, device=device), + torch.arange(width, device=device), + indexing="ij", + ) + key_h, key_w = query_h.flatten(), query_w.flatten() + radius = kernel_size // 2 + start_h = (query_h.flatten() - radius).clamp(min=0, max=height - kernel_size) + start_w = (query_w.flatten() - radius).clamp(min=0, max=width - kernel_size) + inside_h = (key_h[None, :] >= start_h[:, None]) & (key_h[None, :] < start_h[:, None] + kernel_size) + inside_w = (key_w[None, :] >= start_w[:, None]) & (key_w[None, :] < start_w[:, None] + kernel_size) + return inside_h & inside_w + + +def split_rope_dims(head_dim: int) -> tuple[int, int, int]: + """Split a head dimension into temporal, height, and width chunks. + + Leftover rotary pairs go to the spatial axes: correspondence is a spatial + matching task, while temporal positions only span the short offset window. + """ + if head_dim < 6 or head_dim % 2 != 0: + msg = f"3D RoPE requires an even head_dim >= 6, got {head_dim}" + raise ValueError(msg) + pairs, remainder = divmod(head_dim // 2, 3) + pair_counts = tuple(pairs + int(axis >= 3 - remainder) for axis in range(3)) + return 2 * pair_counts[0], 2 * pair_counts[1], 2 * pair_counts[2] + + +class AxialRoPE3D(nn.Module): + """Axial rotary embeddings for ``(time, height, width)`` coordinates.""" + + def __init__(self, head_dim: int, max_period: float = 10000.0) -> None: + """Initialize per-axis rotary frequencies.""" + super().__init__() + self.axis_dims = split_rope_dims(head_dim) + for axis_name, axis_dim in zip(("time", "height", "width"), self.axis_dims, strict=True): + self.register_buffer( + f"{axis_name}_frequencies", + get_freqs(axis_dim // 2, max_period), + persistent=False, + ) + + @torch.autocast(device_type="cuda", enabled=False) + def forward(self, x: Tensor) -> Tensor: + """Apply rotary embeddings and return the original working dtype.""" + working_dtype = x.dtype + chunks = x.to(torch.float32).split(self.axis_dims, dim=-1) + axes = (1, 2, 3) + frequencies = (self.time_frequencies, self.height_frequencies, self.width_frequencies) + rotated = [ + self.rotate_axis(chunk, axis, frequency) + for chunk, axis, frequency in zip(chunks, axes, frequencies, strict=True) + ] + return torch.cat(rotated, dim=-1).to(working_dtype) + + @staticmethod + def rotate_axis(x: Tensor, axis: int, frequencies: Tensor) -> Tensor: + """Rotate one channel chunk according to positions along one axis.""" + positions = torch.arange(x.shape[axis], device=x.device, dtype=torch.float32) + angles = torch.outer(positions, frequencies.to(device=x.device)) + broadcast_shape = [1] * x.ndim + broadcast_shape[axis] = positions.numel() + broadcast_shape[-1] = frequencies.numel() + cosine = angles.cos().reshape(broadcast_shape) + sine = angles.sin().reshape(broadcast_shape) + pairs = x.reshape(*x.shape[:-1], -1, 2) + first = pairs[..., 0] * cosine - pairs[..., 1] * sine + second = pairs[..., 0] * sine + pairs[..., 1] * cosine + return torch.stack((first, second), dim=-1).flatten(-2) + + +class MotionCorrespondenceBlock(nn.Module): + """Cross-frame spatial correspondence with an explicit no-match token.""" + + def __init__(self, spec: MotionCorrespondenceSpec) -> None: + """Initialize a motion-correspondence residual block.""" + super().__init__() + if spec.channels != spec.num_heads * spec.head_dim: + msg = f"channels must equal num_heads * head_dim, got {spec.channels} != {spec.num_heads} * {spec.head_dim}" + raise ValueError(msg) + if spec.spatial_kernel_size < 3 or spec.spatial_kernel_size % 2 == 0: + msg = "spatial_kernel_size must be odd and at least 3" + raise ValueError(msg) + if ( + not spec.temporal_offsets + or 0 in spec.temporal_offsets + or tuple(sorted(set(spec.temporal_offsets))) != spec.temporal_offsets + or set(spec.temporal_offsets) != {-offset for offset in spec.temporal_offsets} + ): + msg = "temporal_offsets must be non-zero, unique, sorted, and symmetric" + raise ValueError(msg) + self.channels = spec.channels + self.num_heads = spec.num_heads + self.head_dim = spec.head_dim + self.spatial_kernel_size = spec.spatial_kernel_size + self.temporal_offsets = spec.temporal_offsets + self.backend = spec.backend + self.natten_backend = spec.natten_backend + self.merge_compile = spec.merge_compile + + self.norm = nn.RMSNorm(spec.channels) + self.to_qkv = nn.Linear(spec.channels, 3 * spec.channels) + self.q_norm = nn.RMSNorm(spec.head_dim) + self.k_norm = nn.RMSNorm(spec.head_dim) + self.rope = AxialRoPE3D(spec.head_dim) + self.null_key = nn.Parameter(torch.randn(spec.num_heads, spec.head_dim) * 0.02) + self.null_value = nn.Parameter(torch.zeros(spec.num_heads, spec.head_dim)) + self.to_out = nn.Linear(spec.channels, spec.channels) + nn.init.zeros_(self.to_out.weight) + nn.init.zeros_(self.to_out.bias) + + def forward(self, x: Tensor) -> Tensor: + """Apply cross-frame correspondence to ``(B,C,T,H,W)`` features.""" + if x.shape[2] == 1: + return x + features = rearrange(x, "b c t h w -> b t h w c") + # fp32 norm inputs keep fused rms_norm kernels under bf16 autocast (same as DiT QK-norm). + qkv = self.to_qkv(self.norm(features.float()).type_as(features)) + q, k, v = [ + rearrange(part, "b t h w (nh d) -> b t h w nh d", nh=self.num_heads) for part in qkv.chunk(3, dim=-1) + ] + q_unrotated = self.q_norm(q.float()).type_as(q) + q = self.rope(q_unrotated) + k = self.rope(self.k_norm(k.float()).type_as(k)) + attention = self.compute_sdpa_attention if self.backend == "sdpa" else self.compute_natten_attention + context = attention(q, k, v, q_unrotated) + output = self.to_out(rearrange(context, "b t h w nh d -> b t h w (nh d)")) + return x + rearrange(output, "b t h w c -> b c t h w") + + def null_match_logits(self, q_unrotated: Tensor) -> Tensor: + """Score queries against the null token in the shared attention scale. + + Uses pre-RoPE queries: the null token has no position, so pairing it with + a rotated query would make the no-match threshold depend on where the + query sits in the clip. + """ + logits = torch.einsum("...nd,nd->...n", q_unrotated.float(), self.null_key.float()) + return logits * self.head_dim**-0.5 + + def compute_sdpa_attention(self, q: Tensor, k: Tensor, v: Tensor, q_unrotated: Tensor) -> Tensor: + """Compute the reference operation with PyTorch SDPA.""" + batch, frames, height, width, _heads, dim = q.shape + spatial_mask = spatial_neighborhood_mask(height, width, self.spatial_kernel_size, q.device) + null_logits = self.null_match_logits(q_unrotated) + frame_outputs: list[Tensor] = [] + for query_index in range(frames): + neighbor_indices = temporal_neighbor_indices(query_index, frames, self.temporal_offsets) + q_frame = rearrange(q[:, query_index], "b h w nh d -> b nh (h w) d") + k_frames = rearrange(k[:, neighbor_indices], "b t h w nh d -> b nh (t h w) d") + v_frames = rearrange(v[:, neighbor_indices], "b t h w nh d -> b nh (t h w) d") + # The null column carries a zero key, so its logit comes solely from the + # additive bias; SDPA applies the bias after q·k scaling, hence pre-scaled logits. + null_key = torch.zeros(batch, self.num_heads, 1, dim, dtype=q.dtype, device=q.device) + null_value = self.null_value.to(q)[None, :, None, :].expand(batch, -1, 1, -1) + keys = torch.cat((k_frames, null_key), dim=2) + values = torch.cat((v_frames, null_value), dim=2) + bias = self.sdpa_attention_bias(spatial_mask, null_logits[:, query_index], len(neighbor_indices)) + attended = functional.scaled_dot_product_attention(q_frame, keys, values, attn_mask=bias.to(q.dtype)) + frame_outputs.append(rearrange(attended, "b nh (h w) d -> b h w nh d", h=height, w=width)) + return torch.stack(frame_outputs, dim=1) + + def sdpa_attention_bias(self, spatial_mask: Tensor, frame_null_logits: Tensor, num_neighbors: int) -> Tensor: + """Build the additive bias: window gating for real keys, null logits in the last column.""" + window_bias = torch.where(spatial_mask.repeat(1, num_neighbors), 0.0, float("-inf")) + null_bias = rearrange(frame_null_logits, "b h w nh -> b nh (h w) 1") + window_bias = window_bias[None, None].expand(*null_bias.shape[:2], -1, -1) + return torch.cat((window_bias, null_bias), dim=-1) + + def compute_natten_attention(self, q: Tensor, k: Tensor, v: Tensor, q_unrotated: Tensor) -> Tensor: + """Compute correspondence with fused 2D neighborhood-attention kernels.""" + if q.device.type != "cuda": + msg = "NATTEN motion attention requires CUDA" + raise RuntimeError(msg) + natten = load_natten() + batch, frames, height, width, _heads, _dim = q.shape + if self.spatial_kernel_size > min(height, width): + msg = f"spatial kernel {self.spatial_kernel_size} exceeds feature map {height}x{width}" + raise ValueError(msg) + null_logits = self.null_match_logits(q_unrotated) + + grouped_queries: dict[int, list[tuple[int, tuple[int, ...]]]] = defaultdict(list) + for query_index in range(frames): + neighbors = temporal_neighbor_indices(query_index, frames, self.temporal_offsets) + grouped_queries[len(neighbors)].append((query_index, neighbors)) + + frame_outputs: list[Tensor | None] = [None] * frames + for group in grouped_queries.values(): + query_indices = tuple(query_index for query_index, _neighbors in group) + q_group = rearrange(q[:, query_indices], "b t h w nh d -> (b t) h w nh d").contiguous() + group_outputs: list[Tensor] = [] + group_lse: list[Tensor] = [] + for neighbor_position in range(len(group[0][1])): + neighbor_indices = tuple(neighbors[neighbor_position] for _query_index, neighbors in group) + k_group = rearrange(k[:, neighbor_indices], "b t h w nh d -> (b t) h w nh d").contiguous() + v_group = rearrange(v[:, neighbor_indices], "b t h w nh d -> (b t) h w nh d").contiguous() + output, lse = natten.na2d( + q_group, + k_group, + v_group, + kernel_size=(self.spatial_kernel_size, self.spatial_kernel_size), + backend=self.natten_backend, + return_lse=True, + ) + group_outputs.append(output.flatten(1, 2)) + group_lse.append(lse.flatten(1, 2)) + merged = merge_attention_branches( + group_outputs, + group_lse, + self.null_value.to(q), + rearrange(null_logits[:, query_indices], "b t h w nh -> (b t) (h w) nh"), + torch_compile=self.merge_compile, + ) + merged = rearrange( + merged, + "(b t) (h w) nh d -> b t h w nh d", + b=batch, + t=len(query_indices), + h=height, + w=width, + ) + for group_index, query_index in enumerate(query_indices): + frame_outputs[query_index] = merged[:, group_index] + if any(output is None for output in frame_outputs): + msg = "NATTEN correspondence did not produce every query frame" + raise RuntimeError(msg) + return torch.stack([output for output in frame_outputs if output is not None], dim=1) + + +def load_natten() -> ModuleType: + """Import NATTEN lazily and require its fused CUDA extension.""" + try: + natten = importlib.import_module("natten") + except ImportError as error: + msg = "NATTEN backend requested, but the natten package is not installed" + raise RuntimeError(msg) from error + if not getattr(natten, "HAS_LIBNATTEN", False): + msg = "NATTEN backend requested, but fused libnatten kernels are unavailable" + raise RuntimeError(msg) + return natten + + +"""Construction and execution helpers for the cascaded multi-scale model.""" + + +STAGE_FACTOR = 2 + + +class UpsampleSpec(BaseModel): + """Configuration shared by the two spatial upsample stages.""" + + mode: Literal["pixel_shuffle", "bilinear", "pxs_v2", "pxs_v2_hybrid"] + dims: int + temporal_mix: bool + icnr: bool + temporal_padding: TemporalPadding = "zeros" + padding_mode: UpsamplePaddingMode = "reflect" + + +class ResidualStackSpec(BaseModel): + """Configuration shared by residual stacks at all model resolutions. + + in_channels: input width of the stack's first block; None means use channels. + """ + + channels: int + mid_channels: int + total_blocks: int + stochastic_depth_rate: float + dims: int + layer_scale_init: float | None + grn: bool + depthwise: bool + kernel_size: int + zq_dim: int | None + in_channels: int | None = None + temporal_padding: TemporalPadding = "zeros" + + +class MotionAttentionSpec(BaseModel): + """Construction and placement contract for motion-correspondence blocks.""" + + after_mid_blocks: tuple[int, ...] + block: MotionCorrespondenceSpec + + +class BlockSequence(NamedTuple): + """Non-registering residual and motion-attention execution view.""" + + blocks: nn.Sequential + motion_blocks: nn.ModuleList | None = None + motion_after_blocks: tuple[int, ...] = () + + +def build_stem( + in_channels: int, + hidden_channels: int, + dims: int, + *, + bare: bool, + temporal_padding: TemporalPadding = "zeros", +) -> nn.Sequential: + """Build an input projection: bare conv, or the legacy conv + RMSNorm + SiLU.""" + conv = make_conv(dims, temporal_padding) + projection = conv(in_channels, hidden_channels, kernel_size=3, padding=1) + if bare: + return nn.Sequential(projection) + return nn.Sequential(projection, RMSNorm(hidden_channels, dims=dims), nn.SiLU()) + + +def build_output_head( + hidden_channels: int, + out_channels: int, + dims: int, + *, + zq_dim: int | None, + temporal_padding: TemporalPadding = "zeros", +) -> nn.Sequential: + """Build a norm + SiLU + conv output head, zq-conditioned when ``zq_dim`` is set.""" + conv = make_conv(dims, temporal_padding) + norm: nn.Module = RMSNorm(hidden_channels, dims=dims) + if zq_dim is not None: + norm = ModulatedRMSNorm(hidden_channels, zq_dim, dims=dims) + return nn.Sequential(norm, nn.SiLU(), conv(hidden_channels, out_channels, kernel_size=3, padding=1)) + + +def apply_output_head(head: nn.Sequential, x: Tensor, zq: Tensor | None) -> Tensor: + """Run an output head, threading ``zq`` through a modulated first norm.""" + norm = head[0] + x = norm(x, zq) if isinstance(norm, ModulatedRMSNorm) else norm(x) + x = head[1](x) + return head[2](x) + + +def build_stage_upsample(channels: int, spec: UpsampleSpec) -> nn.Module: + """Build one 2x spatial upsample stage.""" + if spec.mode == "pixel_shuffle": + return PixelShuffleND( + channels, + dims=spec.dims, + factor=STAGE_FACTOR, + temporal_mix=spec.temporal_mix, + icnr=spec.icnr, + temporal_padding=spec.temporal_padding, + ) + if spec.mode == "bilinear": + return BilinearUpsampleND(dims=spec.dims, factor=STAGE_FACTOR) + if spec.mode == "pxs_v2": + return PXSv2UpsampleND(channels, dims=spec.dims, factor=STAGE_FACTOR, padding_mode=spec.padding_mode) + return PXSv2HybridUpsampleND( + channels, + dims=spec.dims, + factor=STAGE_FACTOR, + temporal_mix=spec.temporal_mix, + icnr=spec.icnr, + temporal_padding=spec.temporal_padding, + padding_mode=spec.padding_mode, + ) + + +def build_residual_stack(count: int, start_index: int, spec: ResidualStackSpec) -> nn.Sequential: + """Build residual blocks with linearly scheduled stochastic depth.""" + blocks: list[nn.Module] = [] + for index in range(count): + drop_probability = spec.stochastic_depth_rate * (start_index + index) / max(spec.total_blocks - 1, 1) + block_in = spec.in_channels if index == 0 and spec.in_channels is not None else spec.channels + blocks.append( + ResidualBlock( + block_in, + out_channels=spec.channels, + mid_channels=spec.mid_channels, + dims=spec.dims, + layer_scale_init=spec.layer_scale_init, + grn=spec.grn, + stochastic_depth_prob=drop_probability, + depthwise=spec.depthwise, + kernel_size=spec.kernel_size, + zq_dim=spec.zq_dim, + temporal_padding=spec.temporal_padding, + ) + ) + return nn.Sequential(*blocks) + + +def build_motion_attention_stack(spec: MotionAttentionSpec) -> nn.ModuleList: + """Build one motion block for every configured mid-stage placement.""" + return nn.ModuleList([MotionCorrespondenceBlock(spec.block) for _placement in spec.after_mid_blocks]) + + +def apply_block_sequence( + x: Tensor, + zq: Tensor | None, + sequence: BlockSequence, + *, + use_checkpointing: bool, +) -> Tensor: + """Run residual blocks with scheduled motion correspondence.""" + for index, block in enumerate(sequence.blocks): + inputs = (x, zq) if zq is not None else (x,) + x = forward_with_checkpointing(block, *inputs, use_checkpointing=use_checkpointing) + block_number = index + 1 + if sequence.motion_blocks is not None and block_number in sequence.motion_after_blocks: + motion_index = sequence.motion_after_blocks.index(block_number) + x = forward_with_checkpointing( + sequence.motion_blocks[motion_index], + x, + # NATTEN merge_attentions custom autograd is incompatible with + # PyTorch non-reentrant activation checkpointing. + use_checkpointing=False, + ) + return x + + +"""Convolutional latent-space spatial upsampler with configurable upsampling strategy.""" + + +class ConvLatentUpsampler(nn.Module): + """Lightweight convolutional spatial upsampler for VAE latents. + + Upsampling strategies (controlled by ``upsample_mode``): + pixel_shuffle: Learned channel-to-space rearrangement (default). + bilinear: Parameter-free bilinear interpolation. + """ + + def __init__( + self, + in_channels: int = 16, + hidden_channels: int = 64, + bottleneck_channels: int | None = None, + num_pre_residual_blocks: int = 0, + num_residual_blocks: int = 16, + upscale_factor: int = 2, + *, + input_skip: bool = True, + upsample_mode: UpsampleMode = "pixel_shuffle", + upsample_position: UpsamplePosition = "after_projection", + temporal_mix: bool = True, + icnr: bool = False, + gradient_checkpointing: bool = False, + dims: int = 2, + expand_ratio: int = 4, + kernel_size: int = 3, + layer_scale_init: float | None = None, + grn: bool = False, + stochastic_depth_rate: float = 0.0, + depthwise: bool = False, + stem_channels: int | None = None, + zq_dim: int | None = None, + temporal_padding: TemporalPadding = "zeros", + upsample_padding_mode: UpsamplePaddingMode = "reflect", + ) -> None: + """Initialize the upsampler. + + Args: + in_channels: Number of VAE latent channels. + hidden_channels: Internal width of residual blocks. + bottleneck_channels: Mid-channel bottleneck width for ResidualBlocks (None = hidden_channels). + num_pre_residual_blocks: Number of residual blocks before upsampling (default 0). + num_residual_blocks: Number of post-upsample residual refinement blocks. + upscale_factor: Spatial upscale factor for pixel_shuffle and bilinear modes. + input_skip: Add ``repeat_interleave`` skip connection from input to ``conv_in`` output. + upsample_mode: Upsampling strategy — "pixel_shuffle" or "bilinear" (parameter-free interpolation). + upsample_position: Where to apply upsampling — "before_projection" or "after_projection". + temporal_mix: Whether pixel-shuffle Conv3d mixes across the temporal dimension. + When False, uses (1,3,3) kernel. Only affects dims=3 + pixel_shuffle. + icnr: Apply ICNR initialization to the sub-pixel-conv weight inside the + ``pixel_shuffle`` and ``pxs_v2_hybrid`` upsample modes. Removes + checkerboard at init. Ignored by ``bilinear`` and ``pxs_v2``. + gradient_checkpointing: Enable gradient checkpointing for residual blocks. + dims: Convolution dimensionality — 2 for Conv2d, 3 for Conv3d. + expand_ratio: Expansion ratio for inverted bottleneck (used when ``bottleneck_channels`` is None). + kernel_size: Kernel size for depthwise convolutions. + layer_scale_init: Initial LayerScale gamma value. ``None`` disables LayerScale. + grn: Whether to insert Global Response Normalization in residual blocks. + stochastic_depth_rate: Maximum drop rate for stochastic depth (linearly scheduled). + depthwise: Use depthwise separable convolutions in residual blocks. + stem_channels: When set, pre-residual blocks and input projection operate at this + narrower width, with a ResidualBlock transition to ``hidden_channels`` after upsampling. + zq_dim: When set, every ``ResidualBlock`` uses ``ModulatedRMSNorm`` conditioned on + a side tensor (the LQ latent), re-injecting LQ structure at every block via FiLM. + temporal_padding: How ``dims == 3`` convolutions extend T — ``"zeros"`` + (default) or ``"replicate"``, which repeats the edge frame as K-VAE does. + upsample_padding_mode: Edge handling for the ``pxs_v2`` residual conv — + ``"reflect"`` (default) or ``"zeros"``, matching the production K-VAE. + """ + super().__init__() + self.in_channels = in_channels + self.hidden_channels = hidden_channels + self.num_pre_residual_blocks = num_pre_residual_blocks + self.num_residual_blocks = num_residual_blocks + self.upscale_factor = upscale_factor + self.upsample_mode = upsample_mode + self.upsample_position = upsample_position + self.input_skip = input_skip + self.gradient_checkpointing = gradient_checkpointing + self.dims = dims + self.zq_dim = zq_dim + + # Channel widths: stem (pre-upsample) vs main (post-upsample) + proj_channels = stem_channels if stem_channels is not None else hidden_channels + post_bottleneck = bottleneck_channels if bottleneck_channels is not None else hidden_channels * expand_ratio + pre_bottleneck = bottleneck_channels if bottleneck_channels is not None else proj_channels * expand_ratio + self.bottleneck_channels = post_bottleneck + + # Projection layers + conv = make_conv(dims, temporal_padding) + self.input_proj = nn.Sequential( + conv(in_channels, proj_channels, kernel_size=3, padding=1), + RMSNorm(proj_channels, dims=dims), + nn.SiLU(), + ) + self.output_proj = nn.Sequential( + RMSNorm(hidden_channels, dims=dims), + nn.SiLU(), + conv(hidden_channels, in_channels, kernel_size=3, padding=1), + ) + + # Spatial upsampling + upsample_channels = in_channels if upsample_position == "before_projection" else proj_channels + if upsample_mode == "pixel_shuffle": + self.upsample: nn.Module = PixelShuffleND( + upsample_channels, + dims=dims, + factor=upscale_factor, + temporal_mix=temporal_mix, + icnr=icnr, + temporal_padding=temporal_padding, + ) + elif upsample_mode == "bilinear": + self.upsample = BilinearUpsampleND(dims=dims, factor=upscale_factor) + elif upsample_mode == "pxs_v2": + self.upsample = PXSv2UpsampleND( + upsample_channels, + dims=dims, + factor=upscale_factor, + padding_mode=upsample_padding_mode, + ) + else: + self.upsample = PXSv2HybridUpsampleND( + upsample_channels, + dims=dims, + factor=upscale_factor, + temporal_mix=temporal_mix, + icnr=icnr, + temporal_padding=temporal_padding, + padding_mode=upsample_padding_mode, + ) + + # Channel transition: bridges stem width → hidden width after upsampling + self.channel_transition: ResidualBlock | None = None + if stem_channels is not None: + self.channel_transition = ResidualBlock( + proj_channels, + out_channels=hidden_channels, + mid_channels=post_bottleneck, + dims=dims, + grn=grn, + zq_dim=zq_dim, + temporal_padding=temporal_padding, + ) + + # Residual stacks + total_blocks = num_pre_residual_blocks + num_residual_blocks + shared_block_kwargs = { + "total_blocks": total_blocks, + "stochastic_depth_rate": stochastic_depth_rate, + "dims": dims, + "layer_scale_init": layer_scale_init, + "grn": grn, + "depthwise": depthwise, + "kernel_size": kernel_size, + "zq_dim": zq_dim, + "temporal_padding": temporal_padding, + } + self.pre_residual_blocks = self._build_residual_stack( + num_pre_residual_blocks, + proj_channels, + pre_bottleneck, + start_idx=0, + **shared_block_kwargs, # type: ignore[arg-type] + ) + self.residual_blocks = self._build_residual_stack( + num_residual_blocks, + hidden_channels, + post_bottleneck, + start_idx=num_pre_residual_blocks, + **shared_block_kwargs, # type: ignore[arg-type] + ) + + def _build_residual_stack( + self, + count: int, + channels: int, + mid_channels: int, + *, + start_idx: int, + total_blocks: int, + stochastic_depth_rate: float, + dims: int, + layer_scale_init: float | None, + grn: bool, + depthwise: bool, + kernel_size: int, + zq_dim: int | None = None, + temporal_padding: TemporalPadding = "zeros", + ) -> nn.Sequential: + """Build a stack of ResidualBlocks with linearly scheduled stochastic depth.""" + blocks: list[nn.Module] = [] + for i in range(count): + drop_prob = stochastic_depth_rate * (start_idx + i) / max(total_blocks - 1, 1) + blocks.append( + ResidualBlock( + channels, + mid_channels=mid_channels, + dims=dims, + layer_scale_init=layer_scale_init, + grn=grn, + stochastic_depth_prob=drop_prob, + depthwise=depthwise, + kernel_size=kernel_size, + zq_dim=zq_dim, + temporal_padding=temporal_padding, + ) + ) + return nn.Sequential(*blocks) + + def _apply_blocks( + self, + x: Tensor, + blocks: nn.Sequential, + *, + use_checkpointing: bool, + zq: Tensor | None = None, + ) -> Tensor: + """Apply a residual stack with optional modulation.""" + for block in blocks: + if zq is not None: + x = forward_with_checkpointing(block, x, zq, use_checkpointing=use_checkpointing) + else: + x = forward_with_checkpointing(block, x, use_checkpointing=use_checkpointing) + return x + + def _apply_input_projection(self, x: Tensor) -> Tensor: + z_skip = x + x = self.input_proj(x) + if self.input_skip: + repeats = x.shape[1] // z_skip.shape[1] + x = x + z_skip.repeat_interleave(repeats, dim=1) + return x + + @property + def last_layer_weight(self) -> torch.Tensor: + """Weight of the final output conv (used by adaptive GAN-weight balancing).""" + return self.output_proj[-1].weight # type: ignore + + def forward(self, z: torch.Tensor) -> torch.Tensor: + """Upsample latent tensor spatially. + + Args: + z: Input latent tensor of shape ``(B, C, T, H, W)``. + + Returns: + Upsampled tensor of shape ``(B, C, T, H', W')``. + """ + b, _, t, _, _ = z.shape + x = rearrange(z, "b c t h w -> (b t) c h w") if self.dims == 2 else z + zq = x if self.zq_dim is not None else None + if self.upsample_position == "before_projection": + x = self.upsample(x) + x = self._apply_input_projection(x) + use_ckpt = self.training and self.gradient_checkpointing + x = self._apply_blocks( + x, + self.pre_residual_blocks, + use_checkpointing=use_ckpt, + zq=zq, + ) + if self.upsample_position == "after_projection": + x = self.upsample(x) + if self.channel_transition is not None: + if zq is not None: + x = forward_with_checkpointing(self.channel_transition, x, zq, use_checkpointing=use_ckpt) + else: + x = forward_with_checkpointing(self.channel_transition, x, use_checkpointing=use_ckpt) + x = self._apply_blocks( + x, + self.residual_blocks, + use_checkpointing=use_ckpt, + zq=zq, + ) + x = self.output_proj(x) + return rearrange(x, "(b t) c h w -> b c t h w", b=b, t=t) if self.dims == 2 else x + + +"""Scale-specific x2 adaptation modules for the cascaded latent upscaler.""" + + +X2TailMode = Literal["shared", "scale_specific_norm", "private", "private_full"] +X2FinisherMode = Literal["none", "linear", "pxs_residual"] + + +class SharedSecondStage(NamedTuple): + """Non-registering view of the x4 second-stage modules.""" + + upsample: nn.Module + blocks: nn.Sequential + output_proj: nn.Sequential + mid_blocks: nn.Sequential | None = None + + +class ScaleSpecificBlockNorms(nn.Module): + """Private normalization copies for one shared residual block.""" + + def __init__(self, block: ResidualBlock) -> None: + """Copy the modulated norms from a non-depthwise residual block.""" + super().__init__() + if block.zq_dim is None or block.depthwise: + msg = "scale-specific norms require non-depthwise modulated residual blocks" + raise ValueError(msg) + self.norm1 = copy.deepcopy(block.norm1) + self.norm2 = copy.deepcopy(block.norm2) + + def copy_from(self, block: ResidualBlock) -> None: + """Reset the private norms from their shared x4 counterparts.""" + self.norm1.load_state_dict(block.norm1.state_dict()) + self.norm2.load_state_dict(block.norm2.state_dict()) + + +def require_pxs_upsample(source: nn.Module) -> PXSv2UpsampleND: + """Require a pxs_v2 upsample stage as the finisher donor.""" + if not isinstance(source, PXSv2UpsampleND): + msg = f"x2_finisher requires a PXSv2UpsampleND source, got {type(source).__name__}" + raise TypeError(msg) + return source + + +class X2Finisher(nn.Module): + """The upsample-stage conv pair applied at the adapter grid without the nearest-2x resize.""" + + def __init__(self, source: nn.Module, mode: X2FinisherMode) -> None: + """Clone the finisher convs from a pxs_v2 upsample stage.""" + super().__init__() + if mode == "none": + msg = "X2Finisher must not be built with x2_finisher='none'" + raise ValueError(msg) + donor = require_pxs_upsample(source) + self.mode = mode + self.linear = copy.deepcopy(donor.linear) + self.spatial_conv = copy.deepcopy(donor.spatial_conv) if mode == "pxs_residual" else None + + def copy_from(self, source: nn.Module) -> None: + """Reset the finisher convs from their x4 upsample counterparts.""" + donor = require_pxs_upsample(source) + self.linear.load_state_dict(donor.linear.state_dict()) + if self.spatial_conv is not None: + self.spatial_conv.load_state_dict(donor.spatial_conv.state_dict()) + + def forward(self, x: Tensor) -> Tensor: + """Apply the cloned convs on the unchanged grid.""" + if self.spatial_conv is not None: + x = x + self.spatial_conv(x) + return self.linear(x) + + +class X2Branch(nn.Module): + """Modules exclusive to the frozen-e115 x2 adaptation path.""" + + def __init__( + self, + adapter: nn.Sequential, + tail_mode: X2TailMode, + shared_stage: SharedSecondStage, + *, + dims: int, + finisher: X2Finisher | None = None, + ) -> None: + """Build the selected scale-specific capacity around a shared stage.""" + super().__init__() + self.adapter = adapter + self.tail_mode = tail_mode + self.dims = dims + self.finisher = finisher + + self.post_norms: nn.ModuleList | None = None + self.output_norm: RMSNorm | None = None + self.private_mid_blocks: nn.Sequential | None = None + self.private_upsample: nn.Module | None = None + self.private_blocks: nn.Sequential | None = None + self.private_output_proj: nn.Sequential | None = None + + if tail_mode == "scale_specific_norm": + self.post_norms = nn.ModuleList([ScaleSpecificBlockNorms(block) for block in shared_stage.blocks]) # type: ignore + self.output_norm = copy.deepcopy(shared_stage.output_proj[0]) # type: ignore + elif tail_mode in ("private", "private_full"): + self.private_upsample = copy.deepcopy(shared_stage.upsample) + self.private_blocks = copy.deepcopy(shared_stage.blocks) + self.private_output_proj = copy.deepcopy(shared_stage.output_proj) + if tail_mode == "private_full": + if shared_stage.mid_blocks is None: + msg = "x2_tail_mode='private_full' requires the shared stage to expose mid_blocks" + raise ValueError(msg) + self.private_mid_blocks = copy.deepcopy(shared_stage.mid_blocks) + + def initialize_from_x4(self, shared_stage: SharedSecondStage) -> None: + """Copy scale-specific tail state from the loaded x4 baseline.""" + if self.tail_mode == "scale_specific_norm": + if self.post_norms is None or self.output_norm is None: + msg = "scale-specific normalization modules are not initialized" + raise RuntimeError(msg) + for norms, block in zip(self.post_norms, shared_stage.blocks, strict=True): + norms.copy_from(block) + self.output_norm.load_state_dict(shared_stage.output_proj[0].state_dict()) + elif self.tail_mode in ("private", "private_full"): + if self.private_upsample is None or self.private_blocks is None or self.private_output_proj is None: + msg = "private x2 tail modules are not initialized" + raise RuntimeError(msg) + self.private_upsample.load_state_dict(shared_stage.upsample.state_dict()) + self.private_blocks.load_state_dict(shared_stage.blocks.state_dict()) + self.private_output_proj.load_state_dict(shared_stage.output_proj.state_dict()) + if self.private_mid_blocks is not None: + if shared_stage.mid_blocks is None: + msg = "private mid blocks require the shared stage to expose mid_blocks" + raise RuntimeError(msg) + self.private_mid_blocks.load_state_dict(shared_stage.mid_blocks.state_dict()) + + def private_stage(self, shared_stage: SharedSecondStage) -> SharedSecondStage: + """Select the shared stage or the registered private copy.""" + if self.tail_mode not in ("private", "private_full"): + return shared_stage + if self.private_upsample is None or self.private_blocks is None or self.private_output_proj is None: + msg = "private x2 tail modules are not initialized" + raise RuntimeError(msg) + return SharedSecondStage( + upsample=self.private_upsample, + blocks=self.private_blocks, + output_proj=self.private_output_proj, + ) + + def apply_post_blocks( + self, + x: Tensor, + zq: Tensor | None, + stage: SharedSecondStage, + *, + use_checkpointing: bool, + ) -> Tensor: + """Apply post blocks with shared or x2-specific normalization.""" + for index, block in enumerate(stage.blocks): + if self.tail_mode == "scale_specific_norm": + if zq is None or self.post_norms is None: + msg = "scale-specific normalization requires zq and private norms" + raise RuntimeError(msg) + norms = self.post_norms[index] + block_call = partial(block.forward_with_norms, norm1=norms.norm1, norm2=norms.norm2) # type: ignore + x = forward_with_checkpointing(block_call, x, zq, use_checkpointing=use_checkpointing) + elif zq is None: + x = forward_with_checkpointing(block, x, use_checkpointing=use_checkpointing) + else: + x = forward_with_checkpointing(block, x, zq, use_checkpointing=use_checkpointing) + return x + + def apply_output_projection(self, x: Tensor, zq: Tensor | None, stage: SharedSecondStage) -> Tensor: + """Project features with a shared or scale-specific final norm.""" + if self.tail_mode != "scale_specific_norm": + return apply_output_head(stage.output_proj, x, zq) + if self.output_norm is None: + msg = "scale-specific output norm is not initialized" + raise RuntimeError(msg) + x = self.output_norm(x) + x = stage.output_proj[1](x) + return stage.output_proj[2](x) + + def apply_block_stack( + self, blocks: nn.Sequential, x: Tensor, zq: Tensor | None, *, use_checkpointing: bool + ) -> Tensor: + """Run a residual stack with the branch's zq-threading convention.""" + for block in blocks: + if zq is None: + x = forward_with_checkpointing(block, x, use_checkpointing=use_checkpointing) + else: + x = forward_with_checkpointing(block, x, zq, use_checkpointing=use_checkpointing) + return x + + def apply_adapter(self, x: Tensor, zq: Tensor | None, *, use_checkpointing: bool) -> Tensor: + """Run the x2-only residual adapter.""" + return self.apply_block_stack(self.adapter, x, zq, use_checkpointing=use_checkpointing) + + def apply_mid_blocks(self, x: Tensor, zq: Tensor | None, *, use_checkpointing: bool) -> Tensor: + """Run the private mid blocks ahead of the second-stage upsample (private_full only).""" + if self.private_mid_blocks is None: + return x + return self.apply_block_stack(self.private_mid_blocks, x, zq, use_checkpointing=use_checkpointing) + + def forward_second_stage( + self, + x: Tensor, + zq: Tensor | None, + batch_size: int, + frames: int, + shared_stage: SharedSecondStage, + *, + use_checkpointing: bool, + ) -> Tensor: + """Run the selected x2 second stage without registering shared aliases.""" + stage = self.private_stage(shared_stage) + if self.finisher is not None: + x = self.finisher(x) + x = self.apply_mid_blocks(x, zq, use_checkpointing=use_checkpointing) + x = stage.upsample(x) + x = self.apply_post_blocks(x, zq, stage, use_checkpointing=use_checkpointing) + x = self.apply_output_projection(x, zq, stage) + if self.dims == 2: + return rearrange(x, "(b t) c h w -> b c t h w", b=batch_size, t=frames) + return x + + +"""Multi-scale cascaded upsampler: two 2x stages with intermediate supervision.""" + + +# The two projections back to latent space. They are scale-specific, so a warm +# start from another network cannot supply them and a warm-started run may want +# to fit them before letting gradients reach the rest. +HEAD_PREFIXES = ("output_proj.", "mid_output_head.") + + +class MultiScaleUpsampler(nn.Module): + """Cascaded 2x+2x VAE-latent upsampler with an optional x2 entry. + + The x4 route keeps the original two-stage architecture. The x2 route enters + after ``mid_blocks`` through a private stem, optional residual adapter, and + a shared, scale-normalized, or private second stage. ``x2_tail_mode`` + ``'private_full'`` instead routes through the branch's own warm copies of + ``mid_blocks`` and the second stage, behind an optional finisher cloned + from ``upsample_1``. + """ + + def __init__( + self, + in_channels: int = 16, + hidden_channels: int = 64, + bottleneck_channels: int | None = None, + num_pre_blocks: int = 2, + num_mid_blocks: int = 6, + num_post_blocks: int = 8, + *, + input_skip: bool = True, + upsample_mode: UpsampleMode = "pixel_shuffle", + temporal_mix: bool = True, + icnr: bool = False, + gradient_checkpointing: bool = False, + dims: int = 3, + expand_ratio: int = 4, + kernel_size: int = 3, + layer_scale_init: float | None = None, + grn: bool = False, + stochastic_depth_rate: float = 0.0, + depthwise: bool = False, + motion_attention: MotionAttentionSpec | None = None, + zq_dim: int | None = None, + bare_stem: bool = False, + modulated_output_proj: bool = False, + global_skip: bool = False, + enable_x2_entry: bool = False, + x2_adapter_blocks: int = 0, + x2_tail_mode: X2TailMode = "shared", + x2_adapter_sources: tuple[int, ...] | None = None, + x2_finisher: X2FinisherMode = "none", + stage_channels: tuple[int, int, int] | None = None, + temporal_padding: TemporalPadding = "zeros", + upsample_padding_mode: UpsamplePaddingMode = "reflect", + ) -> None: + """Initialize MultiScaleUpsampler with the given architecture parameters.""" + super().__init__() + self.in_channels = in_channels + self.hidden_channels = hidden_channels + self.upscale_factor = 4 + self.input_skip = input_skip + self.gradient_checkpointing = gradient_checkpointing + self.dims = dims + self.zq_dim = zq_dim + self.global_skip = global_skip + self.enable_x2_entry = enable_x2_entry + self.x2_tail_mode = x2_tail_mode + self.x2_adapter_sources = x2_adapter_sources + + if modulated_output_proj and zq_dim is None: + msg = "modulated_output_proj requires zq_dim (modulated_norm)" + raise ValueError(msg) + if enable_x2_entry and global_skip: + msg = "global_skip does not support the x2 entry" + raise ValueError(msg) + if enable_x2_entry and modulated_output_proj and x2_tail_mode != "private_full": + msg = "modulated_output_proj with the x2 entry requires x2_tail_mode='private_full'" + raise ValueError(msg) + if x2_adapter_sources is not None and len(x2_adapter_sources) != x2_adapter_blocks: + msg = "x2_adapter_sources must list one pre_blocks index per adapter block" + raise ValueError(msg) + if x2_adapter_sources is not None and any(not 0 <= index < num_pre_blocks for index in x2_adapter_sources): + msg = "x2_adapter_sources indices must be within [0, num_pre_blocks)" + raise ValueError(msg) + if stage_channels is not None: + if enable_x2_entry and x2_tail_mode != "private_full": + msg = "stage_channels with enable_x2_entry requires x2_tail_mode='private_full'" + raise ValueError(msg) + if motion_attention is not None: + msg = "stage_channels is incompatible with motion_attention" + raise ValueError(msg) + if input_skip: + msg = "stage_channels is incompatible with input_skip" + raise ValueError(msg) + if bottleneck_channels is not None: + msg = "stage_channels is incompatible with bottleneck_channels" + raise ValueError(msg) + if hidden_channels != stage_channels[0]: + msg = f"hidden_channels ({hidden_channels}) must equal stage_channels[0] ({stage_channels[0]})" + raise ValueError(msg) + self.stage_channels = stage_channels + w1, w2, w3 = stage_channels if stage_channels is not None else (hidden_channels,) * 3 + self.motion_after_mid_blocks = motion_attention.after_mid_blocks if motion_attention is not None else () + + if input_skip and hidden_channels % in_channels != 0: + msg = f"hidden_channels ({hidden_channels}) must be divisible by in_channels ({in_channels}) for input_skip" + raise ValueError(msg) + if motion_attention is not None and dims != 3: + msg = "motion_attention requires dims=3" + raise ValueError(msg) + if motion_attention is not None and enable_x2_entry: + msg = "motion_attention Round 23 is x4-only and cannot enable the x2 entry" + raise ValueError(msg) + if motion_attention is not None and not motion_attention.after_mid_blocks: + msg = "motion_attention requires at least one mid-stage placement" + raise ValueError(msg) + if motion_attention is not None and motion_attention.after_mid_blocks[-1] > num_mid_blocks: + msg = "motion attention placement exceeds the mid residual stack" + raise ValueError(msg) + + # Projection layers + self.input_proj = build_stem(in_channels, w1, dims, bare=bare_stem, temporal_padding=temporal_padding) + self.mid_output_head = build_output_head(w2, in_channels, dims, zq_dim=None, temporal_padding=temporal_padding) + self.output_proj = build_output_head( + w3, + in_channels, + dims, + zq_dim=zq_dim if modulated_output_proj else None, + temporal_padding=temporal_padding, + ) + + # Two 2x upsample stages + upsample_spec = UpsampleSpec( + dims=dims, + mode=upsample_mode, + temporal_mix=temporal_mix, + icnr=icnr, + temporal_padding=temporal_padding, + padding_mode=upsample_padding_mode, + ) + self.upsample_1 = build_stage_upsample(w1, upsample_spec) + self.upsample_2 = build_stage_upsample(w2, upsample_spec) + + # Residual stacks with linear stochastic depth across all stages; the + # first block of a stage narrows from the previous stage width. + total_blocks = num_pre_blocks + num_mid_blocks + num_post_blocks + + def stage_spec(channels: int, first_in: int | None) -> ResidualStackSpec: + return ResidualStackSpec( + channels=channels, + mid_channels=bottleneck_channels if bottleneck_channels is not None else channels * expand_ratio, + total_blocks=total_blocks, + stochastic_depth_rate=stochastic_depth_rate, + dims=dims, + layer_scale_init=layer_scale_init, + grn=grn, + depthwise=depthwise, + kernel_size=kernel_size, + zq_dim=zq_dim, + in_channels=first_in, + temporal_padding=temporal_padding, + ) + + self.pre_blocks = build_residual_stack(num_pre_blocks, 0, stage_spec(w1, None)) + self.mid_blocks = build_residual_stack(num_mid_blocks, num_pre_blocks, stage_spec(w2, w1 if w1 != w2 else None)) + self.post_blocks = build_residual_stack( + num_post_blocks, num_pre_blocks + num_mid_blocks, stage_spec(w3, w2 if w2 != w3 else None) + ) + + self.mid_motion_blocks: nn.ModuleList | None = None + if motion_attention is not None: + self.mid_motion_blocks = build_motion_attention_stack(motion_attention) + + # Build x2-exclusive modules after the complete x4 path so enabling an x2 + # entry does not perturb the base-path RNG initialization. + self.mid_input_proj: nn.Sequential | None = None + self.x2_branch: X2Branch | None = None + if enable_x2_entry: + self.mid_input_proj = build_stem( + in_channels, hidden_channels, dims, bare=bare_stem, temporal_padding=temporal_padding + ) + adapter_spec = stage_spec(w1, None).model_copy( + update={ + "total_blocks": max(x2_adapter_blocks, 1), + "stochastic_depth_rate": 0.0, + } + ) + adapter = build_residual_stack(x2_adapter_blocks, 0, adapter_spec) + for block in adapter: + block.zero_init_residual() # type: ignore[attr-defined] + finisher = X2Finisher(self.upsample_1, x2_finisher) if x2_finisher != "none" else None + self.x2_branch = X2Branch( + adapter, + x2_tail_mode, + self.shared_second_stage(), + dims=dims, + finisher=finisher, + ) + + if global_skip: + # Zero heads so step 0 reproduces the nearest base exactly at both scales. + for head in (self.mid_output_head, self.output_proj): + final_conv = head[-1] + nn.init.zeros_(final_conv.weight) # type: ignore[arg-type] + nn.init.zeros_(final_conv.bias) # type: ignore[arg-type] + + @property + def last_layer_weight(self) -> torch.Tensor: + """Weight of the final 4x-output conv (used by adaptive GAN-weight balancing). + + The GAN-loss runs only on the 4x output, so the relevant last layer is + ``output_proj[-1]`` rather than the 2x ``mid_output_head``. + """ + return self.output_proj[-1].weight # type: ignore + + def shared_second_stage(self) -> SharedSecondStage: + """Return a non-registering view of the e115-compatible second stage.""" + return SharedSecondStage( + upsample=self.upsample_2, + blocks=self.post_blocks, + output_proj=self.output_proj, + mid_blocks=self.mid_blocks, + ) + + @staticmethod + def is_x2_state_key(key: str) -> bool: + """Return whether a state key belongs exclusively to the x2 path.""" + return key.startswith(("mid_input_proj.", "x2_branch.")) + + def initialize_x2_from_x4(self) -> None: + """Initialize x2-exclusive modules from the currently loaded x4 path.""" + if self.mid_input_proj is None or self.x2_branch is None: + msg = "x2 initialization requires enable_x2_entry=True" + raise ValueError(msg) + self.mid_input_proj.load_state_dict(self.input_proj.state_dict()) + self.x2_branch.initialize_from_x4(self.shared_second_stage()) + if self.x2_adapter_sources is not None: + for block, source_index in zip(self.x2_branch.adapter, self.x2_adapter_sources, strict=True): + block.load_state_dict(self.pre_blocks[source_index].state_dict()) + if self.x2_branch.finisher is not None: + self.x2_branch.finisher.copy_from(self.upsample_1) + + def configure_x2_finetuning(self) -> None: + """Freeze the e115 path and expose only x2-exclusive parameters.""" + if self.mid_input_proj is None or self.x2_branch is None: + msg = "x2 fine-tuning requires enable_x2_entry=True" + raise ValueError(msg) + self.requires_grad_(requires_grad=False) + self.mid_input_proj.requires_grad_(requires_grad=True) + self.x2_branch.requires_grad_(requires_grad=True) + + def _project_with_skip(self, x: Tensor, proj: nn.Module) -> Tensor: + """Apply a projection and optionally add the channel-repeated input skip.""" + projected = proj(x) + if self.input_skip: + repeats = projected.shape[1] // x.shape[1] + projected = projected + x.repeat_interleave(repeats, dim=1) + return projected + + def _head_x4(self, x: Tensor, zq: Tensor | None, *, use_ckpt: bool) -> Tensor: + """x4 entry: ``input_proj -> input_skip -> pre_blocks -> upsample_1`` (-> 2x grid).""" + x = self._project_with_skip(x, self.input_proj) + x = apply_block_sequence( + x, + zq, + BlockSequence(self.pre_blocks), + use_checkpointing=use_ckpt, + ) + return self.upsample_1(x) + + def _head_x2(self, x: Tensor, zq: Tensor | None, *, use_ckpt: bool) -> Tensor: + """Project a genuine x2 latent and apply the scale-specific adapter.""" + if self.mid_input_proj is None or self.x2_branch is None: + msg = "x2 head requires enable_x2_entry=True" + raise ValueError(msg) + x = self._project_with_skip(x, self.mid_input_proj) + return self.x2_branch.apply_adapter(x, zq, use_checkpointing=use_ckpt) + + def _tail( + self, + x: Tensor, + zq: Tensor | None, + b: int, + t: int, + *, + use_ckpt: bool, + return_intermediates: bool, + ) -> dict[str, torch.Tensor] | torch.Tensor: + """x4 continuation: ``mid_blocks -> [2x head] -> second stage``.""" + x = apply_block_sequence( + x, + zq, + BlockSequence(self.mid_blocks, self.mid_motion_blocks, self.motion_after_mid_blocks), + use_checkpointing=use_ckpt, + ) + + # 2x supervision branch (only when the caller will use it) + z_2x = None + if return_intermediates: + z_2x_inner = self.mid_output_head(x) + z_2x = rearrange(z_2x_inner, "(b t) c h w -> b c t h w", b=b, t=t) if self.dims == 2 else z_2x_inner + + z_4x = self._second_stage(x, zq, b, t, use_ckpt=use_ckpt) + if return_intermediates: + return {"2x": z_2x, "4x": z_4x} # type: ignore[return-value] + return z_4x + + def _second_stage(self, x: Tensor, zq: Tensor | None, b: int, t: int, *, use_ckpt: bool) -> Tensor: + """Shared second stage: ``upsample_2 -> post_blocks -> output_proj`` — 2x of the current grid.""" + x = self.upsample_2(x) + x = apply_block_sequence(x, zq, BlockSequence(self.post_blocks), use_checkpointing=use_ckpt) + x = apply_output_head(self.output_proj, x, zq) + return rearrange(x, "(b t) c h w -> b c t h w", b=b, t=t) if self.dims == 2 else x + + def forward( + self, + z: torch.Tensor, + *, + entry: Literal["x4", "x2"] = "x4", + return_intermediates: bool | None = None, + ) -> dict[str, torch.Tensor] | torch.Tensor: + """Upsample through the x4 cascade or the configured x2 entry.""" + if entry not in ("x4", "x2"): + msg = f"entry must be 'x4' or 'x2', got {entry!r}" + raise ValueError(msg) + if entry == "x2" and not self.enable_x2_entry: + msg = "forward(entry='x2') requires the model to be built with enable_x2_entry=True" + raise ValueError(msg) + if entry == "x2" and return_intermediates: + msg = "return_intermediates=True is not supported for entry='x2' — the x2 path has no 2x sub-target" + raise ValueError(msg) + if return_intermediates is None: + return_intermediates = self.training + + b, _, t, _, _ = z.shape + x = rearrange(z, "b c t h w -> (b t) c h w") if self.dims == 2 else z + zq = x if self.zq_dim is not None else None + use_ckpt = self.training and self.gradient_checkpointing + + if entry == "x2": + if self.x2_branch is None: + msg = "x2 branch requires enable_x2_entry=True" + raise ValueError(msg) + x = self._head_x2(x, zq, use_ckpt=use_ckpt) + return self.x2_branch.forward_second_stage( + x, + zq, + b, + t, + self.shared_second_stage(), + use_checkpointing=use_ckpt, + ) + x = self._head_x4(x, zq, use_ckpt=use_ckpt) + out = self._tail(x, zq, b, t, use_ckpt=use_ckpt, return_intermediates=return_intermediates) + if not self.global_skip: + return out + return self._add_nearest_base(out, z) + + def _add_nearest_base( + self, + out: dict[str, torch.Tensor] | torch.Tensor, + z: torch.Tensor, + ) -> dict[str, torch.Tensor] | torch.Tensor: + """Add the parameter-free nearest base of the input latent to each output.""" + if isinstance(out, dict): + return { + "2x": out["2x"] + spatial_nearest_2x(z, DIMS_3, STAGE_FACTOR), + "4x": out["4x"] + spatial_nearest_2x(z, DIMS_3, self.upscale_factor), + } + return out + spatial_nearest_2x(z, DIMS_3, self.upscale_factor) + + +"""Configuration models for the latent upscaler.""" + + +UpsampleMode = Literal["pixel_shuffle", "bilinear", "pxs_v2", "pxs_v2_hybrid"] +UpsamplePosition = Literal["before_projection", "after_projection"] +X2TailMode = Literal["shared", "scale_specific_norm", "private", "private_full"] +X2FinisherMode = Literal["none", "linear", "pxs_residual"] +MotionAttentionBackend = Literal["sdpa", "natten"] +NattenBackend = Literal["cutlass-fna", "hopper-fna", "flex-fna"] +SPATIOTEMPORAL_DIMS = 3 +MIN_ROPE_HEAD_DIM = 6 + + +class _BaseModelConfig(BaseModel): + """Common fields shared by all upsampler architectures.""" + + model_config = pydantic.ConfigDict(extra="forbid") + + in_channels: int = 16 + upscale_factor: int = 2 + upsample_mode: UpsampleMode = "pixel_shuffle" + upsample_position: UpsamplePosition = "after_projection" + temporal_mix: bool = True + icnr: bool = False + dims: int = 2 + input_skip: bool = True + gradient_checkpointing: bool = False + modulated_norm: bool = False + temporal_padding: TemporalPadding = "zeros" # "replicate" pads T both sides, "causal" pads the past only (K-VAE) + upsample_padding_mode: UpsamplePaddingMode = "reflect" # "zeros" matches the production K-VAE upsample + + +class FlatModelConfig(_BaseModelConfig): + """Configuration for ConvLatentUpsampler (flat architecture).""" + + architecture: Literal["flat"] = "flat" + hidden_channels: int = 64 + stem_channels: int | None = None + num_residual_blocks: int = 16 + num_pre_residual_blocks: int = 0 + bottleneck_channels: int | None = None + expand_ratio: int = 4 + kernel_size: int = 3 + layer_scale_init: float | None = None + grn: bool = False + stochastic_depth_rate: float = 0.0 + depthwise: bool = False + + +class MotionAttentionConfig(BaseModel): + """Configuration for cross-frame spatial correspondence blocks.""" + + model_config = pydantic.ConfigDict(extra="forbid") + + spatial_kernel_size: int = Field(default=13, ge=3) + temporal_offsets: tuple[int, ...] = (-3, -2, -1, 1, 2, 3) + after_mid_blocks: tuple[int, ...] = (4, 8, 12) + num_heads: int = Field(default=3, gt=0) + backend: MotionAttentionBackend = "natten" + natten_backend: NattenBackend = "hopper-fna" + merge_compile: bool = False + + @pydantic.model_validator(mode="after") + def validate_attention_geometry(self) -> MotionAttentionConfig: + """Require centered spatial windows and symmetric temporal offsets.""" + if self.spatial_kernel_size % 2 == 0: + msg = "motion spatial_kernel_size must be odd" + raise ValueError(msg) + if not self.temporal_offsets or 0 in self.temporal_offsets: + msg = "motion temporal_offsets must be non-empty and exclude 0" + raise ValueError(msg) + if tuple(sorted(set(self.temporal_offsets))) != self.temporal_offsets: + msg = "motion temporal_offsets must be unique and sorted" + raise ValueError(msg) + if set(self.temporal_offsets) != {-offset for offset in self.temporal_offsets}: + msg = "motion temporal_offsets must be symmetric" + raise ValueError(msg) + if ( + not self.after_mid_blocks + or any(block_number <= 0 for block_number in self.after_mid_blocks) + or tuple(sorted(set(self.after_mid_blocks))) != self.after_mid_blocks + ): + msg = "motion after_mid_blocks must contain positive, unique, sorted block numbers" + raise ValueError(msg) + return self + + +class MultiScaleModelConfig(_BaseModelConfig): + """Configuration for MultiScaleUpsampler (cascaded 2x+2x architecture). + + Decomposes a 4x upscale into two supervised 2x stages: coarse structure + at 2x, fine detail at 4x. Each stage gets direct loss supervision. + + Attributes: + hidden_channels: Internal width of residual blocks. + num_pre_blocks: Residual blocks at 1x (before first upsample). + num_mid_blocks: Residual blocks at 2x (between upsamples). + num_post_blocks: Residual blocks at 4x (after second upsample). + bottleneck_channels: Mid-channel width for ResidualBlocks (None = hidden * expand_ratio). + expand_ratio: Inverted bottleneck expansion when bottleneck_channels is None. + kernel_size: Kernel size for depthwise convolutions. + layer_scale_init: Initial LayerScale gamma. None disables LayerScale. + grn: Insert Global Response Normalization in residual blocks. + stochastic_depth_rate: Max drop rate, linearly scheduled across all stages. + depthwise: Use depthwise separable convolutions. + motion_attention: Optional cross-frame spatial correspondence stack. + bare_stem: Reduce input projections to a single conv (no norm/activation). + modulated_output_proj: Use a zq-conditioned ModulatedRMSNorm in ``output_proj``. + global_skip: Add a nearest-upsampled input base to every output and + zero-init the output heads, so the model starts as a nearest upsampler + and learns only the residual correction. + loss_weight_2x: Training loss weight for 2x intermediate output. + loss_weight_4x: Training loss weight for 4x final output. + enable_x2_entry: Build the ``mid_input_proj`` submodule and allow + ``forward(entry="x2")`` — a second entry that injects an already-2x LQ + latent after ``mid_blocks`` for a single 2x upscale through the shared + second stage (LAY-462). Off by default so the x4 module set and its + checkpoints stay unchanged. + x2_adapter_blocks: Number of x2-only residual blocks before the second stage. + x2_tail_mode: Amount of scale-specific capacity in the second stage; + ``private_full`` additionally privatizes ``mid_blocks`` and runs them + in the x2 path ahead of the second-stage upsample. + x2_adapter_sources: pre_blocks indices warm-copied into the adapter (one + per block); None keeps the zero-init identity adapter. + x2_finisher: upsample_1 conv pair applied at the adapter grid without the + nearest-2x resize: "linear" clones the 1x1 only, "pxs_residual" adds + the spatial conv with its residual add. + stage_channels: Per-grid widths (1x, 2x, 4x) for the K-VAE-shaped pyramid; None keeps constant hidden_channels. + """ + + architecture: Literal["multi_scale"] = "multi_scale" + upscale_factor: int = 4 + hidden_channels: int = 64 + num_pre_blocks: int = 2 + num_mid_blocks: int = 6 + num_post_blocks: int = 8 + bottleneck_channels: int | None = None + expand_ratio: int = 4 + kernel_size: int = 3 + layer_scale_init: float | None = None + grn: bool = False + stochastic_depth_rate: float = 0.0 + depthwise: bool = False + motion_attention: MotionAttentionConfig | None = None + bare_stem: bool = False + modulated_output_proj: bool = False + global_skip: bool = False + loss_weight_2x: float = 0.5 + loss_weight_4x: float = 1.0 + enable_x2_entry: bool = False + x2_adapter_blocks: int = Field(default=0, ge=0) + x2_tail_mode: X2TailMode = "shared" + x2_adapter_sources: tuple[int, ...] | None = None + x2_finisher: X2FinisherMode = "none" + stage_channels: tuple[int, int, int] | None = None + + @pydantic.model_validator(mode="after") + def validate_upscale_factor(self) -> MultiScaleModelConfig: + """Multi-scale architecture only supports 4x upscale (two 2x stages).""" + expected = 4 + if self.upscale_factor != expected: + msg = f"multi_scale architecture requires upscale_factor={expected}, got {self.upscale_factor}" + raise ValueError(msg) + return self + + @pydantic.model_validator(mode="after") + def validate_redesign_flags(self) -> MultiScaleModelConfig: + """Reject redesign-flag combinations the model cannot execute.""" + if self.modulated_output_proj and not self.modulated_norm: + msg = "modulated_output_proj requires modulated_norm=true" + raise ValueError(msg) + if self.enable_x2_entry and self.global_skip: + msg = "global_skip is x4-only and incompatible with enable_x2_entry" + raise ValueError(msg) + if self.enable_x2_entry and self.modulated_output_proj and self.x2_tail_mode != "private_full": + msg = "modulated_output_proj with enable_x2_entry requires x2_tail_mode='private_full'" + raise ValueError(msg) + return self + + @pydantic.model_validator(mode="after") + def validate_stage_channels(self) -> MultiScaleModelConfig: + """Pyramid widths must agree with hidden_channels and exclude width-coupled features.""" + if self.stage_channels is None: + return self + if self.hidden_channels != self.stage_channels[0]: + msg = f"hidden_channels ({self.hidden_channels}) must equal stage_channels[0] ({self.stage_channels[0]})" + raise ValueError(msg) + if self.enable_x2_entry and self.x2_tail_mode != "private_full": + msg = "stage_channels with enable_x2_entry requires x2_tail_mode='private_full'" + raise ValueError(msg) + if self.motion_attention is not None: + msg = "stage_channels is incompatible with motion_attention (motion blocks are built at hidden_channels)" + raise ValueError(msg) + if self.input_skip: + msg = "stage_channels is incompatible with input_skip (channel-repeat skip assumes one width)" + raise ValueError(msg) + if self.bottleneck_channels is not None: + msg = "stage_channels derives per-stage mid widths from expand_ratio; bottleneck_channels is incompatible" + raise ValueError(msg) + if self.num_mid_blocks < 1 or self.num_post_blocks < 1: + msg = "stage_channels needs at least one mid and one post block to host the width transitions" + raise ValueError(msg) + return self + + @pydantic.model_validator(mode="after") + def validate_x2_adaptation(self) -> MultiScaleModelConfig: + """Reject x2 adaptation settings that cannot be executed.""" + has_adaptation = ( + self.x2_adapter_blocks > 0 + or self.x2_tail_mode != "shared" + or self.x2_adapter_sources is not None + or self.x2_finisher != "none" + ) + if has_adaptation and not self.enable_x2_entry: + msg = "x2 adaptation requires enable_x2_entry=true" + raise ValueError(msg) + if self.x2_tail_mode != "shared" and self.x2_adapter_blocks == 0: + msg = f"x2_tail_mode={self.x2_tail_mode!r} requires x2_adapter_blocks > 0" + raise ValueError(msg) + if self.x2_tail_mode == "scale_specific_norm" and not self.modulated_norm: + msg = "scale_specific_norm requires modulated_norm=true" + raise ValueError(msg) + if self.x2_tail_mode == "scale_specific_norm" and self.depthwise: + msg = "scale_specific_norm does not support depthwise residual blocks" + raise ValueError(msg) + if self.x2_adapter_sources is not None and len(self.x2_adapter_sources) != self.x2_adapter_blocks: + msg = ( + f"x2_adapter_sources length ({len(self.x2_adapter_sources)}) must equal " + f"x2_adapter_blocks ({self.x2_adapter_blocks})" + ) + raise ValueError(msg) + if self.x2_adapter_sources is not None and any( + not 0 <= index < self.num_pre_blocks for index in self.x2_adapter_sources + ): + msg = f"x2_adapter_sources indices must be within [0, num_pre_blocks={self.num_pre_blocks})" + raise ValueError(msg) + if self.x2_finisher != "none" and self.upsample_mode != "pxs_v2": + msg = "x2_finisher clones the pxs_v2 upsample convs and requires upsample_mode='pxs_v2'" + raise ValueError(msg) + return self + + @pydantic.model_validator(mode="after") + def validate_motion_attention(self) -> MultiScaleModelConfig: + """Reject motion-attention configurations incompatible with this model.""" + motion = self.motion_attention + if motion is None: + return self + if self.dims != SPATIOTEMPORAL_DIMS: + msg = "motion_attention requires dims=3" + raise ValueError(msg) + if self.enable_x2_entry: + msg = "motion_attention Round 23 is x4-only and cannot enable the x2 entry" + raise ValueError(msg) + if motion.after_mid_blocks[-1] > self.num_mid_blocks: + msg = "motion after_mid_blocks must not exceed num_mid_blocks" + raise ValueError(msg) + if self.hidden_channels % motion.num_heads != 0: + msg = "hidden_channels must be divisible by motion num_heads" + raise ValueError(msg) + head_dim = self.hidden_channels // motion.num_heads + if head_dim < MIN_ROPE_HEAD_DIM or head_dim % 2 != 0: + msg = "motion attention head_dim must be even and at least 6 for axial 3D RoPE" + raise ValueError(msg) + return self + + +ModelConfig = Annotated[ + FlatModelConfig | MultiScaleModelConfig, + Field(discriminator="architecture"), +] +"""Factory functions for building upsampler models.""" + + +if TYPE_CHECKING: + from torch import nn + + +def build_upsampler(config: FlatModelConfig | MultiScaleModelConfig) -> nn.Module: + """Create an upsampler module from config. + + Dispatches on config type: ``FlatModelConfig`` builds a + ``ConvLatentUpsampler``; ``MultiScaleModelConfig`` builds a + ``MultiScaleUpsampler``. + + Args: + config: Model configuration (flat or multi-scale). + + Returns: + Initialized upsampler module. + """ + if isinstance(config, MultiScaleModelConfig): + return _build_multi_scale(config) + return _build_flat(config) + + +def _build_flat(config: FlatModelConfig) -> ConvLatentUpsampler: + """Build a flat (single-stage) ConvLatentUpsampler from config.""" + return ConvLatentUpsampler( + in_channels=config.in_channels, + hidden_channels=config.hidden_channels, + bottleneck_channels=config.bottleneck_channels, + num_pre_residual_blocks=config.num_pre_residual_blocks, + num_residual_blocks=config.num_residual_blocks, + upscale_factor=config.upscale_factor, + input_skip=config.input_skip, + upsample_mode=config.upsample_mode, + upsample_position=config.upsample_position, + temporal_mix=config.temporal_mix, + icnr=config.icnr, + gradient_checkpointing=config.gradient_checkpointing, + dims=config.dims, + expand_ratio=config.expand_ratio, + kernel_size=config.kernel_size, + layer_scale_init=config.layer_scale_init, + grn=config.grn, + stochastic_depth_rate=config.stochastic_depth_rate, + depthwise=config.depthwise, + stem_channels=config.stem_channels, + zq_dim=config.in_channels if config.modulated_norm else None, + temporal_padding=config.temporal_padding, + upsample_padding_mode=config.upsample_padding_mode, + ) + + +def _build_multi_scale(config: MultiScaleModelConfig) -> MultiScaleUpsampler: + """Build a cascaded 2x+2x MultiScaleUpsampler from config.""" + motion_attention = build_motion_attention_spec(config) + return MultiScaleUpsampler( + in_channels=config.in_channels, + hidden_channels=config.hidden_channels, + bottleneck_channels=config.bottleneck_channels, + num_pre_blocks=config.num_pre_blocks, + num_mid_blocks=config.num_mid_blocks, + num_post_blocks=config.num_post_blocks, + input_skip=config.input_skip, + upsample_mode=config.upsample_mode, + temporal_mix=config.temporal_mix, + icnr=config.icnr, + gradient_checkpointing=config.gradient_checkpointing, + dims=config.dims, + expand_ratio=config.expand_ratio, + kernel_size=config.kernel_size, + layer_scale_init=config.layer_scale_init, + grn=config.grn, + stochastic_depth_rate=config.stochastic_depth_rate, + depthwise=config.depthwise, + motion_attention=motion_attention, + zq_dim=config.in_channels if config.modulated_norm else None, + bare_stem=config.bare_stem, + modulated_output_proj=config.modulated_output_proj, + global_skip=config.global_skip, + enable_x2_entry=config.enable_x2_entry, + x2_adapter_blocks=config.x2_adapter_blocks, + x2_tail_mode=config.x2_tail_mode, + x2_adapter_sources=config.x2_adapter_sources, + x2_finisher=config.x2_finisher, + stage_channels=config.stage_channels, + temporal_padding=config.temporal_padding, + upsample_padding_mode=config.upsample_padding_mode, + ) + + +def build_motion_attention_spec(config: MultiScaleModelConfig) -> MotionAttentionSpec | None: + """Resolve the nested motion-attention config into a construction spec.""" + motion = config.motion_attention + if motion is None: + return None + return MotionAttentionSpec( + after_mid_blocks=motion.after_mid_blocks, + block=MotionCorrespondenceSpec( + channels=config.hidden_channels, + spatial_kernel_size=motion.spatial_kernel_size, + temporal_offsets=motion.temporal_offsets, + num_heads=motion.num_heads, + head_dim=config.hidden_channels // motion.num_heads, + backend=motion.backend, + natten_backend=motion.natten_backend, + merge_compile=motion.merge_compile, + ), + ) + + +_MODEL_CONFIG_ADAPTER = TypeAdapter(ModelConfig) + + +class Kandinsky6SRLatentUpscalerBank(ModelMixin, ConfigMixin): + """Diffusers wrapper around the self-contained x2/x4 upscaler bank.""" + + @register_to_config + def __init__( + self, + models: Sequence[Mapping[str, Any]], + scaling_factor: float = 1.0, + ) -> None: + super().__init__() + self._models = nn.ModuleDict() + for spec in models: + target_scale = str(spec["target_scale"]) + if not target_scale.endswith("x"): + target_scale = f"{target_scale}x" + if target_scale in self._models: + raise ValueError(f"duplicate latent upscaler target scale: {target_scale}") + model_config = _MODEL_CONFIG_ADAPTER.validate_python(spec["model"]) + upscaler = build_upsampler(model_config) + upscaler.target_scale = target_scale + upscaler.scaling_factor = float(scaling_factor) + self._models[target_scale] = upscaler + + def for_scale(self, scale: int) -> nn.Module | None: + """Return the bank entry matching ``scale``.""" + key = f"{int(scale)}x" + return self._models[key] if key in self._models else None + + @property + def scales(self) -> tuple[int, ...]: + """Return all serialized upscale factors.""" + return tuple(sorted(int(key.removesuffix("x")) for key in self._models)) + + def forward( + self, + latent: Tensor, + *, + entry: str | None = None, + return_intermediates: bool | None = None, + ) -> Any: + """Forward one bank entry, primarily for direct component use.""" + if entry is None: + if len(self._models) != 1: + raise ValueError("entry is required when the latent-upscaler bank has multiple scales") + entry = next(iter(self._models)) + key = entry.removeprefix("x") + if not key.endswith("x"): + key = f"{key}x" + model = self._models[key] + kwargs: dict[str, Any] = {} + if return_intermediates is not None: + kwargs["return_intermediates"] = return_intermediates + if hasattr(model, "forward"): + kwargs["entry"] = f"x{key.removesuffix('x')}" + try: + return model(latent, **kwargs) + except TypeError: + kwargs.pop("entry", None) + return model(latent, **kwargs) + + +__all__ = ["Kandinsky6SRLatentUpscalerBank"] diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py new file mode 100644 index 000000000000..b3974b44cfa3 --- /dev/null +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py @@ -0,0 +1,3129 @@ +"""Kandinsky 6 super-resolution Diffusers pipeline.""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from fractions import Fraction +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Literal, NamedTuple + +import av +import numpy as np +import torch +from ..pipeline_utils import DiffusionPipeline +from pydantic import BaseModel, ConfigDict, Field +from ...schedulers.scheduling_piflow import DXPolicy, policy_rollout_fm, shift_timesteps +from torch import Tensor, nn +from torch.distributed import all_gather +from torch.nn import functional + +from .pipeline_output import Kandinsky6SRPipelineOutput + +# These values are the SR data contract and must stay local to this pipeline. +VAE_SPATIAL_FACTOR = 16 +VAE_TEMPORAL_FACTOR = 4 +RESOLUTIONS: dict[int, list[tuple[int, int]]] = {512: [(512, 512), (512, 768), (768, 512)]} +TARGET_FPS = 24 +MAX_NUM_FRAMES = 121 +RESAMPLE_FPS_TOLERANCE = 1.5 +LATENT_NDIM = 4 +X0_T_CLAMP = 1e-5 +SUPPORTED_TARGET_SCALES = ("2x", "4x") +ResolutionScale = Literal[2, 4] +TotalResolutionScale = Literal[2, 4, 2.25] +constants = SimpleNamespace(VAE_SPATIAL_FACTOR=VAE_SPATIAL_FACTOR) +TileGridMode = Literal["legacy", "even"] +_sr_constants = None + + +def _call_module_method(module: nn.Module, method: str, *args: Any, **kwargs: Any) -> Any: + """Call a non-forward method through a Diffusers offload hook.""" + hook = getattr(module, "_hf_hook", None) + if hook is None: + return getattr(module, method)(*args, **kwargs) + args, kwargs = hook.pre_forward(module, *args, **kwargs) + output = getattr(module, method)(*args, **kwargs) + return hook.post_forward(module, output) + + +@contextmanager +def _execution_device_context(device: torch.device | str): + """Scope SR model execution to Diffusers' resolved CUDA device.""" + resolved_device = torch.device(device) + if resolved_device.type == "cuda": + with torch.cuda.device(resolved_device): + yield + else: + yield + + +"""Shared tiling utilities for spatial video tiling with Hanning-window blending.""" + + +class TileGrid(NamedTuple): + """Spatial tile grid parameters. + + ``tops`` and ``lefts`` are the authoritative per-axis tile start positions. + They cover ``[0, length - tile]`` with the nominal stride; the **last** + position is clamped to ``length - tile`` if a uniform stride would + overshoot. This means consecutive gaps in ``tops`` / ``lefts`` may be + shorter than the nominal stride near the right/bottom edge. + + ``n_h``, ``n_w``, ``total_tiles``, ``stride_h``, ``stride_w`` are derived + properties exposed for backward compatibility with callers that only + read them. To compute a tile position by index, prefer + ``grid.tops[row]`` / ``grid.lefts[col]`` over ``row * grid.stride_h`` — + the latter is wrong for the last clamped tile. + """ + + tile_h: int + tile_w: int + tops: tuple[int, ...] + lefts: tuple[int, ...] + + @property + def n_h(self) -> int: + """Number of tile rows.""" + return len(self.tops) + + @property + def n_w(self) -> int: + """Number of tile columns.""" + return len(self.lefts) + + @property + def total_tiles(self) -> int: + """Total number of tiles (``n_h * n_w``).""" + return self.n_h * self.n_w + + @property + def stride_h(self) -> int: + """Nominal vertical stride (gap between first two rows; ``tile_h`` if only one row).""" + return self.tops[1] - self.tops[0] if len(self.tops) > 1 else self.tile_h + + @property + def stride_w(self) -> int: + """Nominal horizontal stride (gap between first two cols; ``tile_w`` if only one col).""" + return self.lefts[1] - self.lefts[0] if len(self.lefts) > 1 else self.tile_w + + +def axis_positions(length: int, tile: int, stride: int) -> tuple[int, ...]: + """Return tile start positions covering ``[0, length - tile]``. + + Walks at the given ``stride`` from 0 and appends positions while the tile + still fits inside ``length``. If the last walked position is not exactly + ``length - tile``, appends one final position clamped to the right edge. + Single-position result if ``tile >= length``. + + Args: + length: Axis length in pixels. + tile: Tile size on this axis. + stride: Nominal stride between consecutive tiles. + + Returns: + Strictly increasing tuple of start positions. + """ + if tile <= 0 or stride <= 0: + msg = f"tile and stride must be positive (got tile={tile}, stride={stride})" + raise ValueError(msg) + if tile >= length: + return (0,) + + positions: list[int] = [] + p = 0 + while p + tile <= length: + positions.append(p) + p += stride + last = length - tile + if positions[-1] != last: + positions.append(last) + return tuple(positions) + + +def compute_tile_grid( + h: int, + w: int, + resolution_scale: int, + overlap: float = 0.5, + tile_hw: tuple[int, int] | None = None, +) -> TileGrid: + """Compute tile geometry with configurable overlap. + + By default, tile size is ``(h // resolution_scale, w // resolution_scale)``. + Pass ``tile_hw`` to override with explicit tile dimensions (e.g. derived + from a fixed base resolution rather than the source frame size). + + The grid covers the full frame: if the nominal stride does not divide + evenly into ``(h - tile_h)`` / ``(w - tile_w)``, the last tile in each + axis is clamped to the right/bottom edge. ``stitch_tiles_hanning`` + handles the resulting non-uniform overlap correctly via Hanning-window + normalisation. + + Args: + h: Video height in pixels. + w: Video width in pixels. + resolution_scale: Divisor for tile size when ``tile_hw`` is ``None``. + overlap: Fraction of tile overlap in ``[0, 1)``. Default ``0.5`` (50%). + tile_hw: Explicit ``(tile_h, tile_w)`` override. When set, + ``resolution_scale`` is ignored for tile sizing. + + Returns: + ``TileGrid`` with tile sizes and per-axis tile start positions. + """ + if tile_hw is None: + tile_h = h // resolution_scale + tile_w = w // resolution_scale + else: + tile_h, tile_w = tile_hw + stride_h = max(1, int(tile_h * (1.0 - overlap))) + stride_w = max(1, int(tile_w * (1.0 - overlap))) + + tops = axis_positions(h, tile_h, stride_h) + lefts = axis_positions(w, tile_w, stride_w) + return TileGrid(tile_h=tile_h, tile_w=tile_w, tops=tops, lefts=lefts) + + +def extract_all_tiles( + video: torch.Tensor, + grid: TileGrid, +) -> list[torch.Tensor]: + """Extract all spatial tiles from a video tensor. + + Args: + video: ``[T, C, H, W]`` tensor. + grid: Tile grid parameters from ``compute_tile_grid``. + + Returns: + List of ``[T, C, tile_h, tile_w]`` tensors in row-major order + (``tops`` x ``lefts``). + """ + tiles: list[torch.Tensor] = [] + for top in grid.tops: + for left in grid.lefts: + tile = video[:, :, top : top + grid.tile_h, left : left + grid.tile_w] + tiles.append(tile) + return tiles + + +def hanning_window_2d(h: int, w: int, device: torch.device) -> torch.Tensor: + """Create a 2D Hanning window with non-zero endpoints. + + Uses ``hann_window(n + 2)[1:-1]`` to avoid exact zeros at boundaries, + ensuring non-zero weight where only one tile contributes. + + Args: + h: Window height. + w: Window width. + device: Target device. + + Returns: + ``[h, w]`` float tensor with values in ``(0, 1]``. + """ + wy = torch.hann_window(h + 2, device=device)[1:-1] + wx = torch.hann_window(w + 2, device=device)[1:-1] + return wy[:, None] * wx[None, :] + + +def stitch_tiles_hanning( + tiles: list[torch.Tensor], + grid: TileGrid, + original_h: int, + original_w: int, + scale: int = 1, +) -> torch.Tensor: + """Stitch tiles into a full frame using Hanning-window weighted blending. + + Each tile is multiplied by a 2D Hanning window and accumulated into + the output canvas. The final result is normalised by the accumulated + weights so that overlapping regions blend smoothly. The same scheme + works for non-uniform overlap at the right/bottom edge: in the clamped + region both ``pred_acc`` and ``weight_acc`` receive more contributions, + and the per-pixel division cancels it out. + + When ``scale > 1``, tiles are assumed to be at HR resolution + (i.e. each tile covers ``tile_h * scale x tile_w * scale`` pixels) + and the output canvas is ``original_h * scale x original_w * scale``. + Grid positions are scaled accordingly. + + Args: + tiles: List of ``[C, T, th, tw]`` float tensors in row-major order + (same order as ``extract_all_tiles``). When ``scale == 1``, + ``th == grid.tile_h``; when ``scale > 1``, ``th == grid.tile_h * scale``. + grid: Tile grid parameters (at LQ / original resolution). + original_h: LQ frame height. + original_w: LQ frame width. + scale: Upscale factor. Output resolution is + ``(original_h * scale, original_w * scale)``. + + Returns: + ``[C, T, original_h * scale, original_w * scale]`` float tensor. + """ + first = tiles[0] + c, t = first.shape[0], first.shape[1] + hr_tile_h, hr_tile_w = first.shape[2], first.shape[3] + device = first.device + + window = hanning_window_2d(hr_tile_h, hr_tile_w, device) + window = window.unsqueeze(0).unsqueeze(0) # [1, 1, hr_tile_h, hr_tile_w] + + out_h = original_h * scale + out_w = original_w * scale + pred_acc = torch.zeros(c, t, out_h, out_w, device=device) + weight_acc = torch.zeros(1, 1, out_h, out_w, device=device) + + tile_idx = 0 + for top in grid.tops: + for left in grid.lefts: + y = top * scale + x = left * scale + tile = tiles[tile_idx] + pred_acc[:, :, y : y + hr_tile_h, x : x + hr_tile_w] += tile * window + weight_acc[:, :, y : y + hr_tile_h, x : x + hr_tile_w] += window + tile_idx += 1 + + return pred_acc / weight_acc.clamp(min=1e-6) + + +def axis_positions_even(length: int, tile: int, min_overlap: float, snap: int) -> tuple[int, ...]: + """Return evenly distributed, ``snap``-aligned tile start positions. + + Uses the minimal position count whose uniform stride keeps the tile + overlap at or above ``min_overlap``, then spreads the positions evenly + over ``[0, length - tile]`` in integer ``snap`` units. When ``length`` or + ``tile`` is not ``snap``-aligned the same layout is computed at pixel + precision (``snap=1``) — the latent-grid guard downstream still enforces + alignment where it actually matters (the LU path). + + Args: + length: Axis length in pixels. + tile: Tile size on this axis. + min_overlap: Overlap floor as a fraction of ``tile`` in ``[0, 1)``. + snap: Position alignment unit (the VAE spatial factor). + + Returns: + Strictly increasing positions covering ``[0, length - tile]``. + + Raises: + ValueError: On non-positive ``tile``/``snap`` or ``min_overlap`` + outside ``[0, 1)``. + """ + if tile <= 0 or snap <= 0 or not 0 <= min_overlap < 1: + msg = f"invalid axis spec: tile={tile}, snap={snap}, min_overlap={min_overlap}" + raise ValueError(msg) + if tile >= length: + return (0,) + unit = snap if length % snap == 0 and tile % snap == 0 else 1 + span_units = (length - tile) // unit + max_stride_units = max(1, math.floor(tile * (1.0 - min_overlap) / unit)) + count = math.ceil(span_units / max_stride_units) + 1 + return tuple(round(i * span_units / (count - 1)) * unit for i in range(count)) + + +def compute_tile_grid_even( + h: int, + w: int, + tile_hw: tuple[int, int], + min_overlap: float, + snap: int, +) -> TileGrid: + """Build a :class:`TileGrid` with the even per-axis layout. + + Drop-in for ``compute_tile_grid(..., tile_hw=...)``: same ``TileGrid`` + contract (``tops``/``lefts`` are authoritative), only the position layout + differs — see :func:`axis_positions_even`. + + Args: + h: Video height in pixels. + w: Video width in pixels. + tile_hw: Explicit ``(tile_h, tile_w)``. + min_overlap: Overlap floor as a fraction of the tile size. + snap: Position alignment unit (the VAE spatial factor). + + Returns: + ``TileGrid`` with evenly distributed, aligned tile positions. + """ + tile_h, tile_w = tile_hw + return TileGrid( + tile_h=tile_h, + tile_w=tile_w, + tops=axis_positions_even(h, tile_h, min_overlap, snap), + lefts=axis_positions_even(w, tile_w, min_overlap, snap), + ) + + +def resolve_scale_request(scale: float) -> tuple[int, float]: + """Map a requested ``--resolution-scale`` to ``(tiling_scale, pre_upscale)``. + + Args: + scale: The requested total upscale — ``2``, ``4`` or ``2.25``. + + Returns: + ``(tiling_scale, pre_upscale)``; ``pre_upscale`` is ``1.0`` for the + integer scales. + """ + requested = float(scale) + if requested == 2.25: # noqa: PLR2004 — the one supported fractional total + return 2, 1.125 + if requested in (2.0, 4.0): + return int(requested), 1.0 + raise ValueError("SR supports total scales 2, 4, and 2.25") + + +def pre_upscale_video(video: torch.Tensor, factor: float, spatial_multiple: int) -> torch.Tensor: + """Bilinear-upscale a ``[T, C, H, W]`` uint8 video by ``factor`` in pixel space. + + Target dims are rounded to the nearest multiple of ``spatial_multiple`` + (the VAE spatial factor) so the whole-video encode and the latent tile + grid stay integer-aligned (``latent_tile_grid_from_pixel_grid`` rejects + unaligned grids). For sources whose scaled dims already land on the + factor (512x768 x1.125 -> 576x864) the rounding is a no-op and the total + scale is exact. + + Args: + video: ``[T, C, H, W]`` uint8 source video. + factor: Pixel upscale factor (> 1). + spatial_multiple: VAE spatial factor to align the target dims to. + + Returns: + ``[T, C, H', W']`` uint8 video with ``H' ~= H * factor`` aligned. + """ + if video.ndim != 4: # noqa: PLR2004 + raise ValueError(f"video must have rank 4 [T,C,H,W], got {tuple(video.shape)}") + if factor <= 0 or spatial_multiple <= 0: + raise ValueError("factor and spatial_multiple must be positive") + height, width = video.shape[-2:] + target_h = max(spatial_multiple, round(height * factor / spatial_multiple) * spatial_multiple) + target_w = max(spatial_multiple, round(width * factor / spatial_multiple) * spatial_multiple) + resized = functional.interpolate( + video.float(), + size=(target_h, target_w), + mode="bilinear", + align_corners=False, + ) + return resized.round_().clamp_(0, 255).to(torch.uint8) + + +def get_visual_size(x: torch.Tensor) -> int: + """Return the resolution key matching a visual latent tensor's spatial dims. + + Scales the tensor's spatial dims by the VAE spatial factor and looks them + up in the ``RESOLUTIONS`` registry. Reads ``constants.VAE_SPATIAL_FACTOR`` + late-bound so ``set_vae_factors`` (16 for the KVAE) applies — a + ``from``-import would freeze the import-time default. + + Args: + x: Visual latent tensor of shape ``(T, H, W, C)``. + + Returns: + The matching resolution key. + + Raises: + ValueError: If tensor dimensions do not match any known resolution. + """ + actual_size = (x.shape[1] * constants.VAE_SPATIAL_FACTOR, x.shape[2] * constants.VAE_SPATIAL_FACTOR) + for key, value in RESOLUTIONS.items(): + if actual_size in value: + return key + valid_sizes = {size for sizes in RESOLUTIONS.values() for size in sizes} + msg = ( + f"Visual tensor spatial dimensions {actual_size} do not match any known resolution. " + f"Tensor shape: {x.shape}. Valid resolutions: {valid_sizes}" + ) + raise ValueError(msg) + + +def degrade_lq_latent( + lq_latent: torch.Tensor, + noise_scale: float = 0.7, + noise_type: Literal["linear", "ddpm"] = "linear", + generator: torch.Generator | None = None, +) -> torch.Tensor: + """Mix random Gaussian noise into an LQ latent for data augmentation. + + Args: + lq_latent: LQ latent tensor of arbitrary shape. + noise_scale: Noise fraction ``s`` in ``(0, 1)``. When ``0``, the tensor + is returned unchanged. + noise_type: ``"linear"`` for ``(1-s)*lq + s*eps`` or ``"ddpm"`` for + ``sqrt(1-s²)*lq + s*eps`` (variance-preserving). + generator: Optional RNG generator for deterministic noise sampling. + + Returns: + Degraded LQ latent with the same shape and dtype as the input. + """ + if noise_scale <= 0: + return lq_latent + eps = torch.randn(lq_latent.shape, device=lq_latent.device, dtype=lq_latent.dtype, generator=generator) + if noise_type == "ddpm": + return (1 - noise_scale**2) ** 0.5 * lq_latent + noise_scale * eps + return (1 - noise_scale) * lq_latent + noise_scale * eps + + +def get_sparse_params( + model: Any, + visual: torch.Tensor, + visual_cu_seqlens: torch.Tensor, +) -> dict[str, Any] | None: + """Build sparse-attention parameters for the given visual tensor and model. + + Supports ``"nabla"`` and ``"nabla_framewise_causal"`` attention types; + returns ``None`` for dense attention. The target ``P`` is used directly (the + training-time linear warmup schedule does not apply at inference). + + Args: + model: Model with ``patch_size`` and ``attention_params`` attributes. + visual: Visual latent tensor of shape ``(T, H, W, C)``. + visual_cu_seqlens: Cumulative frame counts for each sample in the batch. + + Returns: + A dict of sparse attention parameters, or ``None`` for dense attention. + + Raises: + ValueError: If ``model.patch_size[0]`` is not 1. + """ + if model.patch_size[0] != 1: + msg = f"Expected model.patch_size[0] == 1, got {model.patch_size[0]}" + raise ValueError(msg) + t, h, w, _ = visual.shape + t, h, w = ( + t // model.patch_size[0], + h // model.patch_size[1], + w // model.patch_size[2], + ) + visual_size = get_visual_size(visual) + attention_configs = model.attention_params + try: + attention_params = attention_configs[visual_size] + except KeyError: + # JSON object keys are always strings, while the native YAML config + # uses integer resolution keys. + attention_params = attention_configs[str(visual_size)] + + def config_value(name: str, default: Any = None) -> Any: + if isinstance(attention_params, Mapping): + return attention_params.get(name, default) + return getattr(attention_params, name, default) + + if config_value("type") == "nabla": + return { + "attention_type": config_value("type"), + "to_fractal": True, + "P": config_value("P"), + "wT": config_value("wT"), + "wW": config_value("wW"), + "wH": config_value("wH"), + "add_sta": config_value("add_sta"), + "visual_shape": (t, h, w), + "visual_seqlens": visual_cu_seqlens, + "method": config_value("method", "topcdf"), + } + if config_value("type") == "nabla_framewise_causal": + return { + "attention_type": config_value("type"), + "to_fractal": True, + "P": config_value("P"), + "wT": config_value("wT"), + "wW": config_value("wW"), + "wH": config_value("wH"), + "add_sta": config_value("add_sta"), + "mf": config_value("mf"), + "visual_shape": (t, h, w), + "visual_seqlens": visual_cu_seqlens, + } + return None + + +def replicate_cached_embeds( + cached: dict[str, torch.Tensor], + bs: int, + device: str | int, +) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + """Replicate bs=1 cached text embeddings for an arbitrary batch size. + + Args: + cached: Dict with ``text_embeds``, ``pooled_embed``, ``cu_seqlens`` + produced by ``save_empty_text_embeddings.py`` for a single sample. + bs: Target batch size. + device: Target device. + + Returns: + Tuple of (text_embeds dict, cu_seqlens tensor) for the full batch. + """ + single_text = cached["text_embeds"] + single_pooled = cached["pooled_embed"] + seq_len = int(cached["cu_seqlens"][-1].item()) + + text_embeds = single_text.repeat(bs, 1).to(device) + pooled_embed = single_pooled.repeat(bs, 1).to(device) + cu_seqlens = seq_len * torch.arange(bs + 1, dtype=torch.int32, device=device) + + return {"text_embeds": text_embeds, "pooled_embed": pooled_embed}, cu_seqlens + + +def kvae_weights_path(checkpoint_prefix: str) -> str: + """Return the kvae weights file for a sidecar prefix, preferring safetensors. + + The release format is ``{prefix}.safetensors`` (flat state dict), the + training format ``{prefix}.ckpt`` (pickled ``{"state_dict": ...}``); both + load through ``CachedCausalVAE.init_from_ckpt``, which dispatches on the + suffix. + + Args: + checkpoint_prefix: Local kvae sidecar prefix (``{prefix}.yaml`` sits + next to the weights). + + Returns: + Path to the weights file to load. + """ + safetensors_path = f"{checkpoint_prefix}.safetensors" + return safetensors_path if Path(safetensors_path).exists() else f"{checkpoint_prefix}.ckpt" + + +def module_dtype(module: torch.nn.Module) -> torch.dtype | None: + """Return the dtype used by the first parameterized layer, if any.""" + try: + return next(module.parameters()).dtype + except StopIteration: + return None + + +def cast_to_module_dtype(module: torch.nn.Module, value: torch.Tensor) -> torch.Tensor: + """Cast floating-point inputs to the module's parameter dtype.""" + dtype = module_dtype(module) + if dtype is not None and value.is_floating_point() and value.dtype != dtype: + return value.to(dtype=dtype) + return value + + +def encode_pixels_to_latent(vae: torch.nn.Module, pixels: torch.Tensor, *, sample: bool = True) -> torch.Tensor: + """Encode ``(B, C, T, H, W)`` pixels in ``[0, 255]`` with the KVAE. + + Normalizes with the KVAE's own convention, then returns the latent + (``(latent, split_list)[0]`` — the regularizer mode). + + Args: + vae: Causal video KVAE. + pixels: ``(B, C, T, H, W)`` tensor in ``[0, 255]``. + sample: Unused — the KVAE always returns the regularizer mode. + + Returns: + Latent ``(B, C, T', H', W')`` — not yet scaled by ``scaling_factor``. + """ + del sample # the KVAE always returns the regularizer mode + x = vae.normalize_data(pixels) + x = cast_to_module_dtype(vae, x) + result = vae.encode(x) + return result[0] + + +def decode_latent_to_uint8(vae: torch.nn.Module, latent: torch.Tensor) -> torch.Tensor: + """Decode a ``(B, C, T, H, W)`` latent to ``uint8`` pixels in ``[0, 255]``. + + The latent must already be unscaled (divided by ``scaling_factor``); + denormalization uses the KVAE's ``denormalize_data``. + + Args: + vae: Causal video KVAE. + latent: ``(B, C, T, H, W)`` unscaled latent. + + Returns: + ``(B, C, T_pixel, H, W)`` ``uint8`` pixels in ``[0, 255]``. + """ + latent = cast_to_module_dtype(vae, latent) + return denormalize_to_uint8(vae, vae.decode(latent).sample) + + +def denormalize_to_float(vae: torch.nn.Module, decoded: torch.Tensor) -> torch.Tensor: + """Convert a decoded tensor to float pixels in ``[0, 1]`` without quantization. + + Args: + vae: Causal video KVAE. + decoded: Decode output in the VAE's normalized pixel space. + + Returns: + Float pixels in ``[0, 1]`` with the same shape. + """ + return (vae.denormalize_data(decoded) / 255.0).clamp(0.0, 1.0) + + +def denormalize_to_uint8(vae: torch.nn.Module, decoded: torch.Tensor) -> torch.Tensor: + """Convert an already-decoded tensor in ``[-1, 1]`` to ``uint8`` ``[0, 255]``. + + Uses the KVAE ``denormalize_data`` convention ((x + 1) * 128). For decode + outputs that did not pass through :func:`decode_latent_to_uint8` (e.g. + debug-clip dumps of already-decoded pixels). + + Args: + vae: Causal video KVAE. + decoded: Decode output in ``[-1, 1]`` (any shape). + + Returns: + ``uint8`` pixels in ``[0, 255]`` with the same shape. + """ + return (denormalize_to_float(vae, decoded) * 255.0).to(torch.uint8) + + +def latent_upscaler_scale(upscaler: nn.Module) -> int: + """Return the loaded LU's fixed spatial upscale factor.""" + raw = str(getattr(upscaler, "target_scale", "4x")) + if raw not in SUPPORTED_TARGET_SCALES: + msg = f"Unexpected latent upscaler target_scale={raw!r}; expected one of {SUPPORTED_TARGET_SCALES}." + raise ValueError(msg) + return int(raw.removesuffix("x")) + + +def run_latent_upscaler(upscaler: nn.Module, z: torch.Tensor) -> torch.Tensor: + """Forward a scaled LQ latent through its configured LU entry.""" + target_scale = getattr(upscaler, "target_scale", None) + if target_scale is not None: + entry = f"x{latent_upscaler_scale(upscaler)}" + try: + return upscaler(z, entry=entry, return_intermediates=False) + except TypeError: + # Flat upscalers do not accept the multi-scale keyword arguments. + return upscaler(z) + return upscaler(z) + + +@torch.no_grad() +def get_model_prediction( + dit: torch.nn.Module, + x: torch.Tensor, + t: torch.Tensor, + text_embeds: dict[str, torch.Tensor], + null_text_embeds: dict[str, torch.Tensor], + visual_cu_seqlens: torch.Tensor, + text_cu_seqlens: torch.Tensor, + null_text_cu_seqlens: torch.Tensor, + visual_rope_pos: list[torch.Tensor], + text_rope_pos: torch.Tensor, + null_text_rope_pos: torch.Tensor, + scale_factor: tuple[float, ...], + guidance_weight: float, + sparse_params: Any | None = None, +) -> torch.Tensor: + """Compute classifier-free guidance model prediction. + + The returned tensor has the same semantics as the raw model output + (velocity when ``prediction_target="velocity"``, clean ``x_0`` when + ``prediction_target="x0"``). The caller is responsible for converting + the prediction to a velocity if needed for the Euler step. + + Args: + dit: DiT model used for inference. + x: Visual latent tensor to denoise. + t: Timestep tensor broadcast over the batch. + text_embeds: Dict with ``text_embeds`` and ``pooled_embed`` keys. + null_text_embeds: Same structure as ``text_embeds`` but for null captions. + visual_cu_seqlens: Cumulative sequence lengths for visual tokens. + text_cu_seqlens: Cumulative sequence lengths for text tokens. + null_text_cu_seqlens: Cumulative sequence lengths for null text tokens. + visual_rope_pos: RoPE position indices for visual tokens. + text_rope_pos: RoPE position indices for text tokens. + null_text_rope_pos: RoPE position indices for null text tokens. + scale_factor: Per-axis RoPE frequency scaling. + guidance_weight: CFG strength (1 = no guidance). + sparse_params: Optional sparse attention parameters. + + Returns: + CFG-combined model prediction (velocity or ``x_0``). + """ + model_time = (t * 1000).to(dtype=x.dtype) + motion_score = torch.full((1,), 900.0, device=x.device, dtype=x.dtype) + cond_pred = dit( + x, + text_embeds["text_embeds"], + text_embeds["pooled_embed"], + model_time, + visual_cu_seqlens, + text_cu_seqlens, + visual_rope_pos, + text_rope_pos, + scale_factor=scale_factor, + sparse_params=sparse_params, + motion_score=motion_score, + ) + if guidance_weight == 1.0: + return cond_pred + uncond_pred = dit( + x, + null_text_embeds["text_embeds"], + null_text_embeds["pooled_embed"], + model_time, + visual_cu_seqlens, + null_text_cu_seqlens, + visual_rope_pos, + null_text_rope_pos, + scale_factor=scale_factor, + sparse_params=sparse_params, + motion_score=motion_score, + ) + return uncond_pred + guidance_weight * (cond_pred - uncond_pred) + + +@torch.no_grad() +def generate( + img: torch.Tensor, + model: torch.nn.Module, + device: str | int, + num_steps: int, + text_embeds: dict[str, torch.Tensor], + null_text_embeds: dict[str, torch.Tensor], + visual_cu_seqlens: torch.Tensor, + text_cu_seqlens: torch.Tensor, + null_text_cu_seqlens: torch.Tensor, + visual_rope_pos: list[torch.Tensor], + text_rope_pos: torch.Tensor, + null_text_rope_pos: torch.Tensor, + scale_factor: tuple[float, ...], + guidance_weight: float, + scheduler_scale: float, + first_frames: torch.Tensor | None = None, + tp_mesh: Any | None = None, + visual_cond_scheme: str = "pretrain", + start_timestep: float = 1.0, + prediction_target: str = "velocity", + channelcat_drop_threshold: float = 0.0, + rfg_scale: float = 1.0, + progress_callback: Callable[[int], Any] | None = None, + progress_bar: Any = None, +) -> torch.Tensor: + """Run the full denoising loop from start_timestep to 0. + + When ``prediction_target="x0"`` the model predicts clean ``x_0`` instead + of velocity. The velocity for the Euler step is recovered as + ``v = (x_t - x_0_pred) / clamp(t, min=X0_T_CLAMP)`` to avoid numerical + instability as ``t → 0``. + + Args: + img: Initial latent tensor ``[T, H, W, C]`` (or wider for instruct). + model: DiT model. + device: Target CUDA device. + num_steps: Number of denoising steps. + text_embeds: Conditional text embeddings. + null_text_embeds: Unconditional text embeddings. + visual_cu_seqlens: Cumulative visual sequence lengths. + text_cu_seqlens: Cumulative text sequence lengths. + null_text_cu_seqlens: Cumulative null-text sequence lengths. + visual_rope_pos: RoPE position indices for visual tokens. + text_rope_pos: RoPE position indices for text tokens. + null_text_rope_pos: RoPE position indices for null text tokens. + scale_factor: Per-axis RoPE frequency scaling. + guidance_weight: CFG strength. + scheduler_scale: Timestep scheduler warp factor. + first_frames: Optional I2V conditioning first frames. + tp_mesh: Tensor-parallel device mesh (optional). + visual_cond_scheme: ``"pretrain"`` or ``"i2v"`` injection strategy. + start_timestep: Starting timestep (``1.0`` = full noise). + prediction_target: ``"velocity"`` or ``"x0"``. + channelcat_drop_threshold: When ``> 0``, zero out channel-cat LQ + channels (and mask) once the denoising timestep drops below this + value. ``0`` = disabled (default). + rfg_scale: Reference-Free Guidance strength on the anchor. Only + applied when ``model.instruct_type == "hybrid_anchor"`` and + ``rfg_scale != 1.0`` — in that case a second forward is run with + the anchor and anchor_mask channels zeroed, and predictions blend + as ``v_uncond + rfg_scale * (v_cond - v_uncond)``. ``1.0`` + (default) skips the second forward entirely. + Returns: + Denoised latent tensor (same leading dims as ``img``). + """ + img = img.to(device) + visual_cu_seqlens = visual_cu_seqlens.to(device) + text_cu_seqlens = text_cu_seqlens.to(device) + null_text_cu_seqlens = null_text_cu_seqlens.to(device) + visual_rope_pos = [position.to(device) for position in visual_rope_pos] + text_rope_pos = text_rope_pos.to(device) + null_text_rope_pos = null_text_rope_pos.to(device) + text_embeds = {key: value.to(device) for key, value in text_embeds.items()} + null_text_embeds = {key: value.to(device) for key, value in null_text_embeds.items()} + sparse_params = get_sparse_params(model, img, visual_cu_seqlens) + timesteps = torch.linspace(start_timestep, 0, num_steps, device=device) + timesteps = scheduler_scale * timesteps / (1 + (scheduler_scale - 1) * timesteps) + if tp_mesh: + tp_rank = tp_mesh["tp"].get_local_rank() + tp_world_size = tp_mesh["tp"].size() + img = torch.chunk(img, tp_world_size, dim=1)[tp_rank] + if first_frames is not None: + first_frames = torch.chunk(first_frames, tp_world_size, dim=1)[tp_rank] + if model.visual_cond and first_frames is not None: + first_frames = first_frames.to(device=img.device, dtype=img.dtype) + out_channels: int = 0 + lq_cond_channels: torch.Tensor | None = None + if channelcat_drop_threshold > 0 and img.shape[-1] > model.in_visual_dim: + lq_cond_channels = img[..., model.in_visual_dim :].clone() + for timestep, timestep_diff in list(zip(timesteps[:-1], torch.diff(timesteps), strict=False)): + time = timestep.unsqueeze(0).repeat(visual_cu_seqlens.shape[0] - 1) + if lq_cond_channels is not None: + if timestep.item() >= channelcat_drop_threshold: + img[..., model.in_visual_dim :] = lq_cond_channels + else: + img[..., model.in_visual_dim :] = 0 + if model.visual_cond and img.shape[-1] == model.in_visual_dim: + visual_cond = torch.zeros_like(img) + visual_cond_mask = torch.zeros([*img.shape[:-1], 1], dtype=img.dtype, device=img.device) + if first_frames is not None: + if visual_cond_scheme == "pretrain": + visual_cond[visual_cu_seqlens[:-1]] = first_frames + elif visual_cond_scheme == "i2v": + img[visual_cu_seqlens[:-1]] = first_frames + else: + msg = f"unknown visual_cond_scheme={visual_cond_scheme}" + raise ValueError(msg) + visual_cond_mask[visual_cu_seqlens[:-1]] = 1 + model_input = torch.cat([img, visual_cond, visual_cond_mask], dim=-1) + else: + model_input = img + v_cond = get_model_prediction( + model, + model_input, + time, + text_embeds, + null_text_embeds, + visual_cu_seqlens, + text_cu_seqlens, + null_text_cu_seqlens, + visual_rope_pos, + text_rope_pos, + null_text_rope_pos, + scale_factor, + guidance_weight, + sparse_params=sparse_params, + ) + if model.instruct_type == "hybrid_anchor" and rfg_scale != 1.0: + model_input_uncond = model_input.clone() + model_input_uncond[..., model.in_visual_dim :] = 0 + v_uncond = get_model_prediction( + model, + model_input_uncond, + time, + text_embeds, + null_text_embeds, + visual_cu_seqlens, + text_cu_seqlens, + null_text_cu_seqlens, + visual_rope_pos, + text_rope_pos, + null_text_rope_pos, + scale_factor, + guidance_weight, + sparse_params=sparse_params, + ) + model_pred = v_uncond + rfg_scale * (v_cond - v_uncond) + else: + model_pred = v_cond + out_channels = model_pred.shape[-1] + if prediction_target == "x0": + t_clamped = timestep.clamp(min=X0_T_CLAMP) + velocity = (img[..., :out_channels] - model_pred) / t_clamped + else: + velocity = model_pred + img[..., :out_channels] += timestep_diff * velocity + if progress_callback is not None: + progress_callback() + if progress_bar is not None: + progress_bar.update() + progress_bar.refresh() + if model.visual_cond and first_frames is not None and (visual_cond_scheme == "i2v"): + img[visual_cu_seqlens[:-1]] = first_frames + return img[..., :out_channels] + + +def _encode_lq_videos(lq_videos: list[torch.Tensor], vae: torch.nn.Module, device: str | int) -> torch.Tensor: + """VAE-encode a list of LQ videos into a stacked latent tensor. + + Args: + lq_videos: List of ``[T, H, W, 3]`` float32 tensors in ``[0, 255]``. + vae: Pre-trained VAE model (eval mode). + device: Target CUDA device. + + Returns: + Stacked latent tensor ``[sum(T'), H', W', C]`` scaled by + ``vae.config.scaling_factor``. + """ + latents: list[torch.Tensor] = [] + for lq in lq_videos: + lq_input = lq.permute(3, 0, 1, 2).unsqueeze(0).to(device=device, dtype=torch.bfloat16) + lq_latent = encode_pixels_to_latent(vae, lq_input) + lq_latent = lq_latent.squeeze(0).permute(1, 2, 3, 0).float() + lq_latent *= vae.config.scaling_factor + latents.append(lq_latent) + return torch.cat(latents, dim=0) + + +def _build_initial_latent( + dit: torch.nn.Module, + lq_latent: torch.Tensor, + bs: int, + duration: int, + height: int, + width: int, + device: str | int, + seed: int, + lq_noise_scale: float, + lq_noise_type: Literal["linear", "ddpm"], + lq_channel_noise_scale: float, + anchor_latent: torch.Tensor | None = None, + anchor_mask: torch.Tensor | None = None, + *, + anchor_free: bool = False, +) -> torch.Tensor: + """Build the initial latent tensor for the SR denoising loop. + + Constructs ``[starting | cond | mask]`` depending on ``instruct_type``: + + - ``"noise"``: degraded LQ as starting point, zero LQ cond + zero mask. + - ``"channel"``: random Gaussian noise as starting point, LQ cond + ones mask. + - ``"hybrid"``: degraded LQ as starting point, LQ cond + ones mask. + - ``"hybrid_anchor"``: degraded LQ as starting point, sparse HR anchor + + sparse anchor_mask. Requires ``anchor_latent`` and ``anchor_mask``. + + Args: + dit: DiT model (reads ``instruct_type`` and ``in_visual_dim``). + lq_latent: Scaled LQ latent ``[bs*T, H, W, C]``. + bs: Batch size. + duration: Number of temporal frames per sample. + height: Latent height. + width: Latent width. + device: Target CUDA device. + seed: RNG seed for noise generation. + lq_noise_scale: Noise fraction mixed into the LQ starting point. + lq_noise_type: ``"linear"`` or ``"ddpm"`` noising formula. + lq_channel_noise_scale: Noise fraction mixed into the channel-cat + conditioning (LQ for ``channel``/``hybrid``, anchor for + ``hybrid_anchor``). + anchor_latent: Pre-scaled sparse HR anchor ``[bs*T, H, W, C]`` (zeros + in slots without anchor). Required when + ``instruct_type == "hybrid_anchor"``. + anchor_mask: Sparse binary mask ``[bs*T, H, W, 1]`` flagging anchor + slots. Required when ``instruct_type == "hybrid_anchor"``. + + Returns: + Initial latent tensor ready for the denoising loop (bs*T, H, W, 33). + + Raises: + ValueError: If ``instruct_type`` is not supported, or if + ``hybrid_anchor`` is missing ``anchor_latent`` / ``anchor_mask``. + """ + lq_latent = lq_latent.to(device) + lq_latent = cast_to_module_dtype(dit, lq_latent) + if anchor_latent is not None: + anchor_latent = anchor_latent.to(device) + anchor_latent = cast_to_module_dtype(dit, anchor_latent) + if anchor_mask is not None: + anchor_mask = anchor_mask.to(device) + anchor_mask = cast_to_module_dtype(dit, anchor_mask) + if dit.instruct_type == "noise": + g_noise = torch.Generator(device=device) + g_noise.manual_seed(seed) + degraded_lq = degrade_lq_latent(lq_latent, lq_noise_scale, lq_noise_type, generator=g_noise) + if not getattr(dit, "visual_cond", False): + return degraded_lq + zero_cond = torch.zeros_like(degraded_lq) + zero_mask = torch.zeros([*degraded_lq.shape[:-1], 1], dtype=degraded_lq.dtype, device=device) + return torch.cat([degraded_lq, zero_cond, zero_mask], dim=-1) + if dit.instruct_type in ("channel", "hybrid"): + if dit.instruct_type == "hybrid": + g_noise = torch.Generator(device=device) + g_noise.manual_seed(seed) + starting = degrade_lq_latent(lq_latent.clone(), lq_noise_scale, lq_noise_type, generator=g_noise) + else: + g = torch.Generator(device=device) + g.manual_seed(seed) + starting = torch.randn(bs * duration, height, width, dit.in_visual_dim, device=device, generator=g) + channel_lq = degrade_lq_latent(lq_latent, lq_channel_noise_scale, noise_type="linear") + mask = torch.ones_like(lq_latent[..., :1]) + return torch.cat([starting, channel_lq, mask], dim=-1) + if dit.instruct_type == "hybrid_anchor": + if anchor_latent is None or anchor_mask is None: + if not anchor_free: + msg = "hybrid_anchor requires anchor_latent and anchor_mask" + raise ValueError(msg) + anchor_latent = torch.zeros_like(lq_latent) + anchor_mask = torch.zeros_like(lq_latent[..., :1]) + g_noise = torch.Generator(device=device) + g_noise.manual_seed(seed) + starting = degrade_lq_latent(lq_latent.clone(), lq_noise_scale, lq_noise_type, generator=g_noise) + anchor = degrade_lq_latent(anchor_latent, lq_channel_noise_scale, noise_type="linear") + return torch.cat([starting, anchor, anchor_mask], dim=-1) + msg = f"generate_sample_sr does not support instruct_type={dit.instruct_type!r}" + raise ValueError(msg) + + +def _empty_text_embeds(bs: int, device: str | int) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + """Build placeholder text embeddings for a text-free DiT. + + A ``use_text=False`` DiT ignores every text input (encoder, cross-attention + and pooled-text conditioning are all absent), so no real embeddings are + needed. These zero-size tensors only satisfy the indexing and ``cu_seqlens`` + bookkeeping in ``generate_sample_sr`` / ``get_model_prediction`` without ever + loading the cached empty-caption file. + + Args: + bs: Batch size. + device: Target CUDA device. + + Returns: + Tuple of ``(text_embeds_dict, text_cu_seqlens)`` with empty sequences. + """ + embeds = {"text_embeds": torch.zeros(0, 1, device=device), "pooled_embed": torch.zeros(bs, 1, device=device)} + cu_seqlens = torch.zeros(bs + 1, dtype=torch.int32, device=device) + return (embeds, cu_seqlens) + + +def _encode_text( + bs: int, + device: str | int, + text_embedder: Any | None = None, + cached_text_embeds: dict[str, torch.Tensor] | None = None, + *, + use_text: bool = True, +) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor], torch.Tensor]: + """Encode text or use cached embeddings for SR generation. + + For SR we always use empty captions. Returns both text and null-text + embeddings (identical when using cached embeds). + + Args: + bs: Batch size. + device: Target CUDA device. + text_embedder: Text encoder (Qwen2.5-VL + CLIP). Can be ``None`` + when ``cached_text_embeds`` is provided. + cached_text_embeds: Pre-computed empty-caption embeddings (bs=1) + loaded from disk. + use_text: Whether the DiT consumes text. When ``False`` placeholder + zero embeddings are returned and neither a text encoder nor cached + embeddings are required. + + Returns: + Tuple of ``(text_embed, text_cu_seqlens, null_text_embed, + null_text_cu_seqlens)``. + + Raises: + ValueError: If a text DiT is given neither ``text_embedder`` nor + ``cached_text_embeds``. + """ + if not use_text: + embeds, cu_seqlens = _empty_text_embeds(bs, device) + return (embeds, cu_seqlens, embeds, cu_seqlens) + if cached_text_embeds is not None: + bs_text_embed, text_cu_seqlens = replicate_cached_embeds(cached_text_embeds, bs, device) + return (bs_text_embed, text_cu_seqlens, bs_text_embed, text_cu_seqlens) + if text_embedder is not None: + empty_captions = [""] * bs + bs_text_embed, text_cu_seqlens = text_embedder.encode(empty_captions, images=None, type_of_content="video") + bs_null_text_embed, null_text_cu_seqlens = text_embedder.encode( + empty_captions, images=None, type_of_content="video" + ) + bs_text_embed = {k: v.to(device) for k, v in bs_text_embed.items()} + text_cu_seqlens = text_cu_seqlens.to(device) + bs_null_text_embed = {k: v.to(device) for k, v in bs_null_text_embed.items()} + null_text_cu_seqlens = null_text_cu_seqlens.to(device) + return (bs_text_embed, text_cu_seqlens, bs_null_text_embed, null_text_cu_seqlens) + msg = "generate_sample_sr requires either text_embedder or cached_text_embeds" + raise ValueError(msg) + + +@torch.no_grad() +def generate_sample_sr( + *, + dit: torch.nn.Module, + vae: torch.nn.Module, + scale_factor: tuple[float, ...] = (1.0, 1.0, 1.0), + num_steps: int = 50, + guidance_weight: float = 5.0, + scheduler_scale: float = 5.0, + seed: int = 42, + device: str | int = "cuda", + tp_mesh: dict[str, Any] | None = None, + lq_noise_scale: float = 0.0, + lq_noise_type: Literal["linear", "ddpm"] = "linear", + lq_channel_noise_scale: float = 0.0, + text_embedder: Any | None = None, + cached_text_embeds: dict[str, torch.Tensor] | None = None, + lq_videos: list[torch.Tensor] | None = None, + lq_latents: torch.Tensor | None = None, + n_samples: int | None = None, + cap_noise_timestep: bool = False, + vae_decode_batch: bool = False, + prediction_target: str = "velocity", + channelcat_drop_threshold: float = 0.0, + anchor_latents: torch.Tensor | None = None, + anchor_masks: torch.Tensor | None = None, + anchor_free: bool = False, + rfg_scale: float = 1.0, +) -> torch.Tensor: + """Generate super-resolved video from LQ inputs (pixels or latents). + + Accepts either raw pixel videos (``lq_videos``) or pre-encoded latents + (``lq_latents``). When ``lq_latents`` is provided, VAE encoding is + skipped and the latents are used directly. + + Builds the initial latent tensor, runs the denoising loop via + ``generate()``, and decodes back to pixel space. + + Supports three SR instruct modes: + + - ``instruct_type="channel"``: random noise in HQ channels with LQ latent + and mask concatenated as extra channels → ``[noise(C) | lq(C) | mask(1)]``. + - ``instruct_type="noise"``: LQ latent is used directly as the starting + point (replaces random noise) → ``[lq(C)]``, zero extra channels. + - ``instruct_type="hybrid"``: LQ latent as starting point (like ``"noise"``) + AND LQ latent concatenated as extra channels (like ``"channel"``) → + ``[degraded_lq(C) | lq(C) | mask(1)]``. + + When ``lq_noise_scale > 0``, noise is mixed into the LQ latent used as + the starting point (applies to ``"noise"`` and ``"hybrid"``). + + When ``lq_channel_noise_scale > 0``, noise is mixed into the LQ latent + concatenated as extra channels (applies to ``"channel"`` and ``"hybrid"``). + + Args: + dit: DiT model with ``instruct_type`` in + ``{"channel", "noise", "hybrid"}``. + vae: Pre-trained VAE model. + scale_factor: Per-axis RoPE frequency scaling. + num_steps: Number of diffusion denoising steps. + guidance_weight: Classifier-free guidance strength. + scheduler_scale: Timestep scheduler scaling factor. + seed: RNG seed for noise initialization. + device: Target CUDA device. + tp_mesh: Tensor parallelism mesh (optional). + lq_noise_scale: Noise fraction ``s`` mixed into LQ latent used as + the noise/starting point (``0`` = disabled). Applies to + ``"noise"`` and ``"hybrid"`` modes. + lq_noise_type: Noising formula — ``"linear"`` or ``"ddpm"``. + lq_channel_noise_scale: Noise fraction mixed into LQ latent + concatenated as extra channels (``0`` = disabled). Applies to + ``"channel"`` and ``"hybrid"`` modes. + text_embedder: Text encoder (Qwen2.5-VL + CLIP). Can be ``None`` + when ``cached_text_embeds`` is provided. + cached_text_embeds: Pre-computed empty-caption embeddings (bs=1) + loaded from disk. When provided, ``text_embedder`` is not called. + lq_videos: List of ``[T, H, W, 3]`` float32 tensors in ``[0, 255]``. + lq_latents: Pre-encoded LQ latents (skips VAE encoding). + n_samples: Required when ``lq_latents`` is provided. + cap_noise_timestep: When ``True`` and ``instruct_type`` in + ``("noise", "hybrid")``, starts the denoising loop from raw + ``t = lq_noise_scale`` instead of ``t = 1.0``. Matches the + training-time behaviour when ``cap_noise_timestep`` is also set + during training. Default ``False`` preserves current behaviour. + vae_decode_batch: When ``True``, decode all samples in a single VAE + call instead of one-by-one. Faster but uses more GPU memory. + prediction_target: ``"velocity"`` for velocity prediction (default), + ``"x0"`` for clean-image prediction. + channelcat_drop_threshold: When ``> 0``, zero out channel-cat LQ + channels during denoising once ``t < threshold``. ``0`` = + disabled (default). + anchor_latents: Pre-encoded sparse HR-anchor latents + ``[bs*T', H', W', C]``, **already scaled** by VAE + ``scaling_factor``. Required for ``instruct_type="hybrid_anchor"``. + anchor_masks: Sparse binary anchor mask ``[bs*T', H', W', 1]``. + Required for ``instruct_type="hybrid_anchor"``. + anchor_free: When ``True`` and ``instruct_type="hybrid_anchor"`` with + no anchors provided, run anchor-free (zeroed anchor + mask) — the + tiled path for GT-less real-world LQ videos. + rfg_scale: Reference-Free Guidance strength on the anchor. ``1.0`` + (default) runs a single cond forward per denoising step; values + ``!= 1.0`` add a second uncond forward (anchor + anchor_mask + zeroed) and blend ``v_uncond + rfg_scale * (v_cond - v_uncond)``. + Only applies when ``instruct_type="hybrid_anchor"``. Cost: +1 + DiT forward per denoising step when active. + Returns: + Generated SR videos as ``[bs, 3, T, H, W]`` uint8 tensor. + + Raises: + ValueError: If neither ``text_embedder`` nor ``cached_text_embeds`` + is provided or if ``instruct_type`` is unsupported. + """ + piflow_params = getattr(dit, "piflow_params", None) + if piflow_params is not None: + if rfg_scale != 1.0: + msg = "π-Flow checkpoints do not support Reference-Free Guidance (drop rfg_scale)" + raise ValueError(msg) + return run_stages( + dit=dit, + vae=vae, + scale_factor=scale_factor, + num_steps=num_steps, + guidance_weight=guidance_weight, + scheduler_scale=scheduler_scale, + seed=seed, + device=device, + tp_mesh=tp_mesh, + lq_noise_scale=lq_noise_scale, + lq_noise_type=lq_noise_type, + lq_channel_noise_scale=lq_channel_noise_scale, + text_embedder=text_embedder, + cached_text_embeds=cached_text_embeds, + lq_videos=lq_videos, + lq_latents=lq_latents, + n_samples=n_samples, + cap_noise_timestep=cap_noise_timestep, + vae_decode_batch=vae_decode_batch, + prediction_target=prediction_target, + channelcat_drop_threshold=channelcat_drop_threshold, + anchor_latents=anchor_latents, + anchor_masks=anchor_masks, + anchor_free=anchor_free, + rfg_scale=rfg_scale, + piflow_params=piflow_params, + ) + + +@torch.no_grad() +def piflow_generate( + img: torch.Tensor, + model: torch.nn.Module, + text_embeds: dict[str, torch.Tensor], + visual_cu_seqlens: torch.Tensor, + text_cu_seqlens: torch.Tensor, + visual_rope_pos: list[torch.Tensor], + text_rope_pos: torch.Tensor, + scale_factor: tuple[float, ...], + *, + nfe: int, + num_policy_substeps: int, + final_step_size_scale: float, + shift: float, + n_grid: int, + out_dim: int, + eps: float = 1e-06, + tp_mesh: dict[str, Any] | None = None, + start_timestep: float = 1.0, + device: str | int | None = None, + progress_callback=None, + scheduler: Any | None = None, + progress_bar: Any = None, +) -> torch.Tensor: + """Few-step π-Flow denoise: ``nfe`` network calls, network-free integration between. + + Walks the deterministic ``nfe``-segment schedule from ``raw_t = start_timestep`` + down to 0 (final segment scaled by ``final_step_size_scale``, matching ``sample_t``). + Each segment: one ``model`` forward -> ``DXPolicy`` -> integrate over the segment. + """ + device = device if device is not None else img.device + img = img.to(device) + visual_cu_seqlens = visual_cu_seqlens.to(device) + text_cu_seqlens = text_cu_seqlens.to(device) + visual_rope_pos = [position.to(device) for position in visual_rope_pos] + text_rope_pos = text_rope_pos.to(device) + text_embeds = {key: value.to(device) for key, value in text_embeds.items()} + sparse_params = get_sparse_params(model, img, visual_cu_seqlens) + if tp_mesh: + tp_world_size = tp_mesh["tp"].size() + tp_rank = tp_mesh["tp"].get_local_rank() + img = torch.chunk(img, tp_world_size, dim=1)[tp_rank] + ndim = img.dim() + if nfe < 1: + msg = f"piflow_generate requires nfe >= 1, got {nfe}" + raise ValueError(msg) + if scheduler is not None: + scheduler.set_timesteps(nfe, device=device) + if scheduler.timesteps.shape[0] != nfe: + raise ValueError(f"Piflow scheduler returned {scheduler.timesteps.shape[0]} steps, expected {nfe}") + else: + final_step_size_scale = max(float(final_step_size_scale), eps) + one_minus_final = 1.0 - final_step_size_scale + base_seg = 1.0 / (float(nfe) - one_minus_final) + final_seg = final_step_size_scale * base_seg + x = img + for step_index in range(nfe): + if scheduler is not None: + timestep = scheduler.timesteps[step_index] + sigma_src = scheduler.sigmas[step_index].item() + else: + idx = nfe - step_index + raw_src = min(max((idx - one_minus_final) * base_seg, eps), start_timestep) + seg = final_seg if idx == 1 else base_seg + raw_dst = max(raw_src - seg, 0.0) + sigma_src = shift_timesteps(torch.tensor(raw_src, device=device), shift).item() + n_objects = visual_cu_seqlens.shape[0] - 1 + t_src = torch.full( + (n_objects,), + float(timestep.item()) if scheduler is not None else sigma_src * 1000.0, + device=device, + dtype=x.dtype, + ) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + v0 = model( + x, + text_embeds["text_embeds"], + text_embeds["pooled_embed"], + t_src, + visual_cu_seqlens, + text_cu_seqlens, + visual_rope_pos, + text_rope_pos, + scale_factor=scale_factor, + sparse_params=sparse_params, + ) + if scheduler is not None: + x_pred = scheduler.step(v0, timestep, x[..., :out_dim], return_dict=False)[0] + else: + v0_grid = v0.unsqueeze(1) if n_grid == 1 else v0 + total_tokens = x.shape[0] + sigma_tok = torch.full((total_tokens, *(ndim - 1) * [1]), sigma_src, device=device) + seg_tok = torch.full((total_tokens,), seg, device=device) + policy = DXPolicy(v0_grid, x[..., :out_dim], sigma_tok, seg_tok, shift=shift, mode="grid", eps=eps) + raw_src_tok = torch.full((total_tokens,), raw_src, device=device) + raw_dst_tok = torch.full((total_tokens,), raw_dst, device=device) + x_pred, _, _ = policy_rollout_fm( + x[..., :out_dim], sigma_tok, raw_src_tok, raw_dst_tok, num_policy_substeps, policy + ) + x = torch.cat([x_pred, x[..., out_dim:]], dim=-1) + if progress_callback is not None: + progress_callback() + if progress_bar is not None: + progress_bar.update() + progress_bar.refresh() + return x[..., :out_dim] + + +class RunConfig(BaseModel): + """Parameters for one tiled SR run.""" + + model_config = ConfigDict(extra="forbid") + + device: str + num_steps: int = Field(default=5, ge=2) + seed: int = 42 + overlap: float = Field(default=0.25, ge=0.0, lt=1.0) + tiles_batch_size: int = Field(default=1, gt=0) + resolution_scale: ResolutionScale = 4 + tile_grid_mode: TileGridMode = "even" + tile_grid_min_overlap: float = Field(default=0.20, ge=0.0, lt=1.0) + + +@dataclass +class SRParams: + """SR sampling parameters extracted from the SR training config.""" + + scale_factor: dict[int, list[float]] + visual_size: list[int] + scheduler_scale: float = 5.0 + lq_noise_scale: float = 0.7 + lq_noise_type: str = "ddpm" + lq_channel_noise_scale: float = 0.0 + cap_noise_timestep: bool = False + fps: int = 24 + + +@dataclass +class SRComponents: + """The weighted components and config consumed by the SR sampler.""" + + dit: torch.nn.Module + vae: torch.nn.Module + latent_upscaler: torch.nn.Module | Any | None + sr_params: SRParams + cached_text_embeds: dict[str, torch.Tensor] | None = None + sampler: Callable[..., torch.Tensor] | None = None + spatial_factor: int | None = None + + +def _closest_base_resolution(h: int, w: int, visual_size: int) -> tuple[int, int]: + if visual_size not in RESOLUTIONS: + raise ValueError(f"Unsupported SR visual_size={visual_size}; known sizes: {sorted(RESOLUTIONS)}") + ratio = w / h if h else 1.0 + return min(RESOLUTIONS[visual_size], key=lambda hw: abs(hw[1] / hw[0] - ratio)) + + +def _tile_geometry( # noqa: PLR0913 + h: int, + w: int, + visual_size: int, + scale: int, + overlap: float, + spatial_factor: int = VAE_SPATIAL_FACTOR, + grid_mode: TileGridMode = "even", + grid_min_overlap: float = 0.20, +) -> tuple[tuple[int, int], tuple[int, int], TileGrid]: + base_h, base_w = _closest_base_resolution(h, w, visual_size) + if base_h % scale or base_w % scale: + raise ValueError(f"resolution_scale={scale} does not divide the base resolution {base_h}x{base_w} exactly") + tile_hw = (base_h // scale, base_w // scale) + if grid_mode == "even": + grid = compute_tile_grid_even(h, w, tile_hw, grid_min_overlap, spatial_factor) + else: + grid = compute_tile_grid(h, w, scale, overlap, tile_hw=tile_hw) + return (base_h, base_w), tile_hw, grid + + +def _spatial_factor(components: Any) -> int: + explicit = getattr(components, "spatial_factor", None) + if explicit is not None: + return int(explicit) + vae = getattr(components, "vae", None) + config = getattr(vae, "config", None) + return int( + getattr(vae, "spatial_factor", None) + or getattr(config, "spatial_factor", None) + or getattr(_sr_constants, "VAE_SPATIAL_FACTOR", None) + or VAE_SPATIAL_FACTOR + ) + + +def _component_params(components: Any) -> Any: + params = getattr(components, "sr_params", None) + if params is None: + raise ValueError("SR components must provide sr_params") + return params + + +def _visual_size(params: Any) -> int: + visual_size = getattr(params, "visual_size", None) + if isinstance(visual_size, int): + return visual_size + if not visual_size: + raise ValueError("SR params must provide a non-empty visual_size") + return int(visual_size[0]) + + +def _scale_factor(params: Any, visual_size: int) -> tuple[float, ...]: + values = getattr(params, "scale_factor", None) + if isinstance(values, Mapping): + try: + values = values[visual_size] + except KeyError: + # YAML keeps numeric resolution keys as integers; JSON converts + # object keys to strings when the config is stored in a bundle. + values = values[str(visual_size)] + return tuple(float(value) for value in values) + + +def _extract_latent(result: Any) -> torch.Tensor: + if isinstance(result, tuple): + result = result[0] + distribution = getattr(result, "latent_dist", None) + if distribution is not None: + result = distribution.sample() + if not isinstance(result, torch.Tensor): + raise TypeError("VAE encode must return a Tensor, (Tensor, ...), or latent_dist output") + return result + + +def _validate_equal_shapes(values: list[torch.Tensor], name: str) -> None: + if not values: + raise ValueError(f"{name} batch must not be empty") + shape = tuple(values[0].shape) + if any(tuple(value.shape) != shape for value in values[1:]): + raise ValueError(f"all {name} batch items must have equal shape") + + +@torch.no_grad() +def encode_lq_videos_to_lr_latents( + lq_videos: list[torch.Tensor], + vae: torch.nn.Module, + device: str | torch.device, +) -> torch.Tensor: + """Encode an equal-shaped batch of ``[T,C,H,W]`` videos in one VAE call.""" + _validate_equal_shapes(lq_videos, "video") + pixel = torch.stack(lq_videos).permute(0, 2, 1, 3, 4).to(device=device) + pixel = vae.normalize_data(pixel.float()) if hasattr(vae, "normalize_data") else pixel.float() / 127.5 - 1.0 + pixel = cast_to_module_dtype(vae, pixel) + result = _extract_latent(vae.encode(pixel)) + if result.ndim != VIDEO_BATCH_RANK: + raise ValueError(f"batched VAE encode must return rank 5, got shape {tuple(result.shape)}") + return result.permute(0, 2, 1, 3, 4).float() + + +def latent_tile_grid_from_pixel_grid(pixel_grid: TileGrid, spatial_factor: int) -> TileGrid: + """Convert an aligned pixel grid to the corresponding latent grid.""" + values = (pixel_grid.tile_h, pixel_grid.tile_w, *pixel_grid.tops, *pixel_grid.lefts) + if any(value % spatial_factor for value in values): + raise ValueError( + f"Pixel tile grid (tile {pixel_grid.tile_h}x{pixel_grid.tile_w}, " + f"tops={pixel_grid.tops}, lefts={pixel_grid.lefts}) not aligned to VAE spatial factor " + f"{spatial_factor}." + ) + return TileGrid( + pixel_grid.tile_h // spatial_factor, + pixel_grid.tile_w // spatial_factor, + tuple(top // spatial_factor for top in pixel_grid.tops), + tuple(left // spatial_factor for left in pixel_grid.lefts), + ) + + +@torch.no_grad() +def upscale_lr_latent_tile( + lr_latent_tile: torch.Tensor, + latent_upscaler: torch.nn.Module, + vae: torch.nn.Module, + device: str | torch.device, +) -> torch.Tensor: + """Upscale one raw latent tile into the sampler's channels-last layout.""" + scaling_factor = float(getattr(getattr(vae, "config", None), "scaling_factor", 1.0)) + tile = lr_latent_tile.permute(1, 0, 2, 3).unsqueeze(0).to(device=device, dtype=torch.float32) + tile = cast_to_module_dtype(latent_upscaler, tile) + if str(device).startswith("cuda"): + with torch.autocast("cuda", dtype=torch.bfloat16): + upscaled = run_latent_upscaler(latent_upscaler, tile * scaling_factor) + else: + upscaled = run_latent_upscaler(latent_upscaler, tile * scaling_factor) + return upscaled.squeeze(0).permute(1, 2, 3, 0).float() + + +def _upsample_tiles_to_base(raw_tiles: list[torch.Tensor], base_h: int, base_w: int) -> list[torch.Tensor]: + return [ + functional.interpolate(tile.float(), size=(base_h, base_w), mode="bilinear", align_corners=False).permute( + 0, 2, 3, 1 + ) + for tile in raw_tiles + ] + + +"""Reusable Kandinsky SR inference stages and tiled orchestration. + +The functions in this module deliberately do not depend on a benchmark +runner. A framework can call ``text_encode`` / ``prepare_latents`` / +``denoise`` / ``vae_decode`` directly, or use the tiled functions for the +scale-aware path. Native benchmark integrations may add timing contexts to +the tiled calls; the public :class:`Kandinsky6SRPipeline` does not expose +those callbacks. +""" + + +@dataclass(frozen=True) +class SRLatentState: + """Inputs and initial noisy latent for one SR DiT invocation.""" + + lq_latent: torch.Tensor + image: torch.Tensor + batch_size: int + duration: int + height: int + width: int + + +@dataclass(frozen=True) +class SRTextState: + """Conditional and null text state consumed by SR denoising.""" + + text_embeds: dict[str, torch.Tensor] + text_cu_seqlens: torch.Tensor + null_text_embeds: dict[str, torch.Tensor] + null_text_cu_seqlens: torch.Tensor + + +@torch.no_grad() +def prepare_latents( + *, + dit: torch.nn.Module, + vae: torch.nn.Module, + device: str | int, + lq_videos: list[torch.Tensor] | None = None, + lq_latents: torch.Tensor | None = None, + n_samples: int | None = None, + seed: int = 42, + lq_noise_scale: float = 0.0, + lq_noise_type: str = "linear", + lq_channel_noise_scale: float = 0.0, + anchor_latents: torch.Tensor | None = None, + anchor_masks: torch.Tensor | None = None, + anchor_free: bool = False, +) -> SRLatentState: + """Encode LQ input and build the initial latent for one SR tile batch.""" + if lq_latents is not None: + lq_latent = lq_latents.to(device) + if n_samples is None: + raise ValueError("n_samples is required when lq_latents is provided") + batch_size = n_samples + elif lq_videos is not None: + batch_size = len(lq_videos) + lq_latent = _encode_lq_videos(lq_videos, vae, device) + else: + raise ValueError("Either lq_videos or lq_latents must be provided") + duration = lq_latent.shape[0] // batch_size + height, width = (lq_latent.shape[1], lq_latent.shape[2]) + image = _build_initial_latent( + dit=dit, + lq_latent=lq_latent, + bs=batch_size, + duration=duration, + height=height, + width=width, + device=device, + seed=seed, + lq_noise_scale=lq_noise_scale, + lq_noise_type=lq_noise_type, + lq_channel_noise_scale=lq_channel_noise_scale, + anchor_latent=anchor_latents.to(device) if anchor_latents is not None else None, + anchor_mask=anchor_masks.to(device) if anchor_masks is not None else None, + anchor_free=anchor_free, + ) + return SRLatentState(lq_latent, image, batch_size, duration, height, width) + + +@torch.no_grad() +def text_encode( + *, + dit: torch.nn.Module, + batch_size: int, + device: str | int, + text_embedder: Any | None = None, + cached_text_embeds: dict[str, torch.Tensor] | None = None, +) -> SRTextState: + """Build SR conditional and null text embeddings.""" + text_embeds, text_cu_seqlens, null_text_embeds, null_text_cu_seqlens = _encode_text( + bs=batch_size, + device=device, + text_embedder=text_embedder, + cached_text_embeds=cached_text_embeds, + use_text=getattr(dit, "use_text", True), + ) + return SRTextState(text_embeds, text_cu_seqlens, null_text_embeds, null_text_cu_seqlens) + + +def _positions( + *, dit: torch.nn.Module, latent_state: SRLatentState, text_state: SRTextState +) -> tuple[torch.Tensor, list[torch.Tensor], torch.Tensor, torch.Tensor]: + runtime_device = latent_state.image.device + visual_cu_seqlens = latent_state.duration * torch.arange( + latent_state.batch_size + 1, dtype=torch.int32, device=runtime_device + ) + visual_rope_pos = [ + torch.cat([torch.arange(int(end), device=runtime_device) for end in torch.diff(visual_cu_seqlens).cpu()]), + torch.arange(latent_state.height // dit.patch_size[1], device=runtime_device), + torch.arange(latent_state.width // dit.patch_size[2], device=runtime_device), + ] + text_cu_seqlens = text_state.text_cu_seqlens.to(runtime_device) + null_text_cu_seqlens = text_state.null_text_cu_seqlens.to(runtime_device) + text_rope_pos = torch.cat( + [torch.arange(int(end), device=runtime_device) for end in torch.diff(text_cu_seqlens).cpu()] + ) + null_text_rope_pos = torch.cat( + [torch.arange(int(end), device=runtime_device) for end in torch.diff(null_text_cu_seqlens).cpu()] + ) + return (visual_cu_seqlens, visual_rope_pos, text_rope_pos, null_text_rope_pos) + + +@torch.no_grad() +def denoise( + *, + dit: torch.nn.Module, + latent_state: SRLatentState, + text_state: SRTextState, + scale_factor: tuple[float, ...], + num_steps: int = 50, + guidance_weight: float = 5.0, + scheduler_scale: float = 5.0, + tp_mesh: dict[str, Any] | None = None, + lq_noise_scale: float = 0.0, + cap_noise_timestep: bool = False, + prediction_target: str = "velocity", + channelcat_drop_threshold: float = 0.0, + rfg_scale: float = 1.0, + piflow_params: dict[str, Any] | None = None, + device: str | int | None = None, + progress_callback: Callable[[int], Any] | None = None, + scheduler: Any | None = None, + progress_bar: Any = None, +) -> torch.Tensor: + """Run either the Euler or π-Flow SR denoising stage.""" + piflow_params = piflow_params or _scheduler_piflow_params(scheduler) + piflow_params = piflow_params or getattr(dit, "piflow_params", None) + if piflow_params is not None: + if rfg_scale != 1.0: + raise ValueError("π-Flow checkpoints do not support Reference-Free Guidance") + if cap_noise_timestep and dit.instruct_type in ("noise", "hybrid"): + raise NotImplementedError("piflow sampler does not support cap_noise_timestep for noise/hybrid instruct") + device = device if device is not None else latent_state.image.device + visual_cu_seqlens, visual_rope_pos, text_rope_pos, null_text_rope_pos = _positions( + dit=dit, latent_state=latent_state, text_state=text_state + ) + if piflow_params is not None: + out_dim = int(getattr(dit, "base_out_visual_dim", dit.in_visual_dim)) + start_t = 1.0 + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + latent_visual = piflow_generate( + latent_state.image, + dit, + text_state.text_embeds, + visual_cu_seqlens, + text_state.text_cu_seqlens, + visual_rope_pos, + text_rope_pos, + scale_factor, + **piflow_params, + out_dim=out_dim, + start_timestep=start_t, + device=device, + progress_callback=progress_callback, + scheduler=scheduler, + progress_bar=progress_bar, + ) + else: + start_t = ( + lq_noise_scale if cap_noise_timestep and dit.instruct_type in ("noise", "hybrid", "hybrid_anchor") else 1.0 + ) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + latent_visual = generate( + latent_state.image, + dit, + device, + num_steps, + text_state.text_embeds, + text_state.null_text_embeds, + visual_cu_seqlens, + text_state.text_cu_seqlens, + text_state.null_text_cu_seqlens, + visual_rope_pos, + text_rope_pos, + null_text_rope_pos, + scale_factor, + guidance_weight, + scheduler_scale, + tp_mesh=tp_mesh, + first_frames=None, + start_timestep=start_t, + prediction_target=prediction_target, + channelcat_drop_threshold=channelcat_drop_threshold, + rfg_scale=rfg_scale, + progress_callback=progress_callback, + progress_bar=progress_bar, + ) + if tp_mesh: + tensor_list = [ + torch.zeros_like(latent_visual, device=latent_visual.device) for _ in range(tp_mesh["tp"].size()) + ] + all_gather(tensor_list, latent_visual.contiguous(), group=tp_mesh.get_group(mesh_dim="tp")) + latent_visual = torch.cat(tensor_list, dim=1) + return latent_visual + + +def _scheduler_piflow_params(scheduler: Any | None) -> dict[str, Any] | None: + if scheduler is None or not bool(getattr(scheduler, "is_piflow", False)): + return None + config = scheduler.config + nfe = getattr(config, "nfe", None) + if nfe is None: + raise ValueError("PiflowScheduler used for SR must define nfe in scheduler_config.json") + return { + "nfe": int(nfe), + "num_policy_substeps": int(config.num_policy_substeps), + "final_step_size_scale": float(config.final_step_size_scale), + "shift": float(config.shift), + "n_grid": int(config.n_grid), + "eps": float(config.eps), + } + + +@torch.no_grad() +def vae_decode( + *, + vae: torch.nn.Module, + latent_visual: torch.Tensor, + batch_size: int, + duration: int, + height: int, + width: int, + vae_decode_batch: bool = False, +) -> torch.Tensor: + """Decode denoised SR latents into ``[batch, 3, T, H, W]`` uint8 frames""" + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + all_latents = latent_visual.reshape(batch_size, duration, height, width, -1) + all_latents = (all_latents / vae.config.scaling_factor).permute(0, 4, 1, 2, 3) + if vae_decode_batch: + return decode_latent_to_uint8(vae, all_latents) + decoded = [decode_latent_to_uint8(vae, all_latents[i : i + 1]) for i in range(batch_size)] + return torch.cat(decoded, dim=0) + + +@torch.no_grad() +def run_stages( + *, + dit: torch.nn.Module, + vae: torch.nn.Module, + scale_factor: tuple[float, ...], + num_steps: int = 50, + guidance_weight: float = 5.0, + scheduler_scale: float = 5.0, + seed: int = 42, + device: str | int = "cuda", + tp_mesh: dict[str, Any] | None = None, + lq_noise_scale: float = 0.0, + lq_noise_type: str = "linear", + lq_channel_noise_scale: float = 0.0, + text_embedder: Any | None = None, + cached_text_embeds: dict[str, torch.Tensor] | None = None, + lq_videos: list[torch.Tensor] | None = None, + lq_latents: torch.Tensor | None = None, + n_samples: int | None = None, + cap_noise_timestep: bool = False, + vae_decode_batch: bool = False, + prediction_target: str = "velocity", + channelcat_drop_threshold: float = 0.0, + anchor_latents: torch.Tensor | None = None, + anchor_masks: torch.Tensor | None = None, + anchor_free: bool = False, + rfg_scale: float = 1.0, + piflow_params: dict[str, Any] | None = None, + progress_callback: Callable[[int], Any] | None = None, + scheduler: Any | None = None, + progress_bar: Any = None, +) -> torch.Tensor: + """Run the four reusable SR stages for one tile batch.""" + piflow_params = piflow_params or _scheduler_piflow_params(scheduler) + piflow_params = piflow_params or getattr(dit, "piflow_params", None) + latent_state = prepare_latents( + dit=dit, + vae=vae, + device=device, + lq_videos=lq_videos, + lq_latents=lq_latents, + n_samples=n_samples, + seed=seed, + lq_noise_scale=lq_noise_scale, + lq_noise_type=lq_noise_type, + lq_channel_noise_scale=lq_channel_noise_scale, + anchor_latents=anchor_latents, + anchor_masks=anchor_masks, + anchor_free=anchor_free, + ) + text_state = text_encode( + dit=dit, + batch_size=latent_state.batch_size, + device=device, + text_embedder=text_embedder, + cached_text_embeds=cached_text_embeds, + ) + latent_visual = denoise( + dit=dit, + latent_state=latent_state, + text_state=text_state, + scale_factor=scale_factor, + num_steps=num_steps, + guidance_weight=guidance_weight, + scheduler_scale=scheduler_scale, + device=device, + tp_mesh=tp_mesh, + lq_noise_scale=lq_noise_scale, + cap_noise_timestep=cap_noise_timestep, + prediction_target=prediction_target, + channelcat_drop_threshold=channelcat_drop_threshold, + rfg_scale=rfg_scale, + piflow_params=piflow_params, + progress_callback=progress_callback, + scheduler=scheduler, + progress_bar=progress_bar, + ) + return vae_decode( + vae=vae, + latent_visual=latent_visual, + batch_size=latent_state.batch_size, + duration=latent_state.duration, + height=latent_state.height, + width=latent_state.width, + vae_decode_batch=vae_decode_batch, + ) + + +def _generate_kwargs( + components: Any, run_config: Any, params: Any, scale_factor: tuple[float, ...], **inputs: Any +) -> dict[str, Any]: + return { + **inputs, + "dit": components.dit, + "vae": components.vae, + "scale_factor": scale_factor, + "num_steps": run_config.num_steps, + "anchor_free": True, + "guidance_weight": 1.0, + "scheduler_scale": float(getattr(params, "scheduler_scale", 5.0)), + "seed": run_config.seed, + "device": run_config.device, + "tp_mesh": None, + "lq_noise_scale": float(getattr(params, "lq_noise_scale", 0.7)), + "lq_noise_type": getattr(params, "lq_noise_type", "ddpm"), + "lq_channel_noise_scale": float(getattr(params, "lq_channel_noise_scale", 0.0)), + "cap_noise_timestep": bool(getattr(params, "cap_noise_timestep", False)), + "cached_text_embeds": getattr(components, "cached_text_embeds", None), + } + + +def _as_tiles(sr_batch: torch.Tensor) -> list[torch.Tensor]: + if sr_batch.ndim != 5: + raise ValueError(f"SR sampler must return [batch,C,T,H,W], got {tuple(sr_batch.shape)}") + return [sample.float().cpu() for sample in sr_batch] + + +def _scale_progress_callback(progress_callback: Callable[[int], Any], chunk_size: int) -> Callable[[int], Any]: + + def update(n: int = 1) -> None: + progress_callback(n * chunk_size) + + return update + + +def _loaded_lu_scales(components: Any) -> tuple[int, ...]: + lu = getattr(components, "latent_upscaler", None) + if lu is None: + return () + scales = getattr(lu, "scales", None) + return tuple((int(scale) for scale in scales)) if scales is not None else (latent_upscaler_scale(lu),) + + +def _run_tile_batches( + tile_inputs: list[torch.Tensor], + sr_components: Any, + run_config: Any, + scale_factor: tuple[float, ...], + *, + prepare_chunk: Callable[[list[torch.Tensor]], list[torch.Tensor]] | None = None, + progress_callback: Callable[[int], Any] | None = None, + scheduler: Any | None = None, + sample_batch_size: int = 1, + progress_bar: Any = None, +) -> list[torch.Tensor]: + """Run the shared batched sampler loop for latent or pixel tile inputs. + + DiT and SR VAE stay resident for the complete tile loop. Only the optional + latent-upscaler is streamed per chunk. + """ + outputs: list[torch.Tensor] = [] + sampler = getattr(sr_components, "sampler", None) + params = _component_params(sr_components) + piflow_params = _scheduler_piflow_params(scheduler) or getattr(sr_components.dit, "piflow_params", None) + progress_steps = int(piflow_params["nfe"] if piflow_params is not None else run_config.num_steps - 1) + for start in range(0, len(tile_inputs), run_config.tiles_batch_size): + raw_chunk = tile_inputs[start : start + run_config.tiles_batch_size] + chunk = raw_chunk + if prepare_chunk is not None: + chunk = prepare_chunk(raw_chunk) + input_kwargs: dict[str, Any] + if prepare_chunk is None: + input_kwargs = {"lq_videos": chunk} + else: + input_kwargs = {"lq_latents": torch.cat(chunk, dim=0), "n_samples": len(raw_chunk)} + kwargs = _generate_kwargs(sr_components, run_config, params, scale_factor, **input_kwargs) + kwargs["seed"] = run_config.seed + start // sample_batch_size + chunk_progress = ( + _scale_progress_callback(progress_callback, len(raw_chunk)) if progress_callback is not None else None + ) + if sampler is None: + outputs.extend( + _as_tiles( + run_stages( + **kwargs, + piflow_params=piflow_params, + progress_callback=chunk_progress, + scheduler=scheduler, + progress_bar=progress_bar, + ) + ) + ) + else: + outputs.extend(_as_tiles(sampler(**kwargs))) + if chunk_progress is not None: + chunk_progress(progress_steps) + return outputs + + +def align_to_vae_stride(t: int) -> int: + """Round ``t`` to the nearest valid pixel-frame count ``1 + 8·k`` (ties up). + + SR data uses ``1 + 8·k`` frames so the temporal VAE stride round-trips + cleanly. Rounding to nearest (not truncation) keeps a requested duration + close to the target. + + Args: + t: Raw frame count. + + Returns: + Nearest valid frame count ``>= 1``. + """ + if t <= 1: + return 1 + k_floor = (t - 1) // 8 + t_floor = 1 + 8 * k_floor + t_ceil = 1 + 8 * (k_floor + 1) + return t_ceil if (t - t_floor) >= (t_ceil - t) else t_floor + + +def select_frame_indices(total_frames: int, src_fps: float, target_fps: float) -> list[int]: + """Fixed-stride frame indices that downsample ``src_fps`` to ``target_fps``. + + Mirrors ``data_encoding.frame_utils.select_frame_indices``: with + ``step = src_fps / target_fps >= 1``, ``round(i * step)`` is strictly + non-decreasing, so the selection never duplicates a source frame. + + Args: + total_frames: Number of frames in the source. + src_fps: Source frame rate (must be ``>= target_fps``). + target_fps: Desired frame rate. + + Returns: + Sorted list of integer source-frame indices. + """ + step = src_fps / target_fps + indices = [round(i * step) for i in range(int(total_frames / step))] + return [i for i in indices if i < total_frames] + + +def resample_to_target_fps( + video: torch.Tensor, + src_fps: float, + target_fps: int = TARGET_FPS, +) -> tuple[torch.Tensor, int]: + """Resample a decoded video toward ``target_fps`` (downsample / no-op / keep). + + Three tiers, mirroring the dataset encoding pipeline: + + - ``|src_fps - target_fps| < RESAMPLE_FPS_TOLERANCE``: pass through unchanged. + - ``src_fps > target_fps``: fixed-stride downsample via + :func:`select_frame_indices`; the result is at ``target_fps``. + - ``src_fps < target_fps``: kept at the native rate with a warning (no ffmpeg + ``minterpolate`` upsample); the clip stays mildly out of distribution. + + Args: + video: ``[T, C, H, W]`` decoded source video. + src_fps: Source frame rate. + target_fps: Training frame rate to resample toward. + + Returns: + ``(resampled_video, effective_fps)`` where ``effective_fps`` is the rate + the returned frames play at (``target_fps`` when downsampled, otherwise + ``round(src_fps)``) and is the correct rate to save the SR result at. + """ + if abs(src_fps - target_fps) < RESAMPLE_FPS_TOLERANCE: + return video, round(src_fps) + if src_fps > target_fps: + indices = select_frame_indices(video.shape[0], src_fps, target_fps) + logger.warning( + "Source fps {:.2f} > target {}fps: downsampling {} frames -> {} (fixed-stride); " + "source temporal detail beyond {}fps is discarded.", + src_fps, + target_fps, + video.shape[0], + len(indices), + target_fps, + ) + return video[indices], target_fps + logger.warning( + "Source fps {:.2f} < target {}fps: keeping native frames (no minterpolate upsample); " + "output is mildly out of distribution.", + src_fps, + target_fps, + ) + return video, round(src_fps) + + +def clip_to_aligned_frames(video: torch.Tensor, max_num_frames: int = MAX_NUM_FRAMES) -> torch.Tensor: + """Take the first ``max_num_frames`` and floor-align to ``1 + 8k`` frames. + + Args: + video: ``[T, C, H, W]`` video. + max_num_frames: Hard cap applied before alignment (``<= 0`` disables it). + + Returns: + ``[T', C, H, W]`` with ``T' == 1 + 8k`` and ``T' <= max_num_frames``. + + Raises: + ValueError: If the video has no frames. + """ + if max_num_frames > 0: + video = video[:max_num_frames] + aligned = 1 + 8 * ((video.shape[0] - 1) // 8) if video.shape[0] > 0 else 0 + if aligned == 0: + msg = "Video has no readable frames." + raise ValueError(msg) + return video[:aligned] + + +def read_video_tchw_uint8(path: Path) -> tuple[torch.Tensor, float]: + """Decode an mp4/mkv to ``([T, C, H, W] uint8, src_fps)`` (no resample/cap). + + Resampling to the training fps and frame-count alignment are applied by the + caller via :func:`resample_to_target_fps` and :func:`clip_to_aligned_frames`. + + Args: + path: Source video file. + + Returns: + ``([T, C, H, W] uint8 video, native_fps)``. + + Raises: + ValueError: If the file has no readable frames or reports no fps. + """ + with av.open(str(path), mode="r") as container: + stream = container.streams.video[0] + src_fps = float(stream.average_rate or stream.base_rate or 0.0) + frames = [ + torch.from_numpy(frame.to_ndarray(format="rgb24")).permute(2, 0, 1) for frame in container.decode(stream) + ] + + if not frames: + msg = f"Video {path.name} has no readable frames." + raise ValueError(msg) + if src_fps <= 0: + msg = f"Video {path.name} reports no usable fps (got {src_fps})." + raise ValueError(msg) + return torch.stack(frames).contiguous(), src_fps + + +def load_lr_latent(path: Path) -> torch.Tensor: + """Load a raw, unscaled LR latent ``.pt`` as ``[T, C, H, W]`` float32. + + The tensor must match what ``encode_lq_video_to_lr_latent`` produces for + the model's VAE (no scaling factor applied) — the LU multiplies internally. + + Args: + path: Path to a ``.pt`` file holding a 4D latent tensor. + + Returns: + ``[T, C, H, W]`` float32 latent on CPU. + + Raises: + ValueError: If the loaded object is not a rank-4 float tensor. + """ + obj = torch.load(path, map_location="cpu", weights_only=True) + if not isinstance(obj, torch.Tensor) or obj.ndim != LATENT_NDIM: + shape = getattr(obj, "shape", type(obj).__name__) + msg = f"Expected a rank-4 [T, C, H, W] latent tensor in {path}, got {shape}." + raise ValueError(msg) + logger.info("Loaded LR latent {} from {}", tuple(obj.shape), path) + return obj.float() + + +def _video_np(frames: Tensor) -> np.ndarray: + """(3, T, H, W) uint8 → (T, H, W, 3) uint8.""" + return frames.detach().permute(1, 2, 3, 0).cpu().numpy().astype(np.uint8, copy=False) + + +def _audio_np(audio: list[np.ndarray] | np.ndarray) -> np.ndarray: + """int16 / float waveform → float32 mono in [-1, 1].""" + audio_np = np.asarray(audio[0] if isinstance(audio, list) else audio) + if np.issubdtype(audio_np.dtype, np.integer): + audio_np = audio_np.astype(np.float32) / np.iinfo(audio_np.dtype).max + return np.clip(audio_np, -1.0, 1.0).astype(np.float32, copy=False) + + +def extract_audio_from_video( + source_video: str | Path, + audio_sample_rate: int = 44100, +) -> np.ndarray | None: + """Decode a source video's audio as mono float32 samples. + + SR changes only the video stream. Decoding the source audio here lets the + output mux keep that audio without requiring the SR caller to load the + entire source container itself. Resampling also gives the AAC writer a + stable sample rate for videos whose source audio uses another rate. + """ + if audio_sample_rate <= 0: + raise ValueError("audio_sample_rate must be positive") + + chunks: list[np.ndarray] = [] + with av.open(str(source_video), mode="r") as container: + audio_streams = list(container.streams.audio) + if not audio_streams: + return None + + resampler = av.AudioResampler(format="fltp", layout="mono", rate=audio_sample_rate) + for frame in container.decode(audio=0): + for resampled in resampler.resample(frame): + chunks.append(resampled.to_ndarray().reshape(-1)) + for resampled in resampler.resample(None): + chunks.append(resampled.to_ndarray().reshape(-1)) + + if not chunks: + return None + return np.concatenate(chunks).astype(np.float32, copy=False) + + +def mux_video_audio( # noqa: PLR0913 + frames: Tensor, + audio: list[np.ndarray] | np.ndarray | None, + output_path: str | Path, + fps: int = 24, + audio_sample_rate: int = 44100, + video_crf: int = 18, + source_video: str | Path | None = None, + lossless: bool = False, +) -> Path: + """Mux video (+ optional audio) into ``output_path`` via PyAV — no temp files. + + frames: (3, T, H, W) uint8 — single video, channels-first. + audio: list of (samples,) int16, a single ndarray, or None for video-only. + source_video: optional source container from which audio is copied when + ``audio`` is not supplied. This is useful when SR receives frames + extracted from a generated T2VA video. + lossless: encode video as FFV1 instead of libx264. Audio remains AAC. + """ + if audio is not None and source_video is not None: + raise ValueError("pass either audio or source_video, not both") + if source_video is not None: + audio = extract_audio_from_video(source_video, audio_sample_rate) + + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + video_np = _video_np(frames) + + with av.open(str(output_path), mode="w") as container: + video_stream = container.add_stream("ffv1" if lossless else "libx264", rate=fps) + video_stream.width = video_np.shape[2] + video_stream.height = video_np.shape[1] + if lossless: + video_stream.pix_fmt = "bgra" + video_stream.options = {"level": "3"} + else: + video_stream.pix_fmt = "yuv420p" + video_stream.options = {"crf": str(video_crf)} + + audio_stream = None + if audio is not None: + audio_stream = container.add_stream("aac", rate=audio_sample_rate) + audio_stream.layout = "mono" + audio_stream.bit_rate = 192_000 + + for i, frame_np in enumerate(video_np): + frame = av.VideoFrame.from_ndarray(frame_np, format="rgb24") + frame.pts = i + frame.time_base = Fraction(1, fps) + for packet in video_stream.encode(frame): + container.mux(packet) + + if audio_stream is not None: + audio_np = _audio_np(audio) + audio_frame = av.AudioFrame.from_ndarray(audio_np[np.newaxis, :], format="fltp", layout="mono") + audio_frame.sample_rate = audio_sample_rate + audio_frame.pts = 0 + audio_frame.time_base = Fraction(1, audio_sample_rate) + for packet in audio_stream.encode(audio_frame): + container.mux(packet) + + for packet in video_stream.encode(): + container.mux(packet) + if audio_stream is not None: + for packet in audio_stream.encode(): + container.mux(packet) + + return output_path + + +def latent_upscaler_for_scale(upscaler: Any, scale: int) -> nn.Module | None: + """Keep bank dispatch on the registered parent for Diffusers offload.""" + if upscaler is None: + return None + if hasattr(upscaler, "_models"): + key = f"{int(scale)}x" + if key not in upscaler._models: + return None + # The native helper returns a child module. A Diffusers CPU-offload + # hook is attached to the registered bank, so call its forward method + # with the selected entry instead. + upscaler.target_scale = key + return upscaler + for_scale = getattr(upscaler, "for_scale", None) + if callable(for_scale): + return for_scale(scale) + return upscaler if str(getattr(upscaler, "target_scale", "4x")) == f"{int(scale)}x" else None + + +SRInput = Tensor | np.ndarray | Sequence[Any] | str | Path +SRBatchInput = SRInput | Sequence[SRInput] +SRScale = Literal[2, 4, 2.25] +SROutputType = Literal["pt", "torch", "np", "numpy"] +VIDEO_RANK = 4 +VIDEO_BATCH_RANK = 5 + + +def _is_video_batch(value: Any) -> bool: + if isinstance(value, (Tensor, np.ndarray)): + return value.ndim == VIDEO_BATCH_RANK + if isinstance(value, (str, Path)) or not isinstance(value, Sequence) or not value: + return False + first = value[0] + if isinstance(first, (str, Path)): + return True + if isinstance(first, (Tensor, np.ndarray)): + return first.ndim >= VIDEO_RANK + return False + + +def _video_batch_tensor(value: Tensor | np.ndarray) -> list[Tensor]: + tensor = value.detach() if isinstance(value, Tensor) else torch.as_tensor(value) + if tensor.ndim != VIDEO_BATCH_RANK: + raise ValueError(f"batched video must have rank 5, got shape {tuple(tensor.shape)}") + if tensor.shape[-1] in (1, 3, 4): + return [sample.permute(0, 3, 1, 2) for sample in tensor] + if tensor.shape[1] in (1, 3, 4): + return [sample.permute(1, 0, 2, 3) for sample in tensor] + if tensor.shape[2] in (1, 3, 4): + return [sample for sample in tensor] + raise ValueError("5D video must be [B,C,T,H,W], [B,T,C,H,W], or [B,T,H,W,C]") + + +def _latent_tensor(value: Tensor | np.ndarray, channels: int) -> Tensor: + tensor = value.detach() if isinstance(value, Tensor) else torch.as_tensor(value) + if tensor.ndim != VIDEO_RANK: + raise ValueError(f"latent sample must have rank 4, got shape {tuple(tensor.shape)}") + if tensor.shape[1] == channels: + return tensor + if tensor.shape[-1] == channels: + return tensor.permute(0, 3, 1, 2) + if tensor.shape[0] == channels: + return tensor.permute(1, 0, 2, 3) + raise ValueError("4D latent must be [T,C,H,W], [T,H,W,C], or [C,T,H,W]") + + +def _latent_batch(value: Tensor | np.ndarray | str | Path | Sequence[Any], channels: int) -> list[Tensor]: + if isinstance(value, (str, Path)): + return [_latent_tensor(load_lr_latent(Path(value)), channels)] + if isinstance(value, (Tensor, np.ndarray)): + tensor = value.detach() if isinstance(value, Tensor) else torch.as_tensor(value) + if tensor.ndim == VIDEO_RANK: + return [_latent_tensor(tensor, channels)] + if tensor.ndim == VIDEO_BATCH_RANK: + if tensor.shape[-1] == channels: + return [_latent_tensor(sample, channels) for sample in tensor] + if tensor.shape[1] == channels: + return [_latent_tensor(sample, channels) for sample in tensor] + if tensor.shape[2] == channels: + return [_latent_tensor(sample, channels) for sample in tensor] + raise ValueError("5D latent must be [B,C,T,H,W], [B,T,C,H,W], or [B,T,H,W,C]") + raise ValueError(f"latents must have rank 4 or 5, got shape {tuple(tensor.shape)}") + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + if not value: + raise ValueError("latents batch must not be empty") + result: list[Tensor] = [] + for item in value: + result.extend(_latent_batch(item, channels)) + return result + raise TypeError("latents must be tensors, arrays, paths, or a batch of them") + + +def _validate_batch_shapes(values: list[Tensor], name: str) -> None: + if not values: + raise ValueError(f"{name} batch must not be empty") + shape = tuple(values[0].shape) + if any(tuple(value.shape) != shape for value in values[1:]): + raise ValueError( + f"all {name} batch items must have equal shape, got {[tuple(value.shape) for value in values]}" + ) + + +def _batch_output_values(value: Any, batch_size: int, name: str) -> list[Any]: + if isinstance(value, list): + if len(value) != batch_size: + raise ValueError(f"{name} must contain one value per input ({batch_size}), got {len(value)}") + return value + if name == "save_path" and value is not None and batch_size > 1: + raise ValueError("save_path must be a list when processing multiple SR inputs") + return [value] * batch_size + + +def _stitch_batch( + outputs: list[Tensor], + grid: Any, + batch_size: int, + height: int, + width: int, + scale: float, +) -> Tensor: + """Stitch tile-major output back into ``[B,C,T,H,W]`` samples.""" + tile_count = grid.total_tiles + expected = batch_size * tile_count + if len(outputs) != expected: + raise RuntimeError(f"expected {expected} SR tile outputs, got {len(outputs)}") + return torch.stack( + [ + stitch_tiles_hanning( + [outputs[tile_index * batch_size + sample_index] for tile_index in range(tile_count)], + grid, + height, + width, + scale=scale, + ) + .clamp(0, 255) + .to(torch.uint8) + for sample_index in range(batch_size) + ] + ) + + +def _run_batched_tiles( # noqa: PLR0913 + samples: list[Tensor], + components: Any, + run_config: Any, + *, + latent_input: bool, + progress_bar: Any, + scheduler: Any, +) -> Tensor: + """Run equal-shape samples tile-major and restore the sample batch.""" + params = _component_params(components) + visual_size = int(params.visual_size[0]) if isinstance(params.visual_size, list) else int(params.visual_size) + batch_size = len(samples) + if latent_input: + spatial_factor = _spatial_factor(components) + height = samples[0].shape[-2] * spatial_factor + width = samples[0].shape[-1] * spatial_factor + _, _, pixel_grid = _tile_geometry( + height, + width, + visual_size, + run_config.resolution_scale, + run_config.overlap, + spatial_factor, + run_config.tile_grid_mode, + run_config.tile_grid_min_overlap, + ) + tile_grid = latent_tile_grid_from_pixel_grid(pixel_grid, spatial_factor) + input_grid = tile_grid + stitch_grid = pixel_grid + tile_sets = [extract_all_tiles(sample, tile_grid) for sample in samples] + upscaler = latent_upscaler_for_scale(components.latent_upscaler, run_config.resolution_scale) + if upscaler is None: + raise ValueError( + f"The latent-upscaler path has no {run_config.resolution_scale}x model " + f"(loaded LU scales: {_loaded_lu_scales(components)})" + ) + prepare_chunk = lambda chunk: [ + upscale_lr_latent_tile(tile, upscaler, components.vae, run_config.device) for tile in chunk + ] + else: + height, width = samples[0].shape[-2:] + base, _tile_hw, grid = _tile_geometry( + height, + width, + visual_size, + run_config.resolution_scale, + run_config.overlap, + _spatial_factor(components), + run_config.tile_grid_mode, + run_config.tile_grid_min_overlap, + ) + tile_sets = [_upsample_tiles_to_base(extract_all_tiles(sample, grid), base[0], base[1]) for sample in samples] + input_grid = grid + stitch_grid = grid + prepare_chunk = None + + tile_inputs = [ + tile_sets[sample_index][tile_index] + for tile_index in range(input_grid.total_tiles) + for sample_index in range(batch_size) + ] + outputs = _run_tile_batches( + tile_inputs, + components, + run_config, + _scale_factor(params, visual_size), + scheduler=scheduler, + prepare_chunk=prepare_chunk, + progress_bar=progress_bar, + ) + return _stitch_batch(outputs, stitch_grid, batch_size, height, width, run_config.resolution_scale) + + +def _sr_progress_plan( + value: Tensor, + components: Any, + run_config: Any, + *, + latent_input: bool, + batch_size: int = 1, + scheduler: Any | None = None, +) -> tuple[int, int, int]: + """Return total updates, denoising steps, and tile batches.""" + params = _component_params(components) + visual_size = int(params.visual_size[0]) if isinstance(params.visual_size, list) else int(params.visual_size) + spatial_factor = _spatial_factor(components) + height, width = value.shape[-2:] + if latent_input: + height *= spatial_factor + width *= spatial_factor + _, _, grid = _tile_geometry( + height, + width, + visual_size, + run_config.resolution_scale, + run_config.overlap, + spatial_factor, + ) + tile_count = grid.total_tiles * batch_size + chunk_count = max(1, math.ceil(tile_count / run_config.tiles_batch_size)) + piflow_params = _scheduler_piflow_params(scheduler) or getattr(components.dit, "piflow_params", None) + if isinstance(piflow_params, Mapping): + denoise_steps = int(piflow_params.get("nfe", run_config.num_steps - 1)) + elif piflow_params is not None: + denoise_steps = int(getattr(piflow_params, "nfe", run_config.num_steps - 1)) + else: + denoise_steps = run_config.num_steps - 1 + if getattr(components, "sampler", None) is not None: + denoise_steps = 1 + denoise_steps = max(1, denoise_steps) + return chunk_count * denoise_steps, denoise_steps, chunk_count + + +class Kandinsky6SRPipeline(DiffusionPipeline): + r"""Standalone Diffusers pipeline for Kandinsky 6 video super-resolution. + + Components are supplied to the constructor or loaded by Diffusers through + ``from_pretrained``. This pipeline deliberately has no config-path + factory: model construction belongs to the Diffusers component package, + while this class owns only SR orchestration and I/O. + + ``resolution_scale=2.25`` is the production route. It performs a 1.125x + pixel pre-upscale followed by the x2 latent-upscaler path. The input + components must already be on a device; ``.to(device)`` can be used after + construction. + + Args: + transformer: SR transformer used to denoise latent tiles. + vae: Video VAE used for pixel-input encoding. + latent_upscaler: Optional x2/x4 latent upscaler bank. + source_vae: Optional source VAE used only by the KVAE latent bridge. + """ + + model_cpu_offload_seq = "source_vae->latent_upscaler->transformer->vae" + _optional_components = ["latent_upscaler", "source_vae"] + + def __init__( # noqa: PLR0913 + self, + transformer: nn.Module, + vae: nn.Module, + scheduler: Any, + latent_upscaler: nn.Module | None = None, + source_vae: nn.Module | None = None, + ) -> None: + super().__init__() + self.register_modules( + transformer=transformer, + vae=vae, + scheduler=scheduler, + latent_upscaler=latent_upscaler, + source_vae=source_vae, + ) + transformer_config = getattr(transformer, "config", None) + self._sr_params = _coerce_sr_params(getattr(transformer_config, "sr_params", None)) + if self._sr_params is None: + raise ValueError("SR transformer config must contain 'sr_params'") + + @staticmethod + def _seed_from_generator( + generator: torch.Generator | None, + device: torch.device, + ) -> int: + if generator is None: + return int(torch.randint(0, 2**31, (1,), device=device).item()) + generator_device = torch.device(getattr(generator, "device", "cpu")) + return int( + torch.randint( + 0, + 2**31, + (1,), + generator=generator, + device=generator_device, + ).item() + ) + + def check_inputs( + self, + video: SRBatchInput | None = None, + latents: Tensor | np.ndarray | str | Path | Sequence[Any] | None = None, + *, + resolution_scale: float, + num_inference_steps: int | None = None, + overlap: float | None = None, + tiles_batch_size: int | None = None, + kvae_bridge: bool = False, + save_path: str | Path | list[str | Path] | None = None, + fps: int | None = None, + audio: list[np.ndarray] | np.ndarray | None = None, + source_video: str | Path | list[str | Path] | None = None, + audio_sample_rate: int = 44100, + output_type: SROutputType = "pt", + ) -> None: + """Validate SR input, scale, tiling, audio, and output arguments. + + Args: + video: Pixel video input in a supported frame layout. + latents: Raw latent tensor or local latent file. + resolution_scale: Total output scale, including the 2.25x + fractional route. + num_inference_steps: Number of denoising steps. + overlap: Tile overlap fraction in ``[0, 1)``. + tiles_batch_size: Number of tiles processed together. + save_path: Optional output video path. + fps: Optional output video frame rate. + audio: Optional audio waveform to mux. + source_video: Optional source video whose audio is copied. + audio_sample_rate: Sample rate used when muxing audio. + output_type: ``pt``/``torch`` or ``np``/``numpy``. + + Raises: + ValueError: If inputs conflict or a value is outside the SR + pipeline's supported range. + """ + if (video is None) == (latents is None): + raise ValueError("pass exactly one of video or latents") + if output_type not in ("pt", "torch", "np", "numpy"): + raise ValueError(f"unsupported output_type={output_type!r}") + if audio is not None and source_video is not None: + raise ValueError("pass either audio or source_video, not both") + if audio_sample_rate <= 0: + raise ValueError("audio_sample_rate must be positive") + if fps is not None and fps <= 0: + raise ValueError("fps must be positive") + + requested_scale = float(resolution_scale) + resolve_scale_request(requested_scale) + effective_steps = 5 if num_inference_steps is None else int(num_inference_steps) + effective_overlap = 0.25 if overlap is None else float(overlap) + effective_tiles_batch_size = 1 if tiles_batch_size is None else int(tiles_batch_size) + if effective_steps < 1: + raise ValueError("num_inference_steps must be positive") + if not 0 <= effective_overlap < 1: + raise ValueError("overlap must be in [0, 1)") + if effective_tiles_batch_size < 1: + raise ValueError("tiles_batch_size must be positive") + if latents is not None: + if requested_scale == 2.25 and not kvae_bridge: # noqa: PLR2004 + raise ValueError( + "2.25x latent SR requires kvae_bridge=True and source_vae; " + "ordinary latent input supports only x2/x4" + ) + if kvae_bridge and self.source_vae is None: + raise ValueError("KVAE bridge latent SR requires source_vae") + + @staticmethod + def _load_video(value: SRInput) -> tuple[Tensor, float | None, Path | None]: + if isinstance(value, (str, Path)): + path = Path(value) + video, source_fps = read_video_tchw_uint8(path) + video, output_fps = resample_to_target_fps(video, source_fps) + return clip_to_aligned_frames(video), float(output_fps), path + return _to_tchw(value), None, None + + @classmethod + def _load_videos(cls, value: SRBatchInput) -> tuple[list[Tensor], list[float | None], list[Path | None]]: + if isinstance(value, (Tensor, np.ndarray)) and value.ndim == VIDEO_BATCH_RANK: + values: Sequence[Any] = _video_batch_tensor(value) + elif _is_video_batch(value): + values = value + else: + values = [value] + loaded = [cls._load_video(item) for item in values] + videos = [item[0] for item in loaded] + _validate_batch_shapes(videos, "video") + return videos, [item[1] for item in loaded], [item[2] for item in loaded] + + @staticmethod + def _format_frames(frames: Tensor, output_type: SROutputType) -> Tensor | np.ndarray: + if output_type in ("pt", "torch"): + return frames + if output_type in ("np", "numpy"): + return frames.detach().cpu().numpy() + raise ValueError(f"Unsupported output_type={output_type!r}; use 'pt'/'torch' or 'np'/'numpy'") + + @torch.no_grad() + def _run( # noqa: PLR0913 + self, + *, + video: list[Tensor] | None = None, + latents: list[Tensor] | None = None, + resolution_scale: float, + num_inference_steps: int | None = None, + generator: torch.Generator | None = None, + overlap: float | None = None, + tiles_batch_size: int | None = None, + kvae_bridge: bool = False, + cached_text_embeds: dict[str, Tensor] | None = None, + save_path: str | Path | list[str | Path] | None = None, + fps: int | None = None, + audio: list[np.ndarray] | np.ndarray | None = None, + source_video: str | Path | list[str | Path] | None = None, + audio_sample_rate: int = 44100, + output_type: SROutputType = "pt", + return_dict: bool = True, + ) -> Kandinsky6SRPipelineOutput | tuple[Any, ...]: + requested_scale = float(resolution_scale) + effective_steps = 5 if num_inference_steps is None else int(num_inference_steps) + effective_overlap = 0.25 if overlap is None else float(overlap) + effective_tiles_batch_size = 1 if tiles_batch_size is None else int(tiles_batch_size) + self.check_inputs( + video=video, + latents=latents, + resolution_scale=requested_scale, + num_inference_steps=effective_steps, + overlap=effective_overlap, + tiles_batch_size=effective_tiles_batch_size, + kvae_bridge=kvae_bridge, + save_path=save_path, + fps=fps, + audio=audio, + source_video=source_video, + audio_sample_rate=audio_sample_rate, + output_type=output_type, + ) + tiling_scale, pre_upscale = resolve_scale_request(requested_scale) + components = SRComponents( + dit=self.transformer, + vae=self.vae, + latent_upscaler=self.latent_upscaler, + sr_params=self._sr_params, + cached_text_embeds=cached_text_embeds, + spatial_factor=getattr(self.vae, "spatial_factor", None), + ) + + batch_size = len(latents) if latents is not None else len(video) if video is not None else 0 + save_paths = _batch_output_values(save_path, batch_size, "save_path") + source_videos = _batch_output_values(source_video, batch_size, "source_video") + sr_video = video + try: + if latents is not None and kvae_bridge: + sr_video = [ + _decode_source_latent_video(latent, self.source_vae, self._execution_device) for latent in latents + ] + latents = None + if sr_video is not None and pre_upscale != 1.0: + sr_video = [ + pre_upscale_video(video_item, pre_upscale, _spatial_factor(components)) for video_item in sr_video + ] + + run_config = RunConfig( + device=str(self._execution_device), + num_steps=effective_steps, + seed=self._seed_from_generator(generator, self._execution_device), + overlap=effective_overlap, + tiles_batch_size=effective_tiles_batch_size * batch_size, + resolution_scale=tiling_scale, + ) + progress_input = latents[0] if latents is not None else sr_video[0] if sr_video is not None else None + if progress_input is None: + raise RuntimeError("SR input disappeared during preparation") + progress_total, progress_steps, progress_tiles = _sr_progress_plan( + progress_input, + components, + run_config, + latent_input=latents is not None, + batch_size=batch_size, + scheduler=self.scheduler, + ) + with self.progress_bar(total=progress_total) as progress_bar: + progress_bar.set_description(f"SR [{progress_steps} steps x {progress_tiles} tiles]") + has_latent_upscaler = ( + latent_upscaler_for_scale(self.latent_upscaler, run_config.resolution_scale) is not None + ) + if latents is not None: + frames = _run_batched_tiles( + latents, + components, + run_config, + latent_input=True, + progress_bar=progress_bar, + scheduler=self.scheduler, + ) + elif has_latent_upscaler: + lr_latents = list( + encode_lq_videos_to_lr_latents( + sr_video, + self.vae, + self._execution_device, + ) + ) + frames = _run_batched_tiles( + lr_latents, + components, + run_config, + latent_input=True, + progress_bar=progress_bar, + scheduler=self.scheduler, + ) + else: + frames = _run_batched_tiles( + sr_video, + components, + run_config, + latent_input=False, + progress_bar=progress_bar, + scheduler=self.scheduler, + ) + finally: + # Diffusers expects custom pipelines to restore the offloaded + # modules after a call. This also keeps the next invocation from + # observing a partially resident model chain. + self.maybe_free_model_hooks() + + path: str | list[str] | None = None + if save_path is not None: + paths = [ + str( + mux_video_audio( + frames[index], + audio, + output_path, + fps=fps or int(getattr(self._sr_params, "fps", TARGET_FPS)), + audio_sample_rate=audio_sample_rate, + source_video=source_videos[index], + ) + ) + for index, output_path in enumerate(save_paths) + ] + path = paths if isinstance(save_path, list) else paths[0] + audio_source = None + if source_video is not None: + audio_source = ( + [str(value) for value in source_videos] if isinstance(source_video, list) else str(source_video) + ) + output = Kandinsky6SRPipelineOutput( + frames=self._format_frames(frames, output_type), + audio=None if audio is None else (audio if isinstance(audio, list) else [audio]), + path=path, + metadata={ + "resolution_scale": requested_scale, + "num_inference_steps": run_config.num_steps, + "overlap": run_config.overlap, + "tiles_batch_size": effective_tiles_batch_size, + **({"audio_source": audio_source} if audio_source is not None else {}), + }, + ) + if return_dict: + return output + return output.frames, output.audio, output.path, output.metadata + + @torch.no_grad() + def from_video( # noqa: PLR0913 + self, + video: SRBatchInput, + resolution_scale: float = 2.25, + num_inference_steps: int | None = None, + generator: torch.Generator | None = None, + overlap: float | None = None, + tiles_batch_size: int | None = None, + cached_text_embeds: dict[str, Tensor] | None = None, + save_path: str | Path | list[str | Path] | None = None, + fps: int | None = None, + audio: list[np.ndarray] | np.ndarray | None = None, + source_video: str | Path | list[str | Path] | None = None, + audio_sample_rate: int = 44100, + output_type: SROutputType = "pt", + return_dict: bool = True, + ) -> Kandinsky6SRPipelineOutput | tuple[Any, ...]: + """Super-resolve pixels and optionally copy audio from a source path.""" + frames, input_fps, input_paths = self._load_videos(video) + input_path = input_paths[0] if len(input_paths) == 1 else None + if source_video is None and save_path is not None: + if input_path is not None: + source_video = input_path + elif all(path is not None for path in input_paths): + source_video = [path for path in input_paths if path is not None] + with _execution_device_context(self._execution_device): + return self._run( + video=frames, + resolution_scale=resolution_scale, + num_inference_steps=num_inference_steps, + generator=generator, + overlap=overlap, + tiles_batch_size=tiles_batch_size, + cached_text_embeds=cached_text_embeds, + save_path=save_path, + fps=( + fps + if fps is not None + else round(input_fps[0]) + if len(input_fps) == 1 and input_fps[0] is not None + else None + ), + audio=audio, + source_video=source_video, + audio_sample_rate=audio_sample_rate, + output_type=output_type, + return_dict=return_dict, + ) + + @torch.no_grad() + def from_latents( # noqa: PLR0913 + self, + latents: Tensor | np.ndarray | str | Path | Sequence[Any], + *, + resolution_scale: float = 2.25, + num_inference_steps: int | None = None, + generator: torch.Generator | None = None, + overlap: float | None = None, + tiles_batch_size: int | None = None, + kvae_bridge: bool = False, + cached_text_embeds: dict[str, Tensor] | None = None, + save_path: str | Path | list[str | Path] | None = None, + fps: int | None = None, + audio: list[np.ndarray] | np.ndarray | None = None, + source_video: str | Path | list[str | Path] | None = None, + audio_sample_rate: int = 44100, + output_type: SROutputType = "pt", + return_dict: bool = True, + ) -> Kandinsky6SRPipelineOutput | tuple[Any, ...]: + """Super-resolve raw latent samples or a batch of local ``.pt`` files.""" + channels = int(getattr(self.transformer, "in_visual_dim", 16)) + latent_value = _latent_batch(latents, channels) + with _execution_device_context(self._execution_device): + return self._run( + latents=latent_value, + resolution_scale=resolution_scale, + num_inference_steps=num_inference_steps, + generator=generator, + overlap=overlap, + tiles_batch_size=tiles_batch_size, + kvae_bridge=kvae_bridge, + cached_text_embeds=cached_text_embeds, + save_path=save_path, + fps=fps, + audio=audio, + source_video=source_video, + audio_sample_rate=audio_sample_rate, + output_type=output_type, + return_dict=return_dict, + ) + + @torch.no_grad() + def __call__( + self, + video: SRBatchInput | None = None, + latents: Tensor | np.ndarray | str | Path | Sequence[Any] | None = None, + resolution_scale: float = 2.25, + num_inference_steps: int = 5, + generator: torch.Generator | None = None, + overlap: float = 0.25, + tiles_batch_size: int = 1, + kvae_bridge: bool = False, + cached_text_embeds: dict[str, Tensor] | None = None, + save_path: str | Path | list[str | Path] | None = None, + fps: int | None = None, + audio: list[np.ndarray] | np.ndarray | None = None, + source_video: str | Path | list[str | Path] | None = None, + audio_sample_rate: int = 44100, + output_type: SROutputType = "pt", + return_dict: bool = True, + ) -> Kandinsky6SRPipelineOutput | tuple[Any, ...]: + """Super-resolve one or more equal-size pixel videos or latent inputs. + + Args: + video: Pixel video in ``[T,C,H,W]``/``[T,H,W,C]`` form, a batch + in a supported rank-5 layout, or local video paths. + latents: Raw latent tensor, a supported rank-5 latent batch, or + local ``.pt`` files. + save_path: Optional output video path or one path per input. + fps: Optional output video frame rate. + audio: Optional audio waveform to mux. + source_video: Optional source video whose audio is copied. + audio_sample_rate: Sample rate used when muxing audio. + output_type: ``pt``/``torch`` or ``np``/``numpy``. + return_dict: Whether to return ``Kandinsky6SRPipelineOutput``. + + Returns: + ``Kandinsky6SRPipelineOutput`` or its tuple representation when + ``return_dict=False``. + """ + if (video is None) == (latents is None): + raise ValueError("pass exactly one of video or latents") + + if video is not None: + return self.from_video( + video, + resolution_scale=resolution_scale, + num_inference_steps=num_inference_steps, + generator=generator, + overlap=overlap, + tiles_batch_size=tiles_batch_size, + cached_text_embeds=cached_text_embeds, + save_path=save_path, + fps=fps, + audio=audio, + source_video=source_video, + audio_sample_rate=audio_sample_rate, + output_type=output_type, + return_dict=return_dict, + ) + + return self.from_latents( + latents, + resolution_scale=resolution_scale, + num_inference_steps=num_inference_steps, + generator=generator, + overlap=overlap, + tiles_batch_size=tiles_batch_size, + kvae_bridge=kvae_bridge, + cached_text_embeds=cached_text_embeds, + save_path=save_path, + fps=fps, + audio=audio, + source_video=source_video, + audio_sample_rate=audio_sample_rate, + output_type=output_type, + return_dict=return_dict, + ) # type: ignore[arg-type] + + +def _to_tchw(value: Tensor | np.ndarray | Sequence[Any]) -> Tensor: # noqa: PLR0912, PLR2004 + """Normalize common frame layouts to contiguous uint8 ``[T,C,H,W]``.""" + if isinstance(value, Tensor): + tensor = value.detach() + else: + if isinstance(value, Sequence) and not isinstance(value, np.ndarray): + value = np.stack([np.asarray(frame) for frame in value]) + tensor = torch.as_tensor(value) + if tensor.ndim == VIDEO_RANK: + if tensor.shape[-1] in (1, 3, 4): + tensor = tensor.permute(0, 3, 1, 2) + elif tensor.shape[1] not in (1, 3, 4): + raise ValueError("4D video must be [T,C,H,W] or [T,H,W,C]") + else: + raise ValueError(f"video sample must have rank 4, got shape {tuple(tensor.shape)}") + if tensor.is_floating_point(): + if tensor.numel() and float(tensor.detach().amax()) <= 1.0: + tensor = tensor * 255.0 + tensor = tensor.round() + return tensor.clamp(0, 255).to(torch.uint8).contiguous() + + +def _coerce_sr_params( + value: Any | None, +) -> Any | None: + if value is not None: + if isinstance(value, Mapping) and "sr" in value: + value = value["sr"] + if isinstance(value, Mapping): + values = dict(value) + if "scale_factor" in values and isinstance(values["scale_factor"], Mapping): + values["scale_factor"] = { + int(key): [float(item) for item in items] for key, items in values["scale_factor"].items() + } + if "visual_size" in values and isinstance(values["visual_size"], int): + values["visual_size"] = [values["visual_size"]] + return SimpleNamespace(**values) + return value + return None + + +def _decode_source_latent_video( + raw_latents: Tensor, + source_vae: nn.Module, + device: str | torch.device, +) -> Tensor: + """Decode a base-VAE latent to ``[T,C,H,W]`` pixels for the KVAE bridge.""" + if raw_latents.ndim != 4: # noqa: PLR2004 + raise ValueError(f"raw_latents must have rank 4 [T,C,H,W], got {tuple(raw_latents.shape)}") + latent_5d = raw_latents.permute(1, 0, 2, 3).unsqueeze(0).to(device=device) + latent_5d = cast_to_module_dtype(source_vae, latent_5d) + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=str(device).startswith("cuda")): + try: + decoded = _call_module_method(source_vae, "decode", latent_5d, alternative_fwd=True).sample + except TypeError as exc: + if "alternative_fwd" not in str(exc): + raise + decoded = _call_module_method(source_vae, "decode", latent_5d).sample + return ( + ((decoded.squeeze(0).float().clamp(-1, 1) + 1.0) * 127.5) + .round() + .clamp(0, 255) + .to(torch.uint8) + .permute(1, 0, 2, 3) + .cpu() + ) + + +__all__ = ["Kandinsky6SRPipeline"] diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py new file mode 100644 index 000000000000..f505fd0a8ec4 --- /dev/null +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -0,0 +1,1599 @@ +"""Kandinsky 6 TI2VA Diffusers pipeline.""" + +from __future__ import annotations + +import math +import warnings +from collections.abc import Callable +from contextlib import nullcontext +from copy import deepcopy +from dataclasses import dataclass +from pathlib import Path +from typing import Any, TypedDict + +import numpy as np +import torch +from ..pipeline_utils import DiffusionPipeline +from torch import Tensor, nn + +from .pipeline_output import Kandinsky6TI2VAPipelineOutput + + +@dataclass +class LatentBundle: + """Packed K6 latent state; audio is optional for a video-only call.""" + + video: Tensor | None + audio: Tensor | None + video_cu_seqlens: Tensor | None + audio_cu_seqlens: Tensor | None + + +class TextEmbeds(TypedDict): + """Text adapter output consumed by the K6 DiT.""" + + text_embeds: Tensor + pooled_embed: Tensor + + +_PROMPT_TEMPLATE = "\n".join( + [ + "<|im_start|>system\nYou are a promt engineer. Describe the video in detail.", + "Describe how the camera moves or shakes, describe the zoom and view angle, whether it follows the objects.", + "Describe the location of the video, main characters or objects and their action.", + "Describe the dynamism of the video and presented actions.", + "Name the visual style of the video: whether it is a professional footage, user generated content, " + "some kind of animation, video game or scren content.", + "Describe the visual effects, postprocessing and transitions if they are presented in the video.", + "Pay attention to the order of key actions shown in the scene.<|im_end|>", + "<|im_start|>user\n{}<|im_end|>", + ] +) +_QWEN_CROP_START = 129 +_CLIP_MAX_LENGTH = 77 +_T2VA_EXPANSION_INSTRUCTION = ( + "You are a prompt beautifier that transforms short user video+audio descriptions into rich, detailed English " + "prompts specifically optimized for video+audio generation models. Preserve any direct speech between " + "and exactly as written. If the prompt asks for speech without exact words, add suitable direct speech " + "between those tags. Put general audio descriptions between and . Describe the scene, " + "actions, camera motion, visual style, and audio in detail. Make the prompt dynamic. Answer only with the " + "expanded prompt.\n\n" + "Rewrite Prompt: {}" +) +_I2VA_EXPANSION_INSTRUCTION = ( + "You are a prompt beautifier that transforms a short user video+audio description and a provided reference " + "image into a rich, detailed English prompt optimized for video+audio generation. The reference image is the " + "ground truth for the initial scene. Keep every visible fact that you mention consistent with it. Do not mainly " + "describe the image: focus on the requested actions, motion, interactions, camera changes, speech, and audio, " + "explaining the changes relative to the initial image. Add only details compatible with the reference image; " + "never invent conflicting objects, identities, colors, locations, or actions. Preserve direct speech between " + " and exactly. Put general audio descriptions between and . Make the prompt dynamic " + "and describe how the scene evolves from the provided image. Answer only with the expanded prompt.\n\n" + "Rewrite Prompt: {}" +) + + +def _prepare_expansion_image(image: Any) -> Any: + if isinstance(image, (str, Path)): + from PIL import Image + + with Image.open(image) as pil_image: + return pil_image.convert("RGB").copy() + return image + + +def _build_video_input( + video: Tensor, + has_visual_cond: bool, + first_frames: Tensor | None, + video_cu_seqlens: Tensor | None, + visual_cond_scheme: str, +) -> Tensor: + """Append visual-conditioning channels for batched ``(B,T,H,W,C)`` video.""" + if not has_visual_cond: + return video + + cond = torch.zeros_like(video) + mask = torch.zeros((*video.shape[:-1], 1), dtype=video.dtype, device=video.device) + if first_frames is None: + return torch.cat([video, cond, mask], dim=-1) + if video_cu_seqlens is None: + raise ValueError(f"{visual_cond_scheme} requires video_cu_seqlens") + + first_frames = first_frames.to(device=video.device, dtype=video.dtype) + tail_cond = visual_cond_scheme == "tail_cond_first_frame" + if visual_cond_scheme == "i2v": + video[:, 0] = first_frames + elif tail_cond: + video[:, -1] = first_frames + elif visual_cond_scheme == "pretrain": + cond[:, 0] = first_frames + else: + raise ValueError(f"unknown visual_cond_scheme={visual_cond_scheme!r}") + mask[:, -1 if tail_cond else 0] = 1 + return torch.cat([video, cond, mask], dim=-1) + + +def _resolve_null_embeds( + null_text_embeds: TextEmbeds | list[TextEmbeds | None], + null_text_rope: Tensor | list[Tensor | None], +) -> tuple[Tensor, Tensor, Tensor | list[Tensor]]: + """Returns (text_embed, pooled_embed, text_rope) for the uncond DiT pass. + + When null_text_embeds is a list of two items (T2VA per-modality nulls), + returns list arguments so the DiT can apply different null text per modality. + """ + if isinstance(null_text_embeds, list): + # null_text_embeds[0] = video null, null_text_embeds[1] = audio null (may be None) + null_v = null_text_embeds[0] + null_a = null_text_embeds[1] + null_rope_v = null_text_rope[0] if isinstance(null_text_rope, list) else null_text_rope + null_rope_a = null_text_rope[1] if isinstance(null_text_rope, list) else None + + if null_a is not None and null_rope_a is not None: + # Different null text per modality — pass as list to DiT + text_embed = [null_v["text_embeds"], null_a["text_embeds"]] + pooled_embed = [null_v["pooled_embed"], null_a["pooled_embed"]] + rope = [null_rope_v, null_rope_a] + else: + text_embed = null_v["text_embeds"] + pooled_embed = null_v["pooled_embed"] + rope = null_rope_v + else: + text_embed = null_text_embeds["text_embeds"] + pooled_embed = null_text_embeds["pooled_embed"] + rope = null_text_rope + + return text_embed, pooled_embed, rope + + +def _raw_dit(dit: nn.Module) -> nn.Module: + inner = getattr(dit, "module", None) + return inner if inner is not None and hasattr(dit, "set_cache") else dit + + +def _rebuild_text_rope(dit: nn.Module, template: Tensor | list[Tensor]) -> Tensor | list[Tensor]: + """Rebuild text RoPE using the standalone Diffusers transformer names.""" + raw = getattr(dit, "module", None) + raw = raw if raw is not None and hasattr(raw, "set_cache") else dit + if isinstance(template, list): + if getattr(raw, "is_multimodal", False): + modules = [raw.video_text_rope_embeddings, raw.audio_text_rope_embeddings] + else: + modules = [raw.text_rope_embeddings] * len(template) + return [compute_rope1d(module, int(item.shape[0])) for module, item in zip(modules, template, strict=True)] + module = raw.video_text_rope_embeddings if getattr(raw, "is_multimodal", False) else raw.text_rope_embeddings + return compute_rope1d(module, int(template.shape[0])) + + +def apply_cfg(cond: Tensor, uncond: Tensor, guidance_weight: float) -> Tensor: + """Classifier-free guidance: uncond + w * (cond - uncond).""" + return uncond + guidance_weight * (cond - uncond) + + +@torch.no_grad() +def denoise_loop( # noqa: PLR0912, PLR0913, PLR0915 + bundle: LatentBundle, + dit: nn.Module, + text_embeds: TextEmbeds, + null_text_embeds: TextEmbeds | list[TextEmbeds | None], + visual_rope: Tensor | None, + audio_rope: Tensor | None, + text_rope: Tensor | list[Tensor], + null_text_rope: Tensor | list[Tensor | None], + num_steps: int, + guidance_weight: float, + scheduler: Any, + first_frames: Tensor | None = None, + visual_cond_scheme: str = "pretrain", + sparse_params: dict | None = None, + sample_video: bool = True, + sample_audio: bool = True, + *, + attention_mask: Tensor | None = None, + null_attention_mask: Tensor | None = None, + visual_token_type_ids: Tensor | None = None, + recompute_ropes_each_step: bool = False, + scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), + step_callback: Callable[[int, Tensor, LatentBundle], LatentBundle | Tensor | None] | None = None, + progress_bar: Any | None = None, +) -> LatentBundle: + """Run the native Euler loop and invoke a Diffusers step callback. + + ``step_callback`` is called after each Euler update with the zero-based + step index, the current scheduler timestep, and the mutable latent + bundle. It may return ``None``, a replacement ``LatentBundle``, or a + replacement video ``Tensor``. The callback is deliberately optional so + the generated pipeline has the same numerical behavior when it is not + supplied. + """ + video = bundle.video # (B, T, H, W, C) or (B, T, H, W, C+extra) for instruct + audio = bundle.audio # (B, A, audio_dim) or None + is_multimodal = video is not None and audio is not None + is_piflow = bool(getattr(scheduler, "is_piflow", False)) + guidance_epsilon = 1e-6 + if is_piflow and abs(guidance_weight - 1.0) > guidance_epsilon: + raise ValueError("PiflowScheduler requires guidance_weight=1.0") + + audio_scheduler = deepcopy(scheduler) if is_multimodal else scheduler + device = video.device if video is not None else audio.device + scheduler.set_timesteps(num_steps, device=device) + if audio_scheduler is not scheduler: + audio_scheduler.set_timesteps(num_steps, device=device) + step_timesteps = scheduler.timesteps + if step_timesteps.shape[0] != num_steps: + raise ValueError( + "Diffusers scheduler returned an unexpected number of timesteps: " + f"expected {num_steps}, got {step_timesteps.shape[0]}" + ) + + bs = ( + bundle.video_cu_seqlens.shape[0] - 1 + if bundle.video_cu_seqlens is not None + else bundle.audio_cu_seqlens.shape[0] - 1 + ) + + null_te, null_pe, null_rope = _resolve_null_embeds(null_text_embeds, null_text_rope) + + out_c: int | None = None # determined after first step, used to strip instruct channels + raw = _raw_dit(dit) + vis_shape = ( + (int(visual_rope.shape[0]), int(visual_rope.shape[1]), int(visual_rope.shape[2])) + if visual_rope is not None + else None + ) + scale = (float(scale_factor[0]), float(scale_factor[1]), float(scale_factor[2])) + tail_cond = visual_cond_scheme == "tail_cond_first_frame" + + def _cache_scope(name: str): + cache_context = getattr(dit, "cache_context", None) + return cache_context(name) if callable(cache_context) else nullcontext() + + def _guided(cond: Tensor, uncond: Tensor) -> Tensor: + return cond if is_piflow else apply_cfg(cond, uncond, guidance_weight) + + for step_index, t in enumerate(step_timesteps): + # Diffusers schedulers expose the model-scale timestep directly. + t_step: Tensor | list[Tensor] = t.unsqueeze(0).expand(bs) # (bs,) + + model_input_v = ( + _build_video_input( + video, + dit.visual_cond, + first_frames, + bundle.video_cu_seqlens, + visual_cond_scheme, + ) + if video is not None + else None + ) + + # Freeze one modality at t=0 for partial sampling (T2VA only) + if is_multimodal: + t_frozen = torch.zeros_like(t_step[0]).expand(bs) + if not sample_audio: + t_step = [t_step, t_frozen] + elif not sample_video: + t_step = [t_frozen, t_step] + + def _forward( + te, + pe, + rope, + attn_mask, + *, + _model_input_v=model_input_v, + _audio=audio, + _t_step=t_step, + ): + vr, ar, tr = visual_rope, audio_rope, rope + if recompute_ropes_each_step: + if vis_shape is not None: + vr = compute_visual_rope(raw.visual_rope_embeddings, vis_shape, scale) + if audio_rope is not None: + ar = compute_rope1d(raw.audio_rope_embeddings, int(audio_rope.shape[0])) + tr = _rebuild_text_rope(raw, rope) + return dit( + x_video=_model_input_v, + x_audio=_audio, + text_embed=te, + pooled_text_embed=pe, + time=_t_step, + visual_rope=vr, + audio_rope=ar, + text_rope=tr, + sparse_params=sparse_params, + attention_mask=attn_mask, + visual_token_type_ids=visual_token_type_ids, + ) + + with _cache_scope("cond"): + vel_cond = _forward( + text_embeds["text_embeds"], + text_embeds["pooled_embed"], + text_rope, + attention_mask, + ) + if abs(guidance_weight - 1.0) > guidance_epsilon: + with _cache_scope("uncond"): + vel_uncond = _forward(null_te, null_pe, null_rope, null_attention_mask) + else: + vel_uncond = vel_cond + + # Diffusers schedulers own the update for each modality. Multimodal + # runs use an independent scheduler state for audio because the + # standard scheduler advances its step index on every call. + if isinstance(vel_cond, tuple): + vel_v, vel_a = vel_cond + uvel_v, uvel_a = vel_uncond if vel_uncond is not vel_cond else (vel_v, vel_a) + if out_c is None: + out_c = video.shape[-1] if video is not None else vel_v.shape[-1] + + if video is not None and sample_video: + guided_v = _guided(vel_v, uvel_v) + video = scheduler.step(guided_v, t, video, return_dict=False)[0] + if tail_cond and first_frames is not None: + video[:, -1] = first_frames.to(device=device, dtype=video.dtype) + + if audio is not None and sample_audio: + guided_a = _guided(vel_a, uvel_a) + audio = audio_scheduler.step(guided_a, t, audio, return_dict=False)[0] + else: + if out_c is None: + out_c = video.shape[-1] if video is not None else vel_cond.shape[-1] + vel = _guided(vel_cond, vel_uncond) + if video is not None and sample_video: + video = scheduler.step(vel, t, video, return_dict=False)[0] + if tail_cond and first_frames is not None: + video[:, -1] = first_frames.to(device=device, dtype=video.dtype) + elif audio is not None and sample_audio: + audio = scheduler.step(vel, t, audio, return_dict=False)[0] + + if step_callback is not None: + callback_result = step_callback( + step_index, + t, + LatentBundle( + video=video, + audio=audio, + video_cu_seqlens=bundle.video_cu_seqlens, + audio_cu_seqlens=bundle.audio_cu_seqlens, + ), + ) + if isinstance(callback_result, LatentBundle): + video = callback_result.video + audio = callback_result.audio + elif isinstance(callback_result, Tensor): + video = callback_result + + if progress_bar is not None: + progress_bar.update() + + # I2V / I2VA: keep injected reference frames unchanged + if first_frames is not None and video is not None and bundle.video_cu_seqlens is not None: + ff = first_frames.to(device=device, dtype=video.dtype) + if visual_cond_scheme == "i2v": + video[:, 0] = ff + elif tail_cond: + video[:, -1] = ff + + # Strip instruct extra-channel padding from the video latent + if out_c is not None and video is not None: + video = video[..., :out_c] + + return LatentBundle( + video=video, + audio=audio, + video_cu_seqlens=bundle.video_cu_seqlens, + audio_cu_seqlens=bundle.audio_cu_seqlens, + ) + + +def prepare_video_latents( + bs: int, + duration: int, + H_lat: int, + W_lat: int, + C: int, + seed: int, + device: torch.device | str, + dtype: torch.dtype = torch.bfloat16, +) -> LatentBundle: + """Sample random video noise latent and compute cu_seqlens. + + Returns a LatentBundle with video=(bs*duration, H_lat, W_lat, C) and + uniform video_cu_seqlens=(bs+1,) int32. + """ + g = torch.Generator(device=device) + g.manual_seed(seed) + video = torch.randn(bs * duration, H_lat, W_lat, C, device=device, dtype=dtype, generator=g) + cu = duration * torch.arange(bs + 1, dtype=torch.int32, device=device) + return LatentBundle(video=video, audio=None, video_cu_seqlens=cu, audio_cu_seqlens=None) + + +def audio_latent_duration( + video_latent_frames: int, + *, + fps: float = 24.0, + audio_fps: int = 44100, + downsample_factor: int = 1024, +) -> int: + """Audio latent length matching K5 T2VA: ceil(sample_frames/fps * audio_fps / downsample).""" + sample_frames = (video_latent_frames - 1) * 4 + 1 + return int(math.ceil(sample_frames / fps * audio_fps / downsample_factor)) + + +def prepare_audio_latents( + bundle: LatentBundle, + audio_duration: int, + audio_dim: int, + seed: int, + device: torch.device | str, + dtype: torch.dtype = torch.bfloat16, +) -> LatentBundle: + """Attach random audio noise latent to an existing LatentBundle. + + audio_duration: number of latent audio frames per batch item. + """ + bs = bundle.video_cu_seqlens.shape[0] - 1 if bundle.video_cu_seqlens is not None else 1 + g = torch.Generator(device=device) + g.manual_seed(seed + 1) # offset from video seed for independence + audio = torch.randn(bs * audio_duration, audio_dim, device=device, dtype=dtype, generator=g) + cu = audio_duration * torch.arange(bs + 1, dtype=torch.int32, device=device) + return LatentBundle( + video=bundle.video, + audio=audio, + video_cu_seqlens=bundle.video_cu_seqlens, + audio_cu_seqlens=cu, + ) + + +def compute_rope1d(rope: nn.Module, length: int, device: torch.device | None = None) -> Tensor: + """Build 1-D RoPE and move non-persistent tables for Diffusers offload.""" + if length < 1: + raise ValueError(f"rope length must be positive, got {length}") + if device is None: + device = next(rope.buffers()).device + else: + device = torch.device(device) + # Accelerate's model CPU-offload hook only moves the transformer when + # its forward starts. RoPE is materialized before that forward, so + # explicitly move its non-persistent lookup table first. + rope.to(device) + pos = torch.arange(length, device=device) + return rope(pos) + + +def compute_visual_rope( + rope: nn.Module, + shape: tuple[int, int, int], + scale_factor: tuple[float, float, float], + device: torch.device | None = None, +) -> Tensor: + """Build 3-D RoPE and move non-persistent tables for Diffusers offload.""" + T, H, W = (int(shape[0]), int(shape[1]), int(shape[2])) + if T < 1 or H < 1 or W < 1: + raise ValueError(f"visual rope shape must be positive, got {shape}") + if device is None: + device = next(rope.buffers()).device + else: + device = torch.device(device) + rope.to(device) + pos = [ + torch.arange(T, device=device), + torch.arange(H, device=device), + torch.arange(W, device=device), + ] + scale = (float(scale_factor[0]), float(scale_factor[1]), float(scale_factor[2])) + return rope((T, H, W), pos, scale) + + +@torch.no_grad() +def postprocess_audio( + bundle: LatentBundle, + audio_vae, + normalization_mode: str = "clip", +) -> list[np.ndarray] | None: + """Decode audio latents → list of (samples,) int16 numpy arrays. + + Returns None when bundle.audio is None (T2V mode). + + ``clip`` preserves the decoded amplitude and saturates it to [-1, 1]. + ``normalize`` peak-normalizes each waveform before converting it to int16. + """ + audio = bundle.audio + if audio is None: + return None + + cu = bundle.audio_cu_seqlens + assert cu is not None + bs = cu.shape[0] - 1 + + # Reverse audio VAE normalization + audio_scaled = audio / getattr(audio_vae, "scaling_factor", 1.0) + audio_scaled = audio_scaled + getattr(audio_vae, "mean_value", 0.0) + + segments = [ + audio_scaled[i] if audio_scaled.ndim == 3 else audio_scaled[cu[i].item() : cu[i + 1].item()] for i in range(bs) + ] # (A, audio_dim) per sample + + # Native generation keeps equal-length samples in the batch. Decode them + # together so batching is preserved through the audio VAE and vocoder. + if len({tuple(segment.shape) for segment in segments}) == 1: + decoded = audio_vae.wrapped_decode(torch.stack(segments).transpose(1, 2)) + if decoded.ndim == 1: + decoded = decoded.unsqueeze(0) + elif decoded.ndim == 3 and decoded.shape[1] == 1: + decoded = decoded[:, 0] + waveforms = [waveform.cpu().float().numpy() for waveform in decoded] + else: + # Preserve support for legacy packed bundles with variable lengths. + waveforms = [ + audio_vae.wrapped_decode(segment.transpose(1, 0).unsqueeze(0)).squeeze().cpu().float().numpy() + for segment in segments + ] + + result: list[np.ndarray] = [] + for waveform in waveforms: + if normalization_mode == "normalize": + peak = np.max(np.abs(waveform)) + if peak > 0: + waveform = waveform / peak * 32767 + elif normalization_mode == "clip": + waveform = np.clip(waveform, -1.0, 1.0) * 32767 + else: + raise ValueError(f"unknown audio normalization_mode={normalization_mode}") + + result.append(waveform.astype(np.int16)) + + return result + + +@torch.no_grad() +def postprocess_video( + bundle: LatentBundle, + vae, + bs: int, +) -> Tensor: + """Decode video latents → (bs, 3, T, H, W) uint8 in [0, 255]. + + Input layout: packed ``(sum_T, H_lat, W_lat, C)`` or batched + ``(bs, T, H_lat, W_lat, C)``. + """ + video = bundle.video + assert video is not None + + if video.ndim == 4: + frames = video.reshape(bs, -1, video.shape[-3], video.shape[-2], video.shape[-1]) + elif video.ndim == 5: + if video.shape[0] != bs: + raise ValueError(f"batched video has batch size {video.shape[0]}, expected {bs}") + frames = video + else: + raise ValueError(f"video must have rank 4 or 5, got {video.ndim}") + # (bs, T, H, W, C) → (bs, C, T, H, W) for VAE input + frames = (frames / vae.config.scaling_factor).permute(0, 4, 1, 2, 3) + # Hunyuan VAE loads as fp16; DiT latents are bf16 — match weight dtype. + vae_dtype = next(vae.parameters()).dtype + frames = vae.decode(frames.to(dtype=vae_dtype)).sample + + return ((frames.clamp(-1.0, 1.0) + 1.0) * 127.5).to(torch.uint8) + + +@torch.no_grad() +def resize_image( + image: Tensor, + max_area: int, + divisibility: int = 16, + world_size: int = 1, +) -> tuple[Tensor, float]: + """Aspect-preserving resize to fit ``max_area`` with sides divisible by ``divisibility``. + + K5 ``i2v_pipeline.resize_image`` parity. ``image`` is ``(B, C, H, W)``. + Returns ``(resized, scale_k)``. + """ + from math import sqrt + + h, w = image.shape[2:] + area = h * w + div = divisibility + if div == 16: + if world_size in (2, 4): + div *= 2 + elif world_size == 8: + div *= 4 + k = sqrt(max_area / area) / div + new_h = int(round(h * k) * div) + new_w = int(round(w * k) * div) + import torchvision.transforms.functional as TF + + return TF.resize(image, (new_h, new_w)), k + + +def _load_pil_rgb(image: str | object): + from PIL import Image + + if isinstance(image, str): + try: + pil_image = Image.open(image).convert("RGB") + pil_image.load() + except Exception as exc: + raise ValueError(f"Cannot decode i2va input image {image!r}: {exc}") from exc + elif isinstance(image, Image.Image): + pil_image = image.convert("RGB") + else: + raise TypeError(f"i2va image must be a path or PIL image, got {type(image).__name__}") + return pil_image + + +@torch.no_grad() +def encode_i2va_first_frame( + image: str | object, + vae, + device: torch.device | str, + height: int | None = None, + width: int | None = None, + *, + max_area: int | None = None, + divisibility: int = 16, + world_size: int = 1, +) -> tuple[Tensor, int, int]: + """Encode one image into the K6 packed first-frame latent layout.""" + import torchvision.transforms.functional as TF + + pil_image = _load_pil_rgb(image) + tensor = TF.pil_to_tensor(pil_image).unsqueeze(0) + if max_area is not None and (height is None or width is None): + tensor, _ = resize_image(tensor, max_area, divisibility, world_size) + height, width = int(tensor.shape[-2]), int(tensor.shape[-1]) + elif height is not None and width is not None: + src_h, src_w = tensor.shape[-2:] + scale = min(src_h / height, src_w / width) + tensor = TF.resize(tensor, (int(src_h / scale), int(src_w / scale))) + cur_h, cur_w = tensor.shape[-2:] + tensor = TF.crop(tensor, (cur_h - height) // 2, (cur_w - width) // 2, height, width) + else: + raise ValueError("encode_i2va_first_frame needs height/width or max_area") + + tensor = tensor.to(device=device, dtype=next(vae.parameters()).dtype) / 127.5 - 1.0 + tensor = tensor.transpose(0, 1).unsqueeze(0) + encoded = vae.encode(tensor) + posterior = encoded.latent_dist if hasattr(encoded, "latent_dist") else encoded[0] + latent = posterior.sample() + latent = latent.squeeze(0).permute(1, 2, 3, 0) + return latent * vae.config.scaling_factor, int(height), int(width) + + +def append_i2va_tail_condition( + latent_visual: Tensor, + first_frames: Tensor, + batch_size: int, + video_duration: int, +) -> tuple[Tensor, Tensor, Tensor]: + """Append one clean reference frame to batched ``(B,T,H,W,C)`` latents.""" + if latent_visual.shape[0] != batch_size or latent_visual.shape[1] != video_duration: + raise ValueError( + "generated visual latent shape mismatch: expected " + f"({batch_size}, {video_duration}, H, W, C), got {tuple(latent_visual.shape)}" + ) + + _, _, height, width, dim = latent_visual.shape + first_frames = first_frames.to(device=latent_visual.device, dtype=latent_visual.dtype) + expected = (batch_size, height, width, dim) + if tuple(first_frames.shape) != expected: + raise ValueError(f"first-frame latent shape mismatch: expected {expected}, got {tuple(first_frames.shape)}") + + latent_visual = torch.cat([latent_visual, first_frames[:, None]], dim=1) + token_types = torch.cat( + [ + torch.zeros((batch_size, video_duration), dtype=torch.long, device=latent_visual.device), + torch.ones((batch_size, 1), dtype=torch.long, device=latent_visual.device), + ], + dim=1, + ) + return latent_visual, token_types, token_types == 0 + + +class Kandinsky6TI2VAPipeline(DiffusionPipeline): + r"""Pipeline for text/image-to-video-and-audio generation with Kandinsky 6. + + This pipeline inherits the generic loading, device placement, and progress + handling provided by ``DiffusionPipeline``. The transformer must expose + the K6 multimodal forward contract, while the VAE adapters must implement + the portable video/audio post-processing contracts. + + Args: + transformer: Multimodal K6 transformer used to denoise video and audio + latents. + vae: Video VAE used to decode generated video latents. + text_encoder: Qwen2.5-VL text encoder for token-level embeddings. + audio_vae: Audio VAE used to decode generated audio latents. It may be + ``None`` only when ``sample_audio=False``. + scheduler: Diffusion scheduler used by the denoising loop. + tokenizer: Qwen2.5-VL processor. + text_encoder_2: CLIP text encoder for pooled embeddings. + tokenizer_2: CLIP tokenizer. + """ + + model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae->audio_vae" + _callback_tensor_inputs = ["latents"] + _DEFAULT_NEGATIVE_PROMPT = ( + "Static, 2D cartoon, cartoon, 2d animation, paintings, images, " + "worst quality, low quality, ugly, deformed, walking backwards" + ) + + def __init__( + self, + transformer: nn.Module, + vae: nn.Module, + text_encoder: Any, + audio_vae: nn.Module | None, + scheduler: Any, + tokenizer: Any, + text_encoder_2: Any, + tokenizer_2: Any, + ) -> None: + super().__init__() + self.register_modules( + transformer=transformer, + vae=vae, + text_encoder=text_encoder, + audio_vae=audio_vae, + tokenizer=tokenizer, + text_encoder_2=text_encoder_2, + tokenizer_2=tokenizer_2, + scheduler=scheduler, + ) + transformer_config = self.transformer.config + scale_factor = transformer_config.get("scale_factor", (1.0, 2.0, 2.0)) + text_token_padding = transformer_config.get("text_token_padding", True) + self.scale_factor = tuple(float(value) for value in scale_factor) + self.text_token_padding = bool(text_token_padding) + + audio_config = self.audio_vae.config if self.audio_vae is not None else {} + self.audio_sample_rate = int(audio_config.get("sample_rate", 44_100)) + self.audio_downsample_factor = int(audio_config.get("downsample_factor", 1_024)) + + @staticmethod + def _value_batch_size(value: Any) -> int: + if value is None: + return 1 + if isinstance(value, list): + return len(value) + if isinstance(value, Tensor): + return int(value.shape[0]) if value.ndim in (3, 5) else 1 + if isinstance(value, dict) and isinstance(value.get("text_embeds"), Tensor): + embeds = value["text_embeds"] + return int(embeds.shape[0]) if embeds.ndim == 3 else 1 + return 1 + + @classmethod + def _batch_size( + cls, + prompt: str | list[str] | None, + negative_prompt: str | list[str] | None, + image: str | object | list[str | object] | None, + latents: Tensor | None, + audio_latents: Tensor | None, + prompt_embeds: TextEmbeds | None, + negative_prompt_embeds: TextEmbeds | None, + ) -> int: + sizes = [ + cls._value_batch_size(value) + for value in ( + prompt, + negative_prompt, + image, + latents, + audio_latents, + prompt_embeds, + negative_prompt_embeds, + ) + ] + batch_size = max(sizes) + if batch_size < 1 or any(size not in (1, batch_size) for size in sizes): + raise ValueError(f"all batched inputs must have the same batch size, got {sizes}") + return batch_size + + @staticmethod + def _as_batch(value: Any, batch_size: int, name: str) -> list[Any] | None: + if value is None: + return None + values = value if isinstance(value, list) else [value] + if len(values) == 1: + return values * batch_size + if len(values) != batch_size: + raise ValueError(f"{name} must contain one item or {batch_size} items, got {len(values)}") + return values + + def check_inputs( + self, + prompt: str | list[str] | None, + negative_prompt: str | list[str] | None, + height: int, + width: int, + num_frames: int, + num_inference_steps: int, + max_sequence_length: int, + output_type: str, + prompt_embeds: TextEmbeds | None = None, + negative_prompt_embeds: TextEmbeds | None = None, + prompt_cu_seqlens: Tensor | None = None, + negative_prompt_cu_seqlens: Tensor | None = None, + callback_on_step_end_tensor_inputs: list[str] | None = None, + sample_audio: bool = True, + image: str | object | None = None, + visual_cond_scheme: str = "pretrain", + ) -> None: + """Validate arguments shared by text-to-video and image-to-video calls. + + Args: + prompt: Text prompt or a batch of prompts. + image: Optional reference image or batch of reference images. + negative_prompt: Negative text prompt or a batch of prompts. + height: Requested output height in pixels. + width: Requested output width in pixels. + num_frames: Requested output frame count. + num_inference_steps: Number of denoising steps. + max_sequence_length: Maximum Qwen prompt length after the template. + output_type: One of ``pt``, ``torch``, ``np``, ``numpy``, or + ``latent``. + prompt_embeds: Optional precomputed positive prompt embeddings. + negative_prompt_embeds: Optional precomputed negative embeddings. + prompt_cu_seqlens: Cumulative positive Qwen sequence lengths. + negative_prompt_cu_seqlens: Cumulative negative Qwen sequence + lengths. + callback_on_step_end_tensor_inputs: Tensor names exposed to the + step callback. + sample_audio: Whether to generate and decode audio. + visual_cond_scheme: Image-conditioning scheme when ``image`` is set. + + Raises: + ValueError: If an input combination is unsupported. + """ + for name, value in (("prompt", prompt), ("negative_prompt", negative_prompt)): + if value is not None and not isinstance(value, (str, list)): + raise ValueError(f"`{name}` has to be a `str` or a `list`, but is {type(value)}") + if isinstance(value, list) and (not value or not all(isinstance(item, str) for item in value)): + raise ValueError(f"`{name}` must be a non-empty list of strings") + + if prompt is None and prompt_embeds is None: + raise ValueError("provide either `prompt` or `prompt_embeds`") + if prompt is not None and prompt_embeds is not None: + raise ValueError("provide either `prompt` or `prompt_embeds`, not both") + if (prompt_embeds is None) != (prompt_cu_seqlens is None): + raise ValueError("`prompt_embeds` and `prompt_cu_seqlens` must be provided together") + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError("provide either `negative_prompt` or `negative_prompt_embeds`, not both") + if (negative_prompt_embeds is None) != (negative_prompt_cu_seqlens is None): + raise ValueError("`negative_prompt_embeds` and `negative_prompt_cu_seqlens` must be provided together") + if prompt_embeds is not None: + self._as_text_embeds(prompt_embeds) + if negative_prompt_embeds is not None: + self._as_text_embeds(negative_prompt_embeds) + + if num_frames < 1: + raise ValueError("num_frames must be positive") + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be positive") + if height < 1 or width < 1: + raise ValueError("height and width must be positive") + if height % 8 != 0 or width % 8 != 0: + raise ValueError("height and width must be divisible by 8") + if output_type not in ("pt", "torch", "np", "numpy", "latent"): + raise ValueError("output_type must be 'pt', 'torch', 'np', 'numpy', or 'latent'") + if max_sequence_length < 1 or max_sequence_length > 1024: + raise ValueError("max_sequence_length must be between 1 and 1024") + if sample_audio and self.audio_vae is None: + raise ValueError("sample_audio=True requires an audio_vae") + + if getattr(self.transformer, "is_multimodal", False) is not True: + raise ValueError("Kandinsky6TI2VAPipeline requires a multimodal transformer") + patch_size = tuple(int(value) for value in getattr(self.transformer, "patch_size", (1, 2, 2))) + if len(patch_size) < 3 or any(value <= 0 for value in patch_size): + raise ValueError(f"transformer.patch_size must contain three positive values, got {patch_size}") + if (height // 8) % patch_size[1] != 0 or (width // 8) % patch_size[2] != 0: + raise ValueError( + "height and width must produce whole DiT spatial patches; " + f"got latent size {(height // 8, width // 8)} and patch size {patch_size}" + ) + if callback_on_step_end_tensor_inputs is not None and not all( + name in self._callback_tensor_inputs for name in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, " + f"but found {callback_on_step_end_tensor_inputs}" + ) + if visual_cond_scheme not in ("pretrain", "i2v", "tail_cond_first_frame"): + raise ValueError(f"unknown visual_cond_scheme={visual_cond_scheme!r}") + if image is not None: + if not getattr(self.transformer, "visual_cond", False): + raise ValueError("image conditioning requires transformer.visual_cond=True") + if ( + visual_cond_scheme == "tail_cond_first_frame" + and getattr(self.transformer, "visual_token_type_num_embeddings", 0) < 2 + ): + raise ValueError("tail_cond_first_frame requires transformer.visual_token_type_num_embeddings >= 2") + + @staticmethod + def _as_text_embeds(value: Any) -> TextEmbeds: + if not isinstance(value, dict): + raise TypeError("prompt embeddings must be a mapping with text_embeds and pooled_embed") + missing = {"text_embeds", "pooled_embed"} - set(value) + if missing: + raise ValueError(f"prompt embeddings are missing: {sorted(missing)}") + return value + + def encode_prompt( + self, text: str | list[str], max_sequence_length: int + ) -> tuple[TextEmbeds, Tensor, Tensor | None]: + texts = [text] if isinstance(text, str) else text + full_texts = [_PROMPT_TEMPLATE.format(item) for item in texts] + qwen_device = next(self.text_encoder.parameters()).device + inputs = self.tokenizer( + text=full_texts, + images=None, + videos=None, + max_length=max_sequence_length + _QWEN_CROP_START, + truncation=True, + return_tensors="pt", + padding="max_length", + ).to(qwen_device) + qwen_output = self.text_encoder( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + return_dict=True, + output_hidden_states=True, + ) + embeds = qwen_output["hidden_states"][-1][:, _QWEN_CROP_START:] + attention = inputs["attention_mask"][:, _QWEN_CROP_START:].to(dtype=torch.bool) + batch_size = len(texts) + if self.text_token_padding: + qwen_embeds = embeds if batch_size > 1 else embeds[0] + qwen_attention = attention if batch_size > 1 else attention[0] + lengths = torch.full((batch_size,), embeds.shape[1], dtype=torch.int32, device=embeds.device) + cu_seqlens = torch.cat([torch.zeros(1, dtype=torch.int32, device=embeds.device), lengths.cumsum(0)]) + elif batch_size > 1: + max_tokens = int(attention.sum(dim=1).max().item()) + qwen_embeds = embeds[:, :max_tokens] + qwen_attention = attention[:, :max_tokens] + lengths = torch.full((batch_size,), max_tokens, dtype=torch.int32, device=embeds.device) + cu_seqlens = torch.cat([torch.zeros(1, dtype=torch.int32, device=embeds.device), lengths.cumsum(0)]) + else: + qwen_embeds = embeds[attention] + qwen_attention = None + lengths = attention.sum(dim=1).to(torch.int32) + cu_seqlens = torch.cat([torch.zeros(1, dtype=torch.int32, device=embeds.device), lengths.cumsum(0)]) + + clip_device = next(self.text_encoder_2.parameters()).device + clip_inputs = self.tokenizer_2( + texts, + max_length=_CLIP_MAX_LENGTH, + truncation=True, + add_special_tokens=True, + padding="max_length", + return_tensors="pt", + ).to(clip_device) + pooled_embed = self.text_encoder_2(**clip_inputs)["pooler_output"] + return ( + { + "text_embeds": qwen_embeds, + "pooled_embed": pooled_embed, + }, + cu_seqlens, + qwen_attention, + ) + + @torch.no_grad() + def expand_prompts( + self, + prompt: str | list[str], + image: Any | list[Any] | None = None, + max_sequence_length: int = 1024, + mode: str | None = None, + ) -> str | list[str]: + """Expand T2VA prompts, or image-grounded I2VA prompts when an image is supplied.""" + if isinstance(prompt, list): + images = image if isinstance(image, (list, tuple)) else [image] * len(prompt) + return [ + self.expand_prompts( + item, + image=item_image, + max_sequence_length=max_sequence_length, + mode=mode, + ) + for item, item_image in zip(prompt, images, strict=True) + ] + mode = mode or ("i2va" if image is not None else "t2va") + if mode not in ("t2va", "i2va"): + raise ValueError(f"unknown prompt expansion mode={mode!r}") + image = _prepare_expansion_image(image) if image is not None else None + instruction = (_I2VA_EXPANSION_INSTRUCTION if mode == "i2va" else _T2VA_EXPANSION_INSTRUCTION).format(prompt) + content = [] + if image is not None: + content.append({"type": "image", "image": image}) + content.append({"type": "text", "text": instruction}) + messages = [ + { + "role": "user", + "content": content, + } + ] + text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + device = next(self.text_encoder.parameters()).device + inputs = self.tokenizer( + text=[text], + images=[image] if image is not None else None, + videos=None, + padding=True, + return_tensors="pt", + ).to(device) + generated = self.text_encoder.generate(**inputs, max_new_tokens=max_sequence_length) + qwen_crop_start = inputs["input_ids"].shape[1] + trimmed = [output[qwen_crop_start:] for output in generated] + return self.tokenizer.batch_decode(trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + + def _encode_prompt( + self, + prompt: str | list[str], + negative_prompt: str | list[str], + *, + prompt_embeds: TextEmbeds | None, + negative_prompt_embeds: TextEmbeds | None, + prompt_cu_seqlens: Tensor | None, + negative_prompt_cu_seqlens: Tensor | None, + max_sequence_length: int, + ) -> tuple[TextEmbeds, Tensor, Tensor | None, TextEmbeds, Tensor, Tensor | None]: + + if prompt_embeds is None: + prompt_embeds, prompt_cu_seqlens, prompt_attention_mask = self.encode_prompt(prompt, max_sequence_length) + elif prompt_cu_seqlens is None: + raise ValueError("prompt_cu_seqlens is required with prompt_embeds") + else: + prompt_attention_mask = None + prompt_embeds = self._as_text_embeds(prompt_embeds) + + if negative_prompt_embeds is None: + negative_prompt_embeds, negative_prompt_cu_seqlens, negative_attention_mask = self.encode_prompt( + negative_prompt, max_sequence_length + ) + elif negative_prompt_cu_seqlens is None: + raise ValueError("negative_prompt_cu_seqlens is required with negative_prompt_embeds") + else: + negative_attention_mask = None + negative_prompt_embeds = self._as_text_embeds(negative_prompt_embeds) + return ( + prompt_embeds, + prompt_cu_seqlens, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_cu_seqlens, + negative_attention_mask, + ) + + @staticmethod + def _seed_from_generator( + generator: torch.Generator | None, + device: torch.device, + ) -> int: + if generator is None: + return int(torch.randint(0, 2**31, (1,), device=device).item()) + generator_device = torch.device(getattr(generator, "device", "cpu")) + return int( + torch.randint( + 0, + 2**31, + (1,), + generator=generator, + device=generator_device, + ).item() + ) + + @staticmethod + def _coerce_video_latents( + latents: Tensor, + *, + latent_frames: int, + height: int, + width: int, + channels: int, + batch_size: int, + device: torch.device, + dtype: torch.dtype, + ) -> LatentBundle: + value = latents.to(device=device, dtype=dtype) + if value.ndim == 4: + value = value.unsqueeze(0) + if value.ndim != 5: + raise ValueError("video latents must have shape (T,H,W,C), (B,T,H,W,C), or (B,C,T,H,W)") + if value.shape[0] != batch_size: + raise ValueError(f"latents batch size must be {batch_size}, got {value.shape[0]}") + expected = (batch_size, latent_frames, height // 8, width // 8, channels) + bchtw = (batch_size, channels, latent_frames, height // 8, width // 8) + if tuple(value.shape) == bchtw: + value = value.permute(0, 2, 3, 4, 1) + elif tuple(value.shape) != expected: + raise ValueError("5-D latents must use BCHTW or BTHWC layout with the configured channel count") + cu_seqlens = latent_frames * torch.arange(batch_size + 1, dtype=torch.int32, device=device) + return LatentBundle( + video=value, + audio=None, + video_cu_seqlens=cu_seqlens, + audio_cu_seqlens=None, + ) + + @staticmethod + def _coerce_audio_latents( + bundle: LatentBundle, + latents: Tensor, + *, + audio_duration: int, + audio_dim: int, + batch_size: int, + device: torch.device, + dtype: torch.dtype, + ) -> LatentBundle: + value = latents.to(device=device, dtype=dtype) + if value.ndim == 2: + value = value.unsqueeze(0) + if value.ndim != 3: + raise ValueError("audio latents must have shape (A,D) or (B,A,D)") + if value.shape[0] != batch_size: + raise ValueError(f"audio_latents batch size must be {batch_size}, got {value.shape[0]}") + expected = (batch_size, audio_duration, audio_dim) + if tuple(value.shape) != expected: + raise ValueError(f"audio latent shape mismatch: expected {expected}, got {tuple(value.shape)}") + cu_seqlens = audio_duration * torch.arange(batch_size + 1, dtype=torch.int32, device=device) + return LatentBundle( + video=bundle.video, + audio=value, + video_cu_seqlens=bundle.video_cu_seqlens, + audio_cu_seqlens=cu_seqlens, + ) + + def _prepare_latents( + self, + *, + latents: Tensor | None, + audio_latents: Tensor | None, + latent_frames: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: torch.Generator | None, + sample_fps: float, + sample_audio: bool, + batch_size: int, + ) -> LatentBundle: + channels = int(getattr(self.transformer, "in_visual_dim", 16)) + seed = self._seed_from_generator(generator, device) + if latents is None: + bundle = prepare_video_latents( + bs=batch_size, + duration=latent_frames, + H_lat=height // 8, + W_lat=width // 8, + C=channels, + seed=seed, + device=device, + dtype=dtype, + ) + bundle.video = bundle.video.reshape(batch_size, latent_frames, height // 8, width // 8, channels) + else: + bundle = self._coerce_video_latents( + latents, + latent_frames=latent_frames, + height=height, + width=width, + channels=channels, + batch_size=batch_size, + device=device, + dtype=dtype, + ) + + if not sample_audio: + if audio_latents is not None: + raise ValueError("audio_latents cannot be supplied when sample_audio=False") + return bundle + if self.audio_vae is None: + raise ValueError("sample_audio=True requires an audio_vae") + + audio_dim = int(getattr(self.transformer, "in_audio_dim", 20)) + downsample_factor = int(self.audio_vae.downsample_factor) + audio_duration = audio_latent_duration( + latent_frames, + fps=sample_fps, + audio_fps=self.audio_sample_rate, + downsample_factor=downsample_factor, + ) + if audio_latents is None: + bundle = prepare_audio_latents( + bundle, + audio_duration=audio_duration, + audio_dim=audio_dim, + seed=seed, + device=device, + dtype=dtype, + ) + bundle.audio = bundle.audio.reshape(batch_size, audio_duration, audio_dim) + return bundle + return self._coerce_audio_latents( + bundle, + audio_latents, + audio_duration=audio_duration, + audio_dim=audio_dim, + batch_size=batch_size, + device=device, + dtype=dtype, + ) + + @staticmethod + def _move_text_embeds( + embeds: TextEmbeds, + cu_seqlens: Tensor, + attention_mask: Tensor | None, + *, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[TextEmbeds, Tensor, Tensor | None]: + moved = {key: value.to(device=device, dtype=dtype) for key, value in embeds.items()} + mask = attention_mask.to(device=device) if attention_mask is not None else None + return moved, cu_seqlens.to(device=device), mask + + @staticmethod + def _packed_bundle(bundle: LatentBundle) -> LatentBundle: + return LatentBundle( + video=bundle.video.flatten(0, 1) if bundle.video is not None else None, + audio=bundle.audio.flatten(0, 1) if bundle.audio is not None else None, + video_cu_seqlens=bundle.video_cu_seqlens, + audio_cu_seqlens=bundle.audio_cu_seqlens, + ) + + @staticmethod + def _text_length(embeds: TextEmbeds, cu_seqlens: Tensor) -> int: + value = embeds["text_embeds"] + return int(value.shape[1] if value.ndim == 3 else cu_seqlens[-1].item()) + + def _text_ropes( + self, + text_length: int, + negative_text_length: int, + audio_length: int | None, + device: torch.device, + ) -> tuple[Tensor | list[Tensor], Tensor | list[Tensor], Tensor | None]: + if getattr(self.transformer, "is_multimodal", False): + text_rope = [ + compute_rope1d(self.transformer.video_text_rope_embeddings, text_length, device=device), + compute_rope1d(self.transformer.audio_text_rope_embeddings, text_length, device=device), + ] + negative_text_rope = [ + compute_rope1d( + self.transformer.video_text_rope_embeddings, + negative_text_length, + device=device, + ), + compute_rope1d( + self.transformer.audio_text_rope_embeddings, + negative_text_length, + device=device, + ), + ] + else: + text_rope = compute_rope1d(self.transformer.text_rope_embeddings, text_length, device=device) + negative_text_rope = compute_rope1d( + self.transformer.text_rope_embeddings, + negative_text_length, + device=device, + ) + audio_rope = ( + compute_rope1d(self.transformer.audio_rope_embeddings, audio_length, device=device) + if audio_length is not None + else None + ) + return text_rope, negative_text_rope, audio_rope + + @staticmethod + def _format_video(frames: Tensor, output_type: str) -> Tensor | np.ndarray: + if output_type in ("pt", "torch"): + return frames + if output_type in ("np", "numpy"): + return frames.permute(0, 2, 3, 4, 1).cpu().numpy() + raise ValueError("output_type must be 'pt', 'torch', 'np', 'numpy', or 'latent'") + + @torch.no_grad() + def __call__( + self, + prompt: str | list[str] | None = None, + image: str | object | list[str | object] | None = None, + negative_prompt: str | list[str] | None = None, + height: int = 512, + width: int = 768, + num_frames: int = 121, + sample_fps: float = 24.0, + num_inference_steps: int = 50, + max_sequence_length: int = 1024, + guidance_scale: float = 5.0, + generator: torch.Generator | None = None, + latents: Tensor | None = None, + audio_latents: Tensor | None = None, + prompt_embeds: TextEmbeds | None = None, + negative_prompt_embeds: TextEmbeds | None = None, + cu_seqlens: Tensor | tuple[Tensor, Tensor] | None = None, + prompt_cu_seqlens: Tensor | None = None, + negative_prompt_cu_seqlens: Tensor | None = None, + sample_audio: bool = True, + expand_prompts: bool = False, + output_type: str = "pt", + return_dict: bool = True, + callback_on_step_end: Callable[..., Any] | None = None, + callback_on_step_end_tensor_inputs: list[str] | None = None, + visual_cond_scheme: str | None = None, + ) -> Kandinsky6TI2VAPipelineOutput | tuple[Tensor | np.ndarray, list[np.ndarray] | None]: + r"""Generate synchronized video and, optionally, audio from text or an image. + + Args: + prompt: Text prompt or a batch of prompts. + image: Optional reference image or batch of reference images. + negative_prompt: Optional negative prompt. + height: Output video height in pixels; must be divisible by 8. + width: Output video width in pixels; must be divisible by 8. + num_frames: Number of decoded video frames. + sample_fps: Output video frame rate. Kandinsky 6 is trained for + 24.0 fps; other values may affect audio/video alignment. + num_inference_steps: Number of denoising steps. + max_sequence_length: Maximum Qwen prompt length after the template. + guidance_scale: Classifier-free guidance weight. + generator: Random generator used for latent initialization. + latents: Optional precomputed video latents. + audio_latents: Optional precomputed audio latents. + prompt_embeds: Optional precomputed positive embeddings. + negative_prompt_embeds: Optional precomputed negative embeddings. + cu_seqlens: Optional positive/negative cumulative sequence lengths. + sample_audio: Whether to generate synchronized audio. + expand_prompts: Whether to use the built-in Qwen video+audio prompt expander before encoding. + output_type: ``pt``/``torch`` for tensors, ``np``/``numpy`` for + NumPy arrays, or ``latent`` for undecoded video latents. + return_dict: Whether to return ``Kandinsky6TI2VAPipelineOutput``. + callback_on_step_end: Optional callback invoked after each step. + callback_on_step_end_tensor_inputs: Names passed to the callback. + visual_cond_scheme: Optional image-conditioning scheme. Defaults to + ``pretrain`` without ``image`` and ``tail_cond_first_frame`` + with ``image``. + + Returns: + ``Kandinsky6TI2VAPipelineOutput`` containing video frames and int16 NumPy + audio waveforms, or a tuple when ``return_dict=False``. + """ + if cu_seqlens is not None: + if isinstance(cu_seqlens, tuple): + if len(cu_seqlens) != 2: + raise ValueError("cu_seqlens tuple must contain positive and negative lengths") + if prompt_cu_seqlens is None: + prompt_cu_seqlens = cu_seqlens[0] + if negative_prompt_cu_seqlens is None: + negative_prompt_cu_seqlens = cu_seqlens[1] + else: + if prompt_cu_seqlens is None: + prompt_cu_seqlens = cu_seqlens + if negative_prompt_cu_seqlens is None: + negative_prompt_cu_seqlens = cu_seqlens + + if sample_fps != 24.0: + warnings.warn( + f"Kandinsky 6 was trained for 24.0 fps; received sample_fps={sample_fps}.", + UserWarning, + stacklevel=2, + ) + + batch_size = self._batch_size( + prompt, + negative_prompt, + image, + latents, + audio_latents, + prompt_embeds, + negative_prompt_embeds, + ) + prompt_batch = self._as_batch(prompt, batch_size, "prompt") + negative_prompt_batch = self._as_batch(negative_prompt, batch_size, "negative_prompt") + image_batch = self._as_batch(image, batch_size, "image") + if visual_cond_scheme is None: + visual_cond_scheme = "tail_cond_first_frame" if image is not None else "pretrain" + self.check_inputs( + image=image_batch, + prompt=prompt_batch, + negative_prompt=negative_prompt_batch, + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=num_inference_steps, + max_sequence_length=max_sequence_length, + output_type=output_type, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_cu_seqlens=prompt_cu_seqlens, + negative_prompt_cu_seqlens=negative_prompt_cu_seqlens, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + sample_audio=sample_audio, + visual_cond_scheme=visual_cond_scheme, + ) + + callback_inputs = ( + self._callback_tensor_inputs + if callback_on_step_end_tensor_inputs is None + else callback_on_step_end_tensor_inputs + ) + if expand_prompts and prompt_batch is not None and prompt_embeds is None: + prompt_batch = self.expand_prompts( + prompt_batch, + image=image_batch, + max_sequence_length=max_sequence_length, + mode="i2va" if image_batch is not None else "t2va", + ) + patch_size = tuple(int(value) for value in getattr(self.transformer, "patch_size", (1, 2, 2))) + + device = getattr(self, "_execution_device", self.transformer.device) + dtype = getattr(self.transformer, "dtype", torch.bfloat16) + if not isinstance(dtype, torch.dtype): + dtype = torch.bfloat16 + first_frames = None + if image_batch is not None: + encoded_frames = [ + encode_i2va_first_frame(item, self.vae, device, height=height, width=width) for item in image_batch + ] + first_frames = torch.cat([item[0] for item in encoded_frames], dim=0) + encoded_height, encoded_width = encoded_frames[0][1:] + if any(item[1:] != (encoded_height, encoded_width) for item in encoded_frames): + raise ValueError("all reference images must produce the same resolution") + height, width = encoded_height, encoded_width + negative_prompt_batch = negative_prompt_batch or [self._DEFAULT_NEGATIVE_PROMPT] * batch_size + ( + positive, + positive_cu, + positive_mask, + negative, + negative_cu, + negative_mask, + ) = self._encode_prompt( + prompt_batch or [""] * batch_size, + negative_prompt_batch, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_cu_seqlens=prompt_cu_seqlens, + negative_prompt_cu_seqlens=negative_prompt_cu_seqlens, + max_sequence_length=max_sequence_length, + ) + if positive_cu is None or negative_cu is None: + raise ValueError("positive and negative prompt cumulative lengths are required") + positive, positive_cu, positive_mask = self._move_text_embeds( + positive, + positive_cu, + positive_mask, + device=device, + dtype=dtype, + ) + negative, negative_cu, negative_mask = self._move_text_embeds( + negative, + negative_cu, + negative_mask, + device=device, + dtype=dtype, + ) + + latent_frames = (num_frames - 1) // 4 + 1 + bundle = self._prepare_latents( + latents=latents, + audio_latents=audio_latents, + latent_frames=latent_frames, + height=height, + width=width, + dtype=dtype, + device=device, + generator=generator, + sample_fps=sample_fps, + sample_audio=sample_audio, + batch_size=batch_size, + ) + visual_token_type_ids = None + generated_visual_mask = None + if image is not None and visual_cond_scheme == "tail_cond_first_frame": + if bundle.video is None or first_frames is None: + raise ValueError("image conditioning requires video latents") + video, visual_token_type_ids, generated_visual_mask = append_i2va_tail_condition( + bundle.video, + first_frames, + batch_size=batch_size, + video_duration=latent_frames, + ) + bundle = LatentBundle( + video=video, + audio=bundle.audio, + video_cu_seqlens=(latent_frames + 1) * torch.arange(batch_size + 1, dtype=torch.int32, device=device), + audio_cu_seqlens=bundle.audio_cu_seqlens, + ) + audio_length = int(bundle.audio.shape[1]) if bundle.audio is not None else None + visual_shape = ( + latent_frames // patch_size[0], + (height // 8) // patch_size[1], + (width // 8) // patch_size[2], + ) + if visual_shape[0] < 1: + raise ValueError(f"invalid visual latent shape {visual_shape}") + visual_rope = compute_visual_rope( + self.transformer.visual_rope_embeddings, + visual_shape, + self.scale_factor, + device=device, + ) + if generated_visual_mask is not None: + visual_rope = torch.cat([visual_rope, visual_rope[:1]], dim=0) + text_rope, negative_text_rope, audio_rope = self._text_ropes( + text_length=self._text_length(positive, positive_cu), + negative_text_length=self._text_length(negative, negative_cu), + audio_length=audio_length, + device=device, + ) + with self.progress_bar(total=num_inference_steps) as progress_bar: + + def _step_callback(step_index: int, timestep: Tensor, current: LatentBundle) -> LatentBundle: + if callback_on_step_end is not None: + callback_result = callback_on_step_end( + self, + step_index, + timestep, + {name: current.video for name in callback_inputs}, + ) + if callback_result is None: + callback_result = {} + if not isinstance(callback_result, dict): + raise TypeError("callback_on_step_end must return a dict or None") + if callback_result.get("latents") is not None: + current.video = callback_result["latents"] + return current + + result = denoise_loop( + bundle=bundle, + dit=self.transformer, + text_embeds=positive, + null_text_embeds=negative, + visual_rope=visual_rope, + audio_rope=audio_rope, + text_rope=text_rope, + null_text_rope=negative_text_rope, + num_steps=num_inference_steps, + guidance_weight=guidance_scale, + first_frames=first_frames, + visual_cond_scheme=visual_cond_scheme, + sample_video=True, + sample_audio=sample_audio, + attention_mask=positive_mask, + null_attention_mask=negative_mask, + visual_token_type_ids=visual_token_type_ids, + scale_factor=self.scale_factor, + scheduler=self.scheduler, + step_callback=_step_callback if callback_on_step_end is not None else None, + progress_bar=progress_bar, + ) + + if generated_visual_mask is not None and result.video is not None: + generated_video = result.video[:, generated_visual_mask[0]] + result = LatentBundle( + video=generated_video, + audio=result.audio, + video_cu_seqlens=latent_frames * torch.arange(batch_size + 1, dtype=torch.int32, device=device), + audio_cu_seqlens=result.audio_cu_seqlens, + ) + + if output_type == "latent": + if result.video is None: + raise ValueError("denoising returned no video latent") + frames: Tensor | np.ndarray = result.video.permute(0, 4, 1, 2, 3) + else: + packed_result = self._packed_bundle(result) + decoded_frames = postprocess_video(packed_result, self.vae, bs=batch_size) + frames = self._format_video(decoded_frames, output_type) + audio = postprocess_audio(self._packed_bundle(result), self.audio_vae) if sample_audio else None + if not return_dict: + return frames, audio + return Kandinsky6TI2VAPipelineOutput(frames=frames, audio=audio) diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_output.py b/src/diffusers/pipelines/kandinsky6/pipeline_output.py new file mode 100644 index 000000000000..989b33037a6d --- /dev/null +++ b/src/diffusers/pipelines/kandinsky6/pipeline_output.py @@ -0,0 +1,36 @@ +"""Output containers for the Kandinsky 6 Diffusers pipelines.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import numpy as np +from ...utils import BaseOutput + +Audio = list[np.ndarray] | np.ndarray | None + + +@dataclass +class Kandinsky6TI2VAPipelineOutput(BaseOutput): + """Output of the K6 video-and-audio pipeline.""" + + frames: Any + audio: list[np.ndarray] | None + + +@dataclass +class Kandinsky6SRPipelineOutput(BaseOutput): + """Output of the K6 super-resolution pipeline.""" + + frames: Any + audio: Audio = None + path: str | list[str] | None = None + metadata: dict[str, Any] | None = None + + +__all__ = [ + "Audio", + "Kandinsky6SRPipelineOutput", + "Kandinsky6TI2VAPipelineOutput", +] diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index c0e46ef445df..baad89b768e9 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -61,6 +61,7 @@ _import_structure["scheduling_euler_ancestral_discrete"] = ["EulerAncestralDiscreteScheduler"] _import_structure["scheduling_euler_discrete"] = ["EulerDiscreteScheduler"] _import_structure["scheduling_flow_map_euler_discrete"] = ["FlowMapEulerDiscreteScheduler"] + _import_structure["scheduling_piflow"] = ["PiflowScheduler"] _import_structure["scheduling_flow_match_euler_discrete"] = ["FlowMatchEulerDiscreteScheduler"] _import_structure["scheduling_flow_match_heun_discrete"] = ["FlowMatchHeunDiscreteScheduler"] _import_structure["scheduling_flow_match_lcm"] = ["FlowMatchLCMScheduler"] @@ -145,6 +146,7 @@ from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler from .scheduling_euler_discrete import EulerDiscreteScheduler from .scheduling_flow_map_euler_discrete import FlowMapEulerDiscreteScheduler + from .scheduling_piflow import PiflowScheduler from .scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler from .scheduling_flow_match_heun_discrete import FlowMatchHeunDiscreteScheduler from .scheduling_flow_match_lcm import FlowMatchLCMScheduler diff --git a/src/diffusers/schedulers/scheduling_piflow.py b/src/diffusers/schedulers/scheduling_piflow.py new file mode 100644 index 000000000000..758568faf68e --- /dev/null +++ b/src/diffusers/schedulers/scheduling_piflow.py @@ -0,0 +1,271 @@ +"""Diffusers scheduler for distilled Kandinsky 6 PiFlow checkpoints.""" + +from __future__ import annotations + +import torch +from ..configuration_utils import register_to_config +from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, + FlowMatchEulerDiscreteSchedulerOutput, +) + + +class DXPolicy: + """Network-free DX policy over one flow-matching segment.""" + + def __init__( # noqa: PLR0913 + self, + denoising_output: torch.Tensor, + x_t_src: torch.Tensor, + sigma_t_src: torch.Tensor, + segment_size: float | torch.Tensor = 1.0, + shift: float = 1.0, + mode: str = "grid", + eps: float = 1e-4, + ) -> None: + self.x_t_src = x_t_src + self.ndim = x_t_src.dim() + self.shift = shift + self.eps = eps + if mode not in ("grid", "polynomial"): + raise ValueError(f"Unknown mode: {mode}") + self.mode = mode + + 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) + segment = segment_size + if isinstance(segment, torch.Tensor) and segment.dim() < self.raw_t_src.dim(): + segment = segment.reshape(*segment.size(), *((self.raw_t_src.dim() - segment.dim()) * [1])) + self.raw_t_dst = (self.raw_t_src - segment).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: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + n = x.size(1) + if n < 2: # noqa: PLR2004 + 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 + indices = torch.stack([t0, t1], dim=1) + values = torch.gather(x, dim=1, index=indices.expand(-1, -1, *x.shape[2:])) + return (t1 - t) * values[:, 0] + (t - t0) * values[:, 1] + + def _unwarp_t(self, sigma_t: torch.Tensor) -> torch.Tensor: + return sigma_t / (self.shift + (1 - self.shift) * sigma_t) + + @staticmethod + def _u_to_x_0( + denoising_output: torch.Tensor, + x_t: torch.Tensor, + sigma_t: torch.Tensor, + ) -> torch.Tensor: + return x_t.unsqueeze(1) - sigma_t.unsqueeze(1) * denoising_output + + def pi(self, x_t: torch.Tensor, sigma_t: torch.Tensor) -> torch.Tensor: + sigma_t = sigma_t.reshape(*sigma_t.size(), *((self.ndim - sigma_t.dim()) * [1])) + raw_t = self._unwarp_t(sigma_t) + if self.mode == "grid": + x_0 = self._interpolate( + self.denoising_output_x_0, + (raw_t - self.raw_t_dst) / self.segment_size, + ) + else: + p_order = self.denoising_output_x_0.size(1) + diff_t = self.raw_t_src - raw_t + basis = torch.stack([diff_t**i for i in range(p_order)], dim=1) + x_0 = torch.sum(basis * self.denoising_output_x_0, dim=1) + return (x_t - x_0) / sigma_t.clamp(min=self.eps) + + +def shift_timesteps(t: torch.Tensor, shift: float) -> torch.Tensor: + """Map raw flow-matching time to the shifted DiT time.""" + return shift * t / (1 + (shift - 1) * t) + + +def policy_rollout_fm( # noqa: PLR0913 + x_t_start: torch.Tensor, + sigma_t_start: torch.Tensor, + raw_t_start: torch.Tensor, + raw_t_end: torch.Tensor, + total_substeps: int, + policy: DXPolicy, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Integrate ``policy.pi`` from ``raw_t_start`` to ``raw_t_end``.""" + num_batches = x_t_start.size(0) + ndim = x_t_start.dim() + shape = (num_batches, *((ndim - 1) * [1])) + raw_t_start = raw_t_start.reshape(shape) + raw_t_end = raw_t_end.reshape(shape) + sigma_t = sigma_t_start.reshape(shape) + + 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 + max_num_substeps = num_substeps.max() + + raw_t = raw_t_start + x_t = x_t_start + for substep_id in range(max_num_substeps.item()): + velocity = 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 + velocity * (sigma_t_minus - sigma_t) + + active_mask = num_substeps > substep_id + x_t = torch.where(active_mask, x_t_minus, x_t) + sigma_t = torch.where(active_mask, sigma_t_minus, sigma_t) + raw_t = torch.where(active_mask, raw_t_minus, raw_t) + + return x_t, sigma_t, sigma_t.flatten() * 1_000 + + +class PiflowScheduler(FlowMatchEulerDiscreteScheduler): + """Few-step flow-matching scheduler for widened-output DiTs.""" + + is_piflow = True + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 5.0, + n_grid: int = 10, + nfe: int | None = None, + eps: float = 1e-6, + final_step_size_scale: float = 0.5, + num_policy_substeps: int = 128, + ) -> None: + if n_grid < 2: + raise ValueError(f"PiflowScheduler requires n_grid >= 2, got {n_grid}") + if eps <= 0: + raise ValueError(f"PiflowScheduler requires eps > 0, got {eps}") + if not 0 < final_step_size_scale <= 1: + raise ValueError("PiflowScheduler requires 0 < final_step_size_scale <= 1") + if num_policy_substeps < 1: + raise ValueError("PiflowScheduler requires num_policy_substeps >= 1") + super().__init__( + num_train_timesteps=num_train_timesteps, + shift=shift, + ) + self.n_grid = int(n_grid) + self.eps = float(eps) + 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) + + def set_timesteps( + 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: + if sigmas is not None or mu is not None or timesteps is not None: + raise ValueError("PiflowScheduler only supports its configured distilled timestep 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}") + 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 + sigmas = shift_timesteps(raw, float(self.config.shift)) + self.num_inference_steps = int(num_inference_steps) + self._piflow_raw_timesteps = raw + self.timesteps = sigmas * self.config.num_train_timesteps + self.sigmas = torch.cat([sigmas, sigmas.new_zeros(1)]) + self._step_index = None + self._begin_index = None + + def _to_grid(self, model_output: torch.Tensor, sample: torch.Tensor) -> torch.Tensor: + if model_output.ndim != sample.ndim or model_output.shape[:-1] != sample.shape[:-1]: + raise ValueError( + "Piflow model output must match sample shape except for the output channels: " + f"got {tuple(model_output.shape)} for sample {tuple(sample.shape)}" + ) + if model_output.shape[-1] % self.n_grid != 0: + raise ValueError( + f"Piflow model output channels {model_output.shape[-1]} are not divisible by n_grid={self.n_grid}" + ) + output_dim = model_output.shape[-1] // self.n_grid + if output_dim != sample.shape[-1]: + raise ValueError( + "Piflow model output channels do not match the sample: " + f"expected {sample.shape[-1] * self.n_grid}, got {model_output.shape[-1]}" + ) + return model_output.reshape(*model_output.shape[:-1], self.n_grid, output_dim).movedim(-2, 1) + + 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) + 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]) + policy = DXPolicy( + model_output, + sample, + sigma, + segment, + shift=float(self.config.shift), + mode="grid", + eps=self.eps, + ) + updated, _, _ = policy_rollout_fm( + sample, + sigma, + raw_src.expand(sample.shape[0]), + raw_dst.expand(sample.shape[0]), + self.num_policy_substeps, + policy, + ) + return updated.to(dtype=sample.dtype) + + def _step_index_for(self, timestep: torch.Tensor | float) -> int: + if self.step_index is None: + if self.begin_index is not None: + self._step_index = self.begin_index + else: + schedule_timesteps = self.timesteps + if not isinstance(schedule_timesteps, torch.Tensor): + schedule_timesteps = torch.as_tensor(schedule_timesteps, dtype=torch.float32) + timestep = torch.as_tensor( + timestep, + device=schedule_timesteps.device, + dtype=schedule_timesteps.dtype, + ) + indices = torch.nonzero(schedule_timesteps == timestep).flatten() + if not indices.numel(): + raise ValueError(f"timestep {timestep.item()} is not in the Piflow schedule") + position = 1 if indices.numel() > 1 else 0 + self._step_index = int(indices[position].item()) + if self.step_index is None or self.step_index >= self.num_inference_steps: + raise RuntimeError("PiflowScheduler.step called after the schedule was exhausted") + return int(self.step_index) + + def step( + self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor, + sample: torch.FloatTensor, + return_dict: bool = True, + ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple: + if isinstance(timestep, int) or isinstance(timestep, (torch.IntTensor, torch.LongTensor)): + raise ValueError( + "Passing integer indices as timesteps to PiflowScheduler.step() is not supported; " + "pass a value from scheduler.timesteps instead" + ) + step_index = self._step_index_for(timestep) + # PiFlow's policy rollout performs its update in float32. Keep that + # precision across outer steps, matching the native sampler. + updated = self._policy_step(model_output, sample.to(torch.float32), step_index) + self._step_index += 1 + if return_dict: + return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=updated) + return (updated,) From c892279b663892e39a1239f1688bae4146d34ee6 Mon Sep 17 00:00:00 2001 From: Denis Koposov Date: Wed, 16 Sep 2026 12:47:00 +0000 Subject: [PATCH 02/62] update docs for Kandinsky 6 --- docs/source/en/_toctree.yml | 8 ++ .../en/api/models/kandinsky6_transformer.md | 28 ++++ docs/source/en/api/models/kandinsky6_vae.md | 29 ++++ docs/source/en/api/pipelines/kandinsky6.md | 40 ++++++ docs/source/en/api/schedulers/piflow.md | 23 ++++ .../autoencoders/autoencoder_kandinsky6_sr.py | 78 ++++++++++- .../autoencoders/autoencoder_mmaudio.py | 36 ++++- .../transformers/transformer_kandinsky6.py | 127 +++++++++++++++--- .../transformers/transformer_kandinsky6_sr.py | 29 ++-- .../pipelines/kandinsky6/latent_upscaler.py | 61 +++++++-- .../kandinsky6/pipeline_kandinsky6_sr.py | 69 ++++++++-- .../kandinsky6/pipeline_kandinsky6_ti2va.py | 66 +++++++++ .../pipelines/kandinsky6/pipeline_output.py | 14 ++ src/diffusers/schedulers/scheduling_piflow.py | 58 +++++++- src/diffusers/utils/dummy_pt_objects.py | 75 +++++++++++ 15 files changed, 677 insertions(+), 64 deletions(-) create mode 100644 docs/source/en/api/models/kandinsky6_transformer.md create mode 100644 docs/source/en/api/models/kandinsky6_vae.md create mode 100644 docs/source/en/api/pipelines/kandinsky6.md create mode 100644 docs/source/en/api/schedulers/piflow.md diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index f05667986f11..3941d091fe83 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -359,6 +359,8 @@ title: JoyImageEditPlusTransformer3DModel - local: api/models/transformer_joyimage title: JoyImageEditTransformer3DModel + - local: api/models/kandinsky6_transformer + title: Kandinsky6Transformer3DModel - local: api/models/krea2_transformer2d title: Krea2Transformer2DModel - local: api/models/latte_transformer3d @@ -481,6 +483,8 @@ title: AutoencoderSAME - local: api/models/consistency_decoder_vae title: ConsistencyDecoderVAE + - local: api/models/kandinsky6_vae + title: Kandinsky 6 VAEs - local: api/models/ltx2_diffusion_decoder title: LTX2VideoDiffusionDecoderModel - local: api/models/autoencoder_oobleck @@ -701,6 +705,8 @@ title: HunyuanVideo1.5 - local: api/pipelines/kandinsky5_video title: Kandinsky 5.0 Video + - local: api/pipelines/kandinsky6 + title: Kandinsky 6 - local: api/pipelines/latte title: Latte - local: api/pipelines/ltx2 @@ -792,6 +798,8 @@ title: LMSDiscreteScheduler - local: api/schedulers/minimax_h3 title: MiniMaxH3Scheduler + - local: api/schedulers/piflow + title: PiflowScheduler - local: api/schedulers/pndm title: PNDMScheduler - local: api/schedulers/repaint diff --git a/docs/source/en/api/models/kandinsky6_transformer.md b/docs/source/en/api/models/kandinsky6_transformer.md new file mode 100644 index 000000000000..798f32695f58 --- /dev/null +++ b/docs/source/en/api/models/kandinsky6_transformer.md @@ -0,0 +1,28 @@ + + +# Kandinsky 6 Transformers + +Kandinsky 6 uses Diffusion Transformers (DiTs) for text-to-video and +text-to-video-and-audio generation, as well as video super-resolution. + +## Kandinsky6Transformer3DModel + +[[autodoc]] Kandinsky6Transformer3DModel + - all + - forward + +## Kandinsky6SRTransformer3DModel + +[[autodoc]] Kandinsky6SRTransformer3DModel + - all + - forward diff --git a/docs/source/en/api/models/kandinsky6_vae.md b/docs/source/en/api/models/kandinsky6_vae.md new file mode 100644 index 000000000000..933270d78e5b --- /dev/null +++ b/docs/source/en/api/models/kandinsky6_vae.md @@ -0,0 +1,29 @@ + + +# Kandinsky 6 VAEs + +Kandinsky 6 uses a causal video VAE for super-resolution and an audio VAE +with a vocoder for synchronized audio generation. + +## Kandinsky6SRVAE + +[[autodoc]] Kandinsky6SRVAE + - all + - encode + - decode + +## MMAudioVAE + +[[autodoc]] MMAudioVAE + - wrapped_encode + - wrapped_decode diff --git a/docs/source/en/api/pipelines/kandinsky6.md b/docs/source/en/api/pipelines/kandinsky6.md new file mode 100644 index 000000000000..042526006943 --- /dev/null +++ b/docs/source/en/api/pipelines/kandinsky6.md @@ -0,0 +1,40 @@ + + +# Kandinsky 6 + +Kandinsky 6 provides pipelines for text/image-to-video-and-audio generation +and video super-resolution. + +## Kandinsky6TI2VAPipeline + +[[autodoc]] Kandinsky6TI2VAPipeline + - all + - __call__ + +## Kandinsky6SRPipeline + +[[autodoc]] Kandinsky6SRPipeline + - all + - __call__ + +## Kandinsky6SRLatentUpscalerBank + +[[autodoc]] Kandinsky6SRLatentUpscalerBank + - all + - forward + +## Pipeline outputs + +[[autodoc]] Kandinsky6TI2VAPipelineOutput + +[[autodoc]] Kandinsky6SRPipelineOutput diff --git a/docs/source/en/api/schedulers/piflow.md b/docs/source/en/api/schedulers/piflow.md new file mode 100644 index 000000000000..61da66c15ae5 --- /dev/null +++ b/docs/source/en/api/schedulers/piflow.md @@ -0,0 +1,23 @@ + + +# PiflowScheduler + +`PiflowScheduler` is a few-step flow-matching scheduler for distilled +Kandinsky 6 models (both TI2VA and SR). + +## PiflowScheduler + +[[autodoc]] PiflowScheduler + - all + - set_timesteps + - step diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index d87fd0df4a6d..27de9285020f 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -1,3 +1,17 @@ +# 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 6 SR KVAE Diffusers component.""" from __future__ import annotations @@ -77,7 +91,7 @@ def forward(self, x, write_to=None, transform=None): def nonlinearity(x): - # swish + # Apply the SiLU activation. return x * torch.sigmoid(x) @@ -1014,12 +1028,32 @@ class DecoderOutput: class Kandinsky6SRVAE(ModelMixin, ConfigMixin): + """Causal 3D VAE used by the Kandinsky 6 video super-resolution pipeline. + + The encoder and decoder process videos in temporal segments while reusing + their causal convolution state between segments. This keeps memory usage + bounded for long videos and preserves the checkpoint layout of the KVAE. + + Args: + vae_type (`str`): VAE architecture identifier. Must be ``"video-kvae"``. + encoder_config (`dict`): Configuration for the cached 3D encoder. + decoder_config (`dict`): Configuration for the cached 3D decoder. + scaling_factor (`float`, *optional*, defaults to 1.0): Latent scaling + factor stored in the component configuration. + spatial_factor (`int`, *optional*, defaults to 16): Spatial compression + factor of the VAE. + temporal_factor (`int`, *optional*, defaults to 4): Temporal compression + factor of the VAE. + """ + @staticmethod def normalize_data(data): + """Normalize pixel values to the KVAE input range.""" return data / 128 - 1.0 @staticmethod def denormalize_data(data): + """Convert normalized KVAE outputs back to pixel values.""" return (data + 1) * 128 @register_to_config @@ -1043,6 +1077,7 @@ def __init__( self.temporal_factor = int(temporal_factor) def init_from_ckpt(self, path): + """Load a KVAE checkpoint in safetensors or training-checkpoint format.""" if str(path).endswith(".safetensors"): # Release checkpoints are flat safetensors state dicts already in # the new-style key naming. @@ -1081,6 +1116,8 @@ def init_from_ckpt(self, path): self.load_state_dict(sd, strict=True) def make_empty_cache(self, block: str): + """Create empty causal-convolution and normalization caches.""" + def make_dict(name, p=None): if name == "conv": return {"padding": None} @@ -1119,9 +1156,20 @@ def make_dict(name, p=None): @apply_forward_hook def encode(self, x, seg_len=16): + """Encode a video in temporal segments and return latents and segment sizes. + + Args: + x (`torch.Tensor`): Video tensor in ``(batch, channels, frames, height, width)`` format. + seg_len (`int`, *optional*, defaults to 16): Number of non-initial + frames processed in each segment. + + Returns: + `tuple[torch.Tensor, list[int]]`: Encoded latents and the pixel-space + segment sizes needed by :meth:`decode`. + """ cache = self.make_empty_cache("enc") - ## get segments size + # Compute segment sizes. split_list = [seg_len + 1] n_frames = x.size(2) - (seg_len + 1) while n_frames > 0: @@ -1130,7 +1178,7 @@ def encode(self, x, seg_len=16): split_list[-1] += n_frames - ## encode by segments + # Encode each segment. latent = [] for chunk in torch.split(x, split_list, dim=2): l = self.encoder(chunk, cache) @@ -1142,9 +1190,19 @@ def encode(self, x, seg_len=16): @apply_forward_hook def decode(self, z, split_list=None): + """Decode latent segments while reusing the causal decoder cache. + + Args: + z (`torch.Tensor`): Latent tensor in ``(batch, channels, frames, height, width)`` format. + split_list (`list[int]`, *optional*): Pixel-space segment sizes + returned by :meth:`encode`. + + Returns: + `DecoderOutput`: Decoded video in ``sample``. + """ cache = self.make_empty_cache("dec") - ## get segments size + # Compute latent segment sizes. if split_list is None: default_split_size = 16 // self.conf["enc"]["temporal_compress_times"] time_dim = z.shape[2] @@ -1160,7 +1218,7 @@ def decode(self, z, split_list=None): else: split_list = [math.ceil(size / self.conf["enc"]["temporal_compress_times"]) for size in split_list] - ## decode by segments + # Decode each segment. recs = [] for chunk in torch.split(z, split_list, dim=2): out = self.decoder(chunk, cache) @@ -1170,6 +1228,16 @@ def decode(self, z, split_list=None): return DecoderOutput(sample=recs) def forward(self, x, seg_len: int = 16): + """Encode and decode a video in one call. + + Args: + x (`torch.Tensor`): Video tensor in ``(batch, channels, frames, height, width)`` format. + seg_len (`int`, *optional*, defaults to 16): Number of non-initial + frames processed in each segment. + + Returns: + `DecoderOutput`: Reconstructed video in ``sample``. + """ latent, split_list = self.encode(x, seg_len) recs = self.decode(latent, split_list) return recs diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index 243bd3c91f5d..a9c8f787a574 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -20,9 +20,8 @@ from einops import rearrange from huggingface_hub import PyTorchModelHubMixin, hf_hub_download from librosa.filters import mel as librosa_mel_fn -from torch import nn, pow, sin +from torch import pow, sin from torch.nn import Conv1d, ConvTranspose1d, Parameter -from torch.nn import functional as F from torch.nn.utils import weight_norm as legacy_weight_norm from torch.nn.utils.parametrizations import weight_norm from torch.nn.utils.parametrize import remove_parametrizations @@ -33,8 +32,8 @@ # Attribution-NonCommercial-ShareAlike 4.0 International License. # You should have received a copy of the license along with this # work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/ -"""Improved diffusion model architecture proposed in the paper -"Analyzing and Improving the Training Dynamics of Diffusion Models".""" +# Improved diffusion model architecture proposed in the paper +# "Analyzing and Improving the Training Dynamics of Diffusion Models". # ---------------------------------------------------------------------------- @@ -1530,6 +1529,7 @@ def forward(self, waveform: torch.Tensor, center: bool = False) -> torch.Tensor: def get_mel_converter(mode: Literal["44k"]) -> MelConverter: + """Create the 44.1 kHz mel-spectrogram converter used by MMAudio.""" return MelConverter( sampling_rate=44_100, n_fft=2048, @@ -1542,10 +1542,29 @@ def get_mel_converter(mode: Literal["44k"]) -> MelConverter: ) -"""Audio feature / VAE facade.""" +# Audio feature / VAE facade. class MMAudioVAE(ModelMixin, ConfigMixin): + """Audio VAE and vocoder used to encode and decode audio latents. + + The component converts waveforms to mel-spectrograms, encodes them into + latent representations, and decodes latents back to waveforms through a + BigVGAN vocoder. + + Args: + tod_vae_ckpt (`str` or `os.PathLike`, *optional*): Path to the audio VAE checkpoint. + bigvgan_vocoder_ckpt (`str` or `os.PathLike`, *optional*): Path to the BigVGAN checkpoint. + vocoder_config (`dict`, *optional*): BigVGAN configuration when no checkpoint directory is supplied. + mode (`str`, *optional*, defaults to ``"44k"``): Audio operating mode. + need_vae_encoder (`bool`, *optional*, defaults to True): Whether to load the VAE encoder. + need_vae_decoder (`bool`, *optional*, defaults to True): Whether to load the VAE decoder and vocoder. + scaling_factor (`float`, *optional*, defaults to 1.0): Latent scaling factor. + sample_rate (`int`, *optional*, defaults to 44100): Audio sample rate. + downsample_factor (`int`, *optional*, defaults to 1024): Audio latent downsampling factor. + device (`str` or `torch.device`, *optional*): Device on which to place the component. + """ + ignore_for_config = ["tod_vae_ckpt", "bigvgan_vocoder_ckpt", "device"] @register_to_config @@ -1580,22 +1599,27 @@ def __init__( self.to(device) def compile(self): + """Compile the decode and vocode methods with ``torch.compile``.""" self.decode = torch.compile(self.decode) self.vocode = torch.compile(self.vocode) def train(self, mode: bool = True) -> MMAudioVAE: + """Keep the inference-only audio component in evaluation mode.""" return super().train(False) def encode_audio(self, x) -> DiagonalGaussianDistribution: + """Encode a waveform into an audio latent distribution.""" assert self.tod is not None, "VAE is not loaded" mel = self.mel_converter(x) return self.tod.encode(mel) def vocode(self, mel: torch.Tensor) -> torch.Tensor: + """Convert a decoded mel-spectrogram into a waveform.""" assert self.tod is not None, "VAE is not loaded" return self.tod.vocode(mel) def decode(self, z: torch.Tensor) -> torch.Tensor: + """Decode audio latents into a mel-spectrogram.""" assert self.tod is not None, "VAE is not loaded" return self.tod.decode(z) @@ -1609,11 +1633,13 @@ def dtype(self): @apply_forward_hook def wrapped_decode(self, z): + """Decode latents and vocode them through the Diffusers forward hook.""" mel_decoded = self.decode(z.to(dtype=self.dtype)) return self.vocode(mel_decoded) @apply_forward_hook def wrapped_encode(self, audio): + """Encode audio and return the mean latent through the forward hook.""" dist = self.encode_audio(audio.to(dtype=self.dtype)) return dist.mean diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index cafc554628fb..33fdd3170708 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -1,3 +1,17 @@ +# 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 6 Diffusers transformer.""" from __future__ import annotations @@ -26,6 +40,16 @@ def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: + """Return inverse frequencies for rotary position embeddings. + + Args: + dim (`int`): Number of frequency values to generate. + max_period (`float`, *optional*, defaults to 10000.0): Maximum period + used by the frequency schedule. + + Returns: + `torch.Tensor`: Frequency values in float32. + """ return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) @@ -69,6 +93,7 @@ def _local_merge(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Te def fractal_flatten(x: Tensor, rope: Tensor, shape: tuple, block_mask: bool = False): + """Flatten visual tokens, optionally arranging them into local fractal blocks.""" if block_mask: ps = 8 x = _local_patch(x, shape, (1, ps, ps), dim=0) @@ -78,6 +103,7 @@ def fractal_flatten(x: Tensor, rope: Tensor, shape: tuple, block_mask: bool = Fa def fractal_unflatten(x: Tensor, shape: tuple, block_mask: bool = False) -> Tensor: + """Restore the spatial layout produced by :func:`fractal_flatten`.""" if block_mask: ps = 8 x = x.reshape(-1, ps * ps, x.shape[-1]) @@ -94,7 +120,7 @@ def fast_sta_nabla( wW: int = 3, device: str | torch.device = "cuda", ) -> Tensor: - """Precomputes the Sliding Tile Attention (STA) boolean mask for nabla attention.""" + """Precompute the Sliding Tile Attention (STA) boolean mask for NABLA attention.""" l = max(T, H, W) r = torch.arange(l, dtype=torch.int16, device=device) mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs() @@ -115,7 +141,7 @@ def nabla_block_mask( thr: float = 0.9, block_size: int = 64, ) -> BlockMask: - """Builds a dynamic nabla BlockMask from query/key statistics + STA prior.""" + """Build a dynamic NABLA block mask from query/key statistics and an STA prior.""" B, h, S, D = q.shape s1 = S // block_size qa = q.reshape(B, h, s1, block_size, D).mean(-2) @@ -214,14 +240,13 @@ def reset_parameters(self) -> None: setattr(self, f"args_{i}", torch.outer(torch.arange(mp, dtype=freq.dtype, device=freq.device), freq)) -"""Diffusers MagCache adapter for the K6 multimodal transformer. - -Diffusers provides the public :class:`MagCacheConfig` and the stateful hook -infrastructure. K6's fused visual blocks carry video and audio streams -together, so the stock hook needs a small adapter to preserve both tensors -when a block is skipped. The adapter keeps the Diffusers ``enable_cache`` / -``disable_cache`` API and uses the Diffusers MagCache state and configuration. -""" +# Diffusers MagCache adapter for the K6 multimodal transformer. +# +# Diffusers provides the public :class:`MagCacheConfig` and the stateful hook +# infrastructure. K6's fused visual blocks carry video and audio streams +# together, so the stock hook needs a small adapter to preserve both tensors +# when a block is skipped. The adapter keeps the Diffusers ``enable_cache`` / +# ``disable_cache`` API and uses the Diffusers MagCache state and configuration. _HEAD_HOOK = "kandinsky6_mag_cache_head" @@ -434,14 +459,13 @@ def disable_cache(self) -> None: self._cache_config = None -"""Diffusers-style K6 transformer components for the source port. - -The assembler inlines this module after the portable attention, RoPE, and -tensor helpers. The classes use the same inner-module vocabulary as the -Diffusers Kandinsky5 transformer (for example ``in_layer``, ``modulation``, -``self_attention`` and ``feed_forward``). Checkpoint conversion maps the -native K6 names to this public Diffusers layout. -""" +# Diffusers-style K6 transformer components for the source port. +# +# The assembler inlines this module after the portable attention, RoPE, and +# tensor helpers. The classes use the same inner-module vocabulary as the +# Diffusers Kandinsky5 transformer (for example ``in_layer``, ``modulation``, +# ``self_attention`` and ``feed_forward``). Checkpoint conversion maps the +# native K6 names to this public Diffusers layout. _MASKED_ATTENTION_BACKENDS = { @@ -1060,9 +1084,47 @@ class Kandinsky6Transformer3DModel( CacheMixin, AttentionMixin, ): - """Standalone K6 T2V/T2VA transformer with Diffusers model mixins.""" - - _repeated_blocks = ["Kandinsky6TransformerEncoderBlock", "Kandinsky6TransformerDecoderBlock"] + """Kandinsky 6 transformer for text-to-video and text-to-video-and-audio generation. + + The model uses separate text and visual transformer blocks for video-only + generation and fused video/audio blocks for multimodal generation. Inputs + and outputs use packed token layouts compatible with the native K6 model. + + Args: + in_visual_dim (`int`, *optional*, defaults to 16): Number of input video latent channels. + out_visual_dim (`int`, *optional*, defaults to 16): Number of output video latent channels. + in_text_dim (`int`, *optional*, defaults to 3584): Text token embedding dimension. + in_text_dim2 (`int`, *optional*, defaults to 768): Pooled text embedding dimension. + time_dim (`int`, *optional*, defaults to 1024): Time embedding dimension. + patch_size (`tuple[int, int, int]`, *optional*, defaults to ``(1, 2, 2)``): Video patch size. + model_dim (`int`, *optional*, defaults to 4096): Video transformer hidden dimension. + ff_dim (`int`, *optional*, defaults to 16384): Video feed-forward hidden dimension. + num_text_blocks (`int`, *optional*, defaults to 4): Number of text blocks. + num_visual_blocks (`int`, *optional*, defaults to 60): Number of visual blocks. + axes_dims (`tuple[int, int, int]`, *optional*, defaults to ``(32, 48, 48)``): RoPE dimensions for video. + visual_cond (`bool`, *optional*, defaults to True): Whether video conditioning channels are present. + is_multimodal (`bool`, *optional*, defaults to False): Whether to construct fused video/audio blocks. + in_audio_dim (`int`, *optional*, defaults to 20): Number of input audio latent channels. + out_audio_dim (`int`, *optional*, defaults to 20): Number of output audio latent channels. + model_dim_a (`int`, *optional*): Audio transformer hidden dimension. Defaults to `model_dim`. + time_dim_a (`int`, *optional*): Audio time embedding dimension. Defaults to `time_dim`. + ff_dim_a (`int`, *optional*): Audio feed-forward hidden dimension. Defaults to `ff_dim`. + axes_dims_a (`tuple[int, int, int]`, *optional*): Audio RoPE dimensions. Defaults to `axes_dims`. + audio_freqs_scaling (`float`, *optional*, defaults to 1.0): Audio RoPE frequency scaling. + text_token_padding (`bool`, *optional*, defaults to False): Whether text sequences are padded. + scale_factor (`tuple[float]`, *optional*): Per-axis RoPE frequency scaling. + ca_rope (`bool`, *optional*, defaults to False): Whether to use cross-modal audio RoPE. + cross_gates (`bool`, *optional*, defaults to False): Whether to use cross-modal residual gates. + fix_modulation (`bool`, *optional*, defaults to False): Whether to use the fixed modulation variant. + visual_token_type_num_embeddings (`int`, *optional*, defaults to 0): Number of visual token type embeddings. + magcache (`dict`, *optional*): MagCache configuration. + """ + + _repeated_blocks = [ + "Kandinsky6TransformerEncoderBlock", + "Kandinsky6TransformerDecoderBlock", + "Kandinsky6FusedTransformerDecoderBlock", + ] _keep_in_fp32_modules = ["time_embeddings", "modulation", "visual_modulation", "text_modulation"] _supports_gradient_checkpointing = True @@ -1181,6 +1243,7 @@ def __init__( ) def clear_text_proj_cache(self) -> None: + """Clear cached projected text embeddings.""" self._text_proj_cache.clear() def _project_pooled(self, prefix: str | None, pooled: Tensor) -> Tensor: @@ -1376,6 +1439,28 @@ def forward( return_dict: bool = False, **kwargs: Any, ) -> Tensor | tuple[Tensor, Tensor] | Transformer2DModelOutput: + """Run the K6 transformer on video, audio, or fused video/audio latents. + + Args: + x_video (`torch.Tensor`, *optional*): Packed video latent tokens. + x_audio (`torch.Tensor`, *optional*): Packed audio latent tokens. + text_embed (`torch.Tensor` or `list[torch.Tensor]`): Text token embeddings. + pooled_text_embed (`torch.Tensor` or `list[torch.Tensor]`): Pooled text embeddings. + time (`torch.Tensor` or `list[torch.Tensor]`): Diffusion timesteps. + visual_rope (`torch.Tensor`): Video rotary position embeddings. + audio_rope (`torch.Tensor`, *optional*): Audio rotary position embeddings. + text_rope (`torch.Tensor` or `list[torch.Tensor]`): Text rotary position embeddings. + sparse_params (`dict`, *optional*): NABLA sparse-attention configuration. + attention_mask (`torch.Tensor`, *optional*): Text attention mask. + visual_token_type_ids (`torch.Tensor`, *optional*): Video token type IDs. + return_dict (`bool`, *optional*, defaults to False): Whether to return a + [`Transformer2DModelOutput`] instead of tensors. + + Returns: + `torch.Tensor`, `tuple[torch.Tensor, torch.Tensor]`, or + [`Transformer2DModelOutput`]: Denoised video output, audio and video + outputs, or a model output object. + """ if text_embed is None or pooled_text_embed is None or time is None or text_rope is None: raise ValueError("text_embed, pooled_text_embed, time, and text_rope are required") both = x_video is not None and x_audio is not None and self.is_multimodal diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index a5f4fb35411a..013436b55a49 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -1,3 +1,17 @@ +# 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 6 SR transformer Diffusers component.""" from __future__ import annotations @@ -421,12 +435,11 @@ def framewise_causal_dense( return block_mask_from_bool(doc[None, None]) -"""Pure-torch helpers for the streaming KV-cache in self-attention. - -Kept free of ``flash_attn`` imports so the cache bookkeeping is importable and -testable on CPU. The actual attention kernel call lives in -``kandinsky_sr.model.nn``; this module only assembles its packed inputs. -""" +# Pure-torch helpers for the streaming KV-cache in self-attention. +# +# Kept free of ``flash_attn`` imports so the cache bookkeeping is importable and +# testable on CPU. The actual attention kernel call lives in +# ``kandinsky_sr.model.nn``; this module only assembles its packed inputs. if TYPE_CHECKING: @@ -669,7 +682,7 @@ def evict_kv_slot(slot: KVSlot, max_tokens_per_seq: int) -> KVSlot: return evicted -"""Neural network building blocks for the diffusion transformer.""" +# Neural network building blocks for the diffusion transformer. logger = logging.get_logger(__name__) @@ -1796,7 +1809,7 @@ def reset_parameters(self) -> None: self.modulation.reset_parameters() -"""Diffusion transformer architecture for text-conditioned visual generation.""" +# Diffusion transformer architecture for text-conditioned visual generation. class TransformerEncoderBlock(nn.Module): diff --git a/src/diffusers/pipelines/kandinsky6/latent_upscaler.py b/src/diffusers/pipelines/kandinsky6/latent_upscaler.py index ec81765c4e0a..854baff7fdf2 100644 --- a/src/diffusers/pipelines/kandinsky6/latent_upscaler.py +++ b/src/diffusers/pipelines/kandinsky6/latent_upscaler.py @@ -1,3 +1,17 @@ +# 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 6 SR latent-upscaler bank Diffusers component.""" from __future__ import annotations @@ -28,7 +42,7 @@ def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) -"""Convolution primitives shared by the latent upsampler architectures.""" +# Convolution primitives shared by the latent upsampler architectures. DIMS_2 = 2 @@ -161,7 +175,7 @@ def make_conv(dims: int, temporal_padding: TemporalPadding) -> type[nn.Module]: return nn.Conv3d -"""Reusable building blocks for latent upsampler architectures.""" +# Reusable building blocks for latent upsampler architectures. class RMSNorm(nn.Module): @@ -499,7 +513,7 @@ def zero_init_residual(self) -> None: nn.init.zeros_(final_conv.bias) # type: ignore[arg-type] -"""Exact merging of flash-style attention branches with a closed-form null branch.""" +# Exact merging of flash-style attention branches with a closed-form null branch. if TYPE_CHECKING: @@ -618,7 +632,7 @@ def merge_attention_branches( return MergeWithNullBranch.apply(null_value, null_logits, len(outputs), torch_compile, *outputs, *lses) -"""Runtime helpers shared by latent-upscaler model components.""" +# Runtime helpers shared by latent-upscaler model components. if TYPE_CHECKING: @@ -636,7 +650,7 @@ def forward_with_checkpointing( return module(*inputs) -"""Spatial upsampling operations: pixel-shuffle, bilinear, and K-VAE PXS v2.""" +# Spatial upsampling operations: pixel-shuffle, bilinear, and K-VAE PXS v2. if TYPE_CHECKING: @@ -925,7 +939,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.linear(self.pxs(x) + self.v3(x)) -"""Motion-correspondence attention for latent video features.""" +# Motion-correspondence attention for latent video features. if TYPE_CHECKING: @@ -1228,7 +1242,7 @@ def load_natten() -> ModuleType: return natten -"""Construction and execution helpers for the cascaded multi-scale model.""" +# Construction and execution helpers for the cascaded multi-scale model. STAGE_FACTOR = 2 @@ -1399,7 +1413,7 @@ def apply_block_sequence( return x -"""Convolutional latent-space spatial upsampler with configurable upsampling strategy.""" +# Convolutional latent-space spatial upsampler with configurable upsampling strategy. class ConvLatentUpsampler(nn.Module): @@ -1680,7 +1694,7 @@ def forward(self, z: torch.Tensor) -> torch.Tensor: return rearrange(x, "(b t) c h w -> b c t h w", b=b, t=t) if self.dims == 2 else x -"""Scale-specific x2 adaptation modules for the cascaded latent upscaler.""" +# Scale-specific x2 adaptation modules for the cascaded latent upscaler. X2TailMode = Literal["shared", "scale_specific_norm", "private", "private_full"] @@ -1902,7 +1916,7 @@ def forward_second_stage( return x -"""Multi-scale cascaded upsampler: two 2x stages with intermediate supervision.""" +# Multi-scale cascaded upsampler: two 2x stages with intermediate supervision. # The two projections back to latent space. They are scale-specific, so a warm @@ -2274,7 +2288,7 @@ def _add_nearest_base( return out + spatial_nearest_2x(z, DIMS_3, self.upscale_factor) -"""Configuration models for the latent upscaler.""" +# Configuration models for the latent upscaler. UpsampleMode = Literal["pixel_shuffle", "bilinear", "pxs_v2", "pxs_v2_hybrid"] @@ -2544,7 +2558,7 @@ def validate_motion_attention(self) -> MultiScaleModelConfig: FlatModelConfig | MultiScaleModelConfig, Field(discriminator="architecture"), ] -"""Factory functions for building upsampler models.""" +# Factory functions for building upsampler models. if TYPE_CHECKING: @@ -2660,7 +2674,14 @@ def build_motion_attention_spec(config: MultiScaleModelConfig) -> MotionAttentio class Kandinsky6SRLatentUpscalerBank(ModelMixin, ConfigMixin): - """Diffusers wrapper around the self-contained x2/x4 upscaler bank.""" + """Diffusers wrapper around the self-contained x2/x4 latent-upscaler bank. + + Args: + models (`list[dict]`): Serialized upscaler definitions. Each definition + contains a ``target_scale`` and a validated ``model`` configuration. + scaling_factor (`float`, *optional*, defaults to 1.0): Factor applied to + latents before they are passed to an upscaler. + """ @register_to_config def __init__( @@ -2699,7 +2720,19 @@ def forward( entry: str | None = None, return_intermediates: bool | None = None, ) -> Any: - """Forward one bank entry, primarily for direct component use.""" + """Forward one bank entry, primarily for direct component use. + + Args: + latent (`torch.Tensor`): Latent tensor to upscale. + entry (`str`, *optional*): Entry name such as ``"x2"`` or ``"2x"``. + Required when the bank contains multiple scales. + return_intermediates (`bool`, *optional*): Whether to return + intermediate upscaler outputs when supported by the entry. + + Returns: + `torch.Tensor` or `dict`: Upscaled latent, or the selected entry's + intermediate-output structure. + """ if entry is None: if len(self._models) != 1: raise ValueError("entry is required when the latent-upscaler bank has multiple scales") diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py index b3974b44cfa3..f57afba38041 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py @@ -1,3 +1,17 @@ +# 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 6 super-resolution Diffusers pipeline.""" from __future__ import annotations @@ -15,6 +29,7 @@ import numpy as np import torch from ..pipeline_utils import DiffusionPipeline +from diffusers.utils import replace_example_docstring from pydantic import BaseModel, ConfigDict, Field from ...schedulers.scheduling_piflow import DXPolicy, policy_rollout_fm, shift_timesteps from torch import Tensor, nn @@ -23,6 +38,28 @@ from .pipeline_output import Kandinsky6SRPipelineOutput +EXAMPLE_DOC_STRING = """ + Examples: + + ```python + >>> import torch + >>> from diffusers import Kandinsky6SRPipeline + + >>> model_id = "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers" + >>> pipe = Kandinsky6SRPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) + >>> pipe = pipe.to("cuda") + + >>> output = pipe.from_video( + ... "input.mp4", + ... resolution_scale=2.25, + ... save_path="output.mp4", + ... ) + >>> output.path + 'output.mp4' + ``` +""" + + # These values are the SR data contract and must stay local to this pipeline. VAE_SPATIAL_FACTOR = 16 VAE_TEMPORAL_FACTOR = 4 @@ -61,7 +98,7 @@ def _execution_device_context(device: torch.device | str): yield -"""Shared tiling utilities for spatial video tiling with Hanning-window blending.""" +# Shared tiling utilities for spatial video tiling with Hanning-window blending. class TileGrid(NamedTuple): @@ -1577,15 +1614,14 @@ def _upsample_tiles_to_base(raw_tiles: list[torch.Tensor], base_h: int, base_w: ] -"""Reusable Kandinsky SR inference stages and tiled orchestration. - -The functions in this module deliberately do not depend on a benchmark -runner. A framework can call ``text_encode`` / ``prepare_latents`` / -``denoise`` / ``vae_decode`` directly, or use the tiled functions for the -scale-aware path. Native benchmark integrations may add timing contexts to -the tiled calls; the public :class:`Kandinsky6SRPipeline` does not expose -those callbacks. -""" +# Reusable Kandinsky SR inference stages and tiled orchestration. +# +# The functions in this module deliberately do not depend on a benchmark +# runner. A framework can call ``text_encode`` / ``prepare_latents`` / +# ``denoise`` / ``vae_decode`` directly, or use the tiled functions for the +# scale-aware path. Native benchmark integrations may add timing contexts to +# the tiled calls; the public :class:`Kandinsky6SRPipeline` does not expose +# those callbacks. @dataclass(frozen=True) @@ -2576,6 +2612,7 @@ class Kandinsky6SRPipeline(DiffusionPipeline): Args: transformer: SR transformer used to denoise latent tiles. vae: Video VAE used for pixel-input encoding. + scheduler: Scheduler used for SR denoising. latent_upscaler: Optional x2/x4 latent upscaler bank. source_vae: Optional source VAE used only by the KVAE latent bridge. """ @@ -2980,6 +3017,7 @@ def from_latents( # noqa: PLR0913 ) @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, video: SRBatchInput | None = None, @@ -3006,6 +3044,15 @@ def __call__( in a supported rank-5 layout, or local video paths. latents: Raw latent tensor, a supported rank-5 latent batch, or local ``.pt`` files. + resolution_scale: Total spatial upscale factor. Supported values + are ``2``, ``4``, and ``2.25``. + num_inference_steps: Number of denoising steps for the SR model. + generator: Optional random generator used to derive the sampling seed. + overlap: Fraction of overlap between adjacent tiles. + tiles_batch_size: Number of tiles processed in one sampler batch. + kvae_bridge: Whether latent inputs come from a base VAE and need to + be decoded before SR. + cached_text_embeds: Optional cached empty-caption embeddings. save_path: Optional output video path or one path per input. fps: Optional output video frame rate. audio: Optional audio waveform to mux. @@ -3014,6 +3061,8 @@ def __call__( output_type: ``pt``/``torch`` or ``np``/``numpy``. return_dict: Whether to return ``Kandinsky6SRPipelineOutput``. + Examples: + Returns: ``Kandinsky6SRPipelineOutput`` or its tuple representation when ``return_dict=False``. diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index f505fd0a8ec4..5dcd11b5016c 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -1,3 +1,17 @@ +# 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 6 TI2VA Diffusers pipeline.""" from __future__ import annotations @@ -14,10 +28,45 @@ import numpy as np import torch from ..pipeline_utils import DiffusionPipeline +from diffusers.utils import replace_example_docstring from torch import Tensor, nn from .pipeline_output import Kandinsky6TI2VAPipelineOutput +EXAMPLE_DOC_STRING = """ + Examples: + + ```python + >>> import torch + >>> from diffusers import Kandinsky6TI2VAPipeline + >>> from diffusers.utils import encode_video + + >>> model_id = "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers" + >>> pipe = Kandinsky6TI2VAPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) + >>> pipe.enable_model_cpu_offload() + + >>> output = pipe( + ... prompt="A cat and a dog baking a cake together in a kitchen.", + ... height=480, + ... width=864, + ... num_frames=121, + ... num_inference_steps=16, + ... guidance_scale=1.0, + ... sample_audio=True, + ... ) + + >>> video = output.frames[0].permute(1, 2, 3, 0) + >>> audio = torch.as_tensor(output.audio[0])[:, None].repeat(1, 2) + >>> encode_video( + ... video, + ... fps=24, + ... output_path="output.mp4", + ... audio=audio, + ... audio_sample_rate=pipe.audio_sample_rate, + ... ) + ``` +""" + @dataclass class LatentBundle: @@ -927,6 +976,16 @@ def _as_text_embeds(value: Any) -> TextEmbeds: def encode_prompt( self, text: str | list[str], max_sequence_length: int ) -> tuple[TextEmbeds, Tensor, Tensor | None]: + """Encode text with Qwen2.5-VL and CLIP and build packed token metadata. + + Args: + text (`str` or `list[str]`): Prompt or prompts to encode. + max_sequence_length (`int`): Maximum number of Qwen text tokens. + + Returns: + `tuple`: Text embeddings, cumulative sequence lengths, and an + optional Qwen attention mask. + """ texts = [text] if isinstance(text, str) else text full_texts = [_PROMPT_TEMPLATE.format(item) for item in texts] qwen_device = next(self.text_encoder.parameters()).device @@ -1302,6 +1361,7 @@ def _format_video(frames: Tensor, output_type: str) -> Tensor | np.ndarray: raise ValueError("output_type must be 'pt', 'torch', 'np', 'numpy', or 'latent'") @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, prompt: str | list[str] | None = None, @@ -1350,6 +1410,10 @@ def __call__( prompt_embeds: Optional precomputed positive embeddings. negative_prompt_embeds: Optional precomputed negative embeddings. cu_seqlens: Optional positive/negative cumulative sequence lengths. + prompt_cu_seqlens: Optional cumulative sequence lengths for + precomputed positive prompt embeddings. + negative_prompt_cu_seqlens: Optional cumulative sequence lengths + for precomputed negative prompt embeddings. sample_audio: Whether to generate synchronized audio. expand_prompts: Whether to use the built-in Qwen video+audio prompt expander before encoding. output_type: ``pt``/``torch`` for tensors, ``np``/``numpy`` for @@ -1361,6 +1425,8 @@ def __call__( ``pretrain`` without ``image`` and ``tail_cond_first_frame`` with ``image``. + Examples: + Returns: ``Kandinsky6TI2VAPipelineOutput`` containing video frames and int16 NumPy audio waveforms, or a tuple when ``return_dict=False``. diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_output.py b/src/diffusers/pipelines/kandinsky6/pipeline_output.py index 989b33037a6d..d70e41b830a7 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_output.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_output.py @@ -1,3 +1,17 @@ +# 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. + """Output containers for the Kandinsky 6 Diffusers pipelines.""" from __future__ import annotations diff --git a/src/diffusers/schedulers/scheduling_piflow.py b/src/diffusers/schedulers/scheduling_piflow.py index 758568faf68e..b66c907bc964 100644 --- a/src/diffusers/schedulers/scheduling_piflow.py +++ b/src/diffusers/schedulers/scheduling_piflow.py @@ -1,3 +1,17 @@ +# 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. + """Diffusers scheduler for distilled Kandinsky 6 PiFlow checkpoints.""" from __future__ import annotations @@ -122,7 +136,27 @@ def policy_rollout_fm( # noqa: PLR0913 class PiflowScheduler(FlowMatchEulerDiscreteScheduler): - """Few-step flow-matching scheduler for widened-output DiTs.""" + """Few-step PiFlow scheduler for widened-output diffusion transformers. + + PiFlow evaluates the denoising model at a small number of grid points and + integrates a network-free policy between those evaluations. The scheduler + is intended for distilled Kandinsky 6 checkpoints, including the main + video/audio model and the video super-resolution model. Their model output + contains ``n_grid`` predictions per sample channel. + + Args: + num_train_timesteps (`int`, *optional*, defaults to 1000): Number of + training diffusion steps. + shift (`float`, *optional*, defaults to 5.0): Flow-matching timestep shift. + n_grid (`int`, *optional*, defaults to 10): Number of predictions in the + widened model output. + nfe (`int`, *optional*): Number of model evaluations used at inference. + eps (`float`, *optional*, defaults to 1e-6): Minimum timestep and policy denominator. + final_step_size_scale (`float`, *optional*, defaults to 0.5): Relative + size of the final raw-timestep segment. + num_policy_substeps (`int`, *optional*, defaults to 128): Maximum policy + integration substeps per raw-timestep unit. + """ is_piflow = True @@ -163,6 +197,15 @@ def set_timesteps( mu: float | None = None, timesteps: list[float] | None = None, ) -> None: + """Set the distilled PiFlow timestep schedule. + + Args: + num_inference_steps (`int`): Number of model evaluations. + device (`str` or `torch.device`, *optional*): Device for the schedule. + sigmas (`list[float]`, *optional*): Unsupported custom sigma schedule. + mu (`float`, *optional*): Unsupported dynamic-shift parameter. + timesteps (`list[float]`, *optional*): Unsupported custom timestep schedule. + """ if sigmas is not None or mu is not None or timesteps is not None: raise ValueError("PiflowScheduler only supports its configured distilled timestep schedule") if num_inference_steps is None or num_inference_steps < 1: @@ -256,6 +299,19 @@ def step( sample: torch.FloatTensor, return_dict: bool = True, ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple: + """Advance one step by integrating the PiFlow policy. + + Args: + model_output (`torch.FloatTensor`): Widened model output containing + ``n_grid`` predictions per sample channel. + timestep (`float` or `torch.FloatTensor`): Current scheduler timestep. + sample (`torch.FloatTensor`): Current noisy sample. + return_dict (`bool`, *optional*, defaults to True): Whether to return + a [`FlowMatchEulerDiscreteSchedulerOutput`]. + + Returns: + [`FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`: Updated sample. + """ if isinstance(timestep, int) or isinstance(timestep, (torch.IntTensor, torch.LongTensor)): raise ValueError( "Passing integer indices as timesteps to PiflowScheduler.step() is not supported; " diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index f34008252ab6..6f9b9bf477ae 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1575,6 +1575,66 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class MMAudioVAE(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 Kandinsky6SRTransformer3DModel(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 Kandinsky6SRVAE(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 Kandinsky6Transformer3DModel(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 Kandinsky5Transformer3DModel(metaclass=DummyObject): _backends = ["torch"] @@ -3422,6 +3482,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class PiflowScheduler(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 FlowMatchEulerDiscreteScheduler(metaclass=DummyObject): _backends = ["torch"] From 6b8a080ca20a3a63d34fe9f4faf30eee1740c7c0 Mon Sep 17 00:00:00 2001 From: Denis Koposov Date: Wed, 16 Sep 2026 14:42:39 +0000 Subject: [PATCH 03/62] remove redundant bigvgan code and manual checkpoint loading inside model code --- .../autoencoders/autoencoder_mmaudio.py | 166 +----------------- 1 file changed, 7 insertions(+), 159 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index a9c8f787a574..17cbec32dd6e 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -2,13 +2,9 @@ from __future__ import annotations -import json import math -import os -import shutil from os import PathLike -from pathlib import Path -from typing import Any, Dict, Literal, Optional, Union +from typing import Any, Literal, Optional import numpy as np import torch @@ -18,11 +14,9 @@ from ..modeling_utils import ModelMixin from ...utils.accelerate_utils import apply_forward_hook from einops import rearrange -from huggingface_hub import PyTorchModelHubMixin, hf_hub_download from librosa.filters import mel as librosa_mel_fn from torch import pow, sin from torch.nn import Conv1d, ConvTranspose1d, Parameter -from torch.nn.utils import weight_norm as legacy_weight_norm from torch.nn.utils.parametrizations import weight_norm from torch.nn.utils.parametrize import remove_parametrizations @@ -933,13 +927,6 @@ def __init__(self, *args, **kwargs): self.__dict__ = self -def build_env(config, config_name, path): - t_path = os.path.join(path, config_name) - if config != t_path: - os.makedirs(path, exist_ok=True) - shutil.copyfile(config, os.path.join(path, config_name)) - - # Adapted from https://github.com/jik876/hifi-gan under the MIT license. # LICENSE is in incl_licenses directory. @@ -950,24 +937,10 @@ def init_weights(m, mean=0.0, std=0.01): m.weight.data.normal_(mean, std) -def apply_weight_norm(m): - classname = m.__class__.__name__ - if classname.find("Conv") != -1: - legacy_weight_norm(m) - - def get_padding(kernel_size, dilation=1): return int((kernel_size * dilation - dilation) / 2) -def load_checkpoint(filepath, device): - assert os.path.isfile(filepath) - print(f"Loading '{filepath}'") - checkpoint_dict = torch.load(filepath, map_location=device) - print("Complete.") - return checkpoint_dict - - # Copyright (c) 2024 NVIDIA CORPORATION. # Licensed under the MIT license. @@ -975,12 +948,6 @@ def load_checkpoint(filepath, device): # LICENSE is in incl_licenses directory. -def load_hparams_from_json(path) -> AttrDict: - with open(path) as f: - data = f.read() - return AttrDict(json.loads(data)) - - class AMPBlock1(torch.nn.Module): """ AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer. @@ -1159,16 +1126,7 @@ def remove_weight_norm(self): TorchActivation1d = Activation1d -class BigVGANv2( - torch.nn.Module, - PyTorchModelHubMixin, - library_name="bigvgan", - repo_url="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/NVIDIA/BigVGAN", - docs_url="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/NVIDIA/BigVGAN/blob/main/README.md", - pipeline_tag="audio-to-audio", - license="mit", - tags=["neural-vocoder", "audio-generation", "arxiv:2206.04658"], -): +class BigVGANv2(torch.nn.Module): """ BigVGAN is a neural vocoder model that applies anti-aliased periodic activation for residual blocks (resblocks). @@ -1290,118 +1248,11 @@ def remove_weight_norm(self): print("[INFO] Model already removed weight norm. Skipping!") pass - # Additional methods for huggingface_hub support - def _save_pretrained(self, save_directory: Path) -> None: - """Save weights and config.json from a Pytorch model to a local directory.""" - - model_path = save_directory / "bigvgan_generator.pt" - torch.save({"generator": self.state_dict()}, model_path) - - config_path = save_directory / "config.json" - with open(config_path, "w") as config_file: - json.dump(self.h, config_file, indent=4) - - @classmethod - def _from_pretrained( - cls, - *, - model_id: str, - revision: str, - cache_dir: str, - force_download: bool, - proxies: Optional[Dict], - resume_download: bool, - local_files_only: bool, - token: Union[str, bool, None], - map_location: str = "cpu", # Additional argument - strict: bool = False, # Additional argument - **model_kwargs, - ): - """Load Pytorch pretrained weights and return the loaded model.""" - - # Download and load hyperparameters (h) used by BigVGAN - if os.path.isdir(model_id): - # print("Loading config.json from local directory") - config_file = os.path.join(model_id, "config.json") - else: - config_file = hf_hub_download( - repo_id=model_id, - filename="config.json", - revision=revision, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - resume_download=resume_download, - token=token, - local_files_only=local_files_only, - ) - h = load_hparams_from_json(config_file) - - model = cls(h) - - # Download and load pretrained generator weight - if os.path.isdir(model_id): - # print("Loading weights from local directory") - model_file = os.path.join(model_id, "bigvgan_generator.pt") - else: - print(f"Loading weights from {model_id}") - model_file = hf_hub_download( - repo_id=model_id, - filename="bigvgan_generator.pt", - revision=revision, - cache_dir=cache_dir, - force_download=force_download, - proxies=proxies, - resume_download=resume_download, - token=token, - local_files_only=local_files_only, - ) - - checkpoint_dict = torch.load(model_file, map_location=map_location, weights_only=True) - try: - model.load_state_dict(checkpoint_dict["generator"]) - except RuntimeError: - print( - f"[INFO] the pretrained checkpoint does not contain weight norm. Loading the checkpoint after removing weight norm!" - ) - model.remove_weight_norm() - model.load_state_dict(checkpoint_dict["generator"]) - - return model - - -def _load_bigvgan_v2( - vocoder_ckpt_path: str | None = None, - *, - vocoder_config: dict | None = None, -) -> BigVGANv2: - """Load BigVGAN-v2 from a local directory (avoids HubMixin/huggingface_hub API drift).""" - model_file = None if vocoder_ckpt_path is None else os.path.join(vocoder_ckpt_path, "bigvgan_generator.pt") - if vocoder_config is None: - if vocoder_ckpt_path is None: - raise ValueError("Either vocoder_ckpt_path or vocoder_config is required") - config_file = os.path.join(vocoder_ckpt_path, "config.json") - if not os.path.isfile(config_file) or not os.path.isfile(model_file): - raise FileNotFoundError(f"Expected config.json and bigvgan_generator.pt under {vocoder_ckpt_path}") - h = load_hparams_from_json(config_file) - else: - if model_file is not None and not os.path.isfile(model_file): - raise FileNotFoundError(f"BigVGAN checkpoint does not exist: {model_file}") - h = AttrDict(vocoder_config) - - model = BigVGANv2(h) - if vocoder_ckpt_path is not None: - assert model_file is not None - checkpoint_dict = torch.load(model_file, map_location="cpu", weights_only=True) - try: - model.load_state_dict(checkpoint_dict["generator"]) - except RuntimeError: - model.remove_weight_norm() - model.load_state_dict(checkpoint_dict["generator"]) - model.remove_weight_norm() - else: - model.remove_weight_norm() +def build_bigvgan_v2(vocoder_config: dict | AttrDict) -> BigVGANv2: + """Build an inference-ready BigVGAN-v2 from its serialized configuration.""" + model = BigVGANv2(vocoder_config if isinstance(vocoder_config, AttrDict) else AttrDict(vocoder_config)) + model.remove_weight_norm() return model @@ -1426,10 +1277,7 @@ def __init__( self.vae.remove_weight_norm() if need_vae_decoder: - self.vocoder = _load_bigvgan_v2( - vocoder_ckpt_path, - vocoder_config=vocoder_config, - ).eval() + self.vocoder = build_bigvgan_v2(vocoder_config).eval() else: del self.vae.decoder From bf4158a41d649b7df061fb44530eca9330be79f9 Mon Sep 17 00:00:00 2001 From: Denis Koposov Date: Wed, 16 Sep 2026 16:42:29 +0000 Subject: [PATCH 04/62] add missing _no_split_modules --- src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py | 2 ++ src/diffusers/models/autoencoders/autoencoder_mmaudio.py | 1 + src/diffusers/models/transformers/transformer_kandinsky6.py | 1 + src/diffusers/models/transformers/transformer_kandinsky6_sr.py | 2 ++ 4 files changed, 6 insertions(+) diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index 27de9285020f..4071b3244aa4 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -1046,6 +1046,8 @@ class Kandinsky6SRVAE(ModelMixin, ConfigMixin): factor of the VAE. """ + _no_split_modules = ["CachedEncoder3D", "CachedDecoder3D"] + @staticmethod def normalize_data(data): """Normalize pixel values to the KVAE input range.""" diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index 17cbec32dd6e..a2a749702103 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -1414,6 +1414,7 @@ class MMAudioVAE(ModelMixin, ConfigMixin): """ ignore_for_config = ["tod_vae_ckpt", "bigvgan_vocoder_ckpt", "device"] + _no_split_modules = ["AutoEncoderModule", "BigVGANv2"] @register_to_config def __init__( diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 33fdd3170708..a9b4f848a597 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -1125,6 +1125,7 @@ class Kandinsky6Transformer3DModel( "Kandinsky6TransformerDecoderBlock", "Kandinsky6FusedTransformerDecoderBlock", ] + _no_split_modules = _repeated_blocks _keep_in_fp32_modules = ["time_embeddings", "modulation", "visual_modulation", "text_modulation"] _supports_gradient_checkpointing = True diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index 013436b55a49..ab5f3eefdad4 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -2155,6 +2155,8 @@ class Kandinsky6SRTransformer3DModel(ModelMixin, ConfigMixin): with RoPE, adaptive modulation, and optional fractal attention sparsity. """ + _no_split_modules = ["TransformerEncoderBlock", "TransformerDecoderBlock"] + @register_to_config def __init__( self, From 1f1e6ee32b14425b39654b593cd76573cc848fe0 Mon Sep 17 00:00:00 2001 From: Denis Koposov Date: Wed, 16 Sep 2026 21:17:44 +0000 Subject: [PATCH 05/62] use attention processor in SR --- src/diffusers/models/attention_processor.py | 290 ++++++++++++ .../transformers/transformer_kandinsky6.py | 170 +------ .../transformers/transformer_kandinsky6_sr.py | 434 +----------------- 3 files changed, 308 insertions(+), 586 deletions(-) diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index 1b923e749663..2286b03ad661 100755 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -5677,3 +5677,293 @@ def __new__(cls, *args, **kwargs): | LoRAXFormersAttnProcessor | LoRAAttnAddedKVProcessor ) + +import math +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn.functional as functional +from .attention_dispatch import ( + _CAN_USE_FLEX_ATTN, + AttentionBackendName, + dispatch_attention_fn, +) +from torch import Tensor +from torch.nn.attention.flex_attention import BlockMask + + +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 nabla_block_mask( + q: Tensor, + k: Tensor, + sta: Tensor, + thr: float = 0.9, + block_size: int = 64, +) -> BlockMask: + """Build a dynamic NABLA block mask from query/key statistics and an STA prior.""" + B, h, S, D = q.shape + s1 = S // block_size + qa = q.reshape(B, h, s1, block_size, D).mean(-2) + ka = k.reshape(B, h, s1, block_size, D).mean(-2).transpose(-2, -1) + attn_map = torch.softmax((qa @ ka) / math.sqrt(D), dim=-1) + + vals, inds = attn_map.sort(-1) + mask = (vals.cumsum_(-1) >= 1 - thr).int().gather(-1, inds.argsort(-1)) + mask = torch.logical_or(mask, sta) + + kv_nb = mask.sum(-1).to(torch.int32) + kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) + return BlockMask.from_kv_blocks( + torch.zeros_like(kv_nb), + kv_inds, + kv_nb, + kv_inds, + BLOCK_SIZE=block_size, + mask_mod=None, + ) + + +# Pure-torch helpers for the streaming KV-cache in self-attention. +# +# Kept free of ``flash_attn`` imports so the cache bookkeeping is importable and +# testable on CPU. The actual attention kernel call lives in +# ``kandinsky_sr.model.nn``; this module only assembles its packed inputs. + + +if TYPE_CHECKING: + + pass + + +def assemble_cached_attention_inputs( + new_key: Tensor, + new_value: Tensor, + cached_key: Tensor, + cached_value: Tensor, + cu_seqlens_new: Tensor, + cached_cu_seqlens: Tensor, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """Merge cached and new K/V into packed varlen attention inputs. + + For each packed sequence ``i`` the keys/values become ``[cache_i | new_i]``; + the query packing stays the new-token packing (queries are the freshly + arrived tokens that attend back over the cache plus themselves). + + Args: + new_key: New keys, packed ``(sum(new_lens), num_heads, head_dim)``. + new_value: New values, same packing as ``new_key``. + cached_key: Cached keys, packed ``(sum(cache_lens), num_heads, head_dim)``. + cached_value: Cached values, same packing as ``cached_key``. + cu_seqlens_new: Cumulative sequence lengths of the new tokens, ``(num_seqs + 1,)``. + cached_cu_seqlens: Cumulative sequence lengths of the cache, ``(num_seqs + 1,)``. + + Returns: + Tuple ``(key, value, cu_seqlens_q, cu_seqlens_k)`` where ``key``/``value`` + are the merged packed tensors, ``cu_seqlens_q`` is the query packing + (``= cu_seqlens_new``) and ``cu_seqlens_k`` the merged key packing. + """ + num_seqs = cu_seqlens_new.numel() - 1 + key_parts: list[Tensor] = [] + value_parts: list[Tensor] = [] + for i in range(num_seqs): + cache_start, cache_end = int(cached_cu_seqlens[i]), int(cached_cu_seqlens[i + 1]) + new_start, new_end = int(cu_seqlens_new[i]), int(cu_seqlens_new[i + 1]) + key_parts.append(cached_key[cache_start:cache_end]) + key_parts.append(new_key[new_start:new_end]) + value_parts.append(cached_value[cache_start:cache_end]) + value_parts.append(new_value[new_start:new_end]) + + key = torch.cat(key_parts, dim=0) + value = torch.cat(value_parts, dim=0) + + merged_lens = torch.diff(cached_cu_seqlens) + torch.diff(cu_seqlens_new) + cu_seqlens_k = torch.cat([cu_seqlens_new.new_zeros(1), torch.cumsum(merged_lens, dim=0)]).to(cu_seqlens_new.dtype) + cu_seqlens_q = cu_seqlens_new + return key, value, cu_seqlens_q, cu_seqlens_k + + +# A captured per-step cache "slot": one ``(key, value, cu_seqlens)`` triple per +# self-attention layer (aligned with ``iter_self_attentions`` order), or ``None`` +# for an empty slot (no past frames cached yet at that denoising step). +LayerCache = tuple["Tensor | None", "Tensor | None", "Tensor | None"] +KVSlot = "list[LayerCache] | None" + + +_MASKED_ATTENTION_BACKENDS = { + "flash": "flash_varlen", + "_flash_3": "_flash_varlen_3", + "sage": "sage_varlen", + "native": "native", +} + + +class Kandinsky6AttnProcessor: + """Diffusers attention processor shared by TI2VA and SR attention.""" + + _attention_backend = None + _parallel_config = None + + def __init__(self, attention_backend=None, parallel_config=None): + if not hasattr(functional, "scaled_dot_product_attention"): + raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0 or newer.") + self._masked = False + self._attention_backend = attention_backend + self._parallel_config = parallel_config + + @property + def _attention_backend(self): + return self.__attention_backend + + @_attention_backend.setter + def _attention_backend(self, backend): + if self._masked and backend is not None: + name = getattr(backend, "value", backend).lower() + backend = _MASKED_ATTENTION_BACKENDS.get(name, name) + backend = AttentionBackendName(backend) + self.__attention_backend = backend + + def __call__( + self, + attn: Any, + hidden_states: Tensor, + encoder_hidden_states: Tensor | None = None, + rotary_emb: Tensor | None = None, + rotary_emb_kv: Tensor | None = None, + sparse_params: dict[str, Any] | None = None, + attn_mask: Tensor | None = None, + ) -> Tensor: + query = attn.to_query(hidden_states) + if encoder_hidden_states is None: + key = attn.to_key(hidden_states) + value = attn.to_value(hidden_states) + else: + key = attn.to_key(encoder_hidden_states) + value = attn.to_value(encoder_hidden_states) + + query = query.reshape(*query.shape[:-1], attn.num_heads, -1) + key = key.reshape(*key.shape[:-1], attn.num_heads, -1) + value = value.reshape(*value.shape[:-1], attn.num_heads, -1) + query = attn.query_norm(query) + key = attn.key_norm(key) + + if rotary_emb is not None: + query = apply_rotary(query, rotary_emb).to(dtype=query.dtype) + if rotary_emb_kv is not None: + key = apply_rotary(key, rotary_emb_kv).to(dtype=key.dtype) + + query_was_batched = query.dim() == 4 + if not query_was_batched: + query, key, value = query.unsqueeze(0), key.unsqueeze(0), value.unsqueeze(0) + elif key.dim() == 3: + key, value = key.unsqueeze(0), value.unsqueeze(0) + + if sparse_params is not None: + q = query.transpose(1, 2).contiguous() + k = key.transpose(1, 2).contiguous() + v = value.transpose(1, 2).contiguous() + block_mask = nabla_block_mask( + q, + k, + sparse_params["sta_mask"], + thr=sparse_params["P"], + ) + if not _CAN_USE_FLEX_ATTN: + raise ValueError("Nabla attention requires PyTorch 2.5 or newer") + output = ( + torch.nn.attention.flex_attention.flex_attention( + q, + k, + v, + block_mask=block_mask, + ) + .transpose(1, 2) + .contiguous() + ) + else: + output = dispatch_attention_fn( + query, + key, + value, + attn_mask=attn_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + + if not query_was_batched: + output = output[0] + return attn.out_layer(output.flatten(-2, -1)) + + +class Kandinsky6SRAttentionProcessor(Kandinsky6AttnProcessor): + """Add SR sparse-attention and KV-cache handling to the shared processor.""" + + def __call__( + self, + attn: Any, + query: Tensor, + key: Tensor, + value: Tensor, + cu_seqlens_q: Tensor, + cu_seqlens_k: Tensor, + *, + sparse_params: dict[str, Any] | None = None, + attn_mask: Tensor | None = None, + ) -> Tensor: + if sparse_params is not None: + return attn.attention_flex(query, key, value, sparse_params=sparse_params) + + if getattr(attn, "cached_k", None) is not None: + key, value, cu_seqlens_q, cu_seqlens_k = assemble_cached_attention_inputs( + key, value, attn.cached_k, attn.cached_v, cu_seqlens_q, attn.cached_cu_seqlens + ) + if getattr(attn, "return_kv", False): + attn.cached_k, attn.cached_v, attn.cached_cu_seqlens = key, value, cu_seqlens_k + + outputs = [] + batch_size = cu_seqlens_q.numel() - 1 + for index in range(batch_size): + query_start, query_end = int(cu_seqlens_q[index]), int(cu_seqlens_q[index + 1]) + key_start, key_end = int(cu_seqlens_k[index]), int(cu_seqlens_k[index + 1]) + outputs.append( + dispatch_attention_fn( + query[query_start:query_end].unsqueeze(0), + key[key_start:key_end].unsqueeze(0), + value[key_start:key_end].unsqueeze(0), + attn_mask=_packed_attention_mask( + attn_mask, + index, + batch_size, + query_end - query_start, + key_end - key_start, + ), + backend=self._attention_backend, + parallel_config=self._parallel_config, + )[0] + ) + return torch.cat(outputs, dim=0) + + +def _packed_attention_mask( + attn_mask: Tensor | None, + index: int, + batch_size: int, + query_length: int, + key_length: int, +) -> Tensor | None: + if attn_mask is None: + return None + if attn_mask.ndim == 1: + return attn_mask[:key_length].unsqueeze(0) + if attn_mask.ndim == 2: + if attn_mask.shape[0] in (1, batch_size): + return attn_mask[index : index + 1, :key_length] + return attn_mask[:query_length, :key_length] + if attn_mask.shape[0] in (1, batch_size): + attn_mask = attn_mask[index : index + 1] + return attn_mask[..., :query_length, :key_length] + diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index a9b4f848a597..3574a95e58a4 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -27,16 +27,13 @@ from ...hooks.mag_cache import MagCacheState from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ..attention import AttentionMixin, AttentionModuleMixin -from ..attention_dispatch import ( - _CAN_USE_FLEX_ATTN, - AttentionBackendName, - dispatch_attention_fn, -) +from ..attention_dispatch import AttentionBackendName from ..cache_utils import CacheMixin from ..modeling_outputs import Transformer2DModelOutput from ..modeling_utils import ModelMixin from torch import Tensor, nn -from torch.nn.attention.flex_attention import BlockMask + +from ..attention_processor import Kandinsky6AttnProcessor def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: @@ -111,59 +108,6 @@ def fractal_unflatten(x: Tensor, shape: tuple, block_mask: bool = False) -> Tens return x.reshape(*shape, x.shape[-1]) -def fast_sta_nabla( - T: int, - H: int, - W: int, - wT: int = 3, - wH: int = 3, - wW: int = 3, - device: str | torch.device = "cuda", -) -> Tensor: - """Precompute the Sliding Tile Attention (STA) boolean mask for NABLA attention.""" - l = max(T, H, W) - r = torch.arange(l, dtype=torch.int16, device=device) - mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs() - - sta_t = mat[:T, :T].flatten() <= wT // 2 - sta_h = mat[:H, :H].flatten() <= wH // 2 - sta_w = mat[:W, :W].flatten() <= wW // 2 - - sta_hw = (sta_h.unsqueeze(1) * sta_w.unsqueeze(0)).reshape(H, H, W, W).transpose(1, 2).flatten() - sta = (sta_t.unsqueeze(1) * sta_hw.unsqueeze(0)).reshape(T, T, H * W, H * W).transpose(1, 2) - return sta.reshape(T * H * W, T * H * W) - - -def nabla_block_mask( - q: Tensor, - k: Tensor, - sta: Tensor, - thr: float = 0.9, - block_size: int = 64, -) -> BlockMask: - """Build a dynamic NABLA block mask from query/key statistics and an STA prior.""" - B, h, S, D = q.shape - s1 = S // block_size - qa = q.reshape(B, h, s1, block_size, D).mean(-2) - ka = k.reshape(B, h, s1, block_size, D).mean(-2).transpose(-2, -1) - attn_map = torch.softmax((qa @ ka) / math.sqrt(D), dim=-1) - - vals, inds = attn_map.sort(-1) - mask = (vals.cumsum_(-1) >= 1 - thr).int().gather(-1, inds.argsort(-1)) - mask = torch.logical_or(mask, sta) - - kv_nb = mask.sum(-1).to(torch.int32) - kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) - return BlockMask.from_kv_blocks( - torch.zeros_like(kv_nb), - kv_inds, - kv_nb, - kv_inds, - BLOCK_SIZE=block_size, - mask_mod=None, - ) - - class RoPE1D(nn.Module): """1-D Rotary Position Embedding — used for text and audio sequences.""" @@ -468,14 +412,6 @@ def disable_cache(self) -> None: # native K6 names to this public Diffusers layout. -_MASKED_ATTENTION_BACKENDS = { - "flash": "flash_varlen", - "_flash_3": "_flash_varlen_3", - "sage": "sage_varlen", - "native": "native", -} - - class Kandinsky6TimeEmbeddings(nn.Module): """Sinusoidal timestep embedding with a K6-compatible parameter layout.""" @@ -579,106 +515,6 @@ def forward(self, x: Tensor) -> Tensor: return self.out_layer(self.activation(self.in_layer(x))) -class Kandinsky6AttnProcessor: - """Diffusers attention-processor boundary for K6 QKV projections.""" - - _attention_backend = None - _parallel_config = None - - def __init__(self, attention_backend=None, parallel_config=None): - if not hasattr(functional, "scaled_dot_product_attention"): - raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0 or newer.") - self._masked = False - self._attention_backend = attention_backend - self._parallel_config = parallel_config - - @property - def _attention_backend(self): - return self.__attention_backend - - @_attention_backend.setter - def _attention_backend(self, backend): - if self._masked and backend is not None: - name = getattr(backend, "value", backend).lower() - backend = _MASKED_ATTENTION_BACKENDS.get(name, name) - backend = AttentionBackendName(backend) - self.__attention_backend = backend - - def __call__( - self, - attn: Kandinsky6Attention, - hidden_states: Tensor, - encoder_hidden_states: Tensor | None = None, - rotary_emb: Tensor | None = None, - rotary_emb_kv: Tensor | None = None, - sparse_params: dict[str, Any] | None = None, - attn_mask: Tensor | None = None, - ) -> Tensor: - query = attn.to_query(hidden_states) - if encoder_hidden_states is None: - key = attn.to_key(hidden_states) - value = attn.to_value(hidden_states) - else: - key = attn.to_key(encoder_hidden_states) - value = attn.to_value(encoder_hidden_states) - - query = query.reshape(*query.shape[:-1], attn.num_heads, -1) - key = key.reshape(*key.shape[:-1], attn.num_heads, -1) - value = value.reshape(*value.shape[:-1], attn.num_heads, -1) - query = attn.query_norm(query) - key = attn.key_norm(key) - - if rotary_emb is not None: - query = apply_rotary(query, rotary_emb).to(dtype=query.dtype) - if rotary_emb_kv is not None: - key = apply_rotary(key, rotary_emb_kv).to(dtype=key.dtype) - - # Diffusers attention kernels use [B,S,H,D]. Text and unbatched - # cross-attention conditions are promoted to a singleton batch while - # packed visual streams already carry their batch dimension. - query_was_batched = query.dim() == 4 - if not query_was_batched: - query, key, value = query.unsqueeze(0), key.unsqueeze(0), value.unsqueeze(0) - elif key.dim() == 3: - key, value = key.unsqueeze(0), value.unsqueeze(0) - - if sparse_params is not None: - q = query.transpose(1, 2).contiguous() - k = key.transpose(1, 2).contiguous() - v = value.transpose(1, 2).contiguous() - block_mask = nabla_block_mask( - q, - k, - sparse_params["sta_mask"], - thr=sparse_params["P"], - ) - if not _CAN_USE_FLEX_ATTN: - raise ValueError("Nabla attention requires PyTorch 2.5 or newer") - output = ( - torch.nn.attention.flex_attention.flex_attention( - q, - k, - v, - block_mask=block_mask, - ) - .transpose(1, 2) - .contiguous() - ) - else: - output = dispatch_attention_fn( - query, - key, - value, - attn_mask=attn_mask, - backend=self._attention_backend, - parallel_config=self._parallel_config, - ) - - if not query_was_batched: - output = output[0] - return attn.out_layer(output.flatten(-2, -1)) - - class Kandinsky6Attention(nn.Module, AttentionModuleMixin): """K6 attention with the Diffusers ``set_processor`` contract.""" diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index ab5f3eefdad4..1f17c35a6587 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -19,17 +19,18 @@ import dataclasses import math from collections.abc import Mapping -from contextlib import contextmanager -from typing import TYPE_CHECKING, Any +from typing import Any import torch from ...configuration_utils import ConfigMixin, register_to_config +from ..attention import AttentionModuleMixin from ..modeling_utils import ModelMixin from diffusers.utils import logging -from flash_attn import flash_attn_varlen_func, flash_attn_varlen_qkvpacked_func from torch import Tensor, nn from torch.nn.attention.flex_attention import BlockMask, flex_attention +from ..attention_processor import Kandinsky6SRAttentionProcessor + # Side of the local 8x8 token block used by fractal (NABLA) attention. # Independent of the VAE spatial compression — do not swap for VAE_SPATIAL_FACTOR. FRACTAL_BLOCK_SIZE = 8 @@ -435,265 +436,11 @@ def framewise_causal_dense( return block_mask_from_bool(doc[None, None]) -# Pure-torch helpers for the streaming KV-cache in self-attention. -# -# Kept free of ``flash_attn`` imports so the cache bookkeeping is importable and -# testable on CPU. The actual attention kernel call lives in -# ``kandinsky_sr.model.nn``; this module only assembles its packed inputs. - - -if TYPE_CHECKING: - from collections.abc import Iterator - - from torch import nn - - -def assemble_cached_attention_inputs( - new_key: Tensor, - new_value: Tensor, - cached_key: Tensor, - cached_value: Tensor, - cu_seqlens_new: Tensor, - cached_cu_seqlens: Tensor, -) -> tuple[Tensor, Tensor, Tensor, Tensor]: - """Merge cached and new K/V into packed varlen attention inputs. - - For each packed sequence ``i`` the keys/values become ``[cache_i | new_i]``; - the query packing stays the new-token packing (queries are the freshly - arrived tokens that attend back over the cache plus themselves). - - Args: - new_key: New keys, packed ``(sum(new_lens), num_heads, head_dim)``. - new_value: New values, same packing as ``new_key``. - cached_key: Cached keys, packed ``(sum(cache_lens), num_heads, head_dim)``. - cached_value: Cached values, same packing as ``cached_key``. - cu_seqlens_new: Cumulative sequence lengths of the new tokens, ``(num_seqs + 1,)``. - cached_cu_seqlens: Cumulative sequence lengths of the cache, ``(num_seqs + 1,)``. - - Returns: - Tuple ``(key, value, cu_seqlens_q, cu_seqlens_k)`` where ``key``/``value`` - are the merged packed tensors, ``cu_seqlens_q`` is the query packing - (``= cu_seqlens_new``) and ``cu_seqlens_k`` the merged key packing. - """ - num_seqs = cu_seqlens_new.numel() - 1 - key_parts: list[Tensor] = [] - value_parts: list[Tensor] = [] - for i in range(num_seqs): - cache_start, cache_end = int(cached_cu_seqlens[i]), int(cached_cu_seqlens[i + 1]) - new_start, new_end = int(cu_seqlens_new[i]), int(cu_seqlens_new[i + 1]) - key_parts.append(cached_key[cache_start:cache_end]) - key_parts.append(new_key[new_start:new_end]) - value_parts.append(cached_value[cache_start:cache_end]) - value_parts.append(new_value[new_start:new_end]) - - key = torch.cat(key_parts, dim=0) - value = torch.cat(value_parts, dim=0) - - merged_lens = torch.diff(cached_cu_seqlens) + torch.diff(cu_seqlens_new) - cu_seqlens_k = torch.cat([cu_seqlens_new.new_zeros(1), torch.cumsum(merged_lens, dim=0)]).to(cu_seqlens_new.dtype) - cu_seqlens_q = cu_seqlens_new - return key, value, cu_seqlens_q, cu_seqlens_k - - -def evict_cached_kv( - cached_key: Tensor, - cached_value: Tensor, - cached_cu_seqlens: Tensor, - max_tokens_per_seq: int, -) -> tuple[Tensor, Tensor, Tensor]: - """Trim a packed KV-cache to its last ``max_tokens_per_seq`` tokens per sequence. - - Implements the fixed-length rolling window: for each packed sequence the - oldest tokens are dropped so at most ``max_tokens_per_seq`` most-recent ones - remain. Sequences already within the limit are kept whole. ``cu_seqlens`` is - rebuilt over the trimmed packing. Pure torch (no ``flash_attn``) so it stays - CPU-testable. - - Args: - cached_key: Cached keys, packed ``(sum(cache_lens), num_heads, head_dim)``. - cached_value: Cached values, same packing as ``cached_key``. - cached_cu_seqlens: Cumulative sequence lengths of the cache, ``(num_seqs + 1,)``. - max_tokens_per_seq: Maximum tokens to keep per packed sequence (the rolling - window length in tokens, i.e. ``cache_frames * H * W``). - - Returns: - Tuple ``(key, value, cu_seqlens)`` of the trimmed packed tensors and the - rebuilt cumulative sequence lengths. - """ - num_seqs = cached_cu_seqlens.numel() - 1 - key_parts: list[Tensor] = [] - value_parts: list[Tensor] = [] - kept_lens: list[int] = [] - for i in range(num_seqs): - start, end = int(cached_cu_seqlens[i]), int(cached_cu_seqlens[i + 1]) - keep_start = max(start, end - max_tokens_per_seq) - key_parts.append(cached_key[keep_start:end]) - value_parts.append(cached_value[keep_start:end]) - kept_lens.append(end - keep_start) - - key = torch.cat(key_parts, dim=0) - value = torch.cat(value_parts, dim=0) - lens = torch.tensor(kept_lens, device=cached_cu_seqlens.device) - cu_seqlens = torch.cat([cached_cu_seqlens.new_zeros(1), torch.cumsum(lens, dim=0)]).to(cached_cu_seqlens.dtype) - return key, value, cu_seqlens - - -def iter_self_attentions(model: nn.Module) -> Iterator[nn.Module]: - """Yield every submodule exposing the streaming KV-cache API. - - Duck-typed on the presence of ``reset_kv_cache`` so this module stays free of - ``flash_attn`` (which importing ``MultiheadSelfAttention`` would pull in) and - avoids a circular import with ``kandinsky_sr.model.nn``. - - Args: - model: Module to scan recursively. - - Yields: - Each self-attention submodule with a rolling KV-cache. - """ - for module in model.modules(): - if callable(getattr(module, "reset_kv_cache", None)): - yield module - - -def set_return_kv(model: nn.Module, *, flag: bool) -> None: - """Toggle ``return_kv`` on every self-attention submodule. - - Args: - model: Model whose self-attention layers to update. - flag: When ``True`` the next forward writes its merged K/V into the cache. - """ - for attn in iter_self_attentions(model): - attn.return_kv = flag - - -def reset_kv_caches(model: nn.Module) -> None: - """Clear the rolling KV-cache on every self-attention submodule. - - Args: - model: Model whose self-attention caches to reset. - """ - for attn in iter_self_attentions(model): - attn.reset_kv_cache() - - -@contextmanager -def kv_cache_session(model: nn.Module) -> Iterator[None]: - """Scope a streaming generation: caches are guaranteed clean before and after. - - Entry reset protects against state leaked by a previously crashed session; - the ``finally`` reset guarantees that even if generation dies mid-clip (e.g. - OOM) no layer keeps stale ``cached_k`` or a dangling ``return_kv=True``, so - subsequent non-streaming calls on the same model stay correct. - - Args: - model: Model whose self-attention caches to scope. - - Yields: - Nothing; run the chunked generation loop inside the block. - """ - reset_kv_caches(model) - try: - yield - finally: - reset_kv_caches(model) - - -@contextmanager -def kv_writer(model: nn.Module) -> Iterator[None]: - """Mark the enclosed forward as the cache writer (``return_kv=True``). - - ``return_kv`` is switched off in a ``finally`` so a crash inside the writer - forward cannot leave the model silently appending K/V on later calls. - - Args: - model: Model whose self-attention layers write K/V inside the block. - - Yields: - Nothing; run the writer forward (and its state capture) inside the block. - """ - set_return_kv(model, flag=True) - try: - yield - finally: - set_return_kv(model, flag=False) - - -# A captured per-step cache "slot": one ``(key, value, cu_seqlens)`` triple per -# self-attention layer (aligned with ``iter_self_attentions`` order), or ``None`` -# for an empty slot (no past frames cached yet at that denoising step). -LayerCache = tuple["Tensor | None", "Tensor | None", "Tensor | None"] -KVSlot = "list[LayerCache] | None" - - -def capture_kv_state(model: nn.Module) -> list[LayerCache]: - """Snapshot every self-attention's rolling KV-cache. - - Stores the current ``(cached_k, cached_v, cached_cu_seqlens)`` references for - each self-attention layer. The references are safe to hold: the attention - forward reassigns these attributes to freshly concatenated tensors rather - than mutating them in place, so a snapshot never aliases future writes. - - Args: - model: Model whose self-attention caches to snapshot. - - Returns: - Per-layer ``(key, value, cu_seqlens)`` triples in ``iter_self_attentions`` order. - """ - return [(attn.cached_k, attn.cached_v, attn.cached_cu_seqlens) for attn in iter_self_attentions(model)] - - -def restore_kv_state(model: nn.Module, slot: KVSlot) -> None: - """Load a captured slot (or clear) into every self-attention's cache. - - Args: - model: Model whose self-attention caches to overwrite. - slot: Per-layer ``(key, value, cu_seqlens)`` triples from - :func:`capture_kv_state`, or ``None`` to clear all caches (an empty - slot — no past frames yet). - """ - attns = list(iter_self_attentions(model)) - if slot is None: - for attn in attns: - attn.cached_k, attn.cached_v, attn.cached_cu_seqlens = None, None, None - return - for attn, (key, value, cu_seqlens) in zip(attns, slot, strict=True): - attn.cached_k, attn.cached_v, attn.cached_cu_seqlens = key, value, cu_seqlens - - -def evict_kv_slot(slot: KVSlot, max_tokens_per_seq: int) -> KVSlot: - """Apply fixed-length rolling eviction to a captured slot, per layer. - - Args: - slot: A captured per-step slot, or ``None`` (empty slot passes through). - max_tokens_per_seq: Rolling window length in tokens (``cache_frames * H * W``). - - Returns: - The evicted slot, or ``None`` if the input was ``None``. - """ - if slot is None: - return None - evicted: list[LayerCache] = [] - for key, value, cu_seqlens in slot: - if key is None: - evicted.append((None, None, None)) - else: - evicted.append(evict_cached_kv(key, value, cu_seqlens, max_tokens_per_seq)) - return evicted - - # Neural network building blocks for the diffusion transformer. logger = logging.get_logger(__name__) -try: - import flash_attn_interface # pyright: ignore - - FA3 = True -except ImportError: - FA3 = False - def _ensure_nabla_compatible_flex_bwd_configs() -> None: """Guarantee flex_attention's backward keeps a 64-block-compatible config. @@ -777,7 +524,6 @@ def patched(self: Any, head_dim: int, dtype: Any, device_type: str = "cuda") -> # only the FA call runs in eager. When FA3 is available the dedicated # ``flash_attn_interface.flash_attn_varlen_func`` does register a custom op # and traces cleanly, so this decorator becomes a no-op. -_disable_dynamo_if_fa2 = (lambda fn: fn) if FA3 else torch.compiler.disable @torch.autocast(device_type="cuda", dtype=torch.float32) @@ -1257,7 +1003,9 @@ def reset_parameters(self) -> None: self.out_layer.bias.data.zero_() -class MultiheadSelfAttention(nn.Module): +class MultiheadSelfAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = Kandinsky6SRAttentionProcessor + _available_processors = [Kandinsky6SRAttentionProcessor] """Multi-head self-attention with flash attention and optional sparse flex attention.""" def __init__(self, num_channels: int, head_dim: int) -> None: @@ -1285,6 +1033,7 @@ def __init__(self, num_channels: int, head_dim: int) -> None: self.cached_v: Tensor | None = None self.cached_cu_seqlens: Tensor | None = None self.return_kv = False + self.set_processor(self._default_processor_cls()) def get_qkv(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: """Project input into query, key, and value tensors. @@ -1320,125 +1069,6 @@ def norm_qk(self, q: Tensor, k: Tensor) -> tuple[Tensor, Tensor]: k = self.key_norm(k.float()).type_as(k) return q, k - @_disable_dynamo_if_fa2 - def scaled_dot_product_attention( - self, - query: Tensor, - key: Tensor, - value: Tensor, - cu_seqlens: Tensor, - *, - return_attn_probs: bool = False, - ) -> Tensor | tuple[Tensor, Tensor, None]: - """Compute self-attention using flash attention. - - When a KV-cache is present (``cached_k is not None``) the queries are the - new tokens and attention runs over ``[cache | new]`` keys/values via - :meth:`attend_with_kv_cache`. When ``return_kv`` is set, the (merged) K/V - are stored as the cache for the next streaming step. - - Args: - query: Query tensor. - key: Key tensor. - value: Value tensor. - cu_seqlens: Cumulative sequence lengths for packed sequences. - return_attn_probs: Whether to return attention log-sum-exp values. - - Returns: - Attention output, or tuple of (output, softmax_lse, None) if returning probs. - """ - if self.cached_k is not None: - return self.attend_with_kv_cache(query, key, value, cu_seqlens, return_attn_probs=return_attn_probs) - - if self.return_kv: - self.cached_k, self.cached_v, self.cached_cu_seqlens = key, value, cu_seqlens - - max_seqlen = torch.diff(cu_seqlens).max() - if FA3: - out, softmax_lse = flash_attn_interface.flash_attn_varlen_func( - q=query, - k=key, - v=value, - cu_seqlens_q=cu_seqlens, - cu_seqlens_k=cu_seqlens, - max_seqlen_q=max_seqlen, - max_seqlen_k=max_seqlen, - ) - else: - query_key_value = torch.stack([query, key, value], dim=-3) - out, softmax_lse, _ = flash_attn_varlen_qkvpacked_func( - query_key_value, cu_seqlens, max_seqlen, return_attn_probs=True - ) - out = out.flatten(-2, -1) - - if return_attn_probs: - return out, softmax_lse, None - return out - - @torch.compiler.disable - def attend_with_kv_cache( - self, - query: Tensor, - key: Tensor, - value: Tensor, - cu_seqlens: Tensor, - *, - return_attn_probs: bool = False, - ) -> Tensor | tuple[Tensor, Tensor, None]: - """Dense flash attention consuming the rolling KV-cache. - - Queries are the new tokens; keys/values are the per-sequence concatenation - ``[cache | new]`` (with absolute RoPE already applied upstream). When - ``return_kv`` is set, the merged K/V become the cache for the next - streaming step. Kept out of ``torch.compile`` because the cache assembly - indexes ``cu_seqlens`` host-side; it only runs during streaming inference. - - Args: - query: New-token query tensor. - key: New-token key tensor (post-RoPE). - value: New-token value tensor. - cu_seqlens: Cumulative sequence lengths of the new tokens. - return_attn_probs: Whether to return attention log-sum-exp values. - - Returns: - Attention output for the new tokens, or tuple of (output, softmax_lse, - None) if returning probs. - """ - key, value, cu_seqlens_q, cu_seqlens_k = assemble_cached_attention_inputs( - key, value, self.cached_k, self.cached_v, cu_seqlens, self.cached_cu_seqlens - ) - if self.return_kv: - self.cached_k, self.cached_v, self.cached_cu_seqlens = key, value, cu_seqlens_k - - max_seqlen_q = torch.diff(cu_seqlens_q).max() - max_seqlen_k = torch.diff(cu_seqlens_k).max() - if FA3: - out, softmax_lse = flash_attn_interface.flash_attn_varlen_func( - q=query, - k=key, - v=value, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_k, - ) - else: - out, softmax_lse, _ = flash_attn_varlen_func( - q=query, - k=key, - v=value, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_k, - return_attn_probs=True, - ) - out = out.flatten(-2, -1) - - if return_attn_probs: - return out, softmax_lse, None - return out - def reset_kv_cache(self) -> None: """Clear the rolling KV-cache and stop emitting K/V.""" self.cached_k = None @@ -1559,10 +1189,7 @@ def forward( query = apply_rotary(query, rope).type_as(query) key = apply_rotary(key, rope).type_as(key) - if sparse_params is not None: - out = self.attention_flex(query, key, value, sparse_params=sparse_params) - else: - out = self.scaled_dot_product_attention(query, key, value, cu_seqlens) + out = self.processor(self, query, key, value, cu_seqlens, cu_seqlens, sparse_params=sparse_params) return self.out_layer(out) @@ -1578,7 +1205,9 @@ def reset_parameters(self) -> None: self.key_norm.reset_parameters() -class MultiheadCrossAttention(nn.Module): +class MultiheadCrossAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = Kandinsky6SRAttentionProcessor + _available_processors = [Kandinsky6SRAttentionProcessor] """Multi-head cross-attention with flash attention.""" def __init__(self, num_channels: int, head_dim: int) -> None: @@ -1601,6 +1230,7 @@ def __init__(self, num_channels: int, head_dim: int) -> None: self.key_norm = nn.RMSNorm(head_dim) self.out_layer = nn.Linear(num_channels, num_channels, bias=True) + self.set_processor(self._default_processor_cls()) def get_qkv(self, x: Tensor, cond: Tensor) -> tuple[Tensor, Tensor, Tensor]: """Project input and condition into query, key, and value tensors. @@ -1637,41 +1267,6 @@ def norm_qk(self, q: Tensor, k: Tensor) -> tuple[Tensor, Tensor]: k = self.key_norm(k.float()).type_as(k) return q, k - @_disable_dynamo_if_fa2 - def scaled_dot_product_attention( - self, - query: Tensor, - key: Tensor, - value: Tensor, - cu_seqlens: Tensor, - cond_cu_seqlens: Tensor, - *, - return_attn_probs: bool = False, - ) -> Tensor | tuple[Tensor, Tensor, None]: - """Compute cross-attention using flash attention. - - Args: - query: Query tensor. - key: Key tensor from conditioning input. - value: Value tensor from conditioning input. - cu_seqlens: Cumulative sequence lengths for query sequences. - cond_cu_seqlens: Cumulative sequence lengths for conditioning sequences. - return_attn_probs: Whether to return attention log-sum-exp values. - - Returns: - Attention output, or tuple of (output, softmax_lse, None) if returning probs. - """ - max_seqlen = torch.diff(cu_seqlens).max() - cond_max_seqlen = torch.diff(cond_cu_seqlens).max() - out, softmax_lse, _ = flash_attn_varlen_func( - query, key, value, cu_seqlens, cond_cu_seqlens, max_seqlen, cond_max_seqlen, return_attn_probs=True - ) - out = out.flatten(-2, -1) - - if return_attn_probs: - return out, softmax_lse, None - return out - def forward(self, x: Tensor, cond: Tensor, cu_seqlens: Tensor, cond_cu_seqlens: Tensor) -> Tensor: """Run cross-attention between input and conditioning. @@ -1687,7 +1282,7 @@ def forward(self, x: Tensor, cond: Tensor, cu_seqlens: Tensor, cond_cu_seqlens: query, key, value = self.get_qkv(x, cond) query, key = self.norm_qk(query, key) - out = self.scaled_dot_product_attention(query, key, value, cu_seqlens, cond_cu_seqlens) + out = self.processor(self, query, key, value, cu_seqlens, cond_cu_seqlens) return self.out_layer(out) def reset_parameters(self) -> None: @@ -2156,6 +1751,7 @@ class Kandinsky6SRTransformer3DModel(ModelMixin, ConfigMixin): """ _no_split_modules = ["TransformerEncoderBlock", "TransformerDecoderBlock"] + _repeated_blocks = ["TransformerEncoderBlock", "TransformerDecoderBlock"] @register_to_config def __init__( From efa2b56d1406d271bccd0788c4934522ab17fa6c Mon Sep 17 00:00:00 2001 From: Denis Koposov Date: Wed, 16 Sep 2026 21:19:13 +0000 Subject: [PATCH 06/62] fix unguarded dependencies - removed: einops, pydantic - lazy: torchvision, av, librosa --- docs/source/en/api/models/kandinsky6_vae.md | 5 - .../autoencoders/autoencoder_kandinsky6_sr.py | 14 +- .../autoencoders/autoencoder_mmaudio.py | 25 +- .../pipelines/kandinsky6/latent_upscaler.py | 393 +++++------------- .../kandinsky6/pipeline_kandinsky6_sr.py | 65 ++- .../kandinsky6/pipeline_kandinsky6_ti2va.py | 14 +- 6 files changed, 194 insertions(+), 322 deletions(-) diff --git a/docs/source/en/api/models/kandinsky6_vae.md b/docs/source/en/api/models/kandinsky6_vae.md index 933270d78e5b..571c407a0585 100644 --- a/docs/source/en/api/models/kandinsky6_vae.md +++ b/docs/source/en/api/models/kandinsky6_vae.md @@ -22,8 +22,3 @@ with a vocoder for synchronized audio generation. - encode - decode -## MMAudioVAE - -[[autodoc]] MMAudioVAE - - wrapped_encode - - wrapped_decode diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index 4071b3244aa4..2ff0a11ee011 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -29,8 +29,6 @@ from ...configuration_utils import ConfigMixin, register_to_config from ..modeling_utils import ModelMixin from ...utils.accelerate_utils import apply_forward_hook -from einops import rearrange -from safetensors.torch import load_file as safetensors_load_file from torch.utils.checkpoint import checkpoint @@ -434,7 +432,7 @@ def __init__( def spatial_downsample(self, input_): # PixelShuffle part - pxs_input = rearrange(input_, "b c t h w -> (b t) c h w") + pxs_input = input_.permute(0, 2, 1, 3, 4).reshape(-1, input_.shape[1], input_.shape[3], input_.shape[4]) pxs_interm = self.unshuffle(pxs_input) b, c, h, w = pxs_interm.shape if self.version > 1: @@ -442,7 +440,9 @@ def spatial_downsample(self, input_): else: # pxs_interm_view = pxs_interm.view(b, c // self.factor**2, self.factor**2, h, w) pxs_out = torch.mean(pxs_interm_view, dim=2) - pxs_out = rearrange(pxs_out, "(b t) c h w -> b c t h w", t=input_.size(2)) + pxs_out = pxs_out.reshape( + input_.shape[0], input_.size(2), pxs_out.shape[1], pxs_out.shape[2], pxs_out.shape[3] + ).permute(0, 2, 1, 3, 4) # Downsampling by 3D-convolution conv_out = self.spatial_conv(input_) @@ -452,7 +452,7 @@ def spatial_downsample(self, input_): def temporal_downsample(self, input_, cache): # Interpolation part - permuted = rearrange(input_, "b c t h w -> (b h w) c t") + permuted = input_.permute(0, 3, 4, 1, 2).reshape(-1, input_.shape[1], input_.shape[2]) if cache[0]["padding"] is None: first, rest = permuted[..., :1], permuted[..., 1:] @@ -466,7 +466,9 @@ def temporal_downsample(self, input_, cache): if rest.size(-1) > 0: full_interp = F.avg_pool1d(rest, kernel_size=2, stride=2) - full_interp = rearrange(full_interp, "(b h w) c t -> b c t h w", h=input_.size(-2), w=input_.size(-1)) + full_interp = full_interp.reshape( + input_.shape[0], input_.size(-2), input_.size(-1), full_interp.shape[1], full_interp.shape[2] + ).permute(0, 3, 4, 1, 2) # Downsampling by convolution if self.version == 1: diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index a2a749702103..d16685c53b11 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -13,8 +13,6 @@ from ...configuration_utils import ConfigMixin, register_to_config from ..modeling_utils import ModelMixin from ...utils.accelerate_utils import apply_forward_hook -from einops import rearrange -from librosa.filters import mel as librosa_mel_fn from torch import pow, sin from torch.nn import Conv1d, ConvTranspose1d, Parameter from torch.nn.utils.parametrizations import weight_norm @@ -281,12 +279,12 @@ def forward(self, x): y = y.reshape(y.shape[0], self.num_heads, -1, 3, y.shape[-1]) q, k, v = normalize(y, dim=2).unbind(3) - q = rearrange(q, "b h c l -> b h l c") - k = rearrange(k, "b h c l -> b h l c") - v = rearrange(v, "b h c l -> b h l c") + q = q.permute(0, 1, 3, 2) + k = k.permute(0, 1, 3, 2) + v = v.permute(0, 1, 3, 2) h = F.scaled_dot_product_attention(q, k, v) - h = rearrange(h, "b h l c -> b (h c) l") + h = h.permute(0, 1, 3, 2).reshape(h.shape[0], -1, h.shape[2]) h = self.proj_out(h) @@ -1120,7 +1118,7 @@ def forward(self, x): def remove_weight_norm(self): for l in self.convs: - remove_weight_norm(l) + remove_parametrizations(l, "weight") TorchActivation1d = Activation1d @@ -1332,8 +1330,19 @@ def __init__( self.fmax = fmax self.norm_fn = norm_fn + try: + from librosa.filters import mel as librosa_mel_fn + except (ImportError, OSError) as exc: + raise RuntimeError( + "MMAudio audio features require librosa. Install it with `pip install librosa`." + ) from exc + mel = librosa_mel_fn( - sr=self.sampling_rate, n_fft=self.n_fft, n_mels=self.num_mels, fmin=self.fmin, fmax=self.fmax + sr=self.sampling_rate, + n_fft=self.n_fft, + n_mels=self.num_mels, + fmin=self.fmin, + fmax=self.fmax, ) mel_basis = torch.from_numpy(mel).float() hann_window = torch.hann_window(self.win_size) diff --git a/src/diffusers/pipelines/kandinsky6/latent_upscaler.py b/src/diffusers/pipelines/kandinsky6/latent_upscaler.py index 854baff7fdf2..74e54d57965f 100644 --- a/src/diffusers/pipelines/kandinsky6/latent_upscaler.py +++ b/src/diffusers/pipelines/kandinsky6/latent_upscaler.py @@ -22,19 +22,63 @@ import math from collections import defaultdict from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace from functools import partial -from typing import TYPE_CHECKING, Annotated, Any, Literal, NamedTuple +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any, Literal, NamedTuple -import pydantic import torch from ...configuration_utils import ConfigMixin, register_to_config from ...models.modeling_utils import ModelMixin -from einops import rearrange -from pydantic import BaseModel, Field, TypeAdapter from torch import Tensor, nn from torch.autograd import Function from torch.nn import functional -from torchvision.ops import StochasticDepth + + +def rearrange(*args: Any, **kwargs: Any) -> Any: + """Load einops only when the latent upscaler is actually executed.""" + try: + from einops import rearrange as _rearrange # noqa: PLC0415 + except (ImportError, OSError) as exc: + raise RuntimeError("SR latent upscaler requires einops. Install with `pip install einops`.") from exc + return _rearrange(*args, **kwargs) + + +def _model_config(value: Mapping[str, Any]) -> Any: + """Convert serialized model mappings to attribute-accessible config objects.""" + if not isinstance(value, Mapping): + raise TypeError(f"latent upscaler model config must be a mapping, got {type(value).__name__}") + + def convert(item: Any) -> Any: + if isinstance(item, Mapping): + return SimpleNamespace(**{str(key): convert(val) for key, val in item.items()}) + if isinstance(item, list): + return tuple(convert(val) for val in item) + return item + + config = convert(value) + if getattr(config, "architecture", None) == "multi_scale": + # Older/published multi-scale configs omit fields whose native + # constructor defaults are sufficient. + defaults = { + "icnr": False, + "global_skip": False, + "x2_adapter_blocks": 0, + "x2_tail_mode": "shared", + "x2_adapter_sources": None, + "x2_finisher": "none", + "motion_attention": None, + } + for key, default in defaults.items(): + if not hasattr(config, key): + setattr(config, key, default) + return config + + +UpsampleMode = str +UpsamplePosition = str +FlatModelConfig = Any +MultiScaleModelConfig = Any @torch.autocast(device_type="cuda", enabled=False) @@ -178,6 +222,28 @@ def make_conv(dims: int, temporal_padding: TemporalPadding) -> type[nn.Module]: # Reusable building blocks for latent upsampler architectures. +class StochasticDepth(nn.Module): + """Drop residual paths without requiring torchvision.""" + + def __init__(self, p: float, mode: str = "row") -> None: + super().__init__() + if not 0.0 <= p <= 1.0: + raise ValueError(f"drop probability must be in [0, 1], got {p}") + if mode not in ("batch", "row"): + raise ValueError(f"mode must be 'batch' or 'row', got {mode!r}") + self.p = p + self.mode = mode + + def forward(self, x: Tensor) -> Tensor: + if not self.training or self.p == 0.0: + return x + if self.p == 1.0: + return torch.zeros_like(x) + shape = (x.shape[0],) + (1,) * (x.ndim - 1) if self.mode == "row" else (1,) * x.ndim + noise = torch.empty(shape, dtype=x.dtype, device=x.device).bernoulli_(1.0 - self.p) + return x * noise / (1.0 - self.p) + + class RMSNorm(nn.Module): """Channel-first Root Mean Square normalization with learnable gamma. @@ -949,7 +1015,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: NattenBackend = Literal["cutlass-fna", "hopper-fna", "flex-fna"] -class MotionCorrespondenceSpec(BaseModel): +@dataclass +class MotionCorrespondenceSpec: """Construction parameters for one motion-correspondence block.""" channels: int @@ -1248,7 +1315,8 @@ def load_natten() -> ModuleType: STAGE_FACTOR = 2 -class UpsampleSpec(BaseModel): +@dataclass +class UpsampleSpec: """Configuration shared by the two spatial upsample stages.""" mode: Literal["pixel_shuffle", "bilinear", "pxs_v2", "pxs_v2_hybrid"] @@ -1259,7 +1327,8 @@ class UpsampleSpec(BaseModel): padding_mode: UpsamplePaddingMode = "reflect" -class ResidualStackSpec(BaseModel): +@dataclass +class ResidualStackSpec: """Configuration shared by residual stacks at all model resolutions. in_channels: input width of the stack's first block; None means use channels. @@ -1279,7 +1348,8 @@ class ResidualStackSpec(BaseModel): temporal_padding: TemporalPadding = "zeros" -class MotionAttentionSpec(BaseModel): +@dataclass +class MotionAttentionSpec: """Construction and placement contract for motion-correspondence blocks.""" after_mid_blocks: tuple[int, ...] @@ -2097,11 +2167,10 @@ def stage_spec(channels: int, first_in: int | None) -> ResidualStackSpec: self.mid_input_proj = build_stem( in_channels, hidden_channels, dims, bare=bare_stem, temporal_padding=temporal_padding ) - adapter_spec = stage_spec(w1, None).model_copy( - update={ - "total_blocks": max(x2_adapter_blocks, 1), - "stochastic_depth_rate": 0.0, - } + adapter_spec = replace( + stage_spec(w1, None), + total_blocks=max(x2_adapter_blocks, 1), + stochastic_depth_rate=0.0, ) adapter = build_residual_stack(x2_adapter_blocks, 0, adapter_spec) for block in adapter: @@ -2288,276 +2357,6 @@ def _add_nearest_base( return out + spatial_nearest_2x(z, DIMS_3, self.upscale_factor) -# Configuration models for the latent upscaler. - - -UpsampleMode = Literal["pixel_shuffle", "bilinear", "pxs_v2", "pxs_v2_hybrid"] -UpsamplePosition = Literal["before_projection", "after_projection"] -X2TailMode = Literal["shared", "scale_specific_norm", "private", "private_full"] -X2FinisherMode = Literal["none", "linear", "pxs_residual"] -MotionAttentionBackend = Literal["sdpa", "natten"] -NattenBackend = Literal["cutlass-fna", "hopper-fna", "flex-fna"] -SPATIOTEMPORAL_DIMS = 3 -MIN_ROPE_HEAD_DIM = 6 - - -class _BaseModelConfig(BaseModel): - """Common fields shared by all upsampler architectures.""" - - model_config = pydantic.ConfigDict(extra="forbid") - - in_channels: int = 16 - upscale_factor: int = 2 - upsample_mode: UpsampleMode = "pixel_shuffle" - upsample_position: UpsamplePosition = "after_projection" - temporal_mix: bool = True - icnr: bool = False - dims: int = 2 - input_skip: bool = True - gradient_checkpointing: bool = False - modulated_norm: bool = False - temporal_padding: TemporalPadding = "zeros" # "replicate" pads T both sides, "causal" pads the past only (K-VAE) - upsample_padding_mode: UpsamplePaddingMode = "reflect" # "zeros" matches the production K-VAE upsample - - -class FlatModelConfig(_BaseModelConfig): - """Configuration for ConvLatentUpsampler (flat architecture).""" - - architecture: Literal["flat"] = "flat" - hidden_channels: int = 64 - stem_channels: int | None = None - num_residual_blocks: int = 16 - num_pre_residual_blocks: int = 0 - bottleneck_channels: int | None = None - expand_ratio: int = 4 - kernel_size: int = 3 - layer_scale_init: float | None = None - grn: bool = False - stochastic_depth_rate: float = 0.0 - depthwise: bool = False - - -class MotionAttentionConfig(BaseModel): - """Configuration for cross-frame spatial correspondence blocks.""" - - model_config = pydantic.ConfigDict(extra="forbid") - - spatial_kernel_size: int = Field(default=13, ge=3) - temporal_offsets: tuple[int, ...] = (-3, -2, -1, 1, 2, 3) - after_mid_blocks: tuple[int, ...] = (4, 8, 12) - num_heads: int = Field(default=3, gt=0) - backend: MotionAttentionBackend = "natten" - natten_backend: NattenBackend = "hopper-fna" - merge_compile: bool = False - - @pydantic.model_validator(mode="after") - def validate_attention_geometry(self) -> MotionAttentionConfig: - """Require centered spatial windows and symmetric temporal offsets.""" - if self.spatial_kernel_size % 2 == 0: - msg = "motion spatial_kernel_size must be odd" - raise ValueError(msg) - if not self.temporal_offsets or 0 in self.temporal_offsets: - msg = "motion temporal_offsets must be non-empty and exclude 0" - raise ValueError(msg) - if tuple(sorted(set(self.temporal_offsets))) != self.temporal_offsets: - msg = "motion temporal_offsets must be unique and sorted" - raise ValueError(msg) - if set(self.temporal_offsets) != {-offset for offset in self.temporal_offsets}: - msg = "motion temporal_offsets must be symmetric" - raise ValueError(msg) - if ( - not self.after_mid_blocks - or any(block_number <= 0 for block_number in self.after_mid_blocks) - or tuple(sorted(set(self.after_mid_blocks))) != self.after_mid_blocks - ): - msg = "motion after_mid_blocks must contain positive, unique, sorted block numbers" - raise ValueError(msg) - return self - - -class MultiScaleModelConfig(_BaseModelConfig): - """Configuration for MultiScaleUpsampler (cascaded 2x+2x architecture). - - Decomposes a 4x upscale into two supervised 2x stages: coarse structure - at 2x, fine detail at 4x. Each stage gets direct loss supervision. - - Attributes: - hidden_channels: Internal width of residual blocks. - num_pre_blocks: Residual blocks at 1x (before first upsample). - num_mid_blocks: Residual blocks at 2x (between upsamples). - num_post_blocks: Residual blocks at 4x (after second upsample). - bottleneck_channels: Mid-channel width for ResidualBlocks (None = hidden * expand_ratio). - expand_ratio: Inverted bottleneck expansion when bottleneck_channels is None. - kernel_size: Kernel size for depthwise convolutions. - layer_scale_init: Initial LayerScale gamma. None disables LayerScale. - grn: Insert Global Response Normalization in residual blocks. - stochastic_depth_rate: Max drop rate, linearly scheduled across all stages. - depthwise: Use depthwise separable convolutions. - motion_attention: Optional cross-frame spatial correspondence stack. - bare_stem: Reduce input projections to a single conv (no norm/activation). - modulated_output_proj: Use a zq-conditioned ModulatedRMSNorm in ``output_proj``. - global_skip: Add a nearest-upsampled input base to every output and - zero-init the output heads, so the model starts as a nearest upsampler - and learns only the residual correction. - loss_weight_2x: Training loss weight for 2x intermediate output. - loss_weight_4x: Training loss weight for 4x final output. - enable_x2_entry: Build the ``mid_input_proj`` submodule and allow - ``forward(entry="x2")`` — a second entry that injects an already-2x LQ - latent after ``mid_blocks`` for a single 2x upscale through the shared - second stage (LAY-462). Off by default so the x4 module set and its - checkpoints stay unchanged. - x2_adapter_blocks: Number of x2-only residual blocks before the second stage. - x2_tail_mode: Amount of scale-specific capacity in the second stage; - ``private_full`` additionally privatizes ``mid_blocks`` and runs them - in the x2 path ahead of the second-stage upsample. - x2_adapter_sources: pre_blocks indices warm-copied into the adapter (one - per block); None keeps the zero-init identity adapter. - x2_finisher: upsample_1 conv pair applied at the adapter grid without the - nearest-2x resize: "linear" clones the 1x1 only, "pxs_residual" adds - the spatial conv with its residual add. - stage_channels: Per-grid widths (1x, 2x, 4x) for the K-VAE-shaped pyramid; None keeps constant hidden_channels. - """ - - architecture: Literal["multi_scale"] = "multi_scale" - upscale_factor: int = 4 - hidden_channels: int = 64 - num_pre_blocks: int = 2 - num_mid_blocks: int = 6 - num_post_blocks: int = 8 - bottleneck_channels: int | None = None - expand_ratio: int = 4 - kernel_size: int = 3 - layer_scale_init: float | None = None - grn: bool = False - stochastic_depth_rate: float = 0.0 - depthwise: bool = False - motion_attention: MotionAttentionConfig | None = None - bare_stem: bool = False - modulated_output_proj: bool = False - global_skip: bool = False - loss_weight_2x: float = 0.5 - loss_weight_4x: float = 1.0 - enable_x2_entry: bool = False - x2_adapter_blocks: int = Field(default=0, ge=0) - x2_tail_mode: X2TailMode = "shared" - x2_adapter_sources: tuple[int, ...] | None = None - x2_finisher: X2FinisherMode = "none" - stage_channels: tuple[int, int, int] | None = None - - @pydantic.model_validator(mode="after") - def validate_upscale_factor(self) -> MultiScaleModelConfig: - """Multi-scale architecture only supports 4x upscale (two 2x stages).""" - expected = 4 - if self.upscale_factor != expected: - msg = f"multi_scale architecture requires upscale_factor={expected}, got {self.upscale_factor}" - raise ValueError(msg) - return self - - @pydantic.model_validator(mode="after") - def validate_redesign_flags(self) -> MultiScaleModelConfig: - """Reject redesign-flag combinations the model cannot execute.""" - if self.modulated_output_proj and not self.modulated_norm: - msg = "modulated_output_proj requires modulated_norm=true" - raise ValueError(msg) - if self.enable_x2_entry and self.global_skip: - msg = "global_skip is x4-only and incompatible with enable_x2_entry" - raise ValueError(msg) - if self.enable_x2_entry and self.modulated_output_proj and self.x2_tail_mode != "private_full": - msg = "modulated_output_proj with enable_x2_entry requires x2_tail_mode='private_full'" - raise ValueError(msg) - return self - - @pydantic.model_validator(mode="after") - def validate_stage_channels(self) -> MultiScaleModelConfig: - """Pyramid widths must agree with hidden_channels and exclude width-coupled features.""" - if self.stage_channels is None: - return self - if self.hidden_channels != self.stage_channels[0]: - msg = f"hidden_channels ({self.hidden_channels}) must equal stage_channels[0] ({self.stage_channels[0]})" - raise ValueError(msg) - if self.enable_x2_entry and self.x2_tail_mode != "private_full": - msg = "stage_channels with enable_x2_entry requires x2_tail_mode='private_full'" - raise ValueError(msg) - if self.motion_attention is not None: - msg = "stage_channels is incompatible with motion_attention (motion blocks are built at hidden_channels)" - raise ValueError(msg) - if self.input_skip: - msg = "stage_channels is incompatible with input_skip (channel-repeat skip assumes one width)" - raise ValueError(msg) - if self.bottleneck_channels is not None: - msg = "stage_channels derives per-stage mid widths from expand_ratio; bottleneck_channels is incompatible" - raise ValueError(msg) - if self.num_mid_blocks < 1 or self.num_post_blocks < 1: - msg = "stage_channels needs at least one mid and one post block to host the width transitions" - raise ValueError(msg) - return self - - @pydantic.model_validator(mode="after") - def validate_x2_adaptation(self) -> MultiScaleModelConfig: - """Reject x2 adaptation settings that cannot be executed.""" - has_adaptation = ( - self.x2_adapter_blocks > 0 - or self.x2_tail_mode != "shared" - or self.x2_adapter_sources is not None - or self.x2_finisher != "none" - ) - if has_adaptation and not self.enable_x2_entry: - msg = "x2 adaptation requires enable_x2_entry=true" - raise ValueError(msg) - if self.x2_tail_mode != "shared" and self.x2_adapter_blocks == 0: - msg = f"x2_tail_mode={self.x2_tail_mode!r} requires x2_adapter_blocks > 0" - raise ValueError(msg) - if self.x2_tail_mode == "scale_specific_norm" and not self.modulated_norm: - msg = "scale_specific_norm requires modulated_norm=true" - raise ValueError(msg) - if self.x2_tail_mode == "scale_specific_norm" and self.depthwise: - msg = "scale_specific_norm does not support depthwise residual blocks" - raise ValueError(msg) - if self.x2_adapter_sources is not None and len(self.x2_adapter_sources) != self.x2_adapter_blocks: - msg = ( - f"x2_adapter_sources length ({len(self.x2_adapter_sources)}) must equal " - f"x2_adapter_blocks ({self.x2_adapter_blocks})" - ) - raise ValueError(msg) - if self.x2_adapter_sources is not None and any( - not 0 <= index < self.num_pre_blocks for index in self.x2_adapter_sources - ): - msg = f"x2_adapter_sources indices must be within [0, num_pre_blocks={self.num_pre_blocks})" - raise ValueError(msg) - if self.x2_finisher != "none" and self.upsample_mode != "pxs_v2": - msg = "x2_finisher clones the pxs_v2 upsample convs and requires upsample_mode='pxs_v2'" - raise ValueError(msg) - return self - - @pydantic.model_validator(mode="after") - def validate_motion_attention(self) -> MultiScaleModelConfig: - """Reject motion-attention configurations incompatible with this model.""" - motion = self.motion_attention - if motion is None: - return self - if self.dims != SPATIOTEMPORAL_DIMS: - msg = "motion_attention requires dims=3" - raise ValueError(msg) - if self.enable_x2_entry: - msg = "motion_attention Round 23 is x4-only and cannot enable the x2 entry" - raise ValueError(msg) - if motion.after_mid_blocks[-1] > self.num_mid_blocks: - msg = "motion after_mid_blocks must not exceed num_mid_blocks" - raise ValueError(msg) - if self.hidden_channels % motion.num_heads != 0: - msg = "hidden_channels must be divisible by motion num_heads" - raise ValueError(msg) - head_dim = self.hidden_channels // motion.num_heads - if head_dim < MIN_ROPE_HEAD_DIM or head_dim % 2 != 0: - msg = "motion attention head_dim must be even and at least 6 for axial 3D RoPE" - raise ValueError(msg) - return self - - -ModelConfig = Annotated[ - FlatModelConfig | MultiScaleModelConfig, - Field(discriminator="architecture"), -] # Factory functions for building upsampler models. @@ -2578,7 +2377,7 @@ def build_upsampler(config: FlatModelConfig | MultiScaleModelConfig) -> nn.Modul Returns: Initialized upsampler module. """ - if isinstance(config, MultiScaleModelConfig): + if config.architecture == "multi_scale": return _build_multi_scale(config) return _build_flat(config) @@ -2670,7 +2469,27 @@ def build_motion_attention_spec(config: MultiScaleModelConfig) -> MotionAttentio ) -_MODEL_CONFIG_ADAPTER = TypeAdapter(ModelConfig) +# Published Diffusers configs may omit the optional motion-attention section. +# Keep the native factory unchanged while accepting the same config shape in +# the serialized Diffusers adapter. +def build_motion_attention_spec(config: MultiScaleModelConfig) -> MotionAttentionSpec | None: + """Resolve optional motion-attention settings from a serialized config.""" + motion = getattr(config, "motion_attention", None) + if motion is None: + return None + return MotionAttentionSpec( + after_mid_blocks=motion.after_mid_blocks, + block=MotionCorrespondenceSpec( + channels=config.hidden_channels, + spatial_kernel_size=motion.spatial_kernel_size, + temporal_offsets=motion.temporal_offsets, + num_heads=motion.num_heads, + head_dim=config.hidden_channels // motion.num_heads, + backend=motion.backend, + natten_backend=motion.natten_backend, + merge_compile=motion.merge_compile, + ), + ) class Kandinsky6SRLatentUpscalerBank(ModelMixin, ConfigMixin): @@ -2697,7 +2516,7 @@ def __init__( target_scale = f"{target_scale}x" if target_scale in self._models: raise ValueError(f"duplicate latent upscaler target scale: {target_scale}") - model_config = _MODEL_CONFIG_ADAPTER.validate_python(spec["model"]) + model_config = _model_config(spec["model"]) upscaler = build_upsampler(model_config) upscaler.target_scale = target_scale upscaler.scaling_factor = float(scaling_factor) diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py index f57afba38041..0a93a5817cf7 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py @@ -25,12 +25,10 @@ from types import SimpleNamespace from typing import Any, Literal, NamedTuple -import av import numpy as np import torch from ..pipeline_utils import DiffusionPipeline -from diffusers.utils import replace_example_docstring -from pydantic import BaseModel, ConfigDict, Field +from diffusers.utils import logging, replace_example_docstring from ...schedulers.scheduling_piflow import DXPolicy, policy_rollout_fm, shift_timesteps from torch import Tensor, nn from torch.distributed import all_gather @@ -38,6 +36,9 @@ from .pipeline_output import Kandinsky6SRPipelineOutput +logger = logging.get_logger(__name__) + + EXAMPLE_DOC_STRING = """ Examples: @@ -1421,19 +1422,32 @@ def piflow_generate( return x[..., :out_dim] -class RunConfig(BaseModel): +@dataclass +class RunConfig: """Parameters for one tiled SR run.""" - model_config = ConfigDict(extra="forbid") - device: str - num_steps: int = Field(default=5, ge=2) + num_steps: int = 5 seed: int = 42 - overlap: float = Field(default=0.25, ge=0.0, lt=1.0) - tiles_batch_size: int = Field(default=1, gt=0) + overlap: float = 0.25 + tiles_batch_size: int = 1 resolution_scale: ResolutionScale = 4 tile_grid_mode: TileGridMode = "even" - tile_grid_min_overlap: float = Field(default=0.20, ge=0.0, lt=1.0) + tile_grid_min_overlap: float = 0.20 + + def __post_init__(self) -> None: + if self.num_steps < 2: + raise ValueError("num_steps must be at least 2") + if not 0.0 <= self.overlap < 1.0: + raise ValueError("overlap must satisfy 0 <= overlap < 1") + if self.tiles_batch_size <= 0: + raise ValueError("tiles_batch_size must be positive") + if self.resolution_scale not in (2, 4): + raise ValueError("resolution_scale must be 2 or 4") + if self.tile_grid_mode not in ("legacy", "even"): + raise ValueError("tile_grid_mode must be 'legacy' or 'even'") + if not 0.0 <= self.tile_grid_min_overlap < 1.0: + raise ValueError("tile_grid_min_overlap must satisfy 0 <= value < 1") @dataclass @@ -2060,6 +2074,17 @@ def _run_tile_batches( return outputs +def _require_av(): + """Load PyAV only for file-based video input.""" + try: + import av + except (ImportError, OSError) as exc: + raise RuntimeError( + "SR file-based video input requires PyAV and FFmpeg. Install with `pip install av`." + ) from exc + return av + + def align_to_vae_stride(t: int) -> int: """Round ``t`` to the nearest valid pixel-frame count ``1 + 8·k`` (ties up). @@ -2131,8 +2156,8 @@ def resample_to_target_fps( if src_fps > target_fps: indices = select_frame_indices(video.shape[0], src_fps, target_fps) logger.warning( - "Source fps {:.2f} > target {}fps: downsampling {} frames -> {} (fixed-stride); " - "source temporal detail beyond {}fps is discarded.", + "Source fps %.2f > target %sfps: downsampling %s frames -> %s (fixed-stride); " + "source temporal detail beyond %sfps is discarded.", src_fps, target_fps, video.shape[0], @@ -2141,7 +2166,7 @@ def resample_to_target_fps( ) return video[indices], target_fps logger.warning( - "Source fps {:.2f} < target {}fps: keeping native frames (no minterpolate upsample); " + "Source fps %.2f < target %sfps: keeping native frames (no minterpolate upsample); " "output is mildly out of distribution.", src_fps, target_fps, @@ -2186,6 +2211,7 @@ def read_video_tchw_uint8(path: Path) -> tuple[torch.Tensor, float]: Raises: ValueError: If the file has no readable frames or reports no fps. """ + av = _require_av() with av.open(str(path), mode="r") as container: stream = container.streams.video[0] src_fps = float(stream.average_rate or stream.base_rate or 0.0) @@ -2222,10 +2248,19 @@ def load_lr_latent(path: Path) -> torch.Tensor: shape = getattr(obj, "shape", type(obj).__name__) msg = f"Expected a rank-4 [T, C, H, W] latent tensor in {path}, got {shape}." raise ValueError(msg) - logger.info("Loaded LR latent {} from {}", tuple(obj.shape), path) + logger.info("Loaded LR latent %s from %s", tuple(obj.shape), path) return obj.float() +def _require_av(): + """Load PyAV only when video or audio output is requested.""" + try: + import av + except (ImportError, OSError) as exc: + raise RuntimeError("SR video/audio output requires PyAV and FFmpeg. Install with `pip install av`.") from exc + return av + + def _video_np(frames: Tensor) -> np.ndarray: """(3, T, H, W) uint8 → (T, H, W, 3) uint8.""" return frames.detach().permute(1, 2, 3, 0).cpu().numpy().astype(np.uint8, copy=False) @@ -2253,6 +2288,7 @@ def extract_audio_from_video( if audio_sample_rate <= 0: raise ValueError("audio_sample_rate must be positive") + av = _require_av() chunks: list[np.ndarray] = [] with av.open(str(source_video), mode="r") as container: audio_streams = list(container.streams.audio) @@ -2290,6 +2326,7 @@ def mux_video_audio( # noqa: PLR0913 extracted from a generated T2VA video. lossless: encode video as FFV1 instead of libx264. Audio remains AAC. """ + av = _require_av() if audio is not None and source_video is not None: raise ValueError("pass either audio or source_video, not both") if source_video is not None: diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index 5dcd11b5016c..c7995bc97a55 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -654,7 +654,12 @@ def resize_image( k = sqrt(max_area / area) / div new_h = int(round(h * k) * div) new_w = int(round(w * k) * div) - import torchvision.transforms.functional as TF + try: + import torchvision.transforms.functional as TF + except (ImportError, OSError) as exc: + raise RuntimeError( + "I2VA image processing requires torchvision. Install it with `pip install torchvision`." + ) from exc return TF.resize(image, (new_h, new_w)), k @@ -688,7 +693,12 @@ def encode_i2va_first_frame( world_size: int = 1, ) -> tuple[Tensor, int, int]: """Encode one image into the K6 packed first-frame latent layout.""" - import torchvision.transforms.functional as TF + try: + import torchvision.transforms.functional as TF + except (ImportError, OSError) as exc: + raise RuntimeError( + "I2VA image processing requires torchvision. Install it with `pip install torchvision`." + ) from exc pil_image = _load_pil_rgb(image) tensor = TF.pil_to_tensor(pil_image).unsqueeze(0) From 3dc5959047ac3f6931e545e295eaec7de518f9a0 Mon Sep 17 00:00:00 2001 From: Denis Koposov Date: Wed, 16 Sep 2026 21:54:17 +0000 Subject: [PATCH 07/62] remove dead manual weight loading code; flattent MMAudioVAE structure --- .../autoencoders/autoencoder_kandinsky6_sr.py | 39 --------- .../autoencoders/autoencoder_mmaudio.py | 83 +++---------------- 2 files changed, 11 insertions(+), 111 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index 2ff0a11ee011..e94c8743617c 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -1080,45 +1080,6 @@ def __init__( self.spatial_factor = int(spatial_factor) self.temporal_factor = int(temporal_factor) - def init_from_ckpt(self, path): - """Load a KVAE checkpoint in safetensors or training-checkpoint format.""" - if str(path).endswith(".safetensors"): - # Release checkpoints are flat safetensors state dicts already in - # the new-style key naming. - self.load_state_dict(safetensors_load_file(str(path)), strict=True) - return - sd = torch.load(path, map_location="cpu")["state_dict"] - - # Fix checkpoint if starting from new style - replace_keys = dict() - delete_keys = list() - for k in sd: - if k.startswith("loss."): - delete_keys.append(k) - continue - if "encoder.down" in k and "downsample.temporal_conv.conv" in k: - continue - - if k.startswith("decoder"): - if "upsample" in k: - continue - elif ".conv_b.conv" in k: - replace_keys[k] = k.replace(".conv_b.conv", ".conv_b") - continue - elif ".conv_y.conv" in k: - replace_keys[k] = k.replace(".conv_y.conv", ".conv_y") - continue - - if k.endswith("sample.temporal_conv.conv.weight") or k.endswith("sample.temporal_conv.conv.bias"): - replace_keys[k] = k.replace(".temporal_conv.conv.", ".temporal_conv.") - for old_k, new_k in replace_keys.items(): - sd[new_k] = sd[old_k] - del sd[old_k] - for k in delete_keys: - del sd[k] - - self.load_state_dict(sd, strict=True) - def make_empty_cache(self, block: str): """Create empty causal-convolution and normalization caches.""" diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index d16685c53b11..605fbfdf8a22 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -3,7 +3,6 @@ from __future__ import annotations import math -from os import PathLike from typing import Any, Literal, Optional import numpy as np @@ -1254,47 +1253,6 @@ def build_bigvgan_v2(vocoder_config: dict | AttrDict) -> BigVGANv2: return model -class AutoEncoderModule(nn.Module): - def __init__( - self, - *, - vae_ckpt_path: str | None = None, - vocoder_ckpt_path: str | None = None, - vocoder_config: dict | None = None, - mode: Literal["44k"], - need_vae_encoder: bool = True, - need_vae_decoder: bool = True, - ): - super().__init__() - if mode != "44k": - raise ValueError(f"Unknown model: {mode}") - self.vae: VAE = VAE(data_dim=128, embed_dim=40, hidden_dim=512).eval() - if vae_ckpt_path is not None: - vae_state_dict = torch.load(vae_ckpt_path, weights_only=True, map_location="cpu") - self.vae.load_state_dict(vae_state_dict) - self.vae.remove_weight_norm() - - if need_vae_decoder: - self.vocoder = build_bigvgan_v2(vocoder_config).eval() - else: - del self.vae.decoder - - if not need_vae_encoder: - del self.vae.encoder - - for param in self.parameters(): - param.requires_grad = False - - def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution: - return self.vae.encode(x) - - def decode(self, z: torch.Tensor) -> torch.Tensor: - return self.vae.decode(z) - - def vocode(self, spec: torch.Tensor) -> torch.Tensor: - return self.vocoder(spec) - - # Reference: # https://github.com/bytedance/Make-An-Audio-2 @@ -1410,51 +1368,35 @@ class MMAudioVAE(ModelMixin, ConfigMixin): BigVGAN vocoder. Args: - tod_vae_ckpt (`str` or `os.PathLike`, *optional*): Path to the audio VAE checkpoint. - bigvgan_vocoder_ckpt (`str` or `os.PathLike`, *optional*): Path to the BigVGAN checkpoint. - vocoder_config (`dict`, *optional*): BigVGAN configuration when no checkpoint directory is supplied. + vocoder_config (`dict`): BigVGAN configuration. mode (`str`, *optional*, defaults to ``"44k"``): Audio operating mode. - need_vae_encoder (`bool`, *optional*, defaults to True): Whether to load the VAE encoder. - need_vae_decoder (`bool`, *optional*, defaults to True): Whether to load the VAE decoder and vocoder. scaling_factor (`float`, *optional*, defaults to 1.0): Latent scaling factor. sample_rate (`int`, *optional*, defaults to 44100): Audio sample rate. downsample_factor (`int`, *optional*, defaults to 1024): Audio latent downsampling factor. - device (`str` or `torch.device`, *optional*): Device on which to place the component. """ - ignore_for_config = ["tod_vae_ckpt", "bigvgan_vocoder_ckpt", "device"] - _no_split_modules = ["AutoEncoderModule", "BigVGANv2"] + _no_split_modules = ["VAE", "BigVGANv2"] @register_to_config def __init__( self, *, - tod_vae_ckpt: str | PathLike[str] | None = None, - bigvgan_vocoder_ckpt: str | PathLike[str] | None = None, - vocoder_config: dict[str, Any] | None = None, + vocoder_config: dict[str, Any], mode: str = "44k", - need_vae_encoder: bool = True, - need_vae_decoder: bool = True, scaling_factor: float = 1.0, sample_rate: int = 44_100, downsample_factor: int = 1_024, - device: str | torch.device | None = None, ) -> None: super().__init__() self.mel_converter = get_mel_converter(mode) - self.tod = AutoEncoderModule( - vae_ckpt_path=tod_vae_ckpt, - vocoder_ckpt_path=bigvgan_vocoder_ckpt, - vocoder_config=vocoder_config, - mode=mode, - need_vae_encoder=need_vae_encoder, - need_vae_decoder=need_vae_decoder, - ) + self.vae: VAE = VAE(data_dim=128, embed_dim=40, hidden_dim=512).eval() + self.vae.remove_weight_norm() + self.vocoder = build_bigvgan_v2(vocoder_config).eval() + for param in self.parameters(): + param.requires_grad = False self.scaling_factor = float(scaling_factor) self.sample_rate = int(sample_rate) self.downsample_factor = int(downsample_factor) - if device is not None: - self.to(device) def compile(self): """Compile the decode and vocode methods with ``torch.compile``.""" @@ -1467,19 +1409,16 @@ def train(self, mode: bool = True) -> MMAudioVAE: def encode_audio(self, x) -> DiagonalGaussianDistribution: """Encode a waveform into an audio latent distribution.""" - assert self.tod is not None, "VAE is not loaded" mel = self.mel_converter(x) - return self.tod.encode(mel) + return self.vae.encode(mel) def vocode(self, mel: torch.Tensor) -> torch.Tensor: """Convert a decoded mel-spectrogram into a waveform.""" - assert self.tod is not None, "VAE is not loaded" - return self.tod.vocode(mel) + return self.vocoder(mel) def decode(self, z: torch.Tensor) -> torch.Tensor: """Decode audio latents into a mel-spectrogram.""" - assert self.tod is not None, "VAE is not loaded" - return self.tod.decode(z) + return self.vae.decode(z) @property def device(self): From d5a34b615e325f889c168c8f1eea6a636008720c Mon Sep 17 00:00:00 2001 From: Denis Koposov Date: Thu, 17 Sep 2026 08:13:40 +0000 Subject: [PATCH 08/62] move attention processors to transformer files --- src/diffusers/models/attention_processor.py | 290 ------------------ .../transformers/transformer_kandinsky6.py | 144 ++++++++- .../transformers/transformer_kandinsky6_sr.py | 102 +++++- 3 files changed, 242 insertions(+), 294 deletions(-) diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py index 2286b03ad661..1b923e749663 100755 --- a/src/diffusers/models/attention_processor.py +++ b/src/diffusers/models/attention_processor.py @@ -5677,293 +5677,3 @@ def __new__(cls, *args, **kwargs): | LoRAXFormersAttnProcessor | LoRAAttnAddedKVProcessor ) - -import math -from typing import TYPE_CHECKING, Any - -import torch -import torch.nn.functional as functional -from .attention_dispatch import ( - _CAN_USE_FLEX_ATTN, - AttentionBackendName, - dispatch_attention_fn, -) -from torch import Tensor -from torch.nn.attention.flex_attention import BlockMask - - -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 nabla_block_mask( - q: Tensor, - k: Tensor, - sta: Tensor, - thr: float = 0.9, - block_size: int = 64, -) -> BlockMask: - """Build a dynamic NABLA block mask from query/key statistics and an STA prior.""" - B, h, S, D = q.shape - s1 = S // block_size - qa = q.reshape(B, h, s1, block_size, D).mean(-2) - ka = k.reshape(B, h, s1, block_size, D).mean(-2).transpose(-2, -1) - attn_map = torch.softmax((qa @ ka) / math.sqrt(D), dim=-1) - - vals, inds = attn_map.sort(-1) - mask = (vals.cumsum_(-1) >= 1 - thr).int().gather(-1, inds.argsort(-1)) - mask = torch.logical_or(mask, sta) - - kv_nb = mask.sum(-1).to(torch.int32) - kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) - return BlockMask.from_kv_blocks( - torch.zeros_like(kv_nb), - kv_inds, - kv_nb, - kv_inds, - BLOCK_SIZE=block_size, - mask_mod=None, - ) - - -# Pure-torch helpers for the streaming KV-cache in self-attention. -# -# Kept free of ``flash_attn`` imports so the cache bookkeeping is importable and -# testable on CPU. The actual attention kernel call lives in -# ``kandinsky_sr.model.nn``; this module only assembles its packed inputs. - - -if TYPE_CHECKING: - - pass - - -def assemble_cached_attention_inputs( - new_key: Tensor, - new_value: Tensor, - cached_key: Tensor, - cached_value: Tensor, - cu_seqlens_new: Tensor, - cached_cu_seqlens: Tensor, -) -> tuple[Tensor, Tensor, Tensor, Tensor]: - """Merge cached and new K/V into packed varlen attention inputs. - - For each packed sequence ``i`` the keys/values become ``[cache_i | new_i]``; - the query packing stays the new-token packing (queries are the freshly - arrived tokens that attend back over the cache plus themselves). - - Args: - new_key: New keys, packed ``(sum(new_lens), num_heads, head_dim)``. - new_value: New values, same packing as ``new_key``. - cached_key: Cached keys, packed ``(sum(cache_lens), num_heads, head_dim)``. - cached_value: Cached values, same packing as ``cached_key``. - cu_seqlens_new: Cumulative sequence lengths of the new tokens, ``(num_seqs + 1,)``. - cached_cu_seqlens: Cumulative sequence lengths of the cache, ``(num_seqs + 1,)``. - - Returns: - Tuple ``(key, value, cu_seqlens_q, cu_seqlens_k)`` where ``key``/``value`` - are the merged packed tensors, ``cu_seqlens_q`` is the query packing - (``= cu_seqlens_new``) and ``cu_seqlens_k`` the merged key packing. - """ - num_seqs = cu_seqlens_new.numel() - 1 - key_parts: list[Tensor] = [] - value_parts: list[Tensor] = [] - for i in range(num_seqs): - cache_start, cache_end = int(cached_cu_seqlens[i]), int(cached_cu_seqlens[i + 1]) - new_start, new_end = int(cu_seqlens_new[i]), int(cu_seqlens_new[i + 1]) - key_parts.append(cached_key[cache_start:cache_end]) - key_parts.append(new_key[new_start:new_end]) - value_parts.append(cached_value[cache_start:cache_end]) - value_parts.append(new_value[new_start:new_end]) - - key = torch.cat(key_parts, dim=0) - value = torch.cat(value_parts, dim=0) - - merged_lens = torch.diff(cached_cu_seqlens) + torch.diff(cu_seqlens_new) - cu_seqlens_k = torch.cat([cu_seqlens_new.new_zeros(1), torch.cumsum(merged_lens, dim=0)]).to(cu_seqlens_new.dtype) - cu_seqlens_q = cu_seqlens_new - return key, value, cu_seqlens_q, cu_seqlens_k - - -# A captured per-step cache "slot": one ``(key, value, cu_seqlens)`` triple per -# self-attention layer (aligned with ``iter_self_attentions`` order), or ``None`` -# for an empty slot (no past frames cached yet at that denoising step). -LayerCache = tuple["Tensor | None", "Tensor | None", "Tensor | None"] -KVSlot = "list[LayerCache] | None" - - -_MASKED_ATTENTION_BACKENDS = { - "flash": "flash_varlen", - "_flash_3": "_flash_varlen_3", - "sage": "sage_varlen", - "native": "native", -} - - -class Kandinsky6AttnProcessor: - """Diffusers attention processor shared by TI2VA and SR attention.""" - - _attention_backend = None - _parallel_config = None - - def __init__(self, attention_backend=None, parallel_config=None): - if not hasattr(functional, "scaled_dot_product_attention"): - raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0 or newer.") - self._masked = False - self._attention_backend = attention_backend - self._parallel_config = parallel_config - - @property - def _attention_backend(self): - return self.__attention_backend - - @_attention_backend.setter - def _attention_backend(self, backend): - if self._masked and backend is not None: - name = getattr(backend, "value", backend).lower() - backend = _MASKED_ATTENTION_BACKENDS.get(name, name) - backend = AttentionBackendName(backend) - self.__attention_backend = backend - - def __call__( - self, - attn: Any, - hidden_states: Tensor, - encoder_hidden_states: Tensor | None = None, - rotary_emb: Tensor | None = None, - rotary_emb_kv: Tensor | None = None, - sparse_params: dict[str, Any] | None = None, - attn_mask: Tensor | None = None, - ) -> Tensor: - query = attn.to_query(hidden_states) - if encoder_hidden_states is None: - key = attn.to_key(hidden_states) - value = attn.to_value(hidden_states) - else: - key = attn.to_key(encoder_hidden_states) - value = attn.to_value(encoder_hidden_states) - - query = query.reshape(*query.shape[:-1], attn.num_heads, -1) - key = key.reshape(*key.shape[:-1], attn.num_heads, -1) - value = value.reshape(*value.shape[:-1], attn.num_heads, -1) - query = attn.query_norm(query) - key = attn.key_norm(key) - - if rotary_emb is not None: - query = apply_rotary(query, rotary_emb).to(dtype=query.dtype) - if rotary_emb_kv is not None: - key = apply_rotary(key, rotary_emb_kv).to(dtype=key.dtype) - - query_was_batched = query.dim() == 4 - if not query_was_batched: - query, key, value = query.unsqueeze(0), key.unsqueeze(0), value.unsqueeze(0) - elif key.dim() == 3: - key, value = key.unsqueeze(0), value.unsqueeze(0) - - if sparse_params is not None: - q = query.transpose(1, 2).contiguous() - k = key.transpose(1, 2).contiguous() - v = value.transpose(1, 2).contiguous() - block_mask = nabla_block_mask( - q, - k, - sparse_params["sta_mask"], - thr=sparse_params["P"], - ) - if not _CAN_USE_FLEX_ATTN: - raise ValueError("Nabla attention requires PyTorch 2.5 or newer") - output = ( - torch.nn.attention.flex_attention.flex_attention( - q, - k, - v, - block_mask=block_mask, - ) - .transpose(1, 2) - .contiguous() - ) - else: - output = dispatch_attention_fn( - query, - key, - value, - attn_mask=attn_mask, - backend=self._attention_backend, - parallel_config=self._parallel_config, - ) - - if not query_was_batched: - output = output[0] - return attn.out_layer(output.flatten(-2, -1)) - - -class Kandinsky6SRAttentionProcessor(Kandinsky6AttnProcessor): - """Add SR sparse-attention and KV-cache handling to the shared processor.""" - - def __call__( - self, - attn: Any, - query: Tensor, - key: Tensor, - value: Tensor, - cu_seqlens_q: Tensor, - cu_seqlens_k: Tensor, - *, - sparse_params: dict[str, Any] | None = None, - attn_mask: Tensor | None = None, - ) -> Tensor: - if sparse_params is not None: - return attn.attention_flex(query, key, value, sparse_params=sparse_params) - - if getattr(attn, "cached_k", None) is not None: - key, value, cu_seqlens_q, cu_seqlens_k = assemble_cached_attention_inputs( - key, value, attn.cached_k, attn.cached_v, cu_seqlens_q, attn.cached_cu_seqlens - ) - if getattr(attn, "return_kv", False): - attn.cached_k, attn.cached_v, attn.cached_cu_seqlens = key, value, cu_seqlens_k - - outputs = [] - batch_size = cu_seqlens_q.numel() - 1 - for index in range(batch_size): - query_start, query_end = int(cu_seqlens_q[index]), int(cu_seqlens_q[index + 1]) - key_start, key_end = int(cu_seqlens_k[index]), int(cu_seqlens_k[index + 1]) - outputs.append( - dispatch_attention_fn( - query[query_start:query_end].unsqueeze(0), - key[key_start:key_end].unsqueeze(0), - value[key_start:key_end].unsqueeze(0), - attn_mask=_packed_attention_mask( - attn_mask, - index, - batch_size, - query_end - query_start, - key_end - key_start, - ), - backend=self._attention_backend, - parallel_config=self._parallel_config, - )[0] - ) - return torch.cat(outputs, dim=0) - - -def _packed_attention_mask( - attn_mask: Tensor | None, - index: int, - batch_size: int, - query_length: int, - key_length: int, -) -> Tensor | None: - if attn_mask is None: - return None - if attn_mask.ndim == 1: - return attn_mask[:key_length].unsqueeze(0) - if attn_mask.ndim == 2: - if attn_mask.shape[0] in (1, batch_size): - return attn_mask[index : index + 1, :key_length] - return attn_mask[:query_length, :key_length] - if attn_mask.shape[0] in (1, batch_size): - attn_mask = attn_mask[index : index + 1] - return attn_mask[..., :query_length, :key_length] - diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 3574a95e58a4..c0788ebcb7e0 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -27,13 +27,16 @@ from ...hooks.mag_cache import MagCacheState from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ..attention import AttentionMixin, AttentionModuleMixin -from ..attention_dispatch import AttentionBackendName +from ..attention_dispatch import ( + _CAN_USE_FLEX_ATTN, + AttentionBackendName, + dispatch_attention_fn, +) from ..cache_utils import CacheMixin from ..modeling_outputs import Transformer2DModelOutput from ..modeling_utils import ModelMixin from torch import Tensor, nn - -from ..attention_processor import Kandinsky6AttnProcessor +from torch.nn.attention.flex_attention import BlockMask def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: @@ -108,6 +111,36 @@ def fractal_unflatten(x: Tensor, shape: tuple, block_mask: bool = False) -> Tens return x.reshape(*shape, x.shape[-1]) +def nabla_block_mask( + q: Tensor, + k: Tensor, + sta: Tensor, + thr: float = 0.9, + block_size: int = 64, +) -> BlockMask: + """Build a dynamic NABLA block mask from query/key statistics and an STA prior.""" + B, h, S, D = q.shape + s1 = S // block_size + qa = q.reshape(B, h, s1, block_size, D).mean(-2) + ka = k.reshape(B, h, s1, block_size, D).mean(-2).transpose(-2, -1) + attn_map = torch.softmax((qa @ ka) / math.sqrt(D), dim=-1) + + vals, inds = attn_map.sort(-1) + mask = (vals.cumsum_(-1) >= 1 - thr).int().gather(-1, inds.argsort(-1)) + mask = torch.logical_or(mask, sta) + + kv_nb = mask.sum(-1).to(torch.int32) + kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) + return BlockMask.from_kv_blocks( + torch.zeros_like(kv_nb), + kv_inds, + kv_nb, + kv_inds, + BLOCK_SIZE=block_size, + mask_mod=None, + ) + + class RoPE1D(nn.Module): """1-D Rotary Position Embedding — used for text and audio sequences.""" @@ -412,6 +445,111 @@ def disable_cache(self) -> None: # native K6 names to this public Diffusers layout. +_MASKED_ATTENTION_BACKENDS = { + "flash": "flash_varlen", + "_flash_3": "_flash_varlen_3", + "sage": "sage_varlen", + "native": "native", +} + + +class Kandinsky6AttnProcessor: + """Diffusers attention processor used by the TI2VA transformer.""" + + _attention_backend = None + _parallel_config = None + + def __init__(self, attention_backend=None, parallel_config=None): + if not hasattr(functional, "scaled_dot_product_attention"): + raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0 or newer.") + self._masked = False + self._attention_backend = attention_backend + self._parallel_config = parallel_config + + @property + def _attention_backend(self): + return self.__attention_backend + + @_attention_backend.setter + def _attention_backend(self, backend): + if self._masked and backend is not None: + name = getattr(backend, "value", backend).lower() + backend = _MASKED_ATTENTION_BACKENDS.get(name, name) + backend = AttentionBackendName(backend) + self.__attention_backend = backend + + def __call__( + self, + attn: Any, + hidden_states: Tensor, + encoder_hidden_states: Tensor | None = None, + rotary_emb: Tensor | None = None, + rotary_emb_kv: Tensor | None = None, + sparse_params: dict[str, Any] | None = None, + attn_mask: Tensor | None = None, + ) -> Tensor: + query = attn.to_query(hidden_states) + if encoder_hidden_states is None: + key = attn.to_key(hidden_states) + value = attn.to_value(hidden_states) + else: + key = attn.to_key(encoder_hidden_states) + value = attn.to_value(encoder_hidden_states) + + query = query.reshape(*query.shape[:-1], attn.num_heads, -1) + key = key.reshape(*key.shape[:-1], attn.num_heads, -1) + value = value.reshape(*value.shape[:-1], attn.num_heads, -1) + query = attn.query_norm(query) + key = attn.key_norm(key) + + if rotary_emb is not None: + query = apply_rotary(query, rotary_emb).to(dtype=query.dtype) + if rotary_emb_kv is not None: + key = apply_rotary(key, rotary_emb_kv).to(dtype=key.dtype) + + query_was_batched = query.dim() == 4 + if not query_was_batched: + query, key, value = query.unsqueeze(0), key.unsqueeze(0), value.unsqueeze(0) + elif key.dim() == 3: + key, value = key.unsqueeze(0), value.unsqueeze(0) + + if sparse_params is not None: + q = query.transpose(1, 2).contiguous() + k = key.transpose(1, 2).contiguous() + v = value.transpose(1, 2).contiguous() + block_mask = nabla_block_mask( + q, + k, + sparse_params["sta_mask"], + thr=sparse_params["P"], + ) + if not _CAN_USE_FLEX_ATTN: + raise ValueError("Nabla attention requires PyTorch 2.5 or newer") + output = ( + torch.nn.attention.flex_attention.flex_attention( + q, + k, + v, + block_mask=block_mask, + ) + .transpose(1, 2) + .contiguous() + ) + else: + output = dispatch_attention_fn( + query, + key, + value, + attn_mask=attn_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + + if not query_was_batched: + output = output[0] + return attn.out_layer(output.flatten(-2, -1)) + + class Kandinsky6TimeEmbeddings(nn.Module): """Sinusoidal timestep embedding with a K6-compatible parameter layout.""" diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index 1f17c35a6587..90d10e5f42c4 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -24,12 +24,112 @@ import torch from ...configuration_utils import ConfigMixin, register_to_config from ..attention import AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn from ..modeling_utils import ModelMixin from diffusers.utils import logging from torch import Tensor, nn from torch.nn.attention.flex_attention import BlockMask, flex_attention -from ..attention_processor import Kandinsky6SRAttentionProcessor +from .transformer_kandinsky6 import Kandinsky6AttnProcessor + + +class Kandinsky6SRAttentionProcessor(Kandinsky6AttnProcessor): + """Diffusers attention processor used by the SR transformer.""" + + def __call__( + self, + attn: Any, + query: Tensor, + key: Tensor, + value: Tensor, + cu_seqlens_q: Tensor, + cu_seqlens_k: Tensor, + *, + sparse_params: dict[str, Any] | None = None, + attn_mask: Tensor | None = None, + ) -> Tensor: + if sparse_params is not None: + return attn.attention_flex(query, key, value, sparse_params=sparse_params) + + if getattr(attn, "cached_k", None) is not None: + key, value, cu_seqlens_q, cu_seqlens_k = self.assemble_cached_attention_inputs( + key, value, attn.cached_k, attn.cached_v, cu_seqlens_q, attn.cached_cu_seqlens + ) + if getattr(attn, "return_kv", False): + attn.cached_k, attn.cached_v, attn.cached_cu_seqlens = key, value, cu_seqlens_k + + outputs = [] + batch_size = cu_seqlens_q.numel() - 1 + for index in range(batch_size): + query_start, query_end = int(cu_seqlens_q[index]), int(cu_seqlens_q[index + 1]) + key_start, key_end = int(cu_seqlens_k[index]), int(cu_seqlens_k[index + 1]) + outputs.append( + dispatch_attention_fn( + query[query_start:query_end].unsqueeze(0), + key[key_start:key_end].unsqueeze(0), + value[key_start:key_end].unsqueeze(0), + attn_mask=self._packed_attention_mask( + attn_mask, + index, + batch_size, + query_end - query_start, + key_end - key_start, + ), + backend=self._attention_backend, + parallel_config=self._parallel_config, + )[0] + ) + return torch.cat(outputs, dim=0) + + @staticmethod + def assemble_cached_attention_inputs( + new_key: Tensor, + new_value: Tensor, + cached_key: Tensor, + cached_value: Tensor, + cu_seqlens_new: Tensor, + cached_cu_seqlens: Tensor, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """Merge cached and new K/V into packed varlen attention inputs.""" + num_seqs = cu_seqlens_new.numel() - 1 + key_parts: list[Tensor] = [] + value_parts: list[Tensor] = [] + for index in range(num_seqs): + cache_start, cache_end = int(cached_cu_seqlens[index]), int(cached_cu_seqlens[index + 1]) + new_start, new_end = int(cu_seqlens_new[index]), int(cu_seqlens_new[index + 1]) + key_parts.append(cached_key[cache_start:cache_end]) + key_parts.append(new_key[new_start:new_end]) + value_parts.append(cached_value[cache_start:cache_end]) + value_parts.append(new_value[new_start:new_end]) + + key = torch.cat(key_parts, dim=0) + value = torch.cat(value_parts, dim=0) + merged_lens = torch.diff(cached_cu_seqlens) + torch.diff(cu_seqlens_new) + cu_seqlens_k = torch.cat([cu_seqlens_new.new_zeros(1), torch.cumsum(merged_lens, dim=0)]).to( + cu_seqlens_new.dtype + ) + return key, value, cu_seqlens_new, cu_seqlens_k + + @staticmethod + def _packed_attention_mask( + attn_mask: Tensor | None, + index: int, + batch_size: int, + query_length: int, + key_length: int, + ) -> Tensor | None: + if attn_mask is None: + return None + if attn_mask.ndim == 1: + return attn_mask[:key_length].unsqueeze(0) + if attn_mask.ndim == 2: + if attn_mask.shape[0] in (1, batch_size): + return attn_mask[index : index + 1, :key_length] + return attn_mask[:query_length, :key_length] + if attn_mask.shape[0] in (1, batch_size): + attn_mask = attn_mask[index : index + 1] + return attn_mask[..., :query_length, :key_length] + # Side of the local 8x8 token block used by fractal (NABLA) attention. # Independent of the VAE spatial compression — do not swap for VAE_SPATIAL_FACTOR. From 334d639ba229f8b0ab8353f4e25b38958a3a1590 Mon Sep 17 00:00:00 2001 From: Denis Koposov Date: Thu, 17 Sep 2026 08:42:10 +0000 Subject: [PATCH 09/62] remove duplicate code in K6 latent_upscaler --- .../pipelines/kandinsky6/latent_upscaler.py | 23 ------------------- 1 file changed, 23 deletions(-) diff --git a/src/diffusers/pipelines/kandinsky6/latent_upscaler.py b/src/diffusers/pipelines/kandinsky6/latent_upscaler.py index 74e54d57965f..f177814ef896 100644 --- a/src/diffusers/pipelines/kandinsky6/latent_upscaler.py +++ b/src/diffusers/pipelines/kandinsky6/latent_upscaler.py @@ -2469,29 +2469,6 @@ def build_motion_attention_spec(config: MultiScaleModelConfig) -> MotionAttentio ) -# Published Diffusers configs may omit the optional motion-attention section. -# Keep the native factory unchanged while accepting the same config shape in -# the serialized Diffusers adapter. -def build_motion_attention_spec(config: MultiScaleModelConfig) -> MotionAttentionSpec | None: - """Resolve optional motion-attention settings from a serialized config.""" - motion = getattr(config, "motion_attention", None) - if motion is None: - return None - return MotionAttentionSpec( - after_mid_blocks=motion.after_mid_blocks, - block=MotionCorrespondenceSpec( - channels=config.hidden_channels, - spatial_kernel_size=motion.spatial_kernel_size, - temporal_offsets=motion.temporal_offsets, - num_heads=motion.num_heads, - head_dim=config.hidden_channels // motion.num_heads, - backend=motion.backend, - natten_backend=motion.natten_backend, - merge_compile=motion.merge_compile, - ), - ) - - class Kandinsky6SRLatentUpscalerBank(ModelMixin, ConfigMixin): """Diffusers wrapper around the self-contained x2/x4 latent-upscaler bank. From 9d7dc717dae8cb00aba28194a6b8e9c0a79d98f9 Mon Sep 17 00:00:00 2001 From: leffff Date: Fri, 25 Sep 2026 15:02:47 +0300 Subject: [PATCH 10/62] address 1st batch of comments --- docs/source/en/api/pipelines/kandinsky6.md | 4 + src/diffusers/__init__.py | 9 +- .../autoencoders/autoencoder_kandinsky6_sr.py | 327 +++++++----------- .../autoencoders/autoencoder_mmaudio.py | 225 +++--------- .../transformers/transformer_kandinsky6.py | 74 ++-- .../transformers/transformer_kandinsky6_sr.py | 36 +- .../pipelines/kandinsky6/__init__.py | 28 +- .../pipelines/kandinsky6/latent_upscaler.py | 43 +-- .../kandinsky6/pipeline_kandinsky6_sr.py | 96 +---- .../kandinsky6/pipeline_kandinsky6_ti2va.py | 38 +- .../pipelines/kandinsky6/pipeline_output.py | 1 + src/diffusers/schedulers/scheduling_piflow.py | 27 +- src/diffusers/utils/dummy_pt_objects.py | 75 ++++ 13 files changed, 377 insertions(+), 606 deletions(-) diff --git a/docs/source/en/api/pipelines/kandinsky6.md b/docs/source/en/api/pipelines/kandinsky6.md index 042526006943..be562bfd6183 100644 --- a/docs/source/en/api/pipelines/kandinsky6.md +++ b/docs/source/en/api/pipelines/kandinsky6.md @@ -15,6 +15,10 @@ specific language governing permissions and limitations under the License. Kandinsky 6 provides pipelines for text/image-to-video-and-audio generation and video super-resolution. +[`Kandinsky6TI2VAPipeline`] denoises video and audio latents together with a single multimodal transformer, conditioned on a Qwen2.5-VL text encoder (and, optionally, a reference image) plus a CLIP text encoder for pooled embeddings. Video and audio come out of the same denoising loop, so there is no separate vocoder or post-hoc audio pass; pass `sample_audio=False` to skip generating audio and drop the `audio_vae` component. + +[`Kandinsky6SRPipeline`] upscales the video latents produced by the base pipeline. Its production route (`resolution_scale=2.25`) does a 1.125x pixel pre-upscale followed by an x2 latent-upscale step, using [`Kandinsky6SRLatentUpscalerBank`] as the latent upscaler; a `source_vae` is only needed for the KVAE latent bridge and is otherwise unused. + ## Kandinsky6TI2VAPipeline [[autodoc]] Kandinsky6TI2VAPipeline diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index a5e9286e7d42..6a7db85ac9af 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -698,12 +698,12 @@ "Kandinsky5I2IPipeline", "Kandinsky5I2VPipeline", "Kandinsky5T2IPipeline", + "Kandinsky5T2VPipeline", + "Kandinsky6SRLatentUpscalerBank", "Kandinsky6SRPipeline", "Kandinsky6SRPipelineOutput", - "Kandinsky6SRLatentUpscalerBank", "Kandinsky6TI2VAPipeline", "Kandinsky6TI2VAPipelineOutput", - "Kandinsky5T2VPipeline", "KandinskyCombinedPipeline", "KandinskyImg2ImgCombinedPipeline", "KandinskyImg2ImgPipeline", @@ -1555,6 +1555,11 @@ Kandinsky5I2VPipeline, Kandinsky5T2IPipeline, Kandinsky5T2VPipeline, + Kandinsky6SRLatentUpscalerBank, + Kandinsky6SRPipeline, + Kandinsky6SRPipelineOutput, + Kandinsky6TI2VAPipeline, + Kandinsky6TI2VAPipelineOutput, KandinskyCombinedPipeline, KandinskyImg2ImgCombinedPipeline, KandinskyImg2ImgPipeline, diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index e94c8743617c..6e37844a6638 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -18,9 +18,8 @@ import functools import math -from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, Tuple, Union +from typing import Tuple, Union import numpy as np import torch @@ -29,14 +28,16 @@ from ...configuration_utils import ConfigMixin, register_to_config from ..modeling_utils import ModelMixin from ...utils.accelerate_utils import apply_forward_hook -from torch.utils.checkpoint import checkpoint - - -def cast_tuple(t, length=1): - return t if isinstance(t, tuple) else ((t,) * length) +from .vae import DiagonalGaussianDistribution class SafeConv3d(nn.Conv3d): + """Conv3d that splits its input along the time axis when it would otherwise materialize more than ~2B + elements, running the convolution chunk by chunk instead. Each chunk is padded with the last + ``kernel_size - 1`` frames of the previous chunk before convolving, so the result is identical to running + the convolution on the whole tensor at once; this only bounds peak memory. + """ + def forward(self, x, write_to=None, transform=None): if transform is None: transform = lambda x: x @@ -48,9 +49,8 @@ def forward(self, x, write_to=None, transform=None): input_chunks = torch.chunk(x, part_num, dim=2) # NCTHW if any(ch.size(2) < kernel_size for ch in input_chunks) and kernel_size > 1: - assert input_chunks[0].numel() * (kernel_size / input_chunks[0].size(2)) < (2 * 10**9), ( - "frames are too big for Conv3d" - ) + if input_chunks[0].numel() * (kernel_size / input_chunks[0].size(2)) >= (2 * 10**9): + raise ValueError("frames are too big for Conv3d") t_stride, output = self.stride[0], [] for i in range(0, x.size(2) - kernel_size + 1, t_stride): @@ -88,35 +88,6 @@ def forward(self, x, write_to=None, transform=None): return write_to -def nonlinearity(x): - # Apply the SiLU activation. - return x * torch.sigmoid(x) - - -def custom_forward(module): - def inside_fn(args): - if isinstance(args, tuple): - return module(*args) - else: - return module(args) - - return inside_fn - - -def ckpt_wrap(module, args): - fn = custom_forward(module) - return checkpoint(fn, args=args, use_reentrant=False) - - -def identity(module, args): - if not isinstance(args, tuple): - return module(args) - return module(*args) - - -Conv3d = SafeConv3d - - class CausalConv3d(nn.Module): def __init__( self, @@ -129,11 +100,16 @@ def __init__( **kwargs, ): super().__init__() - kernel_size = cast_tuple(kernel_size, 3) + if not isinstance(kernel_size, tuple): + kernel_size = (kernel_size,) * 3 time_kernel_size, height_kernel_size, width_kernel_size = kernel_size - assert (height_kernel_size % 2) and (width_kernel_size % 2) + if not (height_kernel_size % 2 and width_kernel_size % 2): + raise ValueError( + f"height_kernel_size and width_kernel_size must be odd, got {height_kernel_size} and " + f"{width_kernel_size}" + ) self.height_pad = height_kernel_size // 2 self.width_pad = width_kernel_size // 2 @@ -142,7 +118,7 @@ def __init__( self.temporal_dim = 2 self.stride = stride - self.conv = Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) + self.conv = SafeConv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) self.cache_padding = None self.padding_mode = padding_mode @@ -156,56 +132,46 @@ def forward(self, input_): return output -def RMSNorm(in_channels, *args, **kwargs): - return WanRMS_norm(n_ch=in_channels, bias=False) - - -class WanRMS_norm(nn.Module): +class Kandinsky6VAERMSNorm(nn.Module): r""" - A custom RMS normalization layer. + RMS normalization over the channel dimension of a video tensor (`N, C, T, H, W`). + + `forward` accepts and ignores extra positional/keyword arguments (e.g. the causal-cache dict the + encoder/decoder pass to every normalization layer) so this stays a drop-in alternative to + `CachedGroupNorm`; construction likewise ignores `zq_ch`/`add_conv`, which only apply to the + `CachedSpatialNorm3D` alternative. Args: - dim (int): The number of dimensions to normalize over. + in_channels (int): The number of channels to normalize over. bias (bool, optional): Whether to include a learnable bias term. Default is False. """ - def __init__(self, n_ch: int, bias: bool = False) -> None: + def __init__(self, in_channels: int, bias: bool = False, **kwargs) -> None: super().__init__() - shape = (n_ch, 1, 1, 1) + shape = (in_channels, 1, 1, 1) - self.scale = n_ch**0.5 + self.scale = in_channels**0.5 self.gamma = nn.Parameter(torch.ones(shape)) self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 def forward(self, x, *args, **kwargs): - return F.normalize(x, dim=1) * self.scale * self.gamma + self.bias + needs_fp32_normalize = x.dtype in (torch.float16, torch.bfloat16) or any( + t in str(x.dtype) for t in ("float4_", "float8_") + ) + normalized = F.normalize(x.float() if needs_fp32_normalize else x, dim=1).to(x.dtype) + return normalized * self.scale * self.gamma + self.bias class CachedGroupNorm(nn.GroupNorm): - def group_forward(self, x_input, expectation=None, variance=None, return_stat=False): - input_dtype = x_input.dtype - x = x_input.to(torch.float32) - chunks = torch.chunk(x, self.num_groups, dim=1) - if expectation is None: - ch_mean = [torch.mean(chunk, dim=(1, 2, 3, 4), keepdim=True) for chunk in chunks] - else: - ch_mean = expectation - - if variance is None: - ch_var = [torch.var(chunk, dim=(1, 2, 3, 4), keepdim=True, unbiased=False) for chunk in chunks] - else: - ch_var = variance - - x_norm = [(chunk - mean) / torch.sqrt(var + self.eps) for chunk, mean, var in zip(chunks, ch_mean, ch_var)] - x_norm = torch.cat(x_norm, dim=1) - - x_norm.mul_(self.weight.data.view(1, -1, 1, 1, 1)) - x_norm.add_(self.bias.data.view(1, -1, 1, 1, 1)) + """GroupNorm alternative to `Kandinsky6VAERMSNorm`: same call convention (ignores `zq_ch`/`add_conv`, + accepts a `cache` dict) so the encoder/decoder can pick either one through the same `normalization(...)` + call. `cache` only ever records whether this layer has run before — `CachedSpatialNorm3D` reads that + presence flag to decide whether it's normalizing the first causal segment or a later one; the statistics + themselves are not reused across calls. + """ - x_out = x_norm.to(input_dtype) - if return_stat: - return x_out, ch_mean, ch_var - return x_out + def __init__(self, in_channels, **kwargs): + super().__init__(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) def forward(self, x, cache: dict): out = super().forward(x) @@ -215,10 +181,6 @@ def forward(self, x, cache: dict): return out -def Normalize(in_channels, gather=False, **kwargs): - return CachedGroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) - - class CachedCausalConv3d(CausalConv3d): def forward(self, input_, cache: dict, fixed_stride=False): t_stride = self.stride[0] @@ -295,8 +257,7 @@ def __init__( temb_channels=512, zq_ch=None, add_conv=False, - gather_norm=False, - normalization=Normalize, + normalization=CachedGroupNorm, padding_mode=None, ): super().__init__() @@ -323,7 +284,7 @@ def __init__( chan_in=in_channels, chan_out=out_channels, kernel_size=3, padding_mode=padding_mode ) else: - self.nin_shortcut = Conv3d( + self.nin_shortcut = SafeConv3d( in_channels, out_channels, kernel_size=1, @@ -344,7 +305,6 @@ def forward(self, x, temb, layer_cache, zq=None): if x.size(2) == 17 and x.size(3) == 1080 and zq is not None: torch.cuda.empty_cache() - # h = nonlinearity(h) h = F.silu(h, inplace=True) if x.size(2) == 17 and x.size(3) == 1080 and zq is not None: torch.cuda.empty_cache() @@ -354,14 +314,13 @@ def forward(self, x, temb, layer_cache, zq=None): torch.cuda.empty_cache() if temb is not None: - h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None, None] + h = h + self.temb_proj(F.silu(temb))[:, :, None, None, None] if zq is None: h = self.norm2(h, cache=layer_cache["norm2"]) else: h = self.norm2(h, zq, cache=layer_cache["norm2"]) - # h = nonlinearity(h) h = F.silu(h, inplace=True) # h = self.dropout(h) h = self.conv2(h, cache=layer_cache["conv2"]) @@ -389,7 +348,7 @@ def __init__( self.version = version out_channels = in_channels * 2 if version > 1 else in_channels - self.spatial_conv = Conv3d( + self.spatial_conv = SafeConv3d( in_channels, out_channels, kernel_size=(1, 3, 3), @@ -428,7 +387,7 @@ def __init__( padding_mode=padding_mode, ) - self.linear = nn.Conv3d(out_channels, out_channels, kernel_size=1, stride=1) + self.linear = SafeConv3d(out_channels, out_channels, kernel_size=1, stride=1) def spatial_downsample(self, input_): # PixelShuffle part @@ -491,9 +450,10 @@ def forward(self, x, cache): class CachedSpatialNorm3D(nn.Module): - """ - Looking at `forward`, it seems this class should be renamed to `DecoderSpatialNorm` or something similar, - because it's not usual forward, but a conditional one. + """Normalization used by the decoder's ResnetBlocks: modulates a GroupNorm/RMSNorm output with a + per-channel scale and shift derived from the latent `zq`, the same "spatial norm" conditioning other VAEs + in the library (e.g. CogVideoX) use to inject the encoder's output back into the decoder — adapted here + for causal, segment-by-segment video decoding. """ def __init__( @@ -503,7 +463,7 @@ def __init__( freeze_norm_layer=False, add_conv=False, padding_mode=None, - normalization=Normalize, + normalization=CachedGroupNorm, **norm_layer_params, ): super().__init__() @@ -515,20 +475,18 @@ def __init__( chan_in=zq_channels, chan_out=zq_channels, kernel_size=3, padding_mode=padding_mode ) - self.conv_y = Conv3d( + self.conv_y = SafeConv3d( zq_channels, f_channels, kernel_size=1, ) - self.conv_b = Conv3d( + self.conv_b = SafeConv3d( zq_channels, f_channels, kernel_size=1, ) def forward(self, f, zq, cache): - f_shape = [s for s in f.shape] - if cache["norm"]["mean"] is None and cache["norm"]["var"] is None: f_first, f_rest = f[:, :, :1], f[:, :, 1:] f_first_size, f_rest_size = f_first.shape[-3:], f_rest.shape[-3:] @@ -566,7 +524,7 @@ def forward(self, f, zq, cache): return norm_f -def Normalize3D(in_channels, zq_ch, add_conv, normalization=Normalize): +def Normalize3D(in_channels, zq_ch, add_conv, normalization=CachedGroupNorm): return CachedSpatialNorm3D( in_channels, zq_ch, @@ -585,7 +543,7 @@ def __init__(self, in_channels: int, compress_time: bool, factor: int = 2, paddi self.temporal_compress = compress_time self.factor = factor self.shuffle = nn.PixelShuffle(self.factor) - self.spatial_conv = Conv3d( + self.spatial_conv = SafeConv3d( in_channels, in_channels, kernel_size=(1, 3, 3), @@ -604,7 +562,7 @@ def __init__(self, in_channels: int, compress_time: bool, factor: int = 2, paddi padding_mode=padding_mode, ) - self.linear = Conv3d(in_channels, in_channels, kernel_size=1, stride=1) + self.linear = SafeConv3d(in_channels, in_channels, kernel_size=1, stride=1) def spatial_upsample_NEW(self, input_): def conv_part(x): @@ -646,79 +604,6 @@ def forward(self, x, cache): return lin_out -class DiagonalGaussianDistribution(object): - def __init__(self, parameters, deterministic=False): - self.parameters = parameters - self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) - self.logvar = torch.clamp(self.logvar, -30.0, 20.0) - self.deterministic = deterministic - self.std = torch.exp(0.5 * self.logvar) - self.var = torch.exp(self.logvar) - if self.deterministic: - self.var = self.std = torch.zeros_like(self.mean) # .to(device=self.parameters.device) - - def sample(self): - x = self.mean + self.std * torch.randn_like(self.mean) - return x - - def kl(self, other=None): - if self.deterministic: - return torch.Tensor([0.0]) - else: - if other is None: - return 0.5 * torch.sum( - torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, - dim=[1, 2, 3, 4], - ) - else: - return 0.5 * torch.sum( - torch.pow(self.mean - other.mean, 2) / other.var - + self.var / other.var - - 1.0 - - self.logvar - + other.logvar, - dim=[1, 2, 3, 4], - ) - - def mode(self): - return self.mean - - -class DiagonalGaussianRegularizer(nn.Module): - def __init__(self, sample: bool = True): - super().__init__() - self.sample = sample - - def get_trainable_parameters(self) -> Any: - yield from () - - def get_kl_loss(self, z: torch.Tensor) -> torch.Tensor: - posterior = DiagonalGaussianDistribution(z) - kl_loss = posterior.kl() - kl_loss = torch.mean(kl_loss, dim=0) - return kl_loss - - def get_sample(self, z: torch.Tensor) -> torch.Tensor: - posterior = DiagonalGaussianDistribution(z) - if self.sample: - z = posterior.sample() - else: - z = posterior.mode() - return z - - def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]: - log = dict() - posterior = DiagonalGaussianDistribution(z) - if self.sample: - z = posterior.sample() - else: - z = posterior.mode() - kl_loss = posterior.kl() - kl_loss = torch.sum(kl_loss) / kl_loss.shape[0] - log["kl_loss"] = kl_loss - return z, log - - class CachedEncoder3D(nn.Module): def __init__( self, @@ -727,20 +612,17 @@ def __init__( ch_mult=(1, 2, 4, 8), num_res_blocks, dropout=0.0, - resamp_with_conv=True, in_channels, resolution=0, z_channels, double_z=True, padding_mode=None, temporal_compress_times=4, - gather_norm=False, fix_pxs=False, norm_type="group_norm", downsample_version=1, temporal_compress_start_level=0, skip_last_resolution=False, - **ignore_kwargs, ): super().__init__() self.ch = ch @@ -759,7 +641,7 @@ def __init__( chan_in=in_channels, chan_out=round(in_ch_mult[0] * self.ch), kernel_size=3, padding_mode=padding_mode ) - normalization = Normalize if norm_type == "group_norm" else RMSNorm + normalization = CachedGroupNorm if norm_type == "group_norm" else Kandinsky6VAERMSNorm curr_res = resolution self.down = nn.ModuleList() @@ -780,7 +662,6 @@ def __init__( out_channels=block_out, dropout=dropout, temb_channels=self.temb_ch, - gather_norm=gather_norm, normalization=normalization, padding_mode=padding_mode, ) @@ -817,7 +698,6 @@ def __init__( out_channels=block_in, temb_channels=self.temb_ch, dropout=dropout, - gather_norm=gather_norm, normalization=normalization, padding_mode=padding_mode, ) @@ -827,13 +707,12 @@ def __init__( out_channels=block_in, temb_channels=self.temb_ch, dropout=dropout, - gather_norm=gather_norm, normalization=normalization, padding_mode=padding_mode, ) # end - self.norm_out = normalization(block_in, gather=gather_norm) + self.norm_out = normalization(block_in) self.conv_out = CachedCausalConv3d( chan_in=block_in, @@ -863,7 +742,7 @@ def forward(self, x, cache_dict, use_cp=True): # end h = self.norm_out(h, cache=cache_dict["norm_out"]) - h = nonlinearity(h) + h = F.silu(h) h = self.conv_out(h, cache=cache_dict["conv_out"]) return h @@ -880,7 +759,6 @@ def __init__( ch_mult=(1, 2, 4, 8), num_res_blocks=2, dropout=0.0, - resamp_with_conv=True, resolution=0, z_channels=16, give_pre_end=False, @@ -888,10 +766,8 @@ def __init__( add_conv=False, padding_mode=None, temporal_compress_times=4, - gather_norm=False, norm_type="group_norm", temporal_compress_start_level=0, - **kwargs, ): super().__init__() self.ch = ch @@ -917,7 +793,7 @@ def __init__( ) modulated_norm = functools.partial( - Normalize3D, normalization=Normalize if norm_type == "group_norm" else RMSNorm + Normalize3D, normalization=CachedGroupNorm if norm_type == "group_norm" else Kandinsky6VAERMSNorm ) # middle @@ -930,7 +806,6 @@ def __init__( zq_ch=zq_ch, add_conv=add_conv, normalization=modulated_norm, - gather_norm=gather_norm, padding_mode=padding_mode, ) @@ -942,7 +817,6 @@ def __init__( zq_ch=zq_ch, add_conv=add_conv, normalization=modulated_norm, - gather_norm=gather_norm, padding_mode=padding_mode, ) @@ -962,7 +836,6 @@ def __init__( zq_ch=zq_ch, add_conv=add_conv, normalization=modulated_norm, - gather_norm=gather_norm, padding_mode=padding_mode, ) ) @@ -981,7 +854,7 @@ def __init__( up.upsample = CachedPXSUpsample(block_in, compress_time=False, padding_mode=padding_mode) self.up.insert(0, up) - self.norm_out = modulated_norm(block_in, zq_ch, add_conv=add_conv) # , gather=gather_norm) + self.norm_out = modulated_norm(block_in, zq_ch, add_conv=add_conv) self.conv_out = CachedCausalConv3d(chan_in=block_in, chan_out=out_ch, kernel_size=3, padding_mode=padding_mode) @@ -991,8 +864,6 @@ def forward(self, z, cache_dict): # timestep embedding temb = None - t = z.shape[2] - zq = z h = self.conv_in(z, cache_dict["conv_in"]) @@ -1015,7 +886,7 @@ def forward(self, z, cache_dict): return h h = self.norm_out(h, zq, cache_dict["norm_out"]) - h = nonlinearity(h) + h = F.silu(h) h = self.conv_out(h, cache_dict["conv_out"]) return h @@ -1038,8 +909,6 @@ class Kandinsky6SRVAE(ModelMixin, ConfigMixin): Args: vae_type (`str`): VAE architecture identifier. Must be ``"video-kvae"``. - encoder_config (`dict`): Configuration for the cached 3D encoder. - decoder_config (`dict`): Configuration for the cached 3D decoder. scaling_factor (`float`, *optional*, defaults to 1.0): Latent scaling factor stored in the component configuration. spatial_factor (`int`, *optional*, defaults to 16): Spatial compression @@ -1049,6 +918,7 @@ class Kandinsky6SRVAE(ModelMixin, ConfigMixin): """ _no_split_modules = ["CachedEncoder3D", "CachedDecoder3D"] + _supports_gradient_checkpointing = False @staticmethod def normalize_data(data): @@ -1064,8 +934,25 @@ def denormalize_data(data): def __init__( self, vae_type: str, - encoder_config: Mapping[str, Any], - decoder_config: Mapping[str, Any], + in_channels: int = 3, + out_channels: int = 3, + z_channels: int = 16, + ch: int = 128, + ch_mult: Tuple[int, ...] = (1, 2, 4, 8), + num_res_blocks: int = 2, + dropout: float = 0.0, + resolution: int = 0, + padding_mode: Union[str, None] = None, + temporal_compress_times: int = 4, + temporal_compress_start_level: int = 0, + norm_type: str = "group_norm", + double_z: bool = True, + downsample_version: int = 1, + fix_pxs: bool = False, + skip_last_resolution: bool = False, + give_pre_end: bool = False, + zq_ch: Union[int, None] = None, + add_conv: bool = False, scaling_factor: float = 1.0, spatial_factor: int = 16, temporal_factor: int = 4, @@ -1073,10 +960,39 @@ def __init__( super().__init__() if vae_type != "video-kvae": raise ValueError(f"Kandinsky6SRVAE supports only 'video-kvae', got: {vae_type!r}") - self.conf = {"enc": encoder_config, "dec": decoder_config} - self.encoder = CachedEncoder3D(**encoder_config) - self.decoder = CachedDecoder3D(**decoder_config) - self.regularizer = DiagonalGaussianRegularizer(sample=False) + self.encoder = CachedEncoder3D( + in_channels=in_channels, + z_channels=z_channels, + ch=ch, + ch_mult=ch_mult, + num_res_blocks=num_res_blocks, + dropout=dropout, + resolution=resolution, + padding_mode=padding_mode, + temporal_compress_times=temporal_compress_times, + temporal_compress_start_level=temporal_compress_start_level, + norm_type=norm_type, + double_z=double_z, + downsample_version=downsample_version, + fix_pxs=fix_pxs, + skip_last_resolution=skip_last_resolution, + ) + self.decoder = CachedDecoder3D( + out_ch=out_channels, + z_channels=z_channels, + ch=ch, + ch_mult=ch_mult, + num_res_blocks=num_res_blocks, + dropout=dropout, + resolution=resolution, + padding_mode=padding_mode, + temporal_compress_times=temporal_compress_times, + temporal_compress_start_level=temporal_compress_start_level, + norm_type=norm_type, + give_pre_end=give_pre_end, + zq_ch=zq_ch, + add_conv=add_conv, + ) self.spatial_factor = int(spatial_factor) self.temporal_factor = int(temporal_factor) @@ -1115,8 +1031,8 @@ def make_dict(name, p=None): "norm_out": make_dict(f"norm_{block}"), "conv_out": make_dict("conv"), } - for i in range(len(self.conf[block].get("ch_mult", [1, 2, 4, 8]))): - cache[i] = make_dict(f"{i}_block", p=self.conf[block]["num_res_blocks"] + 1) + for i in range(len(self.config.ch_mult)): + cache[i] = make_dict(f"{i}_block", p=self.config.num_res_blocks + 1) return cache @apply_forward_hook @@ -1147,8 +1063,7 @@ def encode(self, x, seg_len=16): latent = [] for chunk in torch.split(x, split_list, dim=2): l = self.encoder(chunk, cache) - sample = self.regularizer.get_sample(l) - latent.append(sample) + latent.append(DiagonalGaussianDistribution(l).mode()) latent = torch.cat(latent, dim=2) return latent, split_list @@ -1169,7 +1084,7 @@ def decode(self, z, split_list=None): # Compute latent segment sizes. if split_list is None: - default_split_size = 16 // self.conf["enc"]["temporal_compress_times"] + default_split_size = 16 // self.config.temporal_compress_times time_dim = z.shape[2] if time_dim == 1: # image @@ -1181,7 +1096,7 @@ def decode(self, z, split_list=None): split_list.append((time_dim - 1) % default_split_size) split_list[0] += 1 else: - split_list = [math.ceil(size / self.conf["enc"]["temporal_compress_times"]) for size in split_list] + split_list = [math.ceil(size / self.config.temporal_compress_times) for size in split_list] # Decode each segment. recs = [] diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index 605fbfdf8a22..fe2de0f00a86 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -5,7 +5,6 @@ import math from typing import Any, Literal, Optional -import numpy as np import torch import torch.nn as nn import torch.nn.functional as F @@ -17,213 +16,75 @@ from torch.nn.utils.parametrizations import weight_norm from torch.nn.utils.parametrize import remove_parametrizations -# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# This work is licensed under a Creative Commons -# Attribution-NonCommercial-ShareAlike 4.0 International License. -# You should have received a copy of the license along with this -# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/ -# Improved diffusion model architecture proposed in the paper -# "Analyzing and Improving the Training Dynamics of Diffusion Models". +from .vae import DiagonalGaussianDistribution +# The magnitude-preserving building blocks below (`normalize`, `mp_silu`, `mp_sum`, `MPConv1D`) implement +# equations from Karras et al., "Analyzing and Improving the Training Dynamics of Diffusion Models" +# (https://arxiv.org/abs/2312.02696): each layer is designed so that if its inputs have unit variance, its +# output does too, which removes the need for the running activation-magnitude tracking earlier EDM variants +# relied on. MMAudio's VAE reuses this design for its 1D convolutional blocks. -# ---------------------------------------------------------------------------- -# Variant of constant() that inherits dtype and device from the given -# reference tensor by default. -_constant_cache = dict() +def normalize(x: torch.Tensor, dim: "list[int] | None" = None, eps: float = 1e-4) -> torch.Tensor: + """Rescale `x` to unit L2 norm over `dim` (default: every dimension but the first). - -def constant(value, shape=None, dtype=None, device=None, memory_format=None): - value = np.asarray(value) - if shape is not None: - shape = tuple(shape) - if dtype is None: - dtype = torch.get_default_dtype() - if device is None: - device = torch.device("cpu") - if memory_format is None: - memory_format = torch.contiguous_format - - key = (value.shape, value.dtype, value.tobytes(), shape, dtype, device, memory_format) - tensor = _constant_cache.get(key, None) - if tensor is None: - tensor = torch.as_tensor(value.copy(), dtype=dtype, device=device) - if shape is not None: - tensor, _ = torch.broadcast_tensors(tensor, torch.empty(shape)) - tensor = tensor.contiguous(memory_format=memory_format) - _constant_cache[key] = tensor - return tensor - - -def const_like(ref, value, shape=None, dtype=None, device=None, memory_format=None): - if dtype is None: - dtype = ref.dtype - if device is None: - device = ref.device - return constant(value, shape=shape, dtype=dtype, device=device, memory_format=memory_format) - - -# ---------------------------------------------------------------------------- -# Normalize given tensor to unit magnitude with respect to the given -# dimensions. Default = all dimensions except the first. - - -def normalize(x, dim=None, eps=1e-4): + `eps` is scaled by `sqrt(norm.numel() / x.numel())` so it stays a meaningful floor relative to the + norm's typical magnitude regardless of how many elements are reduced over. + """ if dim is None: dim = list(range(1, x.ndim)) norm = torch.linalg.vector_norm(x, dim=dim, keepdim=True, dtype=torch.float32) - norm = torch.add(eps, norm, alpha=np.sqrt(norm.numel() / x.numel())) + norm = eps + norm * math.sqrt(norm.numel() / x.numel()) return x / norm.to(x.dtype) -class Normalize(torch.nn.Module): - def __init__(self, dim=None, eps=1e-4): - super().__init__() - self.dim = dim - self.eps = eps - - def forward(self, x): - return normalize(x, dim=self.dim, eps=self.eps) - - -# ---------------------------------------------------------------------------- -# Upsample or downsample the given tensor with the given filter, -# or keep it as is. - - -def resample(x, f=[1, 1], mode="keep"): - if mode == "keep": - return x - f = np.float32(f) - assert f.ndim == 1 and len(f) % 2 == 0 - pad = (len(f) - 1) // 2 - f = f / f.sum() - f = np.outer(f, f)[np.newaxis, np.newaxis, :, :] - f = const_like(x, f) - c = x.shape[1] - if mode == "down": - return torch.nn.functional.conv2d(x, f.tile([c, 1, 1, 1]), groups=c, stride=2, padding=(pad,)) - assert mode == "up" - return torch.nn.functional.conv_transpose2d(x, (f * 4).tile([c, 1, 1, 1]), groups=c, stride=2, padding=(pad,)) - - -# ---------------------------------------------------------------------------- -# Magnitude-preserving SiLU (Equation 81). - - -def mp_silu(x): - return torch.nn.functional.silu(x) / 0.596 - - -class MPSiLU(torch.nn.Module): - def forward(self, x): - return mp_silu(x) - - -# ---------------------------------------------------------------------------- -# Magnitude-preserving sum (Equation 88). - - -def mp_sum(a, b, t=0.5): - return a.lerp(b, t) / np.sqrt((1 - t) ** 2 + t**2) - - -# ---------------------------------------------------------------------------- -# Magnitude-preserving concatenation (Equation 103). +def mp_silu(x: torch.Tensor) -> torch.Tensor: + """SiLU rescaled by the standard deviation of `silu(z)` for `z ~ N(0, 1)`, so a unit-variance input + stays unit-variance after activation. + """ + return F.silu(x) / 0.596 -def mp_cat(a, b, dim=1, t=0.5): - Na = a.shape[dim] - Nb = b.shape[dim] - C = np.sqrt((Na + Nb) / ((1 - t) ** 2 + t**2)) - wa = C / np.sqrt(Na) * (1 - t) - wb = C / np.sqrt(Nb) * t - return torch.cat([wa * a, wb * b], dim=dim) +def mp_sum(a: torch.Tensor, b: torch.Tensor, t: float = 0.5) -> torch.Tensor: + """Interpolate between `a` and `b` and rescale by `1 / sqrt((1-t)^2 + t^2)`, the factor that keeps the + result unit-variance when `a` and `b` are unit-variance and uncorrelated. + """ + return a.lerp(b, t) / math.sqrt((1 - t) ** 2 + t**2) -# ---------------------------------------------------------------------------- -# Magnitude-preserving convolution or fully-connected layer (Equation 47) -# with force weight normalization (Equation 66). +class MPConv1D(torch.nn.Module): + """1D convolution (or, with `kernel_size=1` on a 2D input, a linear layer) whose weight is normalized to + unit norm per output channel and rescaled by `1 / sqrt(fan_in)`, so its output variance matches its + input variance regardless of `in_channels` or `kernel_size`. + The weight starts as plain random init and only becomes usable after `remove_weight_norm()` folds the + normalization into `self.weight` once; `forward` then applies that folded weight directly rather than + renormalizing on every call, matching how the converted checkpoint stores it. + """ -class MPConv1D(torch.nn.Module): - def __init__(self, in_channels, out_channels, kernel_size): + def __init__(self, in_channels: int, out_channels: int, kernel_size: int): super().__init__() self.out_channels = out_channels self.weight = torch.nn.Parameter(torch.randn(out_channels, in_channels, kernel_size)) - self.weight_norm_removed = False - def forward(self, x, gain=1): - assert self.weight_norm_removed, "call remove_weight_norm() before inference" + def forward(self, x: torch.Tensor, gain: float = 1) -> torch.Tensor: + if not self.weight_norm_removed: + raise RuntimeError("call remove_weight_norm() before inference") w = self.weight * gain - if w.ndim == 2: - return x @ w.t() - assert w.ndim == 3 - return torch.nn.functional.conv1d(x, w, padding=(w.shape[-1] // 2,)) + return F.conv1d(x, w, padding=(w.shape[-1] // 2,)) - def remove_weight_norm(self): + def remove_weight_norm(self) -> "MPConv1D": w = self.weight.to(torch.float32) - w = normalize(w) # traditional weight normalization - w = w / np.sqrt(w[0].numel()) - w = w.to(self.weight.dtype) - self.weight.data.copy_(w) + w = normalize(w) + w = w / math.sqrt(w[0].numel()) + self.weight.data.copy_(w.to(self.weight.dtype)) self.weight_norm_removed = True return self -class DiagonalGaussianDistribution: - def __init__(self, parameters, deterministic=False): - self.parameters = parameters - self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) - self.logvar = torch.clamp(self.logvar, -30.0, 20.0) - self.deterministic = deterministic - self.std = torch.exp(0.5 * self.logvar) - self.var = torch.exp(self.logvar) - if self.deterministic: - self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device) - - def sample(self, rng: Optional[torch.Generator] = None): - # x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device) - - r = torch.empty_like(self.mean).normal_(generator=rng) - x = self.mean + self.std * r - - return x - - def kl(self, other=None): - if self.deterministic: - return torch.Tensor([0.0]) - else: - if other is None: - return 0.5 * torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar - else: - return 0.5 * ( - torch.pow(self.mean - other.mean, 2) / other.var - + self.var / other.var - - 1.0 - - self.logvar - + other.logvar - ) - - def nll(self, sample, dims=[1, 2, 3]): - if self.deterministic: - return torch.Tensor([0.0]) - logtwopi = np.log(2.0 * np.pi) - return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var, dim=dims) - - def mode(self): - return self.mean - - -def nonlinearity(x): - # swish - return mp_silu(x) - - class ResnetBlock1D(nn.Module): def __init__(self, *, in_dim, out_dim=None, conv_shortcut=False, kernel_size=3, use_norm=True): super().__init__() @@ -248,10 +109,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: x = normalize(x, dim=1) h = x - h = nonlinearity(h) + h = mp_silu(h) h = self.conv1(h) - h = nonlinearity(h) + h = mp_silu(h) h = self.conv2(h) if self.in_dim != self.out_dim: @@ -535,7 +396,7 @@ def forward(self, x): h = h.clamp(-self.clip_act, self.clip_act) # end - h = nonlinearity(h) + h = mp_silu(h) h = self.conv_out(h, gain=(self.learnable_gain + 1)) return h @@ -618,7 +479,7 @@ def forward(self, z): if i_level in self.down_layers: h = self.up[i_level].upsample(h) - h = nonlinearity(h) + h = mp_silu(h) h = self.conv_out(h, gain=(self.learnable_gain + 1)) return h diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index c0788ebcb7e0..898fa2a092e1 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -92,25 +92,6 @@ def _local_merge(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Te return x.flatten(dim, dim + 1).flatten(dim + 1, dim + 2).flatten(dim + 2, dim + 3) -def fractal_flatten(x: Tensor, rope: Tensor, shape: tuple, block_mask: bool = False): - """Flatten visual tokens, optionally arranging them into local fractal blocks.""" - if block_mask: - ps = 8 - x = _local_patch(x, shape, (1, ps, ps), dim=0) - rope = _local_patch(rope, shape, (1, ps, ps), dim=0) - return x.flatten(0, 1), rope.flatten(0, 1) - return x.flatten(0, 2), rope.flatten(0, 2) - - -def fractal_unflatten(x: Tensor, shape: tuple, block_mask: bool = False) -> Tensor: - """Restore the spatial layout produced by :func:`fractal_flatten`.""" - if block_mask: - ps = 8 - x = x.reshape(-1, ps * ps, x.shape[-1]) - return _local_merge(x, shape, (1, ps, ps), dim=0) - return x.reshape(*shape, x.shape[-1]) - - def nabla_block_mask( q: Tensor, k: Tensor, @@ -236,10 +217,6 @@ def _streams_from_args(args: tuple[Any, ...], kwargs: dict[str, Any]) -> tuple[T return video, audio -def _pack_streams(video: Tensor | None, audio: Tensor | None) -> Tensor | tuple[Tensor | None, Tensor | None]: - return video if audio is None else (video, audio) - - def _add_residual( video: Tensor | None, audio: Tensor | None, @@ -366,7 +343,7 @@ def new_forward(self, module: nn.Module, *args, **kwargs): if not state.should_compute: if self.is_tail: _advance(state, self.config, self.num_steps) - return _pack_streams(video, audio) + return video if audio is None else (video, audio) output = self.fn_ref.original_forward(*args, **kwargs) if self.is_tail: @@ -436,13 +413,12 @@ def disable_cache(self) -> None: self._cache_config = None -# Diffusers-style K6 transformer components for the source port. +# Diffusers-style K6 transformer components. # -# The assembler inlines this module after the portable attention, RoPE, and -# tensor helpers. The classes use the same inner-module vocabulary as the -# Diffusers Kandinsky5 transformer (for example ``in_layer``, ``modulation``, -# ``self_attention`` and ``feed_forward``). Checkpoint conversion maps the -# native K6 names to this public Diffusers layout. +# The classes use the same inner-module vocabulary as the Diffusers Kandinsky5 +# transformer (for example ``in_layer``, ``modulation``, ``self_attention`` and +# ``feed_forward``). Checkpoint conversion maps the native K6 names to this +# public Diffusers layout. _MASKED_ATTENTION_BACKENDS = { @@ -567,8 +543,7 @@ def forward(self, time: Tensor) -> Tensor: args = torch.outer(time.float(), freqs) embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) h = functional.linear(embed, self.in_layer.weight.float(), self.in_layer.bias.float()) - out = functional.linear(self.activation(h), self.out_layer.weight.float(), self.out_layer.bias.float()) - return out.to(dtype=self.out_layer.weight.dtype) + return functional.linear(self.activation(h), self.out_layer.weight.float(), self.out_layer.bias.float()) class Kandinsky6TextEmbeddings(nn.Module): @@ -580,8 +555,7 @@ def __init__(self, text_dim: int, model_dim: int): self.norm = nn.LayerNorm(model_dim, elementwise_affine=True) def forward(self, x: Tensor) -> Tensor: - x = self.in_layer(x.to(dtype=self.in_layer.weight.dtype)) - return self.norm(x).to(dtype=x.dtype) + return self.norm(self.in_layer(x)) class Kandinsky6VisualEmbeddings(nn.Module): @@ -622,7 +596,7 @@ def forward(self, x: Tensor) -> Tensor: .permute(0, 2, 4, 1, 3, 5, 6) .flatten(3, 6) ) - return self.in_layer(x.to(dtype=self.in_layer.weight.dtype)) + return self.in_layer(x) class Kandinsky6Modulation(nn.Module): @@ -665,7 +639,6 @@ def __init__( head_dim: int, kv_dim: int | None = None, text_token_padding: bool = False, - visual: bool = False, processor: Kandinsky6AttnProcessor | None = None, ): super().__init__() @@ -679,7 +652,6 @@ def __init__( self.query_norm = nn.RMSNorm(head_dim) self.key_norm = nn.RMSNorm(head_dim) self.out_layer = nn.Linear(num_channels, num_channels) - self.visual = visual self.text_token_padding = text_token_padding self.set_processor(processor or self._default_processor_cls()) if self.text_token_padding: @@ -825,7 +797,7 @@ def __init__( super().__init__() self.visual_modulation = Kandinsky6Modulation(time_dim, model_dim, 9) self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) - self.self_attention = Kandinsky6Attention(model_dim, head_dim, visual=True) + self.self_attention = Kandinsky6Attention(model_dim, head_dim) self.cross_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.cross_attention = Kandinsky6Attention( model_dim, @@ -1091,7 +1063,6 @@ class Kandinsky6Transformer3DModel( cross_gates (`bool`, *optional*, defaults to False): Whether to use cross-modal residual gates. fix_modulation (`bool`, *optional*, defaults to False): Whether to use the fixed modulation variant. visual_token_type_num_embeddings (`int`, *optional*, defaults to 0): Number of visual token type embeddings. - magcache (`dict`, *optional*): MagCache configuration. """ _repeated_blocks = [ @@ -1132,7 +1103,6 @@ def __init__( cross_gates: bool = False, fix_modulation: bool = False, visual_token_type_num_embeddings: int = 0, - magcache: dict[str, Any] | None = None, ) -> None: super().__init__() self.patch_size = patch_size @@ -1217,6 +1187,8 @@ def __init__( ] ) + self.gradient_checkpointing = False + def clear_text_proj_cache(self) -> None: """Clear cached projected text embeddings.""" self._text_proj_cache.clear() @@ -1264,7 +1236,10 @@ def _run_text_blocks( ) -> Tensor: blocks = self.text_transformer_blocks if prefix is None else getattr(self, f"{prefix}_text_transformer_blocks") for block in blocks: - te = block(te, tm, text_rope, attn_mask) + if torch.is_grad_enabled() and self.gradient_checkpointing: + te = self._gradient_checkpointing_func(block, te, tm, text_rope, attn_mask) + else: + te = block(te, tm, text_rope, attn_mask) return te def _encode_text( @@ -1336,18 +1311,23 @@ def _run_visual_blocks_single( sparse_params: dict | None, attn_mask: Tensor | None = None, ) -> tuple[Tensor | None, Tensor | None]: + checkpoint = torch.is_grad_enabled() and self.gradient_checkpointing for block in self.visual_transformer_blocks: if self.is_multimodal: if vis_embed is not None and aud_embed is None: - vis_embed, _ = block(vis_embed, None, te, te, (tm, tm), vis_rope, None, sparse_params, attn_mask) + args = (vis_embed, None, te, te, (tm, tm), vis_rope, None, sparse_params, attn_mask) + vis_embed, _ = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) elif aud_embed is not None and vis_embed is None: - _, aud_embed = block(None, aud_embed, te, te, (tm, tm), None, aud_rope, None, attn_mask) + args = (None, aud_embed, te, te, (tm, tm), None, aud_rope, None, attn_mask) + _, aud_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) else: raise RuntimeError("single-modality fused path expects exactly one of video/audio") elif vis_embed is not None: - vis_embed = block(vis_embed, te, tm, vis_rope, sparse_params, attn_mask) + args = (vis_embed, te, tm, vis_rope, sparse_params, attn_mask) + vis_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) else: - aud_embed = block(aud_embed, te, tm, aud_rope, None, attn_mask) + args = (aud_embed, te, tm, aud_rope, None, attn_mask) + aud_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) return vis_embed, aud_embed def _run_visual_blocks_fused( @@ -1363,8 +1343,9 @@ def _run_visual_blocks_fused( sparse_params: dict | None, attn_mask: Tensor | None = None, ) -> tuple[Tensor, Tensor]: + checkpoint = torch.is_grad_enabled() and self.gradient_checkpointing for block in self.visual_transformer_blocks: - vis_embed, aud_embed = block( + args = ( vis_embed, aud_embed, video_te, @@ -1375,6 +1356,7 @@ def _run_visual_blocks_fused( sparse_params, attn_mask, ) + vis_embed, aud_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) return vis_embed, aud_embed def _project_video(self, vis_embed: Tensor, vis_shape: tuple, tm: Tensor, sparse_params: dict | None) -> Tensor: diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index 90d10e5f42c4..d831c3cd8ae1 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -79,7 +79,7 @@ def __call__( parallel_config=self._parallel_config, )[0] ) - return torch.cat(outputs, dim=0) + return torch.cat(outputs, dim=0).flatten(-2, -1) @staticmethod def assemble_cached_attention_inputs( @@ -1104,9 +1104,10 @@ def reset_parameters(self) -> None: class MultiheadSelfAttention(nn.Module, AttentionModuleMixin): + """Multi-head self-attention with flash attention and optional sparse flex attention.""" + _default_processor_cls = Kandinsky6SRAttentionProcessor _available_processors = [Kandinsky6SRAttentionProcessor] - """Multi-head self-attention with flash attention and optional sparse flex attention.""" def __init__(self, num_channels: int, head_dim: int) -> None: """Initialize multi-head self-attention. @@ -1198,7 +1199,7 @@ def attention_flex( Attention output, or tuple of (output, sparsity percentage) if requested. """ if self.cached_k is not None: - msg = "KV-cache on the flex/NABLA path is out of scope (dense flash only in M0.2; see LAY-429)." + msg = "KV-cache is not supported on the flex/NABLA attention path; only dense flash attention supports it." raise NotImplementedError(msg) query = query.unsqueeze(0).transpose(1, 2).contiguous() @@ -1306,9 +1307,10 @@ def reset_parameters(self) -> None: class MultiheadCrossAttention(nn.Module, AttentionModuleMixin): + """Multi-head cross-attention with flash attention.""" + _default_processor_cls = Kandinsky6SRAttentionProcessor _available_processors = [Kandinsky6SRAttentionProcessor] - """Multi-head cross-attention with flash attention.""" def __init__(self, num_channels: int, head_dim: int) -> None: """Initialize multi-head cross-attention. @@ -1785,20 +1787,6 @@ def __init__( [TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim) for _ in range(num_blocks)] ) - def _patchify(self, x: Tensor, visual_cu_seqlens: Tensor) -> tuple[Tensor, Tensor]: - """Patchify LQ latent with the same logic as VisualEmbeddings._patchify.""" - pt, ph, pw = self.patch_size - if pt > 1: - idxs = torch.ones(x.shape[0], dtype=torch.int32, device=visual_cu_seqlens.device) - idxs[visual_cu_seqlens[:-1]] += pt - 1 - x = torch.repeat_interleave(x, idxs, dim=0) - visual_cu_seqlens = visual_cu_seqlens + torch.arange( - visual_cu_seqlens.shape[0], device=visual_cu_seqlens.device, dtype=torch.int32 - ) - T, H, W, C = x.shape - x = x.view(T // pt, pt, H // ph, ph, W // pw, pw, C).permute(0, 2, 4, 1, 3, 5, 6).flatten(3, 6) - return x, visual_cu_seqlens // pt - def forward( self, lq_visual: Tensor, @@ -1813,7 +1801,16 @@ def forward( to_fractal: bool, ) -> list[Tensor]: """Process LQ tokens through adapter blocks, return γ-scaled features.""" - lq_patches, _ = self._patchify(lq_visual, pre_patch_cu_seqlens) + # Patchify LQ latent with the same logic as VisualEmbeddings._patchify. + pt, ph, pw = self.patch_size + if pt > 1: + idxs = torch.ones(lq_visual.shape[0], dtype=torch.int32, device=pre_patch_cu_seqlens.device) + idxs[pre_patch_cu_seqlens[:-1]] += pt - 1 + lq_visual = torch.repeat_interleave(lq_visual, idxs, dim=0) + T, H, W, C = lq_visual.shape + lq_patches = ( + lq_visual.view(T // pt, pt, H // ph, ph, W // pw, pw, C).permute(0, 2, 4, 1, 3, 5, 6).flatten(3, 6) + ) adapter_embed = self.lq_proj(lq_patches) if to_fractal: visual_shape = adapter_embed.shape[:-1] @@ -1880,7 +1877,6 @@ def __init__( use_text: bool = True, use_lq_noise_cond: bool = False, attribute_overrides: Mapping[str, Any] | None = None, - sr_params: Mapping[str, Any] | None = None, ) -> None: """Initialize the 3D diffusion transformer. diff --git a/src/diffusers/pipelines/kandinsky6/__init__.py b/src/diffusers/pipelines/kandinsky6/__init__.py index 61d785b59ea8..18c6f961473c 100644 --- a/src/diffusers/pipelines/kandinsky6/__init__.py +++ b/src/diffusers/pipelines/kandinsky6/__init__.py @@ -2,19 +2,38 @@ from typing import TYPE_CHECKING -from ...utils import DIFFUSERS_SLOW_IMPORT, _LazyModule, is_torch_available +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, +) +_dummy_objects = {} _import_structure = {} -if is_torch_available(): +try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_pt_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_pt_objects)) +else: _import_structure["pipeline_kandinsky6_ti2va"] = ["Kandinsky6TI2VAPipeline"] _import_structure["pipeline_kandinsky6_sr"] = ["Kandinsky6SRPipeline"] _import_structure["latent_upscaler"] = ["Kandinsky6SRLatentUpscalerBank"] _import_structure["pipeline_output"] = ["Kandinsky6SRPipelineOutput", "Kandinsky6TI2VAPipelineOutput"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: - if is_torch_available(): + try: + if not is_torch_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_pt_objects import * + else: from .pipeline_kandinsky6_ti2va import Kandinsky6TI2VAPipeline from .pipeline_kandinsky6_sr import Kandinsky6SRPipeline from .latent_upscaler import Kandinsky6SRLatentUpscalerBank @@ -29,6 +48,9 @@ module_spec=__spec__, ) + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) + __all__ = [ "Kandinsky6TI2VAPipeline", diff --git a/src/diffusers/pipelines/kandinsky6/latent_upscaler.py b/src/diffusers/pipelines/kandinsky6/latent_upscaler.py index f177814ef896..332375e50bfb 100644 --- a/src/diffusers/pipelines/kandinsky6/latent_upscaler.py +++ b/src/diffusers/pipelines/kandinsky6/latent_upscaler.py @@ -35,6 +35,11 @@ from torch.nn import functional +if TYPE_CHECKING: + from collections.abc import Callable + from types import ModuleType + + def rearrange(*args: Any, **kwargs: Any) -> Any: """Load einops only when the latent upscaler is actually executed.""" try: @@ -96,11 +101,6 @@ def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: UpsamplePaddingMode = Literal["reflect", "zeros"] -def as_triple(value: int | tuple[int, int, int]) -> tuple[int, int, int]: - """Expand a scalar kernel/padding spec into an explicit ``(t, h, w)`` triple.""" - return (value, value, value) if isinstance(value, int) else value - - class TemporalReplicateConv3d(nn.Conv3d): """``Conv3d`` that repeats the edge frame along T and zero-pads H/W. @@ -139,7 +139,7 @@ def __init__( groups: int = 1, ) -> None: """Initialize the convolution with H/W padding only, keeping T for forward.""" - pad_t, pad_h, pad_w = as_triple(padding) + pad_t, pad_h, pad_w = (padding, padding, padding) if isinstance(padding, int) else padding super().__init__( in_channels, out_channels, @@ -582,10 +582,6 @@ def zero_init_residual(self) -> None: # Exact merging of flash-style attention branches with a closed-form null branch. -if TYPE_CHECKING: - from collections.abc import Callable, Sequence - - def merge_math( outputs: Sequence[Tensor], lses: Sequence[Tensor], @@ -701,10 +697,6 @@ def merge_attention_branches( # Runtime helpers shared by latent-upscaler model components. -if TYPE_CHECKING: - from collections.abc import Callable - - def forward_with_checkpointing( module: Callable[..., Tensor], *inputs: Tensor, @@ -719,11 +711,6 @@ def forward_with_checkpointing( # Spatial upsampling operations: pixel-shuffle, bilinear, and K-VAE PXS v2. -if TYPE_CHECKING: - from collections.abc import Callable - -DIMS_2 = 2 -DIMS_3 = 3 VALID_DIMS = {DIMS_2, DIMS_3} @@ -1008,9 +995,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # Motion-correspondence attention for latent video features. -if TYPE_CHECKING: - from types import ModuleType - MotionAttentionBackend = Literal["sdpa", "natten"] NattenBackend = Literal["cutlass-fna", "hopper-fna", "flex-fna"] @@ -1454,11 +1438,6 @@ def build_residual_stack(count: int, start_index: int, spec: ResidualStackSpec) return nn.Sequential(*blocks) -def build_motion_attention_stack(spec: MotionAttentionSpec) -> nn.ModuleList: - """Build one motion block for every configured mid-stage placement.""" - return nn.ModuleList([MotionCorrespondenceBlock(spec.block) for _placement in spec.after_mid_blocks]) - - def apply_block_sequence( x: Tensor, zq: Tensor | None, @@ -2157,7 +2136,9 @@ def stage_spec(channels: int, first_in: int | None) -> ResidualStackSpec: self.mid_motion_blocks: nn.ModuleList | None = None if motion_attention is not None: - self.mid_motion_blocks = build_motion_attention_stack(motion_attention) + self.mid_motion_blocks = nn.ModuleList( + [MotionCorrespondenceBlock(motion_attention.block) for _placement in motion_attention.after_mid_blocks] + ) # Build x2-exclusive modules after the complete x4 path so enabling an x2 # entry does not perturb the base-path RNG initialization. @@ -2360,10 +2341,6 @@ def _add_nearest_base( # Factory functions for building upsampler models. -if TYPE_CHECKING: - from torch import nn - - def build_upsampler(config: FlatModelConfig | MultiScaleModelConfig) -> nn.Module: """Create an upsampler module from config. @@ -2479,6 +2456,8 @@ class Kandinsky6SRLatentUpscalerBank(ModelMixin, ConfigMixin): latents before they are passed to an upscaler. """ + _no_split_modules = ["ConvLatentUpsampler", "MultiScaleUpsampler"] + @register_to_config def __init__( self, diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py index 0a93a5817cf7..e30824329fd6 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py @@ -75,7 +75,6 @@ TotalResolutionScale = Literal[2, 4, 2.25] constants = SimpleNamespace(VAE_SPATIAL_FACTOR=VAE_SPATIAL_FACTOR) TileGridMode = Literal["legacy", "even"] -_sr_constants = None def _call_module_method(module: nn.Module, method: str, *args: Any, **kwargs: Any) -> Any: @@ -450,8 +449,9 @@ def get_visual_size(x: torch.Tensor) -> int: Scales the tensor's spatial dims by the VAE spatial factor and looks them up in the ``RESOLUTIONS`` registry. Reads ``constants.VAE_SPATIAL_FACTOR`` - late-bound so ``set_vae_factors`` (16 for the KVAE) applies — a - ``from``-import would freeze the import-time default. + late-bound (via attribute access on the module) so a caller can override + ``constants.VAE_SPATIAL_FACTOR`` for a different VAE — a ``from``-import + would freeze the import-time default instead. Args: x: Visual latent tensor of shape ``(T, H, W, C)``. @@ -602,56 +602,31 @@ def replicate_cached_embeds( return {"text_embeds": text_embeds, "pooled_embed": pooled_embed}, cu_seqlens -def kvae_weights_path(checkpoint_prefix: str) -> str: - """Return the kvae weights file for a sidecar prefix, preferring safetensors. - - The release format is ``{prefix}.safetensors`` (flat state dict), the - training format ``{prefix}.ckpt`` (pickled ``{"state_dict": ...}``); both - load through ``CachedCausalVAE.init_from_ckpt``, which dispatches on the - suffix. - - Args: - checkpoint_prefix: Local kvae sidecar prefix (``{prefix}.yaml`` sits - next to the weights). - - Returns: - Path to the weights file to load. - """ - safetensors_path = f"{checkpoint_prefix}.safetensors" - return safetensors_path if Path(safetensors_path).exists() else f"{checkpoint_prefix}.ckpt" - - -def module_dtype(module: torch.nn.Module) -> torch.dtype | None: - """Return the dtype used by the first parameterized layer, if any.""" - try: - return next(module.parameters()).dtype - except StopIteration: - return None - - def cast_to_module_dtype(module: torch.nn.Module, value: torch.Tensor) -> torch.Tensor: """Cast floating-point inputs to the module's parameter dtype.""" - dtype = module_dtype(module) - if dtype is not None and value.is_floating_point() and value.dtype != dtype: + try: + dtype = next(module.parameters()).dtype + except StopIteration: + return value + if value.is_floating_point() and value.dtype != dtype: return value.to(dtype=dtype) return value -def encode_pixels_to_latent(vae: torch.nn.Module, pixels: torch.Tensor, *, sample: bool = True) -> torch.Tensor: +def encode_pixels_to_latent(vae: torch.nn.Module, pixels: torch.Tensor) -> torch.Tensor: """Encode ``(B, C, T, H, W)`` pixels in ``[0, 255]`` with the KVAE. Normalizes with the KVAE's own convention, then returns the latent - (``(latent, split_list)[0]`` — the regularizer mode). + (``(latent, split_list)[0]`` — the regularizer mode; the KVAE always + returns this mode, there is no sampling toggle). Args: vae: Causal video KVAE. pixels: ``(B, C, T, H, W)`` tensor in ``[0, 255]``. - sample: Unused — the KVAE always returns the regularizer mode. Returns: Latent ``(B, C, T', H', W')`` — not yet scaled by ``scaling_factor``. """ - del sample # the KVAE always returns the regularizer mode x = vae.normalize_data(pixels) x = cast_to_module_dtype(vae, x) result = vae.encode(x) @@ -675,19 +650,6 @@ def decode_latent_to_uint8(vae: torch.nn.Module, latent: torch.Tensor) -> torch. return denormalize_to_uint8(vae, vae.decode(latent).sample) -def denormalize_to_float(vae: torch.nn.Module, decoded: torch.Tensor) -> torch.Tensor: - """Convert a decoded tensor to float pixels in ``[0, 1]`` without quantization. - - Args: - vae: Causal video KVAE. - decoded: Decode output in the VAE's normalized pixel space. - - Returns: - Float pixels in ``[0, 1]`` with the same shape. - """ - return (vae.denormalize_data(decoded) / 255.0).clamp(0.0, 1.0) - - def denormalize_to_uint8(vae: torch.nn.Module, decoded: torch.Tensor) -> torch.Tensor: """Convert an already-decoded tensor in ``[-1, 1]`` to ``uint8`` ``[0, 255]``. @@ -702,7 +664,8 @@ def denormalize_to_uint8(vae: torch.nn.Module, decoded: torch.Tensor) -> torch.T Returns: ``uint8`` pixels in ``[0, 255]`` with the same shape. """ - return (denormalize_to_float(vae, decoded) * 255.0).to(torch.uint8) + float_pixels = (vae.denormalize_data(decoded) / 255.0).clamp(0.0, 1.0) + return (float_pixels * 255.0).to(torch.uint8) def latent_upscaler_scale(upscaler: nn.Module) -> int: @@ -1512,10 +1475,7 @@ def _spatial_factor(components: Any) -> int: vae = getattr(components, "vae", None) config = getattr(vae, "config", None) return int( - getattr(vae, "spatial_factor", None) - or getattr(config, "spatial_factor", None) - or getattr(_sr_constants, "VAE_SPATIAL_FACTOR", None) - or VAE_SPATIAL_FACTOR + getattr(vae, "spatial_factor", None) or getattr(config, "spatial_factor", None) or VAE_SPATIAL_FACTOR ) @@ -1558,14 +1518,6 @@ def _extract_latent(result: Any) -> torch.Tensor: return result -def _validate_equal_shapes(values: list[torch.Tensor], name: str) -> None: - if not values: - raise ValueError(f"{name} batch must not be empty") - shape = tuple(values[0].shape) - if any(tuple(value.shape) != shape for value in values[1:]): - raise ValueError(f"all {name} batch items must have equal shape") - - @torch.no_grad() def encode_lq_videos_to_lr_latents( lq_videos: list[torch.Tensor], @@ -1573,7 +1525,7 @@ def encode_lq_videos_to_lr_latents( device: str | torch.device, ) -> torch.Tensor: """Encode an equal-shaped batch of ``[T,C,H,W]`` videos in one VAE call.""" - _validate_equal_shapes(lq_videos, "video") + _validate_batch_shapes(lq_videos, "video") pixel = torch.stack(lq_videos).permute(0, 2, 1, 3, 4).to(device=device) pixel = vae.normalize_data(pixel.float()) if hasattr(vae, "normalize_data") else pixel.float() / 127.5 - 1.0 pixel = cast_to_module_dtype(vae, pixel) @@ -2074,17 +2026,6 @@ def _run_tile_batches( return outputs -def _require_av(): - """Load PyAV only for file-based video input.""" - try: - import av - except (ImportError, OSError) as exc: - raise RuntimeError( - "SR file-based video input requires PyAV and FFmpeg. Install with `pip install av`." - ) from exc - return av - - def align_to_vae_stride(t: int) -> int: """Round ``t`` to the nearest valid pixel-frame count ``1 + 8·k`` (ties up). @@ -2530,7 +2471,7 @@ def _run_batched_tiles( # noqa: PLR0913 ) -> Tensor: """Run equal-shape samples tile-major and restore the sample batch.""" params = _component_params(components) - visual_size = int(params.visual_size[0]) if isinstance(params.visual_size, list) else int(params.visual_size) + visual_size = _visual_size(params) batch_size = len(samples) if latent_input: spatial_factor = _spatial_factor(components) @@ -2604,7 +2545,7 @@ def _sr_progress_plan( ) -> tuple[int, int, int]: """Return total updates, denoising steps, and tile batches.""" params = _component_params(components) - visual_size = int(params.visual_size[0]) if isinstance(params.visual_size, list) else int(params.visual_size) + visual_size = _visual_size(params) spatial_factor = _spatial_factor(components) height, width = value.shape[-2:] if latent_input: @@ -2706,7 +2647,6 @@ def check_inputs( overlap: float | None = None, tiles_batch_size: int | None = None, kvae_bridge: bool = False, - save_path: str | Path | list[str | Path] | None = None, fps: int | None = None, audio: list[np.ndarray] | np.ndarray | None = None, source_video: str | Path | list[str | Path] | None = None, @@ -2723,7 +2663,6 @@ def check_inputs( num_inference_steps: Number of denoising steps. overlap: Tile overlap fraction in ``[0, 1)``. tiles_batch_size: Number of tiles processed together. - save_path: Optional output video path. fps: Optional output video frame rate. audio: Optional audio waveform to mux. source_video: Optional source video whose audio is copied. @@ -2828,7 +2767,6 @@ def _run( # noqa: PLR0913 overlap=effective_overlap, tiles_batch_size=effective_tiles_batch_size, kvae_bridge=kvae_bridge, - save_path=save_path, fps=fps, audio=audio, source_video=source_video, diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index c7995bc97a55..ac3c2e5c71f8 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -215,11 +215,6 @@ def _rebuild_text_rope(dit: nn.Module, template: Tensor | list[Tensor]) -> Tenso return compute_rope1d(module, int(template.shape[0])) -def apply_cfg(cond: Tensor, uncond: Tensor, guidance_weight: float) -> Tensor: - """Classifier-free guidance: uncond + w * (cond - uncond).""" - return uncond + guidance_weight * (cond - uncond) - - @torch.no_grad() def denoise_loop( # noqa: PLR0912, PLR0913, PLR0915 bundle: LatentBundle, @@ -299,7 +294,8 @@ def _cache_scope(name: str): return cache_context(name) if callable(cache_context) else nullcontext() def _guided(cond: Tensor, uncond: Tensor) -> Tensor: - return cond if is_piflow else apply_cfg(cond, uncond, guidance_weight) + # Classifier-free guidance: uncond + w * (cond - uncond) + return cond if is_piflow else uncond + guidance_weight * (cond - uncond) for step_index, t in enumerate(step_timesteps): # Diffusers schedulers expose the model-scale timestep directly. @@ -555,7 +551,8 @@ def postprocess_audio( return None cu = bundle.audio_cu_seqlens - assert cu is not None + if cu is None: + raise ValueError("postprocess_audio requires bundle.audio_cu_seqlens when bundle.audio is set") bs = cu.shape[0] - 1 # Reverse audio VAE normalization @@ -610,7 +607,8 @@ def postprocess_video( ``(bs, T, H_lat, W_lat, C)``. """ video = bundle.video - assert video is not None + if video is None: + raise ValueError("postprocess_video requires bundle.video") if video.ndim == 4: frames = video.reshape(bs, -1, video.shape[-3], video.shape[-2], video.shape[-1]) @@ -803,7 +801,7 @@ def __init__( tokenizer_2=tokenizer_2, scheduler=scheduler, ) - transformer_config = self.transformer.config + transformer_config = self.transformer.config if getattr(self, "transformer", None) else {} scale_factor = transformer_config.get("scale_factor", (1.0, 2.0, 2.0)) text_token_padding = transformer_config.get("text_token_padding", True) self.scale_factor = tuple(float(value) for value in scale_factor) @@ -1034,6 +1032,12 @@ def encode_prompt( lengths = attention.sum(dim=1).to(torch.int32) cu_seqlens = torch.cat([torch.zeros(1, dtype=torch.int32, device=embeds.device), lengths.cumsum(0)]) + if qwen_attention is not None and qwen_attention.all(): + # An all-True mask carries no information beyond "attend to everything", which is what the + # transformer already does when no mask is passed (see transformer_kandinsky6.py's + # _normalize_attn_mask). Drop it so the dense mask isn't threaded through the denoise loop. + qwen_attention = None + clip_device = next(self.text_encoder_2.parameters()).device clip_inputs = self.tokenizer_2( texts, @@ -1362,14 +1366,6 @@ def _text_ropes( ) return text_rope, negative_text_rope, audio_rope - @staticmethod - def _format_video(frames: Tensor, output_type: str) -> Tensor | np.ndarray: - if output_type in ("pt", "torch"): - return frames - if output_type in ("np", "numpy"): - return frames.permute(0, 2, 3, 4, 1).cpu().numpy() - raise ValueError("output_type must be 'pt', 'torch', 'np', 'numpy', or 'latent'") - @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( @@ -1668,7 +1664,13 @@ def _step_callback(step_index: int, timestep: Tensor, current: LatentBundle) -> else: packed_result = self._packed_bundle(result) decoded_frames = postprocess_video(packed_result, self.vae, bs=batch_size) - frames = self._format_video(decoded_frames, output_type) + # check_inputs already restricted output_type to pt/torch/np/numpy/latent, and "latent" is + # handled above, so only the tensor-vs-numpy choice is left here. + frames = ( + decoded_frames + if output_type in ("pt", "torch") + else decoded_frames.permute(0, 2, 3, 4, 1).cpu().numpy() + ) audio = postprocess_audio(self._packed_bundle(result), self.audio_vae) if sample_audio else None if not return_dict: return frames, audio diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_output.py b/src/diffusers/pipelines/kandinsky6/pipeline_output.py index d70e41b830a7..042f8d1f5b6c 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_output.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_output.py @@ -20,6 +20,7 @@ from typing import Any import numpy as np + from ...utils import BaseOutput Audio = list[np.ndarray] | np.ndarray | None diff --git a/src/diffusers/schedulers/scheduling_piflow.py b/src/diffusers/schedulers/scheduling_piflow.py index b66c907bc964..db367649e427 100644 --- a/src/diffusers/schedulers/scheduling_piflow.py +++ b/src/diffusers/schedulers/scheduling_piflow.py @@ -18,7 +18,7 @@ import torch from ..configuration_utils import register_to_config -from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( +from .scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, FlowMatchEulerDiscreteSchedulerOutput, ) @@ -37,7 +37,6 @@ def __init__( # noqa: PLR0913 mode: str = "grid", eps: float = 1e-4, ) -> None: - self.x_t_src = x_t_src self.ndim = x_t_src.dim() self.shift = shift self.eps = eps @@ -45,14 +44,14 @@ def __init__( # noqa: PLR0913 raise ValueError(f"Unknown mode: {mode}") self.mode = mode - 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) + 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(sigma_t_src) segment = segment_size if isinstance(segment, torch.Tensor) and segment.dim() < self.raw_t_src.dim(): segment = segment.reshape(*segment.size(), *((self.raw_t_src.dim() - segment.dim()) * [1])) self.raw_t_dst = (self.raw_t_src - segment).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) + self.denoising_output_x_0 = x_t_src.unsqueeze(1) - sigma_t_src.unsqueeze(1) * denoising_output @staticmethod def _interpolate(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: @@ -69,14 +68,6 @@ def _interpolate(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: def _unwarp_t(self, sigma_t: torch.Tensor) -> torch.Tensor: return sigma_t / (self.shift + (1 - self.shift) * sigma_t) - @staticmethod - def _u_to_x_0( - denoising_output: torch.Tensor, - x_t: torch.Tensor, - sigma_t: torch.Tensor, - ) -> torch.Tensor: - return x_t.unsqueeze(1) - sigma_t.unsqueeze(1) * denoising_output - def pi(self, x_t: torch.Tensor, sigma_t: torch.Tensor) -> torch.Tensor: sigma_t = sigma_t.reshape(*sigma_t.size(), *((self.ndim - sigma_t.dim()) * [1])) raw_t = self._unwarp_t(sigma_t) @@ -268,7 +259,7 @@ def _policy_step(self, model_output: torch.Tensor, sample: torch.Tensor, step_in self.num_policy_substeps, policy, ) - return updated.to(dtype=sample.dtype) + return updated def _step_index_for(self, timestep: torch.Tensor | float) -> int: if self.step_index is None: @@ -276,8 +267,6 @@ def _step_index_for(self, timestep: torch.Tensor | float) -> int: self._step_index = self.begin_index else: schedule_timesteps = self.timesteps - if not isinstance(schedule_timesteps, torch.Tensor): - schedule_timesteps = torch.as_tensor(schedule_timesteps, dtype=torch.float32) timestep = torch.as_tensor( timestep, device=schedule_timesteps.device, @@ -318,9 +307,11 @@ def step( "pass a value from scheduler.timesteps instead" ) step_index = self._step_index_for(timestep) - # PiFlow's policy rollout performs its update in float32. Keep that - # precision across outer steps, matching the native sampler. + # PiFlow's policy rollout performs its update in float32, matching the + # native sampler; cast back to the input dtype before returning, like + # the base Euler scheduler does. updated = self._policy_step(model_output, sample.to(torch.float32), step_index) + updated = updated.to(dtype=sample.dtype) self._step_index += 1 if return_dict: return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=updated) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 6f9b9bf477ae..4c5400229ef3 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1590,6 +1590,51 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class Kandinsky6SRLatentUpscalerBank(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 Kandinsky6SRPipeline(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 Kandinsky6SRPipelineOutput(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 Kandinsky6SRTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] @@ -1620,6 +1665,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class Kandinsky6TI2VAPipeline(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 Kandinsky6TI2VAPipelineOutput(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 Kandinsky6Transformer3DModel(metaclass=DummyObject): _backends = ["torch"] From 0eb5211be0eb48e42d3ffe94a3c50b0f54ce704b Mon Sep 17 00:00:00 2001 From: leffff Date: Fri, 25 Sep 2026 16:21:14 +0300 Subject: [PATCH 11/62] patch transformer k6 --- src/diffusers/models/transformers/transformer_kandinsky6.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 898fa2a092e1..897937ee9610 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -1063,6 +1063,7 @@ class Kandinsky6Transformer3DModel( cross_gates (`bool`, *optional*, defaults to False): Whether to use cross-modal residual gates. fix_modulation (`bool`, *optional*, defaults to False): Whether to use the fixed modulation variant. visual_token_type_num_embeddings (`int`, *optional*, defaults to 0): Number of visual token type embeddings. + magcache (`dict`, *optional*): MagCache configuration. """ _repeated_blocks = [ @@ -1103,6 +1104,7 @@ def __init__( cross_gates: bool = False, fix_modulation: bool = False, visual_token_type_num_embeddings: int = 0, + magcache: dict[str, Any] | None = None, ) -> None: super().__init__() self.patch_size = patch_size From c3012d08c6bebc3208ae997d798afcbaf126d344 Mon Sep 17 00:00:00 2001 From: leffff Date: Sun, 27 Sep 2026 16:49:44 +0300 Subject: [PATCH 12/62] patch SR --- .../autoencoders/autoencoder_kandinsky6_sr.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index 6e37844a6638..5f710788b069 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -909,6 +909,11 @@ class Kandinsky6SRVAE(ModelMixin, ConfigMixin): Args: vae_type (`str`): VAE architecture identifier. Must be ``"video-kvae"``. + decoder_ch (`int`, *optional*): Base channel count for the decoder, if + different from the encoder's ``ch``. Defaults to ``ch``. + decoder_ch_mult (`tuple[float, ...]`, *optional*): Per-resolution channel + multiplier for the decoder, if different from the encoder's + ``ch_mult``. Defaults to ``ch_mult``. scaling_factor (`float`, *optional*, defaults to 1.0): Latent scaling factor stored in the component configuration. spatial_factor (`int`, *optional*, defaults to 16): Spatial compression @@ -938,7 +943,9 @@ def __init__( out_channels: int = 3, z_channels: int = 16, ch: int = 128, - ch_mult: Tuple[int, ...] = (1, 2, 4, 8), + ch_mult: Tuple[float, ...] = (1, 2, 4, 8), + decoder_ch: Union[int, None] = None, + decoder_ch_mult: Union[Tuple[float, ...], None] = None, num_res_blocks: int = 2, dropout: float = 0.0, resolution: int = 0, @@ -980,8 +987,8 @@ def __init__( self.decoder = CachedDecoder3D( out_ch=out_channels, z_channels=z_channels, - ch=ch, - ch_mult=ch_mult, + ch=ch if decoder_ch is None else decoder_ch, + ch_mult=ch_mult if decoder_ch_mult is None else decoder_ch_mult, num_res_blocks=num_res_blocks, dropout=dropout, resolution=resolution, From 1f7f234cb63e349645eeb10a397ae8f9b620c546 Mon Sep 17 00:00:00 2001 From: leffff Date: Sun, 27 Sep 2026 18:00:53 +0300 Subject: [PATCH 13/62] patch pipeline style --- .../kandinsky6/pipeline_kandinsky6_ti2va.py | 1723 ++++++----------- src/diffusers/schedulers/scheduling_piflow.py | 87 +- 2 files changed, 604 insertions(+), 1206 deletions(-) diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index ac3c2e5c71f8..458976d37051 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -14,25 +14,23 @@ """Kandinsky 6 TI2VA Diffusers pipeline.""" -from __future__ import annotations - +import copy import math import warnings from collections.abc import Callable -from contextlib import nullcontext -from copy import deepcopy -from dataclasses import dataclass from pathlib import Path -from typing import Any, TypedDict +from typing import Any import numpy as np import torch -from ..pipeline_utils import DiffusionPipeline -from diffusers.utils import replace_example_docstring from torch import Tensor, nn +from ...utils import replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline from .pipeline_output import Kandinsky6TI2VAPipelineOutput + EXAMPLE_DOC_STRING = """ Examples: @@ -67,24 +65,6 @@ ``` """ - -@dataclass -class LatentBundle: - """Packed K6 latent state; audio is optional for a video-only call.""" - - video: Tensor | None - audio: Tensor | None - video_cu_seqlens: Tensor | None - audio_cu_seqlens: Tensor | None - - -class TextEmbeds(TypedDict): - """Text adapter output consumed by the K6 DiT.""" - - text_embeds: Tensor - pooled_embed: Tensor - - _PROMPT_TEMPLATE = "\n".join( [ "<|im_start|>system\nYou are a promt engineer. Describe the video in detail.", @@ -122,657 +102,34 @@ class TextEmbeds(TypedDict): ) -def _prepare_expansion_image(image: Any) -> Any: - if isinstance(image, (str, Path)): - from PIL import Image - - with Image.open(image) as pil_image: - return pil_image.convert("RGB").copy() - return image - - -def _build_video_input( - video: Tensor, - has_visual_cond: bool, - first_frames: Tensor | None, - video_cu_seqlens: Tensor | None, - visual_cond_scheme: str, -) -> Tensor: - """Append visual-conditioning channels for batched ``(B,T,H,W,C)`` video.""" - if not has_visual_cond: - return video - - cond = torch.zeros_like(video) - mask = torch.zeros((*video.shape[:-1], 1), dtype=video.dtype, device=video.device) - if first_frames is None: - return torch.cat([video, cond, mask], dim=-1) - if video_cu_seqlens is None: - raise ValueError(f"{visual_cond_scheme} requires video_cu_seqlens") - - first_frames = first_frames.to(device=video.device, dtype=video.dtype) - tail_cond = visual_cond_scheme == "tail_cond_first_frame" - if visual_cond_scheme == "i2v": - video[:, 0] = first_frames - elif tail_cond: - video[:, -1] = first_frames - elif visual_cond_scheme == "pretrain": - cond[:, 0] = first_frames - else: - raise ValueError(f"unknown visual_cond_scheme={visual_cond_scheme!r}") - mask[:, -1 if tail_cond else 0] = 1 - return torch.cat([video, cond, mask], dim=-1) - - -def _resolve_null_embeds( - null_text_embeds: TextEmbeds | list[TextEmbeds | None], - null_text_rope: Tensor | list[Tensor | None], -) -> tuple[Tensor, Tensor, Tensor | list[Tensor]]: - """Returns (text_embed, pooled_embed, text_rope) for the uncond DiT pass. - - When null_text_embeds is a list of two items (T2VA per-modality nulls), - returns list arguments so the DiT can apply different null text per modality. - """ - if isinstance(null_text_embeds, list): - # null_text_embeds[0] = video null, null_text_embeds[1] = audio null (may be None) - null_v = null_text_embeds[0] - null_a = null_text_embeds[1] - null_rope_v = null_text_rope[0] if isinstance(null_text_rope, list) else null_text_rope - null_rope_a = null_text_rope[1] if isinstance(null_text_rope, list) else None - - if null_a is not None and null_rope_a is not None: - # Different null text per modality — pass as list to DiT - text_embed = [null_v["text_embeds"], null_a["text_embeds"]] - pooled_embed = [null_v["pooled_embed"], null_a["pooled_embed"]] - rope = [null_rope_v, null_rope_a] - else: - text_embed = null_v["text_embeds"] - pooled_embed = null_v["pooled_embed"] - rope = null_rope_v - else: - text_embed = null_text_embeds["text_embeds"] - pooled_embed = null_text_embeds["pooled_embed"] - rope = null_text_rope - - return text_embed, pooled_embed, rope - - -def _raw_dit(dit: nn.Module) -> nn.Module: - inner = getattr(dit, "module", None) - return inner if inner is not None and hasattr(dit, "set_cache") else dit - - -def _rebuild_text_rope(dit: nn.Module, template: Tensor | list[Tensor]) -> Tensor | list[Tensor]: - """Rebuild text RoPE using the standalone Diffusers transformer names.""" - raw = getattr(dit, "module", None) - raw = raw if raw is not None and hasattr(raw, "set_cache") else dit - if isinstance(template, list): - if getattr(raw, "is_multimodal", False): - modules = [raw.video_text_rope_embeddings, raw.audio_text_rope_embeddings] - else: - modules = [raw.text_rope_embeddings] * len(template) - return [compute_rope1d(module, int(item.shape[0])) for module, item in zip(modules, template, strict=True)] - module = raw.video_text_rope_embeddings if getattr(raw, "is_multimodal", False) else raw.text_rope_embeddings - return compute_rope1d(module, int(template.shape[0])) - - -@torch.no_grad() -def denoise_loop( # noqa: PLR0912, PLR0913, PLR0915 - bundle: LatentBundle, - dit: nn.Module, - text_embeds: TextEmbeds, - null_text_embeds: TextEmbeds | list[TextEmbeds | None], - visual_rope: Tensor | None, - audio_rope: Tensor | None, - text_rope: Tensor | list[Tensor], - null_text_rope: Tensor | list[Tensor | None], - num_steps: int, - guidance_weight: float, - scheduler: Any, - first_frames: Tensor | None = None, - visual_cond_scheme: str = "pretrain", - sparse_params: dict | None = None, - sample_video: bool = True, - sample_audio: bool = True, - *, - attention_mask: Tensor | None = None, - null_attention_mask: Tensor | None = None, - visual_token_type_ids: Tensor | None = None, - recompute_ropes_each_step: bool = False, - scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), - step_callback: Callable[[int, Tensor, LatentBundle], LatentBundle | Tensor | None] | None = None, - progress_bar: Any | None = None, -) -> LatentBundle: - """Run the native Euler loop and invoke a Diffusers step callback. - - ``step_callback`` is called after each Euler update with the zero-based - step index, the current scheduler timestep, and the mutable latent - bundle. It may return ``None``, a replacement ``LatentBundle``, or a - replacement video ``Tensor``. The callback is deliberately optional so - the generated pipeline has the same numerical behavior when it is not - supplied. - """ - video = bundle.video # (B, T, H, W, C) or (B, T, H, W, C+extra) for instruct - audio = bundle.audio # (B, A, audio_dim) or None - is_multimodal = video is not None and audio is not None - is_piflow = bool(getattr(scheduler, "is_piflow", False)) - guidance_epsilon = 1e-6 - if is_piflow and abs(guidance_weight - 1.0) > guidance_epsilon: - raise ValueError("PiflowScheduler requires guidance_weight=1.0") - - audio_scheduler = deepcopy(scheduler) if is_multimodal else scheduler - device = video.device if video is not None else audio.device - scheduler.set_timesteps(num_steps, device=device) - if audio_scheduler is not scheduler: - audio_scheduler.set_timesteps(num_steps, device=device) - step_timesteps = scheduler.timesteps - if step_timesteps.shape[0] != num_steps: - raise ValueError( - "Diffusers scheduler returned an unexpected number of timesteps: " - f"expected {num_steps}, got {step_timesteps.shape[0]}" - ) - - bs = ( - bundle.video_cu_seqlens.shape[0] - 1 - if bundle.video_cu_seqlens is not None - else bundle.audio_cu_seqlens.shape[0] - 1 - ) - - null_te, null_pe, null_rope = _resolve_null_embeds(null_text_embeds, null_text_rope) - - out_c: int | None = None # determined after first step, used to strip instruct channels - raw = _raw_dit(dit) - vis_shape = ( - (int(visual_rope.shape[0]), int(visual_rope.shape[1]), int(visual_rope.shape[2])) - if visual_rope is not None - else None - ) - scale = (float(scale_factor[0]), float(scale_factor[1]), float(scale_factor[2])) - tail_cond = visual_cond_scheme == "tail_cond_first_frame" - - def _cache_scope(name: str): - cache_context = getattr(dit, "cache_context", None) - return cache_context(name) if callable(cache_context) else nullcontext() - - def _guided(cond: Tensor, uncond: Tensor) -> Tensor: - # Classifier-free guidance: uncond + w * (cond - uncond) - return cond if is_piflow else uncond + guidance_weight * (cond - uncond) - - for step_index, t in enumerate(step_timesteps): - # Diffusers schedulers expose the model-scale timestep directly. - t_step: Tensor | list[Tensor] = t.unsqueeze(0).expand(bs) # (bs,) - - model_input_v = ( - _build_video_input( - video, - dit.visual_cond, - first_frames, - bundle.video_cu_seqlens, - visual_cond_scheme, - ) - if video is not None - else None - ) - - # Freeze one modality at t=0 for partial sampling (T2VA only) - if is_multimodal: - t_frozen = torch.zeros_like(t_step[0]).expand(bs) - if not sample_audio: - t_step = [t_step, t_frozen] - elif not sample_video: - t_step = [t_frozen, t_step] - - def _forward( - te, - pe, - rope, - attn_mask, - *, - _model_input_v=model_input_v, - _audio=audio, - _t_step=t_step, - ): - vr, ar, tr = visual_rope, audio_rope, rope - if recompute_ropes_each_step: - if vis_shape is not None: - vr = compute_visual_rope(raw.visual_rope_embeddings, vis_shape, scale) - if audio_rope is not None: - ar = compute_rope1d(raw.audio_rope_embeddings, int(audio_rope.shape[0])) - tr = _rebuild_text_rope(raw, rope) - return dit( - x_video=_model_input_v, - x_audio=_audio, - text_embed=te, - pooled_text_embed=pe, - time=_t_step, - visual_rope=vr, - audio_rope=ar, - text_rope=tr, - sparse_params=sparse_params, - attention_mask=attn_mask, - visual_token_type_ids=visual_token_type_ids, - ) - - with _cache_scope("cond"): - vel_cond = _forward( - text_embeds["text_embeds"], - text_embeds["pooled_embed"], - text_rope, - attention_mask, - ) - if abs(guidance_weight - 1.0) > guidance_epsilon: - with _cache_scope("uncond"): - vel_uncond = _forward(null_te, null_pe, null_rope, null_attention_mask) - else: - vel_uncond = vel_cond - - # Diffusers schedulers own the update for each modality. Multimodal - # runs use an independent scheduler state for audio because the - # standard scheduler advances its step index on every call. - if isinstance(vel_cond, tuple): - vel_v, vel_a = vel_cond - uvel_v, uvel_a = vel_uncond if vel_uncond is not vel_cond else (vel_v, vel_a) - if out_c is None: - out_c = video.shape[-1] if video is not None else vel_v.shape[-1] - - if video is not None and sample_video: - guided_v = _guided(vel_v, uvel_v) - video = scheduler.step(guided_v, t, video, return_dict=False)[0] - if tail_cond and first_frames is not None: - video[:, -1] = first_frames.to(device=device, dtype=video.dtype) - - if audio is not None and sample_audio: - guided_a = _guided(vel_a, uvel_a) - audio = audio_scheduler.step(guided_a, t, audio, return_dict=False)[0] - else: - if out_c is None: - out_c = video.shape[-1] if video is not None else vel_cond.shape[-1] - vel = _guided(vel_cond, vel_uncond) - if video is not None and sample_video: - video = scheduler.step(vel, t, video, return_dict=False)[0] - if tail_cond and first_frames is not None: - video[:, -1] = first_frames.to(device=device, dtype=video.dtype) - elif audio is not None and sample_audio: - audio = scheduler.step(vel, t, audio, return_dict=False)[0] - - if step_callback is not None: - callback_result = step_callback( - step_index, - t, - LatentBundle( - video=video, - audio=audio, - video_cu_seqlens=bundle.video_cu_seqlens, - audio_cu_seqlens=bundle.audio_cu_seqlens, - ), - ) - if isinstance(callback_result, LatentBundle): - video = callback_result.video - audio = callback_result.audio - elif isinstance(callback_result, Tensor): - video = callback_result - - if progress_bar is not None: - progress_bar.update() - - # I2V / I2VA: keep injected reference frames unchanged - if first_frames is not None and video is not None and bundle.video_cu_seqlens is not None: - ff = first_frames.to(device=device, dtype=video.dtype) - if visual_cond_scheme == "i2v": - video[:, 0] = ff - elif tail_cond: - video[:, -1] = ff - - # Strip instruct extra-channel padding from the video latent - if out_c is not None and video is not None: - video = video[..., :out_c] - - return LatentBundle( - video=video, - audio=audio, - video_cu_seqlens=bundle.video_cu_seqlens, - audio_cu_seqlens=bundle.audio_cu_seqlens, - ) - - -def prepare_video_latents( - bs: int, - duration: int, - H_lat: int, - W_lat: int, - C: int, - seed: int, - device: torch.device | str, - dtype: torch.dtype = torch.bfloat16, -) -> LatentBundle: - """Sample random video noise latent and compute cu_seqlens. - - Returns a LatentBundle with video=(bs*duration, H_lat, W_lat, C) and - uniform video_cu_seqlens=(bs+1,) int32. - """ - g = torch.Generator(device=device) - g.manual_seed(seed) - video = torch.randn(bs * duration, H_lat, W_lat, C, device=device, dtype=dtype, generator=g) - cu = duration * torch.arange(bs + 1, dtype=torch.int32, device=device) - return LatentBundle(video=video, audio=None, video_cu_seqlens=cu, audio_cu_seqlens=None) - - -def audio_latent_duration( - video_latent_frames: int, - *, - fps: float = 24.0, - audio_fps: int = 44100, - downsample_factor: int = 1024, -) -> int: - """Audio latent length matching K5 T2VA: ceil(sample_frames/fps * audio_fps / downsample).""" - sample_frames = (video_latent_frames - 1) * 4 + 1 - return int(math.ceil(sample_frames / fps * audio_fps / downsample_factor)) - - -def prepare_audio_latents( - bundle: LatentBundle, - audio_duration: int, - audio_dim: int, - seed: int, - device: torch.device | str, - dtype: torch.dtype = torch.bfloat16, -) -> LatentBundle: - """Attach random audio noise latent to an existing LatentBundle. - - audio_duration: number of latent audio frames per batch item. - """ - bs = bundle.video_cu_seqlens.shape[0] - 1 if bundle.video_cu_seqlens is not None else 1 - g = torch.Generator(device=device) - g.manual_seed(seed + 1) # offset from video seed for independence - audio = torch.randn(bs * audio_duration, audio_dim, device=device, dtype=dtype, generator=g) - cu = audio_duration * torch.arange(bs + 1, dtype=torch.int32, device=device) - return LatentBundle( - video=bundle.video, - audio=audio, - video_cu_seqlens=bundle.video_cu_seqlens, - audio_cu_seqlens=cu, - ) - - -def compute_rope1d(rope: nn.Module, length: int, device: torch.device | None = None) -> Tensor: - """Build 1-D RoPE and move non-persistent tables for Diffusers offload.""" - if length < 1: - raise ValueError(f"rope length must be positive, got {length}") - if device is None: - device = next(rope.buffers()).device - else: - device = torch.device(device) - # Accelerate's model CPU-offload hook only moves the transformer when - # its forward starts. RoPE is materialized before that forward, so - # explicitly move its non-persistent lookup table first. - rope.to(device) - pos = torch.arange(length, device=device) - return rope(pos) - - -def compute_visual_rope( - rope: nn.Module, - shape: tuple[int, int, int], - scale_factor: tuple[float, float, float], - device: torch.device | None = None, -) -> Tensor: - """Build 3-D RoPE and move non-persistent tables for Diffusers offload.""" - T, H, W = (int(shape[0]), int(shape[1]), int(shape[2])) - if T < 1 or H < 1 or W < 1: - raise ValueError(f"visual rope shape must be positive, got {shape}") - if device is None: - device = next(rope.buffers()).device - else: - device = torch.device(device) - rope.to(device) - pos = [ - torch.arange(T, device=device), - torch.arange(H, device=device), - torch.arange(W, device=device), - ] - scale = (float(scale_factor[0]), float(scale_factor[1]), float(scale_factor[2])) - return rope((T, H, W), pos, scale) - - -@torch.no_grad() -def postprocess_audio( - bundle: LatentBundle, - audio_vae, - normalization_mode: str = "clip", -) -> list[np.ndarray] | None: - """Decode audio latents → list of (samples,) int16 numpy arrays. - - Returns None when bundle.audio is None (T2V mode). - - ``clip`` preserves the decoded amplitude and saturates it to [-1, 1]. - ``normalize`` peak-normalizes each waveform before converting it to int16. - """ - audio = bundle.audio - if audio is None: - return None - - cu = bundle.audio_cu_seqlens - if cu is None: - raise ValueError("postprocess_audio requires bundle.audio_cu_seqlens when bundle.audio is set") - bs = cu.shape[0] - 1 - - # Reverse audio VAE normalization - audio_scaled = audio / getattr(audio_vae, "scaling_factor", 1.0) - audio_scaled = audio_scaled + getattr(audio_vae, "mean_value", 0.0) - - segments = [ - audio_scaled[i] if audio_scaled.ndim == 3 else audio_scaled[cu[i].item() : cu[i + 1].item()] for i in range(bs) - ] # (A, audio_dim) per sample - - # Native generation keeps equal-length samples in the batch. Decode them - # together so batching is preserved through the audio VAE and vocoder. - if len({tuple(segment.shape) for segment in segments}) == 1: - decoded = audio_vae.wrapped_decode(torch.stack(segments).transpose(1, 2)) - if decoded.ndim == 1: - decoded = decoded.unsqueeze(0) - elif decoded.ndim == 3 and decoded.shape[1] == 1: - decoded = decoded[:, 0] - waveforms = [waveform.cpu().float().numpy() for waveform in decoded] - else: - # Preserve support for legacy packed bundles with variable lengths. - waveforms = [ - audio_vae.wrapped_decode(segment.transpose(1, 0).unsqueeze(0)).squeeze().cpu().float().numpy() - for segment in segments - ] - - result: list[np.ndarray] = [] - for waveform in waveforms: - if normalization_mode == "normalize": - peak = np.max(np.abs(waveform)) - if peak > 0: - waveform = waveform / peak * 32767 - elif normalization_mode == "clip": - waveform = np.clip(waveform, -1.0, 1.0) * 32767 - else: - raise ValueError(f"unknown audio normalization_mode={normalization_mode}") - - result.append(waveform.astype(np.int16)) - - return result - - -@torch.no_grad() -def postprocess_video( - bundle: LatentBundle, - vae, - bs: int, -) -> Tensor: - """Decode video latents → (bs, 3, T, H, W) uint8 in [0, 255]. - - Input layout: packed ``(sum_T, H_lat, W_lat, C)`` or batched - ``(bs, T, H_lat, W_lat, C)``. - """ - video = bundle.video - if video is None: - raise ValueError("postprocess_video requires bundle.video") - - if video.ndim == 4: - frames = video.reshape(bs, -1, video.shape[-3], video.shape[-2], video.shape[-1]) - elif video.ndim == 5: - if video.shape[0] != bs: - raise ValueError(f"batched video has batch size {video.shape[0]}, expected {bs}") - frames = video - else: - raise ValueError(f"video must have rank 4 or 5, got {video.ndim}") - # (bs, T, H, W, C) → (bs, C, T, H, W) for VAE input - frames = (frames / vae.config.scaling_factor).permute(0, 4, 1, 2, 3) - # Hunyuan VAE loads as fp16; DiT latents are bf16 — match weight dtype. - vae_dtype = next(vae.parameters()).dtype - frames = vae.decode(frames.to(dtype=vae_dtype)).sample - - return ((frames.clamp(-1.0, 1.0) + 1.0) * 127.5).to(torch.uint8) - - -@torch.no_grad() -def resize_image( - image: Tensor, - max_area: int, - divisibility: int = 16, - world_size: int = 1, -) -> tuple[Tensor, float]: - """Aspect-preserving resize to fit ``max_area`` with sides divisible by ``divisibility``. - - K5 ``i2v_pipeline.resize_image`` parity. ``image`` is ``(B, C, H, W)``. - Returns ``(resized, scale_k)``. - """ - from math import sqrt - - h, w = image.shape[2:] - area = h * w - div = divisibility - if div == 16: - if world_size in (2, 4): - div *= 2 - elif world_size == 8: - div *= 4 - k = sqrt(max_area / area) / div - new_h = int(round(h * k) * div) - new_w = int(round(w * k) * div) - try: - import torchvision.transforms.functional as TF - except (ImportError, OSError) as exc: - raise RuntimeError( - "I2VA image processing requires torchvision. Install it with `pip install torchvision`." - ) from exc - - return TF.resize(image, (new_h, new_w)), k - - -def _load_pil_rgb(image: str | object): - from PIL import Image - - if isinstance(image, str): - try: - pil_image = Image.open(image).convert("RGB") - pil_image.load() - except Exception as exc: - raise ValueError(f"Cannot decode i2va input image {image!r}: {exc}") from exc - elif isinstance(image, Image.Image): - pil_image = image.convert("RGB") - else: - raise TypeError(f"i2va image must be a path or PIL image, got {type(image).__name__}") - return pil_image - - -@torch.no_grad() -def encode_i2va_first_frame( - image: str | object, - vae, - device: torch.device | str, - height: int | None = None, - width: int | None = None, - *, - max_area: int | None = None, - divisibility: int = 16, - world_size: int = 1, -) -> tuple[Tensor, int, int]: - """Encode one image into the K6 packed first-frame latent layout.""" - try: - import torchvision.transforms.functional as TF - except (ImportError, OSError) as exc: - raise RuntimeError( - "I2VA image processing requires torchvision. Install it with `pip install torchvision`." - ) from exc - - pil_image = _load_pil_rgb(image) - tensor = TF.pil_to_tensor(pil_image).unsqueeze(0) - if max_area is not None and (height is None or width is None): - tensor, _ = resize_image(tensor, max_area, divisibility, world_size) - height, width = int(tensor.shape[-2]), int(tensor.shape[-1]) - elif height is not None and width is not None: - src_h, src_w = tensor.shape[-2:] - scale = min(src_h / height, src_w / width) - tensor = TF.resize(tensor, (int(src_h / scale), int(src_w / scale))) - cur_h, cur_w = tensor.shape[-2:] - tensor = TF.crop(tensor, (cur_h - height) // 2, (cur_w - width) // 2, height, width) - else: - raise ValueError("encode_i2va_first_frame needs height/width or max_area") - - tensor = tensor.to(device=device, dtype=next(vae.parameters()).dtype) / 127.5 - 1.0 - tensor = tensor.transpose(0, 1).unsqueeze(0) - encoded = vae.encode(tensor) - posterior = encoded.latent_dist if hasattr(encoded, "latent_dist") else encoded[0] - latent = posterior.sample() - latent = latent.squeeze(0).permute(1, 2, 3, 0) - return latent * vae.config.scaling_factor, int(height), int(width) - - -def append_i2va_tail_condition( - latent_visual: Tensor, - first_frames: Tensor, - batch_size: int, - video_duration: int, -) -> tuple[Tensor, Tensor, Tensor]: - """Append one clean reference frame to batched ``(B,T,H,W,C)`` latents.""" - if latent_visual.shape[0] != batch_size or latent_visual.shape[1] != video_duration: - raise ValueError( - "generated visual latent shape mismatch: expected " - f"({batch_size}, {video_duration}, H, W, C), got {tuple(latent_visual.shape)}" - ) - - _, _, height, width, dim = latent_visual.shape - first_frames = first_frames.to(device=latent_visual.device, dtype=latent_visual.dtype) - expected = (batch_size, height, width, dim) - if tuple(first_frames.shape) != expected: - raise ValueError(f"first-frame latent shape mismatch: expected {expected}, got {tuple(first_frames.shape)}") - - latent_visual = torch.cat([latent_visual, first_frames[:, None]], dim=1) - token_types = torch.cat( - [ - torch.zeros((batch_size, video_duration), dtype=torch.long, device=latent_visual.device), - torch.ones((batch_size, 1), dtype=torch.long, device=latent_visual.device), - ], - dim=1, - ) - return latent_visual, token_types, token_types == 0 - - class Kandinsky6TI2VAPipeline(DiffusionPipeline): r"""Pipeline for text/image-to-video-and-audio generation with Kandinsky 6. - This pipeline inherits the generic loading, device placement, and progress - handling provided by ``DiffusionPipeline``. The transformer must expose - the K6 multimodal forward contract, while the VAE adapters must implement - the portable video/audio post-processing contracts. + This pipeline inherits the generic loading, device placement, and progress handling provided by + [`DiffusionPipeline`]. The transformer must expose the K6 multimodal forward contract, while the VAE adapters + must implement the portable video/audio post-processing contracts. Args: - transformer: Multimodal K6 transformer used to denoise video and audio - latents. - vae: Video VAE used to decode generated video latents. - text_encoder: Qwen2.5-VL text encoder for token-level embeddings. - audio_vae: Audio VAE used to decode generated audio latents. It may be - ``None`` only when ``sample_audio=False``. - scheduler: Diffusion scheduler used by the denoising loop. - tokenizer: Qwen2.5-VL processor. - text_encoder_2: CLIP text encoder for pooled embeddings. - tokenizer_2: CLIP tokenizer. + transformer ([`Kandinsky6Transformer3DModel`]): + Multimodal K6 transformer used to denoise video and audio latents. + vae (`nn.Module`): + Video VAE used to encode and decode video latents. + text_encoder: + Qwen2.5-VL text encoder for token-level embeddings. + audio_vae (`nn.Module`, *optional*): + Audio VAE used to decode generated audio latents. May be `None` only when `sample_audio=False`. + scheduler: + Diffusion scheduler used by the denoising loop. + tokenizer: + Qwen2.5-VL processor. + text_encoder_2: + CLIP text encoder for pooled embeddings. + tokenizer_2: + CLIP tokenizer. """ model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae->audio_vae" + _optional_components = ["audio_vae"] _callback_tensor_inputs = ["latents"] _DEFAULT_NEGATIVE_PROMPT = ( "Static, 2D cartoon, cartoon, 2d animation, paintings, images, " @@ -801,13 +158,14 @@ def __init__( tokenizer_2=tokenizer_2, scheduler=scheduler, ) - transformer_config = self.transformer.config if getattr(self, "transformer", None) else {} - scale_factor = transformer_config.get("scale_factor", (1.0, 2.0, 2.0)) - text_token_padding = transformer_config.get("text_token_padding", True) + scale_factor = ( + self.transformer.config.get("scale_factor", (1.0, 2.0, 2.0)) + if getattr(self, "transformer", None) is not None + else (1.0, 2.0, 2.0) + ) self.scale_factor = tuple(float(value) for value in scale_factor) - self.text_token_padding = bool(text_token_padding) - audio_config = self.audio_vae.config if self.audio_vae is not None else {} + audio_config = self.audio_vae.config if getattr(self, "audio_vae", None) is not None else {} self.audio_sample_rate = int(audio_config.get("sample_rate", 44_100)) self.audio_downsample_factor = int(audio_config.get("downsample_factor", 1_024)) @@ -819,9 +177,6 @@ def _value_batch_size(value: Any) -> int: return len(value) if isinstance(value, Tensor): return int(value.shape[0]) if value.ndim in (3, 5) else 1 - if isinstance(value, dict) and isinstance(value.get("text_embeds"), Tensor): - embeds = value["text_embeds"] - return int(embeds.shape[0]) if embeds.ndim == 3 else 1 return 1 @classmethod @@ -832,20 +187,12 @@ def _batch_size( image: str | object | list[str | object] | None, latents: Tensor | None, audio_latents: Tensor | None, - prompt_embeds: TextEmbeds | None, - negative_prompt_embeds: TextEmbeds | None, + prompt_embeds: Tensor | None, + negative_prompt_embeds: Tensor | None, ) -> int: sizes = [ cls._value_batch_size(value) - for value in ( - prompt, - negative_prompt, - image, - latents, - audio_latents, - prompt_embeds, - negative_prompt_embeds, - ) + for value in (prompt, negative_prompt, image, latents, audio_latents, prompt_embeds, negative_prompt_embeds) ] batch_size = max(sizes) if batch_size < 1 or any(size not in (1, batch_size) for size in sizes): @@ -873,10 +220,10 @@ def check_inputs( num_inference_steps: int, max_sequence_length: int, output_type: str, - prompt_embeds: TextEmbeds | None = None, - negative_prompt_embeds: TextEmbeds | None = None, - prompt_cu_seqlens: Tensor | None = None, - negative_prompt_cu_seqlens: Tensor | None = None, + prompt_embeds: Tensor | None = None, + pooled_prompt_embeds: Tensor | None = None, + negative_prompt_embeds: Tensor | None = None, + negative_pooled_prompt_embeds: Tensor | None = None, callback_on_step_end_tensor_inputs: list[str] | None = None, sample_audio: bool = True, image: str | object | None = None, @@ -886,24 +233,21 @@ def check_inputs( Args: prompt: Text prompt or a batch of prompts. - image: Optional reference image or batch of reference images. negative_prompt: Negative text prompt or a batch of prompts. height: Requested output height in pixels. width: Requested output width in pixels. num_frames: Requested output frame count. num_inference_steps: Number of denoising steps. max_sequence_length: Maximum Qwen prompt length after the template. - output_type: One of ``pt``, ``torch``, ``np``, ``numpy``, or - ``latent``. - prompt_embeds: Optional precomputed positive prompt embeddings. - negative_prompt_embeds: Optional precomputed negative embeddings. - prompt_cu_seqlens: Cumulative positive Qwen sequence lengths. - negative_prompt_cu_seqlens: Cumulative negative Qwen sequence - lengths. - callback_on_step_end_tensor_inputs: Tensor names exposed to the - step callback. + output_type: One of `pt`, `torch`, `np`, `numpy`, or `latent`. + prompt_embeds: Optional precomputed Qwen prompt embeddings. + pooled_prompt_embeds: Optional precomputed CLIP pooled prompt embeddings. + negative_prompt_embeds: Optional precomputed negative Qwen embeddings. + negative_pooled_prompt_embeds: Optional precomputed negative CLIP pooled embeddings. + callback_on_step_end_tensor_inputs: Tensor names exposed to the step callback. sample_audio: Whether to generate and decode audio. - visual_cond_scheme: Image-conditioning scheme when ``image`` is set. + image: Optional reference image or batch of reference images. + visual_cond_scheme: Image-conditioning scheme when `image` is set. Raises: ValueError: If an input combination is unsupported. @@ -918,16 +262,12 @@ def check_inputs( raise ValueError("provide either `prompt` or `prompt_embeds`") if prompt is not None and prompt_embeds is not None: raise ValueError("provide either `prompt` or `prompt_embeds`, not both") - if (prompt_embeds is None) != (prompt_cu_seqlens is None): - raise ValueError("`prompt_embeds` and `prompt_cu_seqlens` must be provided together") + if (prompt_embeds is None) != (pooled_prompt_embeds is None): + raise ValueError("`prompt_embeds` and `pooled_prompt_embeds` must be provided together") if negative_prompt is not None and negative_prompt_embeds is not None: raise ValueError("provide either `negative_prompt` or `negative_prompt_embeds`, not both") - if (negative_prompt_embeds is None) != (negative_prompt_cu_seqlens is None): - raise ValueError("`negative_prompt_embeds` and `negative_prompt_cu_seqlens` must be provided together") - if prompt_embeds is not None: - self._as_text_embeds(prompt_embeds) - if negative_prompt_embeds is not None: - self._as_text_embeds(negative_prompt_embeds) + if (negative_prompt_embeds is None) != (negative_pooled_prompt_embeds is None): + raise ValueError("`negative_prompt_embeds` and `negative_pooled_prompt_embeds` must be provided together") if num_frames < 1: raise ValueError("num_frames must be positive") @@ -972,27 +312,18 @@ def check_inputs( ): raise ValueError("tail_cond_first_frame requires transformer.visual_token_type_num_embeddings >= 2") - @staticmethod - def _as_text_embeds(value: Any) -> TextEmbeds: - if not isinstance(value, dict): - raise TypeError("prompt embeddings must be a mapping with text_embeds and pooled_embed") - missing = {"text_embeds", "pooled_embed"} - set(value) - if missing: - raise ValueError(f"prompt embeddings are missing: {sorted(missing)}") - return value - - def encode_prompt( + def _encode_single_prompt( self, text: str | list[str], max_sequence_length: int - ) -> tuple[TextEmbeds, Tensor, Tensor | None]: - """Encode text with Qwen2.5-VL and CLIP and build packed token metadata. + ) -> tuple[Tensor, Tensor, Tensor | None]: + """Encode one prompt (or batch of prompts) with Qwen2.5-VL and CLIP. Args: - text (`str` or `list[str]`): Prompt or prompts to encode. - max_sequence_length (`int`): Maximum number of Qwen text tokens. + text: Prompt or prompts to encode. + max_sequence_length: Maximum number of Qwen text tokens. Returns: - `tuple`: Text embeddings, cumulative sequence lengths, and an - optional Qwen attention mask. + `(prompt_embeds, pooled_prompt_embeds, attention_mask)`. `attention_mask` is `None` when every prompt + in the batch fills `max_sequence_length` (an all-`True` mask carries no extra information). """ texts = [text] if isinstance(text, str) else text full_texts = [_PROMPT_TEMPLATE.format(item) for item in texts] @@ -1012,31 +343,13 @@ def encode_prompt( return_dict=True, output_hidden_states=True, ) - embeds = qwen_output["hidden_states"][-1][:, _QWEN_CROP_START:] - attention = inputs["attention_mask"][:, _QWEN_CROP_START:].to(dtype=torch.bool) - batch_size = len(texts) - if self.text_token_padding: - qwen_embeds = embeds if batch_size > 1 else embeds[0] - qwen_attention = attention if batch_size > 1 else attention[0] - lengths = torch.full((batch_size,), embeds.shape[1], dtype=torch.int32, device=embeds.device) - cu_seqlens = torch.cat([torch.zeros(1, dtype=torch.int32, device=embeds.device), lengths.cumsum(0)]) - elif batch_size > 1: - max_tokens = int(attention.sum(dim=1).max().item()) - qwen_embeds = embeds[:, :max_tokens] - qwen_attention = attention[:, :max_tokens] - lengths = torch.full((batch_size,), max_tokens, dtype=torch.int32, device=embeds.device) - cu_seqlens = torch.cat([torch.zeros(1, dtype=torch.int32, device=embeds.device), lengths.cumsum(0)]) - else: - qwen_embeds = embeds[attention] - qwen_attention = None - lengths = attention.sum(dim=1).to(torch.int32) - cu_seqlens = torch.cat([torch.zeros(1, dtype=torch.int32, device=embeds.device), lengths.cumsum(0)]) - - if qwen_attention is not None and qwen_attention.all(): + prompt_embeds = qwen_output["hidden_states"][-1][:, _QWEN_CROP_START:] + attention_mask = inputs["attention_mask"][:, _QWEN_CROP_START:].to(dtype=torch.bool) + if attention_mask.all(): # An all-True mask carries no information beyond "attend to everything", which is what the # transformer already does when no mask is passed (see transformer_kandinsky6.py's # _normalize_attn_mask). Drop it so the dense mask isn't threaded through the denoise loop. - qwen_attention = None + attention_mask = None clip_device = next(self.text_encoder_2.parameters()).device clip_inputs = self.tokenizer_2( @@ -1047,14 +360,56 @@ def encode_prompt( padding="max_length", return_tensors="pt", ).to(clip_device) - pooled_embed = self.text_encoder_2(**clip_inputs)["pooler_output"] + pooled_prompt_embeds = self.text_encoder_2(**clip_inputs)["pooler_output"] + return prompt_embeds, pooled_prompt_embeds, attention_mask + + def encode_prompt( + self, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = True, + prompt_embeds: Tensor | None = None, + pooled_prompt_embeds: Tensor | None = None, + prompt_attention_mask: Tensor | None = None, + negative_prompt_embeds: Tensor | None = None, + negative_pooled_prompt_embeds: Tensor | None = None, + negative_prompt_attention_mask: Tensor | None = None, + max_sequence_length: int = 1024, + ) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None, Tensor | None, Tensor | None]: + """Encode the positive and (optionally) negative prompt. + + Args: + prompt: Prompt or prompts to encode. + negative_prompt: Negative prompt or prompts, used when `do_classifier_free_guidance=True` and + `negative_prompt_embeds` is not already provided. + do_classifier_free_guidance: Whether to also encode `negative_prompt`. + prompt_embeds: Optional precomputed Qwen prompt embeddings, skipping encoding of `prompt`. + pooled_prompt_embeds: Optional precomputed CLIP pooled prompt embeddings. + prompt_attention_mask: Optional Qwen attention mask paired with `prompt_embeds`. + negative_prompt_embeds: Optional precomputed negative Qwen embeddings. + negative_pooled_prompt_embeds: Optional precomputed negative CLIP pooled embeddings. + negative_prompt_attention_mask: Optional Qwen attention mask paired with `negative_prompt_embeds`. + max_sequence_length: Maximum Qwen prompt length after the template. + + Returns: + `(prompt_embeds, pooled_prompt_embeds, prompt_attention_mask, negative_prompt_embeds, + negative_pooled_prompt_embeds, negative_prompt_attention_mask)`. + """ + if prompt_embeds is None: + prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self._encode_single_prompt( + prompt, max_sequence_length + ) + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt_embeds, negative_pooled_prompt_embeds, negative_prompt_attention_mask = ( + self._encode_single_prompt(negative_prompt, max_sequence_length) + ) return ( - { - "text_embeds": qwen_embeds, - "pooled_embed": pooled_embed, - }, - cu_seqlens, - qwen_attention, + prompt_embeds, + pooled_prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_pooled_prompt_embeds, + negative_prompt_attention_mask, ) @torch.no_grad() @@ -1069,29 +424,19 @@ def expand_prompts( if isinstance(prompt, list): images = image if isinstance(image, (list, tuple)) else [image] * len(prompt) return [ - self.expand_prompts( - item, - image=item_image, - max_sequence_length=max_sequence_length, - mode=mode, - ) + self.expand_prompts(item, image=item_image, max_sequence_length=max_sequence_length, mode=mode) for item, item_image in zip(prompt, images, strict=True) ] mode = mode or ("i2va" if image is not None else "t2va") if mode not in ("t2va", "i2va"): raise ValueError(f"unknown prompt expansion mode={mode!r}") - image = _prepare_expansion_image(image) if image is not None else None + image = self._prepare_expansion_image(image) if image is not None else None instruction = (_I2VA_EXPANSION_INSTRUCTION if mode == "i2va" else _T2VA_EXPANSION_INSTRUCTION).format(prompt) content = [] if image is not None: content.append({"type": "image", "image": image}) content.append({"type": "text", "text": instruction}) - messages = [ - { - "role": "user", - "content": content, - } - ] + messages = [{"role": "user", "content": content}] text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) device = next(self.text_encoder.parameters()).device inputs = self.tokenizer( @@ -1106,266 +451,268 @@ def expand_prompts( trimmed = [output[qwen_crop_start:] for output in generated] return self.tokenizer.batch_decode(trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - def _encode_prompt( - self, - prompt: str | list[str], - negative_prompt: str | list[str], - *, - prompt_embeds: TextEmbeds | None, - negative_prompt_embeds: TextEmbeds | None, - prompt_cu_seqlens: Tensor | None, - negative_prompt_cu_seqlens: Tensor | None, - max_sequence_length: int, - ) -> tuple[TextEmbeds, Tensor, Tensor | None, TextEmbeds, Tensor, Tensor | None]: + @staticmethod + def _prepare_expansion_image(image: Any) -> Any: + if isinstance(image, (str, Path)): + from PIL import Image - if prompt_embeds is None: - prompt_embeds, prompt_cu_seqlens, prompt_attention_mask = self.encode_prompt(prompt, max_sequence_length) - elif prompt_cu_seqlens is None: - raise ValueError("prompt_cu_seqlens is required with prompt_embeds") - else: - prompt_attention_mask = None - prompt_embeds = self._as_text_embeds(prompt_embeds) + with Image.open(image) as pil_image: + return pil_image.convert("RGB").copy() + return image - if negative_prompt_embeds is None: - negative_prompt_embeds, negative_prompt_cu_seqlens, negative_attention_mask = self.encode_prompt( - negative_prompt, max_sequence_length - ) - elif negative_prompt_cu_seqlens is None: - raise ValueError("negative_prompt_cu_seqlens is required with negative_prompt_embeds") + @staticmethod + def _load_pil_rgb(image: str | object): + from PIL import Image + + if isinstance(image, str): + try: + pil_image = Image.open(image).convert("RGB") + pil_image.load() + except Exception as exc: + raise ValueError(f"Cannot decode i2va input image {image!r}: {exc}") from exc + elif isinstance(image, Image.Image): + pil_image = image.convert("RGB") else: - negative_attention_mask = None - negative_prompt_embeds = self._as_text_embeds(negative_prompt_embeds) - return ( - prompt_embeds, - prompt_cu_seqlens, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_cu_seqlens, - negative_attention_mask, - ) + raise TypeError(f"i2va image must be a path or PIL image, got {type(image).__name__}") + return pil_image - @staticmethod - def _seed_from_generator( - generator: torch.Generator | None, - device: torch.device, - ) -> int: - if generator is None: - return int(torch.randint(0, 2**31, (1,), device=device).item()) - generator_device = torch.device(getattr(generator, "device", "cpu")) - return int( - torch.randint( - 0, - 2**31, - (1,), - generator=generator, - device=generator_device, - ).item() - ) + @torch.no_grad() + def _encode_i2va_first_frame(self, image: str | object, device: torch.device, height: int, width: int) -> Tensor: + """Resize, center-crop, and VAE-encode one reference image into the packed K6 first-frame latent.""" + try: + import torchvision.transforms.functional as TF + except (ImportError, OSError) as exc: + raise RuntimeError( + "I2VA image processing requires torchvision. Install it with `pip install torchvision`." + ) from exc + + pil_image = self._load_pil_rgb(image) + tensor = TF.pil_to_tensor(pil_image).unsqueeze(0) + src_h, src_w = tensor.shape[-2:] + scale = min(src_h / height, src_w / width) + tensor = TF.resize(tensor, (int(src_h / scale), int(src_w / scale))) + cur_h, cur_w = tensor.shape[-2:] + tensor = TF.crop(tensor, (cur_h - height) // 2, (cur_w - width) // 2, height, width) - @staticmethod - def _coerce_video_latents( - latents: Tensor, - *, - latent_frames: int, + tensor = tensor.to(device=device, dtype=next(self.vae.parameters()).dtype) / 127.5 - 1.0 + tensor = tensor.transpose(0, 1).unsqueeze(0) + encoded = self.vae.encode(tensor) + posterior = encoded.latent_dist if hasattr(encoded, "latent_dist") else encoded[0] + latent = posterior.sample() + latent = latent.squeeze(0).permute(1, 2, 3, 0) + return latent * self.vae.config.scaling_factor + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int, + num_latent_frames: int, height: int, width: int, - channels: int, - batch_size: int, - device: torch.device, dtype: torch.dtype, - ) -> LatentBundle: - value = latents.to(device=device, dtype=dtype) - if value.ndim == 4: - value = value.unsqueeze(0) - if value.ndim != 5: - raise ValueError("video latents must have shape (T,H,W,C), (B,T,H,W,C), or (B,C,T,H,W)") - if value.shape[0] != batch_size: - raise ValueError(f"latents batch size must be {batch_size}, got {value.shape[0]}") - expected = (batch_size, latent_frames, height // 8, width // 8, channels) - bchtw = (batch_size, channels, latent_frames, height // 8, width // 8) - if tuple(value.shape) == bchtw: - value = value.permute(0, 2, 3, 4, 1) - elif tuple(value.shape) != expected: - raise ValueError("5-D latents must use BCHTW or BTHWC layout with the configured channel count") - cu_seqlens = latent_frames * torch.arange(batch_size + 1, dtype=torch.int32, device=device) - return LatentBundle( - video=value, - audio=None, - video_cu_seqlens=cu_seqlens, - audio_cu_seqlens=None, - ) - - @staticmethod - def _coerce_audio_latents( - bundle: LatentBundle, - latents: Tensor, - *, - audio_duration: int, - audio_dim: int, - batch_size: int, device: torch.device, - dtype: torch.dtype, - ) -> LatentBundle: - value = latents.to(device=device, dtype=dtype) - if value.ndim == 2: - value = value.unsqueeze(0) - if value.ndim != 3: - raise ValueError("audio latents must have shape (A,D) or (B,A,D)") - if value.shape[0] != batch_size: - raise ValueError(f"audio_latents batch size must be {batch_size}, got {value.shape[0]}") - expected = (batch_size, audio_duration, audio_dim) - if tuple(value.shape) != expected: - raise ValueError(f"audio latent shape mismatch: expected {expected}, got {tuple(value.shape)}") - cu_seqlens = audio_duration * torch.arange(batch_size + 1, dtype=torch.int32, device=device) - return LatentBundle( - video=bundle.video, - audio=value, - video_cu_seqlens=bundle.video_cu_seqlens, - audio_cu_seqlens=cu_seqlens, - ) + generator: torch.Generator | None, + latents: Tensor | None = None, + ) -> Tensor: + shape = (batch_size, num_latent_frames, height // 8, width // 8, num_channels_latents) + if latents is not None: + latents = latents.to(device=device, dtype=dtype) + if latents.ndim == 4: + latents = latents.unsqueeze(0) + bchtw = (batch_size, num_channels_latents, num_latent_frames, height // 8, width // 8) + if tuple(latents.shape) == bchtw: + latents = latents.permute(0, 2, 3, 4, 1) + if tuple(latents.shape) != shape: + raise ValueError( + f"`latents` must have shape {shape} (or the BCHTW equivalent), got {tuple(latents.shape)}" + ) + return latents + return randn_tensor(shape, generator=generator, device=device, dtype=dtype) - def _prepare_latents( + def prepare_audio_latents( self, - *, - latents: Tensor | None, - audio_latents: Tensor | None, - latent_frames: int, - height: int, - width: int, + batch_size: int, + num_channels_latents: int, + audio_duration: int, dtype: torch.dtype, device: torch.device, generator: torch.Generator | None, - sample_fps: float, - sample_audio: bool, - batch_size: int, - ) -> LatentBundle: - channels = int(getattr(self.transformer, "in_visual_dim", 16)) - seed = self._seed_from_generator(generator, device) - if latents is None: - bundle = prepare_video_latents( - bs=batch_size, - duration=latent_frames, - H_lat=height // 8, - W_lat=width // 8, - C=channels, - seed=seed, - device=device, - dtype=dtype, - ) - bundle.video = bundle.video.reshape(batch_size, latent_frames, height // 8, width // 8, channels) - else: - bundle = self._coerce_video_latents( - latents, - latent_frames=latent_frames, - height=height, - width=width, - channels=channels, - batch_size=batch_size, - device=device, - dtype=dtype, - ) + audio_latents: Tensor | None = None, + ) -> Tensor: + shape = (batch_size, audio_duration, num_channels_latents) + if audio_latents is not None: + audio_latents = audio_latents.to(device=device, dtype=dtype) + if audio_latents.ndim == 2: + audio_latents = audio_latents.unsqueeze(0) + if tuple(audio_latents.shape) != shape: + raise ValueError(f"`audio_latents` must have shape {shape}, got {tuple(audio_latents.shape)}") + return audio_latents + return randn_tensor(shape, generator=generator, device=device, dtype=dtype) - if not sample_audio: - if audio_latents is not None: - raise ValueError("audio_latents cannot be supplied when sample_audio=False") - return bundle - if self.audio_vae is None: - raise ValueError("sample_audio=True requires an audio_vae") + @staticmethod + def _audio_latent_duration( + num_video_latent_frames: int, + *, + fps: float = 24.0, + audio_fps: int = 44100, + downsample_factor: int = 1024, + ) -> int: + """Audio latent length matching K5 T2VA: ceil(sample_frames/fps * audio_fps / downsample).""" + sample_frames = (num_video_latent_frames - 1) * 4 + 1 + return int(math.ceil(sample_frames / fps * audio_fps / downsample_factor)) - audio_dim = int(getattr(self.transformer, "in_audio_dim", 20)) - downsample_factor = int(self.audio_vae.downsample_factor) - audio_duration = audio_latent_duration( - latent_frames, - fps=sample_fps, - audio_fps=self.audio_sample_rate, - downsample_factor=downsample_factor, - ) - if audio_latents is None: - bundle = prepare_audio_latents( - bundle, - audio_duration=audio_duration, - audio_dim=audio_dim, - seed=seed, - device=device, - dtype=dtype, + @staticmethod + def _append_i2va_tail_condition( + latents: Tensor, + first_frames: Tensor, + batch_size: int, + num_latent_frames: int, + ) -> tuple[Tensor, Tensor, Tensor]: + """Append one clean reference frame to batched `(B, T, H, W, C)` latents. + + Returns the extended latents, the per-frame token-type ids (`0` = generated, `1` = reference), and a + boolean mask selecting the generated frames. + """ + if latents.shape[0] != batch_size or latents.shape[1] != num_latent_frames: + raise ValueError( + "generated visual latent shape mismatch: expected " + f"({batch_size}, {num_latent_frames}, H, W, C), got {tuple(latents.shape)}" ) - bundle.audio = bundle.audio.reshape(batch_size, audio_duration, audio_dim) - return bundle - return self._coerce_audio_latents( - bundle, - audio_latents, - audio_duration=audio_duration, - audio_dim=audio_dim, - batch_size=batch_size, - device=device, - dtype=dtype, + _, _, height, width, channels = latents.shape + first_frames = first_frames.to(device=latents.device, dtype=latents.dtype) + expected = (batch_size, height, width, channels) + if tuple(first_frames.shape) != expected: + raise ValueError(f"first-frame latent shape mismatch: expected {expected}, got {tuple(first_frames.shape)}") + + latents = torch.cat([latents, first_frames[:, None]], dim=1) + token_types = torch.cat( + [ + torch.zeros((batch_size, num_latent_frames), dtype=torch.long, device=latents.device), + torch.ones((batch_size, 1), dtype=torch.long, device=latents.device), + ], + dim=1, ) + return latents, token_types, token_types == 0 @staticmethod - def _move_text_embeds( - embeds: TextEmbeds, - cu_seqlens: Tensor, - attention_mask: Tensor | None, - *, - device: torch.device, - dtype: torch.dtype, - ) -> tuple[TextEmbeds, Tensor, Tensor | None]: - moved = {key: value.to(device=device, dtype=dtype) for key, value in embeds.items()} - mask = attention_mask.to(device=device) if attention_mask is not None else None - return moved, cu_seqlens.to(device=device), mask + def _apply_visual_conditioning(latents: Tensor, first_frames: Tensor | None, visual_cond_scheme: str) -> Tensor: + """Append visual-conditioning channels for batched `(B, T, H, W, C)` video latents.""" + cond = torch.zeros_like(latents) + mask = torch.zeros((*latents.shape[:-1], 1), dtype=latents.dtype, device=latents.device) + if first_frames is None: + return torch.cat([latents, cond, mask], dim=-1) + + first_frames = first_frames.to(device=latents.device, dtype=latents.dtype) + tail_cond = visual_cond_scheme == "tail_cond_first_frame" + if visual_cond_scheme == "i2v": + latents[:, 0] = first_frames + elif tail_cond: + latents[:, -1] = first_frames + elif visual_cond_scheme == "pretrain": + cond[:, 0] = first_frames + else: + raise ValueError(f"unknown visual_cond_scheme={visual_cond_scheme!r}") + mask[:, -1 if tail_cond else 0] = 1 + return torch.cat([latents, cond, mask], dim=-1) @staticmethod - def _packed_bundle(bundle: LatentBundle) -> LatentBundle: - return LatentBundle( - video=bundle.video.flatten(0, 1) if bundle.video is not None else None, - audio=bundle.audio.flatten(0, 1) if bundle.audio is not None else None, - video_cu_seqlens=bundle.video_cu_seqlens, - audio_cu_seqlens=bundle.audio_cu_seqlens, - ) + def _compute_rope1d(rope: nn.Module, length: int, device: torch.device | None = None) -> Tensor: + """Build 1-D RoPE and move non-persistent tables for Diffusers offload.""" + if length < 1: + raise ValueError(f"rope length must be positive, got {length}") + if device is None: + device = next(rope.buffers()).device + else: + device = torch.device(device) + # Accelerate's model CPU-offload hook only moves the transformer when its forward starts. RoPE is + # materialized before that forward, so explicitly move its non-persistent lookup table first. + rope.to(device) + pos = torch.arange(length, device=device) + return rope(pos) @staticmethod - def _text_length(embeds: TextEmbeds, cu_seqlens: Tensor) -> int: - value = embeds["text_embeds"] - return int(value.shape[1] if value.ndim == 3 else cu_seqlens[-1].item()) + def _compute_visual_rope( + rope: nn.Module, + shape: tuple[int, int, int], + scale_factor: tuple[float, float, float], + device: torch.device | None = None, + ) -> Tensor: + """Build 3-D RoPE and move non-persistent tables for Diffusers offload.""" + t, h, w = (int(shape[0]), int(shape[1]), int(shape[2])) + if t < 1 or h < 1 or w < 1: + raise ValueError(f"visual rope shape must be positive, got {shape}") + if device is None: + device = next(rope.buffers()).device + else: + device = torch.device(device) + rope.to(device) + pos = [torch.arange(t, device=device), torch.arange(h, device=device), torch.arange(w, device=device)] + return rope((t, h, w), pos, tuple(float(value) for value in scale_factor)) - def _text_ropes( + def _prepare_text_ropes( self, text_length: int, negative_text_length: int, audio_length: int | None, device: torch.device, - ) -> tuple[Tensor | list[Tensor], Tensor | list[Tensor], Tensor | None]: - if getattr(self.transformer, "is_multimodal", False): - text_rope = [ - compute_rope1d(self.transformer.video_text_rope_embeddings, text_length, device=device), - compute_rope1d(self.transformer.audio_text_rope_embeddings, text_length, device=device), - ] - negative_text_rope = [ - compute_rope1d( - self.transformer.video_text_rope_embeddings, - negative_text_length, - device=device, - ), - compute_rope1d( - self.transformer.audio_text_rope_embeddings, - negative_text_length, - device=device, - ), - ] - else: - text_rope = compute_rope1d(self.transformer.text_rope_embeddings, text_length, device=device) - negative_text_rope = compute_rope1d( - self.transformer.text_rope_embeddings, - negative_text_length, - device=device, - ) + ) -> tuple[list[Tensor], list[Tensor], Tensor | None]: + """Build per-modality text RoPE for the cond/uncond branches, plus the audio RoPE.""" + text_rope = [ + self._compute_rope1d(self.transformer.video_text_rope_embeddings, text_length, device=device), + self._compute_rope1d(self.transformer.audio_text_rope_embeddings, text_length, device=device), + ] + negative_text_rope = [ + self._compute_rope1d(self.transformer.video_text_rope_embeddings, negative_text_length, device=device), + self._compute_rope1d(self.transformer.audio_text_rope_embeddings, negative_text_length, device=device), + ] audio_rope = ( - compute_rope1d(self.transformer.audio_rope_embeddings, audio_length, device=device) + self._compute_rope1d(self.transformer.audio_rope_embeddings, audio_length, device=device) if audio_length is not None else None ) return text_rope, negative_text_rope, audio_rope + @staticmethod + def _apply_guidance(cond: Tensor, uncond: Tensor, guidance_scale: float) -> Tensor: + return uncond + guidance_scale * (cond - uncond) + + @torch.no_grad() + def _postprocess_video(self, latents: Tensor) -> Tensor: + """Decode video latents `(B, T, H, W, C)` into `(B, 3, T, H, W)` `uint8` frames in `[0, 255]`.""" + frames = (latents / self.vae.config.scaling_factor).permute(0, 4, 1, 2, 3) + # Hunyuan VAE loads as fp16; DiT latents are bf16 — match weight dtype. + vae_dtype = next(self.vae.parameters()).dtype + frames = self.vae.decode(frames.to(dtype=vae_dtype)).sample + return ((frames.clamp(-1.0, 1.0) + 1.0) * 127.5).to(torch.uint8) + + @torch.no_grad() + def _postprocess_audio(self, audio_latents: Tensor | None) -> list[np.ndarray] | None: + """Decode audio latents `(B, A, D)` into a list of int16 `(samples,)` waveforms, or `None` in T2V mode.""" + if audio_latents is None: + return None + audio = audio_latents / getattr(self.audio_vae, "scaling_factor", 1.0) + audio = audio + getattr(self.audio_vae, "mean_value", 0.0) + decoded = self.audio_vae.wrapped_decode(audio.transpose(1, 2)) + if decoded.ndim == 1: + decoded = decoded.unsqueeze(0) + elif decoded.ndim == 3 and decoded.shape[1] == 1: + decoded = decoded[:, 0] + return [ + (np.clip(waveform.cpu().float().numpy(), -1.0, 1.0) * 32767).astype(np.int16) for waveform in decoded + ] + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( @@ -1383,11 +730,10 @@ def __call__( generator: torch.Generator | None = None, latents: Tensor | None = None, audio_latents: Tensor | None = None, - prompt_embeds: TextEmbeds | None = None, - negative_prompt_embeds: TextEmbeds | None = None, - cu_seqlens: Tensor | tuple[Tensor, Tensor] | None = None, - prompt_cu_seqlens: Tensor | None = None, - negative_prompt_cu_seqlens: Tensor | None = None, + prompt_embeds: Tensor | None = None, + pooled_prompt_embeds: Tensor | None = None, + negative_prompt_embeds: Tensor | None = None, + negative_pooled_prompt_embeds: Tensor | None = None, sample_audio: bool = True, expand_prompts: bool = False, output_type: str = "pt", @@ -1405,52 +751,34 @@ def __call__( height: Output video height in pixels; must be divisible by 8. width: Output video width in pixels; must be divisible by 8. num_frames: Number of decoded video frames. - sample_fps: Output video frame rate. Kandinsky 6 is trained for - 24.0 fps; other values may affect audio/video alignment. + sample_fps: Output video frame rate. Kandinsky 6 is trained for 24.0 fps; other values may affect + audio/video alignment. num_inference_steps: Number of denoising steps. max_sequence_length: Maximum Qwen prompt length after the template. guidance_scale: Classifier-free guidance weight. generator: Random generator used for latent initialization. latents: Optional precomputed video latents. audio_latents: Optional precomputed audio latents. - prompt_embeds: Optional precomputed positive embeddings. - negative_prompt_embeds: Optional precomputed negative embeddings. - cu_seqlens: Optional positive/negative cumulative sequence lengths. - prompt_cu_seqlens: Optional cumulative sequence lengths for - precomputed positive prompt embeddings. - negative_prompt_cu_seqlens: Optional cumulative sequence lengths - for precomputed negative prompt embeddings. + prompt_embeds: Optional precomputed Qwen prompt embeddings. + pooled_prompt_embeds: Optional precomputed CLIP pooled prompt embeddings. + negative_prompt_embeds: Optional precomputed negative Qwen embeddings. + negative_pooled_prompt_embeds: Optional precomputed negative CLIP pooled embeddings. sample_audio: Whether to generate synchronized audio. expand_prompts: Whether to use the built-in Qwen video+audio prompt expander before encoding. - output_type: ``pt``/``torch`` for tensors, ``np``/``numpy`` for - NumPy arrays, or ``latent`` for undecoded video latents. - return_dict: Whether to return ``Kandinsky6TI2VAPipelineOutput``. + output_type: `pt`/`torch` for tensors, `np`/`numpy` for NumPy arrays, or `latent` for undecoded video + latents. + return_dict: Whether to return `Kandinsky6TI2VAPipelineOutput`. callback_on_step_end: Optional callback invoked after each step. callback_on_step_end_tensor_inputs: Names passed to the callback. - visual_cond_scheme: Optional image-conditioning scheme. Defaults to - ``pretrain`` without ``image`` and ``tail_cond_first_frame`` - with ``image``. + visual_cond_scheme: Optional image-conditioning scheme. Defaults to `pretrain` without `image` and + `tail_cond_first_frame` with `image`. Examples: Returns: - ``Kandinsky6TI2VAPipelineOutput`` containing video frames and int16 NumPy - audio waveforms, or a tuple when ``return_dict=False``. + `Kandinsky6TI2VAPipelineOutput` containing video frames and int16 NumPy audio waveforms, or a tuple + when `return_dict=False`. """ - if cu_seqlens is not None: - if isinstance(cu_seqlens, tuple): - if len(cu_seqlens) != 2: - raise ValueError("cu_seqlens tuple must contain positive and negative lengths") - if prompt_cu_seqlens is None: - prompt_cu_seqlens = cu_seqlens[0] - if negative_prompt_cu_seqlens is None: - negative_prompt_cu_seqlens = cu_seqlens[1] - else: - if prompt_cu_seqlens is None: - prompt_cu_seqlens = cu_seqlens - if negative_prompt_cu_seqlens is None: - negative_prompt_cu_seqlens = cu_seqlens - if sample_fps != 24.0: warnings.warn( f"Kandinsky 6 was trained for 24.0 fps; received sample_fps={sample_fps}.", @@ -1458,22 +786,18 @@ def __call__( stacklevel=2, ) + # 1. Define call parameters batch_size = self._batch_size( - prompt, - negative_prompt, - image, - latents, - audio_latents, - prompt_embeds, - negative_prompt_embeds, + prompt, negative_prompt, image, latents, audio_latents, prompt_embeds, negative_prompt_embeds ) prompt_batch = self._as_batch(prompt, batch_size, "prompt") negative_prompt_batch = self._as_batch(negative_prompt, batch_size, "negative_prompt") image_batch = self._as_batch(image, batch_size, "image") if visual_cond_scheme is None: visual_cond_scheme = "tail_cond_first_frame" if image is not None else "pretrain" + + # 2. Check inputs. Raise error if not correct self.check_inputs( - image=image_batch, prompt=prompt_batch, negative_prompt=negative_prompt_batch, height=height, @@ -1483,19 +807,32 @@ def __call__( max_sequence_length=max_sequence_length, output_type=output_type, prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, - prompt_cu_seqlens=prompt_cu_seqlens, - negative_prompt_cu_seqlens=negative_prompt_cu_seqlens, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, sample_audio=sample_audio, + image=image_batch, visual_cond_scheme=visual_cond_scheme, ) + callback_on_step_end_tensor_inputs = callback_on_step_end_tensor_inputs or self._callback_tensor_inputs + + self._current_timestep = None + self._interrupt = False - callback_inputs = ( - self._callback_tensor_inputs - if callback_on_step_end_tensor_inputs is None - else callback_on_step_end_tensor_inputs + device = self._execution_device + dtype = ( + self.transformer.dtype if isinstance(getattr(self.transformer, "dtype", None), torch.dtype) else torch.bfloat16 ) + + # 3. Encode the reference image, if any + first_frames = None + if image_batch is not None: + first_frames = torch.cat( + [self._encode_i2va_first_frame(item, device, height, width) for item in image_batch], dim=0 + ) + + # 4. Encode input prompt if expand_prompts and prompt_batch is not None and prompt_embeds is None: prompt_batch = self.expand_prompts( prompt_batch, @@ -1503,175 +840,195 @@ def __call__( max_sequence_length=max_sequence_length, mode="i2va" if image_batch is not None else "t2va", ) - patch_size = tuple(int(value) for value in getattr(self.transformer, "patch_size", (1, 2, 2))) - - device = getattr(self, "_execution_device", self.transformer.device) - dtype = getattr(self.transformer, "dtype", torch.bfloat16) - if not isinstance(dtype, torch.dtype): - dtype = torch.bfloat16 - first_frames = None - if image_batch is not None: - encoded_frames = [ - encode_i2va_first_frame(item, self.vae, device, height=height, width=width) for item in image_batch - ] - first_frames = torch.cat([item[0] for item in encoded_frames], dim=0) - encoded_height, encoded_width = encoded_frames[0][1:] - if any(item[1:] != (encoded_height, encoded_width) for item in encoded_frames): - raise ValueError("all reference images must produce the same resolution") - height, width = encoded_height, encoded_width negative_prompt_batch = negative_prompt_batch or [self._DEFAULT_NEGATIVE_PROMPT] * batch_size + ( - positive, - positive_cu, - positive_mask, - negative, - negative_cu, - negative_mask, - ) = self._encode_prompt( - prompt_batch or [""] * batch_size, - negative_prompt_batch, + prompt_embeds, + pooled_prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_pooled_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt=prompt_batch or [""] * batch_size, + negative_prompt=negative_prompt_batch, + do_classifier_free_guidance=True, prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, - prompt_cu_seqlens=prompt_cu_seqlens, - negative_prompt_cu_seqlens=negative_prompt_cu_seqlens, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, max_sequence_length=max_sequence_length, ) - if positive_cu is None or negative_cu is None: - raise ValueError("positive and negative prompt cumulative lengths are required") - positive, positive_cu, positive_mask = self._move_text_embeds( - positive, - positive_cu, - positive_mask, - device=device, - dtype=dtype, - ) - negative, negative_cu, negative_mask = self._move_text_embeds( - negative, - negative_cu, - negative_mask, - device=device, - dtype=dtype, + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) + pooled_prompt_embeds = pooled_prompt_embeds.to(device=device, dtype=dtype) + negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.to(device=device, dtype=dtype) + if prompt_attention_mask is not None: + prompt_attention_mask = prompt_attention_mask.to(device=device) + if negative_prompt_attention_mask is not None: + negative_prompt_attention_mask = negative_prompt_attention_mask.to(device=device) + + # 5. Prepare latent variables + patch_size = tuple(int(value) for value in self.transformer.patch_size) + num_latent_frames = (num_frames - 1) // 4 + 1 + num_channels_latents = int(self.transformer.in_visual_dim) + latents = self.prepare_latents( + batch_size, num_channels_latents, num_latent_frames, height, width, dtype, device, generator, latents ) + if sample_audio: + audio_duration = self._audio_latent_duration( + num_latent_frames, + fps=sample_fps, + audio_fps=self.audio_sample_rate, + downsample_factor=self.audio_downsample_factor, + ) + audio_latents = self.prepare_audio_latents( + batch_size, + int(self.transformer.in_audio_dim), + audio_duration, + dtype, + device, + generator, + audio_latents, + ) + elif audio_latents is not None: + raise ValueError("`audio_latents` cannot be provided when `sample_audio=False`") + else: + audio_latents = None - latent_frames = (num_frames - 1) // 4 + 1 - bundle = self._prepare_latents( - latents=latents, - audio_latents=audio_latents, - latent_frames=latent_frames, - height=height, - width=width, - dtype=dtype, - device=device, - generator=generator, - sample_fps=sample_fps, - sample_audio=sample_audio, - batch_size=batch_size, - ) visual_token_type_ids = None generated_visual_mask = None - if image is not None and visual_cond_scheme == "tail_cond_first_frame": - if bundle.video is None or first_frames is None: - raise ValueError("image conditioning requires video latents") - video, visual_token_type_ids, generated_visual_mask = append_i2va_tail_condition( - bundle.video, - first_frames, - batch_size=batch_size, - video_duration=latent_frames, + if image_batch is not None and visual_cond_scheme == "tail_cond_first_frame": + latents, visual_token_type_ids, generated_visual_mask = self._append_i2va_tail_condition( + latents, first_frames, batch_size, num_latent_frames ) - bundle = LatentBundle( - video=video, - audio=bundle.audio, - video_cu_seqlens=(latent_frames + 1) * torch.arange(batch_size + 1, dtype=torch.int32, device=device), - audio_cu_seqlens=bundle.audio_cu_seqlens, - ) - audio_length = int(bundle.audio.shape[1]) if bundle.audio is not None else None + + # 6. Prepare rotary position embeddings visual_shape = ( - latent_frames // patch_size[0], + num_latent_frames // patch_size[0], (height // 8) // patch_size[1], (width // 8) // patch_size[2], ) if visual_shape[0] < 1: raise ValueError(f"invalid visual latent shape {visual_shape}") - visual_rope = compute_visual_rope( - self.transformer.visual_rope_embeddings, - visual_shape, - self.scale_factor, - device=device, + visual_rope = self._compute_visual_rope( + self.transformer.visual_rope_embeddings, visual_shape, self.scale_factor, device=device ) if generated_visual_mask is not None: + # The appended reference frame reuses the first frame's rotary position rather than a real T+1 one. visual_rope = torch.cat([visual_rope, visual_rope[:1]], dim=0) - text_rope, negative_text_rope, audio_rope = self._text_ropes( - text_length=self._text_length(positive, positive_cu), - negative_text_length=self._text_length(negative, negative_cu), - audio_length=audio_length, + text_rope, negative_text_rope, audio_rope = self._prepare_text_ropes( + text_length=prompt_embeds.shape[1], + negative_text_length=negative_prompt_embeds.shape[1], + audio_length=audio_latents.shape[1] if audio_latents is not None else None, device=device, ) + + # 7. Denoising loop + is_piflow = bool(getattr(self.scheduler, "is_piflow", False)) + if is_piflow and guidance_scale != 1.0: + raise ValueError("PiflowScheduler requires guidance_scale=1.0") + do_classifier_free_guidance = abs(guidance_scale - 1.0) > 1e-6 + tail_cond = visual_cond_scheme == "tail_cond_first_frame" + + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + # A second scheduler instance keeps audio's step index independent, since the scheduler advances its own + # internal step index on every `.step()` call. + audio_scheduler = copy.deepcopy(self.scheduler) if audio_latents is not None else None + if audio_scheduler is not None: + audio_scheduler.set_timesteps(num_inference_steps, device=device) + self._num_timesteps = len(timesteps) + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + self._current_timestep = t + timestep = t.unsqueeze(0).expand(batch_size) + + video_model_input = ( + self._apply_visual_conditioning(latents, first_frames, visual_cond_scheme) + if self.transformer.visual_cond + else latents + ) - def _step_callback(step_index: int, timestep: Tensor, current: LatentBundle) -> LatentBundle: - if callback_on_step_end is not None: - callback_result = callback_on_step_end( - self, - step_index, - timestep, - {name: current.video for name in callback_inputs}, + with self.transformer.cache_context("cond"): + model_pred = self.transformer( + x_video=video_model_input, + x_audio=audio_latents, + text_embed=prompt_embeds, + pooled_text_embed=pooled_prompt_embeds, + time=timestep, + visual_rope=visual_rope, + audio_rope=audio_rope, + text_rope=text_rope, + attention_mask=prompt_attention_mask, + visual_token_type_ids=visual_token_type_ids, ) - if callback_result is None: - callback_result = {} - if not isinstance(callback_result, dict): - raise TypeError("callback_on_step_end must return a dict or None") - if callback_result.get("latents") is not None: - current.video = callback_result["latents"] - return current - - result = denoise_loop( - bundle=bundle, - dit=self.transformer, - text_embeds=positive, - null_text_embeds=negative, - visual_rope=visual_rope, - audio_rope=audio_rope, - text_rope=text_rope, - null_text_rope=negative_text_rope, - num_steps=num_inference_steps, - guidance_weight=guidance_scale, - first_frames=first_frames, - visual_cond_scheme=visual_cond_scheme, - sample_video=True, - sample_audio=sample_audio, - attention_mask=positive_mask, - null_attention_mask=negative_mask, - visual_token_type_ids=visual_token_type_ids, - scale_factor=self.scale_factor, - scheduler=self.scheduler, - step_callback=_step_callback if callback_on_step_end is not None else None, - progress_bar=progress_bar, - ) - if generated_visual_mask is not None and result.video is not None: - generated_video = result.video[:, generated_visual_mask[0]] - result = LatentBundle( - video=generated_video, - audio=result.audio, - video_cu_seqlens=latent_frames * torch.arange(batch_size + 1, dtype=torch.int32, device=device), - audio_cu_seqlens=result.audio_cu_seqlens, - ) + if do_classifier_free_guidance: + with self.transformer.cache_context("uncond"): + model_pred_uncond = self.transformer( + x_video=video_model_input, + x_audio=audio_latents, + text_embed=negative_prompt_embeds, + pooled_text_embed=negative_pooled_prompt_embeds, + time=timestep, + visual_rope=visual_rope, + audio_rope=audio_rope, + text_rope=negative_text_rope, + attention_mask=negative_prompt_attention_mask, + visual_token_type_ids=visual_token_type_ids, + ) + if audio_latents is not None: + video_pred = self._apply_guidance(model_pred[0], model_pred_uncond[0], guidance_scale) + audio_pred = self._apply_guidance(model_pred[1], model_pred_uncond[1], guidance_scale) + else: + video_pred = self._apply_guidance(model_pred, model_pred_uncond, guidance_scale) + elif audio_latents is not None: + video_pred, audio_pred = model_pred + else: + video_pred = model_pred + + latents = self.scheduler.step(video_pred, t, latents, return_dict=False)[0] + if audio_latents is not None: + audio_latents = audio_scheduler.step(audio_pred, t, audio_latents, return_dict=False)[0] + if tail_cond and first_frames is not None: + latents[:, -1] = first_frames.to(device=device, dtype=latents.dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + latents = callback_outputs.pop("latents", latents) + + progress_bar.update() + + self._current_timestep = None + + # 8. Restore the injected reference frame and drop it from the output + if first_frames is not None: + ff = first_frames.to(device=device, dtype=latents.dtype) + if visual_cond_scheme == "i2v": + latents[:, 0] = ff + elif tail_cond: + latents[:, -1] = ff + if generated_visual_mask is not None: + latents = latents[:, generated_visual_mask[0]] + # 9. Decode outputs if output_type == "latent": - if result.video is None: - raise ValueError("denoising returned no video latent") - frames: Tensor | np.ndarray = result.video.permute(0, 4, 1, 2, 3) + frames: Tensor | np.ndarray = latents.permute(0, 4, 1, 2, 3) else: - packed_result = self._packed_bundle(result) - decoded_frames = postprocess_video(packed_result, self.vae, bs=batch_size) - # check_inputs already restricted output_type to pt/torch/np/numpy/latent, and "latent" is - # handled above, so only the tensor-vs-numpy choice is left here. - frames = ( - decoded_frames - if output_type in ("pt", "torch") - else decoded_frames.permute(0, 2, 3, 4, 1).cpu().numpy() - ) - audio = postprocess_audio(self._packed_bundle(result), self.audio_vae) if sample_audio else None + frames = self._postprocess_video(latents) + if output_type not in ("pt", "torch"): + frames = frames.permute(0, 2, 3, 4, 1).cpu().numpy() + audio = self._postprocess_audio(audio_latents) if sample_audio else None + + self.maybe_free_model_hooks() + if not return_dict: return frames, audio return Kandinsky6TI2VAPipelineOutput(frames=frames, audio=audio) diff --git a/src/diffusers/schedulers/scheduling_piflow.py b/src/diffusers/schedulers/scheduling_piflow.py index db367649e427..59e6be1715dd 100644 --- a/src/diffusers/schedulers/scheduling_piflow.py +++ b/src/diffusers/schedulers/scheduling_piflow.py @@ -14,20 +14,32 @@ """Diffusers scheduler for distilled Kandinsky 6 PiFlow checkpoints.""" -from __future__ import annotations +from dataclasses import dataclass import torch -from ..configuration_utils import register_to_config -from .scheduling_flow_match_euler_discrete import ( - FlowMatchEulerDiscreteScheduler, - FlowMatchEulerDiscreteSchedulerOutput, -) + +from ..configuration_utils import ConfigMixin, register_to_config +from ..utils import BaseOutput +from .scheduling_utils import SchedulerMixin + + +@dataclass +class PiflowSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.Tensor`): + Computed sample at the next PiFlow grid point. Should be used as the next denoising input. + """ + + prev_sample: torch.Tensor class DXPolicy: """Network-free DX policy over one flow-matching segment.""" - def __init__( # noqa: PLR0913 + def __init__( self, denoising_output: torch.Tensor, x_t_src: torch.Tensor, @@ -56,7 +68,7 @@ def __init__( # noqa: PLR0913 @staticmethod def _interpolate(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: n = x.size(1) - if n < 2: # noqa: PLR2004 + if n < 2: 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) @@ -89,7 +101,7 @@ def shift_timesteps(t: torch.Tensor, shift: float) -> torch.Tensor: return shift * t / (1 + (shift - 1) * t) -def policy_rollout_fm( # noqa: PLR0913 +def policy_rollout_fm( x_t_start: torch.Tensor, sigma_t_start: torch.Tensor, raw_t_start: torch.Tensor, @@ -126,14 +138,16 @@ def policy_rollout_fm( # noqa: PLR0913 return x_t, sigma_t, sigma_t.flatten() * 1_000 -class PiflowScheduler(FlowMatchEulerDiscreteScheduler): +class PiflowScheduler(SchedulerMixin, ConfigMixin): """Few-step PiFlow scheduler for widened-output diffusion transformers. - PiFlow evaluates the denoising model at a small number of grid points and - integrates a network-free policy between those evaluations. The scheduler - is intended for distilled Kandinsky 6 checkpoints, including the main - video/audio model and the video super-resolution model. Their model output - contains ``n_grid`` predictions per sample channel. + PiFlow evaluates the denoising model at a small number of grid points and integrates a network-free policy + between those evaluations. The scheduler is intended for distilled Kandinsky 6 checkpoints, including the main + video/audio model and the video super-resolution model. Their model output contains `n_grid` predictions per + sample channel. + + This scheduler inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the + generic methods implemented for all schedulers (loading, saving, etc.). Args: num_train_timesteps (`int`, *optional*, defaults to 1000): Number of @@ -149,6 +163,8 @@ class PiflowScheduler(FlowMatchEulerDiscreteScheduler): integration substeps per raw-timestep unit. """ + _compatibles = [] + order = 1 is_piflow = True @register_to_config @@ -170,16 +186,41 @@ def __init__( raise ValueError("PiflowScheduler requires 0 < final_step_size_scale <= 1") if num_policy_substeps < 1: raise ValueError("PiflowScheduler requires num_policy_substeps >= 1") - super().__init__( - num_train_timesteps=num_train_timesteps, - shift=shift, - ) + + self._step_index = None + self._begin_index = None + self.timesteps = torch.empty(0) + self.sigmas = torch.empty(0) self.n_grid = int(n_grid) self.eps = float(eps) 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) + @property + def step_index(self): + """The index counter for the current timestep. It increases by 1 after each scheduler step.""" + return self._step_index + + @property + def begin_index(self): + """The index for the first timestep. It should be set from the pipeline with `set_begin_index`.""" + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0) -> None: + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`, defaults to `0`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def __len__(self) -> int: + return self.config.num_train_timesteps + def set_timesteps( self, num_inference_steps: int | None = None, @@ -287,7 +328,7 @@ def step( timestep: float | torch.FloatTensor, sample: torch.FloatTensor, return_dict: bool = True, - ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple: + ) -> PiflowSchedulerOutput | tuple: """Advance one step by integrating the PiFlow policy. Args: @@ -296,10 +337,10 @@ def step( timestep (`float` or `torch.FloatTensor`): Current scheduler timestep. sample (`torch.FloatTensor`): Current noisy sample. return_dict (`bool`, *optional*, defaults to True): Whether to return - a [`FlowMatchEulerDiscreteSchedulerOutput`]. + a [`PiflowSchedulerOutput`]. Returns: - [`FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`: Updated sample. + [`PiflowSchedulerOutput`] or `tuple`: Updated sample. """ if isinstance(timestep, int) or isinstance(timestep, (torch.IntTensor, torch.LongTensor)): raise ValueError( @@ -314,5 +355,5 @@ def step( updated = updated.to(dtype=sample.dtype) self._step_index += 1 if return_dict: - return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=updated) + return PiflowSchedulerOutput(prev_sample=updated) return (updated,) From 01bcb6482075d7fbe8b2a547b2b432c1da943191 Mon Sep 17 00:00:00 2001 From: leffff Date: Sun, 27 Sep 2026 23:02:41 +0300 Subject: [PATCH 14/62] replace TimestepEmbedding and fix batching --- .../transformers/transformer_kandinsky6.py | 74 +++++-------------- 1 file changed, 17 insertions(+), 57 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 897937ee9610..4d5ee0af0dde 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -33,6 +33,7 @@ dispatch_attention_fn, ) from ..cache_utils import CacheMixin +from ..embeddings import get_timestep_embedding from ..modeling_outputs import Transformer2DModelOutput from ..modeling_utils import ModelMixin from torch import Tensor, nn @@ -483,12 +484,6 @@ def __call__( if rotary_emb_kv is not None: key = apply_rotary(key, rotary_emb_kv).to(dtype=key.dtype) - query_was_batched = query.dim() == 4 - if not query_was_batched: - query, key, value = query.unsqueeze(0), key.unsqueeze(0), value.unsqueeze(0) - elif key.dim() == 3: - key, value = key.unsqueeze(0), value.unsqueeze(0) - if sparse_params is not None: q = query.transpose(1, 2).contiguous() k = key.transpose(1, 2).contiguous() @@ -521,8 +516,6 @@ def __call__( parallel_config=self._parallel_config, ) - if not query_was_batched: - output = output[0] return attn.out_layer(output.flatten(-2, -1)) @@ -533,15 +526,16 @@ def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): super().__init__() if model_dim % 2: raise ValueError("model_dim must be even") - self.register_buffer("freqs", get_freqs(model_dim // 2, max_period), persistent=False) + self.model_dim = model_dim + self.max_period = max_period self.in_layer = nn.Linear(model_dim, time_dim) self.activation = nn.SiLU() self.out_layer = nn.Linear(time_dim, time_dim) def forward(self, time: Tensor) -> Tensor: - freqs = self.freqs.to(device=time.device, dtype=torch.float32) - args = torch.outer(time.float(), freqs) - embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + embed = get_timestep_embedding( + time, self.model_dim, flip_sin_to_cos=True, downscale_freq_shift=0, max_period=self.max_period + ) h = functional.linear(embed, self.in_layer.weight.float(), self.in_layer.bias.float()) return functional.linear(self.activation(h), self.out_layer.weight.float(), self.out_layer.bias.float()) @@ -559,7 +553,7 @@ def forward(self, x: Tensor) -> Tensor: class Kandinsky6VisualEmbeddings(nn.Module): - """Patch projection for ``[T,H,W,C]`` or ``[B,T,H,W,C]`` visual tokens.""" + """Patch projection for ``[B,T,H,W,C]`` visual tokens.""" def __init__(self, visual_dim: int, model_dim: int, patch_size: tuple[int, int, int]): super().__init__() @@ -567,35 +561,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:] + batch, 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(batch, 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) + ) return self.in_layer(x) @@ -708,26 +680,14 @@ def forward(self, visual_embed: Tensor, time_embed: Tensor) -> Tensor: ) x = self.out_layer(x) - # batch_dim variant - if x.ndim == 5: - batch, duration, height, width = x.shape[:4] - p_t, p_h, p_w = self.patch_size - return ( - x.view(batch, duration, height, width, -1, p_t, p_h, p_w) - .permute(0, 1, 5, 2, 6, 3, 7, 4) - .flatten(1, 2) - .flatten(2, 3) - .flatten(3, 4) - ) - - duration, height, width, _ = x.shape + batch, duration, height, width = x.shape[:4] p_t, p_h, p_w = self.patch_size return ( - x.view(duration, height, width, -1, p_t, p_h, p_w) - .permute(0, 4, 1, 5, 2, 6, 3) - .flatten(0, 1) + x.view(batch, duration, height, width, -1, p_t, p_h, p_w) + .permute(0, 1, 5, 2, 6, 3, 7, 4) .flatten(1, 2) .flatten(2, 3) + .flatten(3, 4) ) From 8a57f5718aa22504f7617dc08bef37a4d6c1ce70 Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 28 Sep 2026 11:19:21 +0300 Subject: [PATCH 15/62] refactor ti2va dit and pipeline --- .../transformers/transformer_kandinsky6.py | 558 ++++++------------ .../kandinsky6/pipeline_kandinsky6_ti2va.py | 282 +++++---- .../test_models_transformer_kandinsky6.py | 175 ++++++ tests/pipelines/kandinsky6/__init__.py | 0 .../test_pipeline_kandinsky6_ti2va.py | 240 ++++++++ 5 files changed, 759 insertions(+), 496 deletions(-) create mode 100644 tests/models/transformers/test_models_transformer_kandinsky6.py create mode 100644 tests/pipelines/kandinsky6/__init__.py create mode 100644 tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 4d5ee0af0dde..242fdfd692aa 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -21,6 +21,9 @@ import torch import torch.nn.functional as functional +from torch import Tensor, nn +from torch.nn.attention.flex_attention import BlockMask + from ...configuration_utils import ConfigMixin, register_to_config from ...hooks import MagCacheConfig from ...hooks.hooks import HookRegistry, ModelHook, StateManager @@ -36,8 +39,6 @@ from ..embeddings import get_timestep_embedding from ..modeling_outputs import Transformer2DModelOutput from ..modeling_utils import ModelMixin -from torch import Tensor, nn -from torch.nn.attention.flex_attention import BlockMask def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: @@ -744,7 +745,12 @@ def forward(self, x: Tensor, time_embed: Tensor, rope: Tensor, attn_mask: Tensor class Kandinsky6TransformerDecoderBlock(nn.Module): - """Visual self-attention, text cross-attention, and feed-forward block.""" + """Visual self-attention, text cross-attention, and feed-forward submodules. + + `Kandinsky6FusedTransformerDecoderBlock` uses this class only as a named submodule container + (`self_attention`/`cross_attention`/`feed_forward` and their norms/modulation), calling those + submodules directly rather than this class's own `forward`. + """ def __init__( self, @@ -768,35 +774,6 @@ def __init__( self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.feed_forward = Kandinsky6FeedForward(model_dim, ff_dim) - def forward( - self, vis: Tensor, text: Tensor, time_embed: Tensor, rope: Tensor, sparse_params: dict | None, attn_mask=None - ) -> Tensor: - sa_params, ca_params, ff_params = torch.chunk(self.visual_modulation(time_embed), 3, dim=-1) - shift, scale, gate = torch.chunk(sa_params, 3, dim=-1) - vis = apply_gate_sum( - vis, - self.self_attention( - apply_scale_shift_norm(self.self_attention_norm, vis, scale, shift), - rotary_emb=rope, - sparse_params=sparse_params, - ), - gate, - ) - shift, scale, gate = torch.chunk(ca_params, 3, dim=-1) - vis = apply_gate_sum( - vis, - self.cross_attention( - apply_scale_shift_norm(self.cross_attention_norm, vis, scale, shift), - encoder_hidden_states=text, - attn_mask=attn_mask, - ), - gate, - ) - shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) - return apply_gate_sum( - vis, self.feed_forward(apply_scale_shift_norm(self.feed_forward_norm, vis, scale, shift)), gate - ) - class Kandinsky6FusedTransformerDecoderBlock(nn.Module): """Fused K6 video/audio block with cross-modal attention.""" @@ -860,12 +837,7 @@ def forward( aud_rope: Tensor | None, sparse_params: dict | None, attn_mask=None, - modality_mask=None, - av_gate_scale: float = 1.0, - va_gate_scale: float = 1.0, ) -> tuple[Tensor | None, Tensor | None]: - fake_audio = modality_mask[0] if modality_mask is not None else 0 - fake_video = modality_mask[1] if modality_mask is not None else 0 t_v, t_a = time_embed if vis is not None: sa_p, ca_p, ff_p = torch.chunk(self.videoT.visual_modulation(t_v), 3, dim=-1) @@ -927,32 +899,20 @@ def forward( aud_for_av = apply_scale_shift_norm(self.av_normalization, aud, av_scale, av_shift) rq_v = vis_rope if self.ca_rope else None rk_a = aud_rope if self.ca_rope else None - vis_from_aud = ( - self.va_cross_attention( - vis_for_va, - encoder_hidden_states=aud_pre_ca, - rope_q=rq_v, - rope_kv=rk_a, - ) - * (1 - fake_audio) - * (1 - fake_video) + vis_from_aud = self.va_cross_attention( + vis_for_va, + encoder_hidden_states=aud_pre_ca, + rope_q=rq_v, + rope_kv=rk_a, ) - aud_from_vis = ( - self.av_cross_attention( - aud_for_av, - encoder_hidden_states=vis_pre_ca, - rope_q=rk_a, - rope_kv=rq_v, - ) - * (1 - fake_audio) - * (1 - fake_video) + aud_from_vis = self.av_cross_attention( + aud_for_av, + encoder_hidden_states=vis_pre_ca, + rope_q=rk_a, + rope_kv=rq_v, ) - vis = apply_gate_sum( - vis, vis_from_aud, (va_gate if not self.cross_gates else av_gate) * va_gate_scale - ).type_as(vis) - aud = apply_gate_sum( - aud, aud_from_vis, (av_gate if not self.cross_gates else va_gate) * av_gate_scale - ).type_as(aud) + vis = apply_gate_sum(vis, vis_from_aud, va_gate if not self.cross_gates else av_gate).type_as(vis) + aud = apply_gate_sum(aud, aud_from_vis, av_gate if not self.cross_gates else va_gate).type_as(aud) elif vis is not None: vis = apply_gate_sum(vis, vis_out_t, gate_v).type_as(vis) @@ -990,11 +950,13 @@ class Kandinsky6Transformer3DModel( CacheMixin, AttentionMixin, ): - """Kandinsky 6 transformer for text-to-video and text-to-video-and-audio generation. + """Kandinsky 6 multimodal transformer for text/image-to-video-and-audio generation. - The model uses separate text and visual transformer blocks for video-only - generation and fused video/audio blocks for multimodal generation. Inputs - and outputs use packed token layouts compatible with the native K6 model. + Video and audio are denoised together through fused, cross-modal transformer blocks, each + conditioned on its own text branch (Qwen2.5-VL tokens + CLIP pooled embedding). Passing only + `x_video` or only `x_audio` denoises a single modality while still running through the fused + block's per-modality self-attention, cross-attention, and feed-forward stages. Inputs and + outputs use packed token layouts compatible with the native K6 model. Args: in_visual_dim (`int`, *optional*, defaults to 16): Number of input video latent channels. @@ -1005,11 +967,10 @@ class Kandinsky6Transformer3DModel( patch_size (`tuple[int, int, int]`, *optional*, defaults to ``(1, 2, 2)``): Video patch size. model_dim (`int`, *optional*, defaults to 4096): Video transformer hidden dimension. ff_dim (`int`, *optional*, defaults to 16384): Video feed-forward hidden dimension. - num_text_blocks (`int`, *optional*, defaults to 4): Number of text blocks. - num_visual_blocks (`int`, *optional*, defaults to 60): Number of visual blocks. + num_text_blocks (`int`, *optional*, defaults to 4): Number of text blocks per modality. + num_visual_blocks (`int`, *optional*, defaults to 60): Number of fused video/audio blocks. axes_dims (`tuple[int, int, int]`, *optional*, defaults to ``(32, 48, 48)``): RoPE dimensions for video. visual_cond (`bool`, *optional*, defaults to True): Whether video conditioning channels are present. - is_multimodal (`bool`, *optional*, defaults to False): Whether to construct fused video/audio blocks. in_audio_dim (`int`, *optional*, defaults to 20): Number of input audio latent channels. out_audio_dim (`int`, *optional*, defaults to 20): Number of output audio latent channels. model_dim_a (`int`, *optional*): Audio transformer hidden dimension. Defaults to `model_dim`. @@ -1028,12 +989,18 @@ class Kandinsky6Transformer3DModel( _repeated_blocks = [ "Kandinsky6TransformerEncoderBlock", - "Kandinsky6TransformerDecoderBlock", "Kandinsky6FusedTransformerDecoderBlock", ] _no_split_modules = _repeated_blocks - _keep_in_fp32_modules = ["time_embeddings", "modulation", "visual_modulation", "text_modulation"] + _keep_in_fp32_modules = ["time_embeddings", "modulation"] _supports_gradient_checkpointing = True + # `Kandinsky6TimeEmbeddings`/`Kandinsky6Modulation` hand their `nn.Linear` weight/bias straight to + # `functional.linear` (upcast to fp32 first) rather than calling the submodule, so the AdaLN math stays + # fp32 even after a later blanket `.to(bfloat16)` cast (`_keep_in_fp32_modules` only protects the + # `from_pretrained(dtype=...)` load, not a subsequent cast). Leaf-level offload hooks only fire on an + # actual submodule call, so they never see these weights — the same class of gap as + # `HunyuanDiTAttentionPool` (see testing.md), which opts out of group offloading for the same reason. + _supports_group_offloading = False @register_to_config def __init__( @@ -1050,7 +1017,6 @@ def __init__( num_visual_blocks: int = 60, axes_dims: tuple = (32, 48, 48), visual_cond: bool = True, - is_multimodal: bool = False, in_audio_dim: int = 20, out_audio_dim: int = 20, model_dim_a: int | None = None, @@ -1069,13 +1035,11 @@ def __init__( super().__init__() self.patch_size = patch_size self.visual_cond = visual_cond - self.is_multimodal = is_multimodal self.in_visual_dim = in_visual_dim self.in_audio_dim = in_audio_dim self.text_token_padding = text_token_padding self.scale_factor = tuple(float(value) for value in scale_factor) self.visual_token_type_num_embeddings = int(visual_token_type_num_embeddings or 0) - self._text_proj_cache: dict[tuple, object] = {} head_dim = sum(axes_dims) model_dim_a = model_dim_a or model_dim time_dim_a = time_dim_a or time_dim @@ -1090,268 +1054,60 @@ def __init__( self.visual_rope_embeddings = Kandinsky6RoPE3D(axes_dims) self.out_layer = Kandinsky6OutLayer(model_dim, time_dim, out_visual_dim, patch_size) - if not is_multimodal: - self.time_embeddings = Kandinsky6TimeEmbeddings(model_dim, time_dim) - self.text_embeddings = Kandinsky6TextEmbeddings(in_text_dim, model_dim) - self.pooled_text_embeddings = Kandinsky6TextEmbeddings(in_text_dim2, time_dim) - self.text_rope_embeddings = Kandinsky6RoPE1D(head_dim) - self.text_transformer_blocks = nn.ModuleList( - [ - Kandinsky6TransformerEncoderBlock(model_dim, time_dim, ff_dim, head_dim, text_token_padding) - for _ in range(num_text_blocks) - ] - ) - self.visual_transformer_blocks = nn.ModuleList( - [ - Kandinsky6TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim, text_token_padding) - for _ in range(num_visual_blocks) - ] + self.audio_embeddings = Kandinsky6TextEmbeddings(in_audio_dim, model_dim_a) + self.audio_rope_embeddings = Kandinsky6RoPE1D(head_dim_a, freqs_scaling=audio_freqs_scaling) + self.audio_out_layer = Kandinsky6OutLayerAudio(model_dim_a, time_dim_a, out_audio_dim) + for prefix, md, td, fd, hd in ( + ("video", model_dim, time_dim, ff_dim, head_dim), + ("audio", model_dim_a, time_dim_a, ff_dim_a, head_dim_a), + ): + setattr(self, f"{prefix}_time_embeddings", Kandinsky6TimeEmbeddings(md, td)) + setattr(self, f"{prefix}_text_embeddings", Kandinsky6TextEmbeddings(in_text_dim, md)) + setattr(self, f"{prefix}_pooled_text_embeddings", Kandinsky6TextEmbeddings(in_text_dim2, td)) + setattr(self, f"{prefix}_text_rope_embeddings", Kandinsky6RoPE1D(hd)) + setattr( + self, + f"{prefix}_text_transformer_blocks", + nn.ModuleList( + [ + Kandinsky6TransformerEncoderBlock(md, td, fd, hd, text_token_padding) + for _ in range(num_text_blocks) + ] + ), ) - else: - self.audio_embeddings = Kandinsky6TextEmbeddings(in_audio_dim, model_dim_a) - self.audio_rope_embeddings = Kandinsky6RoPE1D(head_dim_a, freqs_scaling=audio_freqs_scaling) - self.audio_out_layer = Kandinsky6OutLayerAudio(model_dim_a, time_dim_a, out_audio_dim) - for prefix, md, td, fd, hd in ( - ("video", model_dim, time_dim, ff_dim, head_dim), - ("audio", model_dim_a, time_dim_a, ff_dim_a, head_dim_a), - ): - setattr(self, f"{prefix}_time_embeddings", Kandinsky6TimeEmbeddings(md, td)) - setattr(self, f"{prefix}_text_embeddings", Kandinsky6TextEmbeddings(in_text_dim, md)) - setattr(self, f"{prefix}_pooled_text_embeddings", Kandinsky6TextEmbeddings(in_text_dim2, td)) - setattr(self, f"{prefix}_text_rope_embeddings", Kandinsky6RoPE1D(hd)) - setattr( - self, - f"{prefix}_text_transformer_blocks", - nn.ModuleList( - [ - Kandinsky6TransformerEncoderBlock(md, td, fd, hd, text_token_padding) - for _ in range(num_text_blocks) - ] - ), + self.visual_transformer_blocks = nn.ModuleList( + [ + Kandinsky6FusedTransformerDecoderBlock( + model_dim, + time_dim, + ff_dim, + head_dim, + model_dim_a, + time_dim_a, + ff_dim_a, + head_dim_a, + text_token_padding, + ca_rope, + cross_gates, + fix_modulation, ) - self.visual_transformer_blocks = nn.ModuleList( - [ - Kandinsky6FusedTransformerDecoderBlock( - model_dim, - time_dim, - ff_dim, - head_dim, - model_dim_a, - time_dim_a, - ff_dim_a, - head_dim_a, - text_token_padding, - ca_rope, - cross_gates, - fix_modulation, - ) - for _ in range(num_visual_blocks) - ] - ) - - self.gradient_checkpointing = False - - def clear_text_proj_cache(self) -> None: - """Clear cached projected text embeddings.""" - self._text_proj_cache.clear() - - def _project_pooled(self, prefix: str | None, pooled: Tensor) -> Tensor: - key = ("pe", prefix, pooled.data_ptr(), tuple(pooled.shape)) - hit = self._text_proj_cache.get(key) - if hit is not None: - return hit # type: ignore[return-value] - pe = ( - self.pooled_text_embeddings(pooled) - if prefix is None - else getattr(self, f"{prefix}_pooled_text_embeddings")(pooled) + for _ in range(num_visual_blocks) + ] ) - self._text_proj_cache[key] = pe - return pe - - def _project_text_tokens(self, prefix: str | None, text_embed: Tensor, pooled: Tensor) -> tuple[Tensor, Tensor]: - key = ("te", prefix, text_embed.data_ptr(), pooled.data_ptr(), tuple(text_embed.shape), tuple(pooled.shape)) - hit = self._text_proj_cache.get(key) - if hit is not None: - return hit # type: ignore[return-value] - te = ( - self.text_embeddings(text_embed) - if prefix is None - else getattr(self, f"{prefix}_text_embeddings")(text_embed) - ) - pe = self._project_pooled(prefix, pooled) - self._text_proj_cache[key] = (te, pe) - return te, pe - - def _time_embed(self, prefix: str | None, time: Tensor, pooled_proj: Tensor) -> Tensor: - return ( - self.time_embeddings(time) if prefix is None else getattr(self, f"{prefix}_time_embeddings")(time) - ) + pooled_proj - - @staticmethod - def _normalize_attn_mask(attn_mask: Tensor | None) -> Tensor | None: - if attn_mask is None: - return None - return attn_mask.unsqueeze(0) if attn_mask.dim() == 1 else attn_mask - - def _run_text_blocks( - self, prefix: str | None, te: Tensor, tm: Tensor, text_rope: Tensor, attn_mask: Tensor | None = None - ) -> Tensor: - blocks = self.text_transformer_blocks if prefix is None else getattr(self, f"{prefix}_text_transformer_blocks") - for block in blocks: - if torch.is_grad_enabled() and self.gradient_checkpointing: - te = self._gradient_checkpointing_func(block, te, tm, text_rope, attn_mask) - else: - te = block(te, tm, text_rope, attn_mask) - return te - - def _encode_text( - self, - prefix: str, - text_embed: Tensor, - pooled: Tensor, - time: Tensor, - text_rope: Tensor, - attn_mask: Tensor | None = None, - ) -> tuple[Tensor, Tensor]: - te, pe = self._project_text_tokens(prefix, text_embed, pooled) - tm = self._time_embed(prefix, time, pe) - return self._run_text_blocks(prefix, te, tm, text_rope, attn_mask), tm - - def _encode_t2v( - self, text_embed: Tensor, pooled: Tensor, time: Tensor, text_rope: Tensor, attn_mask: Tensor | None = None - ) -> tuple[Tensor, Tensor]: - te, pe = self._project_text_tokens(None, text_embed, pooled) - tm = self._time_embed(None, time, pe) - return self._run_text_blocks(None, te, tm, text_rope, attn_mask), tm - - def _time_only(self, prefix: str | None, pooled: Tensor, time: Tensor) -> Tensor: - return self._time_embed(prefix, time, self._project_pooled(prefix, pooled)) - - def _embed_visual( - self, - x_video: Tensor, - visual_rope: Tensor, - sparse_params: dict | None, - *, - apply_fractal: bool = True, - visual_token_type_ids: Tensor | None = None, - ) -> tuple[Tensor, tuple, Tensor]: - if x_video.ndim == 4: - x_video = x_video.unsqueeze(0) - - visual_embed = self.visual_embeddings(x_video) - if hasattr(self, "visual_token_type_embeddings") and visual_token_type_ids is not None: - if visual_token_type_ids.ndim == 1: - visual_token_type_ids = visual_token_type_ids.unsqueeze(0) - token_types = self.visual_token_type_embeddings(visual_token_type_ids.to(device=visual_embed.device)) - token_types = token_types[:, :, None, None, :] - visual_embed = visual_embed + token_types - visual_shape = visual_embed.shape[-4:-1] - to_fractal = sparse_params["to_fractal"] if sparse_params and apply_fractal else False - if to_fractal: - visual_embed = _local_patch(visual_embed, visual_shape, (1, 8, 8), dim=1).flatten(1, 2) - visual_rope = _local_patch(visual_rope, visual_shape, (1, 8, 8), dim=0).flatten(0, 1) - else: - visual_embed = visual_embed.flatten(1, 3) - visual_rope = visual_rope.flatten(0, 2) - return visual_embed, visual_shape, visual_rope - - def _embed_audio(self, x_audio: Tensor, audio_rope: Tensor) -> tuple[Tensor, Tensor]: - if x_audio.ndim == 2: - x_audio = x_audio.unsqueeze(0) - embed = self.audio_embeddings(x_audio) - return embed, audio_rope - def _run_visual_blocks_single( - self, - vis_embed: Tensor | None, - aud_embed: Tensor | None, - te: Tensor, - tm: Tensor, - vis_rope: Tensor | None, - aud_rope: Tensor | None, - sparse_params: dict | None, - attn_mask: Tensor | None = None, - ) -> tuple[Tensor | None, Tensor | None]: - checkpoint = torch.is_grad_enabled() and self.gradient_checkpointing - for block in self.visual_transformer_blocks: - if self.is_multimodal: - if vis_embed is not None and aud_embed is None: - args = (vis_embed, None, te, te, (tm, tm), vis_rope, None, sparse_params, attn_mask) - vis_embed, _ = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) - elif aud_embed is not None and vis_embed is None: - args = (None, aud_embed, te, te, (tm, tm), None, aud_rope, None, attn_mask) - _, aud_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) - else: - raise RuntimeError("single-modality fused path expects exactly one of video/audio") - elif vis_embed is not None: - args = (vis_embed, te, tm, vis_rope, sparse_params, attn_mask) - vis_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) - else: - args = (aud_embed, te, tm, aud_rope, None, attn_mask) - aud_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) - return vis_embed, aud_embed - - def _run_visual_blocks_fused( - self, - vis_embed: Tensor, - aud_embed: Tensor, - video_te: Tensor, - audio_te: Tensor, - video_tm: Tensor, - audio_tm: Tensor, - vis_rope: Tensor, - aud_rope: Tensor, - sparse_params: dict | None, - attn_mask: Tensor | None = None, - ) -> tuple[Tensor, Tensor]: - checkpoint = torch.is_grad_enabled() and self.gradient_checkpointing - for block in self.visual_transformer_blocks: - args = ( - vis_embed, - aud_embed, - video_te, - audio_te, - (video_tm, audio_tm), - vis_rope, - aud_rope, - sparse_params, - attn_mask, - ) - vis_embed, aud_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) - return vis_embed, aud_embed - - def _project_video(self, vis_embed: Tensor, vis_shape: tuple, tm: Tensor, sparse_params: dict | None) -> Tensor: - if sparse_params and sparse_params.get("to_fractal"): - vis_embed = _local_merge( - vis_embed.reshape(vis_embed.shape[0], -1, 64, vis_embed.shape[-1]), - vis_shape, - (1, 8, 8), - dim=1, - ) - else: - vis_embed = vis_embed.reshape(-1, *vis_shape, vis_embed.shape[-1]) - return self.out_layer(vis_embed, tm) - - def _project_audio(self, aud_embed: Tensor, tm: Tensor) -> Tensor: - return self.audio_out_layer(aud_embed, tm) - - def _project_fused( - self, vis_embed: Tensor, aud_embed: Tensor, vis_shape: tuple, video_tm: Tensor, audio_tm: Tensor - ) -> tuple[Tensor, Tensor]: - video = self.out_layer(vis_embed.reshape(-1, *vis_shape, vis_embed.shape[-1]), video_tm) - return video, self.audio_out_layer(aud_embed, audio_tm) + self.gradient_checkpointing = False def forward( self, x_video: Tensor | None = None, x_audio: Tensor | None = None, - text_embed: Tensor | list[Tensor] | None = None, - pooled_text_embed: Tensor | list[Tensor] | None = None, - time: Tensor | list[Tensor] | None = None, + text_embed: Tensor | None = None, + pooled_text_embed: Tensor | None = None, + time: Tensor | None = None, visual_rope: Tensor | None = None, audio_rope: Tensor | None = None, - text_rope: Tensor | list[Tensor] | None = None, + video_text_rope: Tensor | None = None, + audio_text_rope: Tensor | None = None, sparse_params: dict | None = None, attention_mask: Tensor | None = None, visual_token_type_ids: Tensor | None = None, @@ -1363,13 +1119,14 @@ def forward( Args: x_video (`torch.Tensor`, *optional*): Packed video latent tokens. x_audio (`torch.Tensor`, *optional*): Packed audio latent tokens. - text_embed (`torch.Tensor` or `list[torch.Tensor]`): Text token embeddings. - pooled_text_embed (`torch.Tensor` or `list[torch.Tensor]`): Pooled text embeddings. - time (`torch.Tensor` or `list[torch.Tensor]`): Diffusion timesteps. - visual_rope (`torch.Tensor`): Video rotary position embeddings. - audio_rope (`torch.Tensor`, *optional*): Audio rotary position embeddings. - text_rope (`torch.Tensor` or `list[torch.Tensor]`): Text rotary position embeddings. - sparse_params (`dict`, *optional*): NABLA sparse-attention configuration. + text_embed (`torch.Tensor`): Text token embeddings, shared by the video and audio text branches. + pooled_text_embed (`torch.Tensor`): Pooled text embedding, shared by the video and audio text branches. + time (`torch.Tensor`): Diffusion timestep, shared by the video and audio branches. + visual_rope (`torch.Tensor`, *optional*): Video rotary position embeddings. Required with `x_video`. + audio_rope (`torch.Tensor`, *optional*): Audio rotary position embeddings. Required with `x_audio`. + video_text_rope (`torch.Tensor`): Rotary position embeddings for the video text branch. + audio_text_rope (`torch.Tensor`): Rotary position embeddings for the audio text branch. + sparse_params (`dict`, *optional*): NABLA sparse-attention configuration for the video self-attention. attention_mask (`torch.Tensor`, *optional*): Text attention mask. visual_token_type_ids (`torch.Tensor`, *optional*): Video token type IDs. return_dict (`bool`, *optional*, defaults to False): Whether to return a @@ -1377,70 +1134,95 @@ def forward( Returns: `torch.Tensor`, `tuple[torch.Tensor, torch.Tensor]`, or - [`Transformer2DModelOutput`]: Denoised video output, audio and video - outputs, or a model output object. + [`Transformer2DModelOutput`]: The denoised single modality, `(video, audio)` when both `x_video` + and `x_audio` are given, or a model output object. """ - if text_embed is None or pooled_text_embed is None or time is None or text_rope is None: - raise ValueError("text_embed, pooled_text_embed, time, and text_rope are required") - both = x_video is not None and x_audio is not None and self.is_multimodal - attn_mask = self._normalize_attn_mask(attention_mask) - if not both: - te_in = text_embed[0] if isinstance(text_embed, list) else text_embed - pe_in = pooled_text_embed[0] if isinstance(pooled_text_embed, list) else pooled_text_embed - rope_in = text_rope[0] if isinstance(text_rope, list) else text_rope - t_in = time[0] if isinstance(time, list) else time - if self.is_multimodal: - prefix = "audio" if x_audio is not None else "video" - if isinstance(text_rope, list): - rope_in = text_rope[1] if prefix == "audio" else text_rope[0] - te, tm = self._encode_text(prefix, te_in, pe_in, t_in, rope_in, attn_mask) - else: - te, tm = self._encode_t2v(te_in, pe_in, t_in, rope_in, attn_mask) - if x_video is not None: - vis_embed, vis_shape, vis_rope = self._embed_visual( - x_video, visual_rope, sparse_params, visual_token_type_ids=visual_token_type_ids - ) - vis_embed, _ = self._run_visual_blocks_single( - vis_embed, None, te, tm, vis_rope, None, sparse_params, attn_mask + if x_video is None and x_audio is None: + raise ValueError("at least one of `x_video`, `x_audio` must be provided") + if text_embed is None or pooled_text_embed is None or time is None: + raise ValueError("`text_embed`, `pooled_text_embed`, and `time` are required") + if video_text_rope is None or audio_text_rope is None: + raise ValueError("`video_text_rope` and `audio_text_rope` are required") + + attn_mask = attention_mask + if attn_mask is not None and attn_mask.dim() == 1: + attn_mask = attn_mask.unsqueeze(0) + checkpoint = torch.is_grad_enabled() and self.gradient_checkpointing + + # 1. Encode text tokens through the video and audio text branches + video_text_embed = self.video_text_embeddings(text_embed) + video_temb = self.video_time_embeddings(time) + self.video_pooled_text_embeddings(pooled_text_embed) + for block in self.video_text_transformer_blocks: + args = (video_text_embed, video_temb, video_text_rope, attn_mask) + video_text_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) + + audio_text_embed = self.audio_text_embeddings(text_embed) + audio_temb = self.audio_time_embeddings(time) + self.audio_pooled_text_embeddings(pooled_text_embed) + for block in self.audio_text_transformer_blocks: + args = (audio_text_embed, audio_temb, audio_text_rope, attn_mask) + audio_text_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) + + # 2. Patchify the video latents and embed the audio latents (only the modalities that were passed) + visual_embed, visual_shape = None, None + to_fractal = bool(sparse_params and sparse_params.get("to_fractal")) + if x_video is not None: + if x_video.ndim == 4: + x_video = x_video.unsqueeze(0) + visual_embed = self.visual_embeddings(x_video) + if visual_token_type_ids is not None and hasattr(self, "visual_token_type_embeddings"): + token_type_ids = ( + visual_token_type_ids.unsqueeze(0) if visual_token_type_ids.ndim == 1 else visual_token_type_ids ) - result: Tensor | tuple[Tensor, Tensor] = self._project_video(vis_embed, vis_shape, tm, sparse_params) + token_types = self.visual_token_type_embeddings(token_type_ids.to(device=visual_embed.device)) + visual_embed = visual_embed + token_types[:, :, None, None, :] + visual_shape = visual_embed.shape[-4:-1] + if to_fractal: + visual_embed = _local_patch(visual_embed, visual_shape, (1, 8, 8), dim=1).flatten(1, 2) + visual_rope = _local_patch(visual_rope, visual_shape, (1, 8, 8), dim=0).flatten(0, 1) else: - aud_embed, aud_rope = self._embed_audio(x_audio, audio_rope) - _, aud_embed = self._run_visual_blocks_single(None, aud_embed, te, tm, None, aud_rope, None, attn_mask) - result = self._project_audio(aud_embed, tm) - else: - te_v, pe_v = ( - (text_embed[0], pooled_text_embed[0]) - if isinstance(text_embed, list) - else (text_embed, pooled_text_embed) - ) - te_a, pe_a = ( - (text_embed[1], pooled_text_embed[1]) - if isinstance(text_embed, list) - else (text_embed, pooled_text_embed) - ) - rope_v, rope_a = (text_rope[0], text_rope[1]) if isinstance(text_rope, list) else (text_rope, text_rope) - t_v, t_a = (time[0], time[1]) if isinstance(time, list) else (time, time) - video_te, video_tm = self._encode_text("video", te_v, pe_v, t_v, rope_v, attn_mask) - audio_te, audio_tm = self._encode_text("audio", te_a, pe_a, t_a, rope_a, attn_mask) - vis_embed, vis_shape, vis_rope = self._embed_visual( - x_video, visual_rope, sparse_params, apply_fractal=False, visual_token_type_ids=visual_token_type_ids - ) - aud_embed, aud_rope = self._embed_audio(x_audio, audio_rope) - vis_embed, aud_embed = self._run_visual_blocks_fused( - vis_embed, - aud_embed, - video_te, - audio_te, - video_tm, - audio_tm, - vis_rope, - aud_rope, + visual_embed = visual_embed.flatten(1, 3) + visual_rope = visual_rope.flatten(0, 2) + + audio_embed = None + if x_audio is not None: + if x_audio.ndim == 2: + x_audio = x_audio.unsqueeze(0) + audio_embed = self.audio_embeddings(x_audio) + + # 3. Run the fused video/audio transformer blocks (a block no-ops on whichever modality is absent) + for block in self.visual_transformer_blocks: + args = ( + visual_embed, + audio_embed, + video_text_embed, + audio_text_embed, + (video_temb, audio_temb), + visual_rope, + audio_rope, sparse_params, attn_mask, ) - result = self._project_fused(vis_embed, aud_embed, vis_shape, video_tm, audio_tm) + visual_embed, audio_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) + + # 4. Project back to patch/latent space + video_out = None + if visual_embed is not None: + if to_fractal: + visual_embed = _local_merge( + visual_embed.reshape(visual_embed.shape[0], -1, 64, visual_embed.shape[-1]), + visual_shape, + (1, 8, 8), + dim=1, + ) + else: + visual_embed = visual_embed.reshape(-1, *visual_shape, visual_embed.shape[-1]) + video_out = self.out_layer(visual_embed, video_temb) + audio_out = self.audio_out_layer(audio_embed, audio_temb) if audio_embed is not None else None + if video_out is not None and audio_out is not None: + result = (video_out, audio_out) + else: + result = video_out if video_out is not None else audio_out return Transformer2DModelOutput(sample=result) if return_dict else result diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index 458976d37051..30a1fe03a977 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -192,7 +192,15 @@ def _batch_size( ) -> int: sizes = [ cls._value_batch_size(value) - for value in (prompt, negative_prompt, image, latents, audio_latents, prompt_embeds, negative_prompt_embeds) + for value in ( + prompt, + negative_prompt, + image, + latents, + audio_latents, + prompt_embeds, + negative_prompt_embeds, + ) ] batch_size = max(sizes) if batch_size < 1 or any(size not in (1, batch_size) for size in sizes): @@ -229,25 +237,26 @@ def check_inputs( image: str | object | None = None, visual_cond_scheme: str = "pretrain", ) -> None: - """Validate arguments shared by text-to-video and image-to-video calls. + r"""Validate arguments shared by text-to-video and image-to-video calls. Args: - prompt: Text prompt or a batch of prompts. - negative_prompt: Negative text prompt or a batch of prompts. - height: Requested output height in pixels. - width: Requested output width in pixels. - num_frames: Requested output frame count. - num_inference_steps: Number of denoising steps. - max_sequence_length: Maximum Qwen prompt length after the template. - output_type: One of `pt`, `torch`, `np`, `numpy`, or `latent`. - prompt_embeds: Optional precomputed Qwen prompt embeddings. - pooled_prompt_embeds: Optional precomputed CLIP pooled prompt embeddings. - negative_prompt_embeds: Optional precomputed negative Qwen embeddings. - negative_pooled_prompt_embeds: Optional precomputed negative CLIP pooled embeddings. - callback_on_step_end_tensor_inputs: Tensor names exposed to the step callback. - sample_audio: Whether to generate and decode audio. - image: Optional reference image or batch of reference images. - visual_cond_scheme: Image-conditioning scheme when `image` is set. + prompt (`str` or `list[str]`, *optional*): Text prompt or a batch of prompts. + negative_prompt (`str` or `list[str]`, *optional*): Negative text prompt or a batch of prompts. + height (`int`): Requested output height in pixels. + width (`int`): Requested output width in pixels. + num_frames (`int`): Requested output frame count. + num_inference_steps (`int`): Number of denoising steps. + max_sequence_length (`int`): Maximum Qwen prompt length after the template. + output_type (`str`): One of `"pt"`, `"torch"`, `"np"`, `"numpy"`, or `"latent"`. + prompt_embeds (`torch.Tensor`, *optional*): Precomputed Qwen prompt embeddings. + pooled_prompt_embeds (`torch.Tensor`, *optional*): Precomputed CLIP pooled prompt embeddings. + negative_prompt_embeds (`torch.Tensor`, *optional*): Precomputed negative Qwen embeddings. + negative_pooled_prompt_embeds (`torch.Tensor`, *optional*): Precomputed negative CLIP pooled embeddings. + callback_on_step_end_tensor_inputs (`list[str]`, *optional*): Tensor names exposed to the step callback. + sample_audio (`bool`, *optional*, defaults to `True`): Whether to generate and decode audio. + image (`str`, `PIL.Image.Image`, or a list thereof, *optional*): Reference image or batch of images. + visual_cond_scheme (`str`, *optional*, defaults to `"pretrain"`): Image-conditioning scheme when `image` + is set. Raises: ValueError: If an input combination is unsupported. @@ -284,8 +293,6 @@ def check_inputs( if sample_audio and self.audio_vae is None: raise ValueError("sample_audio=True requires an audio_vae") - if getattr(self.transformer, "is_multimodal", False) is not True: - raise ValueError("Kandinsky6TI2VAPipeline requires a multimodal transformer") patch_size = tuple(int(value) for value in getattr(self.transformer, "patch_size", (1, 2, 2))) if len(patch_size) < 3 or any(value <= 0 for value in patch_size): raise ValueError(f"transformer.patch_size must contain three positive values, got {patch_size}") @@ -313,21 +320,25 @@ def check_inputs( raise ValueError("tail_cond_first_frame requires transformer.visual_token_type_num_embeddings >= 2") def _encode_single_prompt( - self, text: str | list[str], max_sequence_length: int + self, text: str | list[str], max_sequence_length: int, device: torch.device ) -> tuple[Tensor, Tensor, Tensor | None]: - """Encode one prompt (or batch of prompts) with Qwen2.5-VL and CLIP. + r"""Encode one prompt (or batch of prompts) with Qwen2.5-VL and CLIP. Args: - text: Prompt or prompts to encode. - max_sequence_length: Maximum number of Qwen text tokens. + text (`str` or `list[str]`): Prompt or prompts to encode. + max_sequence_length (`int`): Maximum number of Qwen text tokens. + device (`torch.device`): Execution device to move tokenized inputs to. This must be the pipeline's + `_execution_device`, not a component's own parameter device: under CPU/group offloading a + component's resting parameter device can be `meta` or `cpu` even though it computes on the + accelerator once its forward hook runs. Returns: - `(prompt_embeds, pooled_prompt_embeds, attention_mask)`. `attention_mask` is `None` when every prompt - in the batch fills `max_sequence_length` (an all-`True` mask carries no extra information). + `tuple[torch.Tensor, torch.Tensor, torch.Tensor or None]`: `(prompt_embeds, pooled_prompt_embeds, + attention_mask)`. `attention_mask` is `None` when every prompt in the batch fills + `max_sequence_length` (an all-`True` mask carries no extra information). """ texts = [text] if isinstance(text, str) else text full_texts = [_PROMPT_TEMPLATE.format(item) for item in texts] - qwen_device = next(self.text_encoder.parameters()).device inputs = self.tokenizer( text=full_texts, images=None, @@ -336,7 +347,7 @@ def _encode_single_prompt( truncation=True, return_tensors="pt", padding="max_length", - ).to(qwen_device) + ).to(device) qwen_output = self.text_encoder( input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], @@ -351,7 +362,6 @@ def _encode_single_prompt( # _normalize_attn_mask). Drop it so the dense mask isn't threaded through the denoise loop. attention_mask = None - clip_device = next(self.text_encoder_2.parameters()).device clip_inputs = self.tokenizer_2( texts, max_length=_CLIP_MAX_LENGTH, @@ -359,7 +369,7 @@ def _encode_single_prompt( add_special_tokens=True, padding="max_length", return_tensors="pt", - ).to(clip_device) + ).to(device) pooled_prompt_embeds = self.text_encoder_2(**clip_inputs)["pooler_output"] return prompt_embeds, pooled_prompt_embeds, attention_mask @@ -375,33 +385,41 @@ def encode_prompt( negative_pooled_prompt_embeds: Tensor | None = None, negative_prompt_attention_mask: Tensor | None = None, max_sequence_length: int = 1024, + device: torch.device | None = None, ) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None, Tensor | None, Tensor | None]: - """Encode the positive and (optionally) negative prompt. + r"""Encode the positive and (optionally) negative prompt. Args: - prompt: Prompt or prompts to encode. - negative_prompt: Negative prompt or prompts, used when `do_classifier_free_guidance=True` and - `negative_prompt_embeds` is not already provided. - do_classifier_free_guidance: Whether to also encode `negative_prompt`. - prompt_embeds: Optional precomputed Qwen prompt embeddings, skipping encoding of `prompt`. - pooled_prompt_embeds: Optional precomputed CLIP pooled prompt embeddings. - prompt_attention_mask: Optional Qwen attention mask paired with `prompt_embeds`. - negative_prompt_embeds: Optional precomputed negative Qwen embeddings. - negative_pooled_prompt_embeds: Optional precomputed negative CLIP pooled embeddings. - negative_prompt_attention_mask: Optional Qwen attention mask paired with `negative_prompt_embeds`. - max_sequence_length: Maximum Qwen prompt length after the template. + prompt (`str` or `list[str]`): Prompt or prompts to encode. + negative_prompt (`str` or `list[str]`, *optional*): Negative prompt or prompts, used when + `do_classifier_free_guidance=True` and `negative_prompt_embeds` is not already provided. + do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): Whether to also encode + `negative_prompt`. + prompt_embeds (`torch.Tensor`, *optional*): Precomputed Qwen prompt embeddings, skipping encoding of + `prompt`. + pooled_prompt_embeds (`torch.Tensor`, *optional*): Precomputed CLIP pooled prompt embeddings. + prompt_attention_mask (`torch.Tensor`, *optional*): Qwen attention mask paired with `prompt_embeds`. + negative_prompt_embeds (`torch.Tensor`, *optional*): Precomputed negative Qwen embeddings. + negative_pooled_prompt_embeds (`torch.Tensor`, *optional*): Precomputed negative CLIP pooled embeddings. + negative_prompt_attention_mask (`torch.Tensor`, *optional*): Qwen attention mask paired with + `negative_prompt_embeds`. + max_sequence_length (`int`, *optional*, defaults to 1024): Maximum Qwen prompt length after the + template. + device (`torch.device`, *optional*): Execution device for newly-encoded prompts. Defaults to + `self._execution_device`. Returns: - `(prompt_embeds, pooled_prompt_embeds, prompt_attention_mask, negative_prompt_embeds, - negative_pooled_prompt_embeds, negative_prompt_attention_mask)`. + `tuple[torch.Tensor, ...]`: `(prompt_embeds, pooled_prompt_embeds, prompt_attention_mask, + negative_prompt_embeds, negative_pooled_prompt_embeds, negative_prompt_attention_mask)`. """ + device = device or self._execution_device if prompt_embeds is None: prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self._encode_single_prompt( - prompt, max_sequence_length + prompt, max_sequence_length, device ) if do_classifier_free_guidance and negative_prompt_embeds is None: negative_prompt_embeds, negative_pooled_prompt_embeds, negative_prompt_attention_mask = ( - self._encode_single_prompt(negative_prompt, max_sequence_length) + self._encode_single_prompt(negative_prompt, max_sequence_length, device) ) return ( prompt_embeds, @@ -412,19 +430,32 @@ def encode_prompt( negative_prompt_attention_mask, ) - @torch.no_grad() def expand_prompts( self, prompt: str | list[str], image: Any | list[Any] | None = None, max_sequence_length: int = 1024, mode: str | None = None, + generator: torch.Generator | None = None, ) -> str | list[str]: - """Expand T2VA prompts, or image-grounded I2VA prompts when an image is supplied.""" + """Expand T2VA prompts, or image-grounded I2VA prompts when an image is supplied. + + Args: + prompt: Prompt or a batch of prompts to expand. + image: Optional reference image or batch of reference images, for `mode="i2va"`. + max_sequence_length: Maximum number of new tokens to generate per prompt. + mode: `"t2va"` or `"i2va"`. Defaults to `"i2va"` when `image` is given, else `"t2va"`. + generator: Optional generator whose seed makes the (sampled) expansion reproducible. + + Returns: + The expanded prompt, or a list of expanded prompts when `prompt` is a list. + """ if isinstance(prompt, list): images = image if isinstance(image, (list, tuple)) else [image] * len(prompt) return [ - self.expand_prompts(item, image=item_image, max_sequence_length=max_sequence_length, mode=mode) + self.expand_prompts( + item, image=item_image, max_sequence_length=max_sequence_length, mode=mode, generator=generator + ) for item, item_image in zip(prompt, images, strict=True) ] mode = mode or ("i2va" if image is not None else "t2va") @@ -438,7 +469,7 @@ def expand_prompts( content.append({"type": "text", "text": instruction}) messages = [{"role": "user", "content": content}] text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) - device = next(self.text_encoder.parameters()).device + device = self._execution_device inputs = self.tokenizer( text=[text], images=[image] if image is not None else None, @@ -446,6 +477,12 @@ def expand_prompts( padding=True, return_tensors="pt", ).to(device) + if generator is not None: + # `generate()` samples from the global RNG rather than accepting a `torch.Generator`, so seed + # the global RNG from the pipeline's own generator to make the expansion reproducible. This + # doesn't affect later `randn_tensor(..., generator=generator)` calls, which use `generator` + # directly rather than the global RNG state. + torch.manual_seed(generator.initial_seed()) generated = self.text_encoder.generate(**inputs, max_new_tokens=max_sequence_length) qwen_crop_start = inputs["input_ids"].shape[1] trimmed = [output[qwen_crop_start:] for output in generated] @@ -476,7 +513,6 @@ def _load_pil_rgb(image: str | object): raise TypeError(f"i2va image must be a path or PIL image, got {type(image).__name__}") return pil_image - @torch.no_grad() def _encode_i2va_first_frame(self, image: str | object, device: torch.device, height: int, width: int) -> Tensor: """Resize, center-crop, and VAE-encode one reference image into the packed K6 first-frame latent.""" try: @@ -582,7 +618,9 @@ def _append_i2va_tail_condition( first_frames = first_frames.to(device=latents.device, dtype=latents.dtype) expected = (batch_size, height, width, channels) if tuple(first_frames.shape) != expected: - raise ValueError(f"first-frame latent shape mismatch: expected {expected}, got {tuple(first_frames.shape)}") + raise ValueError( + f"first-frame latent shape mismatch: expected {expected}, got {tuple(first_frames.shape)}" + ) latents = torch.cat([latents, first_frames[:, None]], dim=1) token_types = torch.cat( @@ -655,28 +693,27 @@ def _prepare_text_ropes( negative_text_length: int, audio_length: int | None, device: torch.device, - ) -> tuple[list[Tensor], list[Tensor], Tensor | None]: - """Build per-modality text RoPE for the cond/uncond branches, plus the audio RoPE.""" - text_rope = [ - self._compute_rope1d(self.transformer.video_text_rope_embeddings, text_length, device=device), - self._compute_rope1d(self.transformer.audio_text_rope_embeddings, text_length, device=device), - ] - negative_text_rope = [ - self._compute_rope1d(self.transformer.video_text_rope_embeddings, negative_text_length, device=device), - self._compute_rope1d(self.transformer.audio_text_rope_embeddings, negative_text_length, device=device), - ] + ) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor | None]: + """Build the video/audio text-branch RoPE for the cond/uncond passes, plus the audio latent RoPE.""" + video_text_rope = self._compute_rope1d(self.transformer.video_text_rope_embeddings, text_length, device=device) + audio_text_rope = self._compute_rope1d(self.transformer.audio_text_rope_embeddings, text_length, device=device) + negative_video_text_rope = self._compute_rope1d( + self.transformer.video_text_rope_embeddings, negative_text_length, device=device + ) + negative_audio_text_rope = self._compute_rope1d( + self.transformer.audio_text_rope_embeddings, negative_text_length, device=device + ) audio_rope = ( self._compute_rope1d(self.transformer.audio_rope_embeddings, audio_length, device=device) if audio_length is not None else None ) - return text_rope, negative_text_rope, audio_rope + return video_text_rope, audio_text_rope, negative_video_text_rope, negative_audio_text_rope, audio_rope @staticmethod def _apply_guidance(cond: Tensor, uncond: Tensor, guidance_scale: float) -> Tensor: return uncond + guidance_scale * (cond - uncond) - @torch.no_grad() def _postprocess_video(self, latents: Tensor) -> Tensor: """Decode video latents `(B, T, H, W, C)` into `(B, 3, T, H, W)` `uint8` frames in `[0, 255]`.""" frames = (latents / self.vae.config.scaling_factor).permute(0, 4, 1, 2, 3) @@ -685,7 +722,6 @@ def _postprocess_video(self, latents: Tensor) -> Tensor: frames = self.vae.decode(frames.to(dtype=vae_dtype)).sample return ((frames.clamp(-1.0, 1.0) + 1.0) * 127.5).to(torch.uint8) - @torch.no_grad() def _postprocess_audio(self, audio_latents: Tensor | None) -> list[np.ndarray] | None: """Decode audio latents `(B, A, D)` into a list of int16 `(samples,)` waveforms, or `None` in T2V mode.""" if audio_latents is None: @@ -697,9 +733,11 @@ def _postprocess_audio(self, audio_latents: Tensor | None) -> list[np.ndarray] | decoded = decoded.unsqueeze(0) elif decoded.ndim == 3 and decoded.shape[1] == 1: decoded = decoded[:, 0] - return [ - (np.clip(waveform.cpu().float().numpy(), -1.0, 1.0) * 32767).astype(np.int16) for waveform in decoded - ] + return [(np.clip(waveform.cpu().float().numpy(), -1.0, 1.0) * 32767).astype(np.int16) for waveform in decoded] + + @property + def guidance_scale(self): + return self._guidance_scale @property def num_timesteps(self): @@ -727,6 +765,7 @@ def __call__( num_inference_steps: int = 50, max_sequence_length: int = 1024, guidance_scale: float = 5.0, + num_images_per_prompt: int = 1, generator: torch.Generator | None = None, latents: Tensor | None = None, audio_latents: Tensor | None = None, @@ -745,39 +784,43 @@ def __call__( r"""Generate synchronized video and, optionally, audio from text or an image. Args: - prompt: Text prompt or a batch of prompts. - image: Optional reference image or batch of reference images. - negative_prompt: Optional negative prompt. - height: Output video height in pixels; must be divisible by 8. - width: Output video width in pixels; must be divisible by 8. - num_frames: Number of decoded video frames. - sample_fps: Output video frame rate. Kandinsky 6 is trained for 24.0 fps; other values may affect - audio/video alignment. - num_inference_steps: Number of denoising steps. - max_sequence_length: Maximum Qwen prompt length after the template. - guidance_scale: Classifier-free guidance weight. - generator: Random generator used for latent initialization. - latents: Optional precomputed video latents. - audio_latents: Optional precomputed audio latents. - prompt_embeds: Optional precomputed Qwen prompt embeddings. - pooled_prompt_embeds: Optional precomputed CLIP pooled prompt embeddings. - negative_prompt_embeds: Optional precomputed negative Qwen embeddings. - negative_pooled_prompt_embeds: Optional precomputed negative CLIP pooled embeddings. - sample_audio: Whether to generate synchronized audio. - expand_prompts: Whether to use the built-in Qwen video+audio prompt expander before encoding. - output_type: `pt`/`torch` for tensors, `np`/`numpy` for NumPy arrays, or `latent` for undecoded video - latents. - return_dict: Whether to return `Kandinsky6TI2VAPipelineOutput`. - callback_on_step_end: Optional callback invoked after each step. - callback_on_step_end_tensor_inputs: Names passed to the callback. - visual_cond_scheme: Optional image-conditioning scheme. Defaults to `pretrain` without `image` and - `tail_cond_first_frame` with `image`. + prompt (`str` or `list[str]`, *optional*): Text prompt or a batch of prompts. + image (`str`, `PIL.Image.Image`, or a list thereof, *optional*): Reference image or batch of images. + negative_prompt (`str` or `list[str]`, *optional*): Negative prompt. + height (`int`, *optional*, defaults to 512): Output video height in pixels; must be divisible by 8. + width (`int`, *optional*, defaults to 768): Output video width in pixels; must be divisible by 8. + num_frames (`int`, *optional*, defaults to 121): Number of decoded video frames. + sample_fps (`float`, *optional*, defaults to 24.0): Output video frame rate. Kandinsky 6 is trained + for 24.0 fps; other values may affect audio/video alignment. + num_inference_steps (`int`, *optional*, defaults to 50): Number of denoising steps. + max_sequence_length (`int`, *optional*, defaults to 1024): Maximum Qwen prompt length after the + template. + guidance_scale (`float`, *optional*, defaults to 5.0): Classifier-free guidance weight. + num_images_per_prompt (`int`, *optional*, defaults to 1): Number of videos to generate per prompt. + generator (`torch.Generator`, *optional*): Random generator used for latent initialization and, when + `expand_prompts=True`, to seed prompt expansion. + latents (`torch.Tensor`, *optional*): Precomputed video latents. + audio_latents (`torch.Tensor`, *optional*): Precomputed audio latents. + prompt_embeds (`torch.Tensor`, *optional*): Precomputed Qwen prompt embeddings. + pooled_prompt_embeds (`torch.Tensor`, *optional*): Precomputed CLIP pooled prompt embeddings. + negative_prompt_embeds (`torch.Tensor`, *optional*): Precomputed negative Qwen embeddings. + negative_pooled_prompt_embeds (`torch.Tensor`, *optional*): Precomputed negative CLIP pooled embeddings. + sample_audio (`bool`, *optional*, defaults to `True`): Whether to generate synchronized audio. + expand_prompts (`bool`, *optional*, defaults to `False`): Whether to use the built-in Qwen video+audio + prompt expander before encoding. + output_type (`str`, *optional*, defaults to `"pt"`): `"pt"`/`"torch"` for tensors, `"np"`/`"numpy"` for + NumPy arrays, or `"latent"` for undecoded video latents. + return_dict (`bool`, *optional*, defaults to `True`): Whether to return `Kandinsky6TI2VAPipelineOutput`. + callback_on_step_end (`Callable`, *optional*): Callback invoked after each step. + callback_on_step_end_tensor_inputs (`list[str]`, *optional*): Names passed to the callback. + visual_cond_scheme (`str`, *optional*): Image-conditioning scheme. Defaults to `"pretrain"` without + `image` and `"tail_cond_first_frame"` with `image`. Examples: Returns: - `Kandinsky6TI2VAPipelineOutput` containing video frames and int16 NumPy audio waveforms, or a tuple - when `return_dict=False`. + [`Kandinsky6TI2VAPipelineOutput`] or `tuple`: [`Kandinsky6TI2VAPipelineOutput`] containing video + frames and int16 NumPy audio waveforms if `return_dict=True`, otherwise a `(frames, audio)` tuple. """ if sample_fps != 24.0: warnings.warn( @@ -817,12 +860,15 @@ def __call__( ) callback_on_step_end_tensor_inputs = callback_on_step_end_tensor_inputs or self._callback_tensor_inputs + self._guidance_scale = guidance_scale self._current_timestep = None self._interrupt = False device = self._execution_device dtype = ( - self.transformer.dtype if isinstance(getattr(self.transformer, "dtype", None), torch.dtype) else torch.bfloat16 + self.transformer.dtype + if isinstance(getattr(self.transformer, "dtype", None), torch.dtype) + else torch.bfloat16 ) # 3. Encode the reference image, if any @@ -839,6 +885,7 @@ def __call__( image=image_batch, max_sequence_length=max_sequence_length, mode="i2va" if image_batch is not None else "t2va", + generator=generator, ) negative_prompt_batch = negative_prompt_batch or [self._DEFAULT_NEGATIVE_PROMPT] * batch_size @@ -853,6 +900,7 @@ def __call__( prompt=prompt_batch or [""] * batch_size, negative_prompt=negative_prompt_batch, do_classifier_free_guidance=True, + device=device, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, @@ -868,6 +916,20 @@ def __call__( if negative_prompt_attention_mask is not None: negative_prompt_attention_mask = negative_prompt_attention_mask.to(device=device) + prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) + pooled_prompt_embeds = pooled_prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) + negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) + if prompt_attention_mask is not None: + prompt_attention_mask = prompt_attention_mask.repeat_interleave(num_images_per_prompt, dim=0) + if negative_prompt_attention_mask is not None: + negative_prompt_attention_mask = negative_prompt_attention_mask.repeat_interleave( + num_images_per_prompt, dim=0 + ) + if first_frames is not None: + first_frames = first_frames.repeat_interleave(num_images_per_prompt, dim=0) + batch_size = batch_size * num_images_per_prompt + # 5. Prepare latent variables patch_size = tuple(int(value) for value in self.transformer.patch_size) num_latent_frames = (num_frames - 1) // 4 + 1 @@ -917,11 +979,13 @@ def __call__( if generated_visual_mask is not None: # The appended reference frame reuses the first frame's rotary position rather than a real T+1 one. visual_rope = torch.cat([visual_rope, visual_rope[:1]], dim=0) - text_rope, negative_text_rope, audio_rope = self._prepare_text_ropes( - text_length=prompt_embeds.shape[1], - negative_text_length=negative_prompt_embeds.shape[1], - audio_length=audio_latents.shape[1] if audio_latents is not None else None, - device=device, + video_text_rope, audio_text_rope, negative_video_text_rope, negative_audio_text_rope, audio_rope = ( + self._prepare_text_ropes( + text_length=prompt_embeds.shape[1], + negative_text_length=negative_prompt_embeds.shape[1], + audio_length=audio_latents.shape[1] if audio_latents is not None else None, + device=device, + ) ) # 7. Denoising loop @@ -962,7 +1026,8 @@ def __call__( time=timestep, visual_rope=visual_rope, audio_rope=audio_rope, - text_rope=text_rope, + video_text_rope=video_text_rope, + audio_text_rope=audio_text_rope, attention_mask=prompt_attention_mask, visual_token_type_ids=visual_token_type_ids, ) @@ -977,15 +1042,16 @@ def __call__( time=timestep, visual_rope=visual_rope, audio_rope=audio_rope, - text_rope=negative_text_rope, + video_text_rope=negative_video_text_rope, + audio_text_rope=negative_audio_text_rope, attention_mask=negative_prompt_attention_mask, visual_token_type_ids=visual_token_type_ids, ) if audio_latents is not None: - video_pred = self._apply_guidance(model_pred[0], model_pred_uncond[0], guidance_scale) - audio_pred = self._apply_guidance(model_pred[1], model_pred_uncond[1], guidance_scale) + video_pred = self._apply_guidance(model_pred[0], model_pred_uncond[0], self.guidance_scale) + audio_pred = self._apply_guidance(model_pred[1], model_pred_uncond[1], self.guidance_scale) else: - video_pred = self._apply_guidance(model_pred, model_pred_uncond, guidance_scale) + video_pred = self._apply_guidance(model_pred, model_pred_uncond, self.guidance_scale) elif audio_latents is not None: video_pred, audio_pred = model_pred else: diff --git a/tests/models/transformers/test_models_transformer_kandinsky6.py b/tests/models/transformers/test_models_transformer_kandinsky6.py new file mode 100644 index 000000000000..93d1a150531c --- /dev/null +++ b/tests/models/transformers/test_models_transformer_kandinsky6.py @@ -0,0 +1,175 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import Kandinsky6Transformer3DModel +from diffusers.models.transformers.transformer_kandinsky6 import Kandinsky6RoPE1D, Kandinsky6RoPE3D +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class Kandinsky6TransformerTesterConfig: + @property + def model_class(self): + return Kandinsky6Transformer3DModel + + @property + def pretrained_model_name_or_path(self): + return "" # TODO: Set Hub repository ID + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} + + @property + def main_input_name(self) -> str: + return "x_video" + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict[str, int | list[int]]: + return { + "in_visual_dim": 4, + "out_visual_dim": 4, + "in_text_dim": 8, + "in_text_dim2": 8, + "time_dim": 16, + "patch_size": (1, 2, 2), + "model_dim": 48, + "ff_dim": 64, + "num_text_blocks": 1, + "num_visual_blocks": 2, + "axes_dims": (4, 4, 4), + "visual_cond": True, + "in_audio_dim": 4, + "out_audio_dim": 4, + } + + def _build_dummy_inputs( + self, batch_size: int, num_frames: int, height: int, width: int, device + ) -> dict[str, torch.Tensor]: + init_dict = self.get_init_dict() + head_dim = sum(init_dict["axes_dims"]) + patch_t, patch_h, patch_w = init_dict["patch_size"] + vis_in_dim = 2 * init_dict["in_visual_dim"] + 1 # visual_cond=True prepends cond + mask channels + text_length = 6 + + x_video = randn_tensor( + (batch_size, num_frames, height, width, vis_in_dim), generator=self.generator, device=device + ) + text_embed = randn_tensor( + (batch_size, text_length, init_dict["in_text_dim"]), generator=self.generator, device=device + ) + pooled_text_embed = randn_tensor( + (batch_size, init_dict["in_text_dim2"]), generator=self.generator, device=device + ) + # `self.generator` is always a CPU generator (see below); draw on CPU and move, like `randn_tensor` does + # internally, so this also works with an accelerator (e.g. MPS) `device`. + time = torch.randint(0, 1000, (batch_size,), generator=self.generator).float().to(device) + + visual_shape = (num_frames // patch_t, height // patch_h, width // patch_w) + visual_rope = Kandinsky6RoPE3D(init_dict["axes_dims"]).to(device)( + visual_shape, + [torch.arange(size, device=device) for size in visual_shape], + ) + video_text_rope = Kandinsky6RoPE1D(head_dim).to(device)(torch.arange(text_length, device=device)) + audio_text_rope = Kandinsky6RoPE1D(head_dim).to(device)(torch.arange(text_length, device=device)) + + return { + "x_video": x_video, + "text_embed": text_embed, + "pooled_text_embed": pooled_text_embed, + "time": time, + "visual_rope": visual_rope, + "video_text_rope": video_text_rope, + "audio_text_rope": audio_text_rope, + } + + def get_dummy_inputs(self, batch_size: int = 1, device=torch_device) -> dict[str, torch.Tensor]: + # Only `x_video` is passed (no `x_audio`): the fused block no-ops its audio branch when `aud is None`, + # so this still runs the real dual-stream block weights, while `forward` returns a single tensor + # (matching what the generic mixins below expect) instead of a `(video, audio)` tuple. + return self._build_dummy_inputs(batch_size, num_frames=2, height=4, width=4, device=device) + + @property + def input_shape(self) -> tuple[int, ...]: + # (num_frames, height, width, vis_in_dim), batch dimension excluded. + return (2, 4, 4, 2 * 4 + 1) + + @property + def output_shape(self) -> tuple[int, ...]: + # (num_frames, height, width, out_visual_dim), batch dimension excluded. + return (2, 4, 4, 4) + + +class TestKandinsky6TransformerModel(Kandinsky6TransformerTesterConfig, ModelTesterMixin): + pass + + +class TestKandinsky6TransformerMemory(Kandinsky6TransformerTesterConfig, MemoryTesterMixin): + # `Kandinsky6TimeEmbeddings`/`Kandinsky6Modulation` bypass leaf-level offload hooks by design (see the + # `_supports_group_offloading = False` comment on the model); the classic accelerate cpu/disk offload + # hooks hit the exact same gap and have no model-level opt-out flag, so skip explicitly here instead. + _OFFLOAD_SKIP_REASON = ( + "Kandinsky6TimeEmbeddings/Kandinsky6Modulation read their nn.Linear weight/bias directly for a " + "fp32-upcast functional.linear call rather than calling the submodule, so leaf-level offload hooks " + "never see (or move) those weights." + ) + + def test_cpu_offload(self, *args, **kwargs): + pytest.skip(self._OFFLOAD_SKIP_REASON) + + def test_disk_offload_without_safetensors(self, *args, **kwargs): + pytest.skip(self._OFFLOAD_SKIP_REASON) + + def test_disk_offload_with_safetensors(self, *args, **kwargs): + pytest.skip(self._OFFLOAD_SKIP_REASON) + + +class TestKandinsky6TransformerTorchCompile(Kandinsky6TransformerTesterConfig, TorchCompileTesterMixin): + @property + def different_shapes_for_compilation(self): + return [(4, 4), (4, 8), (8, 8)] + + def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]: + return self._build_dummy_inputs(batch_size=1, num_frames=2, height=height, width=width, device=torch_device) + + +class TestKandinsky6TransformerTraining(Kandinsky6TransformerTesterConfig, TrainingTesterMixin): + pass + + +class TestKandinsky6TransformerAttention(Kandinsky6TransformerTesterConfig, AttentionTesterMixin): + def test_fuse_unfuse_qkv_projections(self, *args, **kwargs): + pytest.skip( + "Kandinsky6Attention names its projections to_query/to_key/to_value/out_layer (matching the " + "native K6 checkpoint layout) rather than to_q/to_k/to_v/to_out, which " + "AttentionModuleMixin.fuse_projections hardcodes." + ) diff --git a/tests/pipelines/kandinsky6/__init__.py b/tests/pipelines/kandinsky6/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py new file mode 100644 index 000000000000..a571bbc11254 --- /dev/null +++ b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py @@ -0,0 +1,240 @@ +# Copyright 2025 The Kandinsky Team and The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import PIL.Image +import pytest +import torch +from transformers import ( + AutoProcessor, + CLIPTextConfig, + CLIPTextModel, + CLIPTokenizer, + Qwen2_5_VLConfig, + Qwen2_5_VLForConditionalGeneration, +) + +from diffusers import ( + AutoencoderKLHunyuanVideo, + FlowMatchEulerDiscreteScheduler, + Kandinsky6TI2VAPipeline, + Kandinsky6Transformer3DModel, +) + +from ...testing_utils import torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin + + +class Kandinsky6TI2VAPipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = Kandinsky6TI2VAPipeline + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "num_frames", "num_inference_steps", "guidance_scale"] + ) + batch_input_params = frozenset(["prompt", "negative_prompt"]) + # (channels, num_frames, height, width) for output_type="pt", matching `get_dummy_inputs()`'s + # (num_frames=5, height=16, width=16) at the tiny VAE's real 8x spatial / 4x temporal compression. + output_shape = (3, 5, 16, 16) + + def get_dummy_components(self): + torch.manual_seed(0) + # 3 down/up levels so the tiny VAE's *realized* spatial/temporal compression actually matches the + # declared spatial_compression_ratio=8 / temporal_compression_ratio=4 (2 levels only realizes 4x/2x) — + # the pipeline itself hardcodes `// 8` / `// 4` when sizing latents rather than reading these configs. + vae = AutoencoderKLHunyuanVideo( + act_fn="silu", + block_out_channels=[8, 8, 8], + down_block_types=[ + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + "HunyuanVideoDownBlock3D", + ], + in_channels=3, + latent_channels=4, + layers_per_block=1, + mid_block_add_attention=False, + norm_num_groups=2, + out_channels=3, + scaling_factor=0.476986, + spatial_compression_ratio=8, + temporal_compression_ratio=4, + up_block_types=[ + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + "HunyuanVideoUpBlock3D", + ], + ) + + scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0) + + # mrope_section must sum to (hidden_size / num_attention_heads) / 2 — matches the known-good + # hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration-compatible config used by Kandinsky5's + # own test suite (tests/pipelines/kandinsky5/test_kandinsky5.py). + qwen_hidden_size = 32 + torch.manual_seed(0) + qwen_config = Qwen2_5_VLConfig( + text_config={ + "hidden_size": qwen_hidden_size, + "intermediate_size": qwen_hidden_size, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "rope_scaling": { + "mrope_section": [2, 2, 4], + "rope_type": "default", + "type": "default", + }, + "rope_theta": 1000000.0, + }, + vision_config={ + "depth": 2, + "hidden_size": qwen_hidden_size, + "intermediate_size": qwen_hidden_size, + "num_heads": 2, + "out_hidden_size": qwen_hidden_size, + }, + hidden_size=qwen_hidden_size, + vocab_size=152064, + vision_end_token_id=151653, + vision_start_token_id=151652, + vision_token_id=151654, + ) + text_encoder = Qwen2_5_VLForConditionalGeneration(qwen_config) + tokenizer = AutoProcessor.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") + + clip_hidden_size = 16 + torch.manual_seed(0) + clip_config = CLIPTextConfig( + bos_token_id=0, + eos_token_id=2, + hidden_size=clip_hidden_size, + intermediate_size=16, + layer_norm_eps=1e-05, + num_attention_heads=2, + num_hidden_layers=2, + pad_token_id=1, + vocab_size=1000, + projection_dim=clip_hidden_size, + ) + text_encoder_2 = CLIPTextModel(clip_config) + tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + + torch.manual_seed(0) + transformer = Kandinsky6Transformer3DModel( + in_visual_dim=4, + out_visual_dim=4, + in_text_dim=qwen_hidden_size, + in_text_dim2=clip_hidden_size, + time_dim=16, + patch_size=(1, 2, 2), + model_dim=12, + ff_dim=24, + num_text_blocks=1, + num_visual_blocks=2, + axes_dims=(2, 1, 1), + visual_cond=True, + in_audio_dim=4, + out_audio_dim=4, + visual_token_type_num_embeddings=2, + ) + + return { + "transformer": transformer, + "vae": vae, + "text_encoder": text_encoder, + "audio_vae": None, + "scheduler": scheduler, + "tokenizer": tokenizer, + "text_encoder_2": text_encoder_2, + "tokenizer_2": tokenizer_2, + } + + def get_dummy_inputs(self): + return { + "prompt": "A cat and a dog baking a cake together in a kitchen.", + "negative_prompt": "static, blurry", + "generator": self.get_generator(0), + "num_inference_steps": 2, + "guidance_scale": 5.0, + "height": 16, + "width": 16, + "num_frames": 5, + "max_sequence_length": 16, + "sample_audio": False, + "output_type": "pt", + } + + +class TestKandinsky6TI2VAPipeline(Kandinsky6TI2VAPipelineTesterConfig, PipelineTesterMixin): + def test_kandinsky6_ti2va_i2va(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs() + inputs["image"] = PIL.Image.fromarray(np.zeros((16, 16, 3), dtype=np.uint8)) + + output = pipe(**inputs) + + assert output.frames.shape == (1, *self.output_shape) + assert not torch.isnan(output.frames.float()).any() + + def test_kandinsky6_ti2va_different_images(self): + # A freshly initialized (untrained) transformer zero-initializes every AdaLN modulation gate + # (`Kandinsky6Modulation.out_layer`), so at init every attention/feed-forward residual is gated to + # exactly zero: neither the text branch nor (with the default `tail_cond_first_frame` scheme) a + # reference image reaches the *kept* output frames yet, since that path is also self-attention + # (equally gated). `visual_cond_scheme="i2v"` is different: it overwrites frame 0's latent with the + # image encoding directly (`_apply_visual_conditioning`), and that frame is kept (not appended and + # dropped like `tail_cond_first_frame`), so it passes through two real (non-zero-init) linears — + # `visual_embeddings` and `out_layer` — around the identity (zero-gated) blocks. That's the one + # conditioning signal guaranteed to already produce different output on a fresh model. + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs() + inputs["visual_cond_scheme"] = "i2v" + + inputs["image"] = PIL.Image.fromarray(np.zeros((16, 16, 3), dtype=np.uint8)) + output_black_image = pipe(**inputs).frames + + inputs = self.get_dummy_inputs() + inputs["visual_cond_scheme"] = "i2v" + inputs["image"] = PIL.Image.fromarray(np.full((16, 16, 3), 255, dtype=np.uint8)) + output_white_image = pipe(**inputs).frames + + max_diff = (output_black_image.float() - output_white_image.float()).abs().max() + assert max_diff > 1e-6, "Outputs should be different for different reference images." + + def test_kandinsky6_ti2va_num_images_per_prompt(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs() + inputs["num_images_per_prompt"] = 2 + + output = pipe(**inputs) + + assert output.frames.shape == (2, *self.output_shape) + + +class TestKandinsky6TI2VAPipelineMemory(Kandinsky6TI2VAPipelineTesterConfig, MemoryTesterMixin): + # `Kandinsky6TimeEmbeddings`/`Kandinsky6Modulation` (in the transformer) read their `nn.Linear` + # weight/bias directly for a fp32-upcast `functional.linear` call rather than calling the submodule + # (see the `_supports_group_offloading = False` comment on `Kandinsky6Transformer3DModel`), so + # leaf-level offload hooks never see those weights. `enable_sequential_cpu_offload` hits the same gap + # and has no pipeline-level opt-out flag; confirmed this isn't an environment issue — Flux's identical + # tests pass on this same machine/backend. + _OFFLOAD_SKIP_REASON = ( + "Kandinsky6TimeEmbeddings/Kandinsky6Modulation bypass leaf-level offload hooks by design; see " + "Kandinsky6Transformer3DModel._supports_group_offloading." + ) + + def test_sequential_cpu_offload_forward_pass(self, *args, **kwargs): + pytest.skip(self._OFFLOAD_SKIP_REASON) + + def test_sequential_offload_forward_pass_twice(self, *args, **kwargs): + pytest.skip(self._OFFLOAD_SKIP_REASON) From d5f8e5803be5d8aecad46279959f2a076da15b9e Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 28 Sep 2026 12:00:58 +0300 Subject: [PATCH 16/62] refactor piflow scheduler --- src/diffusers/schedulers/scheduling_piflow.py | 2 +- tests/schedulers/test_scheduler_piflow.py | 199 ++++++++++++++++++ 2 files changed, 200 insertions(+), 1 deletion(-) create mode 100644 tests/schedulers/test_scheduler_piflow.py diff --git a/src/diffusers/schedulers/scheduling_piflow.py b/src/diffusers/schedulers/scheduling_piflow.py index 59e6be1715dd..4e0ac7ad90f6 100644 --- a/src/diffusers/schedulers/scheduling_piflow.py +++ b/src/diffusers/schedulers/scheduling_piflow.py @@ -1,4 +1,4 @@ -# Copyright 2025 The Kandinsky Team and The HuggingFace Team. All rights reserved. +# 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. diff --git a/tests/schedulers/test_scheduler_piflow.py b/tests/schedulers/test_scheduler_piflow.py new file mode 100644 index 000000000000..6a21ae8b5594 --- /dev/null +++ b/tests/schedulers/test_scheduler_piflow.py @@ -0,0 +1,199 @@ +# 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 unittest + +import torch + +from diffusers import PiflowScheduler +from diffusers.schedulers.scheduling_piflow import PiflowSchedulerOutput + + +class PiflowSchedulerTest(unittest.TestCase): + """ + PiflowScheduler's model output is "widened" (``n_grid`` predictions packed into the channel + dimension) rather than matching the sample shape, so it cannot use ``SchedulerCommonTest`` (see + ``FlowMapEulerDiscreteSchedulerTest`` for the same situation with a different non-standard + contract). These tests exercise the contract `Kandinsky6TI2VAPipeline` actually relies on. + """ + + scheduler_class = PiflowScheduler + + def get_default_config(self, **kwargs): + config = { + "num_train_timesteps": 1000, + "shift": 5.0, + "n_grid": 4, + "eps": 1e-6, + "final_step_size_scale": 0.5, + "num_policy_substeps": 32, + } + config.update(**kwargs) + return config + + def make_widened_output(self, velocity: torch.Tensor, n_grid: int) -> torch.Tensor: + """Pack one velocity prediction into `n_grid` identical grid slots (channel-minor layout).""" + return velocity.repeat(1, n_grid) + + # ---- config validation ---- + + def test_instantiation_with_defaults(self): + scheduler = self.scheduler_class(**self.get_default_config()) + self.assertEqual(scheduler.config.num_train_timesteps, 1000) + self.assertEqual(scheduler.config.n_grid, 4) + self.assertTrue(scheduler.is_piflow) + + def test_invalid_n_grid_raises(self): + with self.assertRaises(ValueError): + self.scheduler_class(**self.get_default_config(n_grid=1)) + + def test_invalid_eps_raises(self): + with self.assertRaises(ValueError): + self.scheduler_class(**self.get_default_config(eps=0.0)) + + def test_invalid_final_step_size_scale_raises(self): + with self.assertRaises(ValueError): + self.scheduler_class(**self.get_default_config(final_step_size_scale=0.0)) + with self.assertRaises(ValueError): + self.scheduler_class(**self.get_default_config(final_step_size_scale=1.5)) + + def test_invalid_num_policy_substeps_raises(self): + with self.assertRaises(ValueError): + self.scheduler_class(**self.get_default_config(num_policy_substeps=0)) + + # ---- set_timesteps ---- + + def test_set_timesteps_shapes_and_monotonic(self): + scheduler = self.scheduler_class(**self.get_default_config()) + for nfe in [1, 2, 4, 8, 16]: + scheduler.set_timesteps(nfe) + self.assertEqual(scheduler.timesteps.shape, (nfe,)) + self.assertEqual(scheduler.sigmas.shape, (nfe + 1,)) + self.assertEqual(scheduler.sigmas[-1].item(), 0.0) + # strictly decreasing: this schedule should never hit the duplicate-timestep branch + # in `_step_index_for` that other flow schedulers need for interpolated sigmas. + diffs = scheduler.timesteps[1:] - scheduler.timesteps[:-1] + self.assertTrue(torch.all(diffs < 0)) + + def test_set_timesteps_rejects_unsupported_args(self): + scheduler = self.scheduler_class(**self.get_default_config()) + with self.assertRaises(ValueError): + scheduler.set_timesteps(4, sigmas=[1.0, 0.5, 0.0]) + with self.assertRaises(ValueError): + scheduler.set_timesteps(4, mu=1.0) + with self.assertRaises(ValueError): + scheduler.set_timesteps(4, timesteps=[900, 500, 100]) + with self.assertRaises(ValueError): + scheduler.set_timesteps(0) + + # ---- step() input validation ---- + + def test_step_rejects_integer_timestep(self): + scheduler = self.scheduler_class(**self.get_default_config()) + scheduler.set_timesteps(4) + sample = torch.randn(1, 3) + model_output = self.make_widened_output(torch.randn(1, 3), scheduler.n_grid) + with self.assertRaises(ValueError): + scheduler.step(model_output, 0, sample) + + def test_step_after_exhausted_schedule_raises(self): + scheduler = self.scheduler_class(**self.get_default_config()) + scheduler.set_timesteps(2) + sample = torch.randn(1, 3) + for t in scheduler.timesteps: + model_output = self.make_widened_output(torch.randn(1, 3), scheduler.n_grid) + sample = scheduler.step(model_output, t, sample).prev_sample + with self.assertRaises(RuntimeError): + scheduler.step( + self.make_widened_output(torch.randn(1, 3), scheduler.n_grid), scheduler.timesteps[-1], sample + ) + + def test_step_return_dict_false_returns_tuple(self): + scheduler = self.scheduler_class(**self.get_default_config()) + scheduler.set_timesteps(2) + sample = torch.randn(1, 3) + model_output = self.make_widened_output(torch.randn(1, 3), scheduler.n_grid) + output = scheduler.step(model_output, scheduler.timesteps[0], sample, return_dict=False) + self.assertIsInstance(output, tuple) + output_dict = scheduler.step(model_output, scheduler.timesteps[0], sample, return_dict=True) + self.assertIsInstance(output_dict, PiflowSchedulerOutput) + + def test_to_grid_rejects_mismatched_shapes(self): + scheduler = self.scheduler_class(**self.get_default_config()) + scheduler.set_timesteps(2) + sample = torch.randn(1, 3) + with self.assertRaises(ValueError): + # channel count not divisible by n_grid + scheduler.step(torch.randn(1, 5), scheduler.timesteps[0], sample) + with self.assertRaises(ValueError): + # divisible by n_grid, but the per-grid width doesn't match the sample's channel count + scheduler.step(torch.randn(1, 4 * 2), scheduler.timesteps[0], sample) + + # ---- correctness ---- + + def test_step_is_deterministic(self): + # Like every diffusers scheduler with an auto-incrementing `_step_index` (e.g. + # FlowMatchEulerDiscreteScheduler), `step()` only resolves `timestep` into an index on the + # first call after `set_timesteps`; a second call reuses the advanced internal counter + # regardless of the `timestep` passed in. So determinism must be checked across two freshly + # reset schedulers, not two `step()` calls on the same instance. + torch.manual_seed(0) + sample = torch.randn(2, 5) + model_output = self.make_widened_output(torch.randn(2, 5), self.get_default_config()["n_grid"]) + + scheduler1 = self.scheduler_class(**self.get_default_config()) + scheduler1.set_timesteps(4) + out1 = scheduler1.step(model_output, scheduler1.timesteps[0], sample.clone()).prev_sample + + scheduler2 = self.scheduler_class(**self.get_default_config()) + scheduler2.set_timesteps(4) + out2 = scheduler2.step(model_output, scheduler2.timesteps[0], sample.clone()).prev_sample + + torch.testing.assert_close(out1, out2) + + def test_no_nan_across_configs(self): + for n_grid in (2, 4, 8): + for nfe in (1, 2, 8): + scheduler = self.scheduler_class(**self.get_default_config(n_grid=n_grid)) + scheduler.set_timesteps(nfe) + sample = torch.randn(2, 6) + for t in scheduler.timesteps: + model_output = self.make_widened_output(torch.randn(2, 6), n_grid) + sample = scheduler.step(model_output, t, sample).prev_sample + self.assertFalse(torch.isnan(sample).any(), f"NaN with n_grid={n_grid}, nfe={nfe}") + self.assertFalse(torch.isinf(sample).any(), f"Inf with n_grid={n_grid}, nfe={nfe}") + + def test_step_recovers_known_clean_target(self): + """A "perfect" model whose implied x0 prediction is always exactly `target` (i.e. every grid + slot predicts the velocity of the straight line from the current sample to `target`) should + drive the sample to `target` after the full schedule: that straight-line ODE has a velocity + that is exactly constant along the true path, so the scheduler's Euler-style substep + integration has zero discretization error and the only residual comes from stopping at + `eps` instead of sigma=0. This catches sign/broadcasting/indexing bugs that a shape-only or + no-NaN check would miss. + """ + torch.manual_seed(0) + scheduler = self.scheduler_class(**self.get_default_config(n_grid=4, num_policy_substeps=64)) + scheduler.set_timesteps(8) + batch, dim = 2, 3 + target = torch.randn(batch, dim) + sample = torch.randn(batch, dim) * 3.0 + 5.0 # arbitrary, far from target + + for i, t in enumerate(scheduler.timesteps): + sigma = scheduler.sigmas[i] + velocity = (sample - target) / sigma.clamp(min=scheduler.eps) + model_output = self.make_widened_output(velocity, scheduler.n_grid) + sample = scheduler.step(model_output, t, sample).prev_sample + + torch.testing.assert_close(sample, target, atol=1e-3, rtol=1e-3) From 7038c5af65520e0ae77b84ed6f99e620eee7790b Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 28 Sep 2026 13:43:30 +0300 Subject: [PATCH 17/62] refactor SR --- .../autoencoders/autoencoder_kandinsky6_sr.py | 56 +- .../transformers/transformer_kandinsky6_sr.py | 1129 +----- .../pipelines/kandinsky6/latent_upscaler.py | 1201 +----- .../kandinsky6/pipeline_kandinsky6_sr.py | 3584 +++++------------ .../pipelines/kandinsky6/pipeline_output.py | 6 - 5 files changed, 1280 insertions(+), 4696 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index 5f710788b069..587518ac487b 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -19,15 +19,15 @@ import functools import math from dataclasses import dataclass -from typing import Tuple, Union import numpy as np import torch import torch.nn as nn import torch.nn.functional as F + from ...configuration_utils import ConfigMixin, register_to_config -from ..modeling_utils import ModelMixin from ...utils.accelerate_utils import apply_forward_hook +from ..modeling_utils import ModelMixin from .vae import DiagonalGaussianDistribution @@ -40,7 +40,9 @@ class SafeConv3d(nn.Conv3d): def forward(self, x, write_to=None, transform=None): if transform is None: - transform = lambda x: x + + def transform(x): + return x memory_count = x.numel() / (10**9) if memory_count > 2: @@ -93,7 +95,7 @@ def __init__( self, chan_in, chan_out, - kernel_size: Union[int, Tuple[int, int, int]], + kernel_size: int | tuple[int, int, int], stride=(1, 1, 1), dilation=(1, 1, 1), padding_mode=None, @@ -194,13 +196,13 @@ def forward(self, input_, cache: dict, fixed_stride=False): if cache["padding"] is None: first_frame = input_parallel[:, :, :1] - time_pad_shape = [i for i in first_frame.shape] + time_pad_shape = list(first_frame.shape) time_pad_shape[2] = self.time_pad padding = first_frame.expand(time_pad_shape) else: padding = cache["padding"] - out_size = [i for i in input_.shape] + out_size = list(input_.shape) out_size[1] = self.conv.out_channels if t_stride == 2: out_size[2] = (input_.size(2) + 1) // 2 @@ -271,10 +273,7 @@ def __init__( self.conv1 = CachedCausalConv3d( chan_in=in_channels, chan_out=out_channels, kernel_size=3, padding_mode=padding_mode ) - if temb_channels > 0: - self.temb_proj = torch.nn.Linear(temb_channels, out_channels) self.norm2 = normalization(out_channels, zq_ch=zq_ch, add_conv=add_conv) - # self.dropout = torch.nn.Dropout(dropout) self.conv2 = CachedCausalConv3d( chan_in=out_channels, chan_out=out_channels, kernel_size=3, padding_mode=padding_mode ) @@ -293,8 +292,6 @@ def __init__( ) def forward(self, x, temb, layer_cache, zq=None): - if x.size(2) == 17 and x.size(3) == 1080 and zq is not None: - torch.cuda.empty_cache() h = x if zq is None: @@ -302,19 +299,8 @@ def forward(self, x, temb, layer_cache, zq=None): else: h = self.norm1(h, zq, cache=layer_cache["norm1"]) - if x.size(2) == 17 and x.size(3) == 1080 and zq is not None: - torch.cuda.empty_cache() - h = F.silu(h, inplace=True) - if x.size(2) == 17 and x.size(3) == 1080 and zq is not None: - torch.cuda.empty_cache() - h = self.conv1(h, cache=layer_cache["conv1"]) - if x.size(2) == 17 and x.size(3) == 1080 and zq is not None: - torch.cuda.empty_cache() - - if temb is not None: - h = h + self.temb_proj(F.silu(temb))[:, :, None, None, None] if zq is None: h = self.norm2(h, cache=layer_cache["norm2"]) @@ -322,7 +308,6 @@ def forward(self, x, temb, layer_cache, zq=None): h = self.norm2(h, zq, cache=layer_cache["norm2"]) h = F.silu(h, inplace=True) - # h = self.dropout(h) h = self.conv2(h, cache=layer_cache["conv2"]) if self.in_channels != self.out_channels: @@ -647,7 +632,6 @@ def __init__( self.down = nn.ModuleList() for i_level in range(self.num_resolutions): block = nn.ModuleList() - attn = nn.ModuleList() block_in = round(ch * in_ch_mult[i_level]) block_out = round(ch * ch_mult[i_level]) @@ -669,7 +653,6 @@ def __init__( block_in = block_out down = nn.Module() down.block = block - down.attn = attn if i_level != self.num_resolutions - 1: if temporal_compress_start_level <= i_level < temporal_compress_level: down.downsample = CachedPXSDownsample( @@ -731,8 +714,6 @@ def forward(self, x, cache_dict, use_cp=True): if not self.skip_last_resolution or i_level != self.num_resolutions - 1: for i_block in range(self.num_res_blocks): h = self.down[i_level].block[i_block](h, temb, layer_cache=cache_dict[i_level][i_block]) - if len(self.down[i_level].attn) > 0: - h = self.down[i_level].attn[i_block](h) if i_level != self.num_resolutions - 1: h = self.down[i_level].downsample(h, cache=cache_dict[i_level]["down"]) @@ -747,9 +728,6 @@ def forward(self, x, cache_dict, use_cp=True): return h - def get_last_layer(self): - return self.conv_out.conv.weight - class CachedDecoder3D(nn.Module): def __init__( @@ -824,7 +802,6 @@ def __init__( self.up = nn.ModuleList() for i_level in reversed(range(self.num_resolutions)): block = nn.ModuleList() - attn = nn.ModuleList() block_out = round(ch * ch_mult[i_level]) for i_block in range(self.num_res_blocks + 1): block.append( @@ -842,7 +819,6 @@ def __init__( block_in = block_out up = nn.Module() up.block = block - up.attn = attn if i_level != 0: if ( self.num_resolutions - temporal_compress_start_level @@ -875,9 +851,6 @@ def forward(self, z, cache_dict): for i_level in reversed(range(self.num_resolutions)): for i_block in range(self.num_res_blocks + 1): h = self.up[i_level].block[i_block](h, temb, layer_cache=cache_dict[i_level][i_block], zq=zq) - - if len(self.up[i_level].attn) > 0: - h = self.up[i_level].attn[i_block](h, zq) if i_level != 0: h = self.up[i_level].upsample(h, cache_dict[i_level]["up"]) @@ -891,9 +864,6 @@ def forward(self, z, cache_dict): return h - def get_last_layer(self): - return self.conv_out.conv.weight - @dataclass class DecoderOutput: @@ -943,13 +913,13 @@ def __init__( out_channels: int = 3, z_channels: int = 16, ch: int = 128, - ch_mult: Tuple[float, ...] = (1, 2, 4, 8), - decoder_ch: Union[int, None] = None, - decoder_ch_mult: Union[Tuple[float, ...], None] = None, + ch_mult: tuple[float, ...] = (1, 2, 4, 8), + decoder_ch: int | None = None, + decoder_ch_mult: tuple[float, ...] | None = None, num_res_blocks: int = 2, dropout: float = 0.0, resolution: int = 0, - padding_mode: Union[str, None] = None, + padding_mode: str | None = None, temporal_compress_times: int = 4, temporal_compress_start_level: int = 0, norm_type: str = "group_norm", @@ -958,7 +928,7 @@ def __init__( fix_pxs: bool = False, skip_last_resolution: bool = False, give_pre_end: bool = False, - zq_ch: Union[int, None] = None, + zq_ch: int | None = None, add_conv: bool = False, scaling_factor: float = 1.0, spatial_factor: int = 16, diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index d831c3cd8ae1..149ee50236e8 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -16,20 +16,19 @@ from __future__ import annotations -import dataclasses import math from collections.abc import Mapping from typing import Any import torch -from ...configuration_utils import ConfigMixin, register_to_config -from ..attention import AttentionModuleMixin -from ..attention_dispatch import dispatch_attention_fn -from ..modeling_utils import ModelMixin -from diffusers.utils import logging from torch import Tensor, nn from torch.nn.attention.flex_attention import BlockMask, flex_attention +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..modeling_utils import ModelMixin from .transformer_kandinsky6 import Kandinsky6AttnProcessor @@ -136,7 +135,6 @@ def _packed_attention_mask( FRACTAL_BLOCK_SIZE = 8 -@torch.autocast(device_type="cuda", enabled=False) def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: """Compute sinusoidal frequency schedule for rotary embeddings. @@ -305,12 +303,11 @@ def nablaT_v2_doc( wW: int = 3, thr: float = 0.9, add_sta: bool = True, - method: str = "topcdf", ) -> BlockMask: """Build a dynamic sparse attention BlockMask with document boundaries. Estimates an approximate attention map from block-averaged queries and - keys, then binarizes it to select the most relevant blocks per query. + keys, then binarizes it (CDF cutoff) to select the most relevant blocks per query. Args: q: Query tensor of shape ``(B, heads, seq_len, dim)``. @@ -322,16 +319,12 @@ def nablaT_v2_doc( wT: Temporal window for static attention. wH: Height window for static attention. wW: Width window for static attention. - thr: Threshold for CDF cutoff (``topcdf``) or token fraction (``topk``). + thr: CDF threshold for binarization. add_sta: Whether to add static local attention. - method: Binarization method, ``"topcdf"`` or ``"topk"``. Returns: Sparse ``BlockMask`` with block size 64. """ - if method not in {"topcdf", "topk"}: - msg = f"nabla method should be topcdf or topk, got {method}" - raise ValueError(msg) # Q/K are the authoritative execution tensors when this function is # reached through attention. Keep every mask intermediate on their device. device = q.device @@ -351,146 +344,6 @@ def nablaT_v2_doc( ) attn_map += doc.log() attn_map = torch.softmax(attn_map / math.sqrt(D), dim=-1) - if method == "topcdf": - # Map binarization - vals, inds = attn_map.sort(-1) - cvals = vals.cumsum_(-1) - mask = (cvals >= 1 - thr).int() - mask = mask.gather(-1, inds.argsort(-1)) - else: - attn_map = attn_map.reshape(B * h * S // 64, S // 64) - dl = d.tolist() - start_row = 0 - mask = torch.zeros_like(attn_map) - for di in dl: - d_full = di * W * H * h * B - end_row = start_row + d_full - k = max(1, int(thr * di * W * H)) - group = attn_map[start_row:end_row, :] - _, topk_indices = torch.topk(group, k, dim=-1) - row_indices = torch.arange(start_row, end_row, device=mask.device).view(-1, 1) - mask[row_indices, topk_indices] = 1 - start_row = end_row - mask = mask.reshape(B, h, S // 64, S // 64) - - if add_sta: - sta = fast_sta_nabla(T, H, W, wT, wH, wW, device=device).unsqueeze_(0).unsqueeze_(0) - mask = torch.logical_or(mask, sta) - mask = torch.logical_and(mask, doc) - - # BlockMask creation - kv_nb = mask.sum(-1).to(torch.int32) - kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) - return BlockMask.from_kv_blocks(torch.zeros_like(kv_nb), kv_inds, kv_nb, kv_inds, BLOCK_SIZE=64, mask_mod=None) - - -def block_mask_from_bool(mask: Tensor) -> BlockMask: - """Convert a boolean block-attention matrix into a flex-attention ``BlockMask``. - - Args: - mask: Boolean tensor broadcastable to ``(B, heads, S // 64, S // 64)``; - ``True`` at ``(..., i, j)`` lets query block ``i`` attend key block ``j``. - - Returns: - Sparse ``BlockMask`` with block size 64 and no intra-block ``mask_mod``. - """ - kv_nb = mask.sum(-1).to(torch.int32) - kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) - return BlockMask.from_kv_blocks(torch.zeros_like(kv_nb), kv_inds, kv_nb, kv_inds, BLOCK_SIZE=64, mask_mod=None) - - -def build_framewise_causal_block_doc( - seq: Tensor, - H: int, - W: int, - *, - mf: int = 2, -) -> Tensor: - """Build the frame-wise causal block-attention structure (no content sparsity). - - Causal across frames (a frame attends only to current and past frames), fully - bidirectional within each ``mf``-frame group and within a frame, and isolated - across packed documents. Operates at 64-block granularity, so ``H * W`` is the - number of 64-token blocks per frame. Weight-independent: depends only on the - document layout, never on Q/K content. - - Args: - seq: Cumulative per-document frame counts, ``(num_docs + 1,)``. - H: Block-grid height per frame. - W: Block-grid width per frame (``H * W`` = 64-token blocks per frame). - mf: Multi-frame group size; frames in a group attend bidirectionally. - - Returns: - Boolean ``(S // 64, S // 64)`` block-attention matrix. - """ - device = seq.device - d = torch.diff(seq) - doc1 = ( - torch.eye(d.numel(), dtype=torch.bool, device=device) - .repeat_interleave(d, dim=0) - .repeat_interleave(d, dim=1) - .tril() - ) - group_sizes = [[c.sum().item() for c in torch.ones((dd,)).split(mf)] for dd in d] - cl = torch.tensor([x for xs in group_sizes for x in xs], dtype=torch.int32, device=device) - doc2 = ( - torch.eye(cl.numel(), dtype=torch.bool, device=device).repeat_interleave(cl, dim=0).repeat_interleave(cl, dim=1) - ) - return torch.logical_or(doc1, doc2).repeat_interleave(H * W, dim=0).repeat_interleave(H * W, dim=1) - - -@torch.compile(dynamic=True) -@torch.no_grad() -def nablaT_v2_doc_mfcausal( - q: Tensor, - k: Tensor, - seq: Tensor, - T: int, - H: int, - W: int, - *, - wT: int = 3, - wH: int = 3, - wW: int = 3, - thr: float = 0.9, - add_sta: bool = True, - mf: int = 2, -) -> BlockMask: - """Build a dynamic sparse attention BlockMask with multi-frame causal masking. - - Similar to ``nablaT_v2_doc`` but adds causal constraints across - multi-frame groups within each document. - - Args: - q: Query tensor of shape ``(B, heads, seq_len, dim)``. - k: Key tensor of shape ``(B, heads, seq_len, dim)``. - seq: Cumulative document lengths. - T: Temporal dimension. - H: Height dimension. - W: Width dimension. - wT: Temporal window for static attention. - wH: Height window for static attention. - wW: Width window for static attention. - thr: CDF threshold for binarization. - add_sta: Whether to add static local attention. - mf: Multi-frame group size for causal masking. - - Returns: - Sparse ``BlockMask`` with block size 64. - """ - # Q/K are the authoritative execution tensors. - device = q.device - seq = seq.to(device=device) - - # Map estimation - B, h, S, D = q.shape - qa = q.reshape(B, h, S // 64, 64, D).mean(-2) - ka = k.reshape(B, h, S // 64, 64, D).mean(-2).transpose(-2, -1) - attn_map = qa @ ka - - doc = build_framewise_causal_block_doc(seq, H, W, mf=mf) - attn_map += doc.log() - attn_map = torch.softmax(attn_map / math.sqrt(D), dim=-1) # Map binarization vals, inds = attn_map.sort(-1) @@ -503,117 +356,25 @@ def nablaT_v2_doc_mfcausal( mask = torch.logical_or(mask, sta) mask = torch.logical_and(mask, doc) - return block_mask_from_bool(mask) - - -@torch.no_grad() -def framewise_causal_dense( - seq: Tensor, - H: int, - W: int, - *, - mf: int = 2, -) -> BlockMask: - """Build the dense (non-sparse) frame-wise causal ``BlockMask``. - - The exact, weight-independent causal reference: the frame-wise causal block - structure (causal across frames, bidirectional within an ``mf``-group and within - a frame) with no content-based sparsity. Runs through the same ``flex_attention`` - kernel as the NABLA path; only the ``BlockMask`` content differs. Used as the - streaming ground-truth mask and to validate that the NABLA causal selection stays - within this causal envelope. - - Args: - seq: Cumulative per-document frame counts, ``(num_docs + 1,)``. - H: Block-grid height per frame. - W: Block-grid width per frame (``H * W`` = 64-token blocks per frame). - mf: Multi-frame group size; frames in a group attend bidirectionally. - - Returns: - Dense causal ``BlockMask`` with block size 64. - """ - doc = build_framewise_causal_block_doc(seq, H, W, mf=mf) - return block_mask_from_bool(doc[None, None]) + # BlockMask creation + kv_nb = mask.sum(-1).to(torch.int32) + kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) + return BlockMask.from_kv_blocks(torch.zeros_like(kv_nb), kv_inds, kv_nb, kv_inds, BLOCK_SIZE=64, mask_mod=None) # Neural network building blocks for the diffusion transformer. -logger = logging.get_logger(__name__) - +# `mode="max-autotune-no-cudagraphs"` autotunes Triton kernels for CUDA; fall back to eager +# `flex_attention` on CPU/MPS (tests, non-CUDA inference) instead of attempting CUDA autotuning. +_compiled_flex = torch.compile(flex_attention, mode="max-autotune-no-cudagraphs", dynamic=True) -def _ensure_nabla_compatible_flex_bwd_configs() -> None: - """Guarantee flex_attention's backward keeps a 64-block-compatible config. - ``nablaT_v2_doc`` builds a ``BLOCK_SIZE=64`` sparse mask. flex_attention's - backward lowering drops every autotune config whose block sizes don't divide - the mask block size (``SPARSE_BLOCK_SIZE % conf.block_* == 0``), checking each - config's OWN block sizes — ``kernel_options`` cannot inject one. On torch - builds whose head_dim=64 backward configs use ``BLOCK_N=128`` (e.g. the sm90 - default ``FlexBwDConfig(64, 128, 128, 64)``) every choice is filtered out and - backward compilation aborts with ``NoValidChoicesError``. This is hit only by - nabla TRAINING — forward-only validation never compiles the backward, which - is why nabla validation worked but training did not. +def flex(query: Tensor, key: Tensor, value: Tensor, **kwargs: Any) -> Tensor: + if query.device.type == "cuda": + return _compiled_flex(query, key, value, **kwargs) + return flex_attention(query, key, value, **kwargs) - Harmless for non-nabla (flash) runs: ``flex_attention`` is only reached from - ``attention_flex`` under ``sparse_params is not None``, so a flash run never - lowers a flex backward and the wrapper installed here is never invoked — it - only ever fires for a configured nabla run. - - Newer torch fixed this upstream by restricting the backward autotune blocks - to ``{32, 64}``. This shim back-ports the guarantee: it wraps the inductor - config provider and prepends an all-64 config whenever none of the returned - configs divide 64. Fully defensive — any structural mismatch (renamed method, - different config type) makes it a silent no-op rather than breaking startup, - and it never removes existing choices. - """ - try: - from torch._inductor.choices import InductorChoices # noqa: PLC0415 — optional, version-dependent internal - except Exception: - return - - original = getattr(InductorChoices, "get_flex_attention_bwd_configs", None) - if original is None or getattr(original, "nabla_patched", False): - return - - def divides_64(conf: Any) -> bool: - try: - return all(64 % b == 0 for b in (conf.block_m1, conf.block_n1, conf.block_m2, conf.block_n2)) - except Exception: - # Unknown config shape — assume compatible so we don't interfere. - return True - - def with_64_blocks(conf: Any) -> Any | None: - fields = {"block_m1": 64, "block_n1": 64, "block_m2": 64, "block_n2": 64} - try: - return dataclasses.replace(conf, **fields) - except Exception: - try: - return conf._replace(**fields) # NamedTuple fallback - except Exception: - return None - - def patched(self: Any, head_dim: int, dtype: Any, device_type: str = "cuda") -> Any: - configs = original(self, head_dim, dtype, device_type) - try: - if configs and not any(divides_64(c) for c in configs): - compat = with_64_blocks(configs[0]) - if compat is not None: - configs = [compat, *configs] - except Exception as exp: - logger.warning("nabla flex bwd-config shim no-op (config introspection failed): %s", exp) - return configs - - patched.nabla_patched = True # type: ignore[attr-defined] - try: - InductorChoices.get_flex_attention_bwd_configs = patched # type: ignore[method-assign] - except Exception as exp: - logger.warning("Could not install nabla flex bwd-config shim: %s", exp) - - -_ensure_nabla_compatible_flex_bwd_configs() - -flex = torch.compile(flex_attention, mode="max-autotune-no-cudagraphs", dynamic=True) # Without FA3 the only varlen entry points available are FA2's # ``flash_attn_varlen_*_func``, whose C++ binding declares ``max_seqlen_q``/ @@ -626,52 +387,24 @@ def patched(self: Any, head_dim: int, dtype: Any, device_type: str = "cuda") -> # and traces cleanly, so this decorator becomes a no-op. -@torch.autocast(device_type="cuda", dtype=torch.float32) def apply_scale_shift_norm(norm: nn.Module, x: Tensor, scale: Tensor, shift: Tensor, idx: Tensor) -> Tensor: - """Apply adaptive normalization with indexed scale and shift.""" - return norm(x) * (scale.index_select(0, idx) + 1.0) + shift.index_select(0, idx) + """Apply adaptive normalization with indexed scale and shift, forcing fp32 math.""" + with torch.autocast(device_type=x.device.type, dtype=torch.float32, enabled=x.device.type == "cuda"): + return norm(x) * (scale.index_select(0, idx) + 1.0) + shift.index_select(0, idx) -@torch.autocast(device_type="cuda", dtype=torch.float32) def apply_gate_sum(x: Tensor, out: Tensor, gate: Tensor, idx: Tensor) -> Tensor: - """Add gated residual output to input using indexed gate values.""" - return x + gate.index_select(0, idx) * out + """Add gated residual output to input using indexed gate values, forcing fp32 math.""" + with torch.autocast(device_type=x.device.type, dtype=torch.float32, enabled=x.device.type == "cuda"): + return x + gate.index_select(0, idx) * out -@torch.autocast(device_type="cuda", dtype=torch.float32) -def apply_scale_shift_norm_spatial(norm: nn.Module, x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: - """Apply adaptive normalization with per-token scale and shift. - - Unlike ``apply_scale_shift_norm`` which uses ``index_select`` for per-sample - params, this function takes spatially-varying (per-token) params directly. - - Args: - norm: Normalization layer (e.g. LayerNorm). - x: Input tensor of shape ``(tokens, dim)``. - scale: Per-token scale of shape ``(tokens, dim)``. - shift: Per-token shift of shape ``(tokens, dim)``. - """ - return norm(x) * (scale + 1.0) + shift - - -@torch.autocast(device_type="cuda", dtype=torch.float32) -def apply_gate_sum_spatial(x: Tensor, out: Tensor, gate: Tensor) -> Tensor: - """Add gated residual output to input using per-token gate values. - - Args: - x: Input tensor of shape ``(tokens, dim)``. - out: Residual output of shape ``(tokens, dim)``. - gate: Per-token gate of shape ``(tokens, dim)``. - """ - return x + gate * out - - -@torch.autocast(device_type="cuda", enabled=False) def apply_rotary(x: Tensor, rope: Tensor) -> Tensor: """Apply rotary position embeddings to input tensor.""" - x_ = x.reshape(*x.shape[:-1], -1, 1, 2).to(torch.float32) - x_out = rope[..., 0] * x_[..., 0] + rope[..., 1] * x_[..., 1] - return x_out.reshape(*x.shape) + with torch.autocast(device_type=x.device.type, enabled=False): + x_ = x.reshape(*x.shape[:-1], -1, 1, 2).to(torch.float32) + x_out = rope[..., 0] * x_[..., 0] + rope[..., 1] * x_[..., 1] + return x_out.reshape(*x.shape) class TimeEmbeddings(nn.Module): @@ -697,7 +430,6 @@ def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0) - self.activation = nn.SiLU() self.out_layer = nn.Linear(time_dim, time_dim, bias=True) - @torch.autocast(device_type="cuda", dtype=torch.float32) def forward(self, time: Tensor) -> tuple[Tensor, Tensor]: """Encode time steps into embeddings. @@ -707,18 +439,13 @@ def forward(self, time: Tensor) -> tuple[Tensor, Tensor]: Returns: Tuple of (time embeddings, index tensor mapping tokens to embeddings). """ - args = torch.outer(time, self.freqs.to(device=time.device)) - time_embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - time_embed = self.out_layer(self.activation(self.in_layer(time_embed))) + with torch.autocast(device_type=time.device.type, dtype=torch.float32, enabled=time.device.type == "cuda"): + args = torch.outer(time, self.freqs.to(device=time.device)) + time_embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + time_embed = self.out_layer(self.activation(self.in_layer(time_embed))) time_embed_idx = torch.arange(time_embed.shape[0], device=time_embed.device, dtype=torch.int32) return time_embed, time_embed_idx - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - self.in_layer.reset_parameters() - self.out_layer.reset_parameters() - self.freqs = get_freqs(self.model_dim // 2, self.max_period) - class TextEmbeddings(nn.Module): """Linear projection with layer normalization for text encoder outputs.""" @@ -746,57 +473,22 @@ def forward(self, text_embed: Tensor) -> Tensor: text_embed = self.in_layer(text_embed) return self.norm(text_embed).type_as(text_embed) - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - self.in_layer.reset_parameters() - self.norm.reset_parameters() - class VisualEmbeddings(nn.Module): """Patchify and project visual tokens into the transformer hidden space.""" - def __init__( # noqa: PLR0913 - self, - visual_dim: int, - model_dim: int, - patch_size: tuple[int, int, int], - *, - use_lq_modulation: bool = False, - lq_channels: int = 16, - time_dim: int = 512, - zero_lq_in_main_path: bool = False, - ) -> None: + def __init__(self, visual_dim: int, model_dim: int, patch_size: tuple[int, int, int]) -> None: """Initialize visual embeddings. Args: visual_dim: Number of input visual channels. model_dim: Output dimension after projection. patch_size: Patch size as (temporal, height, width). - use_lq_modulation: Whether to extract and encode LQ patches. - lq_channels: Number of LQ channels in the input (used only when - ``use_lq_modulation=True``). - time_dim: Output dimension for the LQ encoder (used only when - ``use_lq_modulation=True``). - zero_lq_in_main_path: Zero out LQ and mask channels in the main - ``in_layer`` input so LQ reaches the model only via per-block - ``ModulationLQ``. Only used when ``use_lq_modulation=True``. """ super().__init__() self.patch_size = patch_size - self.use_lq_modulation = use_lq_modulation - self.zero_lq_in_main_path = zero_lq_in_main_path self.in_layer = nn.Linear(math.prod(patch_size) * visual_dim, model_dim) - if use_lq_modulation: - self.lq_channels = lq_channels - # input will be concat [hq(16) | lq(16) | mask(1)] - self.lq_start = visual_dim - lq_channels - 1 - self.lq_end = visual_dim - 1 - self.lq_encoder = LQEncoder( - input_dim=math.prod(patch_size) * lq_channels, - time_dim=time_dim, - ) - def _patchify(self, x: Tensor, visual_cu_seqlens: Tensor) -> tuple[Tensor, Tensor]: """Shared patchification logic for both full input and LQ-only slices. @@ -834,7 +526,7 @@ def _patchify(self, x: Tensor, visual_cu_seqlens: Tensor) -> tuple[Tensor, Tenso visual_cu_seqlens = visual_cu_seqlens // self.patch_size[0] return x, visual_cu_seqlens - def forward(self, x: Tensor, visual_cu_seqlens: Tensor) -> tuple[Tensor, Tensor] | tuple[Tensor, Tensor, Tensor]: + def forward(self, x: Tensor, visual_cu_seqlens: Tensor) -> tuple[Tensor, Tensor]: """Patchify input and project to model dimension. Args: @@ -842,36 +534,10 @@ def forward(self, x: Tensor, visual_cu_seqlens: Tensor) -> tuple[Tensor, Tensor] visual_cu_seqlens: Cumulative sequence lengths for packed sequences. Returns: - When ``use_lq_modulation=False``: ``(projected_patches, cu_seqlens)``. - When ``use_lq_modulation=True``: ``(projected_patches, cu_seqlens, lq_encoded)``. + ``(projected_patches, cu_seqlens)``. """ - if not self.use_lq_modulation: - # [T, H, W, 33] -> [T, H/2, W/2, 132] - # (T x H/2 x W/2) tokens, additive mix of patched HQ/LQ/mask, no cross-terms - patches, cu = self._patchify(x, visual_cu_seqlens) - # Linear(132 -> 1792) - return self.in_layer(patches), cu - - lq_slice = x[..., self.lq_start : self.lq_end] - if self.zero_lq_in_main_path: - # Zero out LQ+mask so in_layer sees [HQ | zeros | 0]; LQ only via ModulationLQ - x = x.clone() - x[..., self.lq_start :] = 0.0 - # Main path: Linear(132 -> 1792) on all 33 channels (LQ zeroed if zero_lq_in_main_path) - patches, cu = self._patchify(x, visual_cu_seqlens.clone()) - # LQ path: [T, H, W, 16] -> [T, H/2, W/2, 64] - # Per-token LQ encodings for bilinear modulation HQ * f(LQ) in decoder blocks - lq_patches, _ = self._patchify(lq_slice, visual_cu_seqlens.clone()) - # LQEncoder(64 -> 512) - lq_encoded = self.lq_encoder(lq_patches) - # Linear(132 -> 1792) - return self.in_layer(patches), cu, lq_encoded - - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - self.in_layer.reset_parameters() - if self.use_lq_modulation: - self.lq_encoder.reset_parameters() + patches, cu = self._patchify(x, visual_cu_seqlens) + return self.in_layer(patches), cu class RoPE1D(nn.Module): @@ -893,7 +559,6 @@ def __init__(self, dim: int, max_pos: int = 1024, max_period: float = 10000.0) - pos = torch.arange(max_pos, dtype=freq.dtype) self.register_buffer("args", torch.outer(pos, freq), persistent=False) - @torch.autocast(device_type="cuda", enabled=False) def forward(self, pos: Tensor) -> Tensor: """Compute rotary embeddings for given positions. @@ -903,16 +568,11 @@ def forward(self, pos: Tensor) -> Tensor: Returns: Rotary embedding matrix with shape ``(*pos.shape, 1, dim//2, 2, 2)``. """ - args = self.args[pos] - rope = torch.stack([torch.cos(args), -torch.sin(args), torch.sin(args), torch.cos(args)], dim=-1) - rope = rope.view(*rope.shape[:-1], 2, 2) - return rope.unsqueeze(-4) - - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - freq = get_freqs(self.dim // 2, self.max_period).to(self.args.device) - pos = torch.arange(self.max_pos, dtype=freq.dtype, device=freq.device) - self.args = torch.outer(pos, freq) + with torch.autocast(device_type=pos.device.type, enabled=False): + args = self.args[pos] + rope = torch.stack([torch.cos(args), -torch.sin(args), torch.sin(args), torch.cos(args)], dim=-1) + rope = rope.view(*rope.shape[:-1], 2, 2) + return rope.unsqueeze(-4) class RoPE3D(nn.Module): @@ -941,7 +601,6 @@ def __init__( pos = torch.arange(ax_max_pos, dtype=freq.dtype) self.register_buffer(f"args_{i}", torch.outer(pos, freq), persistent=False) - @torch.autocast(device_type="cuda", enabled=False) def forward( self, shape: tuple[int, int, int], @@ -958,29 +617,23 @@ def forward( Returns: Rotary embedding matrix broadcast to (duration, height, width, 1, dim//2, 2, 2). """ - duration, height, width = shape - args_t = self.args_0[pos[0]] / 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, - ) - rope = torch.stack([torch.cos(args), -torch.sin(args), torch.sin(args), torch.cos(args)], dim=-1) - rope = rope.view(*rope.shape[:-1], 2, 2) - return rope.unsqueeze(-4) - - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - for i, (axes_dim, ax_max_pos) in enumerate(zip(self.axes_dims, self.max_pos, strict=False)): - freq = get_freqs(axes_dim // 2, self.max_period).to(self.args_0.device) - pos = torch.arange(ax_max_pos, dtype=freq.dtype, device=freq.device) - setattr(self, f"args_{i}", torch.outer(pos, freq)) + with torch.autocast(device_type=pos[0].device.type, enabled=False): + duration, height, width = shape + args_t = self.args_0[pos[0]] / 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, + ) + rope = torch.stack([torch.cos(args), -torch.sin(args), torch.sin(args), torch.cos(args)], dim=-1) + rope = rope.view(*rope.shape[:-1], 2, 2) + return rope.unsqueeze(-4) class Modulation(nn.Module): @@ -1000,7 +653,6 @@ def __init__(self, time_dim: int, model_dim: int, num_params: int) -> None: self.out_layer.weight.data.zero_() self.out_layer.bias.data.zero_() - @torch.autocast(device_type="cuda", dtype=torch.float32) def forward(self, x: Tensor) -> Tensor: """Compute modulation parameters from time embeddings. @@ -1010,97 +662,8 @@ def forward(self, x: Tensor) -> Tensor: Returns: Concatenated modulation parameters. """ - return self.out_layer(self.activation(x)) - - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - self.out_layer.weight.data.zero_() - self.out_layer.bias.data.zero_() - - -class LQEncoder(nn.Module): - """Lightweight MLP that maps patchified LQ tokens to the time embedding space. - - Architecture: Linear -> SiLU -> Linear. Used to produce per-token conditioning - vectors from LQ patches for spatially-varying modulation in decoder blocks. - - Args: - input_dim: Dimension of patchified LQ tokens (patch_volume * lq_channels). - time_dim: Output dimension matching the time embedding space. - """ - - def __init__(self, input_dim: int, time_dim: int) -> None: - """Initialize LQ encoder layers. - - Args: - input_dim: Dimension of patchified LQ tokens (patch_volume * lq_channels). - time_dim: Output dimension matching the time embedding space. - """ - super().__init__() - self.in_layer = nn.Linear(input_dim, time_dim, bias=True) - self.activation = nn.SiLU() - self.out_layer = nn.Linear(time_dim, time_dim, bias=True) - - def forward(self, x: Tensor) -> Tensor: - """Encode patchified LQ tokens. - - Args: - x: Patchified LQ tokens of shape ``(total_tokens, input_dim)``. - - Returns: - Encoded LQ tokens of shape ``(total_tokens, time_dim)``. - """ - return self.out_layer(self.activation(self.in_layer(x))) - - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - self.in_layer.reset_parameters() - self.out_layer.reset_parameters() - - -class ModulationLQ(nn.Module): - """Per-token modulation layer producing scale, shift, and gate from LQ encodings. - - Identical architecture to ``Modulation`` but designed for per-token (spatially-varying) - input from ``LQEncoder`` rather than per-sample time embeddings. - Zero-initialized so the model starts identically to a pretrained checkpoint - without LQ modulation. - - Args: - time_dim: Input dimension of LQ encodings (same as time embedding dim). - model_dim: Per-parameter output dimension. - num_params: Number of modulation parameters to produce. - """ - - def __init__(self, time_dim: int, model_dim: int, num_params: int) -> None: - """Initialize modulation layer with zero weights. - - Args: - time_dim: Input dimension of LQ encodings (same as time embedding dim). - model_dim: Per-parameter output dimension. - num_params: Number of modulation parameters to produce. - """ - 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: - """Compute per-token modulation parameters from LQ encodings. - - Args: - x: LQ encoding tensor of shape ``(total_tokens, time_dim)``. - - Returns: - Concatenated modulation parameters ``(total_tokens, num_params * model_dim)``. - """ - return self.out_layer(self.activation(x)) - - def reset_parameters(self) -> None: - """Re-initialize with zeros for identity behavior.""" - self.out_layer.weight.data.zero_() - self.out_layer.bias.data.zero_() + with torch.autocast(device_type=x.device.type, dtype=torch.float32, enabled=x.device.type == "cuda"): + return self.out_layer(self.activation(x)) class MultiheadSelfAttention(nn.Module, AttentionModuleMixin): @@ -1211,43 +774,19 @@ def attention_flex( # fractal_flatten and is unrelated to the VAE spatial compression. h, w = h // FRACTAL_BLOCK_SIZE, w // FRACTAL_BLOCK_SIZE visual_seqlens = sparse_params["visual_seqlens"].to(device=query.device) - if sparse_params["attention_type"] == "dense_framewise_causal": - block_mask = framewise_causal_dense( - visual_seqlens, - h, - w, - mf=sparse_params["mf"], - ) - elif "mf" not in sparse_params: - block_mask = nablaT_v2_doc( - query, - key, - visual_seqlens, - t, - h, - w, - wT=sparse_params["wT"], - wH=sparse_params["wH"], - wW=sparse_params["wW"], - thr=sparse_params["P"], - add_sta=sparse_params["add_sta"], - method=sparse_params["method"], - ) - else: - block_mask = nablaT_v2_doc_mfcausal( - query, - key, - visual_seqlens, - t, - h, - w, - wT=sparse_params["wT"], - wH=sparse_params["wH"], - wW=sparse_params["wW"], - thr=sparse_params["P"], - add_sta=sparse_params["add_sta"], - mf=sparse_params["mf"], - ) + block_mask = nablaT_v2_doc( + query, + key, + visual_seqlens, + t, + h, + w, + wT=sparse_params["wT"], + wH=sparse_params["wH"], + wW=sparse_params["wW"], + thr=sparse_params["P"], + add_sta=sparse_params["add_sta"], + ) out = ( flex( query, @@ -1294,17 +833,6 @@ def forward( return self.out_layer(out) - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - self.to_query.reset_parameters() - self.to_key.reset_parameters() - self.to_value.reset_parameters() - - self.out_layer.reset_parameters() - - self.query_norm.reset_parameters() - self.key_norm.reset_parameters() - class MultiheadCrossAttention(nn.Module, AttentionModuleMixin): """Multi-head cross-attention with flash attention.""" @@ -1387,17 +915,6 @@ def forward(self, x: Tensor, cond: Tensor, cu_seqlens: Tensor, cond_cu_seqlens: out = self.processor(self, query, key, value, cu_seqlens, cond_cu_seqlens) return self.out_layer(out) - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - self.to_query.reset_parameters() - self.to_key.reset_parameters() - self.to_value.reset_parameters() - - self.out_layer.reset_parameters() - - self.query_norm.reset_parameters() - self.key_norm.reset_parameters() - class FeedForward(nn.Module): """Two-layer feed-forward network with GELU activation.""" @@ -1425,11 +942,6 @@ def forward(self, x: Tensor) -> Tensor: """ return self.out_layer(self.activation(self.in_layer(x))) - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - self.in_layer.reset_parameters() - self.out_layer.reset_parameters() - class OutLayer(nn.Module): """Final output layer that unpatchifies and projects visual tokens to pixel space.""" @@ -1499,15 +1011,6 @@ def forward( x = torch.repeat_interleave(x, idxs, dim=0) return x - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - self.norm.reset_parameters() - self.out_layer.reset_parameters() - self.modulation.reset_parameters() - - -# Diffusion transformer architecture for text-conditioned visual generation. - class TransformerEncoderBlock(nn.Module): """Transformer encoder block with self-attention, feed-forward, and adaptive modulation.""" @@ -1562,26 +1065,9 @@ def forward( out = self.feed_forward(out) return apply_gate_sum(x, out, gate, time_embed_idx).type_as(x) - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - self.text_modulation.reset_parameters() - - self.self_attention_norm.reset_parameters() - self.self_attention.reset_parameters() - - self.feed_forward_norm.reset_parameters() - self.feed_forward.reset_parameters() - class TransformerDecoderBlock(nn.Module): - """Transformer decoder block with self-attention, cross-attention, and feed-forward. - - When ``use_lq_modulation=True``, a parallel ``ModulationLQ`` produces per-token - scale/shift/gate from LQ encodings, added to the per-sample time modulation. - Zero-initialized so that at init the block behaves identically to the baseline. - """ - - LQ_SUBLAYER_ORDER = ("self_attention", "cross_attention", "ffn") + """Transformer decoder block with self-attention, cross-attention, and feed-forward.""" def __init__( self, @@ -1590,8 +1076,6 @@ def __init__( ff_dim: int, head_dim: int, *, - use_lq_modulation: bool = False, - lq_modulation_sublayers: dict[str, bool] | None = None, use_text: bool = True, ) -> None: """Initialize decoder block layers. @@ -1601,27 +1085,15 @@ def __init__( time_dim: Dimension of time step embeddings. ff_dim: Inner dimension of the feed-forward network. head_dim: Per-head dimension for attention layers. - use_lq_modulation: Whether to add per-token LQ modulation. - lq_modulation_sublayers: Which sub-layers to modulate. Keys: - ``self_attention``, ``cross_attention``, ``ffn``. Defaults to all True. use_text: When False, the block has no text cross-attention: no ``cross_attention``/``cross_attention_norm`` modules and the modulation produces 6 params (self-attn + FFN) instead of 9. Used by the text-free SR variant. """ super().__init__() - self.use_lq_modulation = use_lq_modulation self.use_text = use_text self.visual_modulation = Modulation(time_dim, model_dim, 9 if use_text else 6) - if lq_modulation_sublayers is None: - lq_modulation_sublayers = dict.fromkeys(self.LQ_SUBLAYER_ORDER, True) - self.lq_sublayers = lq_modulation_sublayers - self.lq_num_active = sum(lq_modulation_sublayers[k] for k in self.LQ_SUBLAYER_ORDER) - - if use_lq_modulation and self.lq_num_active > 0: - self.lq_modulation = ModulationLQ(time_dim, model_dim, 3 * self.lq_num_active) - self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.self_attention = MultiheadSelfAttention(model_dim, head_dim) @@ -1632,7 +1104,6 @@ def __init__( self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.feed_forward = FeedForward(model_dim, ff_dim) - @torch.compile(dynamic=True) def forward( self, visual_embed: Tensor, @@ -1643,7 +1114,6 @@ def forward( text_cu_seqlens: Tensor, time_embed_idx: Tensor, sparse_params: dict[str, Any] | None, - lq_tokens: Tensor | None = None, ) -> Tensor: """Run self-attention, cross-attention, and feed-forward with modulation. @@ -1656,8 +1126,6 @@ def forward( text_cu_seqlens: Cumulative sequence lengths for text tokens. time_embed_idx: Index mapping visual tokens to their time embeddings. sparse_params: Optional sparse attention parameters. - lq_tokens: Per-token LQ encodings ``(total_tokens, time_dim)`` for - spatially-varying modulation. ``None`` when LQ modulation is off. Returns: Updated visual embeddings. @@ -1668,187 +1136,104 @@ def forward( self_attn_params, ff_params = torch.chunk(self.visual_modulation(time_embed), 2, dim=-1) cross_attn_params = None - has_lq = self.use_lq_modulation and lq_tokens is not None and self.lq_num_active > 0 - lq_sa_shift = lq_sa_scale = lq_sa_gate = None - lq_ca_shift = lq_ca_scale = lq_ca_gate = None - lq_ff_shift = lq_ff_scale = lq_ff_gate = None - if has_lq: - # lq_modulation output: (total_tokens, 3 * num_active * model_dim) - lq_chunks = torch.chunk(self.lq_modulation(lq_tokens), self.lq_num_active, dim=-1) - # Each chunk: (total_tokens, 3 * model_dim) -> (shift, scale, gate) per-token - idx = 0 - if self.lq_sublayers["self_attention"]: - lq_sa_shift, lq_sa_scale, lq_sa_gate = torch.chunk(lq_chunks[idx], 3, dim=-1) - idx += 1 - if self.lq_sublayers["cross_attention"]: - lq_ca_shift, lq_ca_scale, lq_ca_gate = torch.chunk(lq_chunks[idx], 3, dim=-1) - idx += 1 - if self.lq_sublayers["ffn"]: - lq_ff_shift, lq_ff_scale, lq_ff_gate = torch.chunk(lq_chunks[idx], 3, dim=-1) - # --- Self-attention --- - # Norm -> modulate by (scale_t + scale_lq + 1, shift_t + shift_lq) -> SelfAttn -> gate by (gate_t + gate_lq) + # Norm -> modulate by (scale + 1, shift) -> SelfAttn -> gate shift_t, scale_t, gate_t = torch.chunk(self_attn_params, 3, dim=-1) - with torch.autocast(device_type="cuda", dtype=torch.float32): + with torch.autocast( + device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" + ): normed = self.self_attention_norm(visual_embed) scale = scale_t.index_select(0, time_embed_idx) shift = shift_t.index_select(0, time_embed_idx) visual_out = normed * (scale + 1.0) + shift - if lq_sa_scale is not None: - visual_out = visual_out + lq_sa_scale * normed + lq_sa_shift visual_out = visual_out.type_as(visual_embed) visual_out = self.self_attention(visual_out, rope, visual_cu_seqlens, sparse_params) - with torch.autocast(device_type="cuda", dtype=torch.float32): + with torch.autocast( + device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" + ): visual_embed = visual_embed + gate_t.index_select(0, time_embed_idx) * visual_out - if lq_sa_gate is not None: - visual_embed = visual_embed + lq_sa_gate * visual_out # --- Cross-attention --- (skipped when text-free) - # Norm -> modulate by (scale_t + scale_lq + 1, shift_t + shift_lq) -> CrossAttn -> gate by (gate_t + gate_lq) + # Norm -> modulate by (scale + 1, shift) -> CrossAttn -> gate if self.use_text: shift_t, scale_t, gate_t = torch.chunk(cross_attn_params, 3, dim=-1) - with torch.autocast(device_type="cuda", dtype=torch.float32): + with torch.autocast( + device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" + ): normed = self.cross_attention_norm(visual_embed) scale = scale_t.index_select(0, time_embed_idx) shift = shift_t.index_select(0, time_embed_idx) visual_out = normed * (scale + 1.0) + shift - if lq_ca_scale is not None: - visual_out = visual_out + lq_ca_scale * normed + lq_ca_shift visual_out = visual_out.type_as(visual_embed) visual_out = self.cross_attention(visual_out, text_embed, visual_cu_seqlens, text_cu_seqlens) - with torch.autocast(device_type="cuda", dtype=torch.float32): + with torch.autocast( + device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" + ): visual_embed = visual_embed + gate_t.index_select(0, time_embed_idx) * visual_out - if lq_ca_gate is not None: - visual_embed = visual_embed + lq_ca_gate * visual_out # --- Feed-forward --- - # Norm -> modulate by (scale_t + scale_lq + 1, shift_t + shift_lq) -> FF -> gate by (gate_t + gate_lq) + # Norm -> modulate by (scale + 1, shift) -> FF -> gate shift_t, scale_t, gate_t = torch.chunk(ff_params, 3, dim=-1) - with torch.autocast(device_type="cuda", dtype=torch.float32): + with torch.autocast( + device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" + ): normed = self.feed_forward_norm(visual_embed) scale = scale_t.index_select(0, time_embed_idx) shift = shift_t.index_select(0, time_embed_idx) visual_out = normed * (scale + 1.0) + shift - if lq_ff_scale is not None: - visual_out = visual_out + lq_ff_scale * normed + lq_ff_shift visual_out = visual_out.type_as(visual_embed) visual_out = self.feed_forward(visual_out) - with torch.autocast(device_type="cuda", dtype=torch.float32): + with torch.autocast( + device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" + ): visual_embed = visual_embed + gate_t.index_select(0, time_embed_idx) * visual_out - if lq_ff_gate is not None: - visual_embed = visual_embed + lq_ff_gate * visual_out return visual_embed - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - self.visual_modulation.reset_parameters() - if self.use_lq_modulation and self.lq_num_active > 0: - self.lq_modulation.reset_parameters() - - self.self_attention_norm.reset_parameters() - self.self_attention.reset_parameters() - - if self.use_text: - self.cross_attention_norm.reset_parameters() - self.cross_attention.reset_parameters() - - self.feed_forward_norm.reset_parameters() - self.feed_forward.reset_parameters() - - -class VideoAdapter(nn.Module): - """Lighter DiT that conditions a frozen backbone on degraded video input. - - The adapter has the same block design as the backbone but half as many blocks. - LQ latent tokens c = P(z̃) are processed through adapter blocks whose outputs - are injected into alternating backbone blocks via learnable γ_ℓ scalars. - - γ_ℓ are zero-initialized so the adapter has no effect at init, enabling - fine-tuning from a pretrained backbone checkpoint. - """ - - def __init__( - self, - in_visual_dim: int, - model_dim: int, - time_dim: int, - ff_dim: int, - head_dim: int, - patch_size: tuple[int, int, int], - num_blocks: int, - gamma: float = 0.0, - ) -> None: - super().__init__() - self.patch_size = patch_size - self.gamma = gamma - patch_vol = math.prod(patch_size) - self.lq_proj = nn.Linear(patch_vol * in_visual_dim, model_dim) - self.blocks = nn.ModuleList( - [TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim) for _ in range(num_blocks)] - ) - - def forward( - self, - lq_visual: Tensor, - pre_patch_cu_seqlens: Tensor, - time_embed: Tensor, - text_embed: Tensor, - visual_rope: Tensor, - visual_cu_seqlens: Tensor, - text_cu_seqlens: Tensor, - visual_time_embed_idx: Tensor, - sparse_params: dict[str, Any] | None, - to_fractal: bool, - ) -> list[Tensor]: - """Process LQ tokens through adapter blocks, return γ-scaled features.""" - # Patchify LQ latent with the same logic as VisualEmbeddings._patchify. - pt, ph, pw = self.patch_size - if pt > 1: - idxs = torch.ones(lq_visual.shape[0], dtype=torch.int32, device=pre_patch_cu_seqlens.device) - idxs[pre_patch_cu_seqlens[:-1]] += pt - 1 - lq_visual = torch.repeat_interleave(lq_visual, idxs, dim=0) - T, H, W, C = lq_visual.shape - lq_patches = ( - lq_visual.view(T // pt, pt, H // ph, ph, W // pw, pw, C).permute(0, 2, 4, 1, 3, 5, 6).flatten(3, 6) - ) - adapter_embed = self.lq_proj(lq_patches) - if to_fractal: - visual_shape = adapter_embed.shape[:-1] - adapter_embed = local_patching(adapter_embed, visual_shape, (1, 8, 8), dim=0) - adapter_embed = adapter_embed.flatten(0, 1) - else: - adapter_embed = adapter_embed.flatten(0, 2) - features = [] - for block in self.blocks: - adapter_embed = block( - adapter_embed, - text_embed, - time_embed, - visual_rope, - visual_cu_seqlens, - text_cu_seqlens, - visual_time_embed_idx, - sparse_params, - ) - features.append(self.gamma * adapter_embed) - return features - - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters.""" - self.lq_proj.reset_parameters() - for block in self.blocks: - block.reset_parameters() +class Kandinsky6SRTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin): + """3D diffusion transformer with text-conditioned visual generation for Kandinsky 6 super-resolution. -class Kandinsky6SRTransformer3DModel(ModelMixin, ConfigMixin): - """3D diffusion transformer with text-conditioned visual generation. + Unlike `Kandinsky6Transformer3DModel` (the TI2VA transformer), which batches over a padded + `(B, T, H, W, C)` tensor, this transformer packs variable-length visual/text sequences into a single + sequence per axis using `cu_seqlens` (the flash-attention varlen convention). This lets the SR pipeline + batch many spatial tiles per forward pass without padding overhead. Visual tokens are processed through + encoder (text) and decoder (visual) transformer blocks with RoPE, adaptive modulation, and optional + NABLA sparse attention. - Processes visual tokens through encoder (text) and decoder (visual) transformer blocks - with RoPE, adaptive modulation, and optional fractal attention sparsity. + Args: + in_visual_dim (`int`, *optional*, defaults to 4): Number of input visual channels. + in_text_dim (`int`, *optional*, defaults to 3584): Dimension of text encoder hidden states. + in_text_dim2 (`int`, *optional*, defaults to 768): Dimension of pooled text embeddings. + time_dim (`int`, *optional*, defaults to 512): Dimension of time step embeddings. + out_visual_dim (`int`, *optional*, defaults to 4): Number of output visual channels. + patch_size (`tuple[int, int, int]`, *optional*, defaults to `(1, 2, 2)`): Patch size as + (temporal, height, width). + model_dim (`int`, *optional*, defaults to 2048): Hidden dimension of the transformer. + ff_dim (`int`, *optional*, defaults to 5120): Inner dimension of feed-forward networks. + num_text_blocks (`int`, *optional*, defaults to 2): Number of text encoder blocks. + num_visual_blocks (`int`, *optional*, defaults to 32): Number of visual decoder blocks. + axes_dims (`tuple[int, int, int]`, *optional*, defaults to `(16, 24, 24)`): RoPE dimension split + per axis (temporal, height, width). + visual_cond (`bool`, *optional*, defaults to `False`): Whether to use visual conditioning input. + instruct_type (`str`, *optional*): Conditioning mode: `"noise"`, `"channel"`, `"hybrid"`, or + `"hybrid_anchor"` (the mode Kandinsky 6 SR ships with). + attention_params (`dict`, *optional*): Per-resolution attention configuration, keyed by base tile + resolution; each entry selects dense (`type="flash"`) or sparse NABLA (`type="nabla"`) attention. + use_motion_score (`bool`, *optional*, defaults to `False`): Whether to add motion score embeddings + to time conditioning. + use_text (`bool`, *optional*, defaults to `True`): When `False`, builds a text-free model with no + text encoder, text RoPE/blocks, pooled-text conditioning, or per-block cross-attention. + sr_params (`dict`, *optional*): SR-specific configuration: `scheduler_scale`, `lq_noise_scale`, + `lq_noise_type`, `lq_channel_noise_scale`, `cap_noise_timestep`, `fps`, `scale_factor` (RoPE + frequency scaling per base resolution), and `visual_size` (base tile resolution). Read by + [`Kandinsky6SRPipeline`] via the attributes this unpacks into. + attribute_overrides (`dict`, *optional*): Arbitrary extra attributes set on the model after + construction. """ _no_split_modules = ["TransformerEncoderBlock", "TransformerDecoderBlock"] _repeated_blocks = ["TransformerEncoderBlock", "TransformerDecoderBlock"] + _supports_gradient_checkpointing = True @register_to_config def __init__( @@ -1869,54 +1254,10 @@ def __init__( instruct_type: str | None = None, attention_params: dict[str, Any] | None = None, use_motion_score: bool = False, - use_lq_modulation: bool = False, - lq_modulation_sublayers: dict[str, bool] | None = None, - zero_lq_in_main_path: bool = False, - use_adapter: bool = False, - adapter_gamma: float = 0.0, use_text: bool = True, - use_lq_noise_cond: bool = False, + sr_params: Mapping[str, Any] | None = None, attribute_overrides: Mapping[str, Any] | None = None, ) -> None: - """Initialize the 3D diffusion transformer. - - Args: - in_visual_dim: Number of input visual channels. - in_text_dim: Dimension of text encoder hidden states. - in_text_dim2: Dimension of pooled text embeddings. - time_dim: Dimension of time step embeddings. - out_visual_dim: Number of output visual channels. - patch_size: Patch size as (temporal, height, width). - model_dim: Hidden dimension of the transformer. - ff_dim: Inner dimension of feed-forward networks. - num_text_blocks: Number of text encoder blocks. - num_visual_blocks: Number of visual decoder blocks. - axes_dims: RoPE dimension split per axis (temporal, height, width). - visual_cond: Whether to use visual conditioning input. - instruct_type: Conditioning mode (e.g. ``"channel"``). - attention_params: Extra attention configuration. - use_motion_score: Whether to add motion score embeddings to time conditioning. - use_lq_modulation: Whether to enable per-block spatially-varying LQ modulation. - lq_modulation_sublayers: Which sub-layers to modulate with LQ. Keys: - ``self_attention``, ``cross_attention``, ``ffn``. Defaults to all True. - zero_lq_in_main_path: Zero out LQ+mask channels in ``in_layer`` input so - LQ reaches the model only via ``ModulationLQ``. Ablation experiment flag. - use_adapter: Whether to attach a VideoAdapter that conditions the backbone - on the degraded LQ input. The adapter has ``num_visual_blocks // 2`` - blocks and injects γ-scaled features into alternating backbone blocks. - adapter_gamma: Fixed scalar γ (not a trainable parameter) controlling - adapter influence on the backbone. Small values keep the backbone - dominant; larger values allow stronger structural correction. - Only used when ``use_adapter=True``. - use_text: When False, builds a text-free model with no text encoder, - text RoPE/blocks, pooled-text conditioning, or per-block cross- - attention. Default True keeps the standard architecture. - use_lq_noise_cond: When True, adds a second ``TimeEmbeddings`` that - encodes the LQ-endpoint noise fraction ``s`` (``lq_noise_scale``) - and sums it into the adaLN conditioning bus (like motion_score). - Its ``out_layer`` is zero-initialised so a checkpoint without the - module warm-starts as an exact no-op. - """ super().__init__() head_dim = sum(axes_dims) self.in_visual_dim = in_visual_dim @@ -1925,12 +1266,27 @@ def __init__( self.patch_size = patch_size self.visual_cond = visual_cond self.attention_params = attention_params - self.use_lq_modulation = use_lq_modulation - if use_lq_modulation and instruct_type not in ("channel", "hybrid"): - msg = "use_lq_modulation requires instruct_type='channel' or 'hybrid'" - raise ValueError(msg) - self.use_adapter = use_adapter + sr_params = dict(sr_params) if sr_params is not None else {} + self.sr_params = sr_params + self.scheduler_scale = float(sr_params.get("scheduler_scale", 5.0)) + self.lq_noise_scale = float(sr_params.get("lq_noise_scale", 0.7)) + self.lq_noise_type = sr_params.get("lq_noise_type", "ddpm") + self.lq_channel_noise_scale = float(sr_params.get("lq_channel_noise_scale", 0.0)) + self.cap_noise_timestep = bool(sr_params.get("cap_noise_timestep", False)) + self.fps = int(sr_params.get("fps", 24)) + visual_size = sr_params.get("visual_size") + self.visual_size = ( + int(visual_size[0] if isinstance(visual_size, (list, tuple)) else visual_size) + if visual_size is not None + else None + ) + scale_factor = sr_params.get("scale_factor") + self.scale_factor = ( + {int(key): tuple(float(value) for value in values) for key, values in scale_factor.items()} + if scale_factor is not None + else None + ) visual_embed_dim = ( 2 * in_visual_dim + 1 @@ -1940,16 +1296,9 @@ def __init__( self.time_embeddings = TimeEmbeddings(model_dim, time_dim) if use_motion_score: self.motion_embeddings = TimeEmbeddings(model_dim, time_dim) - if use_lq_noise_cond: - self.lq_noise_embeddings = TimeEmbeddings(model_dim, time_dim) - self.zero_init_lq_noise_out() - self.use_motion_score = use_motion_score - self.use_lq_noise_cond = use_lq_noise_cond + self.use_text = use_text - if use_adapter and not use_text: - msg = "use_adapter requires use_text=True: VideoAdapter blocks use text cross-attention" - raise ValueError(msg) if use_text: self.text_embeddings = TextEmbeddings(in_text_dim, model_dim) self.pooled_text_embeddings = TextEmbeddings(in_text_dim2, time_dim) @@ -1959,15 +1308,7 @@ def __init__( # build_text_free_dit.py bakes it into this bias; dropping it instead corrupts # the time conditioning, so the model still needs it (cross-attention does not). self.pooled_bias = nn.Parameter(torch.zeros(time_dim)) - self.visual_embeddings = VisualEmbeddings( - visual_embed_dim, - model_dim, - patch_size, - use_lq_modulation=use_lq_modulation, - lq_channels=in_visual_dim, - time_dim=time_dim, - zero_lq_in_main_path=zero_lq_in_main_path, - ) + self.visual_embeddings = VisualEmbeddings(visual_embed_dim, model_dim, patch_size) if use_text: self.text_rope_embeddings = RoPE1D(head_dim) @@ -1978,35 +1319,17 @@ def __init__( self.visual_rope_embeddings = RoPE3D(axes_dims) self.visual_transformer_blocks = nn.ModuleList( [ - TransformerDecoderBlock( - model_dim, - time_dim, - ff_dim, - head_dim, - use_lq_modulation=use_lq_modulation, - lq_modulation_sublayers=lq_modulation_sublayers, - use_text=use_text, - ) + TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim, use_text=use_text) for _ in range(num_visual_blocks) ] ) - if use_adapter: - self.adapter = VideoAdapter( - in_visual_dim=in_visual_dim, - model_dim=model_dim, - time_dim=time_dim, - ff_dim=ff_dim, - head_dim=head_dim, - patch_size=patch_size, - num_blocks=num_visual_blocks // 2, - gamma=adapter_gamma, - ) - self.out_layer = OutLayer(model_dim, time_dim, out_visual_dim, patch_size) for name, value in (attribute_overrides or {}).items(): setattr(self, str(name), value) + self.gradient_checkpointing = False + def forward( self, x: Tensor, @@ -2020,8 +1343,6 @@ def forward( scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), sparse_params: dict[str, Any] | None = None, motion_score: Tensor | None = None, - lq_latent: Tensor | None = None, - lq_noise_level: Tensor | None = None, ) -> Tensor: """Run the full diffusion transformer forward pass. @@ -2037,49 +1358,34 @@ def forward( scale_factor: RoPE frequency scaling per axis. sparse_params: Optional sparse/fractal attention parameters. motion_score: Optional motion score for temporal conditioning. - lq_latent: Degraded video latent passed exclusively to the - adapter when ``use_adapter=True``. Not noised. ``None`` when adapter - is disabled. - lq_noise_level: Per-sequence LQ-endpoint noise fraction ``s`` in - ``[0, 1]`` (shape ``[N_seqs]`` or broadcastable ``[1]``). Only - used when ``use_lq_noise_cond=True``; embedded like the timestep - (scaled x1000) and summed into the adaLN conditioning bus. Returns: Denoised visual output tensor. """ + checkpoint = torch.is_grad_enabled() and self.gradient_checkpointing + text_embed = self.text_embeddings(text_embed) if self.use_text else None time_embed, time_embed_idx = self.time_embeddings(time) if motion_score is not None and self.use_motion_score: ms_embed, _ = self.motion_embeddings(motion_score) time_embed = time_embed + ms_embed - if lq_noise_level is not None and self.use_lq_noise_cond: - # Same input scale as the timestep (t*1000 at every call site). - nl_embed, _ = self.lq_noise_embeddings(1000.0 * lq_noise_level) - time_embed = time_embed + nl_embed if self.use_text: time_embed = time_embed + self.pooled_text_embeddings(pooled_text_embed) else: time_embed = time_embed + self.pooled_bias - # Save cu_seqlens before patchification (adapter needs it for its own _patchify) - if self.use_adapter: - pre_patch_cu_seqlens = visual_cu_seqlens - - ve_result = self.visual_embeddings(x, visual_cu_seqlens) - if self.use_lq_modulation: - visual_embed, visual_cu_seqlens, lq_tokens = ve_result - else: - visual_embed, visual_cu_seqlens = ve_result - lq_tokens = None + visual_embed, visual_cu_seqlens = self.visual_embeddings(x, visual_cu_seqlens) if self.use_text: text_rope = self.text_rope_embeddings(text_rope_pos) text_time_embed_idx = time_embed_idx.repeat_interleave(torch.diff(text_cu_seqlens), dim=0) for text_transformer_block in self.text_transformer_blocks: - text_embed = text_transformer_block( - text_embed, time_embed, text_rope, text_cu_seqlens, text_time_embed_idx + args = (text_embed, time_embed, text_rope, text_cu_seqlens, text_time_embed_idx) + text_embed = ( + self._gradient_checkpointing_func(text_transformer_block, *args) + if checkpoint + else text_transformer_block(*args) ) visual_shape = visual_embed.shape[:-1] @@ -2089,41 +1395,10 @@ def forward( visual_embed, visual_rope, visual_cu_seqlens = fractal_flatten( visual_embed, visual_rope, visual_cu_seqlens, visual_shape, fractal=to_fractal ) - if lq_tokens is not None: - # Flatten lq_tokens the same way as visual_embed (spatial dims only). - # fractal_flatten already modified visual_cu_seqlens, so we inline the - # flatten here rather than calling fractal_flatten again. - if to_fractal: - # [T, 32, 32, dim] -> [T * (32/8)*(32/8), 8*8, dim] -> [T*16, 64, dim] - lq_tokens = local_patching(lq_tokens, visual_shape, (1, 8, 8), dim=0) - lq_tokens = lq_tokens.flatten(0, 1) - else: - # [T, 32, 32, 512] -> [T*32*32, 512] - lq_tokens = lq_tokens.flatten(0, 2) - visual_time_embed_idx = time_embed_idx.repeat_interleave(torch.diff(visual_cu_seqlens), dim=0) - # Run adapter on LQ tokens and collect γ-scaled features for injection - adapter_features: list[Tensor] | None = None - if self.use_adapter: - adapter_features = self.adapter( - lq_latent, - pre_patch_cu_seqlens, - time_embed, - text_embed, - visual_rope, - visual_cu_seqlens, - text_cu_seqlens, - visual_time_embed_idx, - sparse_params, - to_fractal, - ) - - for i, visual_transformer_block in enumerate(self.visual_transformer_blocks): - # Inject adapter features before alternating backbone blocks (0, 2, 4, ...) - if adapter_features is not None and i % 2 == 0: - visual_embed = visual_embed + adapter_features[i // 2] - visual_embed = visual_transformer_block( + for visual_transformer_block in self.visual_transformer_blocks: + args = ( visual_embed, text_embed, time_embed, @@ -2132,7 +1407,11 @@ def forward( text_cu_seqlens, visual_time_embed_idx, sparse_params, - lq_tokens, + ) + visual_embed = ( + self._gradient_checkpointing_func(visual_transformer_block, *args) + if checkpoint + else visual_transformer_block(*args) ) visual_embed, visual_cu_seqlens = fractal_unflatten( visual_embed, visual_cu_seqlens, visual_shape, fractal=to_fractal @@ -2141,63 +1420,5 @@ def forward( visual_time_embed_idx = time_embed_idx.repeat_interleave(torch.diff(visual_cu_seqlens), dim=0) return self.out_layer(visual_embed, text_embed, time_embed, visual_cu_seqlens, visual_time_embed_idx) - def enable_explicit_forward_prefetch(self, factor: int = 1) -> None: - """Enable FSDP forward prefetching for visual transformer blocks. - - Args: - factor: Number of blocks to prefetch ahead. Skipped if <= 0 or already enabled. - - Raises: - RuntimeError: If the model is not wrapped with FSDP. - """ - if factor <= 0 or getattr(self, "forward_prefetch", False): - return - for i, layer in enumerate(self.visual_transformer_blocks): - if not hasattr(layer, "set_modules_to_forward_prefetch"): - msg = "wrap DiT with FSDP first to use forward_prefetch" - raise RuntimeError(msg) - - start, end = i, i + factor - prefetch_modules = self.visual_transformer_blocks[start:end] - layer.set_modules_to_forward_prefetch(prefetch_modules) - self.forward_prefetch = True - - def zero_init_lq_noise_out(self) -> None: - """Zero the s-embedder output layer so its contribution starts at exactly 0. - - Keeps warm-start from a checkpoint without ``lq_noise_embeddings`` an - exact no-op (adaLN-zero pattern). Must be re-applied after any - ``reset_parameters`` (FSDP meta-init materialisation re-randomises it). - """ - nn.init.zeros_(self.lq_noise_embeddings.out_layer.weight) - nn.init.zeros_(self.lq_noise_embeddings.out_layer.bias) - - def reset_parameters(self) -> None: - """Re-initialize all learnable parameters across every sub-module.""" - self.time_embeddings.reset_parameters() - if self.use_lq_noise_cond: - self.lq_noise_embeddings.reset_parameters() - self.zero_init_lq_noise_out() - if self.use_text: - self.text_embeddings.reset_parameters() - self.pooled_text_embeddings.reset_parameters() - else: - self.pooled_bias.data.zero_() - self.visual_embeddings.reset_parameters() - - if self.use_text: - self.text_rope_embeddings.reset_parameters() - for module in self.text_transformer_blocks: - module.reset_parameters() - - self.visual_rope_embeddings.reset_parameters() - for module in self.visual_transformer_blocks: - module.reset_parameters() - - self.out_layer.reset_parameters() - - if self.use_adapter: - self.adapter.reset_parameters() - __all__ = ["Kandinsky6SRTransformer3DModel"] diff --git a/src/diffusers/pipelines/kandinsky6/latent_upscaler.py b/src/diffusers/pipelines/kandinsky6/latent_upscaler.py index 332375e50bfb..50ef471d7f4e 100644 --- a/src/diffusers/pipelines/kandinsky6/latent_upscaler.py +++ b/src/diffusers/pipelines/kandinsky6/latent_upscaler.py @@ -14,81 +14,32 @@ """Kandinsky 6 SR latent-upscaler bank Diffusers component.""" -from __future__ import annotations - import copy -import functools -import importlib -import math -from collections import defaultdict -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass, replace -from functools import partial -from types import SimpleNamespace -from typing import TYPE_CHECKING, Any, Literal, NamedTuple +from typing import Any, Literal, NamedTuple import torch -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin from torch import Tensor, nn -from torch.autograd import Function from torch.nn import functional +from ...configuration_utils import ConfigMixin, register_to_config +from ...models.modeling_utils import ModelMixin + -if TYPE_CHECKING: - from collections.abc import Callable - from types import ModuleType - - -def rearrange(*args: Any, **kwargs: Any) -> Any: - """Load einops only when the latent upscaler is actually executed.""" - try: - from einops import rearrange as _rearrange # noqa: PLC0415 - except (ImportError, OSError) as exc: - raise RuntimeError("SR latent upscaler requires einops. Install with `pip install einops`.") from exc - return _rearrange(*args, **kwargs) - - -def _model_config(value: Mapping[str, Any]) -> Any: - """Convert serialized model mappings to attribute-accessible config objects.""" - if not isinstance(value, Mapping): - raise TypeError(f"latent upscaler model config must be a mapping, got {type(value).__name__}") - - def convert(item: Any) -> Any: - if isinstance(item, Mapping): - return SimpleNamespace(**{str(key): convert(val) for key, val in item.items()}) - if isinstance(item, list): - return tuple(convert(val) for val in item) - return item - - config = convert(value) - if getattr(config, "architecture", None) == "multi_scale": - # Older/published multi-scale configs omit fields whose native - # constructor defaults are sufficient. - defaults = { - "icnr": False, - "global_skip": False, - "x2_adapter_blocks": 0, - "x2_tail_mode": "shared", - "x2_adapter_sources": None, - "x2_finisher": "none", - "motion_attention": None, - } - for key, default in defaults.items(): - if not hasattr(config, key): - setattr(config, key, default) - return config +def merge_batch_time(x: Tensor) -> Tensor: + """Merge the temporal axis into batch: ``(B, C, T, H, W) -> (B*T, C, H, W)``.""" + b, c, t, h, w = x.shape + return x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) -UpsampleMode = str -UpsamplePosition = str -FlatModelConfig = Any -MultiScaleModelConfig = Any +def split_batch_time(x: Tensor, b: int, t: int) -> Tensor: + """Inverse of :func:`merge_batch_time`: ``(B*T, C, H, W) -> (B, C, T, H, W)``.""" + _, c, h, w = x.shape + return x.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4) -@torch.autocast(device_type="cuda", enabled=False) -def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: - return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) +UpsampleMode = str # Convolution primitives shared by the latent upsampler architectures. @@ -434,7 +385,9 @@ def __init__( kernel_size=kernel_size, zq_dim=zq_dim, ) - self.shortcut = nn.Identity() if in_channels == out_channels else conv(in_channels, out_channels, kernel_size=1) + self.shortcut = ( + nn.Identity() if in_channels == out_channels else conv(in_channels, out_channels, kernel_size=1) + ) self.layer_scale: nn.Module = ( LayerScale(out_channels, init_value=layer_scale_init, dims=dims) if layer_scale_init is not None @@ -579,135 +532,6 @@ def zero_init_residual(self) -> None: nn.init.zeros_(final_conv.bias) # type: ignore[arg-type] -# Exact merging of flash-style attention branches with a closed-form null branch. - - -def merge_math( - outputs: Sequence[Tensor], - lses: Sequence[Tensor], - null_value: Tensor, - null_logits: Tensor, -) -> tuple[Tensor, Tensor]: - """Combine attention branches and the null token into one softmax, in fp32. - - Args: - outputs: Branch outputs shaped ``(batch, seq, heads, dim)``. - lses: Branch log-sum-exps shaped ``(batch, seq, heads)``. - null_value: Null-token value shaped ``(heads, dim)``. - null_logits: Null-token logits shaped ``(batch, seq, heads)``. - - Returns: - Merged output ``(batch, seq, heads, dim)`` and total lse ``(batch, seq, heads)``. - """ - lse_stack = torch.stack([*[lse.float() for lse in lses], null_logits.float()], dim=0) - total_lse = torch.logsumexp(lse_stack, dim=0) - weights = torch.exp(lse_stack - total_lse) - merged = weights[-1].unsqueeze(-1) * null_value.float() - for weight, output in zip(weights[:-1], outputs, strict=True): - merged = merged + weight.unsqueeze(-1) * output.float() - return merged, total_lse - - -@functools.cache -def compiled_merge_math() -> Callable[..., tuple[Tensor, Tensor]]: - """Compile the merge math once and reuse it across calls.""" - return torch.compile(merge_math, fullgraph=True) - - -class MergeWithNullBranch(Function): - """Merge kernel attention branches with the null branch, with exact gradients. - - Flash-style kernels (NATTEN ``na2d``) return ``(output, lse)`` but their - backward ignores incoming lse gradients, so plain autograd through a softmax - merge silently drops part of their q/k/v gradients. This op follows the - ring-attention convention instead: each kernel branch receives the full - upstream gradient while its saved output/lse tensors are overwritten in - place with the merged values, letting the kernel backward reconstruct exact - global gradients. The null branch is single-key attention, so its gradients - have a closed form and are computed here directly. - - The in-place overwrite means branch tensors must not be read after backward, - and the op is incompatible with gradient recomputation (activation - checkpointing) and double backward. - """ - - @staticmethod - def forward(ctx: Any, *args: Any) -> Tensor: - """Merge branches; see :func:`merge_attention_branches` for the layout of ``args``.""" - null_value, null_logits, num_branches, use_compile, *branch_tensors = args - outputs = branch_tensors[:num_branches] - lses = branch_tensors[num_branches:] - math_fn = compiled_merge_math() if use_compile else merge_math - merged, total_lse = math_fn(outputs, lses, null_value, null_logits) - ctx.num_branches = num_branches - ctx.save_for_backward(null_value, null_logits, merged, total_lse, *branch_tensors) - target = outputs[0] if outputs else null_value - return merged.to(target.dtype) - - @staticmethod - def backward(ctx: Any, grad_out: Tensor) -> tuple[Tensor | None, ...]: - """Route full gradients to kernel branches and closed-form ones to the null pair.""" - num_branches = ctx.num_branches - null_value, null_logits, merged, total_lse, *branch_tensors = ctx.saved_tensors - outputs = branch_tensors[:num_branches] - lses = branch_tensors[num_branches:] - grad = grad_out.float() - null_weight = torch.exp(null_logits.float() - total_lse) - d_null_value = torch.einsum("bsh,bshd->hd", null_weight, grad) - d_null_logits = null_weight * ( - torch.einsum("bshd,hd->bsh", grad, null_value.float()) - (grad * merged).sum(dim=-1) - ) - for output, lse in zip(outputs, lses, strict=True): - output.data.copy_(merged.to(output.dtype)) - lse.data.copy_(total_lse.to(lse.dtype)) - return ( - d_null_value.to(null_value.dtype), - d_null_logits.to(null_logits.dtype), - None, - None, - *(grad_out for _ in outputs), - *(None for _ in lses), - ) - - -def merge_attention_branches( - outputs: Sequence[Tensor], - lses: Sequence[Tensor], - null_value: Tensor, - null_logits: Tensor, - *, - torch_compile: bool = False, -) -> Tensor: - """Merge kernel attention branches with the closed-form null branch. - - Args: - outputs: Kernel branch outputs shaped ``(batch, seq, heads, dim)``; may be empty. - lses: Matching branch log-sum-exps shaped ``(batch, seq, heads)``. - null_value: Null-token value shaped ``(heads, dim)``. - null_logits: Null-token logits shaped ``(batch, seq, heads)``. - torch_compile: Compile the merge math with ``torch.compile``. - - Returns: - Merged attention output shaped ``(batch, seq, heads, dim)`` in the branch dtype - (or the null-value dtype when no branches are given). - """ - return MergeWithNullBranch.apply(null_value, null_logits, len(outputs), torch_compile, *outputs, *lses) - - -# Runtime helpers shared by latent-upscaler model components. - - -def forward_with_checkpointing( - module: Callable[..., Tensor], - *inputs: Tensor, - use_checkpointing: bool = False, -) -> Tensor: - """Run a callable directly or through non-reentrant activation checkpointing.""" - if use_checkpointing: - return torch.utils.checkpoint.checkpoint(module, *inputs, use_reentrant=False) - return module(*inputs) - - # Spatial upsampling operations: pixel-shuffle, bilinear, and K-VAE PXS v2. @@ -762,9 +586,9 @@ def spatial_nearest_2x(x: torch.Tensor, dims: int, factor: int) -> torch.Tensor: """ if dims == DIMS_3: b, _c, t, _h, _w = x.shape - x = rearrange(x, "b c t h w -> (b t) c h w") + x = merge_batch_time(x) x = functional.interpolate(x, scale_factor=factor, mode="nearest") - return rearrange(x, "(b t) c h w -> b c t h w", b=b, t=t) + return split_batch_time(x, b, t) return functional.interpolate(x, scale_factor=factor, mode="nearest") @@ -824,11 +648,17 @@ def __init__( nn.init.zeros_(self.conv.bias) def forward(self, x: torch.Tensor) -> torch.Tensor: - """Expand channels via convolution, then rearrange into spatial dimensions.""" + """Expand channels via convolution, then pixel-shuffle into spatial dimensions.""" x = self.conv(x) if self.dims == DIMS_3: - return rearrange(x, "b (c p1 p2) t h w -> b c t (h p1) (w p2)", p1=self.factor, p2=self.factor) - return rearrange(x, "b (c p1 p2) h w -> b c (h p1) (w p2)", p1=self.factor, p2=self.factor) + b, c_by_factor2, t, h, w = x.shape + c = c_by_factor2 // self.factor**2 + return ( + x.reshape(b, c, self.factor, self.factor, t, h, w) + .permute(0, 1, 4, 5, 2, 6, 3) + .reshape(b, c, t, h * self.factor, w * self.factor) + ) + return functional.pixel_shuffle(x, self.factor) class BilinearUpsampleND(nn.Module): @@ -856,9 +686,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: """Apply bilinear interpolation to spatial dimensions.""" if self.dims == DIMS_3: b, _c, t, _h, _w = x.shape - x = rearrange(x, "b c t h w -> (b t) c h w") + x = merge_batch_time(x) x = functional.interpolate(x, scale_factor=self.factor, mode="bilinear", align_corners=False) - return rearrange(x, "(b t) c h w -> b c t h w", b=b, t=t) + return split_batch_time(x, b, t) return functional.interpolate(x, scale_factor=self.factor, mode="bilinear", align_corners=False) @@ -992,307 +822,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.linear(self.pxs(x) + self.v3(x)) -# Motion-correspondence attention for latent video features. - - -MotionAttentionBackend = Literal["sdpa", "natten"] -NattenBackend = Literal["cutlass-fna", "hopper-fna", "flex-fna"] - - -@dataclass -class MotionCorrespondenceSpec: - """Construction parameters for one motion-correspondence block.""" - - channels: int - num_heads: int - head_dim: int - spatial_kernel_size: int - temporal_offsets: tuple[int, ...] - backend: MotionAttentionBackend - natten_backend: NattenBackend = "hopper-fna" - merge_compile: bool = False - - -def temporal_neighbor_indices( - query_index: int, - num_frames: int, - temporal_offsets: tuple[int, ...], -) -> tuple[int, ...]: - """Return valid neighbor-frame indices without shifting or duplication. - - Args: - query_index: Query frame index. - num_frames: Number of frames in the clip. - temporal_offsets: Allowed non-zero offsets from the query frame. - - Returns: - Valid neighbor indices in the configured offset order. - """ - return tuple(query_index + offset for offset in temporal_offsets if 0 <= query_index + offset < num_frames) - - -def spatial_neighborhood_mask( - height: int, - width: int, - kernel_size: int, - device: torch.device, -) -> Tensor: - """Build the shifted sliding-window mask used by neighborhood attention. - - Args: - height: Feature-map height. - width: Feature-map width. - kernel_size: Odd spatial neighborhood size. - device: Device for the returned mask. - - Returns: - Boolean mask shaped ``(height * width, height * width)``. - """ - if kernel_size > min(height, width): - msg = f"spatial kernel {kernel_size} exceeds feature map {height}x{width}" - raise ValueError(msg) - - query_h, query_w = torch.meshgrid( - torch.arange(height, device=device), - torch.arange(width, device=device), - indexing="ij", - ) - key_h, key_w = query_h.flatten(), query_w.flatten() - radius = kernel_size // 2 - start_h = (query_h.flatten() - radius).clamp(min=0, max=height - kernel_size) - start_w = (query_w.flatten() - radius).clamp(min=0, max=width - kernel_size) - inside_h = (key_h[None, :] >= start_h[:, None]) & (key_h[None, :] < start_h[:, None] + kernel_size) - inside_w = (key_w[None, :] >= start_w[:, None]) & (key_w[None, :] < start_w[:, None] + kernel_size) - return inside_h & inside_w - - -def split_rope_dims(head_dim: int) -> tuple[int, int, int]: - """Split a head dimension into temporal, height, and width chunks. - - Leftover rotary pairs go to the spatial axes: correspondence is a spatial - matching task, while temporal positions only span the short offset window. - """ - if head_dim < 6 or head_dim % 2 != 0: - msg = f"3D RoPE requires an even head_dim >= 6, got {head_dim}" - raise ValueError(msg) - pairs, remainder = divmod(head_dim // 2, 3) - pair_counts = tuple(pairs + int(axis >= 3 - remainder) for axis in range(3)) - return 2 * pair_counts[0], 2 * pair_counts[1], 2 * pair_counts[2] - - -class AxialRoPE3D(nn.Module): - """Axial rotary embeddings for ``(time, height, width)`` coordinates.""" - - def __init__(self, head_dim: int, max_period: float = 10000.0) -> None: - """Initialize per-axis rotary frequencies.""" - super().__init__() - self.axis_dims = split_rope_dims(head_dim) - for axis_name, axis_dim in zip(("time", "height", "width"), self.axis_dims, strict=True): - self.register_buffer( - f"{axis_name}_frequencies", - get_freqs(axis_dim // 2, max_period), - persistent=False, - ) - - @torch.autocast(device_type="cuda", enabled=False) - def forward(self, x: Tensor) -> Tensor: - """Apply rotary embeddings and return the original working dtype.""" - working_dtype = x.dtype - chunks = x.to(torch.float32).split(self.axis_dims, dim=-1) - axes = (1, 2, 3) - frequencies = (self.time_frequencies, self.height_frequencies, self.width_frequencies) - rotated = [ - self.rotate_axis(chunk, axis, frequency) - for chunk, axis, frequency in zip(chunks, axes, frequencies, strict=True) - ] - return torch.cat(rotated, dim=-1).to(working_dtype) - - @staticmethod - def rotate_axis(x: Tensor, axis: int, frequencies: Tensor) -> Tensor: - """Rotate one channel chunk according to positions along one axis.""" - positions = torch.arange(x.shape[axis], device=x.device, dtype=torch.float32) - angles = torch.outer(positions, frequencies.to(device=x.device)) - broadcast_shape = [1] * x.ndim - broadcast_shape[axis] = positions.numel() - broadcast_shape[-1] = frequencies.numel() - cosine = angles.cos().reshape(broadcast_shape) - sine = angles.sin().reshape(broadcast_shape) - pairs = x.reshape(*x.shape[:-1], -1, 2) - first = pairs[..., 0] * cosine - pairs[..., 1] * sine - second = pairs[..., 0] * sine + pairs[..., 1] * cosine - return torch.stack((first, second), dim=-1).flatten(-2) - - -class MotionCorrespondenceBlock(nn.Module): - """Cross-frame spatial correspondence with an explicit no-match token.""" - - def __init__(self, spec: MotionCorrespondenceSpec) -> None: - """Initialize a motion-correspondence residual block.""" - super().__init__() - if spec.channels != spec.num_heads * spec.head_dim: - msg = f"channels must equal num_heads * head_dim, got {spec.channels} != {spec.num_heads} * {spec.head_dim}" - raise ValueError(msg) - if spec.spatial_kernel_size < 3 or spec.spatial_kernel_size % 2 == 0: - msg = "spatial_kernel_size must be odd and at least 3" - raise ValueError(msg) - if ( - not spec.temporal_offsets - or 0 in spec.temporal_offsets - or tuple(sorted(set(spec.temporal_offsets))) != spec.temporal_offsets - or set(spec.temporal_offsets) != {-offset for offset in spec.temporal_offsets} - ): - msg = "temporal_offsets must be non-zero, unique, sorted, and symmetric" - raise ValueError(msg) - self.channels = spec.channels - self.num_heads = spec.num_heads - self.head_dim = spec.head_dim - self.spatial_kernel_size = spec.spatial_kernel_size - self.temporal_offsets = spec.temporal_offsets - self.backend = spec.backend - self.natten_backend = spec.natten_backend - self.merge_compile = spec.merge_compile - - self.norm = nn.RMSNorm(spec.channels) - self.to_qkv = nn.Linear(spec.channels, 3 * spec.channels) - self.q_norm = nn.RMSNorm(spec.head_dim) - self.k_norm = nn.RMSNorm(spec.head_dim) - self.rope = AxialRoPE3D(spec.head_dim) - self.null_key = nn.Parameter(torch.randn(spec.num_heads, spec.head_dim) * 0.02) - self.null_value = nn.Parameter(torch.zeros(spec.num_heads, spec.head_dim)) - self.to_out = nn.Linear(spec.channels, spec.channels) - nn.init.zeros_(self.to_out.weight) - nn.init.zeros_(self.to_out.bias) - - def forward(self, x: Tensor) -> Tensor: - """Apply cross-frame correspondence to ``(B,C,T,H,W)`` features.""" - if x.shape[2] == 1: - return x - features = rearrange(x, "b c t h w -> b t h w c") - # fp32 norm inputs keep fused rms_norm kernels under bf16 autocast (same as DiT QK-norm). - qkv = self.to_qkv(self.norm(features.float()).type_as(features)) - q, k, v = [ - rearrange(part, "b t h w (nh d) -> b t h w nh d", nh=self.num_heads) for part in qkv.chunk(3, dim=-1) - ] - q_unrotated = self.q_norm(q.float()).type_as(q) - q = self.rope(q_unrotated) - k = self.rope(self.k_norm(k.float()).type_as(k)) - attention = self.compute_sdpa_attention if self.backend == "sdpa" else self.compute_natten_attention - context = attention(q, k, v, q_unrotated) - output = self.to_out(rearrange(context, "b t h w nh d -> b t h w (nh d)")) - return x + rearrange(output, "b t h w c -> b c t h w") - - def null_match_logits(self, q_unrotated: Tensor) -> Tensor: - """Score queries against the null token in the shared attention scale. - - Uses pre-RoPE queries: the null token has no position, so pairing it with - a rotated query would make the no-match threshold depend on where the - query sits in the clip. - """ - logits = torch.einsum("...nd,nd->...n", q_unrotated.float(), self.null_key.float()) - return logits * self.head_dim**-0.5 - - def compute_sdpa_attention(self, q: Tensor, k: Tensor, v: Tensor, q_unrotated: Tensor) -> Tensor: - """Compute the reference operation with PyTorch SDPA.""" - batch, frames, height, width, _heads, dim = q.shape - spatial_mask = spatial_neighborhood_mask(height, width, self.spatial_kernel_size, q.device) - null_logits = self.null_match_logits(q_unrotated) - frame_outputs: list[Tensor] = [] - for query_index in range(frames): - neighbor_indices = temporal_neighbor_indices(query_index, frames, self.temporal_offsets) - q_frame = rearrange(q[:, query_index], "b h w nh d -> b nh (h w) d") - k_frames = rearrange(k[:, neighbor_indices], "b t h w nh d -> b nh (t h w) d") - v_frames = rearrange(v[:, neighbor_indices], "b t h w nh d -> b nh (t h w) d") - # The null column carries a zero key, so its logit comes solely from the - # additive bias; SDPA applies the bias after q·k scaling, hence pre-scaled logits. - null_key = torch.zeros(batch, self.num_heads, 1, dim, dtype=q.dtype, device=q.device) - null_value = self.null_value.to(q)[None, :, None, :].expand(batch, -1, 1, -1) - keys = torch.cat((k_frames, null_key), dim=2) - values = torch.cat((v_frames, null_value), dim=2) - bias = self.sdpa_attention_bias(spatial_mask, null_logits[:, query_index], len(neighbor_indices)) - attended = functional.scaled_dot_product_attention(q_frame, keys, values, attn_mask=bias.to(q.dtype)) - frame_outputs.append(rearrange(attended, "b nh (h w) d -> b h w nh d", h=height, w=width)) - return torch.stack(frame_outputs, dim=1) - - def sdpa_attention_bias(self, spatial_mask: Tensor, frame_null_logits: Tensor, num_neighbors: int) -> Tensor: - """Build the additive bias: window gating for real keys, null logits in the last column.""" - window_bias = torch.where(spatial_mask.repeat(1, num_neighbors), 0.0, float("-inf")) - null_bias = rearrange(frame_null_logits, "b h w nh -> b nh (h w) 1") - window_bias = window_bias[None, None].expand(*null_bias.shape[:2], -1, -1) - return torch.cat((window_bias, null_bias), dim=-1) - - def compute_natten_attention(self, q: Tensor, k: Tensor, v: Tensor, q_unrotated: Tensor) -> Tensor: - """Compute correspondence with fused 2D neighborhood-attention kernels.""" - if q.device.type != "cuda": - msg = "NATTEN motion attention requires CUDA" - raise RuntimeError(msg) - natten = load_natten() - batch, frames, height, width, _heads, _dim = q.shape - if self.spatial_kernel_size > min(height, width): - msg = f"spatial kernel {self.spatial_kernel_size} exceeds feature map {height}x{width}" - raise ValueError(msg) - null_logits = self.null_match_logits(q_unrotated) - - grouped_queries: dict[int, list[tuple[int, tuple[int, ...]]]] = defaultdict(list) - for query_index in range(frames): - neighbors = temporal_neighbor_indices(query_index, frames, self.temporal_offsets) - grouped_queries[len(neighbors)].append((query_index, neighbors)) - - frame_outputs: list[Tensor | None] = [None] * frames - for group in grouped_queries.values(): - query_indices = tuple(query_index for query_index, _neighbors in group) - q_group = rearrange(q[:, query_indices], "b t h w nh d -> (b t) h w nh d").contiguous() - group_outputs: list[Tensor] = [] - group_lse: list[Tensor] = [] - for neighbor_position in range(len(group[0][1])): - neighbor_indices = tuple(neighbors[neighbor_position] for _query_index, neighbors in group) - k_group = rearrange(k[:, neighbor_indices], "b t h w nh d -> (b t) h w nh d").contiguous() - v_group = rearrange(v[:, neighbor_indices], "b t h w nh d -> (b t) h w nh d").contiguous() - output, lse = natten.na2d( - q_group, - k_group, - v_group, - kernel_size=(self.spatial_kernel_size, self.spatial_kernel_size), - backend=self.natten_backend, - return_lse=True, - ) - group_outputs.append(output.flatten(1, 2)) - group_lse.append(lse.flatten(1, 2)) - merged = merge_attention_branches( - group_outputs, - group_lse, - self.null_value.to(q), - rearrange(null_logits[:, query_indices], "b t h w nh -> (b t) (h w) nh"), - torch_compile=self.merge_compile, - ) - merged = rearrange( - merged, - "(b t) (h w) nh d -> b t h w nh d", - b=batch, - t=len(query_indices), - h=height, - w=width, - ) - for group_index, query_index in enumerate(query_indices): - frame_outputs[query_index] = merged[:, group_index] - if any(output is None for output in frame_outputs): - msg = "NATTEN correspondence did not produce every query frame" - raise RuntimeError(msg) - return torch.stack([output for output in frame_outputs if output is not None], dim=1) - - -def load_natten() -> ModuleType: - """Import NATTEN lazily and require its fused CUDA extension.""" - try: - natten = importlib.import_module("natten") - except ImportError as error: - msg = "NATTEN backend requested, but the natten package is not installed" - raise RuntimeError(msg) from error - if not getattr(natten, "HAS_LIBNATTEN", False): - msg = "NATTEN backend requested, but fused libnatten kernels are unavailable" - raise RuntimeError(msg) - return natten - - # Construction and execution helpers for the cascaded multi-scale model. @@ -1332,20 +861,10 @@ class ResidualStackSpec: temporal_padding: TemporalPadding = "zeros" -@dataclass -class MotionAttentionSpec: - """Construction and placement contract for motion-correspondence blocks.""" - - after_mid_blocks: tuple[int, ...] - block: MotionCorrespondenceSpec - - class BlockSequence(NamedTuple): - """Non-registering residual and motion-attention execution view.""" + """Non-registering residual-block execution view.""" blocks: nn.Sequential - motion_blocks: nn.ModuleList | None = None - motion_after_blocks: tuple[int, ...] = () def build_stem( @@ -1438,311 +957,13 @@ def build_residual_stack(count: int, start_index: int, spec: ResidualStackSpec) return nn.Sequential(*blocks) -def apply_block_sequence( - x: Tensor, - zq: Tensor | None, - sequence: BlockSequence, - *, - use_checkpointing: bool, -) -> Tensor: - """Run residual blocks with scheduled motion correspondence.""" - for index, block in enumerate(sequence.blocks): - inputs = (x, zq) if zq is not None else (x,) - x = forward_with_checkpointing(block, *inputs, use_checkpointing=use_checkpointing) - block_number = index + 1 - if sequence.motion_blocks is not None and block_number in sequence.motion_after_blocks: - motion_index = sequence.motion_after_blocks.index(block_number) - x = forward_with_checkpointing( - sequence.motion_blocks[motion_index], - x, - # NATTEN merge_attentions custom autograd is incompatible with - # PyTorch non-reentrant activation checkpointing. - use_checkpointing=False, - ) +def apply_block_sequence(x: Tensor, zq: Tensor | None, sequence: BlockSequence) -> Tensor: + """Run a residual-block stack.""" + for block in sequence.blocks: + x = block(x, zq) if zq is not None else block(x) return x -# Convolutional latent-space spatial upsampler with configurable upsampling strategy. - - -class ConvLatentUpsampler(nn.Module): - """Lightweight convolutional spatial upsampler for VAE latents. - - Upsampling strategies (controlled by ``upsample_mode``): - pixel_shuffle: Learned channel-to-space rearrangement (default). - bilinear: Parameter-free bilinear interpolation. - """ - - def __init__( - self, - in_channels: int = 16, - hidden_channels: int = 64, - bottleneck_channels: int | None = None, - num_pre_residual_blocks: int = 0, - num_residual_blocks: int = 16, - upscale_factor: int = 2, - *, - input_skip: bool = True, - upsample_mode: UpsampleMode = "pixel_shuffle", - upsample_position: UpsamplePosition = "after_projection", - temporal_mix: bool = True, - icnr: bool = False, - gradient_checkpointing: bool = False, - dims: int = 2, - expand_ratio: int = 4, - kernel_size: int = 3, - layer_scale_init: float | None = None, - grn: bool = False, - stochastic_depth_rate: float = 0.0, - depthwise: bool = False, - stem_channels: int | None = None, - zq_dim: int | None = None, - temporal_padding: TemporalPadding = "zeros", - upsample_padding_mode: UpsamplePaddingMode = "reflect", - ) -> None: - """Initialize the upsampler. - - Args: - in_channels: Number of VAE latent channels. - hidden_channels: Internal width of residual blocks. - bottleneck_channels: Mid-channel bottleneck width for ResidualBlocks (None = hidden_channels). - num_pre_residual_blocks: Number of residual blocks before upsampling (default 0). - num_residual_blocks: Number of post-upsample residual refinement blocks. - upscale_factor: Spatial upscale factor for pixel_shuffle and bilinear modes. - input_skip: Add ``repeat_interleave`` skip connection from input to ``conv_in`` output. - upsample_mode: Upsampling strategy — "pixel_shuffle" or "bilinear" (parameter-free interpolation). - upsample_position: Where to apply upsampling — "before_projection" or "after_projection". - temporal_mix: Whether pixel-shuffle Conv3d mixes across the temporal dimension. - When False, uses (1,3,3) kernel. Only affects dims=3 + pixel_shuffle. - icnr: Apply ICNR initialization to the sub-pixel-conv weight inside the - ``pixel_shuffle`` and ``pxs_v2_hybrid`` upsample modes. Removes - checkerboard at init. Ignored by ``bilinear`` and ``pxs_v2``. - gradient_checkpointing: Enable gradient checkpointing for residual blocks. - dims: Convolution dimensionality — 2 for Conv2d, 3 for Conv3d. - expand_ratio: Expansion ratio for inverted bottleneck (used when ``bottleneck_channels`` is None). - kernel_size: Kernel size for depthwise convolutions. - layer_scale_init: Initial LayerScale gamma value. ``None`` disables LayerScale. - grn: Whether to insert Global Response Normalization in residual blocks. - stochastic_depth_rate: Maximum drop rate for stochastic depth (linearly scheduled). - depthwise: Use depthwise separable convolutions in residual blocks. - stem_channels: When set, pre-residual blocks and input projection operate at this - narrower width, with a ResidualBlock transition to ``hidden_channels`` after upsampling. - zq_dim: When set, every ``ResidualBlock`` uses ``ModulatedRMSNorm`` conditioned on - a side tensor (the LQ latent), re-injecting LQ structure at every block via FiLM. - temporal_padding: How ``dims == 3`` convolutions extend T — ``"zeros"`` - (default) or ``"replicate"``, which repeats the edge frame as K-VAE does. - upsample_padding_mode: Edge handling for the ``pxs_v2`` residual conv — - ``"reflect"`` (default) or ``"zeros"``, matching the production K-VAE. - """ - super().__init__() - self.in_channels = in_channels - self.hidden_channels = hidden_channels - self.num_pre_residual_blocks = num_pre_residual_blocks - self.num_residual_blocks = num_residual_blocks - self.upscale_factor = upscale_factor - self.upsample_mode = upsample_mode - self.upsample_position = upsample_position - self.input_skip = input_skip - self.gradient_checkpointing = gradient_checkpointing - self.dims = dims - self.zq_dim = zq_dim - - # Channel widths: stem (pre-upsample) vs main (post-upsample) - proj_channels = stem_channels if stem_channels is not None else hidden_channels - post_bottleneck = bottleneck_channels if bottleneck_channels is not None else hidden_channels * expand_ratio - pre_bottleneck = bottleneck_channels if bottleneck_channels is not None else proj_channels * expand_ratio - self.bottleneck_channels = post_bottleneck - - # Projection layers - conv = make_conv(dims, temporal_padding) - self.input_proj = nn.Sequential( - conv(in_channels, proj_channels, kernel_size=3, padding=1), - RMSNorm(proj_channels, dims=dims), - nn.SiLU(), - ) - self.output_proj = nn.Sequential( - RMSNorm(hidden_channels, dims=dims), - nn.SiLU(), - conv(hidden_channels, in_channels, kernel_size=3, padding=1), - ) - - # Spatial upsampling - upsample_channels = in_channels if upsample_position == "before_projection" else proj_channels - if upsample_mode == "pixel_shuffle": - self.upsample: nn.Module = PixelShuffleND( - upsample_channels, - dims=dims, - factor=upscale_factor, - temporal_mix=temporal_mix, - icnr=icnr, - temporal_padding=temporal_padding, - ) - elif upsample_mode == "bilinear": - self.upsample = BilinearUpsampleND(dims=dims, factor=upscale_factor) - elif upsample_mode == "pxs_v2": - self.upsample = PXSv2UpsampleND( - upsample_channels, - dims=dims, - factor=upscale_factor, - padding_mode=upsample_padding_mode, - ) - else: - self.upsample = PXSv2HybridUpsampleND( - upsample_channels, - dims=dims, - factor=upscale_factor, - temporal_mix=temporal_mix, - icnr=icnr, - temporal_padding=temporal_padding, - padding_mode=upsample_padding_mode, - ) - - # Channel transition: bridges stem width → hidden width after upsampling - self.channel_transition: ResidualBlock | None = None - if stem_channels is not None: - self.channel_transition = ResidualBlock( - proj_channels, - out_channels=hidden_channels, - mid_channels=post_bottleneck, - dims=dims, - grn=grn, - zq_dim=zq_dim, - temporal_padding=temporal_padding, - ) - - # Residual stacks - total_blocks = num_pre_residual_blocks + num_residual_blocks - shared_block_kwargs = { - "total_blocks": total_blocks, - "stochastic_depth_rate": stochastic_depth_rate, - "dims": dims, - "layer_scale_init": layer_scale_init, - "grn": grn, - "depthwise": depthwise, - "kernel_size": kernel_size, - "zq_dim": zq_dim, - "temporal_padding": temporal_padding, - } - self.pre_residual_blocks = self._build_residual_stack( - num_pre_residual_blocks, - proj_channels, - pre_bottleneck, - start_idx=0, - **shared_block_kwargs, # type: ignore[arg-type] - ) - self.residual_blocks = self._build_residual_stack( - num_residual_blocks, - hidden_channels, - post_bottleneck, - start_idx=num_pre_residual_blocks, - **shared_block_kwargs, # type: ignore[arg-type] - ) - - def _build_residual_stack( - self, - count: int, - channels: int, - mid_channels: int, - *, - start_idx: int, - total_blocks: int, - stochastic_depth_rate: float, - dims: int, - layer_scale_init: float | None, - grn: bool, - depthwise: bool, - kernel_size: int, - zq_dim: int | None = None, - temporal_padding: TemporalPadding = "zeros", - ) -> nn.Sequential: - """Build a stack of ResidualBlocks with linearly scheduled stochastic depth.""" - blocks: list[nn.Module] = [] - for i in range(count): - drop_prob = stochastic_depth_rate * (start_idx + i) / max(total_blocks - 1, 1) - blocks.append( - ResidualBlock( - channels, - mid_channels=mid_channels, - dims=dims, - layer_scale_init=layer_scale_init, - grn=grn, - stochastic_depth_prob=drop_prob, - depthwise=depthwise, - kernel_size=kernel_size, - zq_dim=zq_dim, - temporal_padding=temporal_padding, - ) - ) - return nn.Sequential(*blocks) - - def _apply_blocks( - self, - x: Tensor, - blocks: nn.Sequential, - *, - use_checkpointing: bool, - zq: Tensor | None = None, - ) -> Tensor: - """Apply a residual stack with optional modulation.""" - for block in blocks: - if zq is not None: - x = forward_with_checkpointing(block, x, zq, use_checkpointing=use_checkpointing) - else: - x = forward_with_checkpointing(block, x, use_checkpointing=use_checkpointing) - return x - - def _apply_input_projection(self, x: Tensor) -> Tensor: - z_skip = x - x = self.input_proj(x) - if self.input_skip: - repeats = x.shape[1] // z_skip.shape[1] - x = x + z_skip.repeat_interleave(repeats, dim=1) - return x - - @property - def last_layer_weight(self) -> torch.Tensor: - """Weight of the final output conv (used by adaptive GAN-weight balancing).""" - return self.output_proj[-1].weight # type: ignore - - def forward(self, z: torch.Tensor) -> torch.Tensor: - """Upsample latent tensor spatially. - - Args: - z: Input latent tensor of shape ``(B, C, T, H, W)``. - - Returns: - Upsampled tensor of shape ``(B, C, T, H', W')``. - """ - b, _, t, _, _ = z.shape - x = rearrange(z, "b c t h w -> (b t) c h w") if self.dims == 2 else z - zq = x if self.zq_dim is not None else None - if self.upsample_position == "before_projection": - x = self.upsample(x) - x = self._apply_input_projection(x) - use_ckpt = self.training and self.gradient_checkpointing - x = self._apply_blocks( - x, - self.pre_residual_blocks, - use_checkpointing=use_ckpt, - zq=zq, - ) - if self.upsample_position == "after_projection": - x = self.upsample(x) - if self.channel_transition is not None: - if zq is not None: - x = forward_with_checkpointing(self.channel_transition, x, zq, use_checkpointing=use_ckpt) - else: - x = forward_with_checkpointing(self.channel_transition, x, use_checkpointing=use_ckpt) - x = self._apply_blocks( - x, - self.residual_blocks, - use_checkpointing=use_ckpt, - zq=zq, - ) - x = self.output_proj(x) - return rearrange(x, "(b t) c h w -> b c t h w", b=b, t=t) if self.dims == 2 else x - - # Scale-specific x2 adaptation modules for the cascaded latent upscaler. @@ -1771,11 +992,6 @@ def __init__(self, block: ResidualBlock) -> None: self.norm1 = copy.deepcopy(block.norm1) self.norm2 = copy.deepcopy(block.norm2) - def copy_from(self, block: ResidualBlock) -> None: - """Reset the private norms from their shared x4 counterparts.""" - self.norm1.load_state_dict(block.norm1.state_dict()) - self.norm2.load_state_dict(block.norm2.state_dict()) - def require_pxs_upsample(source: nn.Module) -> PXSv2UpsampleND: """Require a pxs_v2 upsample stage as the finisher donor.""" @@ -1799,13 +1015,6 @@ def __init__(self, source: nn.Module, mode: X2FinisherMode) -> None: self.linear = copy.deepcopy(donor.linear) self.spatial_conv = copy.deepcopy(donor.spatial_conv) if mode == "pxs_residual" else None - def copy_from(self, source: nn.Module) -> None: - """Reset the finisher convs from their x4 upsample counterparts.""" - donor = require_pxs_upsample(source) - self.linear.load_state_dict(donor.linear.state_dict()) - if self.spatial_conv is not None: - self.spatial_conv.load_state_dict(donor.spatial_conv.state_dict()) - def forward(self, x: Tensor) -> Tensor: """Apply the cloned convs on the unchanged grid.""" if self.spatial_conv is not None: @@ -1852,28 +1061,6 @@ def __init__( raise ValueError(msg) self.private_mid_blocks = copy.deepcopy(shared_stage.mid_blocks) - def initialize_from_x4(self, shared_stage: SharedSecondStage) -> None: - """Copy scale-specific tail state from the loaded x4 baseline.""" - if self.tail_mode == "scale_specific_norm": - if self.post_norms is None or self.output_norm is None: - msg = "scale-specific normalization modules are not initialized" - raise RuntimeError(msg) - for norms, block in zip(self.post_norms, shared_stage.blocks, strict=True): - norms.copy_from(block) - self.output_norm.load_state_dict(shared_stage.output_proj[0].state_dict()) - elif self.tail_mode in ("private", "private_full"): - if self.private_upsample is None or self.private_blocks is None or self.private_output_proj is None: - msg = "private x2 tail modules are not initialized" - raise RuntimeError(msg) - self.private_upsample.load_state_dict(shared_stage.upsample.state_dict()) - self.private_blocks.load_state_dict(shared_stage.blocks.state_dict()) - self.private_output_proj.load_state_dict(shared_stage.output_proj.state_dict()) - if self.private_mid_blocks is not None: - if shared_stage.mid_blocks is None: - msg = "private mid blocks require the shared stage to expose mid_blocks" - raise RuntimeError(msg) - self.private_mid_blocks.load_state_dict(shared_stage.mid_blocks.state_dict()) - def private_stage(self, shared_stage: SharedSecondStage) -> SharedSecondStage: """Select the shared stage or the registered private copy.""" if self.tail_mode not in ("private", "private_full"): @@ -1887,14 +1074,7 @@ def private_stage(self, shared_stage: SharedSecondStage) -> SharedSecondStage: output_proj=self.private_output_proj, ) - def apply_post_blocks( - self, - x: Tensor, - zq: Tensor | None, - stage: SharedSecondStage, - *, - use_checkpointing: bool, - ) -> Tensor: + def apply_post_blocks(self, x: Tensor, zq: Tensor | None, stage: SharedSecondStage) -> Tensor: """Apply post blocks with shared or x2-specific normalization.""" for index, block in enumerate(stage.blocks): if self.tail_mode == "scale_specific_norm": @@ -1902,12 +1082,11 @@ def apply_post_blocks( msg = "scale-specific normalization requires zq and private norms" raise RuntimeError(msg) norms = self.post_norms[index] - block_call = partial(block.forward_with_norms, norm1=norms.norm1, norm2=norms.norm2) # type: ignore - x = forward_with_checkpointing(block_call, x, zq, use_checkpointing=use_checkpointing) + x = block.forward_with_norms(x, zq, norm1=norms.norm1, norm2=norms.norm2) # type: ignore elif zq is None: - x = forward_with_checkpointing(block, x, use_checkpointing=use_checkpointing) + x = block(x) else: - x = forward_with_checkpointing(block, x, zq, use_checkpointing=use_checkpointing) + x = block(x, zq) return x def apply_output_projection(self, x: Tensor, zq: Tensor | None, stage: SharedSecondStage) -> Tensor: @@ -1921,26 +1100,21 @@ def apply_output_projection(self, x: Tensor, zq: Tensor | None, stage: SharedSec x = stage.output_proj[1](x) return stage.output_proj[2](x) - def apply_block_stack( - self, blocks: nn.Sequential, x: Tensor, zq: Tensor | None, *, use_checkpointing: bool - ) -> Tensor: + def apply_block_stack(self, blocks: nn.Sequential, x: Tensor, zq: Tensor | None) -> Tensor: """Run a residual stack with the branch's zq-threading convention.""" for block in blocks: - if zq is None: - x = forward_with_checkpointing(block, x, use_checkpointing=use_checkpointing) - else: - x = forward_with_checkpointing(block, x, zq, use_checkpointing=use_checkpointing) + x = block(x, zq) if zq is not None else block(x) return x - def apply_adapter(self, x: Tensor, zq: Tensor | None, *, use_checkpointing: bool) -> Tensor: + def apply_adapter(self, x: Tensor, zq: Tensor | None) -> Tensor: """Run the x2-only residual adapter.""" - return self.apply_block_stack(self.adapter, x, zq, use_checkpointing=use_checkpointing) + return self.apply_block_stack(self.adapter, x, zq) - def apply_mid_blocks(self, x: Tensor, zq: Tensor | None, *, use_checkpointing: bool) -> Tensor: + def apply_mid_blocks(self, x: Tensor, zq: Tensor | None) -> Tensor: """Run the private mid blocks ahead of the second-stage upsample (private_full only).""" if self.private_mid_blocks is None: return x - return self.apply_block_stack(self.private_mid_blocks, x, zq, use_checkpointing=use_checkpointing) + return self.apply_block_stack(self.private_mid_blocks, x, zq) def forward_second_stage( self, @@ -1949,31 +1123,23 @@ def forward_second_stage( batch_size: int, frames: int, shared_stage: SharedSecondStage, - *, - use_checkpointing: bool, ) -> Tensor: """Run the selected x2 second stage without registering shared aliases.""" stage = self.private_stage(shared_stage) if self.finisher is not None: x = self.finisher(x) - x = self.apply_mid_blocks(x, zq, use_checkpointing=use_checkpointing) + x = self.apply_mid_blocks(x, zq) x = stage.upsample(x) - x = self.apply_post_blocks(x, zq, stage, use_checkpointing=use_checkpointing) + x = self.apply_post_blocks(x, zq, stage) x = self.apply_output_projection(x, zq, stage) if self.dims == 2: - return rearrange(x, "(b t) c h w -> b c t h w", b=batch_size, t=frames) + return split_batch_time(x, batch_size, frames) return x # Multi-scale cascaded upsampler: two 2x stages with intermediate supervision. -# The two projections back to latent space. They are scale-specific, so a warm -# start from another network cannot supply them and a warm-started run may want -# to fit them before letting gradients reach the rest. -HEAD_PREFIXES = ("output_proj.", "mid_output_head.") - - class MultiScaleUpsampler(nn.Module): """Cascaded 2x+2x VAE-latent upsampler with an optional x2 entry. @@ -1998,7 +1164,6 @@ def __init__( upsample_mode: UpsampleMode = "pixel_shuffle", temporal_mix: bool = True, icnr: bool = False, - gradient_checkpointing: bool = False, dims: int = 3, expand_ratio: int = 4, kernel_size: int = 3, @@ -2006,7 +1171,6 @@ def __init__( grn: bool = False, stochastic_depth_rate: float = 0.0, depthwise: bool = False, - motion_attention: MotionAttentionSpec | None = None, zq_dim: int | None = None, bare_stem: bool = False, modulated_output_proj: bool = False, @@ -2026,7 +1190,6 @@ def __init__( self.hidden_channels = hidden_channels self.upscale_factor = 4 self.input_skip = input_skip - self.gradient_checkpointing = gradient_checkpointing self.dims = dims self.zq_dim = zq_dim self.global_skip = global_skip @@ -2053,9 +1216,6 @@ def __init__( if enable_x2_entry and x2_tail_mode != "private_full": msg = "stage_channels with enable_x2_entry requires x2_tail_mode='private_full'" raise ValueError(msg) - if motion_attention is not None: - msg = "stage_channels is incompatible with motion_attention" - raise ValueError(msg) if input_skip: msg = "stage_channels is incompatible with input_skip" raise ValueError(msg) @@ -2067,22 +1227,11 @@ def __init__( raise ValueError(msg) self.stage_channels = stage_channels w1, w2, w3 = stage_channels if stage_channels is not None else (hidden_channels,) * 3 - self.motion_after_mid_blocks = motion_attention.after_mid_blocks if motion_attention is not None else () if input_skip and hidden_channels % in_channels != 0: - msg = f"hidden_channels ({hidden_channels}) must be divisible by in_channels ({in_channels}) for input_skip" - raise ValueError(msg) - if motion_attention is not None and dims != 3: - msg = "motion_attention requires dims=3" - raise ValueError(msg) - if motion_attention is not None and enable_x2_entry: - msg = "motion_attention Round 23 is x4-only and cannot enable the x2 entry" - raise ValueError(msg) - if motion_attention is not None and not motion_attention.after_mid_blocks: - msg = "motion_attention requires at least one mid-stage placement" - raise ValueError(msg) - if motion_attention is not None and motion_attention.after_mid_blocks[-1] > num_mid_blocks: - msg = "motion attention placement exceeds the mid residual stack" + msg = ( + f"hidden_channels ({hidden_channels}) must be divisible by in_channels ({in_channels}) for input_skip" + ) raise ValueError(msg) # Projection layers @@ -2129,17 +1278,13 @@ def stage_spec(channels: int, first_in: int | None) -> ResidualStackSpec: ) self.pre_blocks = build_residual_stack(num_pre_blocks, 0, stage_spec(w1, None)) - self.mid_blocks = build_residual_stack(num_mid_blocks, num_pre_blocks, stage_spec(w2, w1 if w1 != w2 else None)) + self.mid_blocks = build_residual_stack( + num_mid_blocks, num_pre_blocks, stage_spec(w2, w1 if w1 != w2 else None) + ) self.post_blocks = build_residual_stack( num_post_blocks, num_pre_blocks + num_mid_blocks, stage_spec(w3, w2 if w2 != w3 else None) ) - self.mid_motion_blocks: nn.ModuleList | None = None - if motion_attention is not None: - self.mid_motion_blocks = nn.ModuleList( - [MotionCorrespondenceBlock(motion_attention.block) for _placement in motion_attention.after_mid_blocks] - ) - # Build x2-exclusive modules after the complete x4 path so enabling an x2 # entry does not perturb the base-path RNG initialization. self.mid_input_proj: nn.Sequential | None = None @@ -2172,15 +1317,6 @@ def stage_spec(channels: int, first_in: int | None) -> ResidualStackSpec: nn.init.zeros_(final_conv.weight) # type: ignore[arg-type] nn.init.zeros_(final_conv.bias) # type: ignore[arg-type] - @property - def last_layer_weight(self) -> torch.Tensor: - """Weight of the final 4x-output conv (used by adaptive GAN-weight balancing). - - The GAN-loss runs only on the 4x output, so the relevant last layer is - ``output_proj[-1]`` rather than the 2x ``mid_output_head``. - """ - return self.output_proj[-1].weight # type: ignore - def shared_second_stage(self) -> SharedSecondStage: """Return a non-registering view of the e115-compatible second stage.""" return SharedSecondStage( @@ -2190,33 +1326,6 @@ def shared_second_stage(self) -> SharedSecondStage: mid_blocks=self.mid_blocks, ) - @staticmethod - def is_x2_state_key(key: str) -> bool: - """Return whether a state key belongs exclusively to the x2 path.""" - return key.startswith(("mid_input_proj.", "x2_branch.")) - - def initialize_x2_from_x4(self) -> None: - """Initialize x2-exclusive modules from the currently loaded x4 path.""" - if self.mid_input_proj is None or self.x2_branch is None: - msg = "x2 initialization requires enable_x2_entry=True" - raise ValueError(msg) - self.mid_input_proj.load_state_dict(self.input_proj.state_dict()) - self.x2_branch.initialize_from_x4(self.shared_second_stage()) - if self.x2_adapter_sources is not None: - for block, source_index in zip(self.x2_branch.adapter, self.x2_adapter_sources, strict=True): - block.load_state_dict(self.pre_blocks[source_index].state_dict()) - if self.x2_branch.finisher is not None: - self.x2_branch.finisher.copy_from(self.upsample_1) - - def configure_x2_finetuning(self) -> None: - """Freeze the e115 path and expose only x2-exclusive parameters.""" - if self.mid_input_proj is None or self.x2_branch is None: - msg = "x2 fine-tuning requires enable_x2_entry=True" - raise ValueError(msg) - self.requires_grad_(requires_grad=False) - self.mid_input_proj.requires_grad_(requires_grad=True) - self.x2_branch.requires_grad_(requires_grad=True) - def _project_with_skip(self, x: Tensor, proj: nn.Module) -> Tensor: """Apply a projection and optionally add the channel-repeated input skip.""" projected = proj(x) @@ -2225,24 +1334,19 @@ def _project_with_skip(self, x: Tensor, proj: nn.Module) -> Tensor: projected = projected + x.repeat_interleave(repeats, dim=1) return projected - def _head_x4(self, x: Tensor, zq: Tensor | None, *, use_ckpt: bool) -> Tensor: + def _head_x4(self, x: Tensor, zq: Tensor | None) -> Tensor: """x4 entry: ``input_proj -> input_skip -> pre_blocks -> upsample_1`` (-> 2x grid).""" x = self._project_with_skip(x, self.input_proj) - x = apply_block_sequence( - x, - zq, - BlockSequence(self.pre_blocks), - use_checkpointing=use_ckpt, - ) + x = apply_block_sequence(x, zq, BlockSequence(self.pre_blocks)) return self.upsample_1(x) - def _head_x2(self, x: Tensor, zq: Tensor | None, *, use_ckpt: bool) -> Tensor: + def _head_x2(self, x: Tensor, zq: Tensor | None) -> Tensor: """Project a genuine x2 latent and apply the scale-specific adapter.""" if self.mid_input_proj is None or self.x2_branch is None: msg = "x2 head requires enable_x2_entry=True" raise ValueError(msg) x = self._project_with_skip(x, self.mid_input_proj) - return self.x2_branch.apply_adapter(x, zq, use_checkpointing=use_ckpt) + return self.x2_branch.apply_adapter(x, zq) def _tail( self, @@ -2251,34 +1355,28 @@ def _tail( b: int, t: int, *, - use_ckpt: bool, return_intermediates: bool, ) -> dict[str, torch.Tensor] | torch.Tensor: """x4 continuation: ``mid_blocks -> [2x head] -> second stage``.""" - x = apply_block_sequence( - x, - zq, - BlockSequence(self.mid_blocks, self.mid_motion_blocks, self.motion_after_mid_blocks), - use_checkpointing=use_ckpt, - ) + x = apply_block_sequence(x, zq, BlockSequence(self.mid_blocks)) # 2x supervision branch (only when the caller will use it) z_2x = None if return_intermediates: z_2x_inner = self.mid_output_head(x) - z_2x = rearrange(z_2x_inner, "(b t) c h w -> b c t h w", b=b, t=t) if self.dims == 2 else z_2x_inner + z_2x = split_batch_time(z_2x_inner, b, t) if self.dims == 2 else z_2x_inner - z_4x = self._second_stage(x, zq, b, t, use_ckpt=use_ckpt) + z_4x = self._second_stage(x, zq, b, t) if return_intermediates: return {"2x": z_2x, "4x": z_4x} # type: ignore[return-value] return z_4x - def _second_stage(self, x: Tensor, zq: Tensor | None, b: int, t: int, *, use_ckpt: bool) -> Tensor: + def _second_stage(self, x: Tensor, zq: Tensor | None, b: int, t: int) -> Tensor: """Shared second stage: ``upsample_2 -> post_blocks -> output_proj`` — 2x of the current grid.""" x = self.upsample_2(x) - x = apply_block_sequence(x, zq, BlockSequence(self.post_blocks), use_checkpointing=use_ckpt) + x = apply_block_sequence(x, zq, BlockSequence(self.post_blocks)) x = apply_output_head(self.output_proj, x, zq) - return rearrange(x, "(b t) c h w -> b c t h w", b=b, t=t) if self.dims == 2 else x + return split_batch_time(x, b, t) if self.dims == 2 else x def forward( self, @@ -2301,25 +1399,17 @@ def forward( return_intermediates = self.training b, _, t, _, _ = z.shape - x = rearrange(z, "b c t h w -> (b t) c h w") if self.dims == 2 else z + x = merge_batch_time(z) if self.dims == 2 else z zq = x if self.zq_dim is not None else None - use_ckpt = self.training and self.gradient_checkpointing if entry == "x2": if self.x2_branch is None: msg = "x2 branch requires enable_x2_entry=True" raise ValueError(msg) - x = self._head_x2(x, zq, use_ckpt=use_ckpt) - return self.x2_branch.forward_second_stage( - x, - zq, - b, - t, - self.shared_second_stage(), - use_checkpointing=use_ckpt, - ) - x = self._head_x4(x, zq, use_ckpt=use_ckpt) - out = self._tail(x, zq, b, t, use_ckpt=use_ckpt, return_intermediates=return_intermediates) + x = self._head_x2(x, zq) + return self.x2_branch.forward_second_stage(x, zq, b, t, self.shared_second_stage()) + x = self._head_x4(x, zq) + out = self._tail(x, zq, b, t, return_intermediates=return_intermediates) if not self.global_skip: return out return self._add_nearest_base(out, z) @@ -2338,111 +1428,57 @@ def _add_nearest_base( return out + spatial_nearest_2x(z, DIMS_3, self.upscale_factor) -# Factory functions for building upsampler models. +# Factory function for building an upsampler model. -def build_upsampler(config: FlatModelConfig | MultiScaleModelConfig) -> nn.Module: - """Create an upsampler module from config. - - Dispatches on config type: ``FlatModelConfig`` builds a - ``ConvLatentUpsampler``; ``MultiScaleModelConfig`` builds a - ``MultiScaleUpsampler``. +def build_upsampler(config: Mapping[str, Any]) -> MultiScaleUpsampler: + """Create a `MultiScaleUpsampler` from a serialized model config. Args: - config: Model configuration (flat or multi-scale). + config: Multi-scale model configuration, as stored under `latent_upscaler/config.json`'s + `models[i].model`. Fields absent from older published configs fall back to + `MultiScaleUpsampler`'s own defaults. Returns: Initialized upsampler module. - """ - if config.architecture == "multi_scale": - return _build_multi_scale(config) - return _build_flat(config) - - -def _build_flat(config: FlatModelConfig) -> ConvLatentUpsampler: - """Build a flat (single-stage) ConvLatentUpsampler from config.""" - return ConvLatentUpsampler( - in_channels=config.in_channels, - hidden_channels=config.hidden_channels, - bottleneck_channels=config.bottleneck_channels, - num_pre_residual_blocks=config.num_pre_residual_blocks, - num_residual_blocks=config.num_residual_blocks, - upscale_factor=config.upscale_factor, - input_skip=config.input_skip, - upsample_mode=config.upsample_mode, - upsample_position=config.upsample_position, - temporal_mix=config.temporal_mix, - icnr=config.icnr, - gradient_checkpointing=config.gradient_checkpointing, - dims=config.dims, - expand_ratio=config.expand_ratio, - kernel_size=config.kernel_size, - layer_scale_init=config.layer_scale_init, - grn=config.grn, - stochastic_depth_rate=config.stochastic_depth_rate, - depthwise=config.depthwise, - stem_channels=config.stem_channels, - zq_dim=config.in_channels if config.modulated_norm else None, - temporal_padding=config.temporal_padding, - upsample_padding_mode=config.upsample_padding_mode, - ) - -def _build_multi_scale(config: MultiScaleModelConfig) -> MultiScaleUpsampler: - """Build a cascaded 2x+2x MultiScaleUpsampler from config.""" - motion_attention = build_motion_attention_spec(config) + Raises: + ValueError: If `config["architecture"]` is not `"multi_scale"`. + """ + if config["architecture"] != "multi_scale": + raise ValueError(f"unsupported latent-upscaler architecture={config['architecture']!r}") + stage_channels = config["stage_channels"] + x2_adapter_sources = config.get("x2_adapter_sources") return MultiScaleUpsampler( - in_channels=config.in_channels, - hidden_channels=config.hidden_channels, - bottleneck_channels=config.bottleneck_channels, - num_pre_blocks=config.num_pre_blocks, - num_mid_blocks=config.num_mid_blocks, - num_post_blocks=config.num_post_blocks, - input_skip=config.input_skip, - upsample_mode=config.upsample_mode, - temporal_mix=config.temporal_mix, - icnr=config.icnr, - gradient_checkpointing=config.gradient_checkpointing, - dims=config.dims, - expand_ratio=config.expand_ratio, - kernel_size=config.kernel_size, - layer_scale_init=config.layer_scale_init, - grn=config.grn, - stochastic_depth_rate=config.stochastic_depth_rate, - depthwise=config.depthwise, - motion_attention=motion_attention, - zq_dim=config.in_channels if config.modulated_norm else None, - bare_stem=config.bare_stem, - modulated_output_proj=config.modulated_output_proj, - global_skip=config.global_skip, - enable_x2_entry=config.enable_x2_entry, - x2_adapter_blocks=config.x2_adapter_blocks, - x2_tail_mode=config.x2_tail_mode, - x2_adapter_sources=config.x2_adapter_sources, - x2_finisher=config.x2_finisher, - stage_channels=config.stage_channels, - temporal_padding=config.temporal_padding, - upsample_padding_mode=config.upsample_padding_mode, - ) - - -def build_motion_attention_spec(config: MultiScaleModelConfig) -> MotionAttentionSpec | None: - """Resolve the nested motion-attention config into a construction spec.""" - motion = config.motion_attention - if motion is None: - return None - return MotionAttentionSpec( - after_mid_blocks=motion.after_mid_blocks, - block=MotionCorrespondenceSpec( - channels=config.hidden_channels, - spatial_kernel_size=motion.spatial_kernel_size, - temporal_offsets=motion.temporal_offsets, - num_heads=motion.num_heads, - head_dim=config.hidden_channels // motion.num_heads, - backend=motion.backend, - natten_backend=motion.natten_backend, - merge_compile=motion.merge_compile, - ), + in_channels=config["in_channels"], + hidden_channels=config["hidden_channels"], + bottleneck_channels=config["bottleneck_channels"], + num_pre_blocks=config["num_pre_blocks"], + num_mid_blocks=config["num_mid_blocks"], + num_post_blocks=config["num_post_blocks"], + input_skip=config["input_skip"], + upsample_mode=config["upsample_mode"], + temporal_mix=config["temporal_mix"], + icnr=config.get("icnr", False), + dims=config["dims"], + expand_ratio=config["expand_ratio"], + kernel_size=config["kernel_size"], + layer_scale_init=config["layer_scale_init"], + grn=config["grn"], + stochastic_depth_rate=config["stochastic_depth_rate"], + depthwise=config["depthwise"], + zq_dim=config["in_channels"] if config["modulated_norm"] else None, + bare_stem=config["bare_stem"], + modulated_output_proj=config["modulated_output_proj"], + global_skip=config.get("global_skip", False), + enable_x2_entry=config["enable_x2_entry"], + x2_adapter_blocks=config.get("x2_adapter_blocks", 0), + x2_tail_mode=config.get("x2_tail_mode", "shared"), + x2_adapter_sources=tuple(x2_adapter_sources) if x2_adapter_sources is not None else None, + x2_finisher=config.get("x2_finisher", "none"), + stage_channels=tuple(stage_channels) if stage_channels is not None else None, + temporal_padding=config["temporal_padding"], + upsample_padding_mode=config["upsample_padding_mode"], ) @@ -2456,7 +1492,7 @@ class Kandinsky6SRLatentUpscalerBank(ModelMixin, ConfigMixin): latents before they are passed to an upscaler. """ - _no_split_modules = ["ConvLatentUpsampler", "MultiScaleUpsampler"] + _no_split_modules = ["MultiScaleUpsampler"] @register_to_config def __init__( @@ -2472,8 +1508,7 @@ def __init__( target_scale = f"{target_scale}x" if target_scale in self._models: raise ValueError(f"duplicate latent upscaler target scale: {target_scale}") - model_config = _model_config(spec["model"]) - upscaler = build_upsampler(model_config) + upscaler = build_upsampler(spec["model"]) upscaler.target_scale = target_scale upscaler.scaling_factor = float(scaling_factor) self._models[target_scale] = upscaler @@ -2516,16 +1551,10 @@ def forward( if not key.endswith("x"): key = f"{key}x" model = self._models[key] - kwargs: dict[str, Any] = {} + kwargs: dict[str, Any] = {"entry": f"x{key.removesuffix('x')}"} if return_intermediates is not None: kwargs["return_intermediates"] = return_intermediates - if hasattr(model, "forward"): - kwargs["entry"] = f"x{key.removesuffix('x')}" - try: - return model(latent, **kwargs) - except TypeError: - kwargs.pop("entry", None) - return model(latent, **kwargs) + return model(latent, **kwargs) __all__ = ["Kandinsky6SRLatentUpscalerBank"] diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py index e30824329fd6..3178dc564974 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py @@ -14,28 +14,22 @@ """Kandinsky 6 super-resolution Diffusers pipeline.""" -from __future__ import annotations - import math -from collections.abc import Callable, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass -from fractions import Fraction -from pathlib import Path -from types import SimpleNamespace from typing import Any, Literal, NamedTuple import numpy as np import torch -from ..pipeline_utils import DiffusionPipeline -from diffusers.utils import logging, replace_example_docstring -from ...schedulers.scheduling_piflow import DXPolicy, policy_rollout_fm, shift_timesteps from torch import Tensor, nn -from torch.distributed import all_gather from torch.nn import functional +from ...utils import logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline from .pipeline_output import Kandinsky6SRPipelineOutput + logger = logging.get_logger(__name__) @@ -44,58 +38,31 @@ ```python >>> import torch - >>> from diffusers import Kandinsky6SRPipeline + >>> from diffusers import Kandinsky6SRPipeline, Kandinsky6TI2VAPipeline - >>> model_id = "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers" - >>> pipe = Kandinsky6SRPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) - >>> pipe = pipe.to("cuda") + >>> t2va = Kandinsky6TI2VAPipeline.from_pretrained( + ... "kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers", torch_dtype=torch.bfloat16 + ... ) + >>> t2va.enable_model_cpu_offload() + >>> generated = t2va(prompt="A cat and a dog baking a cake together in a kitchen.", height=480, width=864) - >>> output = pipe.from_video( - ... "input.mp4", - ... resolution_scale=2.25, - ... save_path="output.mp4", + >>> pipe = Kandinsky6SRPipeline.from_pretrained( + ... "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers", torch_dtype=torch.bfloat16 ... ) - >>> output.path - 'output.mp4' + >>> pipe.enable_model_cpu_offload() + + >>> output = pipe(video=generated.frames, resolution_scale=2.25) + >>> output.frames.shape + torch.Size([1, 3, 121, 1080, 1944]) ``` """ # These values are the SR data contract and must stay local to this pipeline. VAE_SPATIAL_FACTOR = 16 -VAE_TEMPORAL_FACTOR = 4 RESOLUTIONS: dict[int, list[tuple[int, int]]] = {512: [(512, 512), (512, 768), (768, 512)]} -TARGET_FPS = 24 MAX_NUM_FRAMES = 121 -RESAMPLE_FPS_TOLERANCE = 1.5 -LATENT_NDIM = 4 -X0_T_CLAMP = 1e-5 -SUPPORTED_TARGET_SCALES = ("2x", "4x") ResolutionScale = Literal[2, 4] -TotalResolutionScale = Literal[2, 4, 2.25] -constants = SimpleNamespace(VAE_SPATIAL_FACTOR=VAE_SPATIAL_FACTOR) -TileGridMode = Literal["legacy", "even"] - - -def _call_module_method(module: nn.Module, method: str, *args: Any, **kwargs: Any) -> Any: - """Call a non-forward method through a Diffusers offload hook.""" - hook = getattr(module, "_hf_hook", None) - if hook is None: - return getattr(module, method)(*args, **kwargs) - args, kwargs = hook.pre_forward(module, *args, **kwargs) - output = getattr(module, method)(*args, **kwargs) - return hook.post_forward(module, output) - - -@contextmanager -def _execution_device_context(device: torch.device | str): - """Scope SR model execution to Diffusers' resolved CUDA device.""" - resolved_device = torch.device(device) - if resolved_device.type == "cuda": - with torch.cuda.device(resolved_device): - yield - else: - yield # Shared tiling utilities for spatial video tiling with Hanning-window blending. @@ -105,16 +72,8 @@ class TileGrid(NamedTuple): """Spatial tile grid parameters. ``tops`` and ``lefts`` are the authoritative per-axis tile start positions. - They cover ``[0, length - tile]`` with the nominal stride; the **last** - position is clamped to ``length - tile`` if a uniform stride would - overshoot. This means consecutive gaps in ``tops`` / ``lefts`` may be - shorter than the nominal stride near the right/bottom edge. - - ``n_h``, ``n_w``, ``total_tiles``, ``stride_h``, ``stride_w`` are derived - properties exposed for backward compatibility with callers that only - read them. To compute a tile position by index, prefer - ``grid.tops[row]`` / ``grid.lefts[col]`` over ``row * grid.stride_h`` — - the latter is wrong for the last clamped tile. + They cover ``[0, length - tile]``, evenly spread so every gap keeps at + least the configured overlap (see :func:`axis_positions`). """ tile_h: int @@ -137,105 +96,73 @@ def total_tiles(self) -> int: """Total number of tiles (``n_h * n_w``).""" return self.n_h * self.n_w - @property - def stride_h(self) -> int: - """Nominal vertical stride (gap between first two rows; ``tile_h`` if only one row).""" - return self.tops[1] - self.tops[0] if len(self.tops) > 1 else self.tile_h - - @property - def stride_w(self) -> int: - """Nominal horizontal stride (gap between first two cols; ``tile_w`` if only one col).""" - return self.lefts[1] - self.lefts[0] if len(self.lefts) > 1 else self.tile_w - -def axis_positions(length: int, tile: int, stride: int) -> tuple[int, ...]: - """Return tile start positions covering ``[0, length - tile]``. +def axis_positions(length: int, tile: int, min_overlap: float, snap: int) -> tuple[int, ...]: + """Return evenly distributed, ``snap``-aligned tile start positions. - Walks at the given ``stride`` from 0 and appends positions while the tile - still fits inside ``length``. If the last walked position is not exactly - ``length - tile``, appends one final position clamped to the right edge. - Single-position result if ``tile >= length``. + Uses the minimal position count whose uniform stride keeps the tile + overlap at or above ``min_overlap``, then spreads the positions evenly + over ``[0, length - tile]`` in integer ``snap`` units. When ``length`` or + ``tile`` is not ``snap``-aligned the same layout is computed at pixel + precision (``snap=1``) — the latent-grid guard downstream still enforces + alignment where it actually matters (the LU path). Args: length: Axis length in pixels. tile: Tile size on this axis. - stride: Nominal stride between consecutive tiles. + min_overlap: Overlap floor as a fraction of ``tile`` in ``[0, 1)``. + snap: Position alignment unit (the VAE spatial factor). Returns: - Strictly increasing tuple of start positions. + Strictly increasing positions covering ``[0, length - tile]``. + + Raises: + ValueError: On non-positive ``tile``/``snap`` or ``min_overlap`` + outside ``[0, 1)``. """ - if tile <= 0 or stride <= 0: - msg = f"tile and stride must be positive (got tile={tile}, stride={stride})" + if tile <= 0 or snap <= 0 or not 0 <= min_overlap < 1: + msg = f"invalid axis spec: tile={tile}, snap={snap}, min_overlap={min_overlap}" raise ValueError(msg) if tile >= length: return (0,) - - positions: list[int] = [] - p = 0 - while p + tile <= length: - positions.append(p) - p += stride - last = length - tile - if positions[-1] != last: - positions.append(last) - return tuple(positions) + unit = snap if length % snap == 0 and tile % snap == 0 else 1 + span_units = (length - tile) // unit + max_stride_units = max(1, math.floor(tile * (1.0 - min_overlap) / unit)) + count = math.ceil(span_units / max_stride_units) + 1 + return tuple(round(i * span_units / (count - 1)) * unit for i in range(count)) -def compute_tile_grid( - h: int, - w: int, - resolution_scale: int, - overlap: float = 0.5, - tile_hw: tuple[int, int] | None = None, -) -> TileGrid: - """Compute tile geometry with configurable overlap. - - By default, tile size is ``(h // resolution_scale, w // resolution_scale)``. - Pass ``tile_hw`` to override with explicit tile dimensions (e.g. derived - from a fixed base resolution rather than the source frame size). - - The grid covers the full frame: if the nominal stride does not divide - evenly into ``(h - tile_h)`` / ``(w - tile_w)``, the last tile in each - axis is clamped to the right/bottom edge. ``stitch_tiles_hanning`` - handles the resulting non-uniform overlap correctly via Hanning-window - normalisation. +def compute_tile_grid(h: int, w: int, tile_hw: tuple[int, int], min_overlap: float, snap: int) -> TileGrid: + """Build a :class:`TileGrid` covering ``(h, w)`` with the given tile size. Args: h: Video height in pixels. w: Video width in pixels. - resolution_scale: Divisor for tile size when ``tile_hw`` is ``None``. - overlap: Fraction of tile overlap in ``[0, 1)``. Default ``0.5`` (50%). - tile_hw: Explicit ``(tile_h, tile_w)`` override. When set, - ``resolution_scale`` is ignored for tile sizing. + tile_hw: Explicit ``(tile_h, tile_w)``. + min_overlap: Overlap floor as a fraction of the tile size. + snap: Position alignment unit (the VAE spatial factor). Returns: - ``TileGrid`` with tile sizes and per-axis tile start positions. + ``TileGrid`` with evenly distributed, aligned tile positions. """ - if tile_hw is None: - tile_h = h // resolution_scale - tile_w = w // resolution_scale - else: - tile_h, tile_w = tile_hw - stride_h = max(1, int(tile_h * (1.0 - overlap))) - stride_w = max(1, int(tile_w * (1.0 - overlap))) - - tops = axis_positions(h, tile_h, stride_h) - lefts = axis_positions(w, tile_w, stride_w) - return TileGrid(tile_h=tile_h, tile_w=tile_w, tops=tops, lefts=lefts) + tile_h, tile_w = tile_hw + return TileGrid( + tile_h=tile_h, + tile_w=tile_w, + tops=axis_positions(h, tile_h, min_overlap, snap), + lefts=axis_positions(w, tile_w, min_overlap, snap), + ) -def extract_all_tiles( - video: torch.Tensor, - grid: TileGrid, -) -> list[torch.Tensor]: +def extract_all_tiles(video: torch.Tensor, grid: TileGrid) -> list[torch.Tensor]: """Extract all spatial tiles from a video tensor. Args: - video: ``[T, C, H, W]`` tensor. + video: ``[C, T, H, W]`` tensor (only the last two dims are read). grid: Tile grid parameters from ``compute_tile_grid``. Returns: - List of ``[T, C, tile_h, tile_w]`` tensors in row-major order + List of ``[C, T, tile_h, tile_w]`` tensors in row-major order (``tops`` x ``lefts``). """ tiles: list[torch.Tensor] = [] @@ -325,75 +252,8 @@ def stitch_tiles_hanning( return pred_acc / weight_acc.clamp(min=1e-6) -def axis_positions_even(length: int, tile: int, min_overlap: float, snap: int) -> tuple[int, ...]: - """Return evenly distributed, ``snap``-aligned tile start positions. - - Uses the minimal position count whose uniform stride keeps the tile - overlap at or above ``min_overlap``, then spreads the positions evenly - over ``[0, length - tile]`` in integer ``snap`` units. When ``length`` or - ``tile`` is not ``snap``-aligned the same layout is computed at pixel - precision (``snap=1``) — the latent-grid guard downstream still enforces - alignment where it actually matters (the LU path). - - Args: - length: Axis length in pixels. - tile: Tile size on this axis. - min_overlap: Overlap floor as a fraction of ``tile`` in ``[0, 1)``. - snap: Position alignment unit (the VAE spatial factor). - - Returns: - Strictly increasing positions covering ``[0, length - tile]``. - - Raises: - ValueError: On non-positive ``tile``/``snap`` or ``min_overlap`` - outside ``[0, 1)``. - """ - if tile <= 0 or snap <= 0 or not 0 <= min_overlap < 1: - msg = f"invalid axis spec: tile={tile}, snap={snap}, min_overlap={min_overlap}" - raise ValueError(msg) - if tile >= length: - return (0,) - unit = snap if length % snap == 0 and tile % snap == 0 else 1 - span_units = (length - tile) // unit - max_stride_units = max(1, math.floor(tile * (1.0 - min_overlap) / unit)) - count = math.ceil(span_units / max_stride_units) + 1 - return tuple(round(i * span_units / (count - 1)) * unit for i in range(count)) - - -def compute_tile_grid_even( - h: int, - w: int, - tile_hw: tuple[int, int], - min_overlap: float, - snap: int, -) -> TileGrid: - """Build a :class:`TileGrid` with the even per-axis layout. - - Drop-in for ``compute_tile_grid(..., tile_hw=...)``: same ``TileGrid`` - contract (``tops``/``lefts`` are authoritative), only the position layout - differs — see :func:`axis_positions_even`. - - Args: - h: Video height in pixels. - w: Video width in pixels. - tile_hw: Explicit ``(tile_h, tile_w)``. - min_overlap: Overlap floor as a fraction of the tile size. - snap: Position alignment unit (the VAE spatial factor). - - Returns: - ``TileGrid`` with evenly distributed, aligned tile positions. - """ - tile_h, tile_w = tile_hw - return TileGrid( - tile_h=tile_h, - tile_w=tile_w, - tops=axis_positions_even(h, tile_h, min_overlap, snap), - lefts=axis_positions_even(w, tile_w, min_overlap, snap), - ) - - def resolve_scale_request(scale: float) -> tuple[int, float]: - """Map a requested ``--resolution-scale`` to ``(tiling_scale, pre_upscale)``. + """Map a requested ``resolution_scale`` to ``(tiling_scale, pre_upscale)``. Args: scale: The requested total upscale — ``2``, ``4`` or ``2.25``. @@ -403,7 +263,7 @@ def resolve_scale_request(scale: float) -> tuple[int, float]: integer scales. """ requested = float(scale) - if requested == 2.25: # noqa: PLR2004 — the one supported fractional total + if requested == 2.25: return 2, 1.125 if requested in (2.0, 4.0): return int(requested), 1.0 @@ -411,7 +271,7 @@ def resolve_scale_request(scale: float) -> tuple[int, float]: def pre_upscale_video(video: torch.Tensor, factor: float, spatial_multiple: int) -> torch.Tensor: - """Bilinear-upscale a ``[T, C, H, W]`` uint8 video by ``factor`` in pixel space. + """Bilinear-upscale a ``[C, T, H, W]`` uint8 video by ``factor`` in pixel space. Target dims are rounded to the nearest multiple of ``spatial_multiple`` (the VAE spatial factor) so the whole-video encode and the latent tile @@ -421,15 +281,15 @@ def pre_upscale_video(video: torch.Tensor, factor: float, spatial_multiple: int) scale is exact. Args: - video: ``[T, C, H, W]`` uint8 source video. + video: ``[C, T, H, W]`` uint8 source video (only the last two dims are read). factor: Pixel upscale factor (> 1). spatial_multiple: VAE spatial factor to align the target dims to. Returns: - ``[T, C, H', W']`` uint8 video with ``H' ~= H * factor`` aligned. + ``[C, T, H', W']`` uint8 video with ``H' ~= H * factor`` aligned. """ - if video.ndim != 4: # noqa: PLR2004 - raise ValueError(f"video must have rank 4 [T,C,H,W], got {tuple(video.shape)}") + if video.ndim != 4: + raise ValueError(f"video must have rank 4 [C,T,H,W], got {tuple(video.shape)}") if factor <= 0 or spatial_multiple <= 0: raise ValueError("factor and spatial_multiple must be positive") height, width = video.shape[-2:] @@ -444,17 +304,12 @@ def pre_upscale_video(video: torch.Tensor, factor: float, spatial_multiple: int) return resized.round_().clamp_(0, 255).to(torch.uint8) -def get_visual_size(x: torch.Tensor) -> int: +def get_visual_size(x: torch.Tensor, spatial_factor: int) -> int: """Return the resolution key matching a visual latent tensor's spatial dims. - Scales the tensor's spatial dims by the VAE spatial factor and looks them - up in the ``RESOLUTIONS`` registry. Reads ``constants.VAE_SPATIAL_FACTOR`` - late-bound (via attribute access on the module) so a caller can override - ``constants.VAE_SPATIAL_FACTOR`` for a different VAE — a ``from``-import - would freeze the import-time default instead. - Args: x: Visual latent tensor of shape ``(T, H, W, C)``. + spatial_factor: The VAE's spatial compression factor. Returns: The matching resolution key. @@ -462,7 +317,7 @@ def get_visual_size(x: torch.Tensor) -> int: Raises: ValueError: If tensor dimensions do not match any known resolution. """ - actual_size = (x.shape[1] * constants.VAE_SPATIAL_FACTOR, x.shape[2] * constants.VAE_SPATIAL_FACTOR) + actual_size = (x.shape[1] * spatial_factor, x.shape[2] * spatial_factor) for key, value in RESOLUTIONS.items(): if actual_size in value: return key @@ -495,113 +350,12 @@ def degrade_lq_latent( """ if noise_scale <= 0: return lq_latent - eps = torch.randn(lq_latent.shape, device=lq_latent.device, dtype=lq_latent.dtype, generator=generator) + eps = randn_tensor(lq_latent.shape, generator=generator, device=lq_latent.device, dtype=lq_latent.dtype) if noise_type == "ddpm": return (1 - noise_scale**2) ** 0.5 * lq_latent + noise_scale * eps return (1 - noise_scale) * lq_latent + noise_scale * eps -def get_sparse_params( - model: Any, - visual: torch.Tensor, - visual_cu_seqlens: torch.Tensor, -) -> dict[str, Any] | None: - """Build sparse-attention parameters for the given visual tensor and model. - - Supports ``"nabla"`` and ``"nabla_framewise_causal"`` attention types; - returns ``None`` for dense attention. The target ``P`` is used directly (the - training-time linear warmup schedule does not apply at inference). - - Args: - model: Model with ``patch_size`` and ``attention_params`` attributes. - visual: Visual latent tensor of shape ``(T, H, W, C)``. - visual_cu_seqlens: Cumulative frame counts for each sample in the batch. - - Returns: - A dict of sparse attention parameters, or ``None`` for dense attention. - - Raises: - ValueError: If ``model.patch_size[0]`` is not 1. - """ - if model.patch_size[0] != 1: - msg = f"Expected model.patch_size[0] == 1, got {model.patch_size[0]}" - raise ValueError(msg) - t, h, w, _ = visual.shape - t, h, w = ( - t // model.patch_size[0], - h // model.patch_size[1], - w // model.patch_size[2], - ) - visual_size = get_visual_size(visual) - attention_configs = model.attention_params - try: - attention_params = attention_configs[visual_size] - except KeyError: - # JSON object keys are always strings, while the native YAML config - # uses integer resolution keys. - attention_params = attention_configs[str(visual_size)] - - def config_value(name: str, default: Any = None) -> Any: - if isinstance(attention_params, Mapping): - return attention_params.get(name, default) - return getattr(attention_params, name, default) - - if config_value("type") == "nabla": - return { - "attention_type": config_value("type"), - "to_fractal": True, - "P": config_value("P"), - "wT": config_value("wT"), - "wW": config_value("wW"), - "wH": config_value("wH"), - "add_sta": config_value("add_sta"), - "visual_shape": (t, h, w), - "visual_seqlens": visual_cu_seqlens, - "method": config_value("method", "topcdf"), - } - if config_value("type") == "nabla_framewise_causal": - return { - "attention_type": config_value("type"), - "to_fractal": True, - "P": config_value("P"), - "wT": config_value("wT"), - "wW": config_value("wW"), - "wH": config_value("wH"), - "add_sta": config_value("add_sta"), - "mf": config_value("mf"), - "visual_shape": (t, h, w), - "visual_seqlens": visual_cu_seqlens, - } - return None - - -def replicate_cached_embeds( - cached: dict[str, torch.Tensor], - bs: int, - device: str | int, -) -> tuple[dict[str, torch.Tensor], torch.Tensor]: - """Replicate bs=1 cached text embeddings for an arbitrary batch size. - - Args: - cached: Dict with ``text_embeds``, ``pooled_embed``, ``cu_seqlens`` - produced by ``save_empty_text_embeddings.py`` for a single sample. - bs: Target batch size. - device: Target device. - - Returns: - Tuple of (text_embeds dict, cu_seqlens tensor) for the full batch. - """ - single_text = cached["text_embeds"] - single_pooled = cached["pooled_embed"] - seq_len = int(cached["cu_seqlens"][-1].item()) - - text_embeds = single_text.repeat(bs, 1).to(device) - pooled_embed = single_pooled.repeat(bs, 1).to(device) - cu_seqlens = seq_len * torch.arange(bs + 1, dtype=torch.int32, device=device) - - return {"text_embeds": text_embeds, "pooled_embed": pooled_embed}, cu_seqlens - - def cast_to_module_dtype(module: torch.nn.Module, value: torch.Tensor) -> torch.Tensor: """Cast floating-point inputs to the module's parameter dtype.""" try: @@ -668,1772 +422,108 @@ def denormalize_to_uint8(vae: torch.nn.Module, decoded: torch.Tensor) -> torch.T return (float_pixels * 255.0).to(torch.uint8) -def latent_upscaler_scale(upscaler: nn.Module) -> int: - """Return the loaded LU's fixed spatial upscale factor.""" - raw = str(getattr(upscaler, "target_scale", "4x")) - if raw not in SUPPORTED_TARGET_SCALES: - msg = f"Unexpected latent upscaler target_scale={raw!r}; expected one of {SUPPORTED_TARGET_SCALES}." - raise ValueError(msg) - return int(raw.removesuffix("x")) - - -def run_latent_upscaler(upscaler: nn.Module, z: torch.Tensor) -> torch.Tensor: - """Forward a scaled LQ latent through its configured LU entry.""" - target_scale = getattr(upscaler, "target_scale", None) - if target_scale is not None: - entry = f"x{latent_upscaler_scale(upscaler)}" - try: - return upscaler(z, entry=entry, return_intermediates=False) - except TypeError: - # Flat upscalers do not accept the multi-scale keyword arguments. - return upscaler(z) - return upscaler(z) - - -@torch.no_grad() -def get_model_prediction( - dit: torch.nn.Module, - x: torch.Tensor, - t: torch.Tensor, - text_embeds: dict[str, torch.Tensor], - null_text_embeds: dict[str, torch.Tensor], - visual_cu_seqlens: torch.Tensor, - text_cu_seqlens: torch.Tensor, - null_text_cu_seqlens: torch.Tensor, - visual_rope_pos: list[torch.Tensor], - text_rope_pos: torch.Tensor, - null_text_rope_pos: torch.Tensor, - scale_factor: tuple[float, ...], - guidance_weight: float, - sparse_params: Any | None = None, -) -> torch.Tensor: - """Compute classifier-free guidance model prediction. +def encode_lq_video_to_lr_latent(lq_video: torch.Tensor, vae: torch.nn.Module, device: str | torch.device) -> Tensor: + """Encode one whole ``[C, T, H, W]`` pixel video into a raw ``[C, T', H', W']`` latent. - The returned tensor has the same semantics as the raw model output - (velocity when ``prediction_target="velocity"``, clean ``x_0`` when - ``prediction_target="x0"``). The caller is responsible for converting - the prediction to a velocity if needed for the Euler step. + Unlike :func:`Kandinsky6SRPipeline._encode_lq_videos` (a batch of spatial + tiles, channels-last, pre-scaled for the latent upscaler), this encodes a + single un-tiled video and leaves the result unscaled — the caller tiles it + in latent space and applies ``scaling_factor`` per tile (see + ``_upscale_lr_latent_tile``). Args: - dit: DiT model used for inference. - x: Visual latent tensor to denoise. - t: Timestep tensor broadcast over the batch. - text_embeds: Dict with ``text_embeds`` and ``pooled_embed`` keys. - null_text_embeds: Same structure as ``text_embeds`` but for null captions. - visual_cu_seqlens: Cumulative sequence lengths for visual tokens. - text_cu_seqlens: Cumulative sequence lengths for text tokens. - null_text_cu_seqlens: Cumulative sequence lengths for null text tokens. - visual_rope_pos: RoPE position indices for visual tokens. - text_rope_pos: RoPE position indices for text tokens. - null_text_rope_pos: RoPE position indices for null text tokens. - scale_factor: Per-axis RoPE frequency scaling. - guidance_weight: CFG strength (1 = no guidance). - sparse_params: Optional sparse attention parameters. + lq_video: ``[C, T, H, W]`` uint8 video in ``[0, 255]``. + vae: Causal video KVAE. + device: Target device. Returns: - CFG-combined model prediction (velocity or ``x_0``). - """ - model_time = (t * 1000).to(dtype=x.dtype) - motion_score = torch.full((1,), 900.0, device=x.device, dtype=x.dtype) - cond_pred = dit( - x, - text_embeds["text_embeds"], - text_embeds["pooled_embed"], - model_time, - visual_cu_seqlens, - text_cu_seqlens, - visual_rope_pos, - text_rope_pos, - scale_factor=scale_factor, - sparse_params=sparse_params, - motion_score=motion_score, - ) - if guidance_weight == 1.0: - return cond_pred - uncond_pred = dit( - x, - null_text_embeds["text_embeds"], - null_text_embeds["pooled_embed"], - model_time, - visual_cu_seqlens, - null_text_cu_seqlens, - visual_rope_pos, - null_text_rope_pos, - scale_factor=scale_factor, - sparse_params=sparse_params, - motion_score=motion_score, - ) - return uncond_pred + guidance_weight * (cond_pred - uncond_pred) - - -@torch.no_grad() -def generate( - img: torch.Tensor, - model: torch.nn.Module, - device: str | int, - num_steps: int, - text_embeds: dict[str, torch.Tensor], - null_text_embeds: dict[str, torch.Tensor], - visual_cu_seqlens: torch.Tensor, - text_cu_seqlens: torch.Tensor, - null_text_cu_seqlens: torch.Tensor, - visual_rope_pos: list[torch.Tensor], - text_rope_pos: torch.Tensor, - null_text_rope_pos: torch.Tensor, - scale_factor: tuple[float, ...], - guidance_weight: float, - scheduler_scale: float, - first_frames: torch.Tensor | None = None, - tp_mesh: Any | None = None, - visual_cond_scheme: str = "pretrain", - start_timestep: float = 1.0, - prediction_target: str = "velocity", - channelcat_drop_threshold: float = 0.0, - rfg_scale: float = 1.0, - progress_callback: Callable[[int], Any] | None = None, - progress_bar: Any = None, -) -> torch.Tensor: - """Run the full denoising loop from start_timestep to 0. + Unscaled latent ``[C, T', H', W']``. - When ``prediction_target="x0"`` the model predicts clean ``x_0`` instead - of velocity. The velocity for the Euler step is recovered as - ``v = (x_t - x_0_pred) / clamp(t, min=X0_T_CLAMP)`` to avoid numerical - instability as ``t → 0``. - - Args: - img: Initial latent tensor ``[T, H, W, C]`` (or wider for instruct). - model: DiT model. - device: Target CUDA device. - num_steps: Number of denoising steps. - text_embeds: Conditional text embeddings. - null_text_embeds: Unconditional text embeddings. - visual_cu_seqlens: Cumulative visual sequence lengths. - text_cu_seqlens: Cumulative text sequence lengths. - null_text_cu_seqlens: Cumulative null-text sequence lengths. - visual_rope_pos: RoPE position indices for visual tokens. - text_rope_pos: RoPE position indices for text tokens. - null_text_rope_pos: RoPE position indices for null text tokens. - scale_factor: Per-axis RoPE frequency scaling. - guidance_weight: CFG strength. - scheduler_scale: Timestep scheduler warp factor. - first_frames: Optional I2V conditioning first frames. - tp_mesh: Tensor-parallel device mesh (optional). - visual_cond_scheme: ``"pretrain"`` or ``"i2v"`` injection strategy. - start_timestep: Starting timestep (``1.0`` = full noise). - prediction_target: ``"velocity"`` or ``"x0"``. - channelcat_drop_threshold: When ``> 0``, zero out channel-cat LQ - channels (and mask) once the denoising timestep drops below this - value. ``0`` = disabled (default). - rfg_scale: Reference-Free Guidance strength on the anchor. Only - applied when ``model.instruct_type == "hybrid_anchor"`` and - ``rfg_scale != 1.0`` — in that case a second forward is run with - the anchor and anchor_mask channels zeroed, and predictions blend - as ``v_uncond + rfg_scale * (v_cond - v_uncond)``. ``1.0`` - (default) skips the second forward entirely. - Returns: - Denoised latent tensor (same leading dims as ``img``). + Raises: + ValueError: If ``lq_video`` is not rank 4. """ - img = img.to(device) - visual_cu_seqlens = visual_cu_seqlens.to(device) - text_cu_seqlens = text_cu_seqlens.to(device) - null_text_cu_seqlens = null_text_cu_seqlens.to(device) - visual_rope_pos = [position.to(device) for position in visual_rope_pos] - text_rope_pos = text_rope_pos.to(device) - null_text_rope_pos = null_text_rope_pos.to(device) - text_embeds = {key: value.to(device) for key, value in text_embeds.items()} - null_text_embeds = {key: value.to(device) for key, value in null_text_embeds.items()} - sparse_params = get_sparse_params(model, img, visual_cu_seqlens) - timesteps = torch.linspace(start_timestep, 0, num_steps, device=device) - timesteps = scheduler_scale * timesteps / (1 + (scheduler_scale - 1) * timesteps) - if tp_mesh: - tp_rank = tp_mesh["tp"].get_local_rank() - tp_world_size = tp_mesh["tp"].size() - img = torch.chunk(img, tp_world_size, dim=1)[tp_rank] - if first_frames is not None: - first_frames = torch.chunk(first_frames, tp_world_size, dim=1)[tp_rank] - if model.visual_cond and first_frames is not None: - first_frames = first_frames.to(device=img.device, dtype=img.dtype) - out_channels: int = 0 - lq_cond_channels: torch.Tensor | None = None - if channelcat_drop_threshold > 0 and img.shape[-1] > model.in_visual_dim: - lq_cond_channels = img[..., model.in_visual_dim :].clone() - for timestep, timestep_diff in list(zip(timesteps[:-1], torch.diff(timesteps), strict=False)): - time = timestep.unsqueeze(0).repeat(visual_cu_seqlens.shape[0] - 1) - if lq_cond_channels is not None: - if timestep.item() >= channelcat_drop_threshold: - img[..., model.in_visual_dim :] = lq_cond_channels - else: - img[..., model.in_visual_dim :] = 0 - if model.visual_cond and img.shape[-1] == model.in_visual_dim: - visual_cond = torch.zeros_like(img) - visual_cond_mask = torch.zeros([*img.shape[:-1], 1], dtype=img.dtype, device=img.device) - if first_frames is not None: - if visual_cond_scheme == "pretrain": - visual_cond[visual_cu_seqlens[:-1]] = first_frames - elif visual_cond_scheme == "i2v": - img[visual_cu_seqlens[:-1]] = first_frames - else: - msg = f"unknown visual_cond_scheme={visual_cond_scheme}" - raise ValueError(msg) - visual_cond_mask[visual_cu_seqlens[:-1]] = 1 - model_input = torch.cat([img, visual_cond, visual_cond_mask], dim=-1) - else: - model_input = img - v_cond = get_model_prediction( - model, - model_input, - time, - text_embeds, - null_text_embeds, - visual_cu_seqlens, - text_cu_seqlens, - null_text_cu_seqlens, - visual_rope_pos, - text_rope_pos, - null_text_rope_pos, - scale_factor, - guidance_weight, - sparse_params=sparse_params, - ) - if model.instruct_type == "hybrid_anchor" and rfg_scale != 1.0: - model_input_uncond = model_input.clone() - model_input_uncond[..., model.in_visual_dim :] = 0 - v_uncond = get_model_prediction( - model, - model_input_uncond, - time, - text_embeds, - null_text_embeds, - visual_cu_seqlens, - text_cu_seqlens, - null_text_cu_seqlens, - visual_rope_pos, - text_rope_pos, - null_text_rope_pos, - scale_factor, - guidance_weight, - sparse_params=sparse_params, - ) - model_pred = v_uncond + rfg_scale * (v_cond - v_uncond) - else: - model_pred = v_cond - out_channels = model_pred.shape[-1] - if prediction_target == "x0": - t_clamped = timestep.clamp(min=X0_T_CLAMP) - velocity = (img[..., :out_channels] - model_pred) / t_clamped - else: - velocity = model_pred - img[..., :out_channels] += timestep_diff * velocity - if progress_callback is not None: - progress_callback() - if progress_bar is not None: - progress_bar.update() - progress_bar.refresh() - if model.visual_cond and first_frames is not None and (visual_cond_scheme == "i2v"): - img[visual_cu_seqlens[:-1]] = first_frames - return img[..., :out_channels] + if lq_video.ndim != 4: + raise ValueError(f"lq_video must have rank 4 [C,T,H,W], got {tuple(lq_video.shape)}") + pixel = lq_video.unsqueeze(0).to(device=device) + latent = encode_pixels_to_latent(vae, pixel) + return latent.squeeze(0).float() -def _encode_lq_videos(lq_videos: list[torch.Tensor], vae: torch.nn.Module, device: str | int) -> torch.Tensor: - """VAE-encode a list of LQ videos into a stacked latent tensor. - - Args: - lq_videos: List of ``[T, H, W, 3]`` float32 tensors in ``[0, 255]``. - vae: Pre-trained VAE model (eval mode). - device: Target CUDA device. +@dataclass +class RunConfig: + """Parameters for one tiled SR run.""" - Returns: - Stacked latent tensor ``[sum(T'), H', W', C]`` scaled by - ``vae.config.scaling_factor``. - """ - latents: list[torch.Tensor] = [] - for lq in lq_videos: - lq_input = lq.permute(3, 0, 1, 2).unsqueeze(0).to(device=device, dtype=torch.bfloat16) - lq_latent = encode_pixels_to_latent(vae, lq_input) - lq_latent = lq_latent.squeeze(0).permute(1, 2, 3, 0).float() - lq_latent *= vae.config.scaling_factor - latents.append(lq_latent) - return torch.cat(latents, dim=0) - - -def _build_initial_latent( - dit: torch.nn.Module, - lq_latent: torch.Tensor, - bs: int, - duration: int, - height: int, - width: int, - device: str | int, - seed: int, - lq_noise_scale: float, - lq_noise_type: Literal["linear", "ddpm"], - lq_channel_noise_scale: float, - anchor_latent: torch.Tensor | None = None, - anchor_mask: torch.Tensor | None = None, - *, - anchor_free: bool = False, -) -> torch.Tensor: - """Build the initial latent tensor for the SR denoising loop. + device: str + num_steps: int = 5 + seed: int = 42 + min_overlap: float = 0.20 + tiles_batch_size: int = 1 + resolution_scale: ResolutionScale = 4 - Constructs ``[starting | cond | mask]`` depending on ``instruct_type``: + def __post_init__(self) -> None: + if self.num_steps < 2: + raise ValueError("num_steps must be at least 2") + if not 0.0 <= self.min_overlap < 1.0: + raise ValueError("min_overlap must satisfy 0 <= min_overlap < 1") + if self.tiles_batch_size <= 0: + raise ValueError("tiles_batch_size must be positive") + if self.resolution_scale not in (2, 4): + raise ValueError("resolution_scale must be 2 or 4") - - ``"noise"``: degraded LQ as starting point, zero LQ cond + zero mask. - - ``"channel"``: random Gaussian noise as starting point, LQ cond + ones mask. - - ``"hybrid"``: degraded LQ as starting point, LQ cond + ones mask. - - ``"hybrid_anchor"``: degraded LQ as starting point, sparse HR anchor + - sparse anchor_mask. Requires ``anchor_latent`` and ``anchor_mask``. - Args: - dit: DiT model (reads ``instruct_type`` and ``in_visual_dim``). - lq_latent: Scaled LQ latent ``[bs*T, H, W, C]``. - bs: Batch size. - duration: Number of temporal frames per sample. - height: Latent height. - width: Latent width. - device: Target CUDA device. - seed: RNG seed for noise generation. - lq_noise_scale: Noise fraction mixed into the LQ starting point. - lq_noise_type: ``"linear"`` or ``"ddpm"`` noising formula. - lq_channel_noise_scale: Noise fraction mixed into the channel-cat - conditioning (LQ for ``channel``/``hybrid``, anchor for - ``hybrid_anchor``). - anchor_latent: Pre-scaled sparse HR anchor ``[bs*T, H, W, C]`` (zeros - in slots without anchor). Required when - ``instruct_type == "hybrid_anchor"``. - anchor_mask: Sparse binary mask ``[bs*T, H, W, 1]`` flagging anchor - slots. Required when ``instruct_type == "hybrid_anchor"``. +def _closest_base_resolution(h: int, w: int, visual_size: int) -> tuple[int, int]: + if visual_size not in RESOLUTIONS: + raise ValueError(f"Unsupported SR visual_size={visual_size}; known sizes: {sorted(RESOLUTIONS)}") + ratio = w / h if h else 1.0 + return min(RESOLUTIONS[visual_size], key=lambda hw: abs(hw[1] / hw[0] - ratio)) - Returns: - Initial latent tensor ready for the denoising loop (bs*T, H, W, 33). - Raises: - ValueError: If ``instruct_type`` is not supported, or if - ``hybrid_anchor`` is missing ``anchor_latent`` / ``anchor_mask``. - """ - lq_latent = lq_latent.to(device) - lq_latent = cast_to_module_dtype(dit, lq_latent) - if anchor_latent is not None: - anchor_latent = anchor_latent.to(device) - anchor_latent = cast_to_module_dtype(dit, anchor_latent) - if anchor_mask is not None: - anchor_mask = anchor_mask.to(device) - anchor_mask = cast_to_module_dtype(dit, anchor_mask) - if dit.instruct_type == "noise": - g_noise = torch.Generator(device=device) - g_noise.manual_seed(seed) - degraded_lq = degrade_lq_latent(lq_latent, lq_noise_scale, lq_noise_type, generator=g_noise) - if not getattr(dit, "visual_cond", False): - return degraded_lq - zero_cond = torch.zeros_like(degraded_lq) - zero_mask = torch.zeros([*degraded_lq.shape[:-1], 1], dtype=degraded_lq.dtype, device=device) - return torch.cat([degraded_lq, zero_cond, zero_mask], dim=-1) - if dit.instruct_type in ("channel", "hybrid"): - if dit.instruct_type == "hybrid": - g_noise = torch.Generator(device=device) - g_noise.manual_seed(seed) - starting = degrade_lq_latent(lq_latent.clone(), lq_noise_scale, lq_noise_type, generator=g_noise) - else: - g = torch.Generator(device=device) - g.manual_seed(seed) - starting = torch.randn(bs * duration, height, width, dit.in_visual_dim, device=device, generator=g) - channel_lq = degrade_lq_latent(lq_latent, lq_channel_noise_scale, noise_type="linear") - mask = torch.ones_like(lq_latent[..., :1]) - return torch.cat([starting, channel_lq, mask], dim=-1) - if dit.instruct_type == "hybrid_anchor": - if anchor_latent is None or anchor_mask is None: - if not anchor_free: - msg = "hybrid_anchor requires anchor_latent and anchor_mask" - raise ValueError(msg) - anchor_latent = torch.zeros_like(lq_latent) - anchor_mask = torch.zeros_like(lq_latent[..., :1]) - g_noise = torch.Generator(device=device) - g_noise.manual_seed(seed) - starting = degrade_lq_latent(lq_latent.clone(), lq_noise_scale, lq_noise_type, generator=g_noise) - anchor = degrade_lq_latent(anchor_latent, lq_channel_noise_scale, noise_type="linear") - return torch.cat([starting, anchor, anchor_mask], dim=-1) - msg = f"generate_sample_sr does not support instruct_type={dit.instruct_type!r}" - raise ValueError(msg) +def _tile_geometry( + h: int, + w: int, + visual_size: int, + scale: int, + min_overlap: float, + spatial_factor: int = VAE_SPATIAL_FACTOR, +) -> tuple[tuple[int, int], tuple[int, int], TileGrid]: + base_h, base_w = _closest_base_resolution(h, w, visual_size) + if base_h % scale or base_w % scale: + raise ValueError(f"resolution_scale={scale} does not divide the base resolution {base_h}x{base_w} exactly") + tile_hw = (base_h // scale, base_w // scale) + grid = compute_tile_grid(h, w, tile_hw, min_overlap, spatial_factor) + return (base_h, base_w), tile_hw, grid -def _empty_text_embeds(bs: int, device: str | int) -> tuple[dict[str, torch.Tensor], torch.Tensor]: - """Build placeholder text embeddings for a text-free DiT. +def latent_tile_grid_from_pixel_grid(pixel_grid: TileGrid, spatial_factor: int) -> TileGrid: + """Convert an aligned pixel grid to the corresponding latent grid.""" + values = (pixel_grid.tile_h, pixel_grid.tile_w, *pixel_grid.tops, *pixel_grid.lefts) + if any(value % spatial_factor for value in values): + raise ValueError( + f"Pixel tile grid (tile {pixel_grid.tile_h}x{pixel_grid.tile_w}, " + f"tops={pixel_grid.tops}, lefts={pixel_grid.lefts}) not aligned to VAE spatial factor " + f"{spatial_factor}." + ) + return TileGrid( + pixel_grid.tile_h // spatial_factor, + pixel_grid.tile_w // spatial_factor, + tuple(top // spatial_factor for top in pixel_grid.tops), + tuple(left // spatial_factor for left in pixel_grid.lefts), + ) - A ``use_text=False`` DiT ignores every text input (encoder, cross-attention - and pooled-text conditioning are all absent), so no real embeddings are - needed. These zero-size tensors only satisfy the indexing and ``cu_seqlens`` - bookkeeping in ``generate_sample_sr`` / ``get_model_prediction`` without ever - loading the cached empty-caption file. - Args: - bs: Batch size. - device: Target CUDA device. - - Returns: - Tuple of ``(text_embeds_dict, text_cu_seqlens)`` with empty sequences. - """ - embeds = {"text_embeds": torch.zeros(0, 1, device=device), "pooled_embed": torch.zeros(bs, 1, device=device)} - cu_seqlens = torch.zeros(bs + 1, dtype=torch.int32, device=device) - return (embeds, cu_seqlens) - - -def _encode_text( - bs: int, - device: str | int, - text_embedder: Any | None = None, - cached_text_embeds: dict[str, torch.Tensor] | None = None, - *, - use_text: bool = True, -) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor], torch.Tensor]: - """Encode text or use cached embeddings for SR generation. - - For SR we always use empty captions. Returns both text and null-text - embeddings (identical when using cached embeds). - - Args: - bs: Batch size. - device: Target CUDA device. - text_embedder: Text encoder (Qwen2.5-VL + CLIP). Can be ``None`` - when ``cached_text_embeds`` is provided. - cached_text_embeds: Pre-computed empty-caption embeddings (bs=1) - loaded from disk. - use_text: Whether the DiT consumes text. When ``False`` placeholder - zero embeddings are returned and neither a text encoder nor cached - embeddings are required. - - Returns: - Tuple of ``(text_embed, text_cu_seqlens, null_text_embed, - null_text_cu_seqlens)``. - - Raises: - ValueError: If a text DiT is given neither ``text_embedder`` nor - ``cached_text_embeds``. - """ - if not use_text: - embeds, cu_seqlens = _empty_text_embeds(bs, device) - return (embeds, cu_seqlens, embeds, cu_seqlens) - if cached_text_embeds is not None: - bs_text_embed, text_cu_seqlens = replicate_cached_embeds(cached_text_embeds, bs, device) - return (bs_text_embed, text_cu_seqlens, bs_text_embed, text_cu_seqlens) - if text_embedder is not None: - empty_captions = [""] * bs - bs_text_embed, text_cu_seqlens = text_embedder.encode(empty_captions, images=None, type_of_content="video") - bs_null_text_embed, null_text_cu_seqlens = text_embedder.encode( - empty_captions, images=None, type_of_content="video" - ) - bs_text_embed = {k: v.to(device) for k, v in bs_text_embed.items()} - text_cu_seqlens = text_cu_seqlens.to(device) - bs_null_text_embed = {k: v.to(device) for k, v in bs_null_text_embed.items()} - null_text_cu_seqlens = null_text_cu_seqlens.to(device) - return (bs_text_embed, text_cu_seqlens, bs_null_text_embed, null_text_cu_seqlens) - msg = "generate_sample_sr requires either text_embedder or cached_text_embeds" - raise ValueError(msg) - - -@torch.no_grad() -def generate_sample_sr( - *, - dit: torch.nn.Module, - vae: torch.nn.Module, - scale_factor: tuple[float, ...] = (1.0, 1.0, 1.0), - num_steps: int = 50, - guidance_weight: float = 5.0, - scheduler_scale: float = 5.0, - seed: int = 42, - device: str | int = "cuda", - tp_mesh: dict[str, Any] | None = None, - lq_noise_scale: float = 0.0, - lq_noise_type: Literal["linear", "ddpm"] = "linear", - lq_channel_noise_scale: float = 0.0, - text_embedder: Any | None = None, - cached_text_embeds: dict[str, torch.Tensor] | None = None, - lq_videos: list[torch.Tensor] | None = None, - lq_latents: torch.Tensor | None = None, - n_samples: int | None = None, - cap_noise_timestep: bool = False, - vae_decode_batch: bool = False, - prediction_target: str = "velocity", - channelcat_drop_threshold: float = 0.0, - anchor_latents: torch.Tensor | None = None, - anchor_masks: torch.Tensor | None = None, - anchor_free: bool = False, - rfg_scale: float = 1.0, -) -> torch.Tensor: - """Generate super-resolved video from LQ inputs (pixels or latents). - - Accepts either raw pixel videos (``lq_videos``) or pre-encoded latents - (``lq_latents``). When ``lq_latents`` is provided, VAE encoding is - skipped and the latents are used directly. - - Builds the initial latent tensor, runs the denoising loop via - ``generate()``, and decodes back to pixel space. - - Supports three SR instruct modes: - - - ``instruct_type="channel"``: random noise in HQ channels with LQ latent - and mask concatenated as extra channels → ``[noise(C) | lq(C) | mask(1)]``. - - ``instruct_type="noise"``: LQ latent is used directly as the starting - point (replaces random noise) → ``[lq(C)]``, zero extra channels. - - ``instruct_type="hybrid"``: LQ latent as starting point (like ``"noise"``) - AND LQ latent concatenated as extra channels (like ``"channel"``) → - ``[degraded_lq(C) | lq(C) | mask(1)]``. - - When ``lq_noise_scale > 0``, noise is mixed into the LQ latent used as - the starting point (applies to ``"noise"`` and ``"hybrid"``). - - When ``lq_channel_noise_scale > 0``, noise is mixed into the LQ latent - concatenated as extra channels (applies to ``"channel"`` and ``"hybrid"``). - - Args: - dit: DiT model with ``instruct_type`` in - ``{"channel", "noise", "hybrid"}``. - vae: Pre-trained VAE model. - scale_factor: Per-axis RoPE frequency scaling. - num_steps: Number of diffusion denoising steps. - guidance_weight: Classifier-free guidance strength. - scheduler_scale: Timestep scheduler scaling factor. - seed: RNG seed for noise initialization. - device: Target CUDA device. - tp_mesh: Tensor parallelism mesh (optional). - lq_noise_scale: Noise fraction ``s`` mixed into LQ latent used as - the noise/starting point (``0`` = disabled). Applies to - ``"noise"`` and ``"hybrid"`` modes. - lq_noise_type: Noising formula — ``"linear"`` or ``"ddpm"``. - lq_channel_noise_scale: Noise fraction mixed into LQ latent - concatenated as extra channels (``0`` = disabled). Applies to - ``"channel"`` and ``"hybrid"`` modes. - text_embedder: Text encoder (Qwen2.5-VL + CLIP). Can be ``None`` - when ``cached_text_embeds`` is provided. - cached_text_embeds: Pre-computed empty-caption embeddings (bs=1) - loaded from disk. When provided, ``text_embedder`` is not called. - lq_videos: List of ``[T, H, W, 3]`` float32 tensors in ``[0, 255]``. - lq_latents: Pre-encoded LQ latents (skips VAE encoding). - n_samples: Required when ``lq_latents`` is provided. - cap_noise_timestep: When ``True`` and ``instruct_type`` in - ``("noise", "hybrid")``, starts the denoising loop from raw - ``t = lq_noise_scale`` instead of ``t = 1.0``. Matches the - training-time behaviour when ``cap_noise_timestep`` is also set - during training. Default ``False`` preserves current behaviour. - vae_decode_batch: When ``True``, decode all samples in a single VAE - call instead of one-by-one. Faster but uses more GPU memory. - prediction_target: ``"velocity"`` for velocity prediction (default), - ``"x0"`` for clean-image prediction. - channelcat_drop_threshold: When ``> 0``, zero out channel-cat LQ - channels during denoising once ``t < threshold``. ``0`` = - disabled (default). - anchor_latents: Pre-encoded sparse HR-anchor latents - ``[bs*T', H', W', C]``, **already scaled** by VAE - ``scaling_factor``. Required for ``instruct_type="hybrid_anchor"``. - anchor_masks: Sparse binary anchor mask ``[bs*T', H', W', 1]``. - Required for ``instruct_type="hybrid_anchor"``. - anchor_free: When ``True`` and ``instruct_type="hybrid_anchor"`` with - no anchors provided, run anchor-free (zeroed anchor + mask) — the - tiled path for GT-less real-world LQ videos. - rfg_scale: Reference-Free Guidance strength on the anchor. ``1.0`` - (default) runs a single cond forward per denoising step; values - ``!= 1.0`` add a second uncond forward (anchor + anchor_mask - zeroed) and blend ``v_uncond + rfg_scale * (v_cond - v_uncond)``. - Only applies when ``instruct_type="hybrid_anchor"``. Cost: +1 - DiT forward per denoising step when active. - Returns: - Generated SR videos as ``[bs, 3, T, H, W]`` uint8 tensor. - - Raises: - ValueError: If neither ``text_embedder`` nor ``cached_text_embeds`` - is provided or if ``instruct_type`` is unsupported. - """ - piflow_params = getattr(dit, "piflow_params", None) - if piflow_params is not None: - if rfg_scale != 1.0: - msg = "π-Flow checkpoints do not support Reference-Free Guidance (drop rfg_scale)" - raise ValueError(msg) - return run_stages( - dit=dit, - vae=vae, - scale_factor=scale_factor, - num_steps=num_steps, - guidance_weight=guidance_weight, - scheduler_scale=scheduler_scale, - seed=seed, - device=device, - tp_mesh=tp_mesh, - lq_noise_scale=lq_noise_scale, - lq_noise_type=lq_noise_type, - lq_channel_noise_scale=lq_channel_noise_scale, - text_embedder=text_embedder, - cached_text_embeds=cached_text_embeds, - lq_videos=lq_videos, - lq_latents=lq_latents, - n_samples=n_samples, - cap_noise_timestep=cap_noise_timestep, - vae_decode_batch=vae_decode_batch, - prediction_target=prediction_target, - channelcat_drop_threshold=channelcat_drop_threshold, - anchor_latents=anchor_latents, - anchor_masks=anchor_masks, - anchor_free=anchor_free, - rfg_scale=rfg_scale, - piflow_params=piflow_params, - ) - - -@torch.no_grad() -def piflow_generate( - img: torch.Tensor, - model: torch.nn.Module, - text_embeds: dict[str, torch.Tensor], - visual_cu_seqlens: torch.Tensor, - text_cu_seqlens: torch.Tensor, - visual_rope_pos: list[torch.Tensor], - text_rope_pos: torch.Tensor, - scale_factor: tuple[float, ...], - *, - nfe: int, - num_policy_substeps: int, - final_step_size_scale: float, - shift: float, - n_grid: int, - out_dim: int, - eps: float = 1e-06, - tp_mesh: dict[str, Any] | None = None, - start_timestep: float = 1.0, - device: str | int | None = None, - progress_callback=None, - scheduler: Any | None = None, - progress_bar: Any = None, -) -> torch.Tensor: - """Few-step π-Flow denoise: ``nfe`` network calls, network-free integration between. - - Walks the deterministic ``nfe``-segment schedule from ``raw_t = start_timestep`` - down to 0 (final segment scaled by ``final_step_size_scale``, matching ``sample_t``). - Each segment: one ``model`` forward -> ``DXPolicy`` -> integrate over the segment. - """ - device = device if device is not None else img.device - img = img.to(device) - visual_cu_seqlens = visual_cu_seqlens.to(device) - text_cu_seqlens = text_cu_seqlens.to(device) - visual_rope_pos = [position.to(device) for position in visual_rope_pos] - text_rope_pos = text_rope_pos.to(device) - text_embeds = {key: value.to(device) for key, value in text_embeds.items()} - sparse_params = get_sparse_params(model, img, visual_cu_seqlens) - if tp_mesh: - tp_world_size = tp_mesh["tp"].size() - tp_rank = tp_mesh["tp"].get_local_rank() - img = torch.chunk(img, tp_world_size, dim=1)[tp_rank] - ndim = img.dim() - if nfe < 1: - msg = f"piflow_generate requires nfe >= 1, got {nfe}" - raise ValueError(msg) - if scheduler is not None: - scheduler.set_timesteps(nfe, device=device) - if scheduler.timesteps.shape[0] != nfe: - raise ValueError(f"Piflow scheduler returned {scheduler.timesteps.shape[0]} steps, expected {nfe}") - else: - final_step_size_scale = max(float(final_step_size_scale), eps) - one_minus_final = 1.0 - final_step_size_scale - base_seg = 1.0 / (float(nfe) - one_minus_final) - final_seg = final_step_size_scale * base_seg - x = img - for step_index in range(nfe): - if scheduler is not None: - timestep = scheduler.timesteps[step_index] - sigma_src = scheduler.sigmas[step_index].item() - else: - idx = nfe - step_index - raw_src = min(max((idx - one_minus_final) * base_seg, eps), start_timestep) - seg = final_seg if idx == 1 else base_seg - raw_dst = max(raw_src - seg, 0.0) - sigma_src = shift_timesteps(torch.tensor(raw_src, device=device), shift).item() - n_objects = visual_cu_seqlens.shape[0] - 1 - t_src = torch.full( - (n_objects,), - float(timestep.item()) if scheduler is not None else sigma_src * 1000.0, - device=device, - dtype=x.dtype, - ) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16): - v0 = model( - x, - text_embeds["text_embeds"], - text_embeds["pooled_embed"], - t_src, - visual_cu_seqlens, - text_cu_seqlens, - visual_rope_pos, - text_rope_pos, - scale_factor=scale_factor, - sparse_params=sparse_params, - ) - if scheduler is not None: - x_pred = scheduler.step(v0, timestep, x[..., :out_dim], return_dict=False)[0] - else: - v0_grid = v0.unsqueeze(1) if n_grid == 1 else v0 - total_tokens = x.shape[0] - sigma_tok = torch.full((total_tokens, *(ndim - 1) * [1]), sigma_src, device=device) - seg_tok = torch.full((total_tokens,), seg, device=device) - policy = DXPolicy(v0_grid, x[..., :out_dim], sigma_tok, seg_tok, shift=shift, mode="grid", eps=eps) - raw_src_tok = torch.full((total_tokens,), raw_src, device=device) - raw_dst_tok = torch.full((total_tokens,), raw_dst, device=device) - x_pred, _, _ = policy_rollout_fm( - x[..., :out_dim], sigma_tok, raw_src_tok, raw_dst_tok, num_policy_substeps, policy - ) - x = torch.cat([x_pred, x[..., out_dim:]], dim=-1) - if progress_callback is not None: - progress_callback() - if progress_bar is not None: - progress_bar.update() - progress_bar.refresh() - return x[..., :out_dim] - - -@dataclass -class RunConfig: - """Parameters for one tiled SR run.""" - - device: str - num_steps: int = 5 - seed: int = 42 - overlap: float = 0.25 - tiles_batch_size: int = 1 - resolution_scale: ResolutionScale = 4 - tile_grid_mode: TileGridMode = "even" - tile_grid_min_overlap: float = 0.20 - - def __post_init__(self) -> None: - if self.num_steps < 2: - raise ValueError("num_steps must be at least 2") - if not 0.0 <= self.overlap < 1.0: - raise ValueError("overlap must satisfy 0 <= overlap < 1") - if self.tiles_batch_size <= 0: - raise ValueError("tiles_batch_size must be positive") - if self.resolution_scale not in (2, 4): - raise ValueError("resolution_scale must be 2 or 4") - if self.tile_grid_mode not in ("legacy", "even"): - raise ValueError("tile_grid_mode must be 'legacy' or 'even'") - if not 0.0 <= self.tile_grid_min_overlap < 1.0: - raise ValueError("tile_grid_min_overlap must satisfy 0 <= value < 1") - - -@dataclass -class SRParams: - """SR sampling parameters extracted from the SR training config.""" - - scale_factor: dict[int, list[float]] - visual_size: list[int] - scheduler_scale: float = 5.0 - lq_noise_scale: float = 0.7 - lq_noise_type: str = "ddpm" - lq_channel_noise_scale: float = 0.0 - cap_noise_timestep: bool = False - fps: int = 24 - - -@dataclass -class SRComponents: - """The weighted components and config consumed by the SR sampler.""" - - dit: torch.nn.Module - vae: torch.nn.Module - latent_upscaler: torch.nn.Module | Any | None - sr_params: SRParams - cached_text_embeds: dict[str, torch.Tensor] | None = None - sampler: Callable[..., torch.Tensor] | None = None - spatial_factor: int | None = None - - -def _closest_base_resolution(h: int, w: int, visual_size: int) -> tuple[int, int]: - if visual_size not in RESOLUTIONS: - raise ValueError(f"Unsupported SR visual_size={visual_size}; known sizes: {sorted(RESOLUTIONS)}") - ratio = w / h if h else 1.0 - return min(RESOLUTIONS[visual_size], key=lambda hw: abs(hw[1] / hw[0] - ratio)) - - -def _tile_geometry( # noqa: PLR0913 - h: int, - w: int, - visual_size: int, - scale: int, - overlap: float, - spatial_factor: int = VAE_SPATIAL_FACTOR, - grid_mode: TileGridMode = "even", - grid_min_overlap: float = 0.20, -) -> tuple[tuple[int, int], tuple[int, int], TileGrid]: - base_h, base_w = _closest_base_resolution(h, w, visual_size) - if base_h % scale or base_w % scale: - raise ValueError(f"resolution_scale={scale} does not divide the base resolution {base_h}x{base_w} exactly") - tile_hw = (base_h // scale, base_w // scale) - if grid_mode == "even": - grid = compute_tile_grid_even(h, w, tile_hw, grid_min_overlap, spatial_factor) - else: - grid = compute_tile_grid(h, w, scale, overlap, tile_hw=tile_hw) - return (base_h, base_w), tile_hw, grid - - -def _spatial_factor(components: Any) -> int: - explicit = getattr(components, "spatial_factor", None) - if explicit is not None: - return int(explicit) - vae = getattr(components, "vae", None) - config = getattr(vae, "config", None) - return int( - getattr(vae, "spatial_factor", None) or getattr(config, "spatial_factor", None) or VAE_SPATIAL_FACTOR - ) - - -def _component_params(components: Any) -> Any: - params = getattr(components, "sr_params", None) - if params is None: - raise ValueError("SR components must provide sr_params") - return params - - -def _visual_size(params: Any) -> int: - visual_size = getattr(params, "visual_size", None) - if isinstance(visual_size, int): - return visual_size - if not visual_size: - raise ValueError("SR params must provide a non-empty visual_size") - return int(visual_size[0]) - - -def _scale_factor(params: Any, visual_size: int) -> tuple[float, ...]: - values = getattr(params, "scale_factor", None) - if isinstance(values, Mapping): - try: - values = values[visual_size] - except KeyError: - # YAML keeps numeric resolution keys as integers; JSON converts - # object keys to strings when the config is stored in a bundle. - values = values[str(visual_size)] - return tuple(float(value) for value in values) - - -def _extract_latent(result: Any) -> torch.Tensor: - if isinstance(result, tuple): - result = result[0] - distribution = getattr(result, "latent_dist", None) - if distribution is not None: - result = distribution.sample() - if not isinstance(result, torch.Tensor): - raise TypeError("VAE encode must return a Tensor, (Tensor, ...), or latent_dist output") - return result - - -@torch.no_grad() -def encode_lq_videos_to_lr_latents( - lq_videos: list[torch.Tensor], - vae: torch.nn.Module, - device: str | torch.device, -) -> torch.Tensor: - """Encode an equal-shaped batch of ``[T,C,H,W]`` videos in one VAE call.""" - _validate_batch_shapes(lq_videos, "video") - pixel = torch.stack(lq_videos).permute(0, 2, 1, 3, 4).to(device=device) - pixel = vae.normalize_data(pixel.float()) if hasattr(vae, "normalize_data") else pixel.float() / 127.5 - 1.0 - pixel = cast_to_module_dtype(vae, pixel) - result = _extract_latent(vae.encode(pixel)) - if result.ndim != VIDEO_BATCH_RANK: - raise ValueError(f"batched VAE encode must return rank 5, got shape {tuple(result.shape)}") - return result.permute(0, 2, 1, 3, 4).float() - - -def latent_tile_grid_from_pixel_grid(pixel_grid: TileGrid, spatial_factor: int) -> TileGrid: - """Convert an aligned pixel grid to the corresponding latent grid.""" - values = (pixel_grid.tile_h, pixel_grid.tile_w, *pixel_grid.tops, *pixel_grid.lefts) - if any(value % spatial_factor for value in values): - raise ValueError( - f"Pixel tile grid (tile {pixel_grid.tile_h}x{pixel_grid.tile_w}, " - f"tops={pixel_grid.tops}, lefts={pixel_grid.lefts}) not aligned to VAE spatial factor " - f"{spatial_factor}." - ) - return TileGrid( - pixel_grid.tile_h // spatial_factor, - pixel_grid.tile_w // spatial_factor, - tuple(top // spatial_factor for top in pixel_grid.tops), - tuple(left // spatial_factor for left in pixel_grid.lefts), - ) - - -@torch.no_grad() -def upscale_lr_latent_tile( - lr_latent_tile: torch.Tensor, - latent_upscaler: torch.nn.Module, - vae: torch.nn.Module, - device: str | torch.device, -) -> torch.Tensor: - """Upscale one raw latent tile into the sampler's channels-last layout.""" - scaling_factor = float(getattr(getattr(vae, "config", None), "scaling_factor", 1.0)) - tile = lr_latent_tile.permute(1, 0, 2, 3).unsqueeze(0).to(device=device, dtype=torch.float32) - tile = cast_to_module_dtype(latent_upscaler, tile) - if str(device).startswith("cuda"): - with torch.autocast("cuda", dtype=torch.bfloat16): - upscaled = run_latent_upscaler(latent_upscaler, tile * scaling_factor) - else: - upscaled = run_latent_upscaler(latent_upscaler, tile * scaling_factor) - return upscaled.squeeze(0).permute(1, 2, 3, 0).float() - - -def _upsample_tiles_to_base(raw_tiles: list[torch.Tensor], base_h: int, base_w: int) -> list[torch.Tensor]: - return [ - functional.interpolate(tile.float(), size=(base_h, base_w), mode="bilinear", align_corners=False).permute( - 0, 2, 3, 1 - ) - for tile in raw_tiles - ] - - -# Reusable Kandinsky SR inference stages and tiled orchestration. -# -# The functions in this module deliberately do not depend on a benchmark -# runner. A framework can call ``text_encode`` / ``prepare_latents`` / -# ``denoise`` / ``vae_decode`` directly, or use the tiled functions for the -# scale-aware path. Native benchmark integrations may add timing contexts to -# the tiled calls; the public :class:`Kandinsky6SRPipeline` does not expose -# those callbacks. - - -@dataclass(frozen=True) -class SRLatentState: - """Inputs and initial noisy latent for one SR DiT invocation.""" - - lq_latent: torch.Tensor - image: torch.Tensor - batch_size: int - duration: int - height: int - width: int - - -@dataclass(frozen=True) -class SRTextState: - """Conditional and null text state consumed by SR denoising.""" - - text_embeds: dict[str, torch.Tensor] - text_cu_seqlens: torch.Tensor - null_text_embeds: dict[str, torch.Tensor] - null_text_cu_seqlens: torch.Tensor - - -@torch.no_grad() -def prepare_latents( - *, - dit: torch.nn.Module, - vae: torch.nn.Module, - device: str | int, - lq_videos: list[torch.Tensor] | None = None, - lq_latents: torch.Tensor | None = None, - n_samples: int | None = None, - seed: int = 42, - lq_noise_scale: float = 0.0, - lq_noise_type: str = "linear", - lq_channel_noise_scale: float = 0.0, - anchor_latents: torch.Tensor | None = None, - anchor_masks: torch.Tensor | None = None, - anchor_free: bool = False, -) -> SRLatentState: - """Encode LQ input and build the initial latent for one SR tile batch.""" - if lq_latents is not None: - lq_latent = lq_latents.to(device) - if n_samples is None: - raise ValueError("n_samples is required when lq_latents is provided") - batch_size = n_samples - elif lq_videos is not None: - batch_size = len(lq_videos) - lq_latent = _encode_lq_videos(lq_videos, vae, device) - else: - raise ValueError("Either lq_videos or lq_latents must be provided") - duration = lq_latent.shape[0] // batch_size - height, width = (lq_latent.shape[1], lq_latent.shape[2]) - image = _build_initial_latent( - dit=dit, - lq_latent=lq_latent, - bs=batch_size, - duration=duration, - height=height, - width=width, - device=device, - seed=seed, - lq_noise_scale=lq_noise_scale, - lq_noise_type=lq_noise_type, - lq_channel_noise_scale=lq_channel_noise_scale, - anchor_latent=anchor_latents.to(device) if anchor_latents is not None else None, - anchor_mask=anchor_masks.to(device) if anchor_masks is not None else None, - anchor_free=anchor_free, - ) - return SRLatentState(lq_latent, image, batch_size, duration, height, width) - - -@torch.no_grad() -def text_encode( - *, - dit: torch.nn.Module, - batch_size: int, - device: str | int, - text_embedder: Any | None = None, - cached_text_embeds: dict[str, torch.Tensor] | None = None, -) -> SRTextState: - """Build SR conditional and null text embeddings.""" - text_embeds, text_cu_seqlens, null_text_embeds, null_text_cu_seqlens = _encode_text( - bs=batch_size, - device=device, - text_embedder=text_embedder, - cached_text_embeds=cached_text_embeds, - use_text=getattr(dit, "use_text", True), - ) - return SRTextState(text_embeds, text_cu_seqlens, null_text_embeds, null_text_cu_seqlens) - - -def _positions( - *, dit: torch.nn.Module, latent_state: SRLatentState, text_state: SRTextState -) -> tuple[torch.Tensor, list[torch.Tensor], torch.Tensor, torch.Tensor]: - runtime_device = latent_state.image.device - visual_cu_seqlens = latent_state.duration * torch.arange( - latent_state.batch_size + 1, dtype=torch.int32, device=runtime_device - ) - visual_rope_pos = [ - torch.cat([torch.arange(int(end), device=runtime_device) for end in torch.diff(visual_cu_seqlens).cpu()]), - torch.arange(latent_state.height // dit.patch_size[1], device=runtime_device), - torch.arange(latent_state.width // dit.patch_size[2], device=runtime_device), - ] - text_cu_seqlens = text_state.text_cu_seqlens.to(runtime_device) - null_text_cu_seqlens = text_state.null_text_cu_seqlens.to(runtime_device) - text_rope_pos = torch.cat( - [torch.arange(int(end), device=runtime_device) for end in torch.diff(text_cu_seqlens).cpu()] - ) - null_text_rope_pos = torch.cat( - [torch.arange(int(end), device=runtime_device) for end in torch.diff(null_text_cu_seqlens).cpu()] - ) - return (visual_cu_seqlens, visual_rope_pos, text_rope_pos, null_text_rope_pos) - - -@torch.no_grad() -def denoise( - *, - dit: torch.nn.Module, - latent_state: SRLatentState, - text_state: SRTextState, - scale_factor: tuple[float, ...], - num_steps: int = 50, - guidance_weight: float = 5.0, - scheduler_scale: float = 5.0, - tp_mesh: dict[str, Any] | None = None, - lq_noise_scale: float = 0.0, - cap_noise_timestep: bool = False, - prediction_target: str = "velocity", - channelcat_drop_threshold: float = 0.0, - rfg_scale: float = 1.0, - piflow_params: dict[str, Any] | None = None, - device: str | int | None = None, - progress_callback: Callable[[int], Any] | None = None, - scheduler: Any | None = None, - progress_bar: Any = None, -) -> torch.Tensor: - """Run either the Euler or π-Flow SR denoising stage.""" - piflow_params = piflow_params or _scheduler_piflow_params(scheduler) - piflow_params = piflow_params or getattr(dit, "piflow_params", None) - if piflow_params is not None: - if rfg_scale != 1.0: - raise ValueError("π-Flow checkpoints do not support Reference-Free Guidance") - if cap_noise_timestep and dit.instruct_type in ("noise", "hybrid"): - raise NotImplementedError("piflow sampler does not support cap_noise_timestep for noise/hybrid instruct") - device = device if device is not None else latent_state.image.device - visual_cu_seqlens, visual_rope_pos, text_rope_pos, null_text_rope_pos = _positions( - dit=dit, latent_state=latent_state, text_state=text_state - ) - if piflow_params is not None: - out_dim = int(getattr(dit, "base_out_visual_dim", dit.in_visual_dim)) - start_t = 1.0 - with torch.autocast(device_type="cuda", dtype=torch.bfloat16): - latent_visual = piflow_generate( - latent_state.image, - dit, - text_state.text_embeds, - visual_cu_seqlens, - text_state.text_cu_seqlens, - visual_rope_pos, - text_rope_pos, - scale_factor, - **piflow_params, - out_dim=out_dim, - start_timestep=start_t, - device=device, - progress_callback=progress_callback, - scheduler=scheduler, - progress_bar=progress_bar, - ) - else: - start_t = ( - lq_noise_scale if cap_noise_timestep and dit.instruct_type in ("noise", "hybrid", "hybrid_anchor") else 1.0 - ) - with torch.autocast(device_type="cuda", dtype=torch.bfloat16): - latent_visual = generate( - latent_state.image, - dit, - device, - num_steps, - text_state.text_embeds, - text_state.null_text_embeds, - visual_cu_seqlens, - text_state.text_cu_seqlens, - text_state.null_text_cu_seqlens, - visual_rope_pos, - text_rope_pos, - null_text_rope_pos, - scale_factor, - guidance_weight, - scheduler_scale, - tp_mesh=tp_mesh, - first_frames=None, - start_timestep=start_t, - prediction_target=prediction_target, - channelcat_drop_threshold=channelcat_drop_threshold, - rfg_scale=rfg_scale, - progress_callback=progress_callback, - progress_bar=progress_bar, - ) - if tp_mesh: - tensor_list = [ - torch.zeros_like(latent_visual, device=latent_visual.device) for _ in range(tp_mesh["tp"].size()) - ] - all_gather(tensor_list, latent_visual.contiguous(), group=tp_mesh.get_group(mesh_dim="tp")) - latent_visual = torch.cat(tensor_list, dim=1) - return latent_visual - - -def _scheduler_piflow_params(scheduler: Any | None) -> dict[str, Any] | None: - if scheduler is None or not bool(getattr(scheduler, "is_piflow", False)): - return None - config = scheduler.config - nfe = getattr(config, "nfe", None) - if nfe is None: - raise ValueError("PiflowScheduler used for SR must define nfe in scheduler_config.json") - return { - "nfe": int(nfe), - "num_policy_substeps": int(config.num_policy_substeps), - "final_step_size_scale": float(config.final_step_size_scale), - "shift": float(config.shift), - "n_grid": int(config.n_grid), - "eps": float(config.eps), - } - - -@torch.no_grad() -def vae_decode( - *, - vae: torch.nn.Module, - latent_visual: torch.Tensor, - batch_size: int, - duration: int, - height: int, - width: int, - vae_decode_batch: bool = False, -) -> torch.Tensor: - """Decode denoised SR latents into ``[batch, 3, T, H, W]`` uint8 frames""" - with torch.autocast(device_type="cuda", dtype=torch.bfloat16): - all_latents = latent_visual.reshape(batch_size, duration, height, width, -1) - all_latents = (all_latents / vae.config.scaling_factor).permute(0, 4, 1, 2, 3) - if vae_decode_batch: - return decode_latent_to_uint8(vae, all_latents) - decoded = [decode_latent_to_uint8(vae, all_latents[i : i + 1]) for i in range(batch_size)] - return torch.cat(decoded, dim=0) - - -@torch.no_grad() -def run_stages( - *, - dit: torch.nn.Module, - vae: torch.nn.Module, - scale_factor: tuple[float, ...], - num_steps: int = 50, - guidance_weight: float = 5.0, - scheduler_scale: float = 5.0, - seed: int = 42, - device: str | int = "cuda", - tp_mesh: dict[str, Any] | None = None, - lq_noise_scale: float = 0.0, - lq_noise_type: str = "linear", - lq_channel_noise_scale: float = 0.0, - text_embedder: Any | None = None, - cached_text_embeds: dict[str, torch.Tensor] | None = None, - lq_videos: list[torch.Tensor] | None = None, - lq_latents: torch.Tensor | None = None, - n_samples: int | None = None, - cap_noise_timestep: bool = False, - vae_decode_batch: bool = False, - prediction_target: str = "velocity", - channelcat_drop_threshold: float = 0.0, - anchor_latents: torch.Tensor | None = None, - anchor_masks: torch.Tensor | None = None, - anchor_free: bool = False, - rfg_scale: float = 1.0, - piflow_params: dict[str, Any] | None = None, - progress_callback: Callable[[int], Any] | None = None, - scheduler: Any | None = None, - progress_bar: Any = None, -) -> torch.Tensor: - """Run the four reusable SR stages for one tile batch.""" - piflow_params = piflow_params or _scheduler_piflow_params(scheduler) - piflow_params = piflow_params or getattr(dit, "piflow_params", None) - latent_state = prepare_latents( - dit=dit, - vae=vae, - device=device, - lq_videos=lq_videos, - lq_latents=lq_latents, - n_samples=n_samples, - seed=seed, - lq_noise_scale=lq_noise_scale, - lq_noise_type=lq_noise_type, - lq_channel_noise_scale=lq_channel_noise_scale, - anchor_latents=anchor_latents, - anchor_masks=anchor_masks, - anchor_free=anchor_free, - ) - text_state = text_encode( - dit=dit, - batch_size=latent_state.batch_size, - device=device, - text_embedder=text_embedder, - cached_text_embeds=cached_text_embeds, - ) - latent_visual = denoise( - dit=dit, - latent_state=latent_state, - text_state=text_state, - scale_factor=scale_factor, - num_steps=num_steps, - guidance_weight=guidance_weight, - scheduler_scale=scheduler_scale, - device=device, - tp_mesh=tp_mesh, - lq_noise_scale=lq_noise_scale, - cap_noise_timestep=cap_noise_timestep, - prediction_target=prediction_target, - channelcat_drop_threshold=channelcat_drop_threshold, - rfg_scale=rfg_scale, - piflow_params=piflow_params, - progress_callback=progress_callback, - scheduler=scheduler, - progress_bar=progress_bar, - ) - return vae_decode( - vae=vae, - latent_visual=latent_visual, - batch_size=latent_state.batch_size, - duration=latent_state.duration, - height=latent_state.height, - width=latent_state.width, - vae_decode_batch=vae_decode_batch, - ) - - -def _generate_kwargs( - components: Any, run_config: Any, params: Any, scale_factor: tuple[float, ...], **inputs: Any -) -> dict[str, Any]: - return { - **inputs, - "dit": components.dit, - "vae": components.vae, - "scale_factor": scale_factor, - "num_steps": run_config.num_steps, - "anchor_free": True, - "guidance_weight": 1.0, - "scheduler_scale": float(getattr(params, "scheduler_scale", 5.0)), - "seed": run_config.seed, - "device": run_config.device, - "tp_mesh": None, - "lq_noise_scale": float(getattr(params, "lq_noise_scale", 0.7)), - "lq_noise_type": getattr(params, "lq_noise_type", "ddpm"), - "lq_channel_noise_scale": float(getattr(params, "lq_channel_noise_scale", 0.0)), - "cap_noise_timestep": bool(getattr(params, "cap_noise_timestep", False)), - "cached_text_embeds": getattr(components, "cached_text_embeds", None), - } - - -def _as_tiles(sr_batch: torch.Tensor) -> list[torch.Tensor]: - if sr_batch.ndim != 5: - raise ValueError(f"SR sampler must return [batch,C,T,H,W], got {tuple(sr_batch.shape)}") - return [sample.float().cpu() for sample in sr_batch] - - -def _scale_progress_callback(progress_callback: Callable[[int], Any], chunk_size: int) -> Callable[[int], Any]: - - def update(n: int = 1) -> None: - progress_callback(n * chunk_size) - - return update - - -def _loaded_lu_scales(components: Any) -> tuple[int, ...]: - lu = getattr(components, "latent_upscaler", None) - if lu is None: - return () - scales = getattr(lu, "scales", None) - return tuple((int(scale) for scale in scales)) if scales is not None else (latent_upscaler_scale(lu),) - - -def _run_tile_batches( - tile_inputs: list[torch.Tensor], - sr_components: Any, - run_config: Any, - scale_factor: tuple[float, ...], - *, - prepare_chunk: Callable[[list[torch.Tensor]], list[torch.Tensor]] | None = None, - progress_callback: Callable[[int], Any] | None = None, - scheduler: Any | None = None, - sample_batch_size: int = 1, - progress_bar: Any = None, -) -> list[torch.Tensor]: - """Run the shared batched sampler loop for latent or pixel tile inputs. - - DiT and SR VAE stay resident for the complete tile loop. Only the optional - latent-upscaler is streamed per chunk. - """ - outputs: list[torch.Tensor] = [] - sampler = getattr(sr_components, "sampler", None) - params = _component_params(sr_components) - piflow_params = _scheduler_piflow_params(scheduler) or getattr(sr_components.dit, "piflow_params", None) - progress_steps = int(piflow_params["nfe"] if piflow_params is not None else run_config.num_steps - 1) - for start in range(0, len(tile_inputs), run_config.tiles_batch_size): - raw_chunk = tile_inputs[start : start + run_config.tiles_batch_size] - chunk = raw_chunk - if prepare_chunk is not None: - chunk = prepare_chunk(raw_chunk) - input_kwargs: dict[str, Any] - if prepare_chunk is None: - input_kwargs = {"lq_videos": chunk} - else: - input_kwargs = {"lq_latents": torch.cat(chunk, dim=0), "n_samples": len(raw_chunk)} - kwargs = _generate_kwargs(sr_components, run_config, params, scale_factor, **input_kwargs) - kwargs["seed"] = run_config.seed + start // sample_batch_size - chunk_progress = ( - _scale_progress_callback(progress_callback, len(raw_chunk)) if progress_callback is not None else None - ) - if sampler is None: - outputs.extend( - _as_tiles( - run_stages( - **kwargs, - piflow_params=piflow_params, - progress_callback=chunk_progress, - scheduler=scheduler, - progress_bar=progress_bar, - ) - ) - ) - else: - outputs.extend(_as_tiles(sampler(**kwargs))) - if chunk_progress is not None: - chunk_progress(progress_steps) - return outputs - - -def align_to_vae_stride(t: int) -> int: - """Round ``t`` to the nearest valid pixel-frame count ``1 + 8·k`` (ties up). - - SR data uses ``1 + 8·k`` frames so the temporal VAE stride round-trips - cleanly. Rounding to nearest (not truncation) keeps a requested duration - close to the target. - - Args: - t: Raw frame count. - - Returns: - Nearest valid frame count ``>= 1``. - """ - if t <= 1: - return 1 - k_floor = (t - 1) // 8 - t_floor = 1 + 8 * k_floor - t_ceil = 1 + 8 * (k_floor + 1) - return t_ceil if (t - t_floor) >= (t_ceil - t) else t_floor - - -def select_frame_indices(total_frames: int, src_fps: float, target_fps: float) -> list[int]: - """Fixed-stride frame indices that downsample ``src_fps`` to ``target_fps``. - - Mirrors ``data_encoding.frame_utils.select_frame_indices``: with - ``step = src_fps / target_fps >= 1``, ``round(i * step)`` is strictly - non-decreasing, so the selection never duplicates a source frame. - - Args: - total_frames: Number of frames in the source. - src_fps: Source frame rate (must be ``>= target_fps``). - target_fps: Desired frame rate. - - Returns: - Sorted list of integer source-frame indices. - """ - step = src_fps / target_fps - indices = [round(i * step) for i in range(int(total_frames / step))] - return [i for i in indices if i < total_frames] - - -def resample_to_target_fps( - video: torch.Tensor, - src_fps: float, - target_fps: int = TARGET_FPS, -) -> tuple[torch.Tensor, int]: - """Resample a decoded video toward ``target_fps`` (downsample / no-op / keep). - - Three tiers, mirroring the dataset encoding pipeline: - - - ``|src_fps - target_fps| < RESAMPLE_FPS_TOLERANCE``: pass through unchanged. - - ``src_fps > target_fps``: fixed-stride downsample via - :func:`select_frame_indices`; the result is at ``target_fps``. - - ``src_fps < target_fps``: kept at the native rate with a warning (no ffmpeg - ``minterpolate`` upsample); the clip stays mildly out of distribution. - - Args: - video: ``[T, C, H, W]`` decoded source video. - src_fps: Source frame rate. - target_fps: Training frame rate to resample toward. - - Returns: - ``(resampled_video, effective_fps)`` where ``effective_fps`` is the rate - the returned frames play at (``target_fps`` when downsampled, otherwise - ``round(src_fps)``) and is the correct rate to save the SR result at. - """ - if abs(src_fps - target_fps) < RESAMPLE_FPS_TOLERANCE: - return video, round(src_fps) - if src_fps > target_fps: - indices = select_frame_indices(video.shape[0], src_fps, target_fps) - logger.warning( - "Source fps %.2f > target %sfps: downsampling %s frames -> %s (fixed-stride); " - "source temporal detail beyond %sfps is discarded.", - src_fps, - target_fps, - video.shape[0], - len(indices), - target_fps, - ) - return video[indices], target_fps - logger.warning( - "Source fps %.2f < target %sfps: keeping native frames (no minterpolate upsample); " - "output is mildly out of distribution.", - src_fps, - target_fps, - ) - return video, round(src_fps) - - -def clip_to_aligned_frames(video: torch.Tensor, max_num_frames: int = MAX_NUM_FRAMES) -> torch.Tensor: - """Take the first ``max_num_frames`` and floor-align to ``1 + 8k`` frames. - - Args: - video: ``[T, C, H, W]`` video. - max_num_frames: Hard cap applied before alignment (``<= 0`` disables it). - - Returns: - ``[T', C, H, W]`` with ``T' == 1 + 8k`` and ``T' <= max_num_frames``. - - Raises: - ValueError: If the video has no frames. - """ - if max_num_frames > 0: - video = video[:max_num_frames] - aligned = 1 + 8 * ((video.shape[0] - 1) // 8) if video.shape[0] > 0 else 0 - if aligned == 0: - msg = "Video has no readable frames." - raise ValueError(msg) - return video[:aligned] - - -def read_video_tchw_uint8(path: Path) -> tuple[torch.Tensor, float]: - """Decode an mp4/mkv to ``([T, C, H, W] uint8, src_fps)`` (no resample/cap). - - Resampling to the training fps and frame-count alignment are applied by the - caller via :func:`resample_to_target_fps` and :func:`clip_to_aligned_frames`. - - Args: - path: Source video file. - - Returns: - ``([T, C, H, W] uint8 video, native_fps)``. - - Raises: - ValueError: If the file has no readable frames or reports no fps. - """ - av = _require_av() - with av.open(str(path), mode="r") as container: - stream = container.streams.video[0] - src_fps = float(stream.average_rate or stream.base_rate or 0.0) - frames = [ - torch.from_numpy(frame.to_ndarray(format="rgb24")).permute(2, 0, 1) for frame in container.decode(stream) - ] - - if not frames: - msg = f"Video {path.name} has no readable frames." - raise ValueError(msg) - if src_fps <= 0: - msg = f"Video {path.name} reports no usable fps (got {src_fps})." - raise ValueError(msg) - return torch.stack(frames).contiguous(), src_fps - - -def load_lr_latent(path: Path) -> torch.Tensor: - """Load a raw, unscaled LR latent ``.pt`` as ``[T, C, H, W]`` float32. - - The tensor must match what ``encode_lq_video_to_lr_latent`` produces for - the model's VAE (no scaling factor applied) — the LU multiplies internally. - - Args: - path: Path to a ``.pt`` file holding a 4D latent tensor. - - Returns: - ``[T, C, H, W]`` float32 latent on CPU. - - Raises: - ValueError: If the loaded object is not a rank-4 float tensor. - """ - obj = torch.load(path, map_location="cpu", weights_only=True) - if not isinstance(obj, torch.Tensor) or obj.ndim != LATENT_NDIM: - shape = getattr(obj, "shape", type(obj).__name__) - msg = f"Expected a rank-4 [T, C, H, W] latent tensor in {path}, got {shape}." - raise ValueError(msg) - logger.info("Loaded LR latent %s from %s", tuple(obj.shape), path) - return obj.float() - - -def _require_av(): - """Load PyAV only when video or audio output is requested.""" - try: - import av - except (ImportError, OSError) as exc: - raise RuntimeError("SR video/audio output requires PyAV and FFmpeg. Install with `pip install av`.") from exc - return av - - -def _video_np(frames: Tensor) -> np.ndarray: - """(3, T, H, W) uint8 → (T, H, W, 3) uint8.""" - return frames.detach().permute(1, 2, 3, 0).cpu().numpy().astype(np.uint8, copy=False) - - -def _audio_np(audio: list[np.ndarray] | np.ndarray) -> np.ndarray: - """int16 / float waveform → float32 mono in [-1, 1].""" - audio_np = np.asarray(audio[0] if isinstance(audio, list) else audio) - if np.issubdtype(audio_np.dtype, np.integer): - audio_np = audio_np.astype(np.float32) / np.iinfo(audio_np.dtype).max - return np.clip(audio_np, -1.0, 1.0).astype(np.float32, copy=False) - - -def extract_audio_from_video( - source_video: str | Path, - audio_sample_rate: int = 44100, -) -> np.ndarray | None: - """Decode a source video's audio as mono float32 samples. - - SR changes only the video stream. Decoding the source audio here lets the - output mux keep that audio without requiring the SR caller to load the - entire source container itself. Resampling also gives the AAC writer a - stable sample rate for videos whose source audio uses another rate. - """ - if audio_sample_rate <= 0: - raise ValueError("audio_sample_rate must be positive") - - av = _require_av() - chunks: list[np.ndarray] = [] - with av.open(str(source_video), mode="r") as container: - audio_streams = list(container.streams.audio) - if not audio_streams: - return None - - resampler = av.AudioResampler(format="fltp", layout="mono", rate=audio_sample_rate) - for frame in container.decode(audio=0): - for resampled in resampler.resample(frame): - chunks.append(resampled.to_ndarray().reshape(-1)) - for resampled in resampler.resample(None): - chunks.append(resampled.to_ndarray().reshape(-1)) - - if not chunks: - return None - return np.concatenate(chunks).astype(np.float32, copy=False) - - -def mux_video_audio( # noqa: PLR0913 - frames: Tensor, - audio: list[np.ndarray] | np.ndarray | None, - output_path: str | Path, - fps: int = 24, - audio_sample_rate: int = 44100, - video_crf: int = 18, - source_video: str | Path | None = None, - lossless: bool = False, -) -> Path: - """Mux video (+ optional audio) into ``output_path`` via PyAV — no temp files. - - frames: (3, T, H, W) uint8 — single video, channels-first. - audio: list of (samples,) int16, a single ndarray, or None for video-only. - source_video: optional source container from which audio is copied when - ``audio`` is not supplied. This is useful when SR receives frames - extracted from a generated T2VA video. - lossless: encode video as FFV1 instead of libx264. Audio remains AAC. - """ - av = _require_av() - if audio is not None and source_video is not None: - raise ValueError("pass either audio or source_video, not both") - if source_video is not None: - audio = extract_audio_from_video(source_video, audio_sample_rate) - - output_path = Path(output_path) - output_path.parent.mkdir(parents=True, exist_ok=True) - video_np = _video_np(frames) - - with av.open(str(output_path), mode="w") as container: - video_stream = container.add_stream("ffv1" if lossless else "libx264", rate=fps) - video_stream.width = video_np.shape[2] - video_stream.height = video_np.shape[1] - if lossless: - video_stream.pix_fmt = "bgra" - video_stream.options = {"level": "3"} - else: - video_stream.pix_fmt = "yuv420p" - video_stream.options = {"crf": str(video_crf)} - - audio_stream = None - if audio is not None: - audio_stream = container.add_stream("aac", rate=audio_sample_rate) - audio_stream.layout = "mono" - audio_stream.bit_rate = 192_000 - - for i, frame_np in enumerate(video_np): - frame = av.VideoFrame.from_ndarray(frame_np, format="rgb24") - frame.pts = i - frame.time_base = Fraction(1, fps) - for packet in video_stream.encode(frame): - container.mux(packet) - - if audio_stream is not None: - audio_np = _audio_np(audio) - audio_frame = av.AudioFrame.from_ndarray(audio_np[np.newaxis, :], format="fltp", layout="mono") - audio_frame.sample_rate = audio_sample_rate - audio_frame.pts = 0 - audio_frame.time_base = Fraction(1, audio_sample_rate) - for packet in audio_stream.encode(audio_frame): - container.mux(packet) - - for packet in video_stream.encode(): - container.mux(packet) - if audio_stream is not None: - for packet in audio_stream.encode(): - container.mux(packet) - - return output_path - - -def latent_upscaler_for_scale(upscaler: Any, scale: int) -> nn.Module | None: - """Keep bank dispatch on the registered parent for Diffusers offload.""" - if upscaler is None: - return None - if hasattr(upscaler, "_models"): - key = f"{int(scale)}x" - if key not in upscaler._models: - return None - # The native helper returns a child module. A Diffusers CPU-offload - # hook is attached to the registered bank, so call its forward method - # with the selected entry instead. - upscaler.target_scale = key - return upscaler - for_scale = getattr(upscaler, "for_scale", None) - if callable(for_scale): - return for_scale(scale) - return upscaler if str(getattr(upscaler, "target_scale", "4x")) == f"{int(scale)}x" else None - - -SRInput = Tensor | np.ndarray | Sequence[Any] | str | Path -SRBatchInput = SRInput | Sequence[SRInput] -SRScale = Literal[2, 4, 2.25] -SROutputType = Literal["pt", "torch", "np", "numpy"] -VIDEO_RANK = 4 -VIDEO_BATCH_RANK = 5 - - -def _is_video_batch(value: Any) -> bool: - if isinstance(value, (Tensor, np.ndarray)): - return value.ndim == VIDEO_BATCH_RANK - if isinstance(value, (str, Path)) or not isinstance(value, Sequence) or not value: - return False - first = value[0] - if isinstance(first, (str, Path)): - return True - if isinstance(first, (Tensor, np.ndarray)): - return first.ndim >= VIDEO_RANK - return False - - -def _video_batch_tensor(value: Tensor | np.ndarray) -> list[Tensor]: - tensor = value.detach() if isinstance(value, Tensor) else torch.as_tensor(value) - if tensor.ndim != VIDEO_BATCH_RANK: - raise ValueError(f"batched video must have rank 5, got shape {tuple(tensor.shape)}") - if tensor.shape[-1] in (1, 3, 4): - return [sample.permute(0, 3, 1, 2) for sample in tensor] - if tensor.shape[1] in (1, 3, 4): - return [sample.permute(1, 0, 2, 3) for sample in tensor] - if tensor.shape[2] in (1, 3, 4): - return [sample for sample in tensor] - raise ValueError("5D video must be [B,C,T,H,W], [B,T,C,H,W], or [B,T,H,W,C]") - - -def _latent_tensor(value: Tensor | np.ndarray, channels: int) -> Tensor: - tensor = value.detach() if isinstance(value, Tensor) else torch.as_tensor(value) - if tensor.ndim != VIDEO_RANK: - raise ValueError(f"latent sample must have rank 4, got shape {tuple(tensor.shape)}") - if tensor.shape[1] == channels: - return tensor - if tensor.shape[-1] == channels: - return tensor.permute(0, 3, 1, 2) - if tensor.shape[0] == channels: - return tensor.permute(1, 0, 2, 3) - raise ValueError("4D latent must be [T,C,H,W], [T,H,W,C], or [C,T,H,W]") - - -def _latent_batch(value: Tensor | np.ndarray | str | Path | Sequence[Any], channels: int) -> list[Tensor]: - if isinstance(value, (str, Path)): - return [_latent_tensor(load_lr_latent(Path(value)), channels)] - if isinstance(value, (Tensor, np.ndarray)): - tensor = value.detach() if isinstance(value, Tensor) else torch.as_tensor(value) - if tensor.ndim == VIDEO_RANK: - return [_latent_tensor(tensor, channels)] - if tensor.ndim == VIDEO_BATCH_RANK: - if tensor.shape[-1] == channels: - return [_latent_tensor(sample, channels) for sample in tensor] - if tensor.shape[1] == channels: - return [_latent_tensor(sample, channels) for sample in tensor] - if tensor.shape[2] == channels: - return [_latent_tensor(sample, channels) for sample in tensor] - raise ValueError("5D latent must be [B,C,T,H,W], [B,T,C,H,W], or [B,T,H,W,C]") - raise ValueError(f"latents must have rank 4 or 5, got shape {tuple(tensor.shape)}") - if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): - if not value: - raise ValueError("latents batch must not be empty") - result: list[Tensor] = [] - for item in value: - result.extend(_latent_batch(item, channels)) - return result - raise TypeError("latents must be tensors, arrays, paths, or a batch of them") - - -def _validate_batch_shapes(values: list[Tensor], name: str) -> None: - if not values: - raise ValueError(f"{name} batch must not be empty") - shape = tuple(values[0].shape) - if any(tuple(value.shape) != shape for value in values[1:]): - raise ValueError( - f"all {name} batch items must have equal shape, got {[tuple(value.shape) for value in values]}" - ) - - -def _batch_output_values(value: Any, batch_size: int, name: str) -> list[Any]: - if isinstance(value, list): - if len(value) != batch_size: - raise ValueError(f"{name} must contain one value per input ({batch_size}), got {len(value)}") - return value - if name == "save_path" and value is not None and batch_size > 1: - raise ValueError("save_path must be a list when processing multiple SR inputs") - return [value] * batch_size +def _upsample_tiles_to_base(raw_tiles: list[torch.Tensor], base_h: int, base_w: int) -> list[torch.Tensor]: + """Bilinear-upscale ``[C,T,h,w]`` pixel tiles to ``[T,base_h,base_w,C]`` (channels-last, for ``_encode_lq_videos``).""" + return [ + functional.interpolate(tile.float(), size=(base_h, base_w), mode="bilinear", align_corners=False).permute( + 1, 2, 3, 0 + ) + for tile in raw_tiles + ] def _stitch_batch( outputs: list[Tensor], - grid: Any, + grid: TileGrid, batch_size: int, height: int, width: int, @@ -2460,118 +550,43 @@ def _stitch_batch( ) -def _run_batched_tiles( # noqa: PLR0913 - samples: list[Tensor], - components: Any, - run_config: Any, - *, - latent_input: bool, - progress_bar: Any, - scheduler: Any, -) -> Tensor: - """Run equal-shape samples tile-major and restore the sample batch.""" - params = _component_params(components) - visual_size = _visual_size(params) - batch_size = len(samples) - if latent_input: - spatial_factor = _spatial_factor(components) - height = samples[0].shape[-2] * spatial_factor - width = samples[0].shape[-1] * spatial_factor - _, _, pixel_grid = _tile_geometry( - height, - width, - visual_size, - run_config.resolution_scale, - run_config.overlap, - spatial_factor, - run_config.tile_grid_mode, - run_config.tile_grid_min_overlap, - ) - tile_grid = latent_tile_grid_from_pixel_grid(pixel_grid, spatial_factor) - input_grid = tile_grid - stitch_grid = pixel_grid - tile_sets = [extract_all_tiles(sample, tile_grid) for sample in samples] - upscaler = latent_upscaler_for_scale(components.latent_upscaler, run_config.resolution_scale) - if upscaler is None: - raise ValueError( - f"The latent-upscaler path has no {run_config.resolution_scale}x model " - f"(loaded LU scales: {_loaded_lu_scales(components)})" - ) - prepare_chunk = lambda chunk: [ - upscale_lr_latent_tile(tile, upscaler, components.vae, run_config.device) for tile in chunk - ] - else: - height, width = samples[0].shape[-2:] - base, _tile_hw, grid = _tile_geometry( - height, - width, - visual_size, - run_config.resolution_scale, - run_config.overlap, - _spatial_factor(components), - run_config.tile_grid_mode, - run_config.tile_grid_min_overlap, - ) - tile_sets = [_upsample_tiles_to_base(extract_all_tiles(sample, grid), base[0], base[1]) for sample in samples] - input_grid = grid - stitch_grid = grid - prepare_chunk = None - - tile_inputs = [ - tile_sets[sample_index][tile_index] - for tile_index in range(input_grid.total_tiles) - for sample_index in range(batch_size) - ] - outputs = _run_tile_batches( - tile_inputs, - components, - run_config, - _scale_factor(params, visual_size), - scheduler=scheduler, - prepare_chunk=prepare_chunk, - progress_bar=progress_bar, - ) - return _stitch_batch(outputs, stitch_grid, batch_size, height, width, run_config.resolution_scale) - - -def _sr_progress_plan( - value: Tensor, - components: Any, - run_config: Any, - *, - latent_input: bool, - batch_size: int = 1, - scheduler: Any | None = None, -) -> tuple[int, int, int]: - """Return total updates, denoising steps, and tile batches.""" - params = _component_params(components) - visual_size = _visual_size(params) - spatial_factor = _spatial_factor(components) - height, width = value.shape[-2:] - if latent_input: - height *= spatial_factor - width *= spatial_factor - _, _, grid = _tile_geometry( - height, - width, - visual_size, - run_config.resolution_scale, - run_config.overlap, - spatial_factor, - ) - tile_count = grid.total_tiles * batch_size - chunk_count = max(1, math.ceil(tile_count / run_config.tiles_batch_size)) - piflow_params = _scheduler_piflow_params(scheduler) or getattr(components.dit, "piflow_params", None) - if isinstance(piflow_params, Mapping): - denoise_steps = int(piflow_params.get("nfe", run_config.num_steps - 1)) - elif piflow_params is not None: - denoise_steps = int(getattr(piflow_params, "nfe", run_config.num_steps - 1)) +def _prepare_batch(value: Tensor, name: str, channels: int, *, dtype: torch.dtype | None = None) -> Tensor: + """Add a batch dim to a single ``(C,T,H,W)`` sample; validate a ``(B,C,T,H,W)`` batch.""" + if value.ndim == 4: + value = value.unsqueeze(0) + if value.ndim != 5: + raise ValueError(f"`{name}` must have shape (C,T,H,W) or (B,C,T,H,W), got {tuple(value.shape)}") + if value.shape[1] != channels: + raise ValueError(f"`{name}` must have {channels} channels at dim 1, got shape {tuple(value.shape)}") + if dtype is not None and value.dtype != dtype: + raise ValueError(f"`{name}` must be {dtype}, got {value.dtype}") + return value + + +@contextmanager +def _execution_device_context(device: torch.device | str): + """Scope SR model execution to Diffusers' resolved CUDA device.""" + resolved_device = torch.device(device) + if resolved_device.type == "cuda": + with torch.cuda.device(resolved_device): + yield else: - denoise_steps = run_config.num_steps - 1 - if getattr(components, "sampler", None) is not None: - denoise_steps = 1 - denoise_steps = max(1, denoise_steps) - return chunk_count * denoise_steps, denoise_steps, chunk_count + yield + + +def _decode_source_latent_video(raw_latents: Tensor, source_vae: nn.Module, device: str | torch.device) -> Tensor: + """Decode a base-VAE latent to ``[C,T,H,W]`` pixels for the KVAE bridge.""" + if raw_latents.ndim != 4: + raise ValueError(f"raw_latents must have rank 4 [C,T,H,W], got {tuple(raw_latents.shape)}") + latent_5d = raw_latents.unsqueeze(0).to(device=device) + latent_5d = cast_to_module_dtype(source_vae, latent_5d) + # `source_vae.decode` is decorated with `@apply_forward_hook` (the standard convention on every + # AutoencoderKL-family `encode`/`decode` in this repo), so a Diffusers CPU-offload hook on + # `source_vae` already fires correctly here without any manual hook dispatch. + device_type = torch.device(device).type + with torch.autocast(device_type, dtype=torch.bfloat16, enabled=device_type == "cuda"): + decoded = source_vae.decode(latent_5d).sample + return ((decoded.squeeze(0).float().clamp(-1, 1) + 1.0) * 127.5).round().clamp(0, 255).to(torch.uint8).cpu() class Kandinsky6SRPipeline(DiffusionPipeline): @@ -2580,7 +595,7 @@ class Kandinsky6SRPipeline(DiffusionPipeline): Components are supplied to the constructor or loaded by Diffusers through ``from_pretrained``. This pipeline deliberately has no config-path factory: model construction belongs to the Diffusers component package, - while this class owns only SR orchestration and I/O. + while this class owns only SR orchestration. ``resolution_scale=2.25`` is the production route. It performs a 1.125x pixel pre-upscale followed by the x2 latent-upscaler path. The input @@ -2598,7 +613,7 @@ class Kandinsky6SRPipeline(DiffusionPipeline): model_cpu_offload_seq = "source_vae->latent_upscaler->transformer->vae" _optional_components = ["latent_upscaler", "source_vae"] - def __init__( # noqa: PLR0913 + def __init__( self, transformer: nn.Module, vae: nn.Module, @@ -2614,59 +629,37 @@ def __init__( # noqa: PLR0913 latent_upscaler=latent_upscaler, source_vae=source_vae, ) - transformer_config = getattr(transformer, "config", None) - self._sr_params = _coerce_sr_params(getattr(transformer_config, "sr_params", None)) - if self._sr_params is None: + if not getattr(transformer, "sr_params", None): raise ValueError("SR transformer config must contain 'sr_params'") - - @staticmethod - def _seed_from_generator( - generator: torch.Generator | None, - device: torch.device, - ) -> int: - if generator is None: - return int(torch.randint(0, 2**31, (1,), device=device).item()) - generator_device = torch.device(getattr(generator, "device", "cpu")) - return int( - torch.randint( - 0, - 2**31, - (1,), - generator=generator, - device=generator_device, - ).item() + vae_config = getattr(vae, "config", None) + self.spatial_factor = int( + getattr(vae, "spatial_factor", None) or getattr(vae_config, "spatial_factor", None) or VAE_SPATIAL_FACTOR ) def check_inputs( self, - video: SRBatchInput | None = None, - latents: Tensor | np.ndarray | str | Path | Sequence[Any] | None = None, + video: Tensor | None = None, + latents: Tensor | None = None, *, resolution_scale: float, - num_inference_steps: int | None = None, - overlap: float | None = None, - tiles_batch_size: int | None = None, + num_inference_steps: int, + min_overlap: float, + tiles_batch_size: int, kvae_bridge: bool = False, - fps: int | None = None, - audio: list[np.ndarray] | np.ndarray | None = None, - source_video: str | Path | list[str | Path] | None = None, - audio_sample_rate: int = 44100, - output_type: SROutputType = "pt", + output_type: str = "pt", ) -> None: - """Validate SR input, scale, tiling, audio, and output arguments. + """Validate SR input, scale, tiling, and output arguments. Args: - video: Pixel video input in a supported frame layout. - latents: Raw latent tensor or local latent file. + video: Pixel video batch, already normalized to rank 5. + latents: Latent batch, already normalized to rank 5. resolution_scale: Total output scale, including the 2.25x fractional route. num_inference_steps: Number of denoising steps. - overlap: Tile overlap fraction in ``[0, 1)``. + min_overlap: Minimum tile overlap fraction in ``[0, 1)``. tiles_batch_size: Number of tiles processed together. - fps: Optional output video frame rate. - audio: Optional audio waveform to mux. - source_video: Optional source video whose audio is copied. - audio_sample_rate: Sample rate used when muxing audio. + kvae_bridge: Whether ``latents`` come from a different VAE and + must be bridged through ``self.source_vae``. output_type: ``pt``/``torch`` or ``np``/``numpy``. Raises: @@ -2674,29 +667,27 @@ def check_inputs( pipeline's supported range. """ if (video is None) == (latents is None): - raise ValueError("pass exactly one of video or latents") + raise ValueError("pass exactly one of `video` or `latents`") if output_type not in ("pt", "torch", "np", "numpy"): raise ValueError(f"unsupported output_type={output_type!r}") - if audio is not None and source_video is not None: - raise ValueError("pass either audio or source_video, not both") - if audio_sample_rate <= 0: - raise ValueError("audio_sample_rate must be positive") - if fps is not None and fps <= 0: - raise ValueError("fps must be positive") - - requested_scale = float(resolution_scale) - resolve_scale_request(requested_scale) - effective_steps = 5 if num_inference_steps is None else int(num_inference_steps) - effective_overlap = 0.25 if overlap is None else float(overlap) - effective_tiles_batch_size = 1 if tiles_batch_size is None else int(tiles_batch_size) - if effective_steps < 1: + if video is not None: + num_frames = video.shape[2] + if (num_frames - 1) % 8 != 0: + raise ValueError( + f"`video` must have 1 + 8k frames (1, 9, 17, ..., {MAX_NUM_FRAMES}, ...), got {num_frames}" + ) + if num_frames > MAX_NUM_FRAMES: + raise ValueError(f"`video` must have at most {MAX_NUM_FRAMES} frames (~5s at 24fps), got {num_frames}") + + resolve_scale_request(float(resolution_scale)) + if num_inference_steps < 1: raise ValueError("num_inference_steps must be positive") - if not 0 <= effective_overlap < 1: - raise ValueError("overlap must be in [0, 1)") - if effective_tiles_batch_size < 1: + if not 0.0 <= min_overlap < 1.0: + raise ValueError("min_overlap must be in [0, 1)") + if tiles_batch_size < 1: raise ValueError("tiles_batch_size must be positive") if latents is not None: - if requested_scale == 2.25 and not kvae_bridge: # noqa: PLR2004 + if float(resolution_scale) == 2.25 and not kvae_bridge: raise ValueError( "2.25x latent SR requires kvae_bridge=True and source_vae; " "ordinary latent input supports only x2/x4" @@ -2705,334 +696,733 @@ def check_inputs( raise ValueError("KVAE bridge latent SR requires source_vae") @staticmethod - def _load_video(value: SRInput) -> tuple[Tensor, float | None, Path | None]: - if isinstance(value, (str, Path)): - path = Path(value) - video, source_fps = read_video_tchw_uint8(path) - video, output_fps = resample_to_target_fps(video, source_fps) - return clip_to_aligned_frames(video), float(output_fps), path - return _to_tchw(value), None, None - - @classmethod - def _load_videos(cls, value: SRBatchInput) -> tuple[list[Tensor], list[float | None], list[Path | None]]: - if isinstance(value, (Tensor, np.ndarray)) and value.ndim == VIDEO_BATCH_RANK: - values: Sequence[Any] = _video_batch_tensor(value) - elif _is_video_batch(value): - values = value - else: - values = [value] - loaded = [cls._load_video(item) for item in values] - videos = [item[0] for item in loaded] - _validate_batch_shapes(videos, "video") - return videos, [item[1] for item in loaded], [item[2] for item in loaded] - - @staticmethod - def _format_frames(frames: Tensor, output_type: SROutputType) -> Tensor | np.ndarray: + def _format_frames(frames: Tensor, output_type: str) -> Tensor | np.ndarray: if output_type in ("pt", "torch"): return frames if output_type in ("np", "numpy"): return frames.detach().cpu().numpy() raise ValueError(f"Unsupported output_type={output_type!r}; use 'pt'/'torch' or 'np'/'numpy'") + def _sparse_params(self, visual: torch.Tensor, visual_cu_seqlens: torch.Tensor) -> dict[str, Any] | None: + """Build NABLA sparse-attention parameters for ``self.transformer`` and the given visual tensor. + + Returns ``None`` for dense (``"flash"``) attention. The target ``P`` is used directly (the + training-time linear warmup schedule does not apply at inference). + """ + model = self.transformer + if model.patch_size[0] != 1: + raise ValueError(f"Expected transformer.patch_size[0] == 1, got {model.patch_size[0]}") + t, h, w, _ = visual.shape + t, h, w = (t // model.patch_size[0], h // model.patch_size[1], w // model.patch_size[2]) + visual_size = get_visual_size(visual, self.spatial_factor) + attention_configs = model.attention_params + try: + attention_params = attention_configs[visual_size] + except KeyError: + # JSON object keys are always strings, while the native YAML config + # uses integer resolution keys. + attention_params = attention_configs[str(visual_size)] + + if attention_params.get("type") != "nabla": + return None + return { + "to_fractal": True, + "P": attention_params.get("P"), + "wT": attention_params.get("wT"), + "wW": attention_params.get("wW"), + "wH": attention_params.get("wH"), + "add_sta": attention_params.get("add_sta"), + "visual_shape": (t, h, w), + "visual_seqlens": visual_cu_seqlens, + } + @torch.no_grad() - def _run( # noqa: PLR0913 + def _encode_lq_videos(self, lq_videos: list[torch.Tensor], device: str | torch.device) -> torch.Tensor: + """VAE-encode a batch of ``[T,H,W,C]`` spatial tiles into a stacked ``[sum(T'),H',W',C]`` latent. + + Each tile is encoded independently (rather than as one batched VAE + call) to mirror the reference implementation exactly; ``self.vae``'s + ``scaling_factor`` is applied here, since these latents feed the + denoising loop directly (unlike ``encode_lq_video_to_lr_latent``, + whose unscaled output is scaled later by ``_upscale_lr_latent_tile``). + """ + vae = self.vae + scaling_factor = float(vae.config.scaling_factor) + latents = [] + for lq in lq_videos: + pixel = lq.permute(3, 0, 1, 2).unsqueeze(0).to(device=device) + latent = encode_pixels_to_latent(vae, pixel) + latents.append(latent.squeeze(0).permute(1, 2, 3, 0).float() * scaling_factor) + return torch.cat(latents, dim=0) + + @torch.no_grad() + def _upscale_lr_latent_tile( + self, lr_latent_tile: torch.Tensor, scale: int, device: str | torch.device + ) -> torch.Tensor: + """Upscale one raw ``[C,T,H,W]`` latent tile via ``self.latent_upscaler``'s ``scale``x entry.""" + scaling_factor = float(getattr(getattr(self.vae, "config", None), "scaling_factor", 1.0)) + tile = lr_latent_tile.unsqueeze(0).to(device=device, dtype=torch.float32) + tile = cast_to_module_dtype(self.latent_upscaler, tile) + device_type = torch.device(device).type + with torch.autocast(device_type, dtype=torch.bfloat16, enabled=device_type == "cuda"): + upscaled = self.latent_upscaler(tile * scaling_factor, entry=f"{scale}x", return_intermediates=False) + return upscaled.squeeze(0).permute(1, 2, 3, 0).float() + + def _build_initial_latent( + self, + lq_latent: torch.Tensor, + bs: int, + duration: int, + height: int, + width: int, + device: str | torch.device, + generator: torch.Generator | None, + lq_noise_scale: float, + lq_noise_type: Literal["linear", "ddpm"], + lq_channel_noise_scale: float, + ) -> torch.Tensor: + """Build the initial latent tensor for the SR denoising loop, per ``self.transformer.instruct_type``. + + Constructs ``[starting | cond | mask]`` depending on ``instruct_type``: + + - ``"noise"``: degraded LQ as starting point, zero LQ cond + zero mask. + - ``"channel"``: random Gaussian noise as starting point, LQ cond + ones mask. + - ``"hybrid"``: degraded LQ as starting point, LQ cond + ones mask. + - ``"hybrid_anchor"``: degraded LQ as starting point, zeroed (anchor-free) + HR anchor + zeroed anchor_mask — the only mode Kandinsky 6 SR ships with. + + Returns: + Initial latent tensor ready for the denoising loop (bs*T, H, W, 33). + + Raises: + ValueError: If ``instruct_type`` is not supported. + """ + dit = self.transformer + lq_latent = lq_latent.to(device) + lq_latent = cast_to_module_dtype(dit, lq_latent) + if dit.instruct_type == "noise": + degraded_lq = degrade_lq_latent(lq_latent, lq_noise_scale, lq_noise_type, generator=generator) + if not getattr(dit, "visual_cond", False): + return degraded_lq + zero_cond = torch.zeros_like(degraded_lq) + zero_mask = torch.zeros([*degraded_lq.shape[:-1], 1], dtype=degraded_lq.dtype, device=device) + return torch.cat([degraded_lq, zero_cond, zero_mask], dim=-1) + if dit.instruct_type in ("channel", "hybrid"): + if dit.instruct_type == "hybrid": + starting = degrade_lq_latent(lq_latent.clone(), lq_noise_scale, lq_noise_type, generator=generator) + else: + starting = randn_tensor( + (bs * duration, height, width, dit.in_visual_dim), generator=generator, device=device + ) + channel_lq = degrade_lq_latent(lq_latent, lq_channel_noise_scale, noise_type="linear") + mask = torch.ones_like(lq_latent[..., :1]) + return torch.cat([starting, channel_lq, mask], dim=-1) + if dit.instruct_type == "hybrid_anchor": + # Kandinsky 6 SR ships anchor-free: the public API has no way to supply a real HR + # anchor, so the anchor and its mask are always zeroed. + anchor_latent = torch.zeros_like(lq_latent) + anchor_mask = torch.zeros_like(lq_latent[..., :1]) + starting = degrade_lq_latent(lq_latent.clone(), lq_noise_scale, lq_noise_type, generator=generator) + anchor = degrade_lq_latent(anchor_latent, lq_channel_noise_scale, noise_type="linear") + return torch.cat([starting, anchor, anchor_mask], dim=-1) + raise ValueError(f"Kandinsky6SRPipeline does not support instruct_type={dit.instruct_type!r}") + + def _prepare_sr_latents( self, *, - video: list[Tensor] | None = None, - latents: list[Tensor] | None = None, - resolution_scale: float, - num_inference_steps: int | None = None, - generator: torch.Generator | None = None, - overlap: float | None = None, - tiles_batch_size: int | None = None, - kvae_bridge: bool = False, - cached_text_embeds: dict[str, Tensor] | None = None, - save_path: str | Path | list[str | Path] | None = None, - fps: int | None = None, - audio: list[np.ndarray] | np.ndarray | None = None, - source_video: str | Path | list[str | Path] | None = None, - audio_sample_rate: int = 44100, - output_type: SROutputType = "pt", - return_dict: bool = True, - ) -> Kandinsky6SRPipelineOutput | tuple[Any, ...]: - requested_scale = float(resolution_scale) - effective_steps = 5 if num_inference_steps is None else int(num_inference_steps) - effective_overlap = 0.25 if overlap is None else float(overlap) - effective_tiles_batch_size = 1 if tiles_batch_size is None else int(tiles_batch_size) - self.check_inputs( - video=video, - latents=latents, - resolution_scale=requested_scale, - num_inference_steps=effective_steps, - overlap=effective_overlap, - tiles_batch_size=effective_tiles_batch_size, - kvae_bridge=kvae_bridge, - fps=fps, - audio=audio, - source_video=source_video, - audio_sample_rate=audio_sample_rate, - output_type=output_type, - ) - tiling_scale, pre_upscale = resolve_scale_request(requested_scale) - components = SRComponents( - dit=self.transformer, - vae=self.vae, - latent_upscaler=self.latent_upscaler, - sr_params=self._sr_params, - cached_text_embeds=cached_text_embeds, - spatial_factor=getattr(self.vae, "spatial_factor", None), + lq_videos: list[torch.Tensor] | None, + lq_latents: torch.Tensor | None, + n_samples: int | None, + device: str | torch.device, + generator: torch.Generator | None, + lq_noise_scale: float, + lq_noise_type: str, + lq_channel_noise_scale: float, + ) -> tuple[torch.Tensor, torch.Tensor, int, int, int, int]: + """Encode LQ input and build the initial latent for one SR tile batch. + + Returns: + ``(lq_latent, image, batch_size, duration, height, width)``. + """ + if lq_latents is not None: + lq_latent = lq_latents.to(device) + if n_samples is None: + raise ValueError("n_samples is required when lq_latents is provided") + batch_size = n_samples + elif lq_videos is not None: + batch_size = len(lq_videos) + lq_latent = self._encode_lq_videos(lq_videos, device) + else: + raise ValueError("Either lq_videos or lq_latents must be provided") + duration = lq_latent.shape[0] // batch_size + height, width = (lq_latent.shape[1], lq_latent.shape[2]) + image = self._build_initial_latent( + lq_latent, + bs=batch_size, + duration=duration, + height=height, + width=width, + device=device, + generator=generator, + lq_noise_scale=lq_noise_scale, + lq_noise_type=lq_noise_type, + lq_channel_noise_scale=lq_channel_noise_scale, ) + return lq_latent, image, batch_size, duration, height, width - batch_size = len(latents) if latents is not None else len(video) if video is not None else 0 - save_paths = _batch_output_values(save_path, batch_size, "save_path") - source_videos = _batch_output_values(source_video, batch_size, "source_video") - sr_video = video - try: - if latents is not None and kvae_bridge: - sr_video = [ - _decode_source_latent_video(latent, self.source_vae, self._execution_device) for latent in latents - ] - latents = None - if sr_video is not None and pre_upscale != 1.0: - sr_video = [ - pre_upscale_video(video_item, pre_upscale, _spatial_factor(components)) for video_item in sr_video - ] + def _encode_sr_text( + self, + batch_size: int, + device: str | torch.device, + cached_text_embeds: dict[str, torch.Tensor] | None, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: + """Build the SR text state: empty placeholders, or cached empty-caption embeddings. - run_config = RunConfig( - device=str(self._execution_device), - num_steps=effective_steps, - seed=self._seed_from_generator(generator, self._execution_device), - overlap=effective_overlap, - tiles_batch_size=effective_tiles_batch_size * batch_size, - resolution_scale=tiling_scale, + Kandinsky 6 SR ships with a text-free transformer (``use_text=False``) and has no + text encoder registered on the pipeline, so this never runs a real text encoder; + it only replicates precomputed ``cached_text_embeds`` to the batch size, or returns + placeholders when the transformer ignores text entirely. + + Raises: + ValueError: If the transformer needs text but no ``cached_text_embeds`` were given. + """ + if not getattr(self.transformer, "use_text", True): + # A use_text=False DiT ignores every text input (encoder, cross-attention, and pooled-text + # conditioning are all absent), so these zero-size tensors only satisfy the indexing and + # cu_seqlens bookkeeping in the denoising loop without ever loading a cached empty-caption file. + embeds = { + "text_embeds": torch.zeros(0, 1, device=device), + "pooled_embed": torch.zeros(batch_size, 1, device=device), + } + cu_seqlens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + return embeds, cu_seqlens + if cached_text_embeds is None: + raise ValueError("this SR transformer requires text; pass cached_text_embeds") + single_text = cached_text_embeds["text_embeds"] + single_pooled = cached_text_embeds["pooled_embed"] + seq_len = int(cached_text_embeds["cu_seqlens"][-1].item()) + text_embeds = { + "text_embeds": single_text.repeat(batch_size, 1).to(device), + "pooled_embed": single_pooled.repeat(batch_size, 1).to(device), + } + cu_seqlens = seq_len * torch.arange(batch_size + 1, dtype=torch.int32, device=device) + return text_embeds, cu_seqlens + + def _text_positions( + self, + image: torch.Tensor, + batch_size: int, + duration: int, + height: int, + width: int, + text_cu_seqlens: torch.Tensor, + ) -> tuple[torch.Tensor, list[torch.Tensor], torch.Tensor]: + """Build the packed visual/text RoPE positions ``self.transformer`` needs for one SR call.""" + device = image.device + visual_cu_seqlens = duration * torch.arange(batch_size + 1, dtype=torch.int32, device=device) + visual_rope_pos = [ + torch.cat([torch.arange(int(end), device=device) for end in torch.diff(visual_cu_seqlens).cpu()]), + torch.arange(height // self.transformer.patch_size[1], device=device), + torch.arange(width // self.transformer.patch_size[2], device=device), + ] + text_cu_seqlens = text_cu_seqlens.to(device) + text_rope_pos = torch.cat([torch.arange(int(end), device=device) for end in torch.diff(text_cu_seqlens).cpu()]) + return visual_cu_seqlens, visual_rope_pos, text_rope_pos + + def _predict( + self, + x: torch.Tensor, + t: torch.Tensor, + text_embeds: dict[str, torch.Tensor], + visual_cu_seqlens: torch.Tensor, + text_cu_seqlens: torch.Tensor, + visual_rope_pos: list[torch.Tensor], + text_rope_pos: torch.Tensor, + scale_factor: tuple[float, ...], + sparse_params: dict | None, + ) -> torch.Tensor: + """Run one ``self.transformer`` forward. SR has no guidance: a single conditional pass.""" + model_time = (t * 1000).to(dtype=x.dtype) + motion_score = torch.full((1,), 900.0, device=x.device, dtype=x.dtype) + return self.transformer( + x, + text_embeds["text_embeds"], + text_embeds["pooled_embed"], + model_time, + visual_cu_seqlens, + text_cu_seqlens, + visual_rope_pos, + text_rope_pos, + scale_factor=scale_factor, + sparse_params=sparse_params, + motion_score=motion_score, + ) + + @torch.no_grad() + def _denoise_native( + self, + img: torch.Tensor, + device: str | torch.device, + num_steps: int, + text_embeds: dict[str, torch.Tensor], + visual_cu_seqlens: torch.Tensor, + text_cu_seqlens: torch.Tensor, + visual_rope_pos: list[torch.Tensor], + text_rope_pos: torch.Tensor, + scale_factor: tuple[float, ...], + scheduler_scale: float, + start_timestep: float = 1.0, + progress_bar: Any = None, + ) -> torch.Tensor: + """Run the native Euler loop from ``start_timestep`` to 0 for the non-piflow SR transformer. + + This intentionally does not go through ``self.scheduler``: the SR transformer predicts + velocity directly and this loop applies the same shift formula + ``FlowMatchEulerDiscreteScheduler`` uses (``scheduler_scale`` mirrors its ``shift``), matching + the reference implementation bit-for-bit. + """ + img = img.to(device) + visual_cu_seqlens = visual_cu_seqlens.to(device) + text_cu_seqlens = text_cu_seqlens.to(device) + visual_rope_pos = [position.to(device) for position in visual_rope_pos] + text_rope_pos = text_rope_pos.to(device) + text_embeds = {key: value.to(device) for key, value in text_embeds.items()} + sparse_params = self._sparse_params(img, visual_cu_seqlens) + timesteps = torch.linspace(start_timestep, 0, num_steps, device=device) + timesteps = scheduler_scale * timesteps / (1 + (scheduler_scale - 1) * timesteps) + out_channels = self.transformer.in_visual_dim + for timestep, timestep_diff in zip(timesteps[:-1], torch.diff(timesteps), strict=False): + time = timestep.unsqueeze(0).repeat(visual_cu_seqlens.shape[0] - 1) + velocity = self._predict( + img, + time, + text_embeds, + visual_cu_seqlens, + text_cu_seqlens, + visual_rope_pos, + text_rope_pos, + scale_factor, + sparse_params, ) - progress_input = latents[0] if latents is not None else sr_video[0] if sr_video is not None else None - if progress_input is None: - raise RuntimeError("SR input disappeared during preparation") - progress_total, progress_steps, progress_tiles = _sr_progress_plan( - progress_input, - components, - run_config, - latent_input=latents is not None, - batch_size=batch_size, - scheduler=self.scheduler, + out_channels = velocity.shape[-1] + img[..., :out_channels] += timestep_diff * velocity + if progress_bar is not None: + progress_bar.update() + progress_bar.refresh() + return img[..., :out_channels] + + @torch.no_grad() + def _denoise_piflow( + self, + img: torch.Tensor, + text_embeds: dict[str, torch.Tensor], + visual_cu_seqlens: torch.Tensor, + text_cu_seqlens: torch.Tensor, + visual_rope_pos: list[torch.Tensor], + text_rope_pos: torch.Tensor, + scale_factor: tuple[float, ...], + *, + nfe: int, + out_dim: int, + device: str | torch.device | None = None, + progress_bar: Any = None, + ) -> torch.Tensor: + """Few-step π-Flow denoise: ``nfe`` transformer calls, via ``self.scheduler``. + + ``self.scheduler`` (a ``PiflowScheduler``) owns the segment schedule and the + network-free policy integration between transformer calls entirely from its own config. + """ + device = device if device is not None else img.device + img = img.to(device) + visual_cu_seqlens = visual_cu_seqlens.to(device) + text_cu_seqlens = text_cu_seqlens.to(device) + visual_rope_pos = [position.to(device) for position in visual_rope_pos] + text_rope_pos = text_rope_pos.to(device) + text_embeds = {key: value.to(device) for key, value in text_embeds.items()} + sparse_params = self._sparse_params(img, visual_cu_seqlens) + if nfe < 1: + raise ValueError(f"piflow denoising requires nfe >= 1, got {nfe}") + self.scheduler.set_timesteps(nfe, device=device) + if self.scheduler.timesteps.shape[0] != nfe: + raise ValueError(f"Piflow scheduler returned {self.scheduler.timesteps.shape[0]} steps, expected {nfe}") + x = img + for step_index in range(nfe): + timestep = self.scheduler.timesteps[step_index] + n_objects = visual_cu_seqlens.shape[0] - 1 + t_src = torch.full((n_objects,), float(timestep.item()), device=device, dtype=x.dtype) + v0 = self._predict( + x, + t_src, + text_embeds, + visual_cu_seqlens, + text_cu_seqlens, + visual_rope_pos, + text_rope_pos, + scale_factor, + sparse_params, ) - with self.progress_bar(total=progress_total) as progress_bar: - progress_bar.set_description(f"SR [{progress_steps} steps x {progress_tiles} tiles]") - has_latent_upscaler = ( - latent_upscaler_for_scale(self.latent_upscaler, run_config.resolution_scale) is not None - ) - if latents is not None: - frames = _run_batched_tiles( - latents, - components, - run_config, - latent_input=True, - progress_bar=progress_bar, - scheduler=self.scheduler, - ) - elif has_latent_upscaler: - lr_latents = list( - encode_lq_videos_to_lr_latents( - sr_video, - self.vae, - self._execution_device, - ) - ) - frames = _run_batched_tiles( - lr_latents, - components, - run_config, - latent_input=True, - progress_bar=progress_bar, - scheduler=self.scheduler, - ) - else: - frames = _run_batched_tiles( - sr_video, - components, - run_config, - latent_input=False, - progress_bar=progress_bar, - scheduler=self.scheduler, - ) - finally: - # Diffusers expects custom pipelines to restore the offloaded - # modules after a call. This also keeps the next invocation from - # observing a partially resident model chain. - self.maybe_free_model_hooks() - - path: str | list[str] | None = None - if save_path is not None: - paths = [ - str( - mux_video_audio( - frames[index], - audio, - output_path, - fps=fps or int(getattr(self._sr_params, "fps", TARGET_FPS)), - audio_sample_rate=audio_sample_rate, - source_video=source_videos[index], - ) + x_pred = self.scheduler.step(v0, timestep, x[..., :out_dim], return_dict=False)[0] + x = torch.cat([x_pred, x[..., out_dim:]], dim=-1) + if progress_bar is not None: + progress_bar.update() + progress_bar.refresh() + return x[..., :out_dim] + + def _scheduler_nfe(self) -> int | None: + """Return the configured π-Flow step count, or ``None`` for a regular scheduler.""" + if not bool(getattr(self.scheduler, "is_piflow", False)): + return None + nfe = getattr(self.scheduler.config, "nfe", None) + if nfe is None: + raise ValueError("PiflowScheduler used for SR must define nfe in scheduler_config.json") + return int(nfe) + + def _denoise( + self, + image: torch.Tensor, + text_embeds: dict[str, torch.Tensor], + text_cu_seqlens: torch.Tensor, + batch_size: int, + duration: int, + height: int, + width: int, + *, + scale_factor: tuple[float, ...], + num_steps: int, + scheduler_scale: float, + lq_noise_scale: float, + cap_noise_timestep: bool, + device: str | torch.device, + progress_bar: Any = None, + ) -> torch.Tensor: + """Run either the Euler or π-Flow SR denoising stage on ``self.transformer``/``self.scheduler``.""" + nfe = self._scheduler_nfe() + visual_cu_seqlens, visual_rope_pos, text_rope_pos = self._text_positions( + image, batch_size, duration, height, width, text_cu_seqlens + ) + if nfe is not None: + if cap_noise_timestep and self.transformer.instruct_type in ("noise", "hybrid"): + raise NotImplementedError( + "piflow sampler does not support cap_noise_timestep for noise/hybrid instruct" ) - for index, output_path in enumerate(save_paths) - ] - path = paths if isinstance(save_path, list) else paths[0] - audio_source = None - if source_video is not None: - audio_source = ( - [str(value) for value in source_videos] if isinstance(source_video, list) else str(source_video) + out_dim = int(getattr(self.transformer, "base_out_visual_dim", self.transformer.in_visual_dim)) + return self._denoise_piflow( + image, + text_embeds, + visual_cu_seqlens, + text_cu_seqlens, + visual_rope_pos, + text_rope_pos, + scale_factor, + nfe=nfe, + out_dim=out_dim, + device=device, + progress_bar=progress_bar, ) - output = Kandinsky6SRPipelineOutput( - frames=self._format_frames(frames, output_type), - audio=None if audio is None else (audio if isinstance(audio, list) else [audio]), - path=path, - metadata={ - "resolution_scale": requested_scale, - "num_inference_steps": run_config.num_steps, - "overlap": run_config.overlap, - "tiles_batch_size": effective_tiles_batch_size, - **({"audio_source": audio_source} if audio_source is not None else {}), - }, + start_timestep = ( + lq_noise_scale + if cap_noise_timestep and self.transformer.instruct_type in ("noise", "hybrid", "hybrid_anchor") + else 1.0 + ) + return self._denoise_native( + image, + device, + num_steps, + text_embeds, + visual_cu_seqlens, + text_cu_seqlens, + visual_rope_pos, + text_rope_pos, + scale_factor, + scheduler_scale, + start_timestep=start_timestep, + progress_bar=progress_bar, ) - if return_dict: - return output - return output.frames, output.audio, output.path, output.metadata @torch.no_grad() - def from_video( # noqa: PLR0913 + def _vae_decode( self, - video: SRBatchInput, - resolution_scale: float = 2.25, - num_inference_steps: int | None = None, - generator: torch.Generator | None = None, - overlap: float | None = None, - tiles_batch_size: int | None = None, - cached_text_embeds: dict[str, Tensor] | None = None, - save_path: str | Path | list[str | Path] | None = None, - fps: int | None = None, - audio: list[np.ndarray] | np.ndarray | None = None, - source_video: str | Path | list[str | Path] | None = None, - audio_sample_rate: int = 44100, - output_type: SROutputType = "pt", - return_dict: bool = True, - ) -> Kandinsky6SRPipelineOutput | tuple[Any, ...]: - """Super-resolve pixels and optionally copy audio from a source path.""" - frames, input_fps, input_paths = self._load_videos(video) - input_path = input_paths[0] if len(input_paths) == 1 else None - if source_video is None and save_path is not None: - if input_path is not None: - source_video = input_path - elif all(path is not None for path in input_paths): - source_video = [path for path in input_paths if path is not None] - with _execution_device_context(self._execution_device): - return self._run( - video=frames, - resolution_scale=resolution_scale, - num_inference_steps=num_inference_steps, - generator=generator, - overlap=overlap, - tiles_batch_size=tiles_batch_size, - cached_text_embeds=cached_text_embeds, - save_path=save_path, - fps=( - fps - if fps is not None - else round(input_fps[0]) - if len(input_fps) == 1 and input_fps[0] is not None - else None - ), - audio=audio, - source_video=source_video, - audio_sample_rate=audio_sample_rate, - output_type=output_type, - return_dict=return_dict, - ) + latent_visual: torch.Tensor, + batch_size: int, + duration: int, + height: int, + width: int, + ) -> torch.Tensor: + """Decode denoised SR latents into ``[batch, 3, T, H, W]`` uint8 frames with ``self.vae``.""" + device_type = latent_visual.device.type + with torch.autocast(device_type, dtype=torch.bfloat16, enabled=device_type == "cuda"): + all_latents = latent_visual.reshape(batch_size, duration, height, width, -1) + all_latents = (all_latents / self.vae.config.scaling_factor).permute(0, 4, 1, 2, 3) + decoded = [decode_latent_to_uint8(self.vae, all_latents[i : i + 1]) for i in range(batch_size)] + return torch.cat(decoded, dim=0) @torch.no_grad() - def from_latents( # noqa: PLR0913 + def _run_stages( self, - latents: Tensor | np.ndarray | str | Path | Sequence[Any], *, - resolution_scale: float = 2.25, - num_inference_steps: int | None = None, - generator: torch.Generator | None = None, - overlap: float | None = None, - tiles_batch_size: int | None = None, - kvae_bridge: bool = False, - cached_text_embeds: dict[str, Tensor] | None = None, - save_path: str | Path | list[str | Path] | None = None, - fps: int | None = None, - audio: list[np.ndarray] | np.ndarray | None = None, - source_video: str | Path | list[str | Path] | None = None, - audio_sample_rate: int = 44100, - output_type: SROutputType = "pt", - return_dict: bool = True, - ) -> Kandinsky6SRPipelineOutput | tuple[Any, ...]: - """Super-resolve raw latent samples or a batch of local ``.pt`` files.""" - channels = int(getattr(self.transformer, "in_visual_dim", 16)) - latent_value = _latent_batch(latents, channels) - with _execution_device_context(self._execution_device): - return self._run( - latents=latent_value, - resolution_scale=resolution_scale, - num_inference_steps=num_inference_steps, - generator=generator, - overlap=overlap, - tiles_batch_size=tiles_batch_size, - kvae_bridge=kvae_bridge, - cached_text_embeds=cached_text_embeds, - save_path=save_path, - fps=fps, - audio=audio, - source_video=source_video, - audio_sample_rate=audio_sample_rate, - output_type=output_type, - return_dict=return_dict, + scale_factor: tuple[float, ...], + num_steps: int, + scheduler_scale: float, + device: str | torch.device, + generator: torch.Generator | None, + lq_noise_scale: float, + lq_noise_type: str, + lq_channel_noise_scale: float, + cap_noise_timestep: bool, + cached_text_embeds: dict[str, torch.Tensor] | None, + lq_videos: list[torch.Tensor] | None = None, + lq_latents: torch.Tensor | None = None, + n_samples: int | None = None, + progress_bar: Any = None, + ) -> torch.Tensor: + """Run the four SR stages — prepare latents, encode text, denoise, decode — for one tile batch.""" + lq_latent, image, batch_size, duration, height, width = self._prepare_sr_latents( + lq_videos=lq_videos, + lq_latents=lq_latents, + n_samples=n_samples, + device=device, + generator=generator, + lq_noise_scale=lq_noise_scale, + lq_noise_type=lq_noise_type, + lq_channel_noise_scale=lq_channel_noise_scale, + ) + text_embeds, text_cu_seqlens = self._encode_sr_text(batch_size, device, cached_text_embeds) + latent_visual = self._denoise( + image, + text_embeds, + text_cu_seqlens, + batch_size, + duration, + height, + width, + scale_factor=scale_factor, + num_steps=num_steps, + scheduler_scale=scheduler_scale, + lq_noise_scale=lq_noise_scale, + cap_noise_timestep=cap_noise_timestep, + device=device, + progress_bar=progress_bar, + ) + return self._vae_decode(latent_visual, batch_size, duration, height, width) + + def _tile_batch_kwargs( + self, + run_config: RunConfig, + scale_factor: tuple[float, ...], + cached_text_embeds: dict[str, torch.Tensor] | None, + seed: int, + **inputs: Any, + ) -> dict[str, Any]: + # A fresh per-tile Generator, seeded deterministically from the run's base seed plus this + # tile batch's offset, so a tile's noise never depends on how many other tiles ran before it. + generator = torch.Generator(device=run_config.device) + generator.manual_seed(seed) + transformer = self.transformer + return { + **inputs, + "scale_factor": scale_factor, + "num_steps": run_config.num_steps, + "scheduler_scale": transformer.scheduler_scale, + "generator": generator, + "device": run_config.device, + "lq_noise_scale": transformer.lq_noise_scale, + "lq_noise_type": transformer.lq_noise_type, + "lq_channel_noise_scale": transformer.lq_channel_noise_scale, + "cap_noise_timestep": transformer.cap_noise_timestep, + "cached_text_embeds": cached_text_embeds, + } + + def _loaded_lu_scales(self) -> tuple[int, ...]: + if self.latent_upscaler is None: + return () + return tuple(self.latent_upscaler.scales) + + def _run_tile_batches( + self, + tile_inputs: list[torch.Tensor], + run_config: RunConfig, + scale_factor: tuple[float, ...], + *, + cached_text_embeds: dict[str, torch.Tensor] | None, + prepare_chunk: Any | None = None, + progress_bar: Any = None, + ) -> list[torch.Tensor]: + """Run the batched sampler loop for latent or pixel tile inputs. + + The transformer and SR VAE stay resident for the complete tile loop. Only the + optional latent-upscaler is streamed per chunk. + """ + outputs: list[torch.Tensor] = [] + for start in range(0, len(tile_inputs), run_config.tiles_batch_size): + raw_chunk = tile_inputs[start : start + run_config.tiles_batch_size] + chunk = prepare_chunk(raw_chunk) if prepare_chunk is not None else raw_chunk + input_kwargs: dict[str, Any] = ( + {"lq_videos": chunk} + if prepare_chunk is None + else {"lq_latents": torch.cat(chunk, dim=0), "n_samples": len(raw_chunk)} + ) + kwargs = self._tile_batch_kwargs( + run_config, + scale_factor, + cached_text_embeds, + run_config.seed + start, + **input_kwargs, + ) + sr_batch = self._run_stages(**kwargs, progress_bar=progress_bar) + if sr_batch.ndim != 5: + raise ValueError(f"SR sampler must return [batch,C,T,H,W], got {tuple(sr_batch.shape)}") + outputs.extend(sample.float().cpu() for sample in sr_batch) + return outputs + + def _run_batched_tiles( + self, + samples: list[Tensor], + run_config: RunConfig, + *, + latent_input: bool, + cached_text_embeds: dict[str, torch.Tensor] | None, + progress_bar: Any, + ) -> Tensor: + """Run equal-shape samples tile-major and restore the sample batch.""" + visual_size = self.transformer.visual_size + batch_size = len(samples) + scale_factor = self.transformer.scale_factor[visual_size] + if latent_input: + spatial_factor = self.spatial_factor + height = samples[0].shape[-2] * spatial_factor + width = samples[0].shape[-1] * spatial_factor + _, _, pixel_grid = _tile_geometry( + height, + width, + visual_size, + run_config.resolution_scale, + run_config.min_overlap, + spatial_factor, + ) + tile_grid = latent_tile_grid_from_pixel_grid(pixel_grid, spatial_factor) + input_grid = tile_grid + stitch_grid = pixel_grid + tile_sets = [extract_all_tiles(sample, tile_grid) for sample in samples] + if run_config.resolution_scale not in self._loaded_lu_scales(): + raise ValueError( + f"The latent-upscaler path has no {run_config.resolution_scale}x model " + f"(loaded LU scales: {self._loaded_lu_scales()})" + ) + + def prepare_chunk(chunk: list[Tensor]) -> list[Tensor]: + return [ + self._upscale_lr_latent_tile(tile, run_config.resolution_scale, run_config.device) + for tile in chunk + ] + else: + height, width = samples[0].shape[-2:] + spatial_factor = self.spatial_factor + base, _tile_hw, grid = _tile_geometry( + height, + width, + visual_size, + run_config.resolution_scale, + run_config.min_overlap, + spatial_factor, ) + tile_sets = [ + _upsample_tiles_to_base(extract_all_tiles(sample, grid), base[0], base[1]) for sample in samples + ] + input_grid = grid + stitch_grid = grid + prepare_chunk = None + + tile_inputs = [ + tile_sets[sample_index][tile_index] + for tile_index in range(input_grid.total_tiles) + for sample_index in range(batch_size) + ] + outputs = self._run_tile_batches( + tile_inputs, + run_config, + scale_factor, + cached_text_embeds=cached_text_embeds, + prepare_chunk=prepare_chunk, + progress_bar=progress_bar, + ) + return _stitch_batch(outputs, stitch_grid, batch_size, height, width, run_config.resolution_scale) + + def _sr_progress_plan( + self, + value: Tensor, + run_config: RunConfig, + *, + latent_input: bool, + batch_size: int = 1, + ) -> tuple[int, int, int]: + """Return total updates, denoising steps, and tile batches.""" + visual_size = self.transformer.visual_size + spatial_factor = self.spatial_factor + height, width = value.shape[-2:] + if latent_input: + height *= spatial_factor + width *= spatial_factor + _, _, grid = _tile_geometry( + height, width, visual_size, run_config.resolution_scale, run_config.min_overlap, spatial_factor + ) + tile_count = grid.total_tiles * batch_size + chunk_count = max(1, math.ceil(tile_count / run_config.tiles_batch_size)) + nfe = self._scheduler_nfe() + denoise_steps = nfe if nfe is not None else run_config.num_steps - 1 + denoise_steps = max(1, denoise_steps) + return chunk_count * denoise_steps, denoise_steps, chunk_count + + @staticmethod + def _seed_from_generator(generator: torch.Generator | None, device: torch.device) -> int: + if generator is None: + return int(torch.randint(0, 2**31, (1,), device=device).item()) + generator_device = torch.device(getattr(generator, "device", "cpu")) + return int(torch.randint(0, 2**31, (1,), generator=generator, device=generator_device).item()) @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - video: SRBatchInput | None = None, - latents: Tensor | np.ndarray | str | Path | Sequence[Any] | None = None, + video: Tensor | None = None, + latents: Tensor | None = None, resolution_scale: float = 2.25, num_inference_steps: int = 5, generator: torch.Generator | None = None, - overlap: float = 0.25, + min_overlap: float = 0.20, tiles_batch_size: int = 1, kvae_bridge: bool = False, cached_text_embeds: dict[str, Tensor] | None = None, - save_path: str | Path | list[str | Path] | None = None, - fps: int | None = None, - audio: list[np.ndarray] | np.ndarray | None = None, - source_video: str | Path | list[str | Path] | None = None, - audio_sample_rate: int = 44100, - output_type: SROutputType = "pt", + output_type: str = "pt", return_dict: bool = True, - ) -> Kandinsky6SRPipelineOutput | tuple[Any, ...]: - """Super-resolve one or more equal-size pixel videos or latent inputs. + ) -> Kandinsky6SRPipelineOutput | tuple[Tensor]: + r"""Super-resolve one or more equal-size pixel videos or latent inputs. Args: - video: Pixel video in ``[T,C,H,W]``/``[T,H,W,C]`` form, a batch - in a supported rank-5 layout, or local video paths. - latents: Raw latent tensor, a supported rank-5 latent batch, or - local ``.pt`` files. + video: LQ pixel video(s) to super-resolve, ``uint8`` in ``[0, 255]``, + shape ``(3, T, H, W)`` or a batch ``(B, 3, T, H, W)`` — the same + layout [`Kandinsky6TI2VAPipeline`] returns with + ``output_type="pt"``. ``T`` must be ``1 + 8k`` frames (align or + truncate longer/misaligned clips yourself) and at most + ``MAX_NUM_FRAMES`` (121, ~5s at 24fps). Mutually exclusive with + ``latents``. + latents: Pre-encoded SR-VAE latent(s) already at the target + (post-upscale) resolution, shape ``(C, T, H, W)`` or a batch + ``(B, C, T, H, W)``, **unscaled** (not yet multiplied by + ``self.vae.config.scaling_factor``). Mutually exclusive with + ``video``. Pass ``kvae_bridge=True`` when these latents + instead come from a different (base) VAE at the *source* + resolution — they are then decoded through ``self.source_vae`` + and re-enter the normal ``video`` path (pre-upscale, encode, + latent-upscale) rather than being used directly. resolution_scale: Total spatial upscale factor. Supported values are ``2``, ``4``, and ``2.25``. num_inference_steps: Number of denoising steps for the SR model. generator: Optional random generator used to derive the sampling seed. - overlap: Fraction of overlap between adjacent tiles. + min_overlap: Minimum fraction of overlap between adjacent tiles. tiles_batch_size: Number of tiles processed in one sampler batch. kvae_bridge: Whether latent inputs come from a base VAE and need to be decoded before SR. cached_text_embeds: Optional cached empty-caption embeddings. - save_path: Optional output video path or one path per input. - fps: Optional output video frame rate. - audio: Optional audio waveform to mux. - source_video: Optional source video whose audio is copied. - audio_sample_rate: Sample rate used when muxing audio. output_type: ``pt``/``torch`` or ``np``/``numpy``. return_dict: Whether to return ``Kandinsky6SRPipelineOutput``. @@ -3043,111 +1433,91 @@ def __call__( ``return_dict=False``. """ if (video is None) == (latents is None): - raise ValueError("pass exactly one of video or latents") - + raise ValueError("pass exactly one of `video` or `latents`") if video is not None: - return self.from_video( - video, - resolution_scale=resolution_scale, - num_inference_steps=num_inference_steps, - generator=generator, - overlap=overlap, - tiles_batch_size=tiles_batch_size, - cached_text_embeds=cached_text_embeds, - save_path=save_path, - fps=fps, - audio=audio, - source_video=source_video, - audio_sample_rate=audio_sample_rate, - output_type=output_type, - return_dict=return_dict, - ) + video = _prepare_batch(video, "video", 3, dtype=torch.uint8) + if latents is not None: + latents = _prepare_batch(latents, "latents", int(self.transformer.in_visual_dim)) - return self.from_latents( - latents, + self.check_inputs( + video=video, + latents=latents, resolution_scale=resolution_scale, num_inference_steps=num_inference_steps, - generator=generator, - overlap=overlap, + min_overlap=min_overlap, tiles_batch_size=tiles_batch_size, kvae_bridge=kvae_bridge, - cached_text_embeds=cached_text_embeds, - save_path=save_path, - fps=fps, - audio=audio, - source_video=source_video, - audio_sample_rate=audio_sample_rate, output_type=output_type, - return_dict=return_dict, - ) # type: ignore[arg-type] - - -def _to_tchw(value: Tensor | np.ndarray | Sequence[Any]) -> Tensor: # noqa: PLR0912, PLR2004 - """Normalize common frame layouts to contiguous uint8 ``[T,C,H,W]``.""" - if isinstance(value, Tensor): - tensor = value.detach() - else: - if isinstance(value, Sequence) and not isinstance(value, np.ndarray): - value = np.stack([np.asarray(frame) for frame in value]) - tensor = torch.as_tensor(value) - if tensor.ndim == VIDEO_RANK: - if tensor.shape[-1] in (1, 3, 4): - tensor = tensor.permute(0, 3, 1, 2) - elif tensor.shape[1] not in (1, 3, 4): - raise ValueError("4D video must be [T,C,H,W] or [T,H,W,C]") - else: - raise ValueError(f"video sample must have rank 4, got shape {tuple(tensor.shape)}") - if tensor.is_floating_point(): - if tensor.numel() and float(tensor.detach().amax()) <= 1.0: - tensor = tensor * 255.0 - tensor = tensor.round() - return tensor.clamp(0, 255).to(torch.uint8).contiguous() - - -def _coerce_sr_params( - value: Any | None, -) -> Any | None: - if value is not None: - if isinstance(value, Mapping) and "sr" in value: - value = value["sr"] - if isinstance(value, Mapping): - values = dict(value) - if "scale_factor" in values and isinstance(values["scale_factor"], Mapping): - values["scale_factor"] = { - int(key): [float(item) for item in items] for key, items in values["scale_factor"].items() - } - if "visual_size" in values and isinstance(values["visual_size"], int): - values["visual_size"] = [values["visual_size"]] - return SimpleNamespace(**values) - return value - return None - - -def _decode_source_latent_video( - raw_latents: Tensor, - source_vae: nn.Module, - device: str | torch.device, -) -> Tensor: - """Decode a base-VAE latent to ``[T,C,H,W]`` pixels for the KVAE bridge.""" - if raw_latents.ndim != 4: # noqa: PLR2004 - raise ValueError(f"raw_latents must have rank 4 [T,C,H,W], got {tuple(raw_latents.shape)}") - latent_5d = raw_latents.permute(1, 0, 2, 3).unsqueeze(0).to(device=device) - latent_5d = cast_to_module_dtype(source_vae, latent_5d) - with torch.autocast("cuda", dtype=torch.bfloat16, enabled=str(device).startswith("cuda")): - try: - decoded = _call_module_method(source_vae, "decode", latent_5d, alternative_fwd=True).sample - except TypeError as exc: - if "alternative_fwd" not in str(exc): - raise - decoded = _call_module_method(source_vae, "decode", latent_5d).sample - return ( - ((decoded.squeeze(0).float().clamp(-1, 1) + 1.0) * 127.5) - .round() - .clamp(0, 255) - .to(torch.uint8) - .permute(1, 0, 2, 3) - .cpu() - ) - + ) + tiling_scale, pre_upscale = resolve_scale_request(float(resolution_scale)) + batch_size = latents.shape[0] if latents is not None else video.shape[0] + sr_video = list(video.unbind(0)) if video is not None else None + sr_latents = list(latents.unbind(0)) if latents is not None else None -__all__ = ["Kandinsky6SRPipeline"] + with _execution_device_context(self._execution_device): + try: + if sr_latents is not None and kvae_bridge: + sr_video = [ + _decode_source_latent_video(latent, self.source_vae, self._execution_device) + for latent in sr_latents + ] + sr_latents = None + if sr_video is not None and pre_upscale != 1.0: + spatial_factor = self.spatial_factor + sr_video = [pre_upscale_video(item, pre_upscale, spatial_factor) for item in sr_video] + + run_config = RunConfig( + device=str(self._execution_device), + num_steps=num_inference_steps, + seed=self._seed_from_generator(generator, self._execution_device), + min_overlap=min_overlap, + tiles_batch_size=tiles_batch_size * batch_size, + resolution_scale=tiling_scale, + ) + progress_input = sr_latents[0] if sr_latents is not None else sr_video[0] + progress_total, progress_steps, progress_tiles = self._sr_progress_plan( + progress_input, + run_config, + latent_input=sr_latents is not None, + batch_size=batch_size, + ) + with self.progress_bar(total=progress_total) as progress_bar: + progress_bar.set_description(f"SR [{progress_steps} steps x {progress_tiles} tiles]") + has_latent_upscaler = run_config.resolution_scale in self._loaded_lu_scales() + if sr_latents is not None: + frames = self._run_batched_tiles( + sr_latents, + run_config, + latent_input=True, + cached_text_embeds=cached_text_embeds, + progress_bar=progress_bar, + ) + elif has_latent_upscaler: + lr_latents = [ + encode_lq_video_to_lr_latent(item, self.vae, self._execution_device) for item in sr_video + ] + frames = self._run_batched_tiles( + lr_latents, + run_config, + latent_input=True, + cached_text_embeds=cached_text_embeds, + progress_bar=progress_bar, + ) + else: + frames = self._run_batched_tiles( + sr_video, + run_config, + latent_input=False, + cached_text_embeds=cached_text_embeds, + progress_bar=progress_bar, + ) + finally: + # Diffusers expects custom pipelines to restore the offloaded + # modules after a call. This also keeps the next invocation from + # observing a partially resident model chain. + self.maybe_free_model_hooks() + + output = Kandinsky6SRPipelineOutput(frames=self._format_frames(frames, output_type)) + if not return_dict: + return (output.frames,) + return output diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_output.py b/src/diffusers/pipelines/kandinsky6/pipeline_output.py index 042f8d1f5b6c..1f47e4d5818a 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_output.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_output.py @@ -23,8 +23,6 @@ from ...utils import BaseOutput -Audio = list[np.ndarray] | np.ndarray | None - @dataclass class Kandinsky6TI2VAPipelineOutput(BaseOutput): @@ -39,13 +37,9 @@ class Kandinsky6SRPipelineOutput(BaseOutput): """Output of the K6 super-resolution pipeline.""" frames: Any - audio: Audio = None - path: str | list[str] | None = None - metadata: dict[str, Any] | None = None __all__ = [ - "Audio", "Kandinsky6SRPipelineOutput", "Kandinsky6TI2VAPipelineOutput", ] From 089a5ae97d8deb96edbb6eb38f95c5376e73f5e9 Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 28 Sep 2026 14:33:09 +0300 Subject: [PATCH 18/62] patch sr transformer autocast --- .../models/transformers/transformer_kandinsky6_sr.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index 149ee50236e8..edab68e52872 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -1130,6 +1130,7 @@ def forward( Returns: Updated visual embeddings. """ + dtype = visual_embed.dtype if self.use_text: self_attn_params, cross_attn_params, ff_params = torch.chunk(self.visual_modulation(time_embed), 3, dim=-1) else: @@ -1152,6 +1153,7 @@ def forward( device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" ): visual_embed = visual_embed + gate_t.index_select(0, time_embed_idx) * visual_out + visual_embed = visual_embed.to(dtype=dtype) # --- Cross-attention --- (skipped when text-free) # Norm -> modulate by (scale + 1, shift) -> CrossAttn -> gate @@ -1170,6 +1172,7 @@ def forward( device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" ): visual_embed = visual_embed + gate_t.index_select(0, time_embed_idx) * visual_out + visual_embed = visual_embed.to(dtype=dtype) # --- Feed-forward --- # Norm -> modulate by (scale + 1, shift) -> FF -> gate @@ -1187,7 +1190,7 @@ def forward( device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" ): visual_embed = visual_embed + gate_t.index_select(0, time_embed_idx) * visual_out - return visual_embed + return visual_embed.to(dtype=dtype) class Kandinsky6SRTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin): From 3cc8f887db87f5c861bae032421a8dcca090a2b3 Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 28 Sep 2026 17:52:59 +0300 Subject: [PATCH 19/62] major refactor --- K6_DIFFUSERS_LAUNCH_GUIDE.md | 402 ++++ docs/source/en/_toctree.yml | 2 +- .../en/api/models/kandinsky6_transformer.md | 8 +- docs/source/en/api/models/kandinsky6_vae.md | 12 +- docs/source/en/api/pipelines/kandinsky6.md | 81 +- docs/source/en/api/schedulers/piflow.md | 11 +- notebooks/kandinsky6_sr.ipynb | 96 + notebooks/kandinsky6_sr.py | 105 + notebooks/kandinsky6_ti2va.ipynb | 240 +++ notebooks/kandinsky6_ti2va.py | 182 ++ notebooks/kandinsky6_ti2va_sr.ipynb | 130 ++ notebooks/kandinsky6_ti2va_sr.py | 113 ++ param_counts.json | 11 + src/diffusers/__init__.py | 14 +- src/diffusers/models/__init__.py | 14 +- src/diffusers/models/autoencoders/__init__.py | 4 +- .../autoencoders/autoencoder_kandinsky6_sr.py | 1394 +++++-------- .../autoencoders/autoencoder_mmaudio.py | 1565 +++++---------- src/diffusers/models/transformers/__init__.py | 4 +- .../transformers/transformer_kandinsky6.py | 314 ++- .../transformers/transformer_kandinsky6_sr.py | 1628 ++++------------ src/diffusers/pipelines/__init__.py | 2 - .../pipelines/kandinsky6/__init__.py | 30 +- .../pipelines/kandinsky6/latent_upscaler.py | 1560 --------------- .../kandinsky6/modeling_latent_upscaler.py | 296 +++ .../kandinsky6/pipeline_kandinsky6_sr.py | 1730 +++-------------- .../kandinsky6/pipeline_kandinsky6_ti2va.py | 1362 ++++++------- .../pipelines/kandinsky6/pipeline_output.py | 42 +- src/diffusers/schedulers/__init__.py | 2 +- src/diffusers/schedulers/scheduling_piflow.py | 10 +- src/diffusers/utils/dummy_pt_objects.py | 135 +- .../dummy_torch_and_transformers_objects.py | 60 + .../test_models_autoencoder_kandinsky6_sr.py | 105 + .../test_models_autoencoder_mmaudio.py | 99 + .../test_models_transformer_kandinsky6.py | 121 +- .../test_models_transformer_kandinsky6_sr.py | 122 ++ .../kandinsky6/test_models_latent_upscaler.py | 96 + .../kandinsky6/test_pipeline_kandinsky6_sr.py | 133 ++ .../test_pipeline_kandinsky6_ti2va.py | 96 +- tests/schedulers/test_scheduler_piflow.py | 1 - 40 files changed, 4787 insertions(+), 7545 deletions(-) create mode 100644 K6_DIFFUSERS_LAUNCH_GUIDE.md create mode 100644 notebooks/kandinsky6_sr.ipynb create mode 100644 notebooks/kandinsky6_sr.py create mode 100644 notebooks/kandinsky6_ti2va.ipynb create mode 100644 notebooks/kandinsky6_ti2va.py create mode 100644 notebooks/kandinsky6_ti2va_sr.ipynb create mode 100644 notebooks/kandinsky6_ti2va_sr.py create mode 100644 param_counts.json delete mode 100644 src/diffusers/pipelines/kandinsky6/latent_upscaler.py create mode 100644 src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py create mode 100644 tests/models/autoencoders/test_models_autoencoder_kandinsky6_sr.py create mode 100644 tests/models/autoencoders/test_models_autoencoder_mmaudio.py create mode 100644 tests/models/transformers/test_models_transformer_kandinsky6_sr.py create mode 100644 tests/pipelines/kandinsky6/test_models_latent_upscaler.py create mode 100644 tests/pipelines/kandinsky6/test_pipeline_kandinsky6_sr.py diff --git a/K6_DIFFUSERS_LAUNCH_GUIDE.md b/K6_DIFFUSERS_LAUNCH_GUIDE.md new file mode 100644 index 000000000000..a204b3eb1753 --- /dev/null +++ b/K6_DIFFUSERS_LAUNCH_GUIDE.md @@ -0,0 +1,402 @@ +# Launching Kandinsky 6 through Diffusers + +How to run every Kandinsky 6 (K6) model with the Diffusers code in +`diffusers-new-model-addition-kandinskyX` (branch `Kandinsky6`). The launches follow the notebooks in +`diffusers-new-model-addition-kandinskyX/notebooks/`. Code and Hub configs were checked as described in +[Verification status](#verification-status). + +## 1. What you can launch + +| # | Launch | Hub repo | Pipeline | Steps / guidance | Notes | +|---|--------|----------|----------|------------------|-------| +| 1 | Text → video + audio (T2VA), quality | `kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers` | `Kandinsky6TI2VAPipeline` | 50 steps, `guidance_scale=5.0` | `FlowMatchEulerDiscreteScheduler`; MagCache available | +| 2 | Text → video + audio, fast | `kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers` | `Kandinsky6TI2VAPipeline` | 16 steps, `guidance_scale=1.0` (mandatory) | `PiflowScheduler` rejects any other guidance; no MagCache | +| 3 | Image → video + audio (I2VA) | either repo above | `Kandinsky6TI2VAPipeline` | as 1 or 2 | pass `image=` | +| 4 | Video super-resolution, flow matching | `kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers` | `Kandinsky6SRPipeline` | `num_inference_steps=5` (= 4 Euler steps per tile) | x2, x2.25 (default), x4 | +| 5 | Video super-resolution, distilled | `kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers` | `Kandinsky6SRPipeline` | 2 model evaluations per tile, `num_inference_steps` ignored | `PiflowScheduler` | +| 6 | T2VA → SR chain | one of 1/2 + one of 4/5 | both | as above | video tensor handed over directly; audio muxed in afterward | + +Component layout (from each repo's `model_index.json`): + +- **T2VA repos:** `transformer` (`Kandinsky6Transformer3DModel`), `vae` (`AutoencoderKLHunyuanVideo`), `scheduler`, + `text_encoder` (Qwen2.5-VL), `tokenizer`, `text_encoder_2` (CLIP), `tokenizer_2`, `audio_vae` (`Kandinsky6AudioVAE`). + Audio is 44.1 kHz, video is 24 fps. +- **VSR repos:** `transformer` (`Kandinsky6SRTransformer3DModel`, 1.41B), `vae` (`Kandinsky6SRVAE`, 1.74B), + `latent_upscaler` (`Kandinsky6SRLatentUpscalerBank`, x2 + x4, 3.65B), `scheduler`. + +All K6 repos are private: the Hub answers 401 without a login (section 2.3). `Kandinsky-6.0-Lite-Diffusers` appears +only in the FastVideo and SGLang docs, not in the Diffusers code or notebooks, so it is not covered here. + +## 2. Setup + +### 2.1 Install this branch + +Always use the uv `.venv` of the checkout. Do not use conda and do not use `uv run`/`uv sync`, which can swap your +torch build. + +```bash +BASE_PATH="something" # folder that contains the K6 repos +cd "$BASE_PATH/diffusers-new-model-addition-kandinskyX" +uv venv && source .venv/bin/activate +uv pip install -e . +uv pip install av ipykernel +python -c "from diffusers import Kandinsky6TI2VAPipeline, Kandinsky6SRPipeline; print('K6 OK')" +``` + +- Stock Diffusers from PyPI has no `Kandinsky6*` classes, so it fails with an `ImportError`. Use this branch. +- Install a torch build that matches the GPU driver first (torch >= 2.6). Then run `uv pip install -e .`. +- `av` (PyAV) is required by `diffusers.utils.encode_video` to write mp4 files (with or without audio). Both + pipelines take/return tensors only — neither reads nor writes video files itself. +- `natten` and `flash_attn` are not required. The default `native` backend needs neither. +- On a remote GPU box, run the same block from the checkout on that box after `git pull` there. + +### 2.2 GPU memory + +The pipelines are meant to run with `enable_model_cpu_offload()`. That keeps one component on the GPU at a time. The +T2VA order is `text_encoder -> text_encoder_2 -> transformer -> vae -> audio_vae`. The SR order is +`source_vae -> latent_upscaler -> transformer -> vae`. Without offload, the full T2VA pipeline (Qwen2.5-VL plus a +60-block, 4096-wide transformer) needs a large GPU. + +### 2.3 Hugging Face access + +Log in once, interactively, and never paste a token into a script or notebook: + +```bash +hf auth login +hf auth whoami +``` + +For non-interactive shells use `export HF_TOKEN=...` in the shell, never in a file. After the first download the +snapshot is cached in `~/.cache/huggingface/hub`. Set `HF_HUB_OFFLINE=1` to run without network. + +Note: the notebooks contain `os.environ["HF_TOKEN"] = "TOKEN"`. Delete that line rather than putting a real token there. + +## 3. Text/image → video + audio (`Kandinsky6TI2VAPipeline`) + +Defaults of `__call__`: `height=512, width=768, num_frames=121, sample_fps=24.0, num_inference_steps=50, +guidance_scale=5.0, max_sequence_length=1024, sample_audio=True, expand_prompts=False, output_type="pt"`. +`height` and `width` must be divisible by 8, and `max_sequence_length` must be between 1 and 1024. The notebooks use +`480 x 864`, which is the size the `Pro` checkpoints were configured for (`k6_pro_125_480_864_*` in `k6_video`). + +`output_type="pt"` returns `frames` as a **uint8 tensor `[B, 3, T, H, W]` in 0..255**. `audio` is a list with one 1-D +mono waveform per video. Prompts can carry speech as `spoken words`. + +### 3.1 Launch 1: Pro-sft, quality (50 steps, CFG 5.0, MagCache) + +Save as `run_t2va_sft.py` and run it from the venv: + +```python +import torch +from diffusers import Kandinsky6TI2VAPipeline +from diffusers.hooks import MagCacheConfig +from diffusers.utils import encode_video + +MODEL = "kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers" +GPU_ID = 0 +STEPS, GUIDANCE, FPS, SEED = 50, 5.0, 24.0, 1137 +USE_MAGCACHE = True +MODE = "t2va" # "t2va" or "i2va": selects the MagCache coefficients + +pipe = Kandinsky6TI2VAPipeline.from_pretrained(MODEL, dtype=torch.bfloat16) +pipe.transformer.set_attention_backend("native") # "_flash_3" needs FlashAttention 3 (Hopper) + +cfg = pipe.transformer.config.get("magcache") +if USE_MAGCACHE and cfg is not None: + ratios = cfg["mag_ratios"] + ratios = ratios[MODE] if isinstance(ratios, dict) else ratios + pipe.transformer.enable_cache( + MagCacheConfig( + threshold=cfg["threshold"], + max_skip_steps=cfg["max_skip_steps"], + retention_ratio=cfg["retention_ratio"], + num_inference_steps=2 * STEPS, # coefficients interleave cond and uncond, so 2 forwards per step + mag_ratios=list(ratios), + ) + ) + +pipe.enable_model_cpu_offload(gpu_id=GPU_ID) + +result = pipe( + prompt="A news anchor delivers the evening news in a clear, neutral tone, saying: " + "Good evening, ladies and gentlemen, and welcome to the evening news.", + negative_prompt="blurry, distorted, low quality, noisy audio", + height=480, + width=864, + num_frames=121, + sample_fps=FPS, + num_inference_steps=STEPS, + guidance_scale=GUIDANCE, + sample_audio=True, + generator=torch.Generator().manual_seed(SEED), + output_type="pt", +) + +video = result.frames[0].permute(1, 2, 3, 0) # [T,H,W,3] uint8 +audio = torch.as_tensor(result.audio[0])[:, None].repeat(1, 2) # mono -> stereo [samples, 2] +encode_video(video, fps=int(FPS), output_path="t2va_sft.mp4", + audio=audio, audio_sample_rate=int(pipe.audio_sample_rate)) +print("saved t2va_sft.mp4") +``` + +Notes: + +- The MagCache coefficients live in the transformer config under `magcache`, with separate lists for `t2va` and `i2va`. + Set `MODE` to match the generation you run. +- To turn MagCache off, set `USE_MAGCACHE = False`. You can also call `pipe.transformer.disable_cache()`. +- MagCache is used only with the sft checkpoint. The authors' notebook disables it for the distilled one. + +### 3.2 Launch 2: Pro-distill, fast (PiFlow, 16 steps, guidance 1.0) + +Same script with these changes: + +```python +MODEL = "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers" +STEPS, GUIDANCE = 16, 1.0 +USE_MAGCACHE = False +``` + +The output path can become `t2va_distill.mp4`. `guidance_scale` must stay `1.0` here: the denoise loop raises +`ValueError("PiflowScheduler requires guidance_weight=1.0")` for any other value. With guidance 1.0 there is no +unconditional pass, so no `negative_prompt` is needed. + +You can choose steps and guidance automatically, as the notebook does: + +```python +is_piflow = pipe.scheduler.__class__.__name__ == "PiflowScheduler" +STEPS = 16 if is_piflow else 50 +GUIDANCE = 1.0 if is_piflow else 5.0 +``` + +### 3.3 Launch 3: image → video + audio (I2VA) + +Pass a reference image path (or a PIL image) as `image=`. Set `MODE = "i2va"` for MagCache. For a batch, pass one image +per prompt: + +```python +BASE_PATH = "something" # folder that contains the K6 repos + +result = pipe( + prompt=["A woman turns to the camera and says Hello there!"], + image=[f"{BASE_PATH}/k6_video/assets/i2va_input.png"], + negative_prompt=["blurry, distorted, low quality, noisy audio"], + height=480, width=864, num_frames=121, sample_fps=24.0, + num_inference_steps=STEPS, guidance_scale=GUIDANCE, + sample_audio=True, generator=torch.Generator().manual_seed(1137), output_type="pt", +) +``` + +The repo has a sample image at `k6_video/assets/i2va_input.png`. It is not in the diffusers checkout. + +### 3.4 Optional switches + +| Switch | How | Comment | +|--------|-----|---------| +| Video only | `sample_audio=False` | `result.audio` is then `None`; skip the `audio=` arguments | +| Batching | lists for `prompt`, `negative_prompt`, `image` | the notebook uses batch size 2 | +| Prompt expansion | `expand_prompts=True` | rewrites the prompt with the Qwen2.5-VL model first; slower. Default is `False`. The notebook hard-codes `True` in the call | +| NumPy output | `output_type="np"` | `[B, T, H, W, 3]` | +| Faster attention | `pipe.transformer.set_attention_backend("_flash_3")` | needs FlashAttention 3 (Hopper). Other valid names: `flash`, `native`, `sage`, `xformers`, ... | +| `torch.compile` | see the block below | the notebook uses `native` attention when compiling | + +Notebook-style block compile (only for repeated runs; call it after loading and before `enable_model_cpu_offload`): + +```python +from diffusers.models.transformers.transformer_kandinsky6 import ( + Kandinsky6FusedTransformerDecoderBlock, + Kandinsky6TransformerDecoderBlock, + Kandinsky6TransformerEncoderBlock, +) + +torch.set_float32_matmul_precision("high") +pipe.transformer.set_attention_backend("native") +for block in pipe.transformer.modules(): + if isinstance(block, Kandinsky6FusedTransformerDecoderBlock): + block.forward = torch.compile(block.forward) + elif isinstance(block, (Kandinsky6TransformerEncoderBlock, Kandinsky6TransformerDecoderBlock)): + block.forward = torch.compile(block.forward, mode="max-autotune-no-cudagraphs", dynamic=True) +``` + +## 4. Video super-resolution (`Kandinsky6SRPipeline`) + +Like `Kandinsky6TI2VAPipeline`, this pipeline only takes and returns tensors — no file I/O. Use +`diffusers.utils.load_video` / `encode_video` (or your own decoder) to bridge to mp4 files. + +Inputs and constants: + +- `video`: `uint8` pixels in `[0, 255]`, shape `(3, T, H, W)` or a batch `(B, 3, T, H, W)` — the same layout + `Kandinsky6TI2VAPipeline` returns with `output_type="pt"`. `T` must be `1 + 8k` (the SR model's temporal chunk + size; align/truncate yourself — see `notebooks/kandinsky6_sr.py` for a `load_video`-based helper). A batch requires + every clip to already share resolution and frame count. +- `latents`: a pre-encoded SR-VAE latent, shape `(C, T', H', W')` or a batch `(B, C, T', H', W')`. Mutually exclusive + with `video`. Pass `kvae_bridge=True` when the latents come from a *different* (base T2VA) VAE — the pipeline + decodes them through `source_vae` and re-encodes in its own latent space before tiling. +- `resolution_scale` is `2`, `2.25` (default), or `4`. The `2.25` route is a x1.125 pixel pre-upscale followed by the + x2 model. It is the route used after K6 generation. +- Other knobs are `min_overlap` (minimum tile overlap fraction), `tiles_batch_size`, and `output_type` (`"pt"` or + `"np"`). + +### 4.1 Launch 4: VSR flow-matching + +```python +import numpy as np +import torch +from diffusers import Kandinsky6SRPipeline +from diffusers.utils import encode_video, load_video + +pipe = Kandinsky6SRPipeline.from_pretrained( + "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers", dtype=torch.bfloat16 +) +pipe.enable_model_cpu_offload(gpu_id=0) # or: pipe.to("cuda") + +frames, fps = load_video("input.mp4", return_fps=True) +video = torch.stack([torch.as_tensor(np.array(f)) for f in frames]).permute(3, 0, 1, 2) # [3,T,H,W] uint8 +video = video[:, : 1 + 8 * ((video.shape[1] - 1) // 8)] # align to 1+8k frames + +result = pipe( + video=video, + resolution_scale=2.25, # 2, 2.25 or 4 + num_inference_steps=5, # 4 Euler steps per tile (default) + generator=torch.Generator("cuda:0").manual_seed(1137), +) +print(result.frames.shape) # [1, 3, T, H*scale, W*scale] uint8 +encode_video(result.frames[0].permute(1, 2, 3, 0), fps=round(fps), output_path="output_sr.mp4") +``` + +Batch call (equal-size clips, stacked into one tensor): + +```python +result = pipe( + video=torch.stack([video_a, video_b]), # both [3,T,H,W], same shape + resolution_scale=2.25, + output_type="pt", + generator=torch.Generator("cuda:0").manual_seed(1137), +) +``` + +### 4.2 Launch 5: VSR distilled (2 evaluations per tile) + +The same code with the other repo. Do not pass `num_inference_steps`; it has no effect here. + +```python +pipe = Kandinsky6SRPipeline.from_pretrained( + "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", dtype=torch.bfloat16 +) +pipe.enable_model_cpu_offload(gpu_id=0) +result = pipe(video=video, resolution_scale=2, generator=torch.Generator("cuda:0").manual_seed(1137)) +``` + +## 5. Launch 6: T2VA → SR chain + +Both pipelines stay independent. Stage 2 takes the stage-1 video tensor directly, with no intermediate mp4; the +source audio is kept aside and muxed in only at the very end. This version uses the Hub repos instead of the local +bundles in the stacked notebook. + +```python +import torch +from diffusers import Kandinsky6SRPipeline, Kandinsky6TI2VAPipeline +from diffusers.utils import encode_video + +GPU_ID, FPS = 0, 24.0 + +t2va = Kandinsky6TI2VAPipeline.from_pretrained( + "kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers", dtype=torch.bfloat16) +t2va.enable_model_cpu_offload(gpu_id=GPU_ID) + +sr = Kandinsky6SRPipeline.from_pretrained( + "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers", dtype=torch.bfloat16) +sr.enable_model_cpu_offload(gpu_id=GPU_ID) + +gen = t2va( + prompt="A cinematic shot of a small red sailboat crossing a calm blue lake at sunset, " + "with gentle waves and natural ambient sound.", + negative_prompt="blurry, distorted, low quality, noisy audio", + height=480, width=864, + num_frames=121, # not 125: SR keeps at most 121 frames + sample_fps=FPS, num_inference_steps=50, guidance_scale=5.0, + sample_audio=True, output_type="pt", +) + +out = sr(video=gen.frames, resolution_scale=2.25, output_type="pt") # uint8 [B,3,T,H,W] in, out +print(tuple(out.frames.shape)) + +audio = torch.as_tensor(gen.audio[0])[:, None].repeat(1, 2) +encode_video( + out.frames[0].permute(1, 2, 3, 0), + fps=int(FPS), + output_path="t2va_sr.mp4", + audio=audio, + audio_sample_rate=int(t2va.audio_sample_rate), +) +``` + +The output is 2.25x the generation size: `480 x 864` becomes about `1080 x 1944`. For the fast variant swap in the two +distilled repos, `STEPS=16`, `guidance_scale=1.0`, and drop `negative_prompt`. + +## 6. Making the notebooks run + +The three notebooks in `diffusers-new-model-addition-kandinskyX/notebooks/` were copied from a `k6_video` layout. +Three edits are needed: + +1. **Repo detection.** Each notebook walks up until `dev/exports/diffusers` exists. That folder is in `k6_video`, not in + the diffusers checkout, so the first cell raises `RuntimeError`. Replace that block with + `from pathlib import Path; REPO = Path.cwd()`. +2. **Inputs.** `assets/vsr_input.mp4` does not exist anywhere in the workspace. Point `INPUT_VIDEO` at a real mp4. + `assets/i2va_input.png` exists only in `k6_video/assets/`. +3. **Stacked notebook.** It loads `outputs/diffusers_bundle/kandinsky6_diffusers` and `.../kandinsky6_sr_diffusers`, + which are local bundles from `k6_video`. Replace them with the Hub ids from section 1. It also generates + `num_frames=125`; `Kandinsky6SRPipeline` requires `1 + 8k` frames and raises `ValueError` on anything else + (it no longer silently truncates), so use 121. + +Also remove the `HF_TOKEN = "TOKEN"` line in `kandinsky6_ti2va.ipynb`, and register the kernel with +`python -m ipykernel install --user --name k6-diffusers` from the activated venv. + +## 7. Building a local Diffusers bundle from a native checkpoint (only if needed) + +Official Hub repos exist for every launch above, so this is not needed for normal use. To make a bundle from a native +checkpoint, the `k6_video` repo (its own `.venv` with `just setup`) provides: + +```bash +BASE_PATH="something" # folder that contains the K6 repos +cd "$BASE_PATH/k6_video" +just convert-diffusers --config src/kandinsky/configs/k6_pro_125_480_864_mCache_mOffload.yaml --output-dir outputs/kandinsky6_diffusers +just convert-diffusers-sr --config src/kandinsky/configs/k6_pro_125_480_864_mCache_mOffload.yaml --output-dir outputs/kandinsky6_sr_diffusers +``` + +Bundles created this way are loaded with the same `from_pretrained("")` call. The SR recipe above needs +private checkpoints and extra flags (`--checkpoint-path`, `--vae-path`, `--latent-upscaler-config`); see the `JUSTFILE` +comments. `just integrate-diffusers ` overlays the generated K6 code onto a Diffusers checkout. It does not +commit or push. + +## 8. Troubleshooting + +| Symptom | Cause / fix | +|---------|-------------| +| `ImportError: cannot import name 'Kandinsky6TI2VAPipeline'` | stock Diffusers is installed; activate the venv where this branch is installed with `-e` | +| `401` / `RepositoryNotFoundError` on a K6 repo | not logged in, or the account has no access; `hf auth login` | +| `ValueError: PiflowScheduler requires guidance_weight=1.0` | distilled repo with `guidance_scale != 1.0` | +| `ValueError: height and width must be divisible by 8` | choose sizes divisible by 8 (480x864 works) | +| `ValueError: Expected samples with 2 channels` in `encode_video` | audio must be stereo: `torch.as_tensor(a)[:, None].repeat(1, 2)` | +| `No module named 'av'` | `uv pip install av` (needed by `encode_video` for mp4 output, with or without audio) | +| Error setting `_flash_3` | FlashAttention 3 not available; use `native` | +| CUDA OOM | make sure `enable_model_cpu_offload()` is on; lower `height`/`width`/`num_frames`; use the distilled model | +| `ValueError: \`video\` must have 1 + 8k frames` / `... at most 121 frames` | align/split the input yourself: 121 frames at 24 fps is about 5 s | + +## Verification status + +Checked on the Mac (CPU only, no GPU): + +- The branch imports; every keyword used above exists in the current signatures of `Kandinsky6TI2VAPipeline.__call__`, + `Kandinsky6SRPipeline.__call__`, `encode_video`, `load_video`, `MagCacheConfig` and + `transformer.set_attention_backend`/`enable_cache`. `dtype=` and `torch_dtype=` are both accepted by `from_pretrained`. +- `_flash_3` and `native` are valid attention-backend names, and `PiflowScheduler` guidance is enforced as stated. +- Defaults, output layouts (`uint8 [B,3,T,H,W]`), the SR tensor layouts, and the 121-frame clip come from the code. +- The three Hub repos `Pro-sft-5s`, `VSR-5s` and `VSR-distilled2steps-5s` (all `-Diffusers`) were read from the local HF + cache. Their `model_index.json`, scheduler and component configs match what is written here. + +**Not verified:** + +- No end-to-end run happened: no model weights were loaded and nothing was generated. Runtime, peak VRAM and output + quality are untested. +- `Kandinsky-6.0-Pro-distill-5s-Diffusers` is not in the local cache and cannot be checked without access. Its id and + the 16-step / guidance-1.0 settings come from your notebook and the class docstring. +- The MagCache path was checked against `MagCacheConfig`'s signature, not run. diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 3941d091fe83..528a385953e7 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -360,7 +360,7 @@ - local: api/models/transformer_joyimage title: JoyImageEditTransformer3DModel - local: api/models/kandinsky6_transformer - title: Kandinsky6Transformer3DModel + title: Kandinsky 6 Transformers - local: api/models/krea2_transformer2d title: Krea2Transformer2DModel - local: api/models/latte_transformer3d diff --git a/docs/source/en/api/models/kandinsky6_transformer.md b/docs/source/en/api/models/kandinsky6_transformer.md index 798f32695f58..b712daf88cfa 100644 --- a/docs/source/en/api/models/kandinsky6_transformer.md +++ b/docs/source/en/api/models/kandinsky6_transformer.md @@ -12,8 +12,8 @@ specific language governing permissions and limitations under the License. # Kandinsky 6 Transformers -Kandinsky 6 uses Diffusion Transformers (DiTs) for text-to-video and -text-to-video-and-audio generation, as well as video super-resolution. +Kandinsky 6 uses a multimodal diffusion transformer that denoises video and audio latents together for +text/image-to-video-and-audio generation, and a text-free diffusion transformer for video super-resolution. ## Kandinsky6Transformer3DModel @@ -21,6 +21,10 @@ text-to-video-and-audio generation, as well as video super-resolution. - all - forward +## Kandinsky6Transformer3DModelOutput + +[[autodoc]] models.transformers.transformer_kandinsky6.Kandinsky6Transformer3DModelOutput + ## Kandinsky6SRTransformer3DModel [[autodoc]] Kandinsky6SRTransformer3DModel diff --git a/docs/source/en/api/models/kandinsky6_vae.md b/docs/source/en/api/models/kandinsky6_vae.md index 571c407a0585..ba550b3384a7 100644 --- a/docs/source/en/api/models/kandinsky6_vae.md +++ b/docs/source/en/api/models/kandinsky6_vae.md @@ -12,13 +12,19 @@ specific language governing permissions and limitations under the License. # Kandinsky 6 VAEs -Kandinsky 6 uses a causal video VAE for super-resolution and an audio VAE -with a vocoder for synchronized audio generation. +Kandinsky 6 uses a causal 3D K-VAE for video super-resolution and the MMAudio mel-spectrogram VAE, paired with a +BigVGAN vocoder, for synchronized audio generation. ## Kandinsky6SRVAE [[autodoc]] Kandinsky6SRVAE - - all - encode - decode + - all + +## MMAudioVAE +[[autodoc]] MMAudioVAE + - encode + - decode + - all diff --git a/docs/source/en/api/pipelines/kandinsky6.md b/docs/source/en/api/pipelines/kandinsky6.md index be562bfd6183..aff11abc55ef 100644 --- a/docs/source/en/api/pipelines/kandinsky6.md +++ b/docs/source/en/api/pipelines/kandinsky6.md @@ -12,12 +12,70 @@ specific language governing permissions and limitations under the License. # Kandinsky 6 -Kandinsky 6 provides pipelines for text/image-to-video-and-audio generation -and video super-resolution. - -[`Kandinsky6TI2VAPipeline`] denoises video and audio latents together with a single multimodal transformer, conditioned on a Qwen2.5-VL text encoder (and, optionally, a reference image) plus a CLIP text encoder for pooled embeddings. Video and audio come out of the same denoising loop, so there is no separate vocoder or post-hoc audio pass; pass `sample_audio=False` to skip generating audio and drop the `audio_vae` component. - -[`Kandinsky6SRPipeline`] upscales the video latents produced by the base pipeline. Its production route (`resolution_scale=2.25`) does a 1.125x pixel pre-upscale followed by an x2 latent-upscale step, using [`Kandinsky6SRLatentUpscalerBank`] as the latent upscaler; a `source_vae` is only needed for the KVAE latent bridge and is otherwise unused. +Kandinsky 6 is a family of video generation models from [Kandinsky Lab](https://huggingface.co/kandinskylab). The +main model generates video and synchronized audio from text or a reference image with a single multimodal diffusion +transformer: video and audio latents are denoised together through fused blocks that cross-attend between the two +modalities, each conditioned on its own Qwen2.5-VL text branch and a CLIP pooled embedding. A separate +super-resolution model upscales the generated video tile by tile in the latent space of a causal 3D K-VAE. + +The distilled checkpoints use the few-step [`PiflowScheduler`] and must be run with `guidance_scale=1.0`. + +## Available models + +| Model | Pipeline | Notes | +|---|---|---| +| [`kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers) | [`Kandinsky6TI2VAPipeline`] | Flow matching, `guidance_scale=5.0`, 50 steps | +| [`kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers) | [`Kandinsky6TI2VAPipeline`] | Distilled, `guidance_scale=1.0`, 16 steps | +| [`kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers) | [`Kandinsky6SRPipeline`] | Flow matching super-resolution | +| [`kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers) | [`Kandinsky6SRPipeline`] | Distilled super-resolution, 2 steps | + +## Text/image-to-video-and-audio + +```python +import torch +from diffusers import Kandinsky6TI2VAPipeline +from diffusers.utils import encode_video + +pipe = Kandinsky6TI2VAPipeline.from_pretrained( + "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers", torch_dtype=torch.bfloat16 +) +pipe.enable_model_cpu_offload() + +output = pipe( + prompt="A cat and a dog baking a cake together in a kitchen.", + height=480, + width=864, + num_frames=121, + num_inference_steps=16, + guidance_scale=1.0, +) +encode_video( + output.frames[0], + fps=24, + output_path="output.mp4", + audio=output.audio[0][None], + audio_sample_rate=pipe.audio_sample_rate, +) +``` + +Pass `image=` to condition the first frame on a reference image, `sample_audio=False` to generate video only, and +`expand_prompts=True` to let the Qwen2.5-VL text encoder rewrite short prompts into detailed ones first. + +## Video super-resolution + +[`Kandinsky6SRPipeline`] takes the frames produced by [`Kandinsky6TI2VAPipeline`] and upscales them by `2`, `4`, or +`2.25` (a 1.125x bilinear pre-upscale followed by the 2x path). The video is split into overlapping tiles, every tile +is refined at one of the tile sizes the SR transformer was trained on, and the tiles are blended back with Hann +windows. + +```python +sr_pipe = Kandinsky6SRPipeline.from_pretrained( + "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", torch_dtype=torch.bfloat16 +) +sr_pipe.enable_model_cpu_offload() + +upscaled = sr_pipe(video=output.frames[0], resolution_scale=2.25, num_inference_steps=2).frames[0] +``` ## Kandinsky6TI2VAPipeline @@ -33,12 +91,13 @@ and video super-resolution. ## Kandinsky6SRLatentUpscalerBank -[[autodoc]] Kandinsky6SRLatentUpscalerBank - - all +[[autodoc]] pipelines.kandinsky6.modeling_latent_upscaler.Kandinsky6SRLatentUpscalerBank - forward -## Pipeline outputs +## Kandinsky6TI2VAPipelineOutput + +[[autodoc]] pipelines.kandinsky6.pipeline_output.Kandinsky6TI2VAPipelineOutput -[[autodoc]] Kandinsky6TI2VAPipelineOutput +## Kandinsky6SRPipelineOutput -[[autodoc]] Kandinsky6SRPipelineOutput +[[autodoc]] pipelines.kandinsky6.pipeline_output.Kandinsky6SRPipelineOutput diff --git a/docs/source/en/api/schedulers/piflow.md b/docs/source/en/api/schedulers/piflow.md index 61da66c15ae5..03ff161a2eba 100644 --- a/docs/source/en/api/schedulers/piflow.md +++ b/docs/source/en/api/schedulers/piflow.md @@ -12,12 +12,17 @@ specific language governing permissions and limitations under the License. # PiflowScheduler -`PiflowScheduler` is a few-step flow-matching scheduler for distilled -Kandinsky 6 models (both TI2VA and SR). +`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 +[π-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. ## PiflowScheduler [[autodoc]] PiflowScheduler - - all - set_timesteps - step + +## PiflowSchedulerOutput + +[[autodoc]] schedulers.scheduling_piflow.PiflowSchedulerOutput diff --git a/notebooks/kandinsky6_sr.ipynb b/notebooks/kandinsky6_sr.ipynb new file mode 100644 index 000000000000..b5f51cc2fe13 --- /dev/null +++ b/notebooks/kandinsky6_sr.ipynb @@ -0,0 +1,96 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import torch\n", + "from diffusers import Kandinsky6SRPipeline\n", + "from diffusers.utils import encode_video, load_video\n", + "\n", + "REPO = Path.cwd()\n", + "while REPO != REPO.parent and not (REPO / \"dev\" / \"exports\" / \"diffusers\").exists():\n", + " REPO = REPO.parent\n", + "if not (REPO / \"dev\" / \"exports\" / \"diffusers\").exists():\n", + " raise RuntimeError(\"Run this notebook from the K6 repository or one of its subdirectories.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "MODEL_REPO = \"kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers\"\n", + "GPU_ID = 0\n", + "\n", + "pipe = Kandinsky6SRPipeline.from_pretrained(MODEL_REPO, dtype=torch.bfloat16)\n", + "pipe.enable_model_cpu_offload(gpu_id=GPU_ID)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "BS = 2\n", + "INPUT_VIDEO = os.path.join(REPO, \"assets/vsr_input.mp4\")\n", + "OUTPUT_VIDEOS = [\n", + " os.path.join(REPO, f\"outputs/diffusers/k6_vsr/bs{BS}/vsr_x2.25_batchid{i}.mp4\")\n", + " for i in range(BS)\n", + "]\n", + "SEED = 1137\n", + "GENERATOR = torch.Generator(device=f\"cuda:{GPU_ID}\").manual_seed(SEED)\n", + "\n", + "# The pipeline only takes/returns tensors. Decode the video once, align it to\n", + "# 1 + 8k frames (the SR model's temporal chunk size), and repeat it BS times.\n", + "frames, fps = load_video(INPUT_VIDEO, return_fps=True)\n", + "video = torch.from_numpy(np.stack([np.array(frame) for frame in frames])).permute(3, 0, 1, 2)\n", + "aligned = 1 + 8 * ((video.shape[1] - 1) // 8)\n", + "video_batch = video[:, :aligned].unsqueeze(0).repeat(BS, 1, 1, 1, 1)\n", + "\n", + "# 2.25x means x1.125 pixel pre-upscale followed by the x2 KVAE path.\n", + "result = pipe(\n", + " video=video_batch,\n", + " resolution_scale=2.25,\n", + " output_type=\"pt\",\n", + " generator=GENERATOR,\n", + ")\n", + "print(f\"frames: {tuple(result.frames.shape)}\")\n", + "\n", + "for sr_frames, output_path in zip(result.frames, OUTPUT_VIDEOS, strict=True):\n", + " os.makedirs(os.path.dirname(output_path), exist_ok=True)\n", + " encode_video(sr_frames.permute(1, 2, 3, 0), fps=round(fps), output_path=output_path)\n", + " print(f\"saved: {output_path}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "kandinsky-6-video-pro", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/kandinsky6_sr.py b/notebooks/kandinsky6_sr.py new file mode 100644 index 000000000000..0c2e3f20f5a6 --- /dev/null +++ b/notebooks/kandinsky6_sr.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +"""Kandinsky 6 video super-resolution. + +Script version of kandinsky6_sr.ipynb, for running on a headless server. + +The pipeline itself only accepts/returns tensors (no file I/O): this script +does the mp4 <-> tensor conversion with the shared `diffusers.utils` +video helpers, the same way a user integrating the pipeline into their own +code would. + +Usage: + python kandinsky6_sr.py --input-video assets/vsr_input.mp4 + python kandinsky6_sr.py --input-video a.mp4 b.mp4 --gpu-id 1 --resolution-scale 2.25 + +Authenticate with the Hub before running if the model repo is gated: + huggingface-cli login + or: + export HF_TOKEN=hf_... +""" + +from __future__ import annotations + +import argparse +from os import pipe +from pathlib import Path + +import numpy as np +import torch + +from diffusers import Kandinsky6SRPipeline +from diffusers.utils import encode_video, load_video + +DEFAULT_MODEL = "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--gpu-id", type=int, default=0) + parser.add_argument( + "--input-video", + type=Path, + nargs="+", + required=True, + help="One or more input video paths. A single path is replicated --batch-size times.", + ) + parser.add_argument("--batch-size", type=int, default=1, help="Only used when --input-video has one path.") + parser.add_argument("--resolution-scale", type=float, default=2.25) + parser.add_argument("--seed", type=int, default=1137) + parser.add_argument("--output-dir", type=Path, default=Path("outputs/diffusers/k6_vsr")) + return parser.parse_args() + + +def load_video_tensor(path: Path) -> tuple[torch.Tensor, float]: + """Decode ``path`` into a ``[3, T, H, W]`` uint8 tensor, aligned to ``1 + 8k`` frames.""" + frames, fps = load_video(str(path), return_fps=True) + video = torch.from_numpy(np.stack([np.array(frame) for frame in frames])).permute(3, 0, 1, 2) + aligned = 1 + 8 * ((video.shape[1] - 1) // 8) + if aligned == 0: + raise SystemExit(f"{path}: no readable frames") + return video[:, :aligned], fps + + +def main() -> None: + args = parse_args() + + input_videos = args.input_video + for path in input_videos: + if not path.exists(): + raise SystemExit(f"input video not found: {path}") + if len(input_videos) == 1 and args.batch_size > 1: + input_videos = input_videos * args.batch_size + batch_size = len(input_videos) + + loaded = [load_video_tensor(path) for path in input_videos] + videos, fps_values = [item[0] for item in loaded], [item[1] for item in loaded] + if any(video.shape != videos[0].shape for video in videos[1:]): + raise SystemExit("all --input-video clips must share resolution and frame count") + video_batch = torch.stack(videos) + + pipe = Kandinsky6SRPipeline.from_pretrained(args.model, dtype=torch.bfloat16) + pipe.enable_model_cpu_offload(gpu_id=args.gpu_id) + + output_dir = args.output_dir / f"bs{batch_size}" + output_dir.mkdir(parents=True, exist_ok=True) + output_paths = [output_dir / f"vsr_x{args.resolution_scale}_batchid{i}.mp4" for i in range(batch_size)] + + generator = torch.Generator(device=f"cuda:{args.gpu_id}").manual_seed(args.seed) + + # resolution_scale=2.25 means x1.125 pixel pre-upscale followed by the x2 KVAE path + # (the production route). + result = pipe( + video=video_batch, + resolution_scale=args.resolution_scale, + output_type="pt", + generator=generator, + ) + print(f"frames: {tuple(result.frames.shape)}") + for frames, output_path, fps in zip(result.frames, output_paths, fps_values, strict=True): + encode_video(frames.permute(1, 2, 3, 0), fps=round(fps), output_path=str(output_path)) + print(f"saved: {output_path}") + + +if __name__ == "__main__": + main() diff --git a/notebooks/kandinsky6_ti2va.ipynb b/notebooks/kandinsky6_ti2va.ipynb new file mode 100644 index 000000000000..7ca2eca7fc7d --- /dev/null +++ b/notebooks/kandinsky6_ti2va.ipynb @@ -0,0 +1,240 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "7dead615", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from pathlib import Path\n", + "\n", + "os.environ[\"HF_TOKEN\"] = \"TOKEN\"\n", + "\n", + "import torch\n", + "from diffusers import Kandinsky6TI2VAPipeline\n", + "from diffusers.hooks import MagCacheConfig\n", + "from diffusers.models.transformers.transformer_kandinsky6 import (\n", + " Kandinsky6FusedTransformerDecoderBlock,\n", + " Kandinsky6TransformerDecoderBlock,\n", + " Kandinsky6TransformerEncoderBlock,\n", + ")\n", + "from diffusers.utils import encode_video\n", + "\n", + "REPO = Path.cwd()\n", + "while REPO != REPO.parent and not (REPO / \"dev\" / \"exports\" / \"diffusers\").exists():\n", + " REPO = REPO.parent\n", + "if not (REPO / \"dev\" / \"exports\" / \"diffusers\").exists():\n", + " raise RuntimeError(\"Run this notebook from the K6 repository or one of its subdirectories.\")" + ] + }, + { + "cell_type": "markdown", + "id": "b6a61ec1", + "metadata": {}, + "source": [ + "### Initialize pipeline" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "821a098a", + "metadata": {}, + "outputs": [], + "source": [ + "# Keep generation and compile settings in one K6 YAML config. MagCache defaults are exported with the transformer.\n", + "GPU_ID = 2\n", + "\n", + "MODELS = {\n", + " \"k6pro\": \"kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers\",\n", + " \"k6pro_distill\": \"kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers\",\n", + "}\n", + "\n", + "model_name = \"k6pro_distill\"\n", + "model_path = MODELS[model_name]\n", + "expand_prompts = True\n", + "use_magcache = bool(\"distill\" not in model_name)\n", + "use_compile = False\n", + "max_sequence_length = 1024\n", + "sample_fps = 24.0\n", + "seed = 1137\n", + "\n", + "backend = \"native\" if use_compile else \"_flash_3\"\n", + "\n", + "pipe = Kandinsky6TI2VAPipeline.from_pretrained(model_path, dtype=torch.bfloat16)\n", + "pipe.transformer.set_attention_backend(backend)\n", + "\n", + "# Compile after loading the weights and before CPU offload. This applies the\n", + "# configured K6 block-level torch.compile strategy to the exported transformer.\n", + "if use_compile:\n", + " torch.set_float32_matmul_precision('high')\n", + " for block in pipe.transformer.modules():\n", + " if isinstance(block, Kandinsky6FusedTransformerDecoderBlock):\n", + " block.forward = torch.compile(block.forward)\n", + " elif isinstance(block, (Kandinsky6TransformerEncoderBlock, Kandinsky6TransformerDecoderBlock)):\n", + " block.forward = torch.compile(\n", + " block.forward, mode=\"max-autotune-no-cudagraphs\", dynamic=True\n", + " )\n", + "\n", + "is_piflow = pipe.scheduler.__class__.__name__ == \"PiflowScheduler\"\n", + "num_inference_steps = 16 if is_piflow else 50\n", + "guidance_scale = 1.0 if is_piflow else 5.0\n", + "print(f\"Scheduler: {pipe.scheduler.__class__.__name__}; steps: {num_inference_steps}; guidance: {guidance_scale}\")" + ] + }, + { + "cell_type": "markdown", + "id": "4ed6df0c", + "metadata": {}, + "source": [ + "### Run generation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0c31961d", + "metadata": {}, + "outputs": [], + "source": [ + "# NOTE: uses different mag_ratios\n", + "gen_mode = \"t2va\" # \"t2va\"\n", + "\n", + "output_suffix = \"\"\n", + "output_suffix += \"_piflow\" if is_piflow else \"\"\n", + "output_suffix += \"_expand\" if expand_prompts else \"\"\n", + "output_suffix += \"_compile\" if use_compile else \"\"\n", + "output_suffix += \"_magcache\" if use_magcache else \"\"\n", + "output_path = REPO / \"outputs\" / \"diffusers\" / model_name\n", + "output_path.mkdir(parents=True, exist_ok=True)\n", + "\n", + "# Use Diffusers' official MagCacheConfig and model enable_cache API. K6\n", + "# stores conditional/unconditional coefficients interleaved, so the\n", + "# configuration counts both CFG forwards.\n", + "magcache = pipe.transformer.config.get(\"magcache\")\n", + "if use_magcache and magcache is not None:\n", + " mag_ratios = magcache[\"mag_ratios\"]\n", + " if isinstance(mag_ratios, dict):\n", + " mag_ratios = mag_ratios[gen_mode]\n", + " diffusers_magcache = MagCacheConfig(\n", + " threshold=magcache[\"threshold\"],\n", + " max_skip_steps=magcache[\"max_skip_steps\"],\n", + " retention_ratio=magcache[\"retention_ratio\"],\n", + " num_inference_steps=2 * num_inference_steps,\n", + " mag_ratios=list(mag_ratios),\n", + " )\n", + " if pipe.transformer.is_cache_enabled:\n", + " pipe.transformer.disable_cache()\n", + " pipe.transformer.enable_cache(diffusers_magcache)\n", + " print(f\"Diffusers MagCache enabled: {pipe.transformer.is_cache_enabled}\")\n", + " print(f\"MagCache coefficients ({gen_mode}): {len(mag_ratios)}; threshold: {magcache['threshold']}\")\n", + "elif getattr(pipe.transformer, \"is_cache_enabled\", False):\n", + " pipe.transformer.disable_cache()\n", + "\n", + "# Model CPU offload keeps one active component on GPU at a time.\n", + "pipe.enable_model_cpu_offload(gpu_id=GPU_ID)\n", + "\n", + "print(f\"Qwen context length: {max_sequence_length}\")\n", + "print(f\"Audio sample rate: {pipe.audio_sample_rate}; video FPS: {sample_fps}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cfc1e973", + "metadata": {}, + "outputs": [], + "source": [ + "batch_size = 2\n", + "prompt = (\n", + " \"A news anchor delivers the evening news in a clear, standard, and neutral tone, saying: \"\n", + " \"Good evening, ladies and gentlemen, and welcome to the evening news.\"\n", + ")\n", + "negative_prompt = \"blurry, distorted, low quality, noisy audio\"\n", + "\n", + "prompts = [prompt] * batch_size\n", + "negative_prompts = [negative_prompt] * batch_size\n", + "images = None\n", + "\n", + "# NOTE: use example image for I2VA mode\n", + "if gen_mode == \"i2va\":\n", + " image = REPO / \"assets\" / \"i2va_input.png\"\n", + " print(f\"reference image: {image}\")\n", + " images = [str(image)] * batch_size\n", + "\n", + "\n", + "result = pipe(\n", + " prompt=prompts,\n", + " image=images,\n", + " negative_prompt=negative_prompts,\n", + " height=480,\n", + " width=864,\n", + " num_frames=121,\n", + " sample_fps=sample_fps,\n", + " num_inference_steps=num_inference_steps,\n", + " max_sequence_length=max_sequence_length,\n", + " guidance_scale=guidance_scale,\n", + " sample_audio=True,\n", + " expand_prompts=True,\n", + " generator=torch.Generator().manual_seed(seed),\n", + " output_type=\"pt\",\n", + ")\n", + "frames = result.frames\n", + "audios = result.audio\n", + "print(f\"decoded frames: {tuple(frames.shape)}\")\n", + "print(f\"audio streams: {len(audios)}, samples in first stream: {audios[0].shape[0]}\")\n", + "\n", + "if use_magcache and magcache is not None:\n", + " print(\"MagCache run complete\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f0879c75", + "metadata": {}, + "outputs": [], + "source": [ + "final_output_dir = output_path / f\"bs{batch_size}\"\n", + "final_output_dir.mkdir(parents=True, exist_ok=True)\n", + "\n", + "for i, (video, audio) in enumerate(zip(frames, audios)):\n", + " batch_suffix = f\"_batchid{i}\" if batch_size > 1 else \"\"\n", + " video_path = final_output_dir / f\"kandinsky6_{gen_mode}_{num_inference_steps}steps{output_suffix}{batch_suffix}.mp4\"\n", + " video = video.permute(1, 2, 3, 0)\n", + " audio_samples = torch.as_tensor(audio)[:, None].repeat(1, 2)\n", + " encode_video(\n", + " video,\n", + " fps=int(sample_fps),\n", + " output_path=str(video_path),\n", + " audio=audio_samples,\n", + " audio_sample_rate=int(pipe.audio_sample_rate),\n", + " )\n", + " print(f\"saved: {video_path}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "kandinsky-6-video-pro", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/kandinsky6_ti2va.py b/notebooks/kandinsky6_ti2va.py new file mode 100644 index 000000000000..6b423ad85abd --- /dev/null +++ b/notebooks/kandinsky6_ti2va.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +"""Kandinsky 6 text/image-to-video-and-audio generation. + +Script version of kandinsky6_ti2va.ipynb, for running on a headless server. + +Usage: + python kandinsky6_ti2va.py --gpu-id 0 + python kandinsky6_ti2va.py --model k6pro --batch-size 1 --no-magcache + python kandinsky6_ti2va.py --gen-mode i2va --image assets/i2va_input.png + +Authenticate with the Hub before running if the model repo is gated: + huggingface-cli login + or: + export HF_TOKEN=hf_... +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch + +from diffusers import Kandinsky6TI2VAPipeline +from diffusers.hooks import MagCacheConfig +from diffusers.models.transformers.transformer_kandinsky6 import ( + Kandinsky6FusedTransformerDecoderBlock, + Kandinsky6TransformerDecoderBlock, + Kandinsky6TransformerEncoderBlock, +) +from diffusers.utils import encode_video + +MODELS = { + "k6pro": "kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers", + "k6pro_distill": "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers", +} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model", choices=sorted(MODELS), default="k6pro_distill") + parser.add_argument("--gpu-id", type=int, default=0) + parser.add_argument("--batch-size", type=int, default=2) + parser.add_argument("--gen-mode", choices=["t2va", "i2va"], default="t2va") + parser.add_argument("--image", type=Path, default=None, help="Reference image, required for --gen-mode i2va.") + parser.add_argument( + "--prompt", + default=( + "A news anchor delivers the evening news in a clear, standard, and neutral tone, saying: " + "Good evening, ladies and gentlemen, and welcome to the evening news." + ), + ) + parser.add_argument("--negative-prompt", default="blurry, distorted, low quality, noisy audio") + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=864) + parser.add_argument("--num-frames", type=int, default=121) + parser.add_argument("--sample-fps", type=float, default=24.0) + parser.add_argument("--max-sequence-length", type=int, default=1024) + parser.add_argument("--seed", type=int, default=1137) + parser.add_argument("--expand-prompts", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument( + "--magcache", + action=argparse.BooleanOptionalAction, + default=None, + help="Enable MagCache. Defaults to on for non-distilled models, off for distilled ones.", + ) + parser.add_argument("--use-compile", action="store_true") + parser.add_argument("--output-dir", type=Path, default=Path("outputs/diffusers")) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.gen_mode == "i2va" and args.image is None: + raise SystemExit("--gen-mode i2va requires --image ") + + model_path = MODELS[args.model] + use_magcache = bool("distill" not in args.model) if args.magcache is None else args.magcache + backend = "native" if args.use_compile else "_flash_3" + + pipe = Kandinsky6TI2VAPipeline.from_pretrained(model_path, dtype=torch.bfloat16) + pipe.transformer.set_attention_backend(backend) + + # Compile after loading the weights and before CPU offload. + if args.use_compile: + torch.set_float32_matmul_precision("high") + for block in pipe.transformer.modules(): + if isinstance(block, Kandinsky6FusedTransformerDecoderBlock): + block.forward = torch.compile(block.forward) + elif isinstance(block, (Kandinsky6TransformerEncoderBlock, Kandinsky6TransformerDecoderBlock)): + block.forward = torch.compile(block.forward, mode="max-autotune-no-cudagraphs", dynamic=True) + + is_piflow = pipe.scheduler.__class__.__name__ == "PiflowScheduler" + num_inference_steps = 16 if is_piflow else 50 + guidance_scale = 1.0 if is_piflow else 5.0 + print(f"Scheduler: {pipe.scheduler.__class__.__name__}; steps: {num_inference_steps}; guidance: {guidance_scale}") + + output_suffix = "" + output_suffix += "_piflow" if is_piflow else "" + output_suffix += "_expand" if args.expand_prompts else "" + output_suffix += "_compile" if args.use_compile else "" + output_suffix += "_magcache" if use_magcache else "" + output_path = args.output_dir / args.model + output_path.mkdir(parents=True, exist_ok=True) + + # Diffusers' MagCacheConfig, sourced from the metadata the checkpoint carries in its + # own config. K6 stores conditional/unconditional coefficients interleaved, so the + # configuration counts both CFG forwards. + magcache_meta = pipe.transformer.config.get("magcache") + if use_magcache and magcache_meta is not None: + mag_ratios = magcache_meta["mag_ratios"] + if isinstance(mag_ratios, dict): + mag_ratios = mag_ratios[args.gen_mode] + diffusers_magcache = MagCacheConfig( + threshold=magcache_meta["threshold"], + max_skip_steps=magcache_meta["max_skip_steps"], + retention_ratio=magcache_meta["retention_ratio"], + num_inference_steps=2 * num_inference_steps, + mag_ratios=list(mag_ratios), + ) + if pipe.transformer.is_cache_enabled: + pipe.transformer.disable_cache() + pipe.transformer.enable_cache(diffusers_magcache) + print(f"MagCache enabled: {pipe.transformer.is_cache_enabled}") + print(f"MagCache coefficients ({args.gen_mode}): {len(mag_ratios)}; threshold: {magcache_meta['threshold']}") + else: + use_magcache = False + if getattr(pipe.transformer, "is_cache_enabled", False): + pipe.transformer.disable_cache() + + # Model CPU offload keeps one active component on GPU at a time. + pipe.enable_model_cpu_offload(gpu_id=args.gpu_id) + + print(f"Qwen context length: {args.max_sequence_length}") + print(f"Audio sample rate: {pipe.audio_sample_rate}; video FPS: {args.sample_fps}") + + prompts = [args.prompt] * args.batch_size + negative_prompts = [args.negative_prompt] * args.batch_size + images = [str(args.image)] * args.batch_size if args.gen_mode == "i2va" else None + + result = pipe( + prompt=prompts, + image=images, + negative_prompt=negative_prompts, + height=args.height, + width=args.width, + num_frames=args.num_frames, + sample_fps=args.sample_fps, + num_inference_steps=num_inference_steps, + max_sequence_length=args.max_sequence_length, + guidance_scale=guidance_scale, + sample_audio=True, + expand_prompts=args.expand_prompts, + generator=torch.Generator().manual_seed(args.seed), + output_type="pt", + ) + frames = result.frames + audios = result.audio + print(f"decoded frames: {tuple(frames.shape)}") + print(f"audio streams: {len(audios)}, samples in first stream: {audios[0].shape[0]}") + + final_output_dir = output_path / f"bs{args.batch_size}" + final_output_dir.mkdir(parents=True, exist_ok=True) + for i, (video, audio) in enumerate(zip(frames, audios)): + batch_suffix = f"_batchid{i}" if args.batch_size > 1 else "" + video_path = ( + final_output_dir / f"kandinsky6_{args.gen_mode}_{num_inference_steps}steps{output_suffix}{batch_suffix}.mp4" + ) + video = video.permute(1, 2, 3, 0) + audio_samples = torch.as_tensor(audio)[:, None].repeat(1, 2) + encode_video( + video, + fps=int(args.sample_fps), + output_path=str(video_path), + audio=audio_samples, + audio_sample_rate=int(pipe.audio_sample_rate), + ) + print(f"saved: {video_path}") + + +if __name__ == "__main__": + main() diff --git a/notebooks/kandinsky6_ti2va_sr.ipynb b/notebooks/kandinsky6_ti2va_sr.ipynb new file mode 100644 index 000000000000..2f4a42e5daa4 --- /dev/null +++ b/notebooks/kandinsky6_ti2va_sr.ipynb @@ -0,0 +1,130 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Kandinsky 6 T2VA \u2192 SR \u2014 stacked Diffusers example\n", + "\n", + "This notebook follows the two-stage pattern used by LTX-2.5: load the T2VA and SR artifacts independently, generate video and audio with T2VA, then pass the decoded video and generated audio to SR. The pipelines remain separate and can be loaded, offloaded, or replaced independently.\n", + "\n", + "The SR artifact is created with `convert_sr_checkpoint.py` and is a separately packaged Diffusers bundle containing `Kandinsky6SRPipeline`. The production SR route is `resolution_scale=2.25` (1.125x pixel pre-upscale followed by the x2 KVAE path)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "import torch\n", + "from diffusers import Kandinsky6SRPipeline, Kandinsky6TI2VAPipeline\n", + "from diffusers.utils import encode_video\n", + "\n", + "REPO = Path.cwd()\n", + "while REPO != REPO.parent and not (REPO / \"dev\" / \"exports\" / \"diffusers\").exists():\n", + " REPO = REPO.parent\n", + "if not (REPO / \"dev\" / \"exports\" / \"diffusers\").exists():\n", + " raise RuntimeError(\"Run this notebook from the K6 repository or one of its subdirectories.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Stage 1 and stage 2 are independent Diffusers artifacts.\n", + "T2VA_MODEL_PATH = REPO / \"outputs\" / \"diffusers_bundle\" / \"kandinsky6_diffusers\"\n", + "SR_MODEL_PATH = REPO / \"outputs\" / \"diffusers_bundle\" / \"kandinsky6_sr_diffusers\"\n", + "GPU_ID = 0 # Change this to the CUDA device used for inference.\n", + "\n", + "t2va_pipe = Kandinsky6TI2VAPipeline.from_pretrained(\n", + " T2VA_MODEL_PATH,\n", + " torch_dtype=torch.bfloat16,\n", + ")\n", + "t2va_pipe.enable_model_cpu_offload(gpu_id=GPU_ID)\n", + "\n", + "sr_pipe = Kandinsky6SRPipeline.from_pretrained(\n", + " SR_MODEL_PATH,\n", + " torch_dtype=torch.bfloat16,\n", + ")\n", + "sr_pipe.enable_model_cpu_offload(gpu_id=GPU_ID)\n", + "\n", + "print(f\"T2VA artifact: {T2VA_MODEL_PATH}\")\n", + "print(f\"SR artifact: {SR_MODEL_PATH}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prompt = \"A cinematic shot of a small red sailboat crossing a calm blue lake at sunset, with gentle waves and natural ambient sound.\"\n", + "negative_prompt = \"blurry, distorted, low quality, noisy audio\"\n", + "sample_fps = 24.0\n", + "\n", + "t2va_result = t2va_pipe(\n", + " prompt=prompt,\n", + " negative_prompt=negative_prompt,\n", + " height=480,\n", + " width=864,\n", + " num_frames=121, # 1 + 8k: Kandinsky6SRPipeline rejects other values\n", + " sample_fps=sample_fps,\n", + " num_inference_steps=50,\n", + " guidance_scale=5.0,\n", + " sample_audio=True,\n", + " output_type=\"pt\",\n", + ")\n", + "if t2va_result.audio is None:\n", + " raise RuntimeError(\"The T2VA pipeline returned no audio\")\n", + "print(f\"T2VA frames: {tuple(t2va_result.frames.shape)}\")\n", + "print(f\"T2VA audio streams: {len(t2va_result.audio)}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Stage 2 consumes the stage-1 tensors directly; no intermediate video encode/decode is needed.\n", + "output_path = REPO / \"outputs\" / \"diffusers\" / \"kandinsky6_t2va_sr.mp4\"\n", + "output_path.parent.mkdir(parents=True, exist_ok=True)\n", + "\n", + "sr_result = sr_pipe(\n", + " video=t2va_result.frames,\n", + " resolution_scale=2.25,\n", + " output_type=\"pt\",\n", + ")\n", + "print(f\"SR frames: {tuple(sr_result.frames.shape)}\")\n", + "\n", + "audio = torch.as_tensor(t2va_result.audio[0])[:, None].repeat(1, 2)\n", + "encode_video(\n", + " sr_result.frames[0].permute(1, 2, 3, 0),\n", + " fps=int(sample_fps),\n", + " output_path=str(output_path),\n", + " audio=audio,\n", + " audio_sample_rate=int(t2va_pipe.audio_sample_rate),\n", + ")\n", + "print(f\"saved: {output_path}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/kandinsky6_ti2va_sr.py b/notebooks/kandinsky6_ti2va_sr.py new file mode 100644 index 000000000000..2870c3da64c9 --- /dev/null +++ b/notebooks/kandinsky6_ti2va_sr.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +"""Kandinsky 6 T2VA -> SR, stacked. + +Script version of kandinsky6_ti2va_sr.ipynb, for running on a headless server. + +Follows the two-stage pattern used by LTX-2.5: load the T2VA and SR artifacts +independently, generate video and audio with T2VA, then pass the decoded video +and generated audio to SR directly (no intermediate video encode/decode). The +pipelines remain separate and can be loaded, offloaded, or replaced independently. + +The production SR route is --resolution-scale 2.25 (1.125x pixel pre-upscale +followed by the x2 KVAE path). + +Usage: + python kandinsky6_ti2va_sr.py + python kandinsky6_ti2va_sr.py --t2va-model /path/to/local/bundle --sr-model /path/to/local/sr_bundle + +Authenticate with the Hub before running if a model repo is gated: + huggingface-cli login + or: + export HF_TOKEN=hf_... +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch + +from diffusers import Kandinsky6SRPipeline, Kandinsky6TI2VAPipeline +from diffusers.utils import encode_video + +DEFAULT_T2VA_MODEL = "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers" +DEFAULT_SR_MODEL = "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers" + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--t2va-model", default=DEFAULT_T2VA_MODEL, help="Hub repo id or local bundle path.") + parser.add_argument("--sr-model", default=DEFAULT_SR_MODEL, help="Hub repo id or local bundle path.") + parser.add_argument("--gpu-id", type=int, default=0) + parser.add_argument( + "--prompt", + default=( + "A cinematic shot of a small red sailboat crossing a calm blue lake at sunset, " + "with gentle waves and natural ambient sound." + ), + ) + parser.add_argument("--negative-prompt", default="blurry, distorted, low quality, noisy audio") + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=864) + parser.add_argument( + "--num-frames", type=int, default=121, help="Must be 1 + 8k: Kandinsky6SRPipeline rejects other values." + ) + parser.add_argument("--sample-fps", type=float, default=24.0) + parser.add_argument("--num-inference-steps", type=int, default=50) + parser.add_argument("--guidance-scale", type=float, default=5.0) + parser.add_argument("--resolution-scale", type=float, default=2.25) + parser.add_argument("--output", type=Path, default=Path("outputs/diffusers/kandinsky6_t2va_sr.mp4")) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + t2va_pipe = Kandinsky6TI2VAPipeline.from_pretrained(args.t2va_model, torch_dtype=torch.bfloat16) + t2va_pipe.enable_model_cpu_offload(gpu_id=args.gpu_id) + + sr_pipe = Kandinsky6SRPipeline.from_pretrained(args.sr_model, torch_dtype=torch.bfloat16) + sr_pipe.enable_model_cpu_offload(gpu_id=args.gpu_id) + + print(f"T2VA artifact: {args.t2va_model}") + print(f"SR artifact: {args.sr_model}") + + t2va_result = t2va_pipe( + prompt=args.prompt, + negative_prompt=args.negative_prompt, + height=args.height, + width=args.width, + num_frames=args.num_frames, + sample_fps=args.sample_fps, + num_inference_steps=args.num_inference_steps, + guidance_scale=args.guidance_scale, + sample_audio=True, + output_type="pt", + ) + if t2va_result.audio is None: + raise RuntimeError("The T2VA pipeline returned no audio") + print(f"T2VA frames: {tuple(t2va_result.frames.shape)}") + print(f"T2VA audio streams: {len(t2va_result.audio)}") + + sr_result = sr_pipe( + video=t2va_result.frames, + resolution_scale=args.resolution_scale, + output_type="pt", + ) + print(f"SR frames: {tuple(sr_result.frames.shape)}") + + args.output.parent.mkdir(parents=True, exist_ok=True) + audio = torch.as_tensor(t2va_result.audio[0])[:, None].repeat(1, 2) + encode_video( + sr_result.frames[0].permute(1, 2, 3, 0), + fps=int(args.sample_fps), + output_path=str(args.output), + audio=audio, + audio_sample_rate=int(t2va_pipe.audio_sample_rate), + ) + print(f"saved: {args.output}") + + +if __name__ == "__main__": + main() diff --git a/param_counts.json b/param_counts.json new file mode 100644 index 000000000000..694d3a13df5d --- /dev/null +++ b/param_counts.json @@ -0,0 +1,11 @@ +{ + "pro.transformer": 30136285268, + "pro.vae": 246478803, + "pro.audio_vae(MMAudioVAE)": 427632434, + "vsr.transformer": 1413231168, + "vsr.vae": 1738878243, + "vsr.latent_upscaler": 3654021952, + "vsrd.transformer": 1414263936, + "vsrd.vae": 1738878243, + "vsrd.latent_upscaler": 3654021952 +} \ No newline at end of file diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 6a7db85ac9af..1b038bbd2c98 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -302,11 +302,10 @@ "JoyImageEditPlusTransformer3DModel", "JoyImageEditTransformer3DModel", "Kandinsky3UNet", - "MMAudioVAE", + "Kandinsky5Transformer3DModel", "Kandinsky6SRTransformer3DModel", "Kandinsky6SRVAE", "Kandinsky6Transformer3DModel", - "Kandinsky5Transformer3DModel", "Krea2Transformer2DModel", "LatteTransformer3DModel", "LongCatAudioDiTTransformer", @@ -317,6 +316,7 @@ "LTXVideoTransformer3DModel", "Lumina2Transformer2DModel", "LuminaNextDiT2DModel", + "MMAudioVAE", "MiniMaxH3Transformer3DModel", "MiniMaxMusic3ConditionEncoder", "MiniMaxMusic3RVQDepthDecoder", @@ -447,7 +447,6 @@ "EulerAncestralDiscreteScheduler", "EulerDiscreteScheduler", "FlowMapEulerDiscreteScheduler", - "PiflowScheduler", "FlowMatchEulerDiscreteScheduler", "FlowMatchHeunDiscreteScheduler", "FlowMatchLCMScheduler", @@ -462,6 +461,7 @@ "LTXEulerAncestralRFScheduler", "MiniMaxH3Scheduler", "PNDMScheduler", + "PiflowScheduler", "RePaintScheduler", "SASolverScheduler", "SchedulerMixin", @@ -699,7 +699,6 @@ "Kandinsky5I2VPipeline", "Kandinsky5T2IPipeline", "Kandinsky5T2VPipeline", - "Kandinsky6SRLatentUpscalerBank", "Kandinsky6SRPipeline", "Kandinsky6SRPipelineOutput", "Kandinsky6TI2VAPipeline", @@ -1183,11 +1182,10 @@ JoyImageEditPlusTransformer3DModel, JoyImageEditTransformer3DModel, Kandinsky3UNet, - MMAudioVAE, + Kandinsky5Transformer3DModel, Kandinsky6SRTransformer3DModel, Kandinsky6SRVAE, Kandinsky6Transformer3DModel, - Kandinsky5Transformer3DModel, Krea2Transformer2DModel, LatteTransformer3DModel, LongCatAudioDiTTransformer, @@ -1203,6 +1201,7 @@ MiniMaxMusic3RVQDepthDecoder, MiniMaxMusic3Transformer1DModel, MiniMaxMusic3Vocoder, + MMAudioVAE, MochiTransformer3DModel, ModelMixin, MotifVideoTransformer3DModel, @@ -1324,7 +1323,6 @@ EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, FlowMapEulerDiscreteScheduler, - PiflowScheduler, FlowMatchEulerDiscreteScheduler, FlowMatchHeunDiscreteScheduler, FlowMatchLCMScheduler, @@ -1338,6 +1336,7 @@ LCMScheduler, LTXEulerAncestralRFScheduler, MiniMaxH3Scheduler, + PiflowScheduler, PNDMScheduler, RePaintScheduler, SASolverScheduler, @@ -1555,7 +1554,6 @@ Kandinsky5I2VPipeline, Kandinsky5T2IPipeline, Kandinsky5T2VPipeline, - Kandinsky6SRLatentUpscalerBank, Kandinsky6SRPipeline, Kandinsky6SRPipelineOutput, Kandinsky6TI2VAPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 64948c22ddcc..2d56ee1cdeb2 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -36,6 +36,7 @@ _import_structure["autoencoders.autoencoder_kl_cogvideox"] = ["AutoencoderKLCogVideoX"] _import_structure["autoencoders.autoencoder_kl_cosmos"] = ["AutoencoderKLCosmos"] _import_structure["autoencoders.autoencoder_kl_flux2"] = ["AutoencoderKLFlux2"] + _import_structure["autoencoders.autoencoder_kandinsky6_sr"] = ["Kandinsky6SRVAE"] _import_structure["autoencoders.autoencoder_kl_hunyuan_video"] = ["AutoencoderKLHunyuanVideo"] _import_structure["autoencoders.autoencoder_kl_hunyuanimage"] = ["AutoencoderKLHunyuanImage"] _import_structure["autoencoders.autoencoder_kl_hunyuanimage_refiner"] = ["AutoencoderKLHunyuanImageRefiner"] @@ -52,6 +53,7 @@ _import_structure["autoencoders.autoencoder_kl_qwenimage"] = ["AutoencoderKLQwenImage"] _import_structure["autoencoders.autoencoder_kl_temporal_decoder"] = ["AutoencoderKLTemporalDecoder"] _import_structure["autoencoders.autoencoder_kl_wan"] = ["AutoencoderKLWan"] + _import_structure["autoencoders.autoencoder_mmaudio"] = ["MMAudioVAE"] _import_structure["autoencoders.autoencoder_longcat_audio_dit"] = ["LongCatAudioDiTVae"] _import_structure["autoencoders.autoencoder_oobleck"] = ["AutoencoderOobleck"] _import_structure["autoencoders.autoencoder_rae"] = ["AutoencoderRAE"] @@ -128,11 +130,9 @@ _import_structure["transformers.transformer_ideogram4"] = ["Ideogram4Transformer2DModel"] _import_structure["transformers.transformer_joyimage"] = ["JoyImageEditTransformer3DModel"] _import_structure["transformers.transformer_joyimage_edit_plus"] = ["JoyImageEditPlusTransformer3DModel"] - _import_structure["autoencoders.autoencoder_mmaudio"] = ["MMAudioVAE"] - _import_structure["autoencoders.autoencoder_kandinsky6_sr"] = ["Kandinsky6SRVAE"] - _import_structure["transformers.transformer_kandinsky6_sr"] = ["Kandinsky6SRTransformer3DModel"] - _import_structure["transformers.transformer_kandinsky6"] = ["Kandinsky6Transformer3DModel"] _import_structure["transformers.transformer_kandinsky"] = ["Kandinsky5Transformer3DModel"] + _import_structure["transformers.transformer_kandinsky6"] = ["Kandinsky6Transformer3DModel"] + _import_structure["transformers.transformer_kandinsky6_sr"] = ["Kandinsky6SRTransformer3DModel"] _import_structure["transformers.transformer_krea2"] = ["Krea2Transformer2DModel"] _import_structure["transformers.transformer_longcat_audio_dit"] = ["LongCatAudioDiTTransformer"] _import_structure["transformers.transformer_longcat_image"] = ["LongCatImageTransformer2DModel"] @@ -187,8 +187,6 @@ AutoencoderKLFlux2, AutoencoderKLHunyuanImage, AutoencoderKLHunyuanImageRefiner, - MMAudioVAE, - Kandinsky6SRVAE, AutoencoderKLHunyuanVideo, AutoencoderKLHunyuanVideo15, AutoencoderKLKVAE, @@ -210,9 +208,11 @@ AutoencoderVidTok, ConsistencyDecoderVAE, Cosmos3AVAEAudioTokenizer, + Kandinsky6SRVAE, LongCatAudioDiTVae, LTX2VideoDiffusionDecoderModel, MiniMaxMusic3Vocoder, + MMAudioVAE, VQModel, ) from .cache_utils import CacheMixin @@ -273,9 +273,9 @@ Ideogram4Transformer2DModel, JoyImageEditPlusTransformer3DModel, JoyImageEditTransformer3DModel, + Kandinsky5Transformer3DModel, Kandinsky6SRTransformer3DModel, Kandinsky6Transformer3DModel, - Kandinsky5Transformer3DModel, Krea2Transformer2DModel, LatteTransformer3DModel, LongCatAudioDiTTransformer, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index 08f422cbd4fd..22c2416d71b2 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -1,13 +1,12 @@ from .autoencoder_asym_kl import AsymmetricAutoencoderKL from .autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer from .autoencoder_dc import AutoencoderDC +from .autoencoder_kandinsky6_sr import Kandinsky6SRVAE from .autoencoder_kl import AutoencoderKL from .autoencoder_kl_allegro import AutoencoderKLAllegro from .autoencoder_kl_cogvideox import AutoencoderKLCogVideoX from .autoencoder_kl_cosmos import AutoencoderKLCosmos from .autoencoder_kl_flux2 import AutoencoderKLFlux2 -from .autoencoder_mmaudio import MMAudioVAE -from .autoencoder_kandinsky6_sr import Kandinsky6SRVAE from .autoencoder_kl_hunyuan_video import AutoencoderKLHunyuanVideo from .autoencoder_kl_hunyuanimage import AutoencoderKLHunyuanImage from .autoencoder_kl_hunyuanimage_refiner import AutoencoderKLHunyuanImageRefiner @@ -25,6 +24,7 @@ from .autoencoder_kl_temporal_decoder import AutoencoderKLTemporalDecoder from .autoencoder_kl_wan import AutoencoderKLWan from .autoencoder_longcat_audio_dit import LongCatAudioDiTVae +from .autoencoder_mmaudio import MMAudioVAE from .autoencoder_oobleck import AutoencoderOobleck from .autoencoder_rae import AutoencoderRAE from .autoencoder_same import AutoencoderSAME diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index 587518ac487b..c2165cea25aa 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -12,1092 +12,586 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Kandinsky 6 SR KVAE Diffusers component.""" +"""Causal 3D K-VAE used by the Kandinsky 6 video super-resolution pipeline.""" from __future__ import annotations -import functools import math -from dataclasses import dataclass -import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config from ...utils.accelerate_utils import apply_forward_hook +from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin -from .vae import DiagonalGaussianDistribution +from .vae import DecoderOutput, DiagonalGaussianDistribution -class SafeConv3d(nn.Conv3d): - """Conv3d that splits its input along the time axis when it would otherwise materialize more than ~2B - elements, running the convolution chunk by chunk instead. Each chunk is padded with the last - ``kernel_size - 1`` frames of the previous chunk before convolving, so the result is identical to running - the convolution on the whole tensor at once; this only bounds peak memory. +# Number of pixel frames encoded or decoded per causal segment. Segmentation only bounds peak memory: the causal +# convolutions carry their padding state from one segment to the next, so the result matches a single-pass call. +SEGMENT_FRAMES = 16 +# Element count above which a convolution input is processed in temporal chunks. +CONV_CHUNK_ELEMENTS = 2 * 10**9 + + +class Kandinsky6SRConv3d(nn.Conv3d): + """`Conv3d` that splits very large inputs along the time axis and convolves them chunk by chunk. + + Each chunk is padded with the last `kernel_size - 1` frames of the previous chunk before convolving, so the result + is identical to running the convolution on the whole tensor at once; this only bounds peak memory. """ - def forward(self, x, write_to=None, transform=None): - if transform is None: - - def transform(x): - return x - - memory_count = x.numel() / (10**9) - if memory_count > 2: - kernel_size = self.kernel_size[0] - part_num = math.ceil(memory_count / 2) - input_chunks = torch.chunk(x, part_num, dim=2) # NCTHW - - if any(ch.size(2) < kernel_size for ch in input_chunks) and kernel_size > 1: - if input_chunks[0].numel() * (kernel_size / input_chunks[0].size(2)) >= (2 * 10**9): - raise ValueError("frames are too big for Conv3d") - - t_stride, output = self.stride[0], [] - for i in range(0, x.size(2) - kernel_size + 1, t_stride): - chunk = transform(x[:, :, i : i + kernel_size]) - output.append(super(SafeConv3d, self).forward(chunk)) - output = torch.cat(output, dim=2) - return output - - if write_to is None: - output = [] - for i, chunk in enumerate(input_chunks): - if i == 0 or kernel_size == 1: - z = torch.clone(chunk) - else: - z = torch.cat([z[:, :, -kernel_size + 1 :], chunk], dim=2) - output.append(super(SafeConv3d, self).forward(transform(z))) - output = torch.cat(output, dim=2) - return output + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.numel() <= CONV_CHUNK_ELEMENTS: + return super().forward(x) + + kernel_size = self.kernel_size[0] + chunks = torch.chunk(x, math.ceil(x.numel() / CONV_CHUNK_ELEMENTS), dim=2) + if kernel_size > 1 and any(chunk.size(2) < kernel_size for chunk in chunks): + # Chunks shorter than the kernel: slide one window at a time instead. + if chunks[0].numel() * (kernel_size / chunks[0].size(2)) >= CONV_CHUNK_ELEMENTS: + raise ValueError("frames are too big for Conv3d") + stride = self.stride[0] + windows = range(0, x.size(2) - kernel_size + 1, stride) + return torch.cat([super().forward(x[:, :, i : i + kernel_size]) for i in windows], dim=2) + + outputs = [] + for i, chunk in enumerate(chunks): + if i == 0 or kernel_size == 1: + carried = chunk else: - time_offset = 0 - for i, chunk in enumerate(input_chunks): - if i == 0 or kernel_size == 1: - z = torch.clone(chunk) - else: - z = torch.cat([z[:, :, -kernel_size + 1 :], chunk], dim=2) - z_time = z.size(2) - (kernel_size - 1) - write_to[:, :, time_offset : time_offset + z_time] = super(SafeConv3d, self).forward(transform(z)) - time_offset += z_time - return write_to - else: - if write_to is None: - return super(SafeConv3d, self).forward(transform(x)) - else: - write_to[...] = super(SafeConv3d, self).forward(transform(x)) - return write_to + carried = torch.cat([carried[:, :, -kernel_size + 1 :], chunk], dim=2) + outputs.append(super().forward(carried)) + return torch.cat(outputs, dim=2) + +class Kandinsky6SRCausalConv3d(nn.Module): + """Causal 3D convolution whose temporal padding is carried across segments through a `cache` dict. + + Height and width are zero-padded symmetrically. Along time the first segment is padded by repeating its first frame + `kernel_size - 1` times; later segments are padded with the frames the previous segment left behind in + `cache["padding"]`, so a video processed segment by segment matches a single pass. + """ -class CausalConv3d(nn.Module): def __init__( self, - chan_in, - chan_out, + in_channels: int, + out_channels: int, kernel_size: int | tuple[int, int, int], - stride=(1, 1, 1), - dilation=(1, 1, 1), - padding_mode=None, - **kwargs, - ): + stride: tuple[int, int, int] = (1, 1, 1), + ) -> None: super().__init__() if not isinstance(kernel_size, tuple): kernel_size = (kernel_size,) * 3 - time_kernel_size, height_kernel_size, width_kernel_size = kernel_size - if not (height_kernel_size % 2 and width_kernel_size % 2): raise ValueError( f"height_kernel_size and width_kernel_size must be odd, got {height_kernel_size} and " f"{width_kernel_size}" ) - self.height_pad = height_kernel_size // 2 self.width_pad = width_kernel_size // 2 self.time_pad = time_kernel_size - 1 self.time_kernel_size = time_kernel_size - self.temporal_dim = 2 - - self.stride = stride - self.conv = SafeConv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs) - self.cache_padding = None - self.padding_mode = padding_mode - - def forward(self, input_): - input_parallel = input_ - - padding_3d = (self.width_pad, self.width_pad, self.height_pad, self.height_pad, self.time_pad, 0) - input_parallel = F.pad(input_parallel, padding_3d, mode=self.padding_mode or "replicate") - - output = self.conv(input_parallel) - return output - - -class Kandinsky6VAERMSNorm(nn.Module): - r""" - RMS normalization over the channel dimension of a video tensor (`N, C, T, H, W`). - - `forward` accepts and ignores extra positional/keyword arguments (e.g. the causal-cache dict the - encoder/decoder pass to every normalization layer) so this stays a drop-in alternative to - `CachedGroupNorm`; construction likewise ignores `zq_ch`/`add_conv`, which only apply to the - `CachedSpatialNorm3D` alternative. - - Args: - in_channels (int): The number of channels to normalize over. - bias (bool, optional): Whether to include a learnable bias term. Default is False. - """ - - def __init__(self, in_channels: int, bias: bool = False, **kwargs) -> None: - super().__init__() - shape = (in_channels, 1, 1, 1) - - self.scale = in_channels**0.5 - self.gamma = nn.Parameter(torch.ones(shape)) - self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 - - def forward(self, x, *args, **kwargs): - needs_fp32_normalize = x.dtype in (torch.float16, torch.bfloat16) or any( - t in str(x.dtype) for t in ("float4_", "float8_") - ) - normalized = F.normalize(x.float() if needs_fp32_normalize else x, dim=1).to(x.dtype) - return normalized * self.scale * self.gamma + self.bias + self.time_stride = stride[0] + self.conv = Kandinsky6SRConv3d(in_channels, out_channels, kernel_size, stride=stride) + @staticmethod + def make_cache() -> dict: + return {"padding": None} -class CachedGroupNorm(nn.GroupNorm): - """GroupNorm alternative to `Kandinsky6VAERMSNorm`: same call convention (ignores `zq_ch`/`add_conv`, - accepts a `cache` dict) so the encoder/decoder can pick either one through the same `normalization(...)` - call. `cache` only ever records whether this layer has run before — `CachedSpatialNorm3D` reads that - presence flag to decide whether it's normalizing the first causal segment or a later one; the statistics - themselves are not reused across calls. - """ - - def __init__(self, in_channels, **kwargs): - super().__init__(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) - - def forward(self, x, cache: dict): - out = super().forward(x) - if cache.get("mean") is None and cache.get("var") is None: - cache["mean"] = 1 - cache["var"] = 1 - return out - - -class CachedCausalConv3d(CausalConv3d): - def forward(self, input_, cache: dict, fixed_stride=False): - t_stride = self.stride[0] - padding_3d = (self.height_pad, self.height_pad, self.width_pad, self.width_pad, 0, 0) - input_parallel = F.pad( - input_, - padding_3d, - mode="constant" if self.padding_mode == "zeros" else (self.padding_mode or "replicate"), - value=0 if self.padding_mode in ["constant", "zeros"] else None, - ) + def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: + batch_size, _, num_frames, height, width = hidden_states.shape + hidden_states = F.pad(hidden_states, (self.width_pad, self.width_pad, self.height_pad, self.height_pad)) if cache["padding"] is None: - first_frame = input_parallel[:, :, :1] - time_pad_shape = list(first_frame.shape) - time_pad_shape[2] = self.time_pad - padding = first_frame.expand(time_pad_shape) + first_frame = hidden_states[:, :, :1] + padding = first_frame.expand(-1, -1, self.time_pad, -1, -1) else: padding = cache["padding"] - out_size = list(input_.shape) - out_size[1] = self.conv.out_channels - if t_stride == 2: - out_size[2] = (input_.size(2) + 1) // 2 - output = torch.empty(tuple(out_size), dtype=input_.dtype, device=input_.device) - - offset_out = math.ceil( - padding.size(2) / t_stride - ) # forward on `padding_poisoned` should take exactly this range - offset_in = offset_out * t_stride - padding.size( - 2 - ) # to make forward on `input_parallel` take slice starting with this index + stride = self.time_stride + output_frames = (num_frames + 1) // 2 if stride == 2 else num_frames + output = torch.empty( + (batch_size, self.conv.out_channels, output_frames, height, width), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + # The frames that overlap the carried padding are convolved together with it; the rest run on their own. + offset_out = math.ceil(padding.size(2) / stride) + offset_in = offset_out * stride - padding.size(2) if offset_out > 0: - padding_poisoned = torch.cat( - [padding, input_parallel[:, :, : offset_in + self.time_kernel_size - t_stride]], dim=2 - ) - output[:, :, :offset_out] = self.conv(padding_poisoned) - - if offset_out < output.size(2): - output[:, :, offset_out:] = self.conv(input_parallel[:, :, offset_in:]) + padded = torch.cat([padding, hidden_states[:, :, : offset_in + self.time_kernel_size - stride]], dim=2) + output[:, :, :offset_out] = self.conv(padded) + if offset_out < output_frames: + output[:, :, offset_out:] = self.conv(hidden_states[:, :, offset_in:]) - if t_stride == 2 and not fixed_stride and cache["padding"] is not None: - expected_pad_size = padding.size(2) - 1 - offset_out = math.ceil(expected_pad_size / t_stride) - offset_in = offset_out * t_stride - expected_pad_size - - # exact formula, doesn't depend on size of segments + # Carry the frames the next segment's first window still needs. pad_offset = ( - offset_in - + t_stride * math.trunc((input_parallel.size(2) - offset_in - self.time_kernel_size) / t_stride) - + t_stride + offset_in + stride * math.trunc((num_frames - offset_in - self.time_kernel_size) / stride) + stride ) - - # this condition will be executed ONLY for old models - if t_stride == 2 and not fixed_stride: - pad_offset -= 1 - - if pad_offset < 0: # specific to small chunks (for inference on high resolution videos) - cache["padding"] = torch.cat([padding[:, :, pad_offset:], input_parallel], dim=2) + if pad_offset < 0: + cache["padding"] = torch.cat([padding[:, :, pad_offset:], hidden_states], dim=2) else: - cache["padding"] = torch.clone(input_parallel[:, :, pad_offset:]) - + cache["padding"] = hidden_states[:, :, pad_offset:].clone() return output -class CachedCausalResnetBlock3D(nn.Module): - def __init__( - self, - *, - in_channels, - out_channels=None, - conv_shortcut=False, - dropout, - temb_channels=512, - zq_ch=None, - add_conv=False, - normalization=CachedGroupNorm, - padding_mode=None, - ): - super().__init__() - self.in_channels = in_channels - out_channels = in_channels if out_channels is None else out_channels - self.out_channels = out_channels - self.use_conv_shortcut = conv_shortcut - - self.norm1 = normalization(in_channels, zq_ch=zq_ch, add_conv=add_conv) - - self.conv1 = CachedCausalConv3d( - chan_in=in_channels, chan_out=out_channels, kernel_size=3, padding_mode=padding_mode - ) - self.norm2 = normalization(out_channels, zq_ch=zq_ch, add_conv=add_conv) - self.conv2 = CachedCausalConv3d( - chan_in=out_channels, chan_out=out_channels, kernel_size=3, padding_mode=padding_mode - ) - if self.in_channels != self.out_channels: - if self.use_conv_shortcut: - self.conv_shortcut = CachedCausalConv3d( - chan_in=in_channels, chan_out=out_channels, kernel_size=3, padding_mode=padding_mode - ) - else: - self.nin_shortcut = SafeConv3d( - in_channels, - out_channels, - kernel_size=1, - stride=1, - padding=0, - ) - - def forward(self, x, temb, layer_cache, zq=None): - h = x - - if zq is None: - h = self.norm1(h, cache=layer_cache["norm1"]) - else: - h = self.norm1(h, zq, cache=layer_cache["norm1"]) +class Kandinsky6SRRMSNorm(nn.Module): + """RMS normalization over the channel axis of a `(B, C, T, H, W)` tensor, computed in float32.""" - h = F.silu(h, inplace=True) - h = self.conv1(h, cache=layer_cache["conv1"]) - - if zq is None: - h = self.norm2(h, cache=layer_cache["norm2"]) - else: - h = self.norm2(h, zq, cache=layer_cache["norm2"]) - - h = F.silu(h, inplace=True) - h = self.conv2(h, cache=layer_cache["conv2"]) - - if self.in_channels != self.out_channels: - if self.use_conv_shortcut: - x = self.conv_shortcut(x, cache=layer_cache["conv_shortcut"]) - else: - x = self.nin_shortcut(x) - - return x + h - - -class CachedPXSDownsample(nn.Module): - def __init__( - self, in_channels: int, compress_time: bool, factor: int = 2, version=1, fixed_stride=False, padding_mode=None - ): + def __init__(self, num_channels: int) -> None: super().__init__() - self.temporal_compress = compress_time - self.fixed_stride = fixed_stride - self.factor = factor - self.unshuffle = nn.PixelUnshuffle(self.factor) - self.s_pool = nn.AvgPool3d((1, 2, 2), (1, 2, 2)) - - self.version = version - out_channels = in_channels * 2 if version > 1 else in_channels - - self.spatial_conv = SafeConv3d( - in_channels, - out_channels, - kernel_size=(1, 3, 3), - stride=(1, 2, 2), - padding=(0, 1, 1), - padding_mode=padding_mode or "reflect", - ) + self.scale = num_channels**0.5 + self.gamma = nn.Parameter(torch.ones(num_channels, 1, 1, 1)) - if self.temporal_compress: - if version == 2: - self.temporal_conv = nn.Sequential( - CachedCausalConv3d( - out_channels, - out_channels, - kernel_size=(2, 1, 1), - stride=(1, 1, 1), - dilation=(1, 1, 1), - padding_mode=padding_mode, - ), - CachedCausalConv3d( - out_channels, - out_channels, - kernel_size=(2, 1, 1), - stride=(2, 1, 1), - dilation=(1, 1, 1), - padding_mode=padding_mode, - ), - ) - else: # 1 or 3 - self.temporal_conv = CachedCausalConv3d( - out_channels, - out_channels, - kernel_size=(3, 1, 1), - stride=(2, 1, 1), - dilation=(1, 1, 1), - padding_mode=padding_mode, - ) - - self.linear = SafeConv3d(out_channels, out_channels, kernel_size=1, stride=1) - - def spatial_downsample(self, input_): - # PixelShuffle part - pxs_input = input_.permute(0, 2, 1, 3, 4).reshape(-1, input_.shape[1], input_.shape[3], input_.shape[4]) - pxs_interm = self.unshuffle(pxs_input) - b, c, h, w = pxs_interm.shape - if self.version > 1: - pxs_interm_view = pxs_interm.view(b, c // self.factor, self.factor, h, w) - else: # - pxs_interm_view = pxs_interm.view(b, c // self.factor**2, self.factor**2, h, w) - pxs_out = torch.mean(pxs_interm_view, dim=2) - pxs_out = pxs_out.reshape( - input_.shape[0], input_.size(2), pxs_out.shape[1], pxs_out.shape[2], pxs_out.shape[3] - ).permute(0, 2, 1, 3, 4) - - # Downsampling by 3D-convolution - conv_out = self.spatial_conv(input_) - - # adding it all together - return conv_out + pxs_out - - def temporal_downsample(self, input_, cache): - # Interpolation part - permuted = input_.permute(0, 3, 4, 1, 2).reshape(-1, input_.shape[1], input_.shape[2]) - if cache[0]["padding"] is None: - first, rest = permuted[..., :1], permuted[..., 1:] - - if rest.size(-1) > 0: - rest_interp = F.avg_pool1d(rest, kernel_size=2, stride=2) - full_interp = torch.cat([first, rest_interp], dim=-1) - else: - full_interp = first - else: - rest = permuted - if rest.size(-1) > 0: - full_interp = F.avg_pool1d(rest, kernel_size=2, stride=2) - - full_interp = full_interp.reshape( - input_.shape[0], input_.size(-2), input_.size(-1), full_interp.shape[1], full_interp.shape[2] - ).permute(0, 3, 4, 1, 2) + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + normalized = F.normalize(hidden_states.float(), dim=1).to(hidden_states.dtype) + return normalized * self.scale * self.gamma - # Downsampling by convolution - if self.version == 1: - conv_out = self.temporal_conv(input_, cache[0], fixed_stride=self.fixed_stride) - elif self.version == 2: - conv_out = self.temporal_conv[0](input_, cache[0], fixed_stride=self.fixed_stride) - conv_out = self.temporal_conv[1](conv_out, cache[1], fixed_stride=self.fixed_stride) - return conv_out + full_interp +class Kandinsky6SRSpatialNorm3D(nn.Module): + """RMS norm modulated by the latent `zq`, the same conditioning other video VAEs in the library use to inject the + latent back into the decoder, adapted to causal segment-by-segment decoding. - def forward(self, x, cache): - # SPATIAL DOWNSAMPLE - out = self.spatial_downsample(x) - - if self.temporal_compress: - # TEMPORAL DOWNSAMPLE - out = self.temporal_downsample(out, cache=cache) - - return self.linear(out) - - -class CachedSpatialNorm3D(nn.Module): - """Normalization used by the decoder's ResnetBlocks: modulates a GroupNorm/RMSNorm output with a - per-channel scale and shift derived from the latent `zq`, the same "spatial norm" conditioning other VAEs - in the library (e.g. CogVideoX) use to inject the encoder's output back into the decoder — adapted here - for causal, segment-by-segment video decoding. + `zq` is nearest-upsampled to the feature grid. In the first segment the first frame is upsampled separately, + because the temporal upsampler turns `T + 1` latent frames into `2T + 1` pixel frames. """ - def __init__( - self, - f_channels, - zq_channels, - freeze_norm_layer=False, - add_conv=False, - padding_mode=None, - normalization=CachedGroupNorm, - **norm_layer_params, - ): + def __init__(self, num_channels: int, zq_channels: int) -> None: super().__init__() - self.norm_layer = normalization(in_channels=f_channels, **norm_layer_params) - - self.add_conv = add_conv - if add_conv: - self.conv = CachedCausalConv3d( - chan_in=zq_channels, chan_out=zq_channels, kernel_size=3, padding_mode=padding_mode - ) - - self.conv_y = SafeConv3d( - zq_channels, - f_channels, - kernel_size=1, - ) - self.conv_b = SafeConv3d( - zq_channels, - f_channels, - kernel_size=1, - ) + self.norm_layer = Kandinsky6SRRMSNorm(num_channels) + self.conv_y = Kandinsky6SRConv3d(zq_channels, num_channels, kernel_size=1) + self.conv_b = Kandinsky6SRConv3d(zq_channels, num_channels, kernel_size=1) - def forward(self, f, zq, cache): - if cache["norm"]["mean"] is None and cache["norm"]["var"] is None: - f_first, f_rest = f[:, :, :1], f[:, :, 1:] - f_first_size, f_rest_size = f_first.shape[-3:], f_rest.shape[-3:] - zq_first, zq_rest = zq[:, :, :1], zq[:, :, 1:] - - zq_first = F.interpolate(zq_first, size=f_first_size, mode="nearest") + @staticmethod + def make_cache() -> dict: + return {"is_first_segment": True} + def forward(self, hidden_states: torch.Tensor, zq: torch.Tensor, cache: dict) -> torch.Tensor: + if cache["is_first_segment"]: + zq_first = F.interpolate(zq[:, :, :1], size=hidden_states[:, :, :1].shape[-3:], mode="nearest") if zq.size(2) > 1: - zq_rest_splits = torch.split(zq_rest, 32, dim=1) - interpolated_splits = [ - F.interpolate(split, size=f_rest_size, mode="nearest") for split in zq_rest_splits - ] - - zq_rest = torch.cat(interpolated_splits, dim=1) + # Interpolate in channel chunks to bound the memory of the upsampled conditioning tensor. + zq_rest = torch.cat( + [ + F.interpolate(split, size=hidden_states[:, :, 1:].shape[-3:], mode="nearest") + for split in torch.split(zq[:, :, 1:], 32, dim=1) + ], + dim=1, + ) zq = torch.cat([zq_first, zq_rest], dim=2) else: zq = zq_first + cache["is_first_segment"] = False else: - f_size = f.shape[-3:] - zq_splits = torch.split(zq, 32, dim=1) - interpolated_splits = [F.interpolate(split, size=f_size, mode="nearest") for split in zq_splits] - zq = torch.cat(interpolated_splits, dim=1) - - if self.add_conv: - zq = self.conv(zq, cache["add_conv"]) + zq = torch.cat( + [ + F.interpolate(split, size=hidden_states.shape[-3:], mode="nearest") + for split in torch.split(zq, 32, dim=1) + ], + dim=1, + ) + return self.norm_layer(hidden_states) * self.conv_y(zq) + self.conv_b(zq) - norm_f = self.norm_layer(f, cache["norm"]) - norm_f.mul_(self.conv_y(zq)) - norm_f.add_(self.conv_b(zq)) - if cache["norm"]["mean"] is None and cache["norm"]["var"] is None: - cache["norm"]["mean"] = 1 - cache["norm"]["var"] = 1 +class Kandinsky6SRResnetBlock3D(nn.Module): + """Causal residual block; decoder blocks modulate their norms with the latent `zq`.""" - return norm_f + def __init__(self, in_channels: int, out_channels: int, zq_channels: int | None = None) -> None: + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + if zq_channels is None: + self.norm1 = Kandinsky6SRRMSNorm(in_channels) + self.norm2 = Kandinsky6SRRMSNorm(out_channels) + else: + self.norm1 = Kandinsky6SRSpatialNorm3D(in_channels, zq_channels) + self.norm2 = Kandinsky6SRSpatialNorm3D(out_channels, zq_channels) + self.conv1 = Kandinsky6SRCausalConv3d(in_channels, out_channels, kernel_size=3) + self.conv2 = Kandinsky6SRCausalConv3d(out_channels, out_channels, kernel_size=3) + if in_channels != out_channels: + self.nin_shortcut = Kandinsky6SRConv3d(in_channels, out_channels, kernel_size=1) + + def make_cache(self) -> dict: + cache = {"conv1": self.conv1.make_cache(), "conv2": self.conv2.make_cache()} + if isinstance(self.norm1, Kandinsky6SRSpatialNorm3D): + cache["norm1"] = self.norm1.make_cache() + cache["norm2"] = self.norm2.make_cache() + return cache + def forward(self, hidden_states: torch.Tensor, cache: dict, zq: torch.Tensor | None = None) -> torch.Tensor: + residual = hidden_states + hidden_states = self.norm1(hidden_states) if zq is None else self.norm1(hidden_states, zq, cache["norm1"]) + hidden_states = self.conv1(F.silu(hidden_states), cache["conv1"]) + hidden_states = self.norm2(hidden_states) if zq is None else self.norm2(hidden_states, zq, cache["norm2"]) + hidden_states = self.conv2(F.silu(hidden_states), cache["conv2"]) + if self.in_channels != self.out_channels: + residual = self.nin_shortcut(residual) + return residual + hidden_states -def Normalize3D(in_channels, zq_ch, add_conv, normalization=CachedGroupNorm): - return CachedSpatialNorm3D( - in_channels, - zq_ch, - freeze_norm_layer=False, - add_conv=add_conv, - num_groups=32, - eps=1e-6, - affine=True, - normalization=normalization, - ) +class Kandinsky6SRDownsample(nn.Module): + """Spatial 2x downsample (strided conv plus pixel-unshuffle average) with an optional causal temporal 2x.""" -class CachedPXSUpsample(nn.Module): - def __init__(self, in_channels: int, compress_time: bool, factor: int = 2, padding_mode=None): + def __init__(self, in_channels: int, compress_time: bool) -> None: super().__init__() - self.temporal_compress = compress_time - self.factor = factor - self.shuffle = nn.PixelShuffle(self.factor) - self.spatial_conv = SafeConv3d( - in_channels, - in_channels, - kernel_size=(1, 3, 3), - stride=(1, 1, 1), - padding=(0, 1, 1), - padding_mode=padding_mode or "reflect", + out_channels = 2 * in_channels + self.compress_time = compress_time + self.unshuffle = nn.PixelUnshuffle(2) + self.spatial_conv = Kandinsky6SRConv3d( + in_channels, out_channels, kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1) ) - - if self.temporal_compress: - self.temporal_conv = CachedCausalConv3d( - in_channels, - in_channels, - kernel_size=(3, 1, 1), - stride=(1, 1, 1), - dilation=(1, 1, 1), - padding_mode=padding_mode, + if compress_time: + self.temporal_conv = nn.ModuleList( + [ + Kandinsky6SRCausalConv3d(out_channels, out_channels, kernel_size=(2, 1, 1), stride=(1, 1, 1)), + Kandinsky6SRCausalConv3d(out_channels, out_channels, kernel_size=(2, 1, 1), stride=(2, 1, 1)), + ] + ) + self.linear = Kandinsky6SRConv3d(out_channels, out_channels, kernel_size=1) + + def make_cache(self) -> dict: + if not self.compress_time: + return {} + return {"temporal_conv": [conv.make_cache() for conv in self.temporal_conv]} + + def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: + batch_size, channels, num_frames, height, width = hidden_states.shape + + # Spatial: pixel-unshuffle, average pairs of sub-pixel channels, add the strided convolution. + frames = hidden_states.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames, channels, height, width) + frames = self.unshuffle(frames) + frames = frames.view(batch_size * num_frames, 2 * channels, 2, height // 2, width // 2).mean(dim=2) + frames = frames.view(batch_size, num_frames, 2 * channels, height // 2, width // 2).permute(0, 2, 1, 3, 4) + hidden_states = self.spatial_conv(hidden_states) + frames + + if not self.compress_time: + return self.linear(hidden_states) + + # Temporal: average pairs of frames (the first segment keeps its first frame), add the causal convolutions. + batch_size, channels, num_frames, height, width = hidden_states.shape + sequence = hidden_states.permute(0, 3, 4, 1, 2).reshape(-1, channels, num_frames) + if cache["temporal_conv"][0]["padding"] is None: + first, rest = sequence[..., :1], sequence[..., 1:] + pooled = ( + torch.cat([first, F.avg_pool1d(rest, kernel_size=2, stride=2)], dim=-1) if rest.size(-1) else first ) - - self.linear = SafeConv3d(in_channels, in_channels, kernel_size=1, stride=1) - - def spatial_upsample_NEW(self, input_): - def conv_part(x): - to = torch.empty_like(x) - out = self.spatial_conv(x, write_to=to) - return out - - # 5D interpolate keeps channels_last_3d layout intact; merging dims - # via view() is impossible for channels_last strides (dim 1 is innermost) - input_interp = F.interpolate(input_, scale_factor=(1, 2, 2), mode="nearest") - input_interp.add_(conv_part(input_interp)) - return input_interp - - def temporal_upsample(self, input_, cache): - # input_ : (T + 1) x H x W - - repeated = input_.repeat_interleave(2, dim=2) - # repeated: (2T + 2) x H x W - - if cache["padding"] is None: - tail = repeated[..., 1:, :, :] # tail: (2T + 1) x H x W else: - tail = repeated + pooled = F.avg_pool1d(sequence, kernel_size=2, stride=2) + pooled = pooled.reshape(batch_size, height, width, channels, -1).permute(0, 3, 4, 1, 2) - conv_out = self.temporal_conv(tail, cache) - return conv_out + tail + conv_out = self.temporal_conv[0](hidden_states, cache["temporal_conv"][0]) + conv_out = self.temporal_conv[1](conv_out, cache["temporal_conv"][1]) + return self.linear(conv_out + pooled) - def forward(self, x, cache): - if self.temporal_compress: - # TEMPORAL UPSAMPLE - x = self.temporal_upsample(x, cache) - # SPATIAL UPSAMPLE - s_out = self.spatial_upsample_NEW(x) - to = torch.empty_like(s_out) +class Kandinsky6SRUpsample(nn.Module): + """Spatial 2x nearest upsample with a convolutional residual, preceded by an optional causal temporal 2x.""" - lin_out = self.linear(s_out, write_to=to) - - return lin_out - - -class CachedEncoder3D(nn.Module): + def __init__(self, channels: int, compress_time: bool) -> None: + super().__init__() + self.compress_time = compress_time + self.spatial_conv = Kandinsky6SRConv3d(channels, channels, kernel_size=(1, 3, 3), padding=(0, 1, 1)) + if compress_time: + self.temporal_conv = Kandinsky6SRCausalConv3d(channels, channels, kernel_size=(3, 1, 1)) + self.linear = Kandinsky6SRConv3d(channels, channels, kernel_size=1) + + def make_cache(self) -> dict: + return {"temporal_conv": self.temporal_conv.make_cache()} if self.compress_time else {} + + def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: + if self.compress_time: + # `T + 1` frames become `2T + 1`: every frame is repeated and the first segment drops the extra copy of + # its first frame. + repeated = hidden_states.repeat_interleave(2, dim=2) + if cache["temporal_conv"]["padding"] is None: + repeated = repeated[:, :, 1:] + hidden_states = self.temporal_conv(repeated, cache["temporal_conv"]) + repeated + + hidden_states = F.interpolate(hidden_states, scale_factor=(1, 2, 2), mode="nearest") + hidden_states = hidden_states + self.spatial_conv(hidden_states) + return self.linear(hidden_states) + + +class Kandinsky6SREncoder3D(nn.Module): def __init__( self, - *, - ch=128, - ch_mult=(1, 2, 4, 8), - num_res_blocks, - dropout=0.0, - in_channels, - resolution=0, - z_channels, - double_z=True, - padding_mode=None, - temporal_compress_times=4, - fix_pxs=False, - norm_type="group_norm", - downsample_version=1, - temporal_compress_start_level=0, - skip_last_resolution=False, - ): + in_channels: int, + latent_channels: int, + block_out_channels: tuple[int, ...], + layers_per_block: int, + temporal_compression_ratio: int, + temporal_compression_start_level: int, + ) -> None: super().__init__() - self.ch = ch - self.temb_ch = 0 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.in_channels = in_channels - self.skip_last_resolution = skip_last_resolution - - # log2 of temporal_compress_times - temporal_compress_level = int(np.log2(temporal_compress_times)) + temporal_compress_start_level - - in_ch_mult = (ch_mult[0],) + tuple(ch_mult) - self.conv_in = CachedCausalConv3d( - chan_in=in_channels, chan_out=round(in_ch_mult[0] * self.ch), kernel_size=3, padding_mode=padding_mode - ) - - normalization = CachedGroupNorm if norm_type == "group_norm" else Kandinsky6VAERMSNorm + num_levels = len(block_out_channels) + temporal_compression_end_level = int(math.log2(temporal_compression_ratio)) + temporal_compression_start_level - curr_res = resolution + self.conv_in = Kandinsky6SRCausalConv3d(in_channels, block_out_channels[0], kernel_size=3) self.down = nn.ModuleList() - for i_level in range(self.num_resolutions): - block = nn.ModuleList() - - block_in = round(ch * in_ch_mult[i_level]) - block_out = round(ch * ch_mult[i_level]) - - if downsample_version > 1 and i_level > 0: - block_in *= 2 - - for i_block in range(self.num_res_blocks): - block.append( - CachedCausalResnetBlock3D( - in_channels=block_in, - out_channels=block_out, - dropout=dropout, - temb_channels=self.temb_ch, - normalization=normalization, - padding_mode=padding_mode, - ) - ) + for level in range(num_levels): + # Every downsample doubles the channel count, so the next level starts at twice the previous width. + block_in = block_out_channels[0] if level == 0 else 2 * block_out_channels[level - 1] + block_out = block_out_channels[level] + blocks = nn.ModuleList() + for _ in range(layers_per_block): + blocks.append(Kandinsky6SRResnetBlock3D(block_in, block_out)) block_in = block_out down = nn.Module() - down.block = block - if i_level != self.num_resolutions - 1: - if temporal_compress_start_level <= i_level < temporal_compress_level: - down.downsample = CachedPXSDownsample( - block_in, - compress_time=True, - version=downsample_version, - fixed_stride=fix_pxs, - padding_mode=padding_mode, - ) - else: - down.downsample = CachedPXSDownsample( - block_in, - compress_time=False, - version=downsample_version, - fixed_stride=fix_pxs, - padding_mode=padding_mode, - ) - curr_res = curr_res // 2 - if not skip_last_resolution or i_level != self.num_resolutions - 1: - self.down.append(down) - - # middle - self.mid = nn.Module() - self.mid.block_1 = CachedCausalResnetBlock3D( - in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - normalization=normalization, - padding_mode=padding_mode, - ) - - self.mid.block_2 = CachedCausalResnetBlock3D( - in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - normalization=normalization, - padding_mode=padding_mode, - ) - - # end - self.norm_out = normalization(block_in) + down.block = blocks + if level != num_levels - 1: + compress_time = temporal_compression_start_level <= level < temporal_compression_end_level + down.downsample = Kandinsky6SRDownsample(block_in, compress_time=compress_time) + self.down.append(down) - self.conv_out = CachedCausalConv3d( - chan_in=block_in, - chan_out=2 * z_channels if double_z else z_channels, - kernel_size=3, - padding_mode=padding_mode, - ) - - def forward(self, x, cache_dict, use_cp=True): - # timestep embedding - temb = None - - # downsampling - h = self.conv_in(x, cache=cache_dict["conv_in"]) - for i_level in range(self.num_resolutions): - if not self.skip_last_resolution or i_level != self.num_resolutions - 1: - for i_block in range(self.num_res_blocks): - h = self.down[i_level].block[i_block](h, temb, layer_cache=cache_dict[i_level][i_block]) - if i_level != self.num_resolutions - 1: - h = self.down[i_level].downsample(h, cache=cache_dict[i_level]["down"]) - - # middle - h = self.mid.block_1(h, temb, layer_cache=cache_dict["mid_1"]) - h = self.mid.block_2(h, temb, layer_cache=cache_dict["mid_2"]) - - # end - h = self.norm_out(h, cache=cache_dict["norm_out"]) - h = F.silu(h) - h = self.conv_out(h, cache=cache_dict["conv_out"]) + self.mid = nn.Module() + self.mid.block_1 = Kandinsky6SRResnetBlock3D(block_in, block_in) + self.mid.block_2 = Kandinsky6SRResnetBlock3D(block_in, block_in) + self.norm_out = Kandinsky6SRRMSNorm(block_in) + self.conv_out = Kandinsky6SRCausalConv3d(block_in, 2 * latent_channels, kernel_size=3) + + def make_cache(self) -> dict: + return { + "conv_in": self.conv_in.make_cache(), + "down": [ + { + "block": [block.make_cache() for block in down.block], + "downsample": down.downsample.make_cache() if hasattr(down, "downsample") else {}, + } + for down in self.down + ], + "mid_1": self.mid.block_1.make_cache(), + "mid_2": self.mid.block_2.make_cache(), + "conv_out": self.conv_out.make_cache(), + } - return h + def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: + hidden_states = self.conv_in(hidden_states, cache["conv_in"]) + for down, level_cache in zip(self.down, cache["down"]): + for block, block_cache in zip(down.block, level_cache["block"]): + hidden_states = block(hidden_states, block_cache) + if hasattr(down, "downsample"): + hidden_states = down.downsample(hidden_states, level_cache["downsample"]) + hidden_states = self.mid.block_1(hidden_states, cache["mid_1"]) + hidden_states = self.mid.block_2(hidden_states, cache["mid_2"]) + hidden_states = F.silu(self.norm_out(hidden_states)) + return self.conv_out(hidden_states, cache["conv_out"]) -class CachedDecoder3D(nn.Module): +class Kandinsky6SRDecoder3D(nn.Module): def __init__( self, - ch=128, - out_ch=None, - ch_mult=(1, 2, 4, 8), - num_res_blocks=2, - dropout=0.0, - resolution=0, - z_channels=16, - give_pre_end=False, - zq_ch=None, - add_conv=False, - padding_mode=None, - temporal_compress_times=4, - norm_type="group_norm", - temporal_compress_start_level=0, - ): + out_channels: int, + latent_channels: int, + block_out_channels: tuple[int, ...], + layers_per_block: int, + temporal_compression_ratio: int, + temporal_compression_start_level: int, + ) -> None: super().__init__() - self.ch = ch - self.temb_ch = 0 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.give_pre_end = give_pre_end - - # log2 of temporal_compress_times - temporal_compress_level = int(np.log2(temporal_compress_times)) + temporal_compress_start_level - - if zq_ch is None: - zq_ch = z_channels - - # compute in_ch_mult, block_in and curr_res at lowest res - block_in = round(ch * ch_mult[self.num_resolutions - 1]) - curr_res = resolution // 2 ** (self.num_resolutions - 1) - self.z_shape = (1, z_channels, curr_res, curr_res) - - self.conv_in = CachedCausalConv3d( - chan_in=z_channels, chan_out=block_in, kernel_size=3, padding_mode=padding_mode - ) + num_levels = len(block_out_channels) + temporal_compression_end_level = int(math.log2(temporal_compression_ratio)) + temporal_compression_start_level - modulated_norm = functools.partial( - Normalize3D, normalization=CachedGroupNorm if norm_type == "group_norm" else Kandinsky6VAERMSNorm - ) - - # middle + block_in = block_out_channels[-1] + self.conv_in = Kandinsky6SRCausalConv3d(latent_channels, block_in, kernel_size=3) self.mid = nn.Module() - self.mid.block_1 = CachedCausalResnetBlock3D( - in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - zq_ch=zq_ch, - add_conv=add_conv, - normalization=modulated_norm, - padding_mode=padding_mode, - ) - - self.mid.block_2 = CachedCausalResnetBlock3D( - in_channels=block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout, - zq_ch=zq_ch, - add_conv=add_conv, - normalization=modulated_norm, - padding_mode=padding_mode, - ) + self.mid.block_1 = Kandinsky6SRResnetBlock3D(block_in, block_in, zq_channels=latent_channels) + self.mid.block_2 = Kandinsky6SRResnetBlock3D(block_in, block_in, zq_channels=latent_channels) - # upsampling self.up = nn.ModuleList() - for i_level in reversed(range(self.num_resolutions)): - block = nn.ModuleList() - block_out = round(ch * ch_mult[i_level]) - for i_block in range(self.num_res_blocks + 1): - block.append( - CachedCausalResnetBlock3D( - in_channels=block_in, - out_channels=block_out, - temb_channels=self.temb_ch, - dropout=dropout, - zq_ch=zq_ch, - add_conv=add_conv, - normalization=modulated_norm, - padding_mode=padding_mode, - ) - ) + for level in reversed(range(num_levels)): + block_out = block_out_channels[level] + blocks = nn.ModuleList() + for _ in range(layers_per_block + 1): + blocks.append(Kandinsky6SRResnetBlock3D(block_in, block_out, zq_channels=latent_channels)) block_in = block_out up = nn.Module() - up.block = block - if i_level != 0: - if ( - self.num_resolutions - temporal_compress_start_level - > i_level - >= self.num_resolutions - temporal_compress_level - ): - up.upsample = CachedPXSUpsample(block_in, compress_time=True, padding_mode=padding_mode) - else: - up.upsample = CachedPXSUpsample(block_in, compress_time=False, padding_mode=padding_mode) + up.block = blocks + if level != 0: + compress_time = ( + num_levels - temporal_compression_start_level + > level + >= num_levels - temporal_compression_end_level + ) + up.upsample = Kandinsky6SRUpsample(block_in, compress_time=compress_time) self.up.insert(0, up) - self.norm_out = modulated_norm(block_in, zq_ch, add_conv=add_conv) - - self.conv_out = CachedCausalConv3d(chan_in=block_in, chan_out=out_ch, kernel_size=3, padding_mode=padding_mode) - - def forward(self, z, cache_dict): - self.last_z_shape = z.shape - - # timestep embedding - temb = None - - zq = z - h = self.conv_in(z, cache_dict["conv_in"]) - - # middle - h = self.mid.block_1(h, temb, layer_cache=cache_dict["mid_1"], zq=zq) - h = self.mid.block_2(h, temb, layer_cache=cache_dict["mid_2"], zq=zq) - - # upsampling - for i_level in reversed(range(self.num_resolutions)): - for i_block in range(self.num_res_blocks + 1): - h = self.up[i_level].block[i_block](h, temb, layer_cache=cache_dict[i_level][i_block], zq=zq) - if i_level != 0: - h = self.up[i_level].upsample(h, cache_dict[i_level]["up"]) - - # end - if self.give_pre_end: - return h - - h = self.norm_out(h, zq, cache_dict["norm_out"]) - h = F.silu(h) - h = self.conv_out(h, cache_dict["conv_out"]) - - return h - + self.norm_out = Kandinsky6SRSpatialNorm3D(block_in, latent_channels) + self.conv_out = Kandinsky6SRCausalConv3d(block_in, out_channels, kernel_size=3) + + def make_cache(self) -> dict: + return { + "conv_in": self.conv_in.make_cache(), + "mid_1": self.mid.block_1.make_cache(), + "mid_2": self.mid.block_2.make_cache(), + "up": [ + { + "block": [block.make_cache() for block in up.block], + "upsample": up.upsample.make_cache() if hasattr(up, "upsample") else {}, + } + for up in self.up + ], + "norm_out": self.norm_out.make_cache(), + "conv_out": self.conv_out.make_cache(), + } -@dataclass -class DecoderOutput: - sample: torch.Tensor + def forward(self, latents: torch.Tensor, cache: dict) -> torch.Tensor: + hidden_states = self.conv_in(latents, cache["conv_in"]) + hidden_states = self.mid.block_1(hidden_states, cache["mid_1"], zq=latents) + hidden_states = self.mid.block_2(hidden_states, cache["mid_2"], zq=latents) + for level in reversed(range(len(self.up))): + up, level_cache = self.up[level], cache["up"][level] + for block, block_cache in zip(up.block, level_cache["block"]): + hidden_states = block(hidden_states, block_cache, zq=latents) + if hasattr(up, "upsample"): + hidden_states = up.upsample(hidden_states, level_cache["upsample"]) + hidden_states = F.silu(self.norm_out(hidden_states, latents, cache["norm_out"])) + return self.conv_out(hidden_states, cache["conv_out"]) class Kandinsky6SRVAE(ModelMixin, ConfigMixin): - """Causal 3D VAE used by the Kandinsky 6 video super-resolution pipeline. + r""" + Causal 3D K-VAE used by [`Kandinsky6SRPipeline`] to encode and decode video. - The encoder and decoder process videos in temporal segments while reusing - their causal convolution state between segments. This keeps memory usage - bounded for long videos and preserves the checkpoint layout of the KVAE. + Videos are processed in temporal segments of 16 pixel frames (plus the leading frame). The causal convolutions + carry their padding state between segments, so the segmentation only bounds peak memory and does not change the + result. Args: - vae_type (`str`): VAE architecture identifier. Must be ``"video-kvae"``. - decoder_ch (`int`, *optional*): Base channel count for the decoder, if - different from the encoder's ``ch``. Defaults to ``ch``. - decoder_ch_mult (`tuple[float, ...]`, *optional*): Per-resolution channel - multiplier for the decoder, if different from the encoder's - ``ch_mult``. Defaults to ``ch_mult``. - scaling_factor (`float`, *optional*, defaults to 1.0): Latent scaling - factor stored in the component configuration. - spatial_factor (`int`, *optional*, defaults to 16): Spatial compression - factor of the VAE. - temporal_factor (`int`, *optional*, defaults to 4): Temporal compression - factor of the VAE. + in_channels (`int`, defaults to `3`): + Number of pixel channels. + out_channels (`int`, defaults to `3`): + Number of reconstructed pixel channels. + latent_channels (`int`, defaults to `64`): + Number of latent channels. + encoder_block_out_channels (`tuple[int, ...]`, defaults to `(16, 128, 256, 512, 1024)`): + Output width of the residual blocks at each encoder level; every level but the last halves the spatial size + and doubles the width on the way to the next level. + decoder_block_out_channels (`tuple[int, ...]`, defaults to `(16, 256, 512, 1024, 2048)`): + Output width of the residual blocks at each decoder level. + layers_per_block (`int`, defaults to `2`): + Number of residual blocks per encoder level; the decoder uses one more per level. + temporal_compression_ratio (`int`, defaults to `4`): + Temporal compression factor; `log2` of it consecutive levels also compress time. + temporal_compression_start_level (`int`, defaults to `1`): + First level that compresses time. + scaling_factor (`float`, defaults to `0.910344`): + Scale applied to the latents before they enter the diffusion transformer. """ - _no_split_modules = ["CachedEncoder3D", "CachedDecoder3D"] - _supports_gradient_checkpointing = False - - @staticmethod - def normalize_data(data): - """Normalize pixel values to the KVAE input range.""" - return data / 128 - 1.0 - - @staticmethod - def denormalize_data(data): - """Convert normalized KVAE outputs back to pixel values.""" - return (data + 1) * 128 + _no_split_modules = ["Kandinsky6SREncoder3D", "Kandinsky6SRDecoder3D"] @register_to_config def __init__( self, - vae_type: str, in_channels: int = 3, out_channels: int = 3, - z_channels: int = 16, - ch: int = 128, - ch_mult: tuple[float, ...] = (1, 2, 4, 8), - decoder_ch: int | None = None, - decoder_ch_mult: tuple[float, ...] | None = None, - num_res_blocks: int = 2, - dropout: float = 0.0, - resolution: int = 0, - padding_mode: str | None = None, - temporal_compress_times: int = 4, - temporal_compress_start_level: int = 0, - norm_type: str = "group_norm", - double_z: bool = True, - downsample_version: int = 1, - fix_pxs: bool = False, - skip_last_resolution: bool = False, - give_pre_end: bool = False, - zq_ch: int | None = None, - add_conv: bool = False, - scaling_factor: float = 1.0, - spatial_factor: int = 16, - temporal_factor: int = 4, + latent_channels: int = 64, + encoder_block_out_channels: tuple[int, ...] = (16, 128, 256, 512, 1024), + decoder_block_out_channels: tuple[int, ...] = (16, 256, 512, 1024, 2048), + layers_per_block: int = 2, + temporal_compression_ratio: int = 4, + temporal_compression_start_level: int = 1, + scaling_factor: float = 0.910344004631042, ) -> None: super().__init__() - if vae_type != "video-kvae": - raise ValueError(f"Kandinsky6SRVAE supports only 'video-kvae', got: {vae_type!r}") - self.encoder = CachedEncoder3D( + if len(encoder_block_out_channels) != len(decoder_block_out_channels): + raise ValueError("`encoder_block_out_channels` and `decoder_block_out_channels` must have the same length") + + self.encoder = Kandinsky6SREncoder3D( in_channels=in_channels, - z_channels=z_channels, - ch=ch, - ch_mult=ch_mult, - num_res_blocks=num_res_blocks, - dropout=dropout, - resolution=resolution, - padding_mode=padding_mode, - temporal_compress_times=temporal_compress_times, - temporal_compress_start_level=temporal_compress_start_level, - norm_type=norm_type, - double_z=double_z, - downsample_version=downsample_version, - fix_pxs=fix_pxs, - skip_last_resolution=skip_last_resolution, + latent_channels=latent_channels, + block_out_channels=encoder_block_out_channels, + layers_per_block=layers_per_block, + temporal_compression_ratio=temporal_compression_ratio, + temporal_compression_start_level=temporal_compression_start_level, ) - self.decoder = CachedDecoder3D( - out_ch=out_channels, - z_channels=z_channels, - ch=ch if decoder_ch is None else decoder_ch, - ch_mult=ch_mult if decoder_ch_mult is None else decoder_ch_mult, - num_res_blocks=num_res_blocks, - dropout=dropout, - resolution=resolution, - padding_mode=padding_mode, - temporal_compress_times=temporal_compress_times, - temporal_compress_start_level=temporal_compress_start_level, - norm_type=norm_type, - give_pre_end=give_pre_end, - zq_ch=zq_ch, - add_conv=add_conv, + self.decoder = Kandinsky6SRDecoder3D( + out_channels=out_channels, + latent_channels=latent_channels, + block_out_channels=decoder_block_out_channels, + layers_per_block=layers_per_block, + temporal_compression_ratio=temporal_compression_ratio, + temporal_compression_start_level=temporal_compression_start_level, ) - self.spatial_factor = int(spatial_factor) - self.temporal_factor = int(temporal_factor) - - def make_empty_cache(self, block: str): - """Create empty causal-convolution and normalization caches.""" - - def make_dict(name, p=None): - if name == "conv": - return {"padding": None} - - layer, module = name.split("_") - if layer == "norm": - if module == "enc": - return {"mean": None, "var": None} - else: - return {"norm": make_dict("norm_enc"), "add_conv": make_dict("conv")} - elif layer == "resblock": - return { - "norm1": make_dict(f"norm_{module}"), - "norm2": make_dict(f"norm_{module}"), - "conv1": make_dict("conv"), - "conv2": make_dict("conv"), - "conv_shortcut": make_dict("conv"), - } - elif layer.isdigit(): - out_dict = {"down": [make_dict("conv"), make_dict("conv")], "up": make_dict("conv")} - for i in range(p): - out_dict[i] = make_dict(f"resblock_{module}") - - return out_dict - - cache = { - "conv_in": make_dict("conv"), - "mid_1": make_dict(f"resblock_{block}"), - "mid_2": make_dict(f"resblock_{block}"), - "norm_out": make_dict(f"norm_{block}"), - "conv_out": make_dict("conv"), - } - for i in range(len(self.config.ch_mult)): - cache[i] = make_dict(f"{i}_block", p=self.config.num_res_blocks + 1) - return cache + + self.spatial_compression_ratio = 2 ** (len(encoder_block_out_channels) - 1) + self.temporal_compression_ratio = temporal_compression_ratio @apply_forward_hook - def encode(self, x, seg_len=16): - """Encode a video in temporal segments and return latents and segment sizes. + def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutput | tuple: + r""" + Encode a video into its latent distribution. Args: - x (`torch.Tensor`): Video tensor in ``(batch, channels, frames, height, width)`` format. - seg_len (`int`, *optional*, defaults to 16): Number of non-initial - frames processed in each segment. - - Returns: - `tuple[torch.Tensor, list[int]]`: Encoded latents and the pixel-space - segment sizes needed by :meth:`decode`. + x (`torch.Tensor` of shape `(batch_size, channels, num_frames, height, width)`): + Pixel video in `[-1, 1]`. `num_frames` should be `1 + k * temporal_compression_ratio`. + return_dict (`bool`, defaults to `True`): + Whether to return an [`~models.modeling_outputs.AutoencoderKLOutput`] instead of a plain tuple. """ - cache = self.make_empty_cache("enc") - - # Compute segment sizes. - split_list = [seg_len + 1] - n_frames = x.size(2) - (seg_len + 1) - while n_frames > 0: - split_list.append(seg_len) - n_frames -= seg_len - - split_list[-1] += n_frames - - # Encode each segment. - latent = [] - for chunk in torch.split(x, split_list, dim=2): - l = self.encoder(chunk, cache) - latent.append(DiagonalGaussianDistribution(l).mode()) - - latent = torch.cat(latent, dim=2) - return latent, split_list + cache = self.encoder.make_cache() + segment_lengths = [min(SEGMENT_FRAMES + 1, x.size(2))] + remaining = x.size(2) - segment_lengths[0] + while remaining > 0: + segment_lengths.append(min(SEGMENT_FRAMES, remaining)) + remaining -= SEGMENT_FRAMES + + moments = torch.cat( + [self.encoder(segment, cache) for segment in torch.split(x, segment_lengths, dim=2)], dim=2 + ) + posterior = DiagonalGaussianDistribution(moments) + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) @apply_forward_hook - def decode(self, z, split_list=None): - """Decode latent segments while reusing the causal decoder cache. + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple: + r""" + Decode latents into a video. Args: - z (`torch.Tensor`): Latent tensor in ``(batch, channels, frames, height, width)`` format. - split_list (`list[int]`, *optional*): Pixel-space segment sizes - returned by :meth:`encode`. - - Returns: - `DecoderOutput`: Decoded video in ``sample``. + z (`torch.Tensor` of shape `(batch_size, latent_channels, num_latent_frames, height, width)`): + Latents, already divided by `scaling_factor`. + return_dict (`bool`, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.DecoderOutput`] instead of a plain tuple. """ - cache = self.make_empty_cache("dec") - - # Compute latent segment sizes. - if split_list is None: - default_split_size = 16 // self.config.temporal_compress_times - time_dim = z.shape[2] - if time_dim == 1: - # image - split_list = [1] - else: - splits_num = (time_dim - 1) // default_split_size - split_list = [default_split_size] * splits_num - if (time_dim - 1) % default_split_size != 0: - split_list.append((time_dim - 1) % default_split_size) - split_list[0] += 1 + cache = self.decoder.make_cache() + latent_segment = SEGMENT_FRAMES // self.temporal_compression_ratio + num_latent_frames = z.size(2) + if num_latent_frames == 1: + segment_lengths = [1] else: - split_list = [math.ceil(size / self.config.temporal_compress_times) for size in split_list] - - # Decode each segment. - recs = [] - for chunk in torch.split(z, split_list, dim=2): - out = self.decoder(chunk, cache) - recs.append(out) - - recs = torch.cat(recs, dim=2) - return DecoderOutput(sample=recs) - - def forward(self, x, seg_len: int = 16): - """Encode and decode a video in one call. - - Args: - x (`torch.Tensor`): Video tensor in ``(batch, channels, frames, height, width)`` format. - seg_len (`int`, *optional*, defaults to 16): Number of non-initial - frames processed in each segment. + # The leading latent frame decodes to a single pixel frame; every following latent frame decodes to + # `temporal_compression_ratio` pixel frames. + segment_lengths = [latent_segment] * ((num_latent_frames - 1) // latent_segment) + if (num_latent_frames - 1) % latent_segment: + segment_lengths.append((num_latent_frames - 1) % latent_segment) + segment_lengths[0] += 1 + + decoded = torch.cat( + [self.decoder(segment, cache) for segment in torch.split(z, segment_lengths, dim=2)], dim=2 + ) + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) - Returns: - `DecoderOutput`: Reconstructed video in ``sample``. - """ - latent, split_list = self.encode(x, seg_len) - recs = self.decode(latent, split_list) - return recs + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: torch.Generator | None = None, + ) -> DecoderOutput | tuple: + posterior = self.encode(sample).latent_dist + z = posterior.sample(generator=generator) if sample_posterior else posterior.mode() + decoded = self.decode(z).sample + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) __all__ = ["Kandinsky6SRVAE"] diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index fe2de0f00a86..d51200034d49 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -1,36 +1,79 @@ -"""MMAudio VAE Diffusers component.""" +# 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. + +"""MMAudio mel-spectrogram VAE and BigVGAN vocoder used by the Kandinsky 6 TI2VA pipeline.""" from __future__ import annotations import math -from typing import Any, Literal, Optional import torch import torch.nn as nn import torch.nn.functional as F + from ...configuration_utils import ConfigMixin, register_to_config -from ..modeling_utils import ModelMixin from ...utils.accelerate_utils import apply_forward_hook -from torch import pow, sin -from torch.nn import Conv1d, ConvTranspose1d, Parameter -from torch.nn.utils.parametrizations import weight_norm -from torch.nn.utils.parametrize import remove_parametrizations +from ..attention_dispatch import dispatch_attention_fn +from ..modeling_outputs import AutoencoderKLOutput +from ..modeling_utils import ModelMixin +from .vae import DecoderOutput, DiagonalGaussianDistribution -from .vae import DiagonalGaussianDistribution -# The magnitude-preserving building blocks below (`normalize`, `mp_silu`, `mp_sum`, `MPConv1D`) implement -# equations from Karras et al., "Analyzing and Improving the Training Dynamics of Diffusion Models" -# (https://arxiv.org/abs/2312.02696): each layer is designed so that if its inputs have unit variance, its -# output does too, which removes the need for the running activation-magnitude tracking earlier EDM variants -# relied on. MMAudio's VAE reuses this design for its 1D convolutional blocks. +# Activations of the magnitude-preserving blocks are clipped to this range. +ACTIVATION_CLIP = 256.0 -def normalize(x: torch.Tensor, dim: "list[int] | None" = None, eps: float = 1e-4) -> torch.Tensor: - """Rescale `x` to unit L2 norm over `dim` (default: every dimension but the first). +def mel_filterbank(sample_rate: int, n_fft: int, num_mels: int, f_min: float, f_max: float) -> torch.Tensor: + """Slaney-scale mel filterbank of shape `(num_mels, n_fft // 2 + 1)` with Slaney area normalization, the + `librosa.filters.mel` default that the MMAudio front end was trained with.""" - `eps` is scaled by `sqrt(norm.numel() / x.numel())` so it stays a meaningful floor relative to the - norm's typical magnitude regardless of how many elements are reduced over. - """ + def hz_to_mel(freq: torch.Tensor) -> torch.Tensor: + linear_step = 200.0 / 3 + min_log_hz = 1000.0 + log_step = math.log(6.4) / 27.0 + return torch.where( + freq >= min_log_hz, min_log_hz / linear_step + torch.log(freq / min_log_hz) / log_step, freq / linear_step + ) + + def mel_to_hz(mel: torch.Tensor) -> torch.Tensor: + linear_step = 200.0 / 3 + min_log_hz = 1000.0 + log_step = math.log(6.4) / 27.0 + min_log_mel = min_log_hz / linear_step + return torch.where( + mel >= min_log_mel, min_log_hz * torch.exp(log_step * (mel - min_log_mel)), linear_step * mel + ) + + fft_freqs = torch.linspace(0, sample_rate / 2, 1 + n_fft // 2, dtype=torch.float64) + mel_limits = hz_to_mel(torch.tensor([f_min, f_max], dtype=torch.float64)) + mel_freqs = mel_to_hz(torch.linspace(mel_limits[0], mel_limits[1], num_mels + 2, dtype=torch.float64)) + freq_diff = torch.diff(mel_freqs) + ramps = mel_freqs[:, None] - fft_freqs[None, :] + lower = -ramps[:-2] / freq_diff[:-1, None] + upper = ramps[2:] / freq_diff[1:, None] + weights = torch.clamp(torch.minimum(lower, upper), min=0) + weights = weights * (2.0 / (mel_freqs[2 : num_mels + 2] - mel_freqs[:num_mels]))[:, None] + return weights.float() + + +# The magnitude-preserving building blocks below follow Karras et al., "Analyzing and Improving the Training +# Dynamics of Diffusion Models" (https://arxiv.org/abs/2312.02696): each layer keeps a unit-variance input +# unit-variance, which MMAudio's VAE relies on instead of normalization layers. + + +def normalize(x: torch.Tensor, dim: list[int] | None = None, eps: float = 1e-4) -> torch.Tensor: + """Rescale `x` to unit L2 norm over `dim` (default: every dimension but the first).""" if dim is None: dim = list(range(1, x.ndim)) norm = torch.linalg.vector_norm(x, dim=dim, keepdim=True, dtype=torch.float32) @@ -39,1267 +82,535 @@ def normalize(x: torch.Tensor, dim: "list[int] | None" = None, eps: float = 1e-4 def mp_silu(x: torch.Tensor) -> torch.Tensor: - """SiLU rescaled by the standard deviation of `silu(z)` for `z ~ N(0, 1)`, so a unit-variance input - stays unit-variance after activation. - """ + """SiLU rescaled so that a unit-variance input stays unit-variance.""" return F.silu(x) / 0.596 def mp_sum(a: torch.Tensor, b: torch.Tensor, t: float = 0.5) -> torch.Tensor: - """Interpolate between `a` and `b` and rescale by `1 / sqrt((1-t)^2 + t^2)`, the factor that keeps the - result unit-variance when `a` and `b` are unit-variance and uncorrelated. - """ + """Interpolate `a` and `b` and rescale so that the result stays unit-variance.""" return a.lerp(b, t) / math.sqrt((1 - t) ** 2 + t**2) -class MPConv1D(torch.nn.Module): - """1D convolution (or, with `kernel_size=1` on a 2D input, a linear layer) whose weight is normalized to - unit norm per output channel and rescaled by `1 / sqrt(fan_in)`, so its output variance matches its - input variance regardless of `in_channels` or `kernel_size`. +class MMAudioMPConv1d(nn.Conv1d): + """Magnitude-preserving 1D convolution with an optional per-call gain. The released weights are already + normalized to unit norm per output channel and scaled by `1 / sqrt(fan_in)`, so the weight is applied directly.""" - The weight starts as plain random init and only becomes usable after `remove_weight_norm()` folds the - normalization into `self.weight` once; `forward` then applies that folded weight directly rather than - renormalizing on every call, matching how the converted checkpoint stores it. - """ + def __init__(self, in_channels: int, out_channels: int, kernel_size: int) -> None: + super().__init__(in_channels, out_channels, kernel_size, padding=kernel_size // 2, bias=False) - def __init__(self, in_channels: int, out_channels: int, kernel_size: int): - super().__init__() - self.out_channels = out_channels - self.weight = torch.nn.Parameter(torch.randn(out_channels, in_channels, kernel_size)) - self.weight_norm_removed = False - - def forward(self, x: torch.Tensor, gain: float = 1) -> torch.Tensor: - if not self.weight_norm_removed: - raise RuntimeError("call remove_weight_norm() before inference") - - w = self.weight * gain - return F.conv1d(x, w, padding=(w.shape[-1] // 2,)) + def forward(self, x: torch.Tensor, gain: float | torch.Tensor = 1.0) -> torch.Tensor: + return F.conv1d(x, (self.weight * gain).to(x.dtype), padding=self.padding) - def remove_weight_norm(self) -> "MPConv1D": - w = self.weight.to(torch.float32) - w = normalize(w) - w = w / math.sqrt(w[0].numel()) - self.weight.data.copy_(w.to(self.weight.dtype)) - self.weight_norm_removed = True - return self - - -class ResnetBlock1D(nn.Module): - def __init__(self, *, in_dim, out_dim=None, conv_shortcut=False, kernel_size=3, use_norm=True): +class MMAudioResnetBlock1D(nn.Module): + def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3) -> None: super().__init__() - self.in_dim = in_dim - out_dim = in_dim if out_dim is None else out_dim - self.out_dim = out_dim - self.use_conv_shortcut = conv_shortcut - self.use_norm = use_norm - - self.conv1 = MPConv1D(in_dim, out_dim, kernel_size=kernel_size) - self.conv2 = MPConv1D(out_dim, out_dim, kernel_size=kernel_size) - if self.in_dim != self.out_dim: - if self.use_conv_shortcut: - self.conv_shortcut = MPConv1D(in_dim, out_dim, kernel_size=kernel_size) - else: - self.nin_shortcut = MPConv1D(in_dim, out_dim, kernel_size=1) + self.in_channels = in_channels + self.out_channels = out_channels + self.conv1 = MMAudioMPConv1d(in_channels, out_channels, kernel_size) + self.conv2 = MMAudioMPConv1d(out_channels, out_channels, kernel_size) + if in_channels != out_channels: + self.nin_shortcut = MMAudioMPConv1d(in_channels, out_channels, kernel_size=1) def forward(self, x: torch.Tensor) -> torch.Tensor: + x = normalize(x, dim=1) + hidden_states = self.conv1(mp_silu(x)) + hidden_states = self.conv2(mp_silu(hidden_states)) + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return mp_sum(x, hidden_states, t=0.3) - # pixel norm - if self.use_norm: - x = normalize(x, dim=1) - h = x - h = mp_silu(h) - h = self.conv1(h) - - h = mp_silu(h) - h = self.conv2(h) - - if self.in_dim != self.out_dim: - if self.use_conv_shortcut: - x = self.conv_shortcut(x) - else: - x = self.nin_shortcut(x) - - return mp_sum(x, h, t=0.3) - - -class AttnBlock1D(nn.Module): - def __init__(self, in_channels, num_heads=1): +class MMAudioAttnBlock1D(nn.Module): + def __init__(self, channels: int, num_heads: int = 1) -> None: super().__init__() - self.in_channels = in_channels - self.num_heads = num_heads - self.qkv = MPConv1D(in_channels, in_channels * 3, kernel_size=1) - self.proj_out = MPConv1D(in_channels, in_channels, kernel_size=1) - - def forward(self, x): - h = x - y = self.qkv(h) - y = y.reshape(y.shape[0], self.num_heads, -1, 3, y.shape[-1]) - q, k, v = normalize(y, dim=2).unbind(3) + self.qkv = MMAudioMPConv1d(channels, channels * 3, kernel_size=1) + self.proj_out = MMAudioMPConv1d(channels, channels, kernel_size=1) - q = q.permute(0, 1, 3, 2) - k = k.permute(0, 1, 3, 2) - v = v.permute(0, 1, 3, 2) - - h = F.scaled_dot_product_attention(q, k, v) - h = h.permute(0, 1, 3, 2).reshape(h.shape[0], -1, h.shape[2]) - - h = self.proj_out(h) - - return mp_sum(x, h, t=0.3) - - -class Upsample1D(nn.Module): - def __init__(self, in_channels, with_conv): - super().__init__() - self.with_conv = with_conv - if self.with_conv: - self.conv = MPConv1D(in_channels, in_channels, kernel_size=3) - - def forward(self, x): - x = F.interpolate(x, scale_factor=2.0, mode="nearest-exact") # support 3D tensor(B,C,T) - if self.with_conv: - x = self.conv(x) - return x - - -class Downsample1D(nn.Module): - def __init__(self, in_channels, with_conv): + def forward(self, x: torch.Tensor) -> torch.Tensor: + batch_size, channels, length = x.shape + qkv = self.qkv(x).reshape(batch_size, self.num_heads, -1, 3, length) + query, key, value = normalize(qkv, dim=2).unbind(3) + # `(B, heads, D, T)` -> `(B, T, heads, D)` for the attention dispatcher. + query, key, value = (t.permute(0, 3, 1, 2) for t in (query, key, value)) + hidden_states = dispatch_attention_fn(query, key, value) + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, channels, length) + return mp_sum(x, self.proj_out(hidden_states), t=0.3) + + +class MMAudioUpsample1D(nn.Module): + def __init__(self, channels: int) -> None: super().__init__() - self.with_conv = with_conv - if self.with_conv: - # no asymmetric padding in torch conv, must do it ourselves - self.conv1 = MPConv1D(in_channels, in_channels, kernel_size=1) - self.conv2 = MPConv1D(in_channels, in_channels, kernel_size=1) - - def forward(self, x): - - if self.with_conv: - x = self.conv1(x) + self.conv = MMAudioMPConv1d(channels, channels, kernel_size=3) - x = F.avg_pool1d(x, kernel_size=2, stride=2) - - if self.with_conv: - x = self.conv2(x) - - return x + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(F.interpolate(x, scale_factor=2.0, mode="nearest-exact")) -# fmt: off -DATA_MEAN_128D = [ - -3.3462, -2.6723, -2.4893, -2.3143, -2.2664, -2.3317, -2.1802, -2.4006, -2.2357, -2.4597, - -2.3717, -2.4690, -2.5142, -2.4919, -2.6610, -2.5047, -2.7483, -2.5926, -2.7462, -2.7033, - -2.7386, -2.8112, -2.7502, -2.9594, -2.7473, -3.0035, -2.8891, -2.9922, -2.9856, -3.0157, - -3.1191, -2.9893, -3.1718, -3.0745, -3.1879, -3.2310, -3.1424, -3.2296, -3.2791, -3.2782, - -3.2756, -3.3134, -3.3509, -3.3750, -3.3951, -3.3698, -3.4505, -3.4509, -3.5089, -3.4647, - -3.5536, -3.5788, -3.5867, -3.6036, -3.6400, -3.6747, -3.7072, -3.7279, -3.7283, -3.7795, - -3.8259, -3.8447, -3.8663, -3.9182, -3.9605, -3.9861, -4.0105, -4.0373, -4.0762, -4.1121, - -4.1488, -4.1874, -4.2461, -4.3170, -4.3639, -4.4452, -4.5282, -4.6297, -4.7019, -4.7960, - -4.8700, -4.9507, -5.0303, -5.0866, -5.1634, -5.2342, -5.3242, -5.4053, -5.4927, -5.5712, - -5.6464, -5.7052, -5.7619, -5.8410, -5.9188, -6.0103, -6.0955, -6.1673, -6.2362, -6.3120, - -6.3926, -6.4797, -6.5565, -6.6511, -6.8130, -6.9961, -7.1275, -7.2457, -7.3576, -7.4663, - -7.6136, -7.7469, -7.8815, -8.0132, -8.1515, -8.3071, -8.4722, -8.7418, -9.3975, -9.6628, - -9.7671, -9.8863, -9.9992, -10.0860, -10.1709, -10.5418, -11.2795, -11.3861 -] - -DATA_STD_128D = [ - 2.3804, 2.4368, 2.3772, 2.3145, 2.2803, 2.2510, 2.2316, 2.2083, 2.1996, 2.1835, 2.1769, 2.1659, - 2.1631, 2.1618, 2.1540, 2.1606, 2.1571, 2.1567, 2.1612, 2.1579, 2.1679, 2.1683, 2.1634, 2.1557, - 2.1668, 2.1518, 2.1415, 2.1449, 2.1406, 2.1350, 2.1313, 2.1415, 2.1281, 2.1352, 2.1219, 2.1182, - 2.1327, 2.1195, 2.1137, 2.1080, 2.1179, 2.1036, 2.1087, 2.1036, 2.1015, 2.1068, 2.0975, 2.0991, - 2.0902, 2.1015, 2.0857, 2.0920, 2.0893, 2.0897, 2.0910, 2.0881, 2.0925, 2.0873, 2.0960, 2.0900, - 2.0957, 2.0958, 2.0978, 2.0936, 2.0886, 2.0905, 2.0845, 2.0855, 2.0796, 2.0840, 2.0813, 2.0817, - 2.0838, 2.0840, 2.0917, 2.1061, 2.1431, 2.1976, 2.2482, 2.3055, 2.3700, 2.4088, 2.4372, 2.4609, - 2.4731, 2.4847, 2.5072, 2.5451, 2.5772, 2.6147, 2.6529, 2.6596, 2.6645, 2.6726, 2.6803, 2.6812, - 2.6899, 2.6916, 2.6931, 2.6998, 2.7062, 2.7262, 2.7222, 2.7158, 2.7041, 2.7485, 2.7491, 2.7451, - 2.7485, 2.7233, 2.7297, 2.7233, 2.7145, 2.6958, 2.6788, 2.6439, 2.6007, 2.4786, 2.2469, 2.1877, - 2.1392, 2.0717, 2.0107, 1.9676, 1.9140, 1.7102, 0.9101, 0.7164 -] -# fmt: on - - -class VAE(nn.Module): - def __init__( - self, - *, - data_dim: int, - embed_dim: int, - hidden_dim: int, - ): +class MMAudioDownsample1D(nn.Module): + def __init__(self, channels: int) -> None: super().__init__() + self.conv1 = MMAudioMPConv1d(channels, channels, kernel_size=1) + self.conv2 = MMAudioMPConv1d(channels, channels, kernel_size=1) - if data_dim != 128: - raise ValueError(f"Only 44k audio VAE data is supported, got data_dim={data_dim}") - self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_128D, dtype=torch.float32)) - self.data_std = nn.Buffer(torch.tensor(DATA_STD_128D, dtype=torch.float32)) - - self.data_mean = self.data_mean.view(1, -1, 1) - self.data_std = self.data_std.view(1, -1, 1) - - self.encoder = Encoder1D( - dim=hidden_dim, - ch_mult=(1, 2, 4), - num_res_blocks=2, - attn_layers=[3], - down_layers=[0], - in_dim=data_dim, - embed_dim=embed_dim, - ) - self.decoder = Decoder1D( - dim=hidden_dim, - ch_mult=(1, 2, 4), - num_res_blocks=2, - attn_layers=[3], - down_layers=[0], - in_dim=data_dim, - out_dim=data_dim, - embed_dim=embed_dim, - ) - - self.embed_dim = embed_dim - # self.quant_conv = nn.Conv1d(2 * embed_dim, 2 * embed_dim, 1) - # self.post_quant_conv = nn.Conv1d(embed_dim, embed_dim, 1) - - def encode(self, x: torch.Tensor, normalize: bool = True) -> DiagonalGaussianDistribution: - if normalize: - x = self.normalize(x) - moments = self.encoder(x) - posterior = DiagonalGaussianDistribution(moments) - return posterior - - def decode(self, z: torch.Tensor, unnormalize: bool = True) -> torch.Tensor: - dec = self.decoder(z) - if unnormalize: - dec = self.unnormalize(dec) - return dec - - def normalize(self, x: torch.Tensor) -> torch.Tensor: - return (x - self.data_mean) / self.data_std - - def unnormalize(self, x: torch.Tensor) -> torch.Tensor: - return x * self.data_std + self.data_mean - - def forward( - self, - x: torch.Tensor, - sample_posterior: bool = True, - rng: Optional[torch.Generator] = None, - normalize: bool = True, - unnormalize: bool = True, - ) -> tuple[torch.Tensor, DiagonalGaussianDistribution]: - - posterior = self.encode(x, normalize=normalize) - if sample_posterior: - z = posterior.sample(rng) - else: - z = posterior.mode() - dec = self.decode(z, unnormalize=unnormalize) - return dec, posterior - - def load_weights(self, src_dict) -> None: - self.load_state_dict(src_dict, strict=True) - - @property - def device(self) -> torch.device: - return next(self.parameters()).device - - def get_last_layer(self): - return self.decoder.conv_out.weight - - def remove_weight_norm(self): - for name, m in self.named_modules(): - if isinstance(m, MPConv1D): - m.remove_weight_norm() + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv2(F.avg_pool1d(self.conv1(x), kernel_size=2, stride=2)) - return self +class MMAudioEncoder1D(nn.Module): + """Mel-spectrogram encoder: residual blocks over `channel_multipliers` levels, a single 2x temporal downsample + after the first level, and an attention block in the middle.""" -class Encoder1D(nn.Module): def __init__( self, - *, - dim: int, - ch_mult: tuple[int] = (1, 2, 4, 8), - num_res_blocks: int, - attn_layers: list[int] = [], - down_layers: list[int] = [], - resamp_with_conv: bool = True, - in_dim: int, - embed_dim: int, - double_z: bool = True, - kernel_size: int = 3, - clip_act: float = 256.0, - ): + mel_bins: int, + latent_channels: int, + hidden_channels: int, + channel_multipliers: tuple[int, ...], + layers_per_block: int, + ) -> None: super().__init__() - self.dim = dim - self.num_layers = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.in_channels = in_dim - self.clip_act = clip_act - self.down_layers = down_layers - self.attn_layers = attn_layers - self.conv_in = MPConv1D(in_dim, self.dim, kernel_size=kernel_size) - - in_ch_mult = (1,) + tuple(ch_mult) - self.in_ch_mult = in_ch_mult - # downsampling + self.conv_in = MMAudioMPConv1d(mel_bins, hidden_channels, kernel_size=3) + self.down = nn.ModuleList() - for i_level in range(self.num_layers): - block = nn.ModuleList() - attn = nn.ModuleList() - block_in = dim * in_ch_mult[i_level] - block_out = dim * ch_mult[i_level] - for i_block in range(self.num_res_blocks): - block.append(ResnetBlock1D(in_dim=block_in, out_dim=block_out, kernel_size=kernel_size, use_norm=True)) + block_in = hidden_channels + for level, multiplier in enumerate(channel_multipliers): + block_out = hidden_channels * multiplier + blocks = nn.ModuleList() + for _ in range(layers_per_block): + blocks.append(MMAudioResnetBlock1D(block_in, block_out)) block_in = block_out - if i_level in attn_layers: - attn.append(AttnBlock1D(block_in)) down = nn.Module() - down.block = block - down.attn = attn - if i_level in down_layers: - down.downsample = Downsample1D(block_in, resamp_with_conv) + down.block = blocks + if level == 0: + down.downsample = MMAudioDownsample1D(block_in) self.down.append(down) - # middle self.mid = nn.Module() - self.mid.block_1 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, kernel_size=kernel_size, use_norm=True) - self.mid.attn_1 = AttnBlock1D(block_in) - self.mid.block_2 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, kernel_size=kernel_size, use_norm=True) - - # end - self.conv_out = MPConv1D(block_in, 2 * embed_dim if double_z else embed_dim, kernel_size=kernel_size) + self.mid.block_1 = MMAudioResnetBlock1D(block_in, block_in) + self.mid.attn_1 = MMAudioAttnBlock1D(block_in) + self.mid.block_2 = MMAudioResnetBlock1D(block_in, block_in) + self.conv_out = MMAudioMPConv1d(block_in, 2 * latent_channels, kernel_size=3) self.learnable_gain = nn.Parameter(torch.zeros([])) - def forward(self, x): - - # downsampling - hs = [self.conv_in(x)] - for i_level in range(self.num_layers): - for i_block in range(self.num_res_blocks): - h = self.down[i_level].block[i_block](hs[-1]) - if len(self.down[i_level].attn) > 0: - h = self.down[i_level].attn[i_block](h) - h = h.clamp(-self.clip_act, self.clip_act) - hs.append(h) - if i_level in self.down_layers: - hs.append(self.down[i_level].downsample(hs[-1])) - - # middle - h = hs[-1] - h = self.mid.block_1(h) - h = self.mid.attn_1(h) - h = self.mid.block_2(h) - h = h.clamp(-self.clip_act, self.clip_act) - - # end - h = mp_silu(h) - h = self.conv_out(h, gain=(self.learnable_gain + 1)) - return h - - -class Decoder1D(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_in(x) + for down in self.down: + for block in down.block: + hidden_states = block(hidden_states).clamp(-ACTIVATION_CLIP, ACTIVATION_CLIP) + if hasattr(down, "downsample"): + hidden_states = down.downsample(hidden_states) + + hidden_states = self.mid.block_1(hidden_states) + hidden_states = self.mid.attn_1(hidden_states) + hidden_states = self.mid.block_2(hidden_states).clamp(-ACTIVATION_CLIP, ACTIVATION_CLIP) + return self.conv_out(mp_silu(hidden_states), gain=self.learnable_gain + 1) + + +class MMAudioDecoder1D(nn.Module): + """Mirror of [`MMAudioEncoder1D`]: the 2x temporal upsample sits after the second-to-last level.""" + def __init__( self, - *, - dim: int, - out_dim: int, - ch_mult: tuple[int] = (1, 2, 4, 8), - num_res_blocks: int, - attn_layers: list[int] = [], - down_layers: list[int] = [], - kernel_size: int = 3, - resamp_with_conv: bool = True, - in_dim: int, - embed_dim: int, - clip_act: float = 256.0, - ): + mel_bins: int, + latent_channels: int, + hidden_channels: int, + channel_multipliers: tuple[int, ...], + layers_per_block: int, + ) -> None: super().__init__() - self.ch = dim - self.num_layers = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.in_channels = in_dim - self.clip_act = clip_act - self.down_layers = [i + 1 for i in down_layers] # each downlayer add one - - # compute in_ch_mult, block_in and curr_res at lowest res - block_in = dim * ch_mult[self.num_layers - 1] + block_in = hidden_channels * channel_multipliers[-1] + self.conv_in = MMAudioMPConv1d(latent_channels, block_in, kernel_size=3) - # z to block_in - self.conv_in = MPConv1D(embed_dim, block_in, kernel_size=kernel_size) - - # middle self.mid = nn.Module() - self.mid.block_1 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True) - self.mid.attn_1 = AttnBlock1D(block_in) - self.mid.block_2 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True) + self.mid.block_1 = MMAudioResnetBlock1D(block_in, block_in) + self.mid.attn_1 = MMAudioAttnBlock1D(block_in) + self.mid.block_2 = MMAudioResnetBlock1D(block_in, block_in) - # upsampling self.up = nn.ModuleList() - for i_level in reversed(range(self.num_layers)): - block = nn.ModuleList() - attn = nn.ModuleList() - block_out = dim * ch_mult[i_level] - for i_block in range(self.num_res_blocks + 1): - block.append(ResnetBlock1D(in_dim=block_in, out_dim=block_out, use_norm=True)) + for level in reversed(range(len(channel_multipliers))): + block_out = hidden_channels * channel_multipliers[level] + blocks = nn.ModuleList() + for _ in range(layers_per_block + 1): + blocks.append(MMAudioResnetBlock1D(block_in, block_out)) block_in = block_out - if i_level in attn_layers: - attn.append(AttnBlock1D(block_in)) up = nn.Module() - up.block = block - up.attn = attn - if i_level in self.down_layers: - up.upsample = Upsample1D(block_in, resamp_with_conv) - self.up.insert(0, up) # prepend to get consistent order - - # end - self.conv_out = MPConv1D(block_in, out_dim, kernel_size=kernel_size) - self.learnable_gain = nn.Parameter(torch.zeros([])) - - def forward(self, z): - # z to block_in - h = self.conv_in(z) - - # middle - h = self.mid.block_1(h) - h = self.mid.attn_1(h) - h = self.mid.block_2(h) - h = h.clamp(-self.clip_act, self.clip_act) - - # upsampling - for i_level in reversed(range(self.num_layers)): - for i_block in range(self.num_res_blocks + 1): - h = self.up[i_level].block[i_block](h) - if len(self.up[i_level].attn) > 0: - h = self.up[i_level].attn[i_block](h) - h = h.clamp(-self.clip_act, self.clip_act) - if i_level in self.down_layers: - h = self.up[i_level].upsample(h) - - h = mp_silu(h) - h = self.conv_out(h, gain=(self.learnable_gain + 1)) - return h - - -# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license. -# LICENSE is in incl_licenses directory. - - -class Snake(nn.Module): - """ - Implementation of a sine-based periodic activation function - Shape: - - Input: (B, C, T) - - Output: (B, C, T), same shape as the input - Parameters: - - alpha - trainable parameter - References: - - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: - https://arxiv.org/abs/2006.08195 - Examples: - >>> a1 = snake(256) - >>> x = torch.randn(256) - >>> x = a1(x) - """ - - def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): - """ - Initialization. - INPUT: - - in_features: shape of the input - - alpha: trainable parameter - alpha is initialized to 1 by default, higher values = higher-frequency. - alpha will be trained along with the rest of your model. - """ - super(Snake, self).__init__() - self.in_features = in_features - - # Initialize alpha - self.alpha_logscale = alpha_logscale - if self.alpha_logscale: # Log scale alphas initialized to zeros - self.alpha = Parameter(torch.zeros(in_features) * alpha) - else: # Linear scale alphas initialized to ones - self.alpha = Parameter(torch.ones(in_features) * alpha) - - self.alpha.requires_grad = alpha_trainable - - self.no_div_by_zero = 0.000000001 - - def forward(self, x): - """ - Forward pass of the function. - Applies the function to the input elementwise. - Snake ∶= x + 1/a * sin^2 (xa) - """ - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # Line up with x to [B, C, T] - if self.alpha_logscale: - alpha = torch.exp(alpha) - x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2) - - return x - + up.block = blocks + if level == 1: + up.upsample = MMAudioUpsample1D(block_in) + self.up.insert(0, up) -class SnakeBeta(nn.Module): - """ - A modified Snake function which uses separate parameters for the magnitude of the periodic components - Shape: - - Input: (B, C, T) - - Output: (B, C, T), same shape as the input - Parameters: - - alpha - trainable parameter that controls frequency - - beta - trainable parameter that controls magnitude - References: - - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: - https://arxiv.org/abs/2006.08195 - Examples: - >>> a1 = snakebeta(256) - >>> x = torch.randn(256) - >>> x = a1(x) - """ - - def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False): - """ - Initialization. - INPUT: - - in_features: shape of the input - - alpha - trainable parameter that controls frequency - - beta - trainable parameter that controls magnitude - alpha is initialized to 1 by default, higher values = higher-frequency. - beta is initialized to 1 by default, higher values = higher-magnitude. - alpha will be trained along with the rest of your model. - """ - super(SnakeBeta, self).__init__() - self.in_features = in_features + self.conv_out = MMAudioMPConv1d(block_in, mel_bins, kernel_size=3) + self.learnable_gain = nn.Parameter(torch.zeros([])) - # Initialize alpha - self.alpha_logscale = alpha_logscale - if self.alpha_logscale: # Log scale alphas initialized to zeros - self.alpha = Parameter(torch.zeros(in_features) * alpha) - self.beta = Parameter(torch.zeros(in_features) * alpha) - else: # Linear scale alphas initialized to ones - self.alpha = Parameter(torch.ones(in_features) * alpha) - self.beta = Parameter(torch.ones(in_features) * alpha) + def forward(self, z: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_in(z) + hidden_states = self.mid.block_1(hidden_states) + hidden_states = self.mid.attn_1(hidden_states) + hidden_states = self.mid.block_2(hidden_states).clamp(-ACTIVATION_CLIP, ACTIVATION_CLIP) - self.alpha.requires_grad = alpha_trainable - self.beta.requires_grad = alpha_trainable + for level in reversed(range(len(self.up))): + up = self.up[level] + for block in up.block: + hidden_states = block(hidden_states).clamp(-ACTIVATION_CLIP, ACTIVATION_CLIP) + if hasattr(up, "upsample"): + hidden_states = up.upsample(hidden_states) + return self.conv_out(mp_silu(hidden_states), gain=self.learnable_gain + 1) - self.no_div_by_zero = 0.000000001 - def forward(self, x): - """ - Forward pass of the function. - Applies the function to the input elementwise. - SnakeBeta ∶= x + 1/b * sin^2 (xa) - """ - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # Line up with x to [B, C, T] - beta = self.beta.unsqueeze(0).unsqueeze(-1) - if self.alpha_logscale: - alpha = torch.exp(alpha) - beta = torch.exp(beta) - x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2) +class MMAudioAutoencoder(nn.Module): + """Encoder/decoder pair over standardized log-mel spectrograms. `data_mean` and `data_std` hold the per-bin + statistics the checkpoint was trained with.""" - return x + def __init__( + self, + mel_bins: int, + latent_channels: int, + hidden_channels: int, + channel_multipliers: tuple[int, ...], + layers_per_block: int, + ) -> None: + super().__init__() + self.register_buffer("data_mean", torch.zeros(1, mel_bins, 1)) + self.register_buffer("data_std", torch.ones(1, mel_bins, 1)) + self.encoder = MMAudioEncoder1D( + mel_bins, latent_channels, hidden_channels, channel_multipliers, layers_per_block + ) + self.decoder = MMAudioDecoder1D( + mel_bins, latent_channels, hidden_channels, channel_multipliers, layers_per_block + ) + def encode(self, mel: torch.Tensor) -> torch.Tensor: + return self.encoder((mel - self.data_mean) / self.data_std) -# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 -# LICENSE is in incl_licenses directory. + def decode(self, z: torch.Tensor) -> torch.Tensor: + return self.decoder(z) * self.data_std + self.data_mean -if "sinc" in dir(torch): - sinc = torch.sinc -else: - # This code is adopted from adefossez's julius.core.sinc under the MIT License - # https://adefossez.github.io/julius/julius/core.html - # LICENSE is in incl_licenses directory. - def sinc(x: torch.Tensor): - """ - Implementation of sinc, i.e. sin(pi * x) / (pi * x) - __Warning__: Different to julius.sinc, the input is multiplied by `pi`! - """ - return torch.where( - x == 0, - torch.tensor(1.0, device=x.device, dtype=x.dtype), - torch.sin(math.pi * x) / math.pi / x, - ) +# BigVGAN vocoder (https://github.com/NVIDIA/BigVGAN, MIT license) with the anti-aliased Snake activations of +# https://github.com/junjun3518/alias-free-torch (Apache 2.0). -# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License -# https://adefossez.github.io/julius/julius/lowpass.html -# LICENSE is in incl_licenses directory. -def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): # return filter [1,1,kernel_size] +def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: + """Kaiser-windowed sinc low-pass filter of shape `(1, 1, kernel_size)` normalized to unit sum.""" even = kernel_size % 2 == 0 half_size = kernel_size // 2 - # For kaiser window - delta_f = 4 * half_width - A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 - if A > 50.0: - beta = 0.1102 * (A - 8.7) - elif A >= 21.0: - beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0) + attenuation = 2.285 * (half_size - 1) * math.pi * 4 * half_width + 7.95 + if attenuation > 50.0: + beta = 0.1102 * (attenuation - 8.7) + elif attenuation >= 21.0: + beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0) else: beta = 0.0 window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) - # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio - if even: - time = torch.arange(-half_size, half_size) + 0.5 - else: - time = torch.arange(kernel_size) - half_size - if cutoff == 0: - filter_ = torch.zeros_like(time) - else: - filter_ = 2 * cutoff * window * sinc(2 * cutoff * time) - """ - Normalize filter to have sum = 1, otherwise we will have a small leakage of the constant component in the input signal. - """ - filter_ /= filter_.sum() - filter = filter_.view(1, 1, kernel_size) + time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size + filter_ = 2 * cutoff * window * torch.sinc(2 * cutoff * time) + filter_ = filter_ / filter_.sum() + return filter_.view(1, 1, kernel_size) - return filter +class MMAudioSnakeBeta(nn.Module): + """`x + 1/b * sin^2(a * x)` with per-channel log-scale frequency `a` and magnitude `b`.""" -class LowPassFilter1d(nn.Module): - def __init__( - self, - cutoff=0.5, - half_width=0.6, - stride: int = 1, - padding: bool = True, - padding_mode: str = "replicate", - kernel_size: int = 12, - ): - """ - kernel_size should be even number for stylegan3 setup, in this implementation, odd number is also possible. - """ + def __init__(self, channels: int) -> None: super().__init__() - if cutoff < -0.0: - raise ValueError("Minimum cutoff must be larger than zero.") - if cutoff > 0.5: - raise ValueError("A cutoff above 0.5 does not make sense.") - self.kernel_size = kernel_size - self.even = kernel_size % 2 == 0 - self.pad_left = kernel_size // 2 - int(self.even) - self.pad_right = kernel_size // 2 - self.stride = stride - self.padding = padding - self.padding_mode = padding_mode - filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) - self.register_buffer("filter", filter) - - # Input [B, C, T] - def forward(self, x): - _, C, _ = x.shape + self.alpha = nn.Parameter(torch.zeros(channels)) + self.beta = nn.Parameter(torch.zeros(channels)) - if self.padding: - x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) - out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) + def forward(self, x: torch.Tensor) -> torch.Tensor: + alpha = torch.exp(self.alpha)[None, :, None] + beta = torch.exp(self.beta)[None, :, None] + return x + (1.0 / (beta + 1e-9)) * torch.sin(x * alpha).pow(2) - return out +class MMAudioLowPassFilter1d(nn.Module): + def __init__(self, cutoff: float, half_width: float, stride: int, kernel_size: int) -> None: + super().__init__() + even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size)) -# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 -# LICENSE is in incl_licenses directory. + def forward(self, x: torch.Tensor) -> torch.Tensor: + channels = x.shape[1] + x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate") + return F.conv1d(x, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels) -class UpSample1d(nn.Module): - def __init__(self, ratio=2, kernel_size=None): +class MMAudioUpSample1d(nn.Module): + def __init__(self, ratio: int, kernel_size: int) -> None: super().__init__() self.ratio = ratio - self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size self.stride = ratio - self.pad = self.kernel_size // ratio - 1 - self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 - self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 - filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size) - self.register_buffer("filter", filter) - - # x: [B, C, T] - def forward(self, x): - _, C, _ = x.shape + self.pad = kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (kernel_size - self.stride + 1) // 2 + self.register_buffer("filter", kaiser_sinc_filter1d(0.5 / ratio, 0.6 / ratio, kernel_size)) + def forward(self, x: torch.Tensor) -> torch.Tensor: + channels = x.shape[1] x = F.pad(x, (self.pad, self.pad), mode="replicate") - x = self.ratio * F.conv_transpose1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C) - x = x[..., self.pad_left : -self.pad_right] - - return x - - -class DownSample1d(nn.Module): - def __init__(self, ratio=2, kernel_size=None): - super().__init__() - self.ratio = ratio - self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size - self.lowpass = LowPassFilter1d( - cutoff=0.5 / ratio, - half_width=0.6 / ratio, - stride=ratio, - kernel_size=self.kernel_size, + x = self.ratio * F.conv_transpose1d( + x, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels ) + return x[..., self.pad_left : -self.pad_right] - def forward(self, x): - xx = self.lowpass(x) - - return xx - -# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0 -# LICENSE is in incl_licenses directory. - - -class Activation1d(nn.Module): - def __init__( - self, - activation, - up_ratio: int = 2, - down_ratio: int = 2, - up_kernel_size: int = 12, - down_kernel_size: int = 12, - ): +class MMAudioDownSample1d(nn.Module): + def __init__(self, ratio: int, kernel_size: int) -> None: super().__init__() - self.up_ratio = up_ratio - self.down_ratio = down_ratio - self.act = activation - self.upsample = UpSample1d(up_ratio, up_kernel_size) - self.downsample = DownSample1d(down_ratio, down_kernel_size) - - # x: [B,C,T] - def forward(self, x): - x = self.upsample(x) - x = self.act(x) - x = self.downsample(x) - - return x - - -# Adapted from https://github.com/jik876/hifi-gan under the MIT license. -# LICENSE is in incl_licenses directory. - - -class AttrDict(dict): - def __init__(self, *args, **kwargs): - super(AttrDict, self).__init__(*args, **kwargs) - self.__dict__ = self - - -# Adapted from https://github.com/jik876/hifi-gan under the MIT license. -# LICENSE is in incl_licenses directory. - - -def init_weights(m, mean=0.0, std=0.01): - classname = m.__class__.__name__ - if classname.find("Conv") != -1: - m.weight.data.normal_(mean, std) - + self.lowpass = MMAudioLowPassFilter1d(0.5 / ratio, 0.6 / ratio, stride=ratio, kernel_size=kernel_size) -def get_padding(kernel_size, dilation=1): - return int((kernel_size * dilation - dilation) / 2) + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lowpass(x) -# Copyright (c) 2024 NVIDIA CORPORATION. -# Licensed under the MIT license. +class MMAudioActivation1d(nn.Module): + """Anti-aliased activation: 2x upsample, Snake-beta, 2x downsample.""" -# Adapted from https://github.com/jik876/hifi-gan under the MIT license. -# LICENSE is in incl_licenses directory. + def __init__(self, channels: int, ratio: int = 2, kernel_size: int = 12) -> None: + super().__init__() + self.act = MMAudioSnakeBeta(channels) + self.upsample = MMAudioUpSample1d(ratio, kernel_size) + self.downsample = MMAudioDownSample1d(ratio, kernel_size) + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.downsample(self.act(self.upsample(x))) -class AMPBlock1(torch.nn.Module): - """ - AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer. - AMPBlock1 has additional self.convs2 that contains additional Conv1d layers with a fixed dilation=1 followed by each layer in self.convs1 - Args: - h (AttrDict): Hyperparameters. - channels (int): Number of convolution channels. - kernel_size (int): Size of the convolution kernel. Default is 3. - dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5). - activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Default is None. - """ +class MMAudioAMPBlock(nn.Module): + """Anti-aliased multi-periodicity block: dilated convolutions each followed by a dilation-1 convolution.""" - def __init__( - self, - h: AttrDict, - channels: int, - kernel_size: int = 3, - dilation: tuple = (1, 3, 5), - activation: str = None, - ): + def __init__(self, channels: int, kernel_size: int, dilations: tuple[int, ...]) -> None: super().__init__() - - self.h = h - self.convs1 = nn.ModuleList( [ - weight_norm( - Conv1d( - channels, - channels, - kernel_size, - stride=1, - dilation=d, - padding=get_padding(kernel_size, d), - ) + nn.Conv1d( + channels, + channels, + kernel_size, + dilation=dilation, + padding=(kernel_size * dilation - dilation) // 2, ) - for d in dilation + for dilation in dilations ] ) - self.convs1.apply(init_weights) - self.convs2 = nn.ModuleList( - [ - weight_norm( - Conv1d( - channels, - channels, - kernel_size, - stride=1, - dilation=1, - padding=get_padding(kernel_size, 1), - ) - ) - for _ in range(len(dilation)) - ] + [nn.Conv1d(channels, channels, kernel_size, padding=(kernel_size - 1) // 2) for _ in dilations] ) - self.convs2.apply(init_weights) - - self.num_layers = len(self.convs1) + len(self.convs2) # Total number of conv layers - - # Activation functions - if activation == "snake": - self.activations = nn.ModuleList( - [ - TorchActivation1d(activation=Snake(channels, alpha_logscale=h.snake_logscale)) - for _ in range(self.num_layers) - ] - ) - elif activation == "snakebeta": - self.activations = nn.ModuleList( - [ - TorchActivation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale)) - for _ in range(self.num_layers) - ] - ) - else: - raise NotImplementedError( - "activation incorrectly specified. check the config file and look for 'activation'." - ) - - def forward(self, x): - acts1, acts2 = self.activations[::2], self.activations[1::2] - for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2): - xt = a1(x) - xt = c1(xt) - xt = a2(xt) - xt = c2(xt) - x = xt + x + self.activations = nn.ModuleList([MMAudioActivation1d(channels) for _ in range(2 * len(dilations))]) + def forward(self, x: torch.Tensor) -> torch.Tensor: + activations_1, activations_2 = self.activations[::2], self.activations[1::2] + for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, activations_1, activations_2): + x = conv2(act2(conv1(act1(x)))) + x return x - def remove_weight_norm(self): - for l in self.convs1: - remove_parametrizations(l, "weight") - for l in self.convs2: - remove_parametrizations(l, "weight") - - -class AMPBlock2(torch.nn.Module): - """ - AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer. - Unlike AMPBlock1, AMPBlock2 does not contain extra Conv1d layers with fixed dilation=1 - Args: - h (AttrDict): Hyperparameters. - channels (int): Number of convolution channels. - kernel_size (int): Size of the convolution kernel. Default is 3. - dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5). - activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Default is None. - """ +class MMAudioBigVGAN(nn.Module): + """BigVGAN-v2 vocoder turning a mel spectrogram into a waveform.""" def __init__( self, - h: AttrDict, - channels: int, - kernel_size: int = 3, - dilation: tuple = (1, 3, 5), - activation: str = None, - ): - super().__init__() - - self.h = h - - self.convs = nn.ModuleList( - [ - weight_norm( - Conv1d( - channels, - channels, - kernel_size, - stride=1, - dilation=d, - padding=get_padding(kernel_size, d), - ) - ) - for d in dilation - ] - ) - self.convs.apply(init_weights) - - self.num_layers = len(self.convs) # Total number of conv layers - - # Activation functions - if activation == "snake": - self.activations = nn.ModuleList( - [ - TorchActivation1d(activation=Snake(channels, alpha_logscale=h.snake_logscale)) - for _ in range(self.num_layers) - ] - ) - elif activation == "snakebeta": - self.activations = nn.ModuleList( - [ - TorchActivation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale)) - for _ in range(self.num_layers) - ] - ) - else: - raise NotImplementedError( - "activation incorrectly specified. check the config file and look for 'activation'." - ) - - def forward(self, x): - for c, a in zip(self.convs, self.activations): - xt = a(x) - xt = c(xt) - x = xt + x - return x - - def remove_weight_norm(self): - for l in self.convs: - remove_parametrizations(l, "weight") - - -TorchActivation1d = Activation1d - - -class BigVGANv2(torch.nn.Module): - """ - BigVGAN is a neural vocoder model that applies anti-aliased periodic activation for residual blocks (resblocks). - - Args: - h (AttrDict): Hyperparameters. - - Note: - - Ensure that the activation function is correctly specified in the hyperparameters (h.activation). - """ - - def __init__(self, h: AttrDict): + num_mels: int, + upsample_initial_channel: int, + upsample_rates: tuple[int, ...], + upsample_kernel_sizes: tuple[int, ...], + resblock_kernel_sizes: tuple[int, ...], + resblock_dilation_sizes: tuple[tuple[int, ...], ...], + ) -> None: super().__init__() - self.h = h - - self.num_kernels = len(h.resblock_kernel_sizes) - self.num_upsamples = len(h.upsample_rates) - - # Pre-conv - self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)) - - # Define which AMPBlock to use. BigVGAN uses AMPBlock1 as default - if h.resblock == "1": - resblock_class = AMPBlock1 - elif h.resblock == "2": - resblock_class = AMPBlock2 - else: - raise ValueError(f"Incorrect resblock class specified in hyperparameters. Got {h.resblock}") + self.num_kernels = len(resblock_kernel_sizes) + self.conv_pre = nn.Conv1d(num_mels, upsample_initial_channel, 7, padding=3) - # Transposed conv-based upsamplers. does not apply anti-aliasing self.ups = nn.ModuleList() - for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)): + self.resblocks = nn.ModuleList() + channels = upsample_initial_channel + for rate, kernel_size in zip(upsample_rates, upsample_kernel_sizes): self.ups.append( nn.ModuleList( - [ - weight_norm( - ConvTranspose1d( - h.upsample_initial_channel // (2**i), - h.upsample_initial_channel // (2 ** (i + 1)), - k, - u, - padding=(k - u) // 2, - ) - ) - ] + [nn.ConvTranspose1d(channels, channels // 2, kernel_size, rate, padding=(kernel_size - rate) // 2)] ) ) + channels //= 2 + for block_kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes): + self.resblocks.append(MMAudioAMPBlock(channels, block_kernel_size, tuple(dilations))) - # Residual blocks using anti-aliased multi-periodicity composition modules (AMP) - self.resblocks = nn.ModuleList() - for i in range(len(self.ups)): - ch = h.upsample_initial_channel // (2 ** (i + 1)) - for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)): - self.resblocks.append(resblock_class(h, ch, k, d, activation=h.activation)) - - # Post-conv - activation_post = ( - Snake(ch, alpha_logscale=h.snake_logscale) - if h.activation == "snake" - else (SnakeBeta(ch, alpha_logscale=h.snake_logscale) if h.activation == "snakebeta" else None) - ) - if activation_post is None: - raise NotImplementedError( - "activation incorrectly specified. check the config file and look for 'activation'." - ) - - self.activation_post = TorchActivation1d(activation=activation_post) - - # Whether to use bias for the final conv_post. Default to True for backward compatibility - self.use_bias_at_final = h.get("use_bias_at_final", True) - self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final)) - - # Weight initialization - for i in range(len(self.ups)): - self.ups[i].apply(init_weights) - self.conv_post.apply(init_weights) - - # Final tanh activation. Defaults to True for backward compatibility - self.use_tanh_at_final = h.get("use_tanh_at_final", True) - - def forward(self, x): - # Pre-conv - x = self.conv_pre(x) - - for i in range(self.num_upsamples): - # Upsampling - for i_up in range(len(self.ups[i])): - x = self.ups[i][i_up](x) - # AMP blocks - xs = None - for j in range(self.num_kernels): - if xs is None: - xs = self.resblocks[i * self.num_kernels + j](x) - else: - xs += self.resblocks[i * self.num_kernels + j](x) - x = xs / self.num_kernels - - # Post-conv - x = self.activation_post(x) - x = self.conv_post(x) - # Final tanh activation - if self.use_tanh_at_final: - x = torch.tanh(x) - else: - x = torch.clamp(x, min=-1.0, max=1.0) # Bound the output to [-1, 1] - - return x - - def remove_weight_norm(self): - try: - # print("Removing weight norm...") - for l in self.ups: - for l_i in l: - remove_parametrizations(l_i, "weight") - for l in self.resblocks: - l.remove_weight_norm() - remove_parametrizations(self.conv_pre, "weight") - remove_parametrizations(self.conv_post, "weight") - except ValueError: - print("[INFO] Model already removed weight norm. Skipping!") - pass + self.activation_post = MMAudioActivation1d(channels) + self.conv_post = nn.Conv1d(channels, 1, 7, padding=3, bias=False) + def forward(self, mel: torch.Tensor) -> torch.Tensor: + hidden_states = self.conv_pre(mel) + for stage, up in enumerate(self.ups): + hidden_states = up[0](hidden_states) + blocks = self.resblocks[stage * self.num_kernels : (stage + 1) * self.num_kernels] + hidden_states = sum(block(hidden_states) for block in blocks) / self.num_kernels + hidden_states = self.conv_post(self.activation_post(hidden_states)) + return torch.clamp(hidden_states, min=-1.0, max=1.0) -def build_bigvgan_v2(vocoder_config: dict | AttrDict) -> BigVGANv2: - """Build an inference-ready BigVGAN-v2 from its serialized configuration.""" - model = BigVGANv2(vocoder_config if isinstance(vocoder_config, AttrDict) else AttrDict(vocoder_config)) - model.remove_weight_norm() - return model +class MMAudioMelSpectrogram(nn.Module): + """Log-mel front end of the encoder. The filterbank and window are buffers so they follow the model's device.""" -# Reference: # https://github.com/bytedance/Make-An-Audio-2 - - -def dynamic_range_compression_torch(x, C=1, clip_val=1e-5, *, norm_fn): - return norm_fn(torch.clamp(x, min=clip_val) * C) - - -def spectral_normalize_torch(magnitudes, norm_fn): - output = dynamic_range_compression_torch(magnitudes, norm_fn=norm_fn) - return output - - -class MelConverter(nn.Module): - def __init__( - self, - *, - sampling_rate: float, - n_fft: int, - num_mels: int, - hop_size: int, - win_size: int, - fmin: float, - fmax: float, - norm_fn, - ): + def __init__(self, sample_rate: int, n_fft: int, num_mels: int, hop_length: int) -> None: super().__init__() - self.sampling_rate = sampling_rate self.n_fft = n_fft - self.num_mels = num_mels - self.hop_size = hop_size - self.win_size = win_size - self.fmin = fmin - self.fmax = fmax - self.norm_fn = norm_fn - - try: - from librosa.filters import mel as librosa_mel_fn - except (ImportError, OSError) as exc: - raise RuntimeError( - "MMAudio audio features require librosa. Install it with `pip install librosa`." - ) from exc - - mel = librosa_mel_fn( - sr=self.sampling_rate, - n_fft=self.n_fft, - n_mels=self.num_mels, - fmin=self.fmin, - fmax=self.fmax, - ) - mel_basis = torch.from_numpy(mel).float() - hann_window = torch.hann_window(self.win_size) - - self.register_buffer("mel_basis", mel_basis) - self.register_buffer("hann_window", hann_window) - - @property - def device(self): - return self.mel_basis.device - - def forward(self, waveform: torch.Tensor, center: bool = False) -> torch.Tensor: - waveform = waveform.clamp(min=-1.0, max=1.0).to(self.device) - - waveform = torch.nn.functional.pad( - waveform.unsqueeze(1), - [int((self.n_fft - self.hop_size) / 2), int((self.n_fft - self.hop_size) / 2)], - mode="reflect", - ) - waveform = waveform.squeeze(1) - - spec = torch.stft( + self.hop_length = hop_length + self.register_buffer("mel_basis", mel_filterbank(sample_rate, n_fft, num_mels, 0.0, sample_rate / 2)) + self.register_buffer("hann_window", torch.hann_window(n_fft)) + + def forward(self, waveform: torch.Tensor) -> torch.Tensor: + waveform = waveform.clamp(min=-1.0, max=1.0) + padding = (self.n_fft - self.hop_length) // 2 + waveform = F.pad(waveform.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1) + spectrum = torch.stft( waveform, self.n_fft, - hop_length=self.hop_size, - win_length=self.win_size, + hop_length=self.hop_length, + win_length=self.n_fft, window=self.hann_window, - center=center, + center=False, pad_mode="reflect", normalized=False, onesided=True, return_complex=True, ) - - spec = torch.view_as_real(spec) - spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)).float() - spec = torch.matmul(self.mel_basis, spec) - spec = spectral_normalize_torch(spec, self.norm_fn) - - return spec - - -def get_mel_converter(mode: Literal["44k"]) -> MelConverter: - """Create the 44.1 kHz mel-spectrogram converter used by MMAudio.""" - return MelConverter( - sampling_rate=44_100, - n_fft=2048, - num_mels=128, - hop_size=512, - win_size=2048, - fmin=0, - fmax=44100 / 2, - norm_fn=torch.log, - ) - - -# Audio feature / VAE facade. + magnitude = torch.sqrt(torch.view_as_real(spectrum).pow(2).sum(-1) + 1e-9).float() + return torch.log(torch.clamp(torch.matmul(self.mel_basis, magnitude), min=1e-5)) class MMAudioVAE(ModelMixin, ConfigMixin): - """Audio VAE and vocoder used to encode and decode audio latents. + r""" + Audio VAE of [`Kandinsky6TI2VAPipeline`]: a magnitude-preserving autoencoder over log-mel spectrograms (MMAudio, + https://arxiv.org/abs/2412.15322) followed by a BigVGAN-v2 vocoder. - The component converts waveforms to mel-spectrograms, encodes them into - latent representations, and decodes latents back to waveforms through a - BigVGAN vocoder. + `encode` turns a waveform into a latent distribution; `decode` turns latents into a waveform through the mel + decoder and the vocoder. One latent frame covers `hop_length * 2` samples. Args: - vocoder_config (`dict`): BigVGAN configuration. - mode (`str`, *optional*, defaults to ``"44k"``): Audio operating mode. - scaling_factor (`float`, *optional*, defaults to 1.0): Latent scaling factor. - sample_rate (`int`, *optional*, defaults to 44100): Audio sample rate. - downsample_factor (`int`, *optional*, defaults to 1024): Audio latent downsampling factor. + mel_bins (`int`, defaults to `128`): + Number of mel bins. + latent_channels (`int`, defaults to `40`): + Number of latent channels. + hidden_channels (`int`, defaults to `512`): + Base width of the autoencoder. + channel_multipliers (`tuple[int, ...]`, defaults to `(1, 2, 4)`): + Width multipliers of the autoencoder levels. + layers_per_block (`int`, defaults to `2`): + Residual blocks per encoder level; the decoder uses one more per level. + sample_rate (`int`, defaults to `44100`): + Waveform sample rate. + n_fft (`int`, defaults to `2048`): + FFT size of the mel front end. + hop_length (`int`, defaults to `512`): + Hop length of the mel front end, also the vocoder's total upsampling factor. + vocoder_upsample_initial_channel (`int`, defaults to `1536`): + Width of the vocoder's first layer. + vocoder_upsample_rates (`tuple[int, ...]`, defaults to `(8, 4, 2, 2, 2, 2)`): + Upsampling factors of the vocoder stages. + vocoder_upsample_kernel_sizes (`tuple[int, ...]`, defaults to `(16, 8, 4, 4, 4, 4)`): + Transposed-convolution kernel sizes of the vocoder stages. + vocoder_resblock_kernel_sizes (`tuple[int, ...]`, defaults to `(3, 7, 11)`): + Kernel sizes of the vocoder's residual blocks. + vocoder_resblock_dilation_sizes (`tuple[tuple[int, ...], ...]`, defaults to `((1, 3, 5), (1, 3, 5), (1, 3, 5))`): + Dilations of the vocoder's residual blocks. + scaling_factor (`float`, defaults to `0.417`): + Scale applied to the latents before they enter the diffusion transformer. """ - _no_split_modules = ["VAE", "BigVGANv2"] + _no_split_modules = ["MMAudioResnetBlock1D", "MMAudioAttnBlock1D", "MMAudioAMPBlock"] @register_to_config def __init__( self, - *, - vocoder_config: dict[str, Any], - mode: str = "44k", - scaling_factor: float = 1.0, + mel_bins: int = 128, + latent_channels: int = 40, + hidden_channels: int = 512, + channel_multipliers: tuple[int, ...] = (1, 2, 4), + layers_per_block: int = 2, sample_rate: int = 44_100, - downsample_factor: int = 1_024, + n_fft: int = 2048, + hop_length: int = 512, + vocoder_upsample_initial_channel: int = 1536, + vocoder_upsample_rates: tuple[int, ...] = (8, 4, 2, 2, 2, 2), + vocoder_upsample_kernel_sizes: tuple[int, ...] = (16, 8, 4, 4, 4, 4), + vocoder_resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11), + vocoder_resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)), + scaling_factor: float = 0.417, ) -> None: super().__init__() - self.mel_converter = get_mel_converter(mode) - self.vae: VAE = VAE(data_dim=128, embed_dim=40, hidden_dim=512).eval() - self.vae.remove_weight_norm() - self.vocoder = build_bigvgan_v2(vocoder_config).eval() - for param in self.parameters(): - param.requires_grad = False - self.scaling_factor = float(scaling_factor) - self.sample_rate = int(sample_rate) - self.downsample_factor = int(downsample_factor) - - def compile(self): - """Compile the decode and vocode methods with ``torch.compile``.""" - self.decode = torch.compile(self.decode) - self.vocode = torch.compile(self.vocode) - - def train(self, mode: bool = True) -> MMAudioVAE: - """Keep the inference-only audio component in evaluation mode.""" - return super().train(False) - - def encode_audio(self, x) -> DiagonalGaussianDistribution: - """Encode a waveform into an audio latent distribution.""" - mel = self.mel_converter(x) - return self.vae.encode(mel) - - def vocode(self, mel: torch.Tensor) -> torch.Tensor: - """Convert a decoded mel-spectrogram into a waveform.""" - return self.vocoder(mel) - - def decode(self, z: torch.Tensor) -> torch.Tensor: - """Decode audio latents into a mel-spectrogram.""" - return self.vae.decode(z) + if math.prod(vocoder_upsample_rates) != hop_length: + raise ValueError("the vocoder upsampling rates must multiply to `hop_length`") - @property - def device(self): - return next(self.parameters()).device - - @property - def dtype(self): - return next(self.parameters()).dtype + self.mel_converter = MMAudioMelSpectrogram(sample_rate, n_fft, mel_bins, hop_length) + self.vae = MMAudioAutoencoder( + mel_bins, latent_channels, hidden_channels, channel_multipliers, layers_per_block + ) + self.vocoder = MMAudioBigVGAN( + num_mels=mel_bins, + upsample_initial_channel=vocoder_upsample_initial_channel, + upsample_rates=vocoder_upsample_rates, + upsample_kernel_sizes=vocoder_upsample_kernel_sizes, + resblock_kernel_sizes=vocoder_resblock_kernel_sizes, + resblock_dilation_sizes=vocoder_resblock_dilation_sizes, + ) + # The encoder downsamples the mel frames once by 2. + self.latent_hop_length = hop_length * 2 @apply_forward_hook - def wrapped_decode(self, z): - """Decode latents and vocode them through the Diffusers forward hook.""" - mel_decoded = self.decode(z.to(dtype=self.dtype)) - return self.vocode(mel_decoded) + def encode(self, audio: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutput | tuple: + r""" + Encode a waveform into its latent distribution. + + Args: + audio (`torch.Tensor` of shape `(batch_size, num_samples)`): + Mono waveform in `[-1, 1]` at `sample_rate`. + return_dict (`bool`, defaults to `True`): + Whether to return an [`~models.modeling_outputs.AutoencoderKLOutput`] instead of a plain tuple. + """ + mel = self.mel_converter(audio).to(self.dtype) + posterior = DiagonalGaussianDistribution(self.vae.encode(mel)) + if not return_dict: + return (posterior,) + return AutoencoderKLOutput(latent_dist=posterior) @apply_forward_hook - def wrapped_encode(self, audio): - """Encode audio and return the mean latent through the forward hook.""" - dist = self.encode_audio(audio.to(dtype=self.dtype)) - return dist.mean + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple: + r""" + Decode latents into a waveform. + + Args: + z (`torch.Tensor` of shape `(batch_size, latent_channels, num_latent_frames)`): + Latents, already divided by `scaling_factor`. + return_dict (`bool`, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.DecoderOutput`] instead of a plain tuple. + + Returns: + The waveform of shape `(batch_size, 1, num_samples)` in `[-1, 1]`. + """ + waveform = self.vocoder(self.vae.decode(z)) + if not return_dict: + return (waveform,) + return DecoderOutput(sample=waveform) + + def forward( + self, + sample: torch.Tensor, + sample_posterior: bool = False, + return_dict: bool = True, + generator: torch.Generator | None = None, + ) -> DecoderOutput | tuple: + posterior = self.encode(sample).latent_dist + z = posterior.sample(generator=generator) if sample_posterior else posterior.mode() + waveform = self.decode(z).sample + if not return_dict: + return (waveform,) + return DecoderOutput(sample=waveform) __all__ = ["MMAudioVAE"] diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index acf8f8be3fa4..e6c425d36488 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -44,9 +44,9 @@ from .transformer_ideogram4 import Ideogram4Transformer2DModel from .transformer_joyimage import JoyImageEditTransformer3DModel from .transformer_joyimage_edit_plus import JoyImageEditPlusTransformer3DModel - from .transformer_kandinsky6_sr import Kandinsky6SRTransformer3DModel - from .transformer_kandinsky6 import Kandinsky6Transformer3DModel from .transformer_kandinsky import Kandinsky5Transformer3DModel + from .transformer_kandinsky6 import Kandinsky6Transformer3DModel + from .transformer_kandinsky6_sr import Kandinsky6SRTransformer3DModel 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_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 242fdfd692aa..fb3383cc5296 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -17,10 +17,10 @@ from __future__ import annotations import math +from dataclasses import dataclass from typing import Any import torch -import torch.nn.functional as functional from torch import Tensor, nn from torch.nn.attention.flex_attention import BlockMask @@ -29,16 +29,28 @@ from ...hooks.hooks import HookRegistry, ModelHook, StateManager from ...hooks.mag_cache import MagCacheState from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ...utils import BaseOutput from ..attention import AttentionMixin, AttentionModuleMixin -from ..attention_dispatch import ( - _CAN_USE_FLEX_ATTN, - AttentionBackendName, - dispatch_attention_fn, -) +from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn from ..cache_utils import CacheMixin from ..embeddings import get_timestep_embedding -from ..modeling_outputs import Transformer2DModelOutput -from ..modeling_utils import ModelMixin +from ..modeling_utils import ModelMixin, get_parameter_dtype + + +@dataclass +class Kandinsky6Transformer3DModelOutput(BaseOutput): + r""" + Output of [`Kandinsky6Transformer3DModel`]. + + Args: + sample (`torch.Tensor` of shape `(batch_size, num_frames, height, width, out_visual_dim)`): + The predicted video velocity, in the unpatchified `(B, T, H, W, C)` latent layout. + audio_sample (`torch.Tensor` of shape `(batch_size, audio_length, out_audio_dim)`, *optional*): + The predicted audio velocity. `None` when no `audio_hidden_states` were passed. + """ + + sample: torch.Tensor + audio_sample: torch.Tensor | None = None def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: @@ -148,10 +160,6 @@ def forward(self, pos: 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 reset_parameters(self) -> None: - freq = get_freqs(self.dim // 2, self.max_period).to(self.args.device) * self.freqs_scaling - self.args = torch.outer(torch.arange(self.max_pos, dtype=freq.dtype, device=freq.device), freq) - class RoPE3D(nn.Module): """3-D Rotary Position Embedding — used for video spatial-temporal tokens (T, H, W).""" @@ -172,20 +180,20 @@ def __init__( def forward( self, - shape: tuple, - pos: list[Tensor], + pos: tuple[Tensor, Tensor, Tensor], scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), ) -> Tensor: - T, H, W = shape - args_t = getattr(self, "args_0")[pos[0]] / scale_factor[0] # (T, d//2) - args_h = getattr(self, "args_1")[pos[1]] / scale_factor[1] # (H, d//2) - args_w = getattr(self, "args_2")[pos[2]] / scale_factor[2] # (W, d//2) + # `pos` holds one index tensor per axis; the grid size is implied by their lengths. + num_frames, height, width = (int(axis_pos.shape[0]) for axis_pos in pos) + args_t = self.args_0[pos[0]] / scale_factor[0] # (T, d//2) + args_h = self.args_1[pos[1]] / scale_factor[1] # (H, d//2) + args_w = self.args_2[pos[2]] / scale_factor[2] # (W, d//2) args = torch.cat( [ - args_t.view(T, 1, 1, -1).expand(T, H, W, -1), - args_h.view(1, H, 1, -1).expand(T, H, W, -1), - args_w.view(1, 1, W, -1).expand(T, H, W, -1), + args_t.view(num_frames, 1, 1, -1).expand(num_frames, height, width, -1), + args_h.view(1, height, 1, -1).expand(num_frames, height, width, -1), + args_w.view(1, 1, width, -1).expand(num_frames, height, width, -1), ], dim=-1, ) @@ -194,11 +202,6 @@ def forward( rope = rope.view(*rope.shape[:-1], 2, 2) # (T, H, W, total_dim, 2, 2) return rope.unsqueeze(-4) # (T, H, W, 1, total_dim, 2, 2) - def reset_parameters(self) -> None: - for i, (d, mp) in enumerate(zip(self.axes_dims, self.max_pos)): - freq = get_freqs(d // 2, self.max_period).to(getattr(self, f"args_{i}").device) - setattr(self, f"args_{i}", torch.outer(torch.arange(mp, dtype=freq.dtype, device=freq.device), freq)) - # Diffusers MagCache adapter for the K6 multimodal transformer. # @@ -438,8 +441,6 @@ class Kandinsky6AttnProcessor: _parallel_config = None def __init__(self, attention_backend=None, parallel_config=None): - if not hasattr(functional, "scaled_dot_product_attention"): - raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0 or newer.") self._masked = False self._attention_backend = attention_backend self._parallel_config = parallel_config @@ -486,26 +487,21 @@ def __call__( key = apply_rotary(key, rotary_emb_kv).to(dtype=key.dtype) if sparse_params is not None: - q = query.transpose(1, 2).contiguous() - k = key.transpose(1, 2).contiguous() - v = value.transpose(1, 2).contiguous() + # The block statistics are computed from the `(B, heads, S, D)` layout the mask builder expects; the + # resulting `BlockMask` only runs on the flex backend. block_mask = nabla_block_mask( - q, - k, + query.transpose(1, 2), + key.transpose(1, 2), sparse_params["sta_mask"], thr=sparse_params["P"], ) - if not _CAN_USE_FLEX_ATTN: - raise ValueError("Nabla attention requires PyTorch 2.5 or newer") - output = ( - torch.nn.attention.flex_attention.flex_attention( - q, - k, - v, - block_mask=block_mask, - ) - .transpose(1, 2) - .contiguous() + output = dispatch_attention_fn( + query, + key, + value, + attn_mask=block_mask, + backend=AttentionBackendName.FLEX, + parallel_config=self._parallel_config, ) else: output = dispatch_attention_fn( @@ -533,12 +529,14 @@ def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): self.activation = nn.SiLU() self.out_layer = nn.Linear(time_dim, time_dim) - def forward(self, time: Tensor) -> Tensor: + def forward(self, timestep: Tensor) -> Tensor: + # The sinusoidal embedding is float32; `_keep_in_fp32_modules` keeps these layers float32 under + # `from_pretrained(torch_dtype=...)`, and the cast aligns the input with whatever dtype they hold. embed = get_timestep_embedding( - time, self.model_dim, flip_sin_to_cos=True, downscale_freq_shift=0, max_period=self.max_period + timestep, self.model_dim, flip_sin_to_cos=True, downscale_freq_shift=0, max_period=self.max_period ) - h = functional.linear(embed, self.in_layer.weight.float(), self.in_layer.bias.float()) - return functional.linear(self.activation(h), self.out_layer.weight.float(), self.out_layer.bias.float()) + embed = embed.to(get_parameter_dtype(self.in_layer)) + return self.out_layer(self.activation(self.in_layer(embed))) class Kandinsky6TextEmbeddings(nn.Module): @@ -583,8 +581,7 @@ def __init__(self, time_dim: int, model_dim: int, num_params: int): nn.init.zeros_(self.out_layer.bias) def forward(self, x: Tensor) -> Tensor: - out = functional.linear(self.activation(x.float()), self.out_layer.weight.float(), self.out_layer.bias.float()) - return out.to(dtype=x.dtype) + return self.out_layer(self.activation(x.to(get_parameter_dtype(self.out_layer)))) class Kandinsky6FeedForward(nn.Module): @@ -704,6 +701,8 @@ def __init__(self, model_dim: int, time_dim: int, audio_dim: int): def forward(self, audio_embed: Tensor, time_embed: Tensor) -> Tensor: shift, scale = torch.chunk(self.modulation(time_embed), 2, dim=-1) x = apply_scale_shift_norm(self.norm, audio_embed, scale, shift) + # The reference audio head normalizes a second time after the AdaLN affine; the released weights were + # trained this way, so the double norm is kept for parity even though the video head has none. x = self.norm(x) return self.out_layer(x) @@ -748,8 +747,8 @@ class Kandinsky6TransformerDecoderBlock(nn.Module): """Visual self-attention, text cross-attention, and feed-forward submodules. `Kandinsky6FusedTransformerDecoderBlock` uses this class only as a named submodule container - (`self_attention`/`cross_attention`/`feed_forward` and their norms/modulation), calling those - submodules directly rather than this class's own `forward`. + (`self_attention`/`cross_attention`/`feed_forward` and their norms/modulation), calling those submodules directly + rather than this class's own `forward`. """ def __init__( @@ -952,11 +951,10 @@ class Kandinsky6Transformer3DModel( ): """Kandinsky 6 multimodal transformer for text/image-to-video-and-audio generation. - Video and audio are denoised together through fused, cross-modal transformer blocks, each - conditioned on its own text branch (Qwen2.5-VL tokens + CLIP pooled embedding). Passing only - `x_video` or only `x_audio` denoises a single modality while still running through the fused - block's per-modality self-attention, cross-attention, and feed-forward stages. Inputs and - outputs use packed token layouts compatible with the native K6 model. + Video and audio are denoised together through fused, cross-modal transformer blocks, each conditioned on its own + text branch (Qwen2.5-VL tokens + CLIP pooled embedding). Passing no `audio_hidden_states` denoises video alone + while still running the fused block's video self-attention, cross-attention, and feed-forward stages. Rotary + embeddings are computed inside `forward` from the token grid. Args: in_visual_dim (`int`, *optional*, defaults to 16): Number of input video latent channels. @@ -979,7 +977,8 @@ class Kandinsky6Transformer3DModel( axes_dims_a (`tuple[int, int, int]`, *optional*): Audio RoPE dimensions. Defaults to `axes_dims`. audio_freqs_scaling (`float`, *optional*, defaults to 1.0): Audio RoPE frequency scaling. text_token_padding (`bool`, *optional*, defaults to False): Whether text sequences are padded. - scale_factor (`tuple[float]`, *optional*): Per-axis RoPE frequency scaling. + scale_factor (`tuple[float, float, float]`, *optional*, defaults to `(1.0, 2.0, 2.0)`): Per-axis + `(t, h, w)` RoPE frequency scaling applied to the video positions. ca_rope (`bool`, *optional*, defaults to False): Whether to use cross-modal audio RoPE. cross_gates (`bool`, *optional*, defaults to False): Whether to use cross-modal residual gates. fix_modulation (`bool`, *optional*, defaults to False): Whether to use the fixed modulation variant. @@ -992,15 +991,11 @@ class Kandinsky6Transformer3DModel( "Kandinsky6FusedTransformerDecoderBlock", ] _no_split_modules = _repeated_blocks + _skip_layerwise_casting_patterns = ["norm"] + # Timestep embeddings and every AdaLN modulation projection (`*_modulation`, `out_layer.modulation`) are + # precision-sensitive and stay float32 under `from_pretrained(torch_dtype=...)`. _keep_in_fp32_modules = ["time_embeddings", "modulation"] _supports_gradient_checkpointing = True - # `Kandinsky6TimeEmbeddings`/`Kandinsky6Modulation` hand their `nn.Linear` weight/bias straight to - # `functional.linear` (upcast to fp32 first) rather than calling the submodule, so the AdaLN math stays - # fp32 even after a later blanket `.to(bfloat16)` cast (`_keep_in_fp32_modules` only protects the - # `from_pretrained(dtype=...)` load, not a subsequent cast). Leaf-level offload hooks only fire on an - # actual submodule call, so they never see these weights — the same class of gap as - # `HunyuanDiTAttentionPool` (see testing.md), which opts out of group offloading for the same reason. - _supports_group_offloading = False @register_to_config def __init__( @@ -1099,97 +1094,97 @@ def __init__( def forward( self, - x_video: Tensor | None = None, - x_audio: Tensor | None = None, - text_embed: Tensor | None = None, - pooled_text_embed: Tensor | None = None, - time: Tensor | None = None, - visual_rope: Tensor | None = None, - audio_rope: Tensor | None = None, - video_text_rope: Tensor | None = None, - audio_text_rope: Tensor | None = None, - sparse_params: dict | None = None, - attention_mask: Tensor | None = None, + hidden_states: Tensor, + encoder_hidden_states: Tensor, + pooled_projections: Tensor, + timestep: Tensor, + audio_hidden_states: Tensor | None = None, + visual_rope_pos: tuple[Tensor, Tensor, Tensor] | None = None, + encoder_attention_mask: Tensor | None = None, visual_token_type_ids: Tensor | None = None, - return_dict: bool = False, - **kwargs: Any, - ) -> Tensor | tuple[Tensor, Tensor] | Transformer2DModelOutput: - """Run the K6 transformer on video, audio, or fused video/audio latents. - + sparse_params: dict | None = None, + return_dict: bool = True, + ) -> Kandinsky6Transformer3DModelOutput | tuple[Tensor, ...]: + r""" Args: - x_video (`torch.Tensor`, *optional*): Packed video latent tokens. - x_audio (`torch.Tensor`, *optional*): Packed audio latent tokens. - text_embed (`torch.Tensor`): Text token embeddings, shared by the video and audio text branches. - pooled_text_embed (`torch.Tensor`): Pooled text embedding, shared by the video and audio text branches. - time (`torch.Tensor`): Diffusion timestep, shared by the video and audio branches. - visual_rope (`torch.Tensor`, *optional*): Video rotary position embeddings. Required with `x_video`. - audio_rope (`torch.Tensor`, *optional*): Audio rotary position embeddings. Required with `x_audio`. - video_text_rope (`torch.Tensor`): Rotary position embeddings for the video text branch. - audio_text_rope (`torch.Tensor`): Rotary position embeddings for the audio text branch. - sparse_params (`dict`, *optional*): NABLA sparse-attention configuration for the video self-attention. - attention_mask (`torch.Tensor`, *optional*): Text attention mask. - visual_token_type_ids (`torch.Tensor`, *optional*): Video token type IDs. - return_dict (`bool`, *optional*, defaults to False): Whether to return a - [`Transformer2DModelOutput`] instead of tensors. + hidden_states (`torch.Tensor` of shape `(batch_size, num_frames, height, width, in_channels)`): + Video latents in the `(B, T, H, W, C)` layout. With `visual_cond=True`, `in_channels` is `2 * + in_visual_dim + 1`: the noisy latent, the conditioning latent and a conditioning mask. + encoder_hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, in_text_dim)`): + Text token embeddings, shared by the video and audio text branches. + pooled_projections (`torch.Tensor` of shape `(batch_size, in_text_dim2)`): + Pooled text embedding added to the timestep embedding of both branches. + timestep (`torch.Tensor` of shape `(batch_size,)`): + Diffusion timestep on the `[0, num_train_timesteps]` scale, shared by both modalities. + audio_hidden_states (`torch.Tensor` of shape `(batch_size, audio_length, in_audio_dim)`, *optional*): + Audio latents. When omitted the fused blocks run their video path only. + visual_rope_pos (`tuple[torch.Tensor, torch.Tensor, torch.Tensor]`, *optional*): + Per-axis `(t, h, w)` rotary position indices of the patchified video tokens. Defaults to `arange` over + each axis. The pipeline passes an explicit temporal index when it appends a reference frame that reuses + position `0`. + encoder_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Boolean padding mask for `encoder_hidden_states`. Pass `None` when no prompt is padded. + visual_token_type_ids (`torch.Tensor` of shape `(batch_size, num_frames)`, *optional*): + Per-frame token type ids, embedded through `visual_token_type_embeddings`. Requires + `visual_token_type_num_embeddings > 0`. + sparse_params (`dict`, *optional*): + NABLA sparse-attention configuration for the video self-attention. + return_dict (`bool`, defaults to `True`): + Whether to return a [`Kandinsky6Transformer3DModelOutput`] instead of a plain tuple. Returns: - `torch.Tensor`, `tuple[torch.Tensor, torch.Tensor]`, or - [`Transformer2DModelOutput`]: The denoised single modality, `(video, audio)` when both `x_video` - and `x_audio` are given, or a model output object. + [`Kandinsky6Transformer3DModelOutput`] or `tuple`: + The predicted video velocity in the `(B, T, H, W, out_visual_dim)` layout and, when + `audio_hidden_states` was given, the predicted audio velocity of shape `(B, audio_length, + out_audio_dim)`. """ - if x_video is None and x_audio is None: - raise ValueError("at least one of `x_video`, `x_audio` must be provided") - if text_embed is None or pooled_text_embed is None or time is None: - raise ValueError("`text_embed`, `pooled_text_embed`, and `time` are required") - if video_text_rope is None or audio_text_rope is None: - raise ValueError("`video_text_rope` and `audio_text_rope` are required") - - attn_mask = attention_mask - if attn_mask is not None and attn_mask.dim() == 1: - attn_mask = attn_mask.unsqueeze(0) + if visual_token_type_ids is not None and self.visual_token_type_num_embeddings == 0: + raise ValueError("`visual_token_type_ids` requires `visual_token_type_num_embeddings > 0`") checkpoint = torch.is_grad_enabled() and self.gradient_checkpointing + device = hidden_states.device + + # 1. Rotary embeddings for the two text branches + text_pos = torch.arange(encoder_hidden_states.shape[1], device=device) + video_text_rope = self.video_text_rope_embeddings(text_pos) + audio_text_rope = self.audio_text_rope_embeddings(text_pos) - # 1. Encode text tokens through the video and audio text branches - video_text_embed = self.video_text_embeddings(text_embed) - video_temb = self.video_time_embeddings(time) + self.video_pooled_text_embeddings(pooled_text_embed) + # 2. Encode the text tokens through the video and audio text branches + video_text_embed = self.video_text_embeddings(encoder_hidden_states) + video_temb = self.video_time_embeddings(timestep) + self.video_pooled_text_embeddings(pooled_projections) for block in self.video_text_transformer_blocks: - args = (video_text_embed, video_temb, video_text_rope, attn_mask) + args = (video_text_embed, video_temb, video_text_rope, encoder_attention_mask) video_text_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) - audio_text_embed = self.audio_text_embeddings(text_embed) - audio_temb = self.audio_time_embeddings(time) + self.audio_pooled_text_embeddings(pooled_text_embed) + audio_text_embed = self.audio_text_embeddings(encoder_hidden_states) + audio_temb = self.audio_time_embeddings(timestep) + self.audio_pooled_text_embeddings(pooled_projections) for block in self.audio_text_transformer_blocks: - args = (audio_text_embed, audio_temb, audio_text_rope, attn_mask) + args = (audio_text_embed, audio_temb, audio_text_rope, encoder_attention_mask) audio_text_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) - # 2. Patchify the video latents and embed the audio latents (only the modalities that were passed) - visual_embed, visual_shape = None, None + # 3. Patchify the video latents and build their rotary embeddings + visual_embed = self.visual_embeddings(hidden_states) + if visual_token_type_ids is not None: + token_types = self.visual_token_type_embeddings(visual_token_type_ids) + visual_embed = visual_embed + token_types[:, :, None, None, :] + visual_shape = visual_embed.shape[1:4] + if visual_rope_pos is None: + visual_rope_pos = tuple(torch.arange(size, device=device) for size in visual_shape) + visual_rope = self.visual_rope_embeddings(visual_rope_pos, self.scale_factor) to_fractal = bool(sparse_params and sparse_params.get("to_fractal")) - if x_video is not None: - if x_video.ndim == 4: - x_video = x_video.unsqueeze(0) - visual_embed = self.visual_embeddings(x_video) - if visual_token_type_ids is not None and hasattr(self, "visual_token_type_embeddings"): - token_type_ids = ( - visual_token_type_ids.unsqueeze(0) if visual_token_type_ids.ndim == 1 else visual_token_type_ids - ) - token_types = self.visual_token_type_embeddings(token_type_ids.to(device=visual_embed.device)) - visual_embed = visual_embed + token_types[:, :, None, None, :] - visual_shape = visual_embed.shape[-4:-1] - if to_fractal: - visual_embed = _local_patch(visual_embed, visual_shape, (1, 8, 8), dim=1).flatten(1, 2) - visual_rope = _local_patch(visual_rope, visual_shape, (1, 8, 8), dim=0).flatten(0, 1) - else: - visual_embed = visual_embed.flatten(1, 3) - visual_rope = visual_rope.flatten(0, 2) - - audio_embed = None - if x_audio is not None: - if x_audio.ndim == 2: - x_audio = x_audio.unsqueeze(0) - audio_embed = self.audio_embeddings(x_audio) - - # 3. Run the fused video/audio transformer blocks (a block no-ops on whichever modality is absent) + if to_fractal: + visual_embed = _local_patch(visual_embed, visual_shape, (1, 8, 8), dim=1).flatten(1, 2) + visual_rope = _local_patch(visual_rope, visual_shape, (1, 8, 8), dim=0).flatten(0, 1) + else: + visual_embed = visual_embed.flatten(1, 3) + visual_rope = visual_rope.flatten(0, 2) + + # 4. Embed the audio latents and build their rotary embeddings + audio_embed = audio_rope = None + if audio_hidden_states is not None: + audio_embed = self.audio_embeddings(audio_hidden_states) + audio_rope = self.audio_rope_embeddings(torch.arange(audio_hidden_states.shape[1], device=device)) + + # 5. Run the fused video/audio transformer blocks (a block skips the audio path when it is absent) for block in self.visual_transformer_blocks: args = ( visual_embed, @@ -1200,34 +1195,31 @@ def forward( visual_rope, audio_rope, sparse_params, - attn_mask, + encoder_attention_mask, ) visual_embed, audio_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) - # 4. Project back to patch/latent space - video_out = None - if visual_embed is not None: - if to_fractal: - visual_embed = _local_merge( - visual_embed.reshape(visual_embed.shape[0], -1, 64, visual_embed.shape[-1]), - visual_shape, - (1, 8, 8), - dim=1, - ) - else: - visual_embed = visual_embed.reshape(-1, *visual_shape, visual_embed.shape[-1]) - video_out = self.out_layer(visual_embed, video_temb) + # 6. Project back to the latent space + if to_fractal: + visual_embed = _local_merge( + visual_embed.reshape(visual_embed.shape[0], -1, 64, visual_embed.shape[-1]), + visual_shape, + (1, 8, 8), + dim=1, + ) + else: + visual_embed = visual_embed.reshape(-1, *visual_shape, visual_embed.shape[-1]) + video_out = self.out_layer(visual_embed, video_temb) audio_out = self.audio_out_layer(audio_embed, audio_temb) if audio_embed is not None else None - if video_out is not None and audio_out is not None: - result = (video_out, audio_out) - else: - result = video_out if video_out is not None else audio_out - return Transformer2DModelOutput(sample=result) if return_dict else result + if not return_dict: + return (video_out,) if audio_out is None else (video_out, audio_out) + return Kandinsky6Transformer3DModelOutput(sample=video_out, audio_sample=audio_out) __all__ = [ "Kandinsky6Transformer3DModel", + "Kandinsky6Transformer3DModelOutput", "Kandinsky6TransformerEncoderBlock", "Kandinsky6TransformerDecoderBlock", "Kandinsky6FusedTransformerDecoderBlock", diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index edab68e52872..cbc7095fe2de 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -12,1416 +12,552 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Kandinsky 6 SR transformer Diffusers component.""" +"""Kandinsky 6 video super-resolution transformer.""" from __future__ import annotations import math -from collections.abc import Mapping from typing import Any import torch from torch import Tensor, nn -from torch.nn.attention.flex_attention import BlockMask, flex_attention +from torch.nn.attention.flex_attention import BlockMask from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin from ..attention import AttentionMixin, AttentionModuleMixin -from ..attention_dispatch import dispatch_attention_fn -from ..modeling_utils import ModelMixin -from .transformer_kandinsky6 import Kandinsky6AttnProcessor +from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn +from ..embeddings import get_timestep_embedding +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin, get_parameter_dtype -class Kandinsky6SRAttentionProcessor(Kandinsky6AttnProcessor): - """Diffusers attention processor used by the SR transformer.""" - - def __call__( - self, - attn: Any, - query: Tensor, - key: Tensor, - value: Tensor, - cu_seqlens_q: Tensor, - cu_seqlens_k: Tensor, - *, - sparse_params: dict[str, Any] | None = None, - attn_mask: Tensor | None = None, - ) -> Tensor: - if sparse_params is not None: - return attn.attention_flex(query, key, value, sparse_params=sparse_params) - - if getattr(attn, "cached_k", None) is not None: - key, value, cu_seqlens_q, cu_seqlens_k = self.assemble_cached_attention_inputs( - key, value, attn.cached_k, attn.cached_v, cu_seqlens_q, attn.cached_cu_seqlens - ) - if getattr(attn, "return_kv", False): - attn.cached_k, attn.cached_v, attn.cached_cu_seqlens = key, value, cu_seqlens_k - - outputs = [] - batch_size = cu_seqlens_q.numel() - 1 - for index in range(batch_size): - query_start, query_end = int(cu_seqlens_q[index]), int(cu_seqlens_q[index + 1]) - key_start, key_end = int(cu_seqlens_k[index]), int(cu_seqlens_k[index + 1]) - outputs.append( - dispatch_attention_fn( - query[query_start:query_end].unsqueeze(0), - key[key_start:key_end].unsqueeze(0), - value[key_start:key_end].unsqueeze(0), - attn_mask=self._packed_attention_mask( - attn_mask, - index, - batch_size, - query_end - query_start, - key_end - key_start, - ), - backend=self._attention_backend, - parallel_config=self._parallel_config, - )[0] - ) - return torch.cat(outputs, dim=0).flatten(-2, -1) - - @staticmethod - def assemble_cached_attention_inputs( - new_key: Tensor, - new_value: Tensor, - cached_key: Tensor, - cached_value: Tensor, - cu_seqlens_new: Tensor, - cached_cu_seqlens: Tensor, - ) -> tuple[Tensor, Tensor, Tensor, Tensor]: - """Merge cached and new K/V into packed varlen attention inputs.""" - num_seqs = cu_seqlens_new.numel() - 1 - key_parts: list[Tensor] = [] - value_parts: list[Tensor] = [] - for index in range(num_seqs): - cache_start, cache_end = int(cached_cu_seqlens[index]), int(cached_cu_seqlens[index + 1]) - new_start, new_end = int(cu_seqlens_new[index]), int(cu_seqlens_new[index + 1]) - key_parts.append(cached_key[cache_start:cache_end]) - key_parts.append(new_key[new_start:new_end]) - value_parts.append(cached_value[cache_start:cache_end]) - value_parts.append(new_value[new_start:new_end]) - - key = torch.cat(key_parts, dim=0) - value = torch.cat(value_parts, dim=0) - merged_lens = torch.diff(cached_cu_seqlens) + torch.diff(cu_seqlens_new) - cu_seqlens_k = torch.cat([cu_seqlens_new.new_zeros(1), torch.cumsum(merged_lens, dim=0)]).to( - cu_seqlens_new.dtype - ) - return key, value, cu_seqlens_new, cu_seqlens_k - - @staticmethod - def _packed_attention_mask( - attn_mask: Tensor | None, - index: int, - batch_size: int, - query_length: int, - key_length: int, - ) -> Tensor | None: - if attn_mask is None: - return None - if attn_mask.ndim == 1: - return attn_mask[:key_length].unsqueeze(0) - if attn_mask.ndim == 2: - if attn_mask.shape[0] in (1, batch_size): - return attn_mask[index : index + 1, :key_length] - return attn_mask[:query_length, :key_length] - if attn_mask.shape[0] in (1, batch_size): - attn_mask = attn_mask[index : index + 1] - return attn_mask[..., :query_length, :key_length] - - -# Side of the local 8x8 token block used by fractal (NABLA) attention. -# Independent of the VAE spatial compression — do not swap for VAE_SPATIAL_FACTOR. +# Side of the local token block that NABLA sparse attention groups into one 64-token attention block. FRACTAL_BLOCK_SIZE = 8 +# Copied from diffusers.models.transformers.transformer_kandinsky6.get_freqs def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: - """Compute sinusoidal frequency schedule for rotary embeddings. + """Return inverse frequencies for rotary position embeddings. Args: - dim: Embedding dimension. - max_period: Maximum period of the sinusoidal frequencies. + dim (`int`): Number of frequency values to generate. + max_period (`float`, *optional*, defaults to 10000.0): Maximum period + used by the frequency schedule. Returns: - Frequency tensor of shape ``(dim,)``. + `torch.Tensor`: Frequency values in float32. """ return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) -def fractal_flatten( - x: Tensor, - rope: Tensor, - cu_seqlens: Tensor, - shape: tuple[int, int, int], - *, - fractal: bool = False, -) -> tuple[Tensor, Tensor, Tensor]: - """Flatten spatial dimensions with optional fractal (local block) patching. - - Args: - x: Input tensor with spatial dimensions. - rope: Rotary position embedding tensor matching ``x`` layout. - cu_seqlens: Cumulative sequence lengths (per-frame counts). - shape: Spatial shape as ``(length, height, width)``. - fractal: If True, apply local 8x8 block patching before flattening. +# Copied from diffusers.models.transformers.transformer_kandinsky6.apply_scale_shift_norm +def apply_scale_shift_norm(norm, x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: + """AdaLN-style affine in fp32, cast back to ``x.dtype``.""" + if x.ndim > 2 and scale.ndim == 2: + shape = (scale.shape[0],) + (1,) * (x.ndim - 2) + (scale.shape[-1],) + scale, shift = scale.reshape(shape), shift.reshape(shape) + return (norm(x.float()) * (scale.float() + 1.0) + shift.float()).to(dtype=x.dtype) - Returns: - Tuple of (flattened ``x``, flattened ``rope``, scaled ``cu_seqlens``). - """ - _, height, width = shape - if fractal: - block = FRACTAL_BLOCK_SIZE - x = local_patching(x, shape, (1, block, block), dim=0) - rope = local_patching(rope, shape, (1, block, block), dim=0) - x = x.flatten(0, 1) - rope = rope.flatten(0, 1) - else: - x = x.flatten(0, 2) - rope = rope.flatten(0, 2) - cu_seqlens = cu_seqlens * (height * width) - return x, rope, cu_seqlens - - -def fractal_unflatten( - x: Tensor, - cu_seqlens: Tensor, - shape: tuple[int, int, int], - *, - fractal: bool = False, -) -> tuple[Tensor, Tensor]: - """Unflatten spatial dimensions, inverse of ``fractal_flatten``. - Args: - x: Flattened input tensor. - cu_seqlens: Scaled cumulative sequence lengths. - shape: Target spatial shape as ``(length, height, width)``. - fractal: If True, reverse local 8x8 block patching. +# Copied from diffusers.models.transformers.transformer_kandinsky6.apply_gate_sum +def apply_gate_sum(x: Tensor, out: Tensor, gate: Tensor) -> Tensor: + """Residual gate in fp32, cast back to ``x.dtype``.""" + if x.ndim > 2 and gate.ndim == 2: + gate = gate.reshape((gate.shape[0],) + (1,) * (x.ndim - 2) + (gate.shape[-1],)) + return (x.float() + gate.float() * out.float()).to(dtype=x.dtype) - Returns: - Tuple of (unflattened ``x``, rescaled ``cu_seqlens``). - """ - _, height, width = shape - if fractal: - block = FRACTAL_BLOCK_SIZE - x = x.reshape(-1, block**2, *x.shape[1:]) - x = local_merge(x, shape, (1, block, block), dim=0) - else: - x = x.reshape(*shape, *x.shape[1:]) - cu_seqlens = cu_seqlens // (height * width) - return x, cu_seqlens +# Copied from diffusers.models.transformers.transformer_kandinsky6.apply_rotary +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 local_patching(x: Tensor, shape: tuple[int, int, int], group_size: tuple[int, int, int], dim: int = 0) -> Tensor: - """Rearrange tensor into local spatial patches. - - Groups neighboring elements along each spatial axis into patches, - producing a tensor with patch-count and intra-patch dimensions. - Args: - x: Input tensor with spatial dimensions starting at ``dim``. - shape: Spatial shape as ``(duration, height, width)``. - group_size: Patch size per axis ``(g1, g2, g3)``. - dim: First spatial dimension index. - - Returns: - Patched tensor with shape ``(..., num_patches, patch_elems, ...)``. - """ - duration, height, width = shape +# Copied from diffusers.models.transformers.transformer_kandinsky6._local_patch +def _local_patch(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Tensor: + T, H, W = shape g1, g2, g3 = group_size - x = x.reshape(*x.shape[:dim], duration // g1, g1, height // g2, g2, width // g3, g3, *x.shape[dim + 3 :]) - x = x.permute( - *range(len(x.shape[:dim])), dim, dim + 2, dim + 4, dim + 1, dim + 3, dim + 5, *range(dim + 6, len(x.shape)) - ) + x = x.reshape(*x.shape[:dim], T // g1, g1, H // g2, g2, W // g3, g3, *x.shape[dim + 3 :]) + d = len(x.shape[:dim]) + x = x.permute(*range(d), d, d + 2, d + 4, d + 1, d + 3, d + 5, *range(d + 6, len(x.shape))) return x.flatten(dim, dim + 2).flatten(dim + 1, dim + 3) -def local_merge(x: Tensor, shape: tuple[int, int, int], group_size: tuple[int, int, int], dim: int = 0) -> Tensor: - """Reverse local patching, restoring the original spatial layout. - - Args: - x: Patched tensor with ``(num_patches, patch_elems)`` at ``dim``. - shape: Original spatial shape as ``(duration, height, width)``. - group_size: Patch size per axis ``(g1, g2, g3)``. - dim: First spatial dimension index. - - Returns: - Tensor with restored spatial dimensions. - """ - duration, height, width = shape +# Copied from diffusers.models.transformers.transformer_kandinsky6._local_merge +def _local_merge(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Tensor: + T, H, W = shape g1, g2, g3 = group_size - x = x.reshape(*x.shape[:dim], duration // g1, height // g2, width // g3, g1, g2, g3, *x.shape[dim + 2 :]) - x = x.permute( - *range(len(x.shape[:dim])), dim, dim + 3, dim + 1, dim + 4, dim + 2, dim + 5, *range(dim + 6, len(x.shape)) - ) + x = x.reshape(*x.shape[:dim], T // g1, H // g2, W // g3, g1, g2, g3, *x.shape[dim + 2 :]) + d = len(x.shape[:dim]) + x = x.permute(*range(d), d, d + 3, d + 1, d + 4, d + 2, d + 5, *range(d + 6, len(x.shape))) return x.flatten(dim, dim + 1).flatten(dim + 1, dim + 2).flatten(dim + 2, dim + 3) -@torch.compile() -@torch.no_grad() -def fast_sta_nabla(T: int, H: int, W: int, wT: int = 3, wH: int = 3, wW: int = 3, device: str = "cuda") -> Tensor: - """Build a static spatiotemporal neighborhood attention mask. - - Each token attends to spatial and temporal neighbors within - the given window sizes. - - Args: - T: Temporal dimension (number of frames). - H: Height dimension. - W: Width dimension. - wT: Temporal window size (odd). - wH: Height window size (odd). - wW: Width window size (odd). - device: Device for the output tensor. - - Returns: - Boolean mask of shape ``(T*H*W, T*H*W)``. - """ - max_dim = max(T, H, W) - r = torch.arange(0, max_dim, 1, dtype=torch.int16, device=device) - mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs() - sta_t, sta_h, sta_w = mat[:T, :T].flatten(), mat[:H, :H].flatten(), mat[:W, :W].flatten() - sta_t = sta_t <= wT // 2 - sta_h = sta_h <= wH // 2 - sta_w = sta_w <= wW // 2 - sta_hw = (sta_h.unsqueeze(1) * sta_w.unsqueeze(0)).reshape(H, H, W, W).transpose(1, 2).flatten() - sta = (sta_t.unsqueeze(1) * sta_hw.unsqueeze(0)).reshape(T, T, H * W, H * W).transpose(1, 2) - return sta.reshape(T * H * W, T * H * W) - - -@torch.compile(dynamic=True) -@torch.no_grad() -def nablaT_v2_doc( +# Copied from diffusers.models.transformers.transformer_kandinsky6.nabla_block_mask +def nabla_block_mask( q: Tensor, k: Tensor, - seq: Tensor, - T: int, - H: int, - W: int, - *, - wT: int = 3, - wH: int = 3, - wW: int = 3, + sta: Tensor, thr: float = 0.9, - add_sta: bool = True, + block_size: int = 64, ) -> BlockMask: - """Build a dynamic sparse attention BlockMask with document boundaries. - - Estimates an approximate attention map from block-averaged queries and - keys, then binarizes it (CDF cutoff) to select the most relevant blocks per query. - - Args: - q: Query tensor of shape ``(B, heads, seq_len, dim)``. - k: Key tensor of shape ``(B, heads, seq_len, dim)``. - seq: Cumulative document lengths (e.g. [0, 31, 51, 66, 97] - video boundaries in seq_len). - T: Temporal dimension. - H: Height dimension. - W: Width dimension. - wT: Temporal window for static attention. - wH: Height window for static attention. - wW: Width window for static attention. - thr: CDF threshold for binarization. - add_sta: Whether to add static local attention. - - Returns: - Sparse ``BlockMask`` with block size 64. - """ - # Q/K are the authoritative execution tensors when this function is - # reached through attention. Keep every mask intermediate on their device. - device = q.device - seq = seq.to(device=device) - - # Map estimation + """Build a dynamic NABLA block mask from query/key statistics and an STA prior.""" B, h, S, D = q.shape - qa = q.reshape(B, h, S // 64, 64, D).mean(-2) - ka = k.reshape(B, h, S // 64, 64, D).mean(-2).transpose(-2, -1) - attn_map = qa @ ka - - d = torch.diff(seq) - doc = ( - torch.eye(d.numel(), dtype=torch.bool, device=device) - .repeat_interleave(d * H * W, dim=0) - .repeat_interleave(d * H * W, dim=1) - ) - attn_map += doc.log() - attn_map = torch.softmax(attn_map / math.sqrt(D), dim=-1) + s1 = S // block_size + qa = q.reshape(B, h, s1, block_size, D).mean(-2) + ka = k.reshape(B, h, s1, block_size, D).mean(-2).transpose(-2, -1) + attn_map = torch.softmax((qa @ ka) / math.sqrt(D), dim=-1) - # Map binarization vals, inds = attn_map.sort(-1) - cvals = vals.cumsum_(-1) - mask = (cvals >= 1 - thr).int() - mask = mask.gather(-1, inds.argsort(-1)) + mask = (vals.cumsum_(-1) >= 1 - thr).int().gather(-1, inds.argsort(-1)) + mask = torch.logical_or(mask, sta) - if add_sta: - sta = fast_sta_nabla(T, H, W, wT, wH, wW, device=device).unsqueeze_(0).unsqueeze_(0) - mask = torch.logical_or(mask, sta) - mask = torch.logical_and(mask, doc) - - # BlockMask creation kv_nb = mask.sum(-1).to(torch.int32) kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) - return BlockMask.from_kv_blocks(torch.zeros_like(kv_nb), kv_inds, kv_nb, kv_inds, BLOCK_SIZE=64, mask_mod=None) - - -# Neural network building blocks for the diffusion transformer. - + return BlockMask.from_kv_blocks( + torch.zeros_like(kv_nb), + kv_inds, + kv_nb, + kv_inds, + BLOCK_SIZE=block_size, + mask_mod=None, + ) -# `mode="max-autotune-no-cudagraphs"` autotunes Triton kernels for CUDA; fall back to eager -# `flex_attention` on CPU/MPS (tests, non-CUDA inference) instead of attempting CUDA autotuning. -_compiled_flex = torch.compile(flex_attention, mode="max-autotune-no-cudagraphs", dynamic=True) +def sliding_tile_mask( + num_frames: int, height: int, width: int, window: tuple[int, int, int], device: torch.device +) -> Tensor: + """Build the static sliding-tile-attention prior over the `(t, h, w)` grid of 8x8 token blocks. -def flex(query: Tensor, key: Tensor, value: Tensor, **kwargs: Any) -> Tensor: - if query.device.type == "cuda": - return _compiled_flex(query, key, value, **kwargs) - return flex_attention(query, key, value, **kwargs) + Every block attends to the blocks within `window` (odd `(t, h, w)` extents) of itself. Returned as a boolean + `(num_frames * height * width, num_frames * height * width)` matrix in the flattened block order. + """ + window_t, window_h, window_w = window + positions = torch.arange(max(num_frames, height, width), device=device) + distance = (positions[:, None] - positions[None, :]).abs() + near_t = (distance[:num_frames, :num_frames] <= window_t // 2).flatten() + near_h = (distance[:height, :height] <= window_h // 2).flatten() + near_w = (distance[:width, :width] <= window_w // 2).flatten() + near_hw = (near_h[:, None] & near_w[None, :]).reshape(height, height, width, width).transpose(1, 2).flatten() + near = (near_t[:, None] & near_hw[None, :]).reshape(num_frames, num_frames, height * width, height * width) + return near.transpose(1, 2).reshape(num_frames * height * width, num_frames * height * width) -# Without FA3 the only varlen entry points available are FA2's -# ``flash_attn_varlen_*_func``, whose C++ binding declares ``max_seqlen_q``/ -# ``max_seqlen_k`` as ``SymInt`` but does not support FakeTensors → dynamo -# can't trace through it. Skip dynamo on the affected methods so the -# surrounding code in ``TransformerDecoderBlock.forward`` (norms, QKV -# projections, modulation, FFN) still gets ``@torch.compile``-fused while -# only the FA call runs in eager. When FA3 is available the dedicated -# ``flash_attn_interface.flash_attn_varlen_func`` does register a custom op -# and traces cleanly, so this decorator becomes a no-op. +class Kandinsky6SRAttnProcessor: + """Self-attention processor of the SR transformer: dense by default, NABLA sparse when `sparse_params` is set.""" + _attention_backend = None + _parallel_config = None -def apply_scale_shift_norm(norm: nn.Module, x: Tensor, scale: Tensor, shift: Tensor, idx: Tensor) -> Tensor: - """Apply adaptive normalization with indexed scale and shift, forcing fp32 math.""" - with torch.autocast(device_type=x.device.type, dtype=torch.float32, enabled=x.device.type == "cuda"): - return norm(x) * (scale.index_select(0, idx) + 1.0) + shift.index_select(0, idx) + def __call__( + self, + attn: "Kandinsky6SRAttention", + hidden_states: Tensor, + rotary_emb: Tensor, + sparse_params: dict[str, Any] | None = None, + ) -> Tensor: + query = attn.to_query(hidden_states).unflatten(-1, (attn.num_heads, -1)) + key = attn.to_key(hidden_states).unflatten(-1, (attn.num_heads, -1)) + value = attn.to_value(hidden_states).unflatten(-1, (attn.num_heads, -1)) + query = attn.query_norm(query.float()).type_as(query) + key = attn.key_norm(key.float()).type_as(key) + query = apply_rotary(query, rotary_emb) + key = apply_rotary(key, rotary_emb) + + if sparse_params is None: + hidden_states = dispatch_attention_fn( + query, + key, + value, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + else: + # The block statistics are computed from the `(B, heads, S, D)` layout the mask builder expects; the + # resulting `BlockMask` only runs on the flex backend. + block_mask = nabla_block_mask( + query.transpose(1, 2), + key.transpose(1, 2), + sparse_params["sta_mask"], + thr=sparse_params["threshold"], + ) + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=block_mask, + backend=AttentionBackendName.FLEX, + parallel_config=self._parallel_config, + ) + return attn.out_layer(hidden_states.flatten(2, 3)) -def apply_gate_sum(x: Tensor, out: Tensor, gate: Tensor, idx: Tensor) -> Tensor: - """Add gated residual output to input using indexed gate values, forcing fp32 math.""" - with torch.autocast(device_type=x.device.type, dtype=torch.float32, enabled=x.device.type == "cuda"): - return x + gate.index_select(0, idx) * out +class Kandinsky6SRAttention(nn.Module, AttentionModuleMixin): + """SR self-attention with the Diffusers `set_processor` contract.""" + _default_processor_cls = Kandinsky6SRAttnProcessor + _available_processors = [Kandinsky6SRAttnProcessor] -def apply_rotary(x: Tensor, rope: Tensor) -> Tensor: - """Apply rotary position embeddings to input tensor.""" - with torch.autocast(device_type=x.device.type, enabled=False): - x_ = x.reshape(*x.shape[:-1], -1, 1, 2).to(torch.float32) - x_out = rope[..., 0] * x_[..., 0] + rope[..., 1] * x_[..., 1] - return x_out.reshape(*x.shape) + def __init__(self, num_channels: int, head_dim: int, processor: Kandinsky6SRAttnProcessor | None = None): + super().__init__() + if num_channels % head_dim: + raise ValueError("num_channels must be divisible by head_dim") + self.num_heads = num_channels // head_dim + self.to_query = nn.Linear(num_channels, num_channels) + self.to_key = nn.Linear(num_channels, num_channels) + self.to_value = nn.Linear(num_channels, num_channels) + self.query_norm = nn.RMSNorm(head_dim) + self.key_norm = nn.RMSNorm(head_dim) + self.out_layer = nn.Linear(num_channels, num_channels) + self.set_processor(processor or self._default_processor_cls()) + def forward( + self, + hidden_states: Tensor, + rotary_emb: Tensor, + sparse_params: dict[str, Any] | None = None, + ) -> Tensor: + return self.processor(self, hidden_states, rotary_emb, sparse_params) -class TimeEmbeddings(nn.Module): - """Sinusoidal time step embeddings projected through a two-layer MLP.""" - def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0) -> None: - """Initialize time embeddings. +# Copied from diffusers.models.transformers.transformer_kandinsky6.Kandinsky6TimeEmbeddings with Kandinsky6->Kandinsky6SR +class Kandinsky6SRTimeEmbeddings(nn.Module): + """Sinusoidal timestep embedding with a K6-compatible parameter layout.""" - Args: - model_dim: Dimension of sinusoidal encoding. Must be even. - time_dim: Output dimension after MLP projection. - max_period: Maximum period for frequency computation. - """ + def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): super().__init__() - if model_dim % 2 != 0: - msg = "model_dim must be even" - raise ValueError(msg) + if model_dim % 2: + raise ValueError("model_dim must be even") self.model_dim = model_dim self.max_period = max_period - self.register_buffer("freqs", get_freqs(model_dim // 2, max_period), persistent=False) - - self.in_layer = nn.Linear(model_dim, time_dim, bias=True) + self.in_layer = nn.Linear(model_dim, time_dim) self.activation = nn.SiLU() - self.out_layer = nn.Linear(time_dim, time_dim, bias=True) - - def forward(self, time: Tensor) -> tuple[Tensor, Tensor]: - """Encode time steps into embeddings. - - Args: - time: Diffusion time steps. - - Returns: - Tuple of (time embeddings, index tensor mapping tokens to embeddings). - """ - with torch.autocast(device_type=time.device.type, dtype=torch.float32, enabled=time.device.type == "cuda"): - args = torch.outer(time, self.freqs.to(device=time.device)) - time_embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - time_embed = self.out_layer(self.activation(self.in_layer(time_embed))) - time_embed_idx = torch.arange(time_embed.shape[0], device=time_embed.device, dtype=torch.int32) - return time_embed, time_embed_idx - - -class TextEmbeddings(nn.Module): - """Linear projection with layer normalization for text encoder outputs.""" - - def __init__(self, text_dim: int, model_dim: int) -> None: - """Initialize text embeddings. - - Args: - text_dim: Input dimension of text encoder hidden states. - model_dim: Output dimension after projection. - """ - super().__init__() - self.in_layer = nn.Linear(text_dim, model_dim, bias=True) - self.norm = nn.LayerNorm(model_dim, elementwise_affine=True) - - def forward(self, text_embed: Tensor) -> Tensor: - """Project and normalize text embeddings. - - Args: - text_embed: Raw text encoder hidden states. - - Returns: - Projected and normalized text embeddings. - """ - text_embed = self.in_layer(text_embed) - return self.norm(text_embed).type_as(text_embed) + self.out_layer = nn.Linear(time_dim, time_dim) + def forward(self, timestep: Tensor) -> Tensor: + # The sinusoidal embedding is float32; `_keep_in_fp32_modules` keeps these layers float32 under + # `from_pretrained(torch_dtype=...)`, and the cast aligns the input with whatever dtype they hold. + embed = get_timestep_embedding( + timestep, self.model_dim, flip_sin_to_cos=True, downscale_freq_shift=0, max_period=self.max_period + ) + embed = embed.to(get_parameter_dtype(self.in_layer)) + return self.out_layer(self.activation(self.in_layer(embed))) -class VisualEmbeddings(nn.Module): - """Patchify and project visual tokens into the transformer hidden space.""" - def __init__(self, visual_dim: int, model_dim: int, patch_size: tuple[int, int, int]) -> None: - """Initialize visual embeddings. +# Copied from diffusers.models.transformers.transformer_kandinsky6.Kandinsky6VisualEmbeddings with Kandinsky6->Kandinsky6SR +class Kandinsky6SRVisualEmbeddings(nn.Module): + """Patch projection for ``[B,T,H,W,C]`` visual tokens.""" - Args: - visual_dim: Number of input visual channels. - model_dim: Output dimension after projection. - patch_size: Patch size as (temporal, height, width). - """ + 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 _patchify(self, x: Tensor, visual_cu_seqlens: Tensor) -> tuple[Tensor, Tensor]: - """Shared patchification logic for both full input and LQ-only slices. - - Args: - x: Input tensor of shape ``(duration, height, width, channels)``. - visual_cu_seqlens: Cumulative sequence lengths. - - Returns: - Tuple of (patchified tensor, updated cumulative sequence lengths). - """ - if self.patch_size[0] > 1: - idxs = torch.ones(x.shape[0], dtype=torch.int32, device=visual_cu_seqlens.device) - idxs[visual_cu_seqlens[:-1]] += self.patch_size[0] - 1 - x = torch.repeat_interleave(x, idxs, dim=0) - visual_cu_seqlens = visual_cu_seqlens + torch.arange( - visual_cu_seqlens.shape[0], device=visual_cu_seqlens.device, dtype=torch.int32 - ) - - duration, height, width, dim = x.shape - # [T, H, W, C] -> [T/pt, H/ph, W/pw, pt*ph*pw*C] - # Groups of (pt, ph, pw) neighboring pixels are concatenated into one token vector + def forward(self, x: Tensor) -> Tensor: + batch, duration, height, width, channels = x.shape + p_t, p_h, p_w = self.patch_size x = ( - x.view( - duration // self.patch_size[0], - self.patch_size[0], - height // self.patch_size[1], - self.patch_size[1], - width // self.patch_size[2], - self.patch_size[2], - dim, - ) - .permute(0, 2, 4, 1, 3, 5, 6) - .flatten(3, 6) + x.view(batch, 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) ) - visual_cu_seqlens = visual_cu_seqlens // self.patch_size[0] - return x, visual_cu_seqlens - - def forward(self, x: Tensor, visual_cu_seqlens: Tensor) -> tuple[Tensor, Tensor]: - """Patchify input and project to model dimension. - - Args: - x: Visual input of shape ``(duration, height, width, channels)``. - visual_cu_seqlens: Cumulative sequence lengths for packed sequences. - - Returns: - ``(projected_patches, cu_seqlens)``. - """ - patches, cu = self._patchify(x, visual_cu_seqlens) - return self.in_layer(patches), cu - - -class RoPE1D(nn.Module): - """1D rotary position embeddings for text tokens.""" - - def __init__(self, dim: int, max_pos: int = 1024, max_period: float = 10000.0) -> None: - """Initialize 1D rotary position embeddings. - - Args: - dim: Embedding dimension. Must be even. - max_pos: Maximum number of positions to precompute. - max_period: Maximum period for frequency computation. - """ - super().__init__() - self.max_period = max_period - self.dim = dim - self.max_pos = max_pos - 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: - """Compute rotary embeddings for given positions. - - Args: - pos: Position indices. - - Returns: - Rotary embedding matrix with shape ``(*pos.shape, 1, dim//2, 2, 2)``. - """ - with torch.autocast(device_type=pos.device.type, enabled=False): - args = self.args[pos] - rope = torch.stack([torch.cos(args), -torch.sin(args), torch.sin(args), torch.cos(args)], dim=-1) - rope = rope.view(*rope.shape[:-1], 2, 2) - return rope.unsqueeze(-4) - - -class RoPE3D(nn.Module): - """3D rotary position embeddings for visual tokens along temporal, height, and width axes.""" + return self.in_layer(x) - def __init__( - self, - axes_dims: tuple[int, int, int], - max_pos: tuple[int, int, int] = (128, 128, 128), - max_period: float = 10000.0, - ) -> None: - """Initialize 3D rotary position embeddings. - - Args: - axes_dims: Per-axis embedding dimensions (temporal, height, width). - max_pos: Maximum positions per axis. - max_period: Maximum period for frequency computation. - """ - super().__init__() - self.axes_dims = axes_dims - self.max_pos = max_pos - self.max_period = max_period - for i, (axes_dim, ax_max_pos) in enumerate(zip(axes_dims, max_pos, strict=False)): - freq = get_freqs(axes_dim // 2, max_period) - pos = torch.arange(ax_max_pos, dtype=freq.dtype) - self.register_buffer(f"args_{i}", torch.outer(pos, freq), persistent=False) +# Copied from diffusers.models.transformers.transformer_kandinsky6.Kandinsky6Modulation with Kandinsky6->Kandinsky6SR +class Kandinsky6SRModulation(nn.Module): + """Zero-initialized AdaLN modulation projection.""" - def forward( - self, - shape: tuple[int, int, int], - pos: tuple[Tensor, Tensor, Tensor], - scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), - ) -> Tensor: - """Compute 3D rotary embeddings for given positions. - - Args: - shape: Spatial dimensions as (duration, height, width). - pos: Position indices per axis (temporal, height, width). - scale_factor: Frequency scaling per axis. - - Returns: - Rotary embedding matrix broadcast to (duration, height, width, 1, dim//2, 2, 2). - """ - with torch.autocast(device_type=pos[0].device.type, enabled=False): - duration, height, width = shape - args_t = self.args_0[pos[0]] / 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, - ) - rope = torch.stack([torch.cos(args), -torch.sin(args), torch.sin(args), torch.cos(args)], dim=-1) - rope = rope.view(*rope.shape[:-1], 2, 2) - return rope.unsqueeze(-4) - - -class Modulation(nn.Module): - """Adaptive modulation layer producing scale, shift, and gate parameters from time embeddings.""" - - def __init__(self, time_dim: int, model_dim: int, num_params: int) -> None: - """Initialize modulation layer. - - Args: - time_dim: Input dimension of time embeddings. - model_dim: Per-parameter output dimension. - num_params: Number of modulation parameters to produce. - """ + 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_() + nn.init.zeros_(self.out_layer.weight) + nn.init.zeros_(self.out_layer.bias) def forward(self, x: Tensor) -> Tensor: - """Compute modulation parameters from time embeddings. - - Args: - x: Time embedding tensor. - - Returns: - Concatenated modulation parameters. - """ - with torch.autocast(device_type=x.device.type, dtype=torch.float32, enabled=x.device.type == "cuda"): - return self.out_layer(self.activation(x)) - - -class MultiheadSelfAttention(nn.Module, AttentionModuleMixin): - """Multi-head self-attention with flash attention and optional sparse flex attention.""" - - _default_processor_cls = Kandinsky6SRAttentionProcessor - _available_processors = [Kandinsky6SRAttentionProcessor] - - def __init__(self, num_channels: int, head_dim: int) -> None: - """Initialize multi-head self-attention. - - Args: - num_channels: Total number of channels. Must be divisible by head_dim. - head_dim: Dimension per attention head. - """ - super().__init__() - if num_channels % head_dim != 0: - msg = "num_channels must be divisible by head_dim" - raise ValueError(msg) - 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.cached_k: Tensor | None = None - self.cached_v: Tensor | None = None - self.cached_cu_seqlens: Tensor | None = None - self.return_kv = False - self.set_processor(self._default_processor_cls()) - - def get_qkv(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: - """Project input into query, key, and value tensors. - - Args: - x: Input tensor. - - Returns: - Tuple of (query, key, value) reshaped to ``(seq_len, num_heads, head_dim)``. - """ - query = self.to_query(x) - key = self.to_key(x) - value = self.to_value(x) - - shape = query.shape[:-1] # for TP compatibility - query = query.reshape(*shape, self.num_heads, -1) - key = key.reshape(*shape, self.num_heads, -1) - value = value.reshape(*shape, self.num_heads, -1) - - return query, key, value - - def norm_qk(self, q: Tensor, k: Tensor) -> tuple[Tensor, Tensor]: - """Apply RMS normalization to query and key. - - Args: - q: Query tensor. - k: Key tensor. - - Returns: - Tuple of (normalized query, normalized key). - """ - q = self.query_norm(q.float()).type_as(q) - k = self.key_norm(k.float()).type_as(k) - return q, k - - def reset_kv_cache(self) -> None: - """Clear the rolling KV-cache and stop emitting K/V.""" - self.cached_k = None - self.cached_v = None - self.cached_cu_seqlens = None - self.return_kv = False - - def attention_flex( - self, - query: Tensor, - key: Tensor, - value: Tensor, - sparse_params: dict[str, Any] | None = None, - *, - return_sparsity: bool = False, - ) -> Tensor | tuple[Tensor, float]: - """Compute sparse self-attention using flex attention with block masks. - - Args: - query: Query tensor. - key: Key tensor. - value: Value tensor. - sparse_params: Sparse attention configuration. - return_sparsity: Whether to return sparsity percentage. - - Returns: - Attention output, or tuple of (output, sparsity percentage) if requested. - """ - if self.cached_k is not None: - msg = "KV-cache is not supported on the flex/NABLA attention path; only dense flash attention supports it." - raise NotImplementedError(msg) - - query = query.unsqueeze(0).transpose(1, 2).contiguous() - key = key.unsqueeze(0).transpose(1, 2).contiguous() - value = value.unsqueeze(0).transpose(1, 2).contiguous() - - t, h, w = sparse_params["visual_shape"] - # Token grid -> fractal 8x8 block grid; the block size matches - # fractal_flatten and is unrelated to the VAE spatial compression. - h, w = h // FRACTAL_BLOCK_SIZE, w // FRACTAL_BLOCK_SIZE - visual_seqlens = sparse_params["visual_seqlens"].to(device=query.device) - block_mask = nablaT_v2_doc( - query, - key, - visual_seqlens, - t, - h, - w, - wT=sparse_params["wT"], - wH=sparse_params["wH"], - wW=sparse_params["wW"], - thr=sparse_params["P"], - add_sta=sparse_params["add_sta"], - ) - out = ( - flex( - query, - key, - value, - block_mask=block_mask, - kernel_options={"BLOCK_M": 64, "BLOCK_N": 64}, - ) - .transpose(1, 2) - .squeeze(0) - .contiguous() - ) - out = out.flatten(-2, -1) - - if return_sparsity: - sparsity = 100.0 * (1 - (1 - block_mask.sparsity() / 100) * (sparse_params["visual_seqlens"].shape[0] - 1)) - return out, sparsity - return out - - def forward( - self, - x: Tensor, - rope: Tensor, - cu_seqlens: Tensor, - sparse_params: dict[str, Any] | None = None, - ) -> Tensor: - """Run self-attention with rotary embeddings. - - Args: - x: Input tensor. - rope: Rotary position embeddings. - cu_seqlens: Cumulative sequence lengths for packed sequences. - sparse_params: Optional sparse attention parameters. - - Returns: - Self-attention output. - """ - query, key, value = self.get_qkv(x) - query, key = self.norm_qk(query, key) - query = apply_rotary(query, rope).type_as(query) - key = apply_rotary(key, rope).type_as(key) + return self.out_layer(self.activation(x.to(get_parameter_dtype(self.out_layer)))) - out = self.processor(self, query, key, value, cu_seqlens, cu_seqlens, sparse_params=sparse_params) - return self.out_layer(out) +# Copied from diffusers.models.transformers.transformer_kandinsky6.Kandinsky6FeedForward with Kandinsky6->Kandinsky6SR +class Kandinsky6SRFeedForward(nn.Module): + """K6 bias-free GELU feed-forward network.""" - -class MultiheadCrossAttention(nn.Module, AttentionModuleMixin): - """Multi-head cross-attention with flash attention.""" - - _default_processor_cls = Kandinsky6SRAttentionProcessor - _available_processors = [Kandinsky6SRAttentionProcessor] - - def __init__(self, num_channels: int, head_dim: int) -> None: - """Initialize multi-head cross-attention. - - Args: - num_channels: Total number of channels. Must be divisible by head_dim. - head_dim: Dimension per attention head. - """ - super().__init__() - if num_channels % head_dim != 0: - msg = "num_channels must be divisible by head_dim" - raise ValueError(msg) - 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(self._default_processor_cls()) - - def get_qkv(self, x: Tensor, cond: Tensor) -> tuple[Tensor, Tensor, Tensor]: - """Project input and condition into query, key, and value tensors. - - Args: - x: Input tensor for query projection. - cond: Conditioning tensor for key and value projections. - - Returns: - Tuple of (query, key, value) reshaped to ``(seq_len, num_heads, head_dim)``. - """ - query = self.to_query(x) - key = self.to_key(cond) - value = self.to_value(cond) - - shape, cond_shape = query.shape[:-1], key.shape[:-1] # for TP compatibility - query = query.reshape(*shape, self.num_heads, -1) - key = key.reshape(*cond_shape, self.num_heads, -1) - value = value.reshape(*cond_shape, self.num_heads, -1) - - return query, key, value - - def norm_qk(self, q: Tensor, k: Tensor) -> tuple[Tensor, Tensor]: - """Apply RMS normalization to query and key. - - Args: - q: Query tensor. - k: Key tensor. - - Returns: - Tuple of (normalized query, normalized key). - """ - q = self.query_norm(q.float()).type_as(q) - k = self.key_norm(k.float()).type_as(k) - return q, k - - def forward(self, x: Tensor, cond: Tensor, cu_seqlens: Tensor, cond_cu_seqlens: Tensor) -> Tensor: - """Run cross-attention between input and conditioning. - - Args: - x: Input tensor for query computation. - cond: Conditioning tensor for key and value computation. - cu_seqlens: Cumulative sequence lengths for input sequences. - cond_cu_seqlens: Cumulative sequence lengths for conditioning sequences. - - Returns: - Cross-attention output. - """ - query, key, value = self.get_qkv(x, cond) - query, key = self.norm_qk(query, key) - - out = self.processor(self, query, key, value, cu_seqlens, cond_cu_seqlens) - return self.out_layer(out) - - -class FeedForward(nn.Module): - """Two-layer feed-forward network with GELU activation.""" - - def __init__(self, dim: int, ff_dim: int) -> None: - """Initialize feed-forward network. - - Args: - dim: Input and output dimension. - ff_dim: Hidden layer dimension. - """ + def __init__(self, dim: int, ff_dim: int): 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: Tensor) -> Tensor: - """Apply feed-forward transformation. - - Args: - x: Input tensor. - - Returns: - Transformed tensor. - """ return self.out_layer(self.activation(self.in_layer(x))) -class OutLayer(nn.Module): - """Final output layer that unpatchifies and projects visual tokens to pixel space.""" - - def __init__(self, model_dim: int, time_dim: int, visual_dim: int, patch_size: tuple[int, int, int]) -> None: - """Initialize output layer. +# Copied from diffusers.models.transformers.transformer_kandinsky6.Kandinsky6OutLayer with Kandinsky6->Kandinsky6SR +class Kandinsky6SROutLayer(nn.Module): + """Projects visual hidden states back to packed latent patches.""" - Args: - model_dim: Input dimension from the transformer. - time_dim: Dimension of time embeddings for modulation. - visual_dim: Number of output visual channels. - patch_size: Patch size as (temporal, height, width). - """ + def __init__(self, model_dim: int, time_dim: int, visual_dim: int, patch_size: tuple[int, int, int]): super().__init__() self.patch_size = patch_size - self.modulation = Modulation(time_dim, model_dim, 2) + self.modulation = Kandinsky6SRModulation(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) + self.out_layer = nn.Linear(model_dim, math.prod(patch_size) * visual_dim) - def forward( - self, - visual_embed: Tensor, - _text_embed: Tensor, - time_embed: Tensor, - visual_cu_seqlens: Tensor, - time_embed_idx: Tensor, - ) -> Tensor: - """Unpatchify and project visual embeddings to output channels. - - Args: - visual_embed: Visual token embeddings from the transformer. - _text_embed: Text embeddings (unused, reserved for interface compatibility). - time_embed: Time step embeddings for adaptive modulation. - visual_cu_seqlens: Cumulative sequence lengths for visual tokens. - time_embed_idx: Index mapping visual tokens to their time embeddings. - - Returns: - Denoised visual output tensor. - """ + def forward(self, visual_embed: Tensor, time_embed: Tensor) -> Tensor: shift, scale = torch.chunk(self.modulation(time_embed), 2, dim=-1) - visual_embed = apply_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) + condition_shape = (scale.shape[0],) + (1,) * (visual_embed.ndim - 2) + (scale.shape[-1],) + x = apply_scale_shift_norm( + self.norm, + visual_embed, + scale.reshape(condition_shape), + shift.reshape(condition_shape), + ) + x = self.out_layer(x) - duration, height, width, _dim = x.shape - x = ( - x.view( - duration, - height, - width, - -1, - self.patch_size[0], - self.patch_size[1], - self.patch_size[2], - ) - .permute(0, 4, 1, 5, 2, 6, 3) - .flatten(0, 1) + batch, duration, height, width = x.shape[:4] + p_t, p_h, p_w = self.patch_size + return ( + x.view(batch, duration, height, width, -1, p_t, p_h, p_w) + .permute(0, 1, 5, 2, 6, 3, 7, 4) .flatten(1, 2) .flatten(2, 3) + .flatten(3, 4) ) - visual_cu_seqlens = visual_cu_seqlens * self.patch_size[0] - - if self.patch_size[0] > 1: - idxs = torch.ones(duration * self.patch_size[0], dtype=torch.int32, device=visual_cu_seqlens.device) - idxs[visual_cu_seqlens[:-1]] -= self.patch_size[0] - 1 - x = torch.repeat_interleave(x, idxs, dim=0) - return x - -class TransformerEncoderBlock(nn.Module): - """Transformer encoder block with self-attention, feed-forward, and adaptive modulation.""" - def __init__(self, model_dim: int, time_dim: int, ff_dim: int, head_dim: int) -> None: - """Initialize encoder block layers. +# Copied from diffusers.models.transformers.transformer_kandinsky6.RoPE3D with RoPE3D->Kandinsky6SRRoPE3D +class Kandinsky6SRRoPE3D(nn.Module): + """3-D Rotary Position Embedding — used for video spatial-temporal tokens (T, H, W).""" - Args: - model_dim: Hidden dimension of the transformer. - time_dim: Dimension of time step embeddings. - ff_dim: Inner dimension of the feed-forward network. - head_dim: Per-head dimension for self-attention. - """ + 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.text_modulation = Modulation(time_dim, model_dim, 6) - - self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) - self.self_attention = MultiheadSelfAttention(model_dim, head_dim) - - self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) - self.feed_forward = FeedForward(model_dim, ff_dim) + self.axes_dims = axes_dims + self.max_pos = max_pos + self.max_period = max_period + for i, (d, mp) in enumerate(zip(axes_dims, max_pos)): + freq = get_freqs(d // 2, max_period) + self.register_buffer(f"args_{i}", torch.outer(torch.arange(mp, dtype=freq.dtype), freq), persistent=False) def forward( self, - x: Tensor, - time_embed: Tensor, - rope: Tensor, - cu_seqlens: Tensor, - time_embed_idx: Tensor, + pos: tuple[Tensor, Tensor, Tensor], + scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), ) -> Tensor: - """Run self-attention and feed-forward with time-conditioned modulation. - - Args: - x: Input text embeddings. - time_embed: Time step embeddings for modulation. - rope: Rotary position embeddings. - cu_seqlens: Cumulative sequence lengths for packed sequences. - time_embed_idx: Index mapping tokens to their time embeddings. - - Returns: - Modulated text embeddings. - """ - 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 = apply_scale_shift_norm(self.self_attention_norm, x, scale, shift, time_embed_idx).type_as(x) - out = self.self_attention(out, rope, cu_seqlens) - x = apply_gate_sum(x, out, gate, time_embed_idx).type_as(x) - - shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) - out = apply_scale_shift_norm(self.feed_forward_norm, x, scale, shift, time_embed_idx).type_as(x) - out = self.feed_forward(out) - return apply_gate_sum(x, out, gate, time_embed_idx).type_as(x) + # `pos` holds one index tensor per axis; the grid size is implied by their lengths. + num_frames, height, width = (int(axis_pos.shape[0]) for axis_pos in pos) + args_t = self.args_0[pos[0]] / scale_factor[0] # (T, d//2) + args_h = self.args_1[pos[1]] / scale_factor[1] # (H, d//2) + args_w = self.args_2[pos[2]] / scale_factor[2] # (W, d//2) + args = torch.cat( + [ + args_t.view(num_frames, 1, 1, -1).expand(num_frames, height, width, -1), + args_h.view(1, height, 1, -1).expand(num_frames, height, width, -1), + args_w.view(1, 1, width, -1).expand(num_frames, height, width, -1), + ], + dim=-1, + ) + cos, sin = torch.cos(args), torch.sin(args) + rope = torch.stack([cos, -sin, sin, cos], dim=-1) # (T, H, W, total_dim, 4) + rope = rope.view(*rope.shape[:-1], 2, 2) # (T, H, W, total_dim, 2, 2) + return rope.unsqueeze(-4) # (T, H, W, 1, total_dim, 2, 2) -class TransformerDecoderBlock(nn.Module): - """Transformer decoder block with self-attention, cross-attention, and feed-forward.""" - def __init__( - self, - model_dim: int, - time_dim: int, - ff_dim: int, - head_dim: int, - *, - use_text: bool = True, - ) -> None: - """Initialize decoder block layers. +class Kandinsky6SRTransformerBlock(nn.Module): + """AdaLN-modulated self-attention + feed-forward block of the SR transformer.""" - Args: - model_dim: Hidden dimension of the transformer. - time_dim: Dimension of time step embeddings. - ff_dim: Inner dimension of the feed-forward network. - head_dim: Per-head dimension for attention layers. - use_text: When False, the block has no text cross-attention: no - ``cross_attention``/``cross_attention_norm`` modules and the - modulation produces 6 params (self-attn + FFN) instead of 9. - Used by the text-free SR variant. - """ + def __init__(self, model_dim: int, time_dim: int, ff_dim: int, head_dim: int): super().__init__() - self.use_text = use_text - self.visual_modulation = Modulation(time_dim, model_dim, 9 if use_text else 6) - + self.visual_modulation = Kandinsky6SRModulation(time_dim, model_dim, 6) self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) - self.self_attention = MultiheadSelfAttention(model_dim, head_dim) - - if use_text: - self.cross_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) - self.cross_attention = MultiheadCrossAttention(model_dim, head_dim) - + self.self_attention = Kandinsky6SRAttention(model_dim, head_dim) self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) - self.feed_forward = FeedForward(model_dim, ff_dim) + self.feed_forward = Kandinsky6SRFeedForward(model_dim, ff_dim) def forward( self, - visual_embed: Tensor, - text_embed: Tensor, - time_embed: Tensor, - rope: Tensor, - visual_cu_seqlens: Tensor, - text_cu_seqlens: Tensor, - time_embed_idx: Tensor, - sparse_params: dict[str, Any] | None, + hidden_states: Tensor, + temb: Tensor, + rotary_emb: Tensor, + sparse_params: dict[str, Any] | None = None, ) -> Tensor: - """Run self-attention, cross-attention, and feed-forward with modulation. - - Args: - visual_embed: Visual token embeddings. - text_embed: Encoded text embeddings used as cross-attention context. - time_embed: Time step embeddings for modulation. - rope: Rotary position embeddings for visual tokens. - visual_cu_seqlens: Cumulative sequence lengths for visual tokens. - text_cu_seqlens: Cumulative sequence lengths for text tokens. - time_embed_idx: Index mapping visual tokens to their time embeddings. - sparse_params: Optional sparse attention parameters. - - Returns: - Updated visual embeddings. - """ - dtype = visual_embed.dtype - if self.use_text: - self_attn_params, cross_attn_params, ff_params = torch.chunk(self.visual_modulation(time_embed), 3, dim=-1) - else: - self_attn_params, ff_params = torch.chunk(self.visual_modulation(time_embed), 2, dim=-1) - cross_attn_params = None - - # --- Self-attention --- - # Norm -> modulate by (scale + 1, shift) -> SelfAttn -> gate - shift_t, scale_t, gate_t = torch.chunk(self_attn_params, 3, dim=-1) - with torch.autocast( - device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" - ): - normed = self.self_attention_norm(visual_embed) - scale = scale_t.index_select(0, time_embed_idx) - shift = shift_t.index_select(0, time_embed_idx) - visual_out = normed * (scale + 1.0) + shift - visual_out = visual_out.type_as(visual_embed) - visual_out = self.self_attention(visual_out, rope, visual_cu_seqlens, sparse_params) - with torch.autocast( - device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" - ): - visual_embed = visual_embed + gate_t.index_select(0, time_embed_idx) * visual_out - visual_embed = visual_embed.to(dtype=dtype) - - # --- Cross-attention --- (skipped when text-free) - # Norm -> modulate by (scale + 1, shift) -> CrossAttn -> gate - if self.use_text: - shift_t, scale_t, gate_t = torch.chunk(cross_attn_params, 3, dim=-1) - with torch.autocast( - device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" - ): - normed = self.cross_attention_norm(visual_embed) - scale = scale_t.index_select(0, time_embed_idx) - shift = shift_t.index_select(0, time_embed_idx) - visual_out = normed * (scale + 1.0) + shift - visual_out = visual_out.type_as(visual_embed) - visual_out = self.cross_attention(visual_out, text_embed, visual_cu_seqlens, text_cu_seqlens) - with torch.autocast( - device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" - ): - visual_embed = visual_embed + gate_t.index_select(0, time_embed_idx) * visual_out - visual_embed = visual_embed.to(dtype=dtype) - - # --- Feed-forward --- - # Norm -> modulate by (scale + 1, shift) -> FF -> gate - shift_t, scale_t, gate_t = torch.chunk(ff_params, 3, dim=-1) - with torch.autocast( - device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" - ): - normed = self.feed_forward_norm(visual_embed) - scale = scale_t.index_select(0, time_embed_idx) - shift = shift_t.index_select(0, time_embed_idx) - visual_out = normed * (scale + 1.0) + shift - visual_out = visual_out.type_as(visual_embed) - visual_out = self.feed_forward(visual_out) - with torch.autocast( - device_type=visual_embed.device.type, dtype=torch.float32, enabled=visual_embed.device.type == "cuda" - ): - visual_embed = visual_embed + gate_t.index_select(0, time_embed_idx) * visual_out - return visual_embed.to(dtype=dtype) + self_attention_params, feed_forward_params = torch.chunk(self.visual_modulation(temb), 2, dim=-1) + shift, scale, gate = torch.chunk(self_attention_params, 3, dim=-1) + hidden_states = apply_gate_sum( + hidden_states, + self.self_attention( + apply_scale_shift_norm(self.self_attention_norm, hidden_states, scale, shift), + rotary_emb, + sparse_params, + ), + gate, + ) + shift, scale, gate = torch.chunk(feed_forward_params, 3, dim=-1) + return apply_gate_sum( + hidden_states, + self.feed_forward(apply_scale_shift_norm(self.feed_forward_norm, hidden_states, scale, shift)), + gate, + ) class Kandinsky6SRTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin): - """3D diffusion transformer with text-conditioned visual generation for Kandinsky 6 super-resolution. + r""" + Text-free diffusion transformer for Kandinsky 6 video super-resolution. - Unlike `Kandinsky6Transformer3DModel` (the TI2VA transformer), which batches over a padded - `(B, T, H, W, C)` tensor, this transformer packs variable-length visual/text sequences into a single - sequence per axis using `cu_seqlens` (the flash-attention varlen convention). This lets the SR pipeline - batch many spatial tiles per forward pass without padding overhead. Visual tokens are processed through - encoder (text) and decoder (visual) transformer blocks with RoPE, adaptive modulation, and optional - NABLA sparse attention. + The model denoises tiles of the K-VAE latent video. Its input concatenates the noisy latent with the anchor latent + and the anchor mask that condition the super-resolution (`2 * in_visual_dim + 1` channels), and its output holds + `out_visual_dim` channels: `in_visual_dim` for a plain velocity checkpoint, or a widened `n_grid * in_visual_dim` + for the [`PiflowScheduler`] distilled checkpoints. Video self-attention runs dense, or through the NABLA sparse + block pattern when `nabla_threshold` is set. Args: - in_visual_dim (`int`, *optional*, defaults to 4): Number of input visual channels. - in_text_dim (`int`, *optional*, defaults to 3584): Dimension of text encoder hidden states. - in_text_dim2 (`int`, *optional*, defaults to 768): Dimension of pooled text embeddings. - time_dim (`int`, *optional*, defaults to 512): Dimension of time step embeddings. - out_visual_dim (`int`, *optional*, defaults to 4): Number of output visual channels. - patch_size (`tuple[int, int, int]`, *optional*, defaults to `(1, 2, 2)`): Patch size as - (temporal, height, width). - model_dim (`int`, *optional*, defaults to 2048): Hidden dimension of the transformer. - ff_dim (`int`, *optional*, defaults to 5120): Inner dimension of feed-forward networks. - num_text_blocks (`int`, *optional*, defaults to 2): Number of text encoder blocks. - num_visual_blocks (`int`, *optional*, defaults to 32): Number of visual decoder blocks. - axes_dims (`tuple[int, int, int]`, *optional*, defaults to `(16, 24, 24)`): RoPE dimension split - per axis (temporal, height, width). - visual_cond (`bool`, *optional*, defaults to `False`): Whether to use visual conditioning input. - instruct_type (`str`, *optional*): Conditioning mode: `"noise"`, `"channel"`, `"hybrid"`, or - `"hybrid_anchor"` (the mode Kandinsky 6 SR ships with). - attention_params (`dict`, *optional*): Per-resolution attention configuration, keyed by base tile - resolution; each entry selects dense (`type="flash"`) or sparse NABLA (`type="nabla"`) attention. - use_motion_score (`bool`, *optional*, defaults to `False`): Whether to add motion score embeddings - to time conditioning. - use_text (`bool`, *optional*, defaults to `True`): When `False`, builds a text-free model with no - text encoder, text RoPE/blocks, pooled-text conditioning, or per-block cross-attention. - sr_params (`dict`, *optional*): SR-specific configuration: `scheduler_scale`, `lq_noise_scale`, - `lq_noise_type`, `lq_channel_noise_scale`, `cap_noise_timestep`, `fps`, `scale_factor` (RoPE - frequency scaling per base resolution), and `visual_size` (base tile resolution). Read by - [`Kandinsky6SRPipeline`] via the attributes this unpacks into. - attribute_overrides (`dict`, *optional*): Arbitrary extra attributes set on the model after - construction. + in_visual_dim (`int`, defaults to `64`): + Number of latent channels of the K-VAE. + out_visual_dim (`int`, defaults to `640`): + Number of output channels. + time_dim (`int`, defaults to `512`): + Dimension of the timestep embedding. + patch_size (`tuple[int, int, int]`, defaults to `(1, 1, 1)`): + Patch size as `(temporal, height, width)`. + model_dim (`int`, defaults to `1792`): + Hidden dimension of the transformer. + ff_dim (`int`, defaults to `7168`): + Inner dimension of the feed-forward networks. + num_visual_blocks (`int`, defaults to `32`): + Number of transformer blocks. + axes_dims (`tuple[int, int, int]`, defaults to `(16, 24, 24)`): + RoPE dimensions per `(t, h, w)` axis; their sum is the attention head dimension. + scale_factor (`tuple[float, float, float]`, defaults to `(1.0, 2.0, 2.0)`): + Per-axis RoPE frequency scaling applied to the token positions. + nabla_threshold (`float`, *optional*, defaults to `0.8`): + Cumulative-attention threshold of the NABLA block selection. `None` runs dense attention. + nabla_window (`tuple[int, int, int]`, defaults to `(11, 7, 7)`): + Odd `(t, h, w)` extents of the sliding-tile prior that every 8x8 token block always attends to. + tile_sizes (`tuple[tuple[int, int], ...]`, defaults to `((512, 512), (512, 768), (768, 512))`): + Pixel `(height, width)` sizes of the video tiles the model was trained on. [`Kandinsky6SRPipeline`] refines + every tile at the size whose aspect ratio is closest to the input video's. """ - _no_split_modules = ["TransformerEncoderBlock", "TransformerDecoderBlock"] - _repeated_blocks = ["TransformerEncoderBlock", "TransformerDecoderBlock"] _supports_gradient_checkpointing = True + _no_split_modules = ["Kandinsky6SRTransformerBlock"] + _repeated_blocks = ["Kandinsky6SRTransformerBlock"] + _skip_layerwise_casting_patterns = ["norm"] + # Timestep embeddings, the pooled-text bias they add to, and every AdaLN modulation projection stay float32 + # under `from_pretrained(torch_dtype=...)`. + _keep_in_fp32_modules = ["time_embeddings", "pooled_bias", "modulation"] @register_to_config def __init__( self, - in_visual_dim: int = 4, - in_text_dim: int = 3584, - in_text_dim2: int = 768, + in_visual_dim: int = 64, + out_visual_dim: int = 640, time_dim: int = 512, - out_visual_dim: int = 4, - patch_size: tuple[int, int, int] = (1, 2, 2), - model_dim: int = 2048, - ff_dim: int = 5120, - num_text_blocks: int = 2, + patch_size: tuple[int, int, int] = (1, 1, 1), + model_dim: int = 1792, + ff_dim: int = 7168, num_visual_blocks: int = 32, axes_dims: tuple[int, int, int] = (16, 24, 24), - *, - visual_cond: bool = False, - instruct_type: str | None = None, - attention_params: dict[str, Any] | None = None, - use_motion_score: bool = False, - use_text: bool = True, - sr_params: Mapping[str, Any] | None = None, - attribute_overrides: Mapping[str, Any] | None = None, + scale_factor: tuple[float, float, float] = (1.0, 2.0, 2.0), + nabla_threshold: float | None = 0.8, + nabla_window: tuple[int, int, int] = (11, 7, 7), + tile_sizes: tuple[tuple[int, int], ...] = ((512, 512), (512, 768), (768, 512)), ) -> None: super().__init__() head_dim = sum(axes_dims) - self.in_visual_dim = in_visual_dim - self.instruct_type = instruct_type - self.model_dim = model_dim - self.patch_size = patch_size - self.visual_cond = visual_cond - self.attention_params = attention_params - - sr_params = dict(sr_params) if sr_params is not None else {} - self.sr_params = sr_params - self.scheduler_scale = float(sr_params.get("scheduler_scale", 5.0)) - self.lq_noise_scale = float(sr_params.get("lq_noise_scale", 0.7)) - self.lq_noise_type = sr_params.get("lq_noise_type", "ddpm") - self.lq_channel_noise_scale = float(sr_params.get("lq_channel_noise_scale", 0.0)) - self.cap_noise_timestep = bool(sr_params.get("cap_noise_timestep", False)) - self.fps = int(sr_params.get("fps", 24)) - visual_size = sr_params.get("visual_size") - self.visual_size = ( - int(visual_size[0] if isinstance(visual_size, (list, tuple)) else visual_size) - if visual_size is not None - else None - ) - scale_factor = sr_params.get("scale_factor") - self.scale_factor = ( - {int(key): tuple(float(value) for value in values) for key, values in scale_factor.items()} - if scale_factor is not None - else None - ) - visual_embed_dim = ( - 2 * in_visual_dim + 1 - if visual_cond or instruct_type in ("channel", "hybrid", "hybrid_anchor") - else in_visual_dim - ) - self.time_embeddings = TimeEmbeddings(model_dim, time_dim) - if use_motion_score: - self.motion_embeddings = TimeEmbeddings(model_dim, time_dim) - self.use_motion_score = use_motion_score - - self.use_text = use_text - if use_text: - self.text_embeddings = TextEmbeddings(in_text_dim, model_dim) - self.pooled_text_embeddings = TextEmbeddings(in_text_dim2, time_dim) - else: - # Text-free: under the empty caption the pooled-text contribution - # `pooled_text_embeddings("")` is a constant added to the time embedding. - # build_text_free_dit.py bakes it into this bias; dropping it instead corrupts - # the time conditioning, so the model still needs it (cross-attention does not). - self.pooled_bias = nn.Parameter(torch.zeros(time_dim)) - self.visual_embeddings = VisualEmbeddings(visual_embed_dim, model_dim, patch_size) - - if use_text: - self.text_rope_embeddings = RoPE1D(head_dim) - self.text_transformer_blocks = nn.ModuleList( - [TransformerEncoderBlock(model_dim, time_dim, ff_dim, head_dim) for _ in range(num_text_blocks)] - ) - - self.visual_rope_embeddings = RoPE3D(axes_dims) + self.time_embeddings = Kandinsky6SRTimeEmbeddings(model_dim, time_dim) + # The SR checkpoints were trained with an empty caption: the pooled-text projection of that caption is a + # constant, folded into this bias by the checkpoint conversion. + self.pooled_bias = nn.Parameter(torch.zeros(time_dim)) + self.visual_embeddings = Kandinsky6SRVisualEmbeddings(2 * in_visual_dim + 1, model_dim, patch_size) + self.visual_rope_embeddings = Kandinsky6SRRoPE3D(axes_dims) self.visual_transformer_blocks = nn.ModuleList( - [ - TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim, use_text=use_text) - for _ in range(num_visual_blocks) - ] + [Kandinsky6SRTransformerBlock(model_dim, time_dim, ff_dim, head_dim) for _ in range(num_visual_blocks)] ) - - self.out_layer = OutLayer(model_dim, time_dim, out_visual_dim, patch_size) - for name, value in (attribute_overrides or {}).items(): - setattr(self, str(name), value) + self.out_layer = Kandinsky6SROutLayer(model_dim, time_dim, out_visual_dim, patch_size) self.gradient_checkpointing = False def forward( self, - x: Tensor, - text_embed: Tensor, - pooled_text_embed: Tensor, - time: Tensor, - visual_cu_seqlens: Tensor, - text_cu_seqlens: Tensor, - visual_rope_pos: tuple[Tensor, Tensor, Tensor], - text_rope_pos: Tensor, - scale_factor: tuple[float, float, float] = (1.0, 1.0, 1.0), - sparse_params: dict[str, Any] | None = None, - motion_score: Tensor | None = None, - ) -> Tensor: - """Run the full diffusion transformer forward pass. - + hidden_states: Tensor, + timestep: Tensor, + return_dict: bool = True, + ) -> Transformer2DModelOutput | tuple[Tensor]: + r""" Args: - x: Noisy visual input tensor. - text_embed: Raw text encoder hidden states. - pooled_text_embed: Pooled text embeddings added to time conditioning. - time: Diffusion time steps. - visual_cu_seqlens: Cumulative sequence lengths for visual tokens. - text_cu_seqlens: Cumulative sequence lengths for text tokens. - visual_rope_pos: 3D positional indices (temporal, height, width) for visual RoPE. - text_rope_pos: 1D positional indices for text RoPE. - scale_factor: RoPE frequency scaling per axis. - sparse_params: Optional sparse/fractal attention parameters. - motion_score: Optional motion score for temporal conditioning. + hidden_states (`torch.Tensor` of shape `(batch_size, num_frames, height, width, 2 * in_visual_dim + 1)`): + Latent tiles in the `(B, T, H, W, C)` layout: the noisy latent, the anchor latent and the anchor mask + concatenated along the channel axis. + timestep (`torch.Tensor` of shape `(batch_size,)`): + Diffusion timestep on the `[0, num_train_timesteps]` scale. + return_dict (`bool`, defaults to `True`): + Whether to return a [`~models.modeling_outputs.Transformer2DModelOutput`] instead of a plain tuple. Returns: - Denoised visual output tensor. + [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple`: + The prediction of shape `(batch_size, num_frames, height, width, out_visual_dim)`. """ checkpoint = torch.is_grad_enabled() and self.gradient_checkpointing - - text_embed = self.text_embeddings(text_embed) if self.use_text else None - time_embed, time_embed_idx = self.time_embeddings(time) - if motion_score is not None and self.use_motion_score: - ms_embed, _ = self.motion_embeddings(motion_score) - time_embed = time_embed + ms_embed - - if self.use_text: - time_embed = time_embed + self.pooled_text_embeddings(pooled_text_embed) - else: - time_embed = time_embed + self.pooled_bias - - visual_embed, visual_cu_seqlens = self.visual_embeddings(x, visual_cu_seqlens) - - if self.use_text: - text_rope = self.text_rope_embeddings(text_rope_pos) - text_time_embed_idx = time_embed_idx.repeat_interleave(torch.diff(text_cu_seqlens), dim=0) - for text_transformer_block in self.text_transformer_blocks: - args = (text_embed, time_embed, text_rope, text_cu_seqlens, text_time_embed_idx) - text_embed = ( - self._gradient_checkpointing_func(text_transformer_block, *args) - if checkpoint - else text_transformer_block(*args) + device = hidden_states.device + + temb = self.time_embeddings(timestep) + self.pooled_bias + + # 1. Patchify the latents and build the rotary embeddings of the token grid + visual_embed = self.visual_embeddings(hidden_states) + visual_shape = visual_embed.shape[1:4] + visual_rope_pos = tuple(torch.arange(size, device=device) for size in visual_shape) + visual_rope = self.visual_rope_embeddings(visual_rope_pos, self.config.scale_factor) + + # 2. Flatten the grid into a token sequence; NABLA groups 8x8 spatial neighbourhoods into attention blocks + sparse_params = None + if self.config.nabla_threshold is not None: + num_frames, height, width = visual_shape + if height % FRACTAL_BLOCK_SIZE or width % FRACTAL_BLOCK_SIZE: + raise ValueError( + f"NABLA attention needs the token grid {(height, width)} to be divisible by {FRACTAL_BLOCK_SIZE}" ) - - visual_shape = visual_embed.shape[:-1] - visual_rope = self.visual_rope_embeddings(visual_shape, visual_rope_pos, scale_factor) - to_fractal = sparse_params["to_fractal"] if sparse_params is not None else False - # [T, 32, 32, 1792] -> [T*32*32, 1792] - visual_embed, visual_rope, visual_cu_seqlens = fractal_flatten( - visual_embed, visual_rope, visual_cu_seqlens, visual_shape, fractal=to_fractal - ) - visual_time_embed_idx = time_embed_idx.repeat_interleave(torch.diff(visual_cu_seqlens), dim=0) - - for visual_transformer_block in self.visual_transformer_blocks: - args = ( - visual_embed, - text_embed, - time_embed, - visual_rope, - visual_cu_seqlens, - text_cu_seqlens, - visual_time_embed_idx, - sparse_params, + group_size = (1, FRACTAL_BLOCK_SIZE, FRACTAL_BLOCK_SIZE) + visual_embed = _local_patch(visual_embed, visual_shape, group_size, dim=1).flatten(1, 2) + visual_rope = _local_patch(visual_rope, visual_shape, group_size, dim=0).flatten(0, 1) + sta_mask = sliding_tile_mask( + num_frames, height // FRACTAL_BLOCK_SIZE, width // FRACTAL_BLOCK_SIZE, self.config.nabla_window, device ) - visual_embed = ( - self._gradient_checkpointing_func(visual_transformer_block, *args) - if checkpoint - else visual_transformer_block(*args) + sparse_params = {"sta_mask": sta_mask, "threshold": self.config.nabla_threshold} + else: + visual_embed = visual_embed.flatten(1, 3) + visual_rope = visual_rope.flatten(0, 2) + + # 3. Transformer blocks + for block in self.visual_transformer_blocks: + if checkpoint: + visual_embed = self._gradient_checkpointing_func(block, visual_embed, temb, visual_rope, sparse_params) + else: + visual_embed = block(visual_embed, temb, visual_rope, sparse_params) + + # 4. Restore the grid and project back to the latent space + if sparse_params is not None: + visual_embed = _local_merge( + visual_embed.reshape(visual_embed.shape[0], -1, FRACTAL_BLOCK_SIZE**2, visual_embed.shape[-1]), + visual_shape, + group_size, + dim=1, ) - visual_embed, visual_cu_seqlens = fractal_unflatten( - visual_embed, visual_cu_seqlens, visual_shape, fractal=to_fractal - ) + else: + visual_embed = visual_embed.reshape(-1, *visual_shape, visual_embed.shape[-1]) + output = self.out_layer(visual_embed, temb) - visual_time_embed_idx = time_embed_idx.repeat_interleave(torch.diff(visual_cu_seqlens), dim=0) - return self.out_layer(visual_embed, text_embed, time_embed, visual_cu_seqlens, visual_time_embed_idx) + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) __all__ = ["Kandinsky6SRTransformer3DModel"] diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index ff4162bdc7a8..fc415647492d 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -452,7 +452,6 @@ "Kandinsky6SRPipeline", "Kandinsky6SRPipelineOutput", "Kandinsky6TI2VAPipeline", - "Kandinsky6SRLatentUpscalerBank", "Kandinsky6TI2VAPipelineOutput", ] _import_structure["z_image"] = [ @@ -788,7 +787,6 @@ Kandinsky6SRPipeline, Kandinsky6SRPipelineOutput, Kandinsky6TI2VAPipeline, - Kandinsky6SRLatentUpscalerBank, Kandinsky6TI2VAPipelineOutput, ) from .krea2 import Krea2Pipeline diff --git a/src/diffusers/pipelines/kandinsky6/__init__.py b/src/diffusers/pipelines/kandinsky6/__init__.py index 18c6f961473c..7e18b45cf44e 100644 --- a/src/diffusers/pipelines/kandinsky6/__init__.py +++ b/src/diffusers/pipelines/kandinsky6/__init__.py @@ -1,5 +1,3 @@ -"""Kandinsky 6 pipelines.""" - from typing import TYPE_CHECKING from ...utils import ( @@ -8,6 +6,7 @@ _LazyModule, get_objects_from_module, is_torch_available, + is_transformers_available, ) @@ -15,28 +14,28 @@ _import_structure = {} try: - if not is_torch_available(): + if not (is_transformers_available() and is_torch_available()): raise OptionalDependencyNotAvailable() except OptionalDependencyNotAvailable: - from ...utils import dummy_pt_objects # noqa F403 + from ...utils import dummy_torch_and_transformers_objects # noqa F403 - _dummy_objects.update(get_objects_from_module(dummy_pt_objects)) + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: - _import_structure["pipeline_kandinsky6_ti2va"] = ["Kandinsky6TI2VAPipeline"] + _import_structure["modeling_latent_upscaler"] = ["Kandinsky6SRLatentUpscalerBank"] _import_structure["pipeline_kandinsky6_sr"] = ["Kandinsky6SRPipeline"] - _import_structure["latent_upscaler"] = ["Kandinsky6SRLatentUpscalerBank"] + _import_structure["pipeline_kandinsky6_ti2va"] = ["Kandinsky6TI2VAPipeline"] _import_structure["pipeline_output"] = ["Kandinsky6SRPipelineOutput", "Kandinsky6TI2VAPipelineOutput"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: - if not is_torch_available(): + if not (is_transformers_available() and is_torch_available()): raise OptionalDependencyNotAvailable() except OptionalDependencyNotAvailable: - from ...utils.dummy_pt_objects import * + from ...utils.dummy_torch_and_transformers_objects import * else: - from .pipeline_kandinsky6_ti2va import Kandinsky6TI2VAPipeline + from .modeling_latent_upscaler import Kandinsky6SRLatentUpscalerBank from .pipeline_kandinsky6_sr import Kandinsky6SRPipeline - from .latent_upscaler import Kandinsky6SRLatentUpscalerBank + from .pipeline_kandinsky6_ti2va import Kandinsky6TI2VAPipeline from .pipeline_output import Kandinsky6SRPipelineOutput, Kandinsky6TI2VAPipelineOutput else: import sys @@ -50,12 +49,3 @@ for name, value in _dummy_objects.items(): setattr(sys.modules[__name__], name, value) - - -__all__ = [ - "Kandinsky6TI2VAPipeline", - "Kandinsky6TI2VAPipelineOutput", - "Kandinsky6SRPipeline", - "Kandinsky6SRLatentUpscalerBank", - "Kandinsky6SRPipelineOutput", -] diff --git a/src/diffusers/pipelines/kandinsky6/latent_upscaler.py b/src/diffusers/pipelines/kandinsky6/latent_upscaler.py deleted file mode 100644 index 50ef471d7f4e..000000000000 --- a/src/diffusers/pipelines/kandinsky6/latent_upscaler.py +++ /dev/null @@ -1,1560 +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. - -"""Kandinsky 6 SR latent-upscaler bank Diffusers component.""" - -import copy -from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass, replace -from typing import Any, Literal, NamedTuple - -import torch -from torch import Tensor, nn -from torch.nn import functional - -from ...configuration_utils import ConfigMixin, register_to_config -from ...models.modeling_utils import ModelMixin - - -def merge_batch_time(x: Tensor) -> Tensor: - """Merge the temporal axis into batch: ``(B, C, T, H, W) -> (B*T, C, H, W)``.""" - b, c, t, h, w = x.shape - return x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) - - -def split_batch_time(x: Tensor, b: int, t: int) -> Tensor: - """Inverse of :func:`merge_batch_time`: ``(B*T, C, H, W) -> (B, C, T, H, W)``.""" - _, c, h, w = x.shape - return x.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4) - - -UpsampleMode = str - - -# Convolution primitives shared by the latent upsampler architectures. - - -DIMS_2 = 2 -DIMS_3 = 3 - -TemporalPadding = Literal["zeros", "replicate", "causal"] -UpsamplePaddingMode = Literal["reflect", "zeros"] - - -class TemporalReplicateConv3d(nn.Conv3d): - """``Conv3d`` that repeats the edge frame along T and zero-pads H/W. - - K-VAE extends the temporal axis by repeating a boundary frame — the encoder - and decoder both seed their causal padding with a replica of frame 0 — while - padding H/W with zeros (``padding_mode: zeros`` in the shipped sidecar). - ``nn.Conv3d`` cannot express that split, because ``padding_mode`` applies to - every padded dim at once; here T is padded explicitly and H/W is left to the - convolution. - - Without this, a zero-padded temporal axis makes the first and last latent - frames see a hole where a neighbour should be — on K-VAE latents that hole - lands on frame 0, which is already the odd one out (it encodes a single - pixel frame, while every later latent aggregates four). - - Parameter names and shapes are those of ``nn.Conv3d``, so checkpoints - trained before this padding change still load. - - Args: - in_channels: Number of input channels. - out_channels: Number of output channels. - kernel_size: Kernel as ``(kt, kh, kw)`` or a single int applied to all dims. - padding: The "same" padding the caller would have passed to ``nn.Conv3d``. - Its H/W entries go to the convolution; its T entry becomes the width - of the replicate pad. - groups: Convolution groups (``in_channels`` for a depthwise conv). - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int | tuple[int, int, int], - padding: int | tuple[int, int, int] = 0, - *, - groups: int = 1, - ) -> None: - """Initialize the convolution with H/W padding only, keeping T for forward.""" - pad_t, pad_h, pad_w = (padding, padding, padding) if isinstance(padding, int) else padding - super().__init__( - in_channels, - out_channels, - kernel_size, - padding=(0, pad_h, pad_w), - groups=groups, - ) - self.temporal_pad = pad_t - - def temporal_pad_lr(self) -> tuple[int, int]: - """Split the temporal pad budget into ``(before, after)`` frame counts.""" - return (self.temporal_pad, self.temporal_pad) - - def forward(self, x: Tensor) -> Tensor: - """Replicate-pad T, then convolve with zero padding on H and W.""" - before, after = self.temporal_pad_lr() - if before or after: - x = functional.pad(x, (0, 0, 0, 0, before, after), mode="replicate") - return super().forward(x) - - -class TemporalCausalConv3d(TemporalReplicateConv3d): - """``TemporalReplicateConv3d`` with K-VAE's causal temporal window. - - K-VAE's ``CausalConv3d`` spends the whole temporal pad budget in front: - ``kT - 1`` replicas of frame 0 precede the clip and nothing follows it, so - each 3-tap kernel reads ``(t-2, t-1, t)``. Moving the symmetric budget - (``pad_t`` per side) to the front reproduces that window exactly, which is - what lets decoder kernels load without recentring — and keeps every frame's - output independent of its future, the property a chunked streaming - inference would rely on. - """ - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int | tuple[int, int, int], - padding: int | tuple[int, int, int] = 0, - *, - groups: int = 1, - ) -> None: - """Validate that the pad budget covers exactly the causal window.""" - super().__init__(in_channels, out_channels, kernel_size, padding, groups=groups) - kernel_t = self.kernel_size[0] - if 2 * self.temporal_pad != kernel_t - 1: - msg = ( - f"a causal window needs the full kT-1 = {kernel_t - 1} pad budget in front, " - f"but padding supplies 2 * {self.temporal_pad}" - ) - raise ValueError(msg) - - def temporal_pad_lr(self) -> tuple[int, int]: - """Put the whole budget before the clip: taps end on the current frame.""" - return (2 * self.temporal_pad, 0) - - -def make_conv(dims: int, temporal_padding: TemporalPadding) -> type[nn.Module]: - """Pick the convolution class for a stack of ``dims``-dimensional convs. - - Args: - dims: 2 for ``Conv2d``, 3 for ``Conv3d``. - temporal_padding: How ``dims == 3`` convolutions extend T. ``"zeros"`` - keeps the stock convolution; ``"replicate"`` repeats the edge frame; - ``"causal"`` pads the past only (K-VAE semantics). - Ignored for ``dims == 2``, which has no temporal axis. - - Returns: - The convolution class to instantiate. - """ - if dims == DIMS_2: - return nn.Conv2d - if temporal_padding == "replicate": - return TemporalReplicateConv3d - if temporal_padding == "causal": - return TemporalCausalConv3d - return nn.Conv3d - - -# Reusable building blocks for latent upsampler architectures. - - -class StochasticDepth(nn.Module): - """Drop residual paths without requiring torchvision.""" - - def __init__(self, p: float, mode: str = "row") -> None: - super().__init__() - if not 0.0 <= p <= 1.0: - raise ValueError(f"drop probability must be in [0, 1], got {p}") - if mode not in ("batch", "row"): - raise ValueError(f"mode must be 'batch' or 'row', got {mode!r}") - self.p = p - self.mode = mode - - def forward(self, x: Tensor) -> Tensor: - if not self.training or self.p == 0.0: - return x - if self.p == 1.0: - return torch.zeros_like(x) - shape = (x.shape[0],) + (1,) * (x.ndim - 1) if self.mode == "row" else (1,) * x.ndim - noise = torch.empty(shape, dtype=x.dtype, device=x.device).bernoulli_(1.0 - self.p) - return x * noise / (1.0 - self.p) - - -class RMSNorm(nn.Module): - """Channel-first Root Mean Square normalization with learnable gamma. - - Args: - dim: Number of channels. - dims: Convolution dimensionality — 2 for ``(C,1,1)``, 3 for ``(C,1,1,1)``. - """ - - def __init__(self, dim: int, dims: int = 2) -> None: - """Initialize RMSNorm. - - Args: - dim: Number of channels. - dims: Convolution dimensionality — 2 for ``(C,1,1)``, 3 for ``(C,1,1,1)``. - """ - super().__init__() - broadcastable_dims = (1, 1) if dims == 2 else (1, 1, 1) - self.scale = dim**0.5 - self.gamma = nn.Parameter(torch.ones(dim, *broadcastable_dims)) - - def forward(self, x: Tensor) -> Tensor: - """Apply RMS normalization along the channel dimension.""" - return functional.normalize(x, dim=1) * self.scale * self.gamma - - -class ModulatedRMSNorm(nn.Module): - """RMSNorm followed by spatial FiLM modulation conditioned on a side tensor. - - Computes ``RMSNorm(x) * conv_y(zq) + conv_b(zq)`` where ``conv_y`` and - ``conv_b`` are 1x1 convolutions of ``zq`` (the conditioning tensor, e.g. - the LQ latent). ``zq`` is aligned to ``x`` by nearest-neighbor - interpolation along the spatial dims, so the same ``zq`` can be reused - across feature pyramids of different resolutions. - - Both convs keep the stock PyTorch initialization, so modulation is active - from the first step — the same regime as K-VAE's ``CachedSpatialNorm3D``, - whose ``conv_y`` / ``conv_b`` are plain 1x1 convs with no custom init. - - Args: - dim: Number of feature channels. - zq_dim: Number of channels in the conditioning tensor. - dims: Convolution dimensionality — 2 for Conv2d, 3 for Conv3d. - """ - - def __init__(self, dim: int, zq_dim: int, dims: int = 2) -> None: - """Initialize ModulatedRMSNorm. - - Args: - dim: Number of feature channels. - zq_dim: Number of channels in the conditioning tensor. - dims: Convolution dimensionality — 2 for Conv2d, 3 for Conv3d. - """ - super().__init__() - conv = nn.Conv2d if dims == 2 else nn.Conv3d - self.norm = RMSNorm(dim, dims=dims) - self.conv_y = conv(zq_dim, dim, kernel_size=1) - self.conv_b = conv(zq_dim, dim, kernel_size=1) - - def forward(self, x: Tensor, zq: Tensor) -> Tensor: - """Apply RMS norm + FiLM modulation by spatially-aligned ``zq``.""" - if zq.shape[2:] != x.shape[2:]: - zq = functional.interpolate(zq, size=x.shape[2:], mode="nearest") - return self.norm(x) * self.conv_y(zq) + self.conv_b(zq) - - -class LayerScale(nn.Module): - """Learnable per-channel scaling applied before residual addition. - - Args: - channels: Number of channels. - init_value: Initial scale value (small for training stability). - dims: Convolution dimensionality — 2 for ``(C,1,1)``, 3 for ``(C,1,1,1)``. - """ - - def __init__(self, channels: int, init_value: float = 1e-6, dims: int = 2) -> None: - """Initialize LayerScale. - - Args: - channels: Number of channels. - init_value: Initial scale value. - dims: Convolution dimensionality — 2 for ``(C,1,1)``, 3 for ``(C,1,1,1)``. - """ - super().__init__() - broadcastable_dims = (1, 1) if dims == 2 else (1, 1, 1) - self.gamma = nn.Parameter(init_value * torch.ones(channels, *broadcastable_dims)) - - def forward(self, x: Tensor) -> Tensor: - """Scale input by learnable per-channel gamma.""" - return x * self.gamma - - -class GRN(nn.Module): - """Global Response Normalization (ConvNeXtV2). - - Aggregates global spatial info per channel and normalizes across channels - to encourage feature diversity and prevent feature collapse. - - Args: - channels: Number of channels. - dims: Convolution dimensionality — 2 for spatial dims ``(2,3)``, 3 for ``(2,3,4)``. - """ - - def __init__(self, channels: int, dims: int = 2) -> None: - """Initialize GRN. - - Args: - channels: Number of channels. - dims: Convolution dimensionality — 2 for spatial dims ``(2,3)``, 3 for ``(2,3,4)``. - """ - super().__init__() - self.spatial_dims: tuple[int, ...] = (2, 3) if dims == 2 else (2, 3, 4) - broadcastable = (1, 1) if dims == 2 else (1, 1, 1) - self.gamma = nn.Parameter(torch.zeros(1, channels, *broadcastable)) - self.beta = nn.Parameter(torch.zeros(1, channels, *broadcastable)) - - def forward(self, x: Tensor) -> Tensor: - """Apply global response normalization.""" - gx = torch.norm(x, p=2, dim=self.spatial_dims, keepdim=True) - nx = gx / (gx.mean(dim=1, keepdim=True) + 1e-6) - return self.gamma * (x * nx) + self.beta + x - - -class ResidualBlock(nn.Module): - """Pre-activation residual block with optional LayerScale, GRN, and StochasticDepth. - - Standard path: ``RMSNorm → SiLU → Conv3x3 → RMSNorm → SiLU → [GRN →] Conv3x3``. - Depthwise path: ``DWConv_KxK → RMSNorm → PWConv1x1 → SiLU → [GRN →] PWConv1x1``. - Skip: identity when ``in_channels == out_channels``, else ``Conv 1x1``. - - When ``zq_dim`` is set, every ``RMSNorm`` is replaced by a - ``ModulatedRMSNorm`` conditioned on a side tensor ``zq``, and the inner - layers are stored as named submodules so ``zq`` can be threaded through - ``forward(x, zq)``. When ``zq_dim is None``, the legacy ``nn.Sequential`` - layout is preserved bit-identically — the parameter names match older - checkpoints, so loading them remains backward compatible. - """ - - def __init__( - self, - in_channels: int, - out_channels: int | None = None, - mid_channels: int | None = None, - dims: int = 2, - layer_scale_init: float | None = None, - *, - grn: bool = False, - stochastic_depth_prob: float = 0.0, - depthwise: bool = False, - kernel_size: int = 3, - zq_dim: int | None = None, - temporal_padding: TemporalPadding = "zeros", - ) -> None: - """Initialize residual block. - - Args: - in_channels: Number of input channels. - out_channels: Number of output channels (defaults to ``in_channels``). - mid_channels: Internal channel width (defaults to ``out_channels``). - dims: Convolution dimensionality — 2 for Conv2d, 3 for Conv3d. - layer_scale_init: Initial LayerScale gamma value. ``None`` disables LayerScale. - grn: Whether to insert Global Response Normalization on mid_channels. - stochastic_depth_prob: Drop probability for stochastic depth (0 = off). - depthwise: Use depthwise separable convolutions instead of standard convs. - kernel_size: Kernel size for the depthwise convolution (only used when ``depthwise=True``). - zq_dim: When set, replace every ``RMSNorm`` with ``ModulatedRMSNorm(zq_dim)``; - the block then accepts a conditioning tensor in ``forward(x, zq)``. - temporal_padding: How ``dims == 3`` convolutions extend T — ``"zeros"`` - (default) or ``"replicate"``, which repeats the edge frame as K-VAE does. - """ - super().__init__() - out_channels = out_channels or in_channels - mid_channels = mid_channels or out_channels - self.depthwise = depthwise - self.has_grn = grn - self.zq_dim = zq_dim - conv = make_conv(dims, temporal_padding) - - if zq_dim is None: - self._build_legacy(in_channels, out_channels, mid_channels, dims, conv, grn=grn, kernel_size=kernel_size) - else: - self._build_modulated( - in_channels, - out_channels, - mid_channels, - dims, - conv, - grn=grn, - kernel_size=kernel_size, - zq_dim=zq_dim, - ) - self.shortcut = ( - nn.Identity() if in_channels == out_channels else conv(in_channels, out_channels, kernel_size=1) - ) - self.layer_scale: nn.Module = ( - LayerScale(out_channels, init_value=layer_scale_init, dims=dims) - if layer_scale_init is not None - else nn.Identity() - ) - self.stochastic_depth: nn.Module = ( - StochasticDepth(stochastic_depth_prob, mode="row") if stochastic_depth_prob > 0 else nn.Identity() - ) - - def _build_legacy( - self, - in_channels: int, - out_channels: int, - mid_channels: int, - dims: int, - conv: type[nn.Module], - *, - grn: bool, - kernel_size: int, - ) -> None: - """Build the original ``nn.Sequential`` layout (parameter names unchanged).""" - if self.depthwise: - ks = kernel_size if dims == 2 else (min(kernel_size, 5), kernel_size, kernel_size) - pad = kernel_size // 2 if dims == 2 else (min(kernel_size, 5) // 2, kernel_size // 2, kernel_size // 2) - layers: list[nn.Module] = [ - conv(in_channels, in_channels, kernel_size=ks, padding=pad, groups=in_channels), - RMSNorm(in_channels, dims=dims), - conv(in_channels, mid_channels, kernel_size=1), - nn.SiLU(), - ] - if grn: - layers.append(GRN(mid_channels, dims=dims)) - layers.append(conv(mid_channels, out_channels, kernel_size=1)) - else: - layers = [ - RMSNorm(in_channels, dims=dims), - nn.SiLU(), - conv(in_channels, mid_channels, kernel_size=3, padding=1), - RMSNorm(mid_channels, dims=dims), - nn.SiLU(), - ] - if grn: - layers.append(GRN(mid_channels, dims=dims)) - layers.append(conv(mid_channels, out_channels, kernel_size=3, padding=1)) - self.block = nn.Sequential(*layers) - - def _build_modulated( - self, - in_channels: int, - out_channels: int, - mid_channels: int, - dims: int, - conv: type[nn.Module], - *, - grn: bool, - kernel_size: int, - zq_dim: int, - ) -> None: - """Build a path with FiLM-modulated RMSNorms accepting ``zq`` in forward.""" - if self.depthwise: - ks = kernel_size if dims == 2 else (min(kernel_size, 5), kernel_size, kernel_size) - pad = kernel_size // 2 if dims == 2 else (min(kernel_size, 5) // 2, kernel_size // 2, kernel_size // 2) - self.dwconv = conv(in_channels, in_channels, kernel_size=ks, padding=pad, groups=in_channels) - self.norm1 = ModulatedRMSNorm(in_channels, zq_dim, dims=dims) - self.pwconv1 = conv(in_channels, mid_channels, kernel_size=1) - self.grn: nn.Module | None = GRN(mid_channels, dims=dims) if grn else None - self.pwconv2 = conv(mid_channels, out_channels, kernel_size=1) - else: - self.norm1 = ModulatedRMSNorm(in_channels, zq_dim, dims=dims) - self.conv1 = conv(in_channels, mid_channels, kernel_size=3, padding=1) - self.norm2 = ModulatedRMSNorm(mid_channels, zq_dim, dims=dims) - self.grn = GRN(mid_channels, dims=dims) if grn else None - self.conv2 = conv(mid_channels, out_channels, kernel_size=3, padding=1) - - def forward(self, x: Tensor, zq: Tensor | None = None) -> Tensor: - """Apply residual block with skip connection. - - Args: - x: Feature tensor. - zq: Optional conditioning tensor for FiLM-modulated norms. Required - when the block was constructed with ``zq_dim`` set; ignored otherwise. - """ - out = self.block(x) if self.zq_dim is None else self.forward_modulated_path(x, zq) # type: ignore - out = self.layer_scale(out) - return self.shortcut(x) + self.stochastic_depth(out) - - def forward_modulated_path( - self, - x: Tensor, - zq: Tensor, - norm1: ModulatedRMSNorm | None = None, - norm2: ModulatedRMSNorm | None = None, - ) -> Tensor: - """Evaluate the residual path with optional scale-specific norms.""" - if self.zq_dim is None: - msg = "forward_modulated_path requires a block with zq_dim" - raise ValueError(msg) - active_norm1 = norm1 if norm1 is not None else self.norm1 - if self.depthwise: - h = self.dwconv(x) - h = active_norm1(h, zq) - h = self.pwconv1(h) - h = functional.silu(h) - if self.grn is not None: - h = self.grn(h) - return self.pwconv2(h) - active_norm2 = norm2 if norm2 is not None else self.norm2 - h = active_norm1(x, zq) - h = functional.silu(h) - h = self.conv1(h) - h = active_norm2(h, zq) - h = functional.silu(h) - if self.grn is not None: - h = self.grn(h) - return self.conv2(h) - - def forward_with_norms( - self, - x: Tensor, - zq: Tensor, - norm1: ModulatedRMSNorm, - norm2: ModulatedRMSNorm, - ) -> Tensor: - """Apply the block with x2-specific modulation and shared transforms.""" - if self.depthwise: - msg = "scale-specific normalization does not support depthwise residual blocks" - raise ValueError(msg) - out = self.forward_modulated_path(x, zq, norm1=norm1, norm2=norm2) - out = self.layer_scale(out) - return self.shortcut(x) + self.stochastic_depth(out) - - def zero_init_residual(self) -> None: - """Initialize the residual path to zero so the block starts as identity.""" - if self.zq_dim is None: - final_conv = self.block[-1] - elif self.depthwise: - final_conv = self.pwconv2 - else: - final_conv = self.conv2 - nn.init.zeros_(final_conv.weight) # type: ignore[arg-type] - if final_conv.bias is not None: - nn.init.zeros_(final_conv.bias) # type: ignore[arg-type] - - -# Spatial upsampling operations: pixel-shuffle, bilinear, and K-VAE PXS v2. - - -VALID_DIMS = {DIMS_2, DIMS_3} - - -def icnr_( - weight: torch.Tensor, - factor: int, - init_fn: Callable[[torch.Tensor], torch.Tensor] = nn.init.kaiming_normal_, -) -> None: - """Initialize sub-pixel-conv weight so initial output equals nearest 2x. - - Implements the ICNR scheme of Aitken et al. (2017): the ``factor**2`` - sub-channel groups consumed by ``PixelShuffle(factor)`` are initialized - with identical kernels, so on the first forward pass the output is a - pure nearest-neighbor upsample. This eliminates checkerboard artifacts - at the start of training. - - Args: - weight: Conv weight tensor of shape ``(C_out * factor**2, C_in, *kernel)``. - Modified in-place. - factor: Upscale factor used by the downstream ``PixelShuffle``. - init_fn: Base initializer applied to the underlying ``C_out`` kernel - before replication. Defaults to Kaiming normal. - - Raises: - ValueError: If ``weight.shape[0]`` is not divisible by ``factor**2``. - """ - out_c = weight.shape[0] // factor**2 - if out_c * factor**2 != weight.shape[0]: - msg = f"weight.shape[0] ({weight.shape[0]}) must be divisible by factor**2 ({factor**2})" - raise ValueError(msg) - base = torch.empty(out_c, *weight.shape[1:], device=weight.device, dtype=weight.dtype) - init_fn(base) - weight.data.copy_(base.repeat_interleave(factor**2, dim=0)) - - -def spatial_nearest_2x(x: torch.Tensor, dims: int, factor: int) -> torch.Tensor: - """Apply nearest-neighbor upsampling to H, W only, preserving B (and T). - - For 5D ``(B, C, T, H, W)`` tensors, T is merged into batch, interpolated, - and unmerged. For 4D tensors, interpolation runs directly. - - Args: - x: Input tensor of shape ``(B, C, H, W)`` or ``(B, C, T, H, W)``. - dims: Tensor dimensionality — 2 for 4D, 3 for 5D. - factor: Spatial upscale factor. - - Returns: - Upsampled tensor with H, W scaled by ``factor``. - """ - if dims == DIMS_3: - b, _c, t, _h, _w = x.shape - x = merge_batch_time(x) - x = functional.interpolate(x, scale_factor=factor, mode="nearest") - return split_batch_time(x, b, t) - return functional.interpolate(x, scale_factor=factor, mode="nearest") - - -class PixelShuffleND(nn.Module): - """Conv + pixel-shuffle that upscales only H and W. - - Internally applies a convolution to expand channels by ``upscale_factor ** 2``, - then rearranges channels into spatial dimensions. - - Supports both 4D ``(B, C, H, W)`` and 5D ``(B, C, T, H, W)`` tensors. - The temporal dimension is always preserved. - - Args: - channels: Number of input (and output) channels. - dims: Tensor dimensionality — 2 for 4D input, 3 for 5D (video) input. - factor: Spatial upscale factor applied to both H and W. - temporal_mix: Whether the 3D convolution mixes across the temporal dimension. - When ``False`` and ``dims == 3``, uses ``kernel_size=(1, 3, 3)`` so that - each frame is convolved independently in H/W. Ignored when ``dims == 2``. - Defaults to ``True`` (full 3x3x3 kernel). - icnr: Apply ICNR initialization (Aitken et al., 2017) to the conv weight - and zero its bias, so the module starts as a nearest-neighbor upsample. - Removes checkerboard at init. Defaults to ``False``. - temporal_padding: How the ``temporal_mix=True`` kernel extends T — - ``"zeros"`` (default) or ``"replicate"``, matching K-VAE's edge repeat. - """ - - def __init__( - self, - channels: int, - dims: int, - factor: int = 2, - *, - temporal_mix: bool = True, - icnr: bool = False, - temporal_padding: TemporalPadding = "zeros", - ) -> None: - """Initialize PixelShuffleND.""" - super().__init__() - if dims not in VALID_DIMS: - msg = f"dims must be one of {VALID_DIMS}, got {dims}" - raise ValueError(msg) - self.dims = dims - self.factor = factor - - if dims == DIMS_2: - self.conv = nn.Conv2d(channels, channels * factor**2, kernel_size=3, padding=1) - elif temporal_mix: - conv3d = make_conv(DIMS_3, temporal_padding) - self.conv = conv3d(channels, channels * factor**2, kernel_size=3, padding=1) - else: - self.conv = nn.Conv3d(channels, channels * factor**2, kernel_size=(1, 3, 3), padding=(0, 1, 1)) - - if icnr: - icnr_(self.conv.weight, factor) - if self.conv.bias is not None: - nn.init.zeros_(self.conv.bias) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Expand channels via convolution, then pixel-shuffle into spatial dimensions.""" - x = self.conv(x) - if self.dims == DIMS_3: - b, c_by_factor2, t, h, w = x.shape - c = c_by_factor2 // self.factor**2 - return ( - x.reshape(b, c, self.factor, self.factor, t, h, w) - .permute(0, 1, 4, 5, 2, 6, 3) - .reshape(b, c, t, h * self.factor, w * self.factor) - ) - return functional.pixel_shuffle(x, self.factor) - - -class BilinearUpsampleND(nn.Module): - """Parameter-free bilinear spatial upsampling for 4D and 5D tensors. - - Upscales only H and W dimensions. For 5D ``(B, C, T, H, W)`` tensors, - the temporal dimension is merged into the batch, interpolated in 2D, - and unmerged. - - Args: - dims: Tensor dimensionality — 2 for 4D input, 3 for 5D (video) input. - factor: Spatial upscale factor applied to both H and W. - """ - - def __init__(self, dims: int, factor: int = 2) -> None: - """Initialize BilinearUpsampleND.""" - super().__init__() - if dims not in VALID_DIMS: - msg = f"dims must be one of {VALID_DIMS}, got {dims}" - raise ValueError(msg) - self.dims = dims - self.factor = factor - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Apply bilinear interpolation to spatial dimensions.""" - if self.dims == DIMS_3: - b, _c, t, _h, _w = x.shape - x = merge_batch_time(x) - x = functional.interpolate(x, scale_factor=self.factor, mode="bilinear", align_corners=False) - return split_batch_time(x, b, t) - return functional.interpolate(x, scale_factor=self.factor, mode="bilinear", align_corners=False) - - -class PXSv2UpsampleND(nn.Module): - """K-VAE PXS v2 spatial upsample: ``nearest 2x + Conv_(1,3,3) residual + Conv_1x1``. - - Replicates ``CachedPXSUpsample.spatial_upsample_NEW`` followed by the - post-``linear`` 1x1x1 conv from K-VAE-3D-2.0. The base path is a - parameter-free nearest upsample; the residual conv adds learnable detail - on top, and a final pointwise conv mixes channels. - - Spatial conv is hard-coded to ``kernel=(1, 3, 3)`` for ``dims == 3``, - matching K-VAE: per-frame, no temporal mixing through the upsample itself. - Temporal correlation is restored by surrounding ResBlocks. - - Args: - channels: Number of input (and output) channels. - dims: Tensor dimensionality — 2 for 4D, 3 for 5D. - factor: Spatial upscale factor for H and W. - with_linear: Append a final 1x1 (or 1x1x1) pointwise conv. Disable - when this module is composed inside a parent that owns its own - post-mixer (e.g. ``PXSv2HybridUpsampleND``). Defaults to ``True``. - padding_mode: Edge handling for the residual conv. ``"reflect"`` is the - historical default; ``"zeros"`` matches the production K-VAE, whose - ``CachedPXSUpsample`` resolves ``padding_mode or 'reflect'`` against a - sidecar that passes ``'zeros'``. - """ - - def __init__( - self, - channels: int, - dims: int, - factor: int = 2, - *, - with_linear: bool = True, - padding_mode: UpsamplePaddingMode = "reflect", - ) -> None: - """Initialize PXSv2UpsampleND.""" - super().__init__() - if dims not in VALID_DIMS: - msg = f"dims must be one of {VALID_DIMS}, got {dims}" - raise ValueError(msg) - self.dims = dims - self.factor = factor - - if dims == DIMS_2: - self.spatial_conv: nn.Module = nn.Conv2d( - channels, - channels, - kernel_size=3, - padding=1, - padding_mode=padding_mode, - ) - self.linear: nn.Module = nn.Conv2d(channels, channels, kernel_size=1) if with_linear else nn.Identity() - else: - self.spatial_conv = nn.Conv3d( - channels, - channels, - kernel_size=(1, 3, 3), - padding=(0, 1, 1), - padding_mode=padding_mode, - ) - self.linear = nn.Conv3d(channels, channels, kernel_size=1) if with_linear else nn.Identity() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Compute ``Linear( NN_2x(x) + Conv·NN_2x(x) )``.""" - up = spatial_nearest_2x(x, self.dims, self.factor) - return self.linear(up + self.spatial_conv(up)) - - -class PXSv2HybridUpsampleND(nn.Module): - """Hybrid upsample: ICNR-PXS branch parallel to PXS v2 branch, mixed by 1x1. - - Combines the learnable sub-pixel-conv branch (``PixelShuffleND``, optionally - ICNR-initialized so it starts as nearest) with the K-VAE PXS v2 branch - (``nearest + Conv_(1,3,3) residual``), then mixes both via a final 1x1 - pointwise conv. - - Args: - channels: Number of input (and output) channels. - dims: Tensor dimensionality — 2 for 4D, 3 for 5D. - factor: Spatial upscale factor for H and W. - temporal_mix: Forwarded to the inner ``PixelShuffleND`` (controls 3x3x3 - vs (1,3,3) kernel for ``dims == 3``). The PXS v2 branch always uses - (1,3,3) regardless. Defaults to ``True``. - icnr: Apply ICNR initialization to the inner ``PixelShuffleND`` so it - starts as nearest, matching the PXS v2 branch and removing - checkerboard at init. Defaults to ``True``. - temporal_padding: Forwarded to the inner ``PixelShuffleND``. - padding_mode: Forwarded to the inner ``PXSv2UpsampleND``. - """ - - def __init__( - self, - channels: int, - dims: int, - factor: int = 2, - *, - temporal_mix: bool = True, - icnr: bool = True, - temporal_padding: TemporalPadding = "zeros", - padding_mode: UpsamplePaddingMode = "reflect", - ) -> None: - """Initialize PXSv2HybridUpsampleND.""" - super().__init__() - if dims not in VALID_DIMS: - msg = f"dims must be one of {VALID_DIMS}, got {dims}" - raise ValueError(msg) - self.pxs = PixelShuffleND( - channels, - dims=dims, - factor=factor, - temporal_mix=temporal_mix, - icnr=icnr, - temporal_padding=temporal_padding, - ) - self.v3 = PXSv2UpsampleND( - channels, - dims=dims, - factor=factor, - with_linear=False, - padding_mode=padding_mode, - ) - if dims == DIMS_2: - self.linear: nn.Module = nn.Conv2d(channels, channels, kernel_size=1) - else: - self.linear = nn.Conv3d(channels, channels, kernel_size=1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Compute ``Linear( PXS(x) + (NN_2x + Conv·NN_2x)(x) )``.""" - return self.linear(self.pxs(x) + self.v3(x)) - - -# Construction and execution helpers for the cascaded multi-scale model. - - -STAGE_FACTOR = 2 - - -@dataclass -class UpsampleSpec: - """Configuration shared by the two spatial upsample stages.""" - - mode: Literal["pixel_shuffle", "bilinear", "pxs_v2", "pxs_v2_hybrid"] - dims: int - temporal_mix: bool - icnr: bool - temporal_padding: TemporalPadding = "zeros" - padding_mode: UpsamplePaddingMode = "reflect" - - -@dataclass -class ResidualStackSpec: - """Configuration shared by residual stacks at all model resolutions. - - in_channels: input width of the stack's first block; None means use channels. - """ - - channels: int - mid_channels: int - total_blocks: int - stochastic_depth_rate: float - dims: int - layer_scale_init: float | None - grn: bool - depthwise: bool - kernel_size: int - zq_dim: int | None - in_channels: int | None = None - temporal_padding: TemporalPadding = "zeros" - - -class BlockSequence(NamedTuple): - """Non-registering residual-block execution view.""" - - blocks: nn.Sequential - - -def build_stem( - in_channels: int, - hidden_channels: int, - dims: int, - *, - bare: bool, - temporal_padding: TemporalPadding = "zeros", -) -> nn.Sequential: - """Build an input projection: bare conv, or the legacy conv + RMSNorm + SiLU.""" - conv = make_conv(dims, temporal_padding) - projection = conv(in_channels, hidden_channels, kernel_size=3, padding=1) - if bare: - return nn.Sequential(projection) - return nn.Sequential(projection, RMSNorm(hidden_channels, dims=dims), nn.SiLU()) - - -def build_output_head( - hidden_channels: int, - out_channels: int, - dims: int, - *, - zq_dim: int | None, - temporal_padding: TemporalPadding = "zeros", -) -> nn.Sequential: - """Build a norm + SiLU + conv output head, zq-conditioned when ``zq_dim`` is set.""" - conv = make_conv(dims, temporal_padding) - norm: nn.Module = RMSNorm(hidden_channels, dims=dims) - if zq_dim is not None: - norm = ModulatedRMSNorm(hidden_channels, zq_dim, dims=dims) - return nn.Sequential(norm, nn.SiLU(), conv(hidden_channels, out_channels, kernel_size=3, padding=1)) - - -def apply_output_head(head: nn.Sequential, x: Tensor, zq: Tensor | None) -> Tensor: - """Run an output head, threading ``zq`` through a modulated first norm.""" - norm = head[0] - x = norm(x, zq) if isinstance(norm, ModulatedRMSNorm) else norm(x) - x = head[1](x) - return head[2](x) - - -def build_stage_upsample(channels: int, spec: UpsampleSpec) -> nn.Module: - """Build one 2x spatial upsample stage.""" - if spec.mode == "pixel_shuffle": - return PixelShuffleND( - channels, - dims=spec.dims, - factor=STAGE_FACTOR, - temporal_mix=spec.temporal_mix, - icnr=spec.icnr, - temporal_padding=spec.temporal_padding, - ) - if spec.mode == "bilinear": - return BilinearUpsampleND(dims=spec.dims, factor=STAGE_FACTOR) - if spec.mode == "pxs_v2": - return PXSv2UpsampleND(channels, dims=spec.dims, factor=STAGE_FACTOR, padding_mode=spec.padding_mode) - return PXSv2HybridUpsampleND( - channels, - dims=spec.dims, - factor=STAGE_FACTOR, - temporal_mix=spec.temporal_mix, - icnr=spec.icnr, - temporal_padding=spec.temporal_padding, - padding_mode=spec.padding_mode, - ) - - -def build_residual_stack(count: int, start_index: int, spec: ResidualStackSpec) -> nn.Sequential: - """Build residual blocks with linearly scheduled stochastic depth.""" - blocks: list[nn.Module] = [] - for index in range(count): - drop_probability = spec.stochastic_depth_rate * (start_index + index) / max(spec.total_blocks - 1, 1) - block_in = spec.in_channels if index == 0 and spec.in_channels is not None else spec.channels - blocks.append( - ResidualBlock( - block_in, - out_channels=spec.channels, - mid_channels=spec.mid_channels, - dims=spec.dims, - layer_scale_init=spec.layer_scale_init, - grn=spec.grn, - stochastic_depth_prob=drop_probability, - depthwise=spec.depthwise, - kernel_size=spec.kernel_size, - zq_dim=spec.zq_dim, - temporal_padding=spec.temporal_padding, - ) - ) - return nn.Sequential(*blocks) - - -def apply_block_sequence(x: Tensor, zq: Tensor | None, sequence: BlockSequence) -> Tensor: - """Run a residual-block stack.""" - for block in sequence.blocks: - x = block(x, zq) if zq is not None else block(x) - return x - - -# Scale-specific x2 adaptation modules for the cascaded latent upscaler. - - -X2TailMode = Literal["shared", "scale_specific_norm", "private", "private_full"] -X2FinisherMode = Literal["none", "linear", "pxs_residual"] - - -class SharedSecondStage(NamedTuple): - """Non-registering view of the x4 second-stage modules.""" - - upsample: nn.Module - blocks: nn.Sequential - output_proj: nn.Sequential - mid_blocks: nn.Sequential | None = None - - -class ScaleSpecificBlockNorms(nn.Module): - """Private normalization copies for one shared residual block.""" - - def __init__(self, block: ResidualBlock) -> None: - """Copy the modulated norms from a non-depthwise residual block.""" - super().__init__() - if block.zq_dim is None or block.depthwise: - msg = "scale-specific norms require non-depthwise modulated residual blocks" - raise ValueError(msg) - self.norm1 = copy.deepcopy(block.norm1) - self.norm2 = copy.deepcopy(block.norm2) - - -def require_pxs_upsample(source: nn.Module) -> PXSv2UpsampleND: - """Require a pxs_v2 upsample stage as the finisher donor.""" - if not isinstance(source, PXSv2UpsampleND): - msg = f"x2_finisher requires a PXSv2UpsampleND source, got {type(source).__name__}" - raise TypeError(msg) - return source - - -class X2Finisher(nn.Module): - """The upsample-stage conv pair applied at the adapter grid without the nearest-2x resize.""" - - def __init__(self, source: nn.Module, mode: X2FinisherMode) -> None: - """Clone the finisher convs from a pxs_v2 upsample stage.""" - super().__init__() - if mode == "none": - msg = "X2Finisher must not be built with x2_finisher='none'" - raise ValueError(msg) - donor = require_pxs_upsample(source) - self.mode = mode - self.linear = copy.deepcopy(donor.linear) - self.spatial_conv = copy.deepcopy(donor.spatial_conv) if mode == "pxs_residual" else None - - def forward(self, x: Tensor) -> Tensor: - """Apply the cloned convs on the unchanged grid.""" - if self.spatial_conv is not None: - x = x + self.spatial_conv(x) - return self.linear(x) - - -class X2Branch(nn.Module): - """Modules exclusive to the frozen-e115 x2 adaptation path.""" - - def __init__( - self, - adapter: nn.Sequential, - tail_mode: X2TailMode, - shared_stage: SharedSecondStage, - *, - dims: int, - finisher: X2Finisher | None = None, - ) -> None: - """Build the selected scale-specific capacity around a shared stage.""" - super().__init__() - self.adapter = adapter - self.tail_mode = tail_mode - self.dims = dims - self.finisher = finisher - - self.post_norms: nn.ModuleList | None = None - self.output_norm: RMSNorm | None = None - self.private_mid_blocks: nn.Sequential | None = None - self.private_upsample: nn.Module | None = None - self.private_blocks: nn.Sequential | None = None - self.private_output_proj: nn.Sequential | None = None - - if tail_mode == "scale_specific_norm": - self.post_norms = nn.ModuleList([ScaleSpecificBlockNorms(block) for block in shared_stage.blocks]) # type: ignore - self.output_norm = copy.deepcopy(shared_stage.output_proj[0]) # type: ignore - elif tail_mode in ("private", "private_full"): - self.private_upsample = copy.deepcopy(shared_stage.upsample) - self.private_blocks = copy.deepcopy(shared_stage.blocks) - self.private_output_proj = copy.deepcopy(shared_stage.output_proj) - if tail_mode == "private_full": - if shared_stage.mid_blocks is None: - msg = "x2_tail_mode='private_full' requires the shared stage to expose mid_blocks" - raise ValueError(msg) - self.private_mid_blocks = copy.deepcopy(shared_stage.mid_blocks) - - def private_stage(self, shared_stage: SharedSecondStage) -> SharedSecondStage: - """Select the shared stage or the registered private copy.""" - if self.tail_mode not in ("private", "private_full"): - return shared_stage - if self.private_upsample is None or self.private_blocks is None or self.private_output_proj is None: - msg = "private x2 tail modules are not initialized" - raise RuntimeError(msg) - return SharedSecondStage( - upsample=self.private_upsample, - blocks=self.private_blocks, - output_proj=self.private_output_proj, - ) - - def apply_post_blocks(self, x: Tensor, zq: Tensor | None, stage: SharedSecondStage) -> Tensor: - """Apply post blocks with shared or x2-specific normalization.""" - for index, block in enumerate(stage.blocks): - if self.tail_mode == "scale_specific_norm": - if zq is None or self.post_norms is None: - msg = "scale-specific normalization requires zq and private norms" - raise RuntimeError(msg) - norms = self.post_norms[index] - x = block.forward_with_norms(x, zq, norm1=norms.norm1, norm2=norms.norm2) # type: ignore - elif zq is None: - x = block(x) - else: - x = block(x, zq) - return x - - def apply_output_projection(self, x: Tensor, zq: Tensor | None, stage: SharedSecondStage) -> Tensor: - """Project features with a shared or scale-specific final norm.""" - if self.tail_mode != "scale_specific_norm": - return apply_output_head(stage.output_proj, x, zq) - if self.output_norm is None: - msg = "scale-specific output norm is not initialized" - raise RuntimeError(msg) - x = self.output_norm(x) - x = stage.output_proj[1](x) - return stage.output_proj[2](x) - - def apply_block_stack(self, blocks: nn.Sequential, x: Tensor, zq: Tensor | None) -> Tensor: - """Run a residual stack with the branch's zq-threading convention.""" - for block in blocks: - x = block(x, zq) if zq is not None else block(x) - return x - - def apply_adapter(self, x: Tensor, zq: Tensor | None) -> Tensor: - """Run the x2-only residual adapter.""" - return self.apply_block_stack(self.adapter, x, zq) - - def apply_mid_blocks(self, x: Tensor, zq: Tensor | None) -> Tensor: - """Run the private mid blocks ahead of the second-stage upsample (private_full only).""" - if self.private_mid_blocks is None: - return x - return self.apply_block_stack(self.private_mid_blocks, x, zq) - - def forward_second_stage( - self, - x: Tensor, - zq: Tensor | None, - batch_size: int, - frames: int, - shared_stage: SharedSecondStage, - ) -> Tensor: - """Run the selected x2 second stage without registering shared aliases.""" - stage = self.private_stage(shared_stage) - if self.finisher is not None: - x = self.finisher(x) - x = self.apply_mid_blocks(x, zq) - x = stage.upsample(x) - x = self.apply_post_blocks(x, zq, stage) - x = self.apply_output_projection(x, zq, stage) - if self.dims == 2: - return split_batch_time(x, batch_size, frames) - return x - - -# Multi-scale cascaded upsampler: two 2x stages with intermediate supervision. - - -class MultiScaleUpsampler(nn.Module): - """Cascaded 2x+2x VAE-latent upsampler with an optional x2 entry. - - The x4 route keeps the original two-stage architecture. The x2 route enters - after ``mid_blocks`` through a private stem, optional residual adapter, and - a shared, scale-normalized, or private second stage. ``x2_tail_mode`` - ``'private_full'`` instead routes through the branch's own warm copies of - ``mid_blocks`` and the second stage, behind an optional finisher cloned - from ``upsample_1``. - """ - - def __init__( - self, - in_channels: int = 16, - hidden_channels: int = 64, - bottleneck_channels: int | None = None, - num_pre_blocks: int = 2, - num_mid_blocks: int = 6, - num_post_blocks: int = 8, - *, - input_skip: bool = True, - upsample_mode: UpsampleMode = "pixel_shuffle", - temporal_mix: bool = True, - icnr: bool = False, - dims: int = 3, - expand_ratio: int = 4, - kernel_size: int = 3, - layer_scale_init: float | None = None, - grn: bool = False, - stochastic_depth_rate: float = 0.0, - depthwise: bool = False, - zq_dim: int | None = None, - bare_stem: bool = False, - modulated_output_proj: bool = False, - global_skip: bool = False, - enable_x2_entry: bool = False, - x2_adapter_blocks: int = 0, - x2_tail_mode: X2TailMode = "shared", - x2_adapter_sources: tuple[int, ...] | None = None, - x2_finisher: X2FinisherMode = "none", - stage_channels: tuple[int, int, int] | None = None, - temporal_padding: TemporalPadding = "zeros", - upsample_padding_mode: UpsamplePaddingMode = "reflect", - ) -> None: - """Initialize MultiScaleUpsampler with the given architecture parameters.""" - super().__init__() - self.in_channels = in_channels - self.hidden_channels = hidden_channels - self.upscale_factor = 4 - self.input_skip = input_skip - self.dims = dims - self.zq_dim = zq_dim - self.global_skip = global_skip - self.enable_x2_entry = enable_x2_entry - self.x2_tail_mode = x2_tail_mode - self.x2_adapter_sources = x2_adapter_sources - - if modulated_output_proj and zq_dim is None: - msg = "modulated_output_proj requires zq_dim (modulated_norm)" - raise ValueError(msg) - if enable_x2_entry and global_skip: - msg = "global_skip does not support the x2 entry" - raise ValueError(msg) - if enable_x2_entry and modulated_output_proj and x2_tail_mode != "private_full": - msg = "modulated_output_proj with the x2 entry requires x2_tail_mode='private_full'" - raise ValueError(msg) - if x2_adapter_sources is not None and len(x2_adapter_sources) != x2_adapter_blocks: - msg = "x2_adapter_sources must list one pre_blocks index per adapter block" - raise ValueError(msg) - if x2_adapter_sources is not None and any(not 0 <= index < num_pre_blocks for index in x2_adapter_sources): - msg = "x2_adapter_sources indices must be within [0, num_pre_blocks)" - raise ValueError(msg) - if stage_channels is not None: - if enable_x2_entry and x2_tail_mode != "private_full": - msg = "stage_channels with enable_x2_entry requires x2_tail_mode='private_full'" - raise ValueError(msg) - if input_skip: - msg = "stage_channels is incompatible with input_skip" - raise ValueError(msg) - if bottleneck_channels is not None: - msg = "stage_channels is incompatible with bottleneck_channels" - raise ValueError(msg) - if hidden_channels != stage_channels[0]: - msg = f"hidden_channels ({hidden_channels}) must equal stage_channels[0] ({stage_channels[0]})" - raise ValueError(msg) - self.stage_channels = stage_channels - w1, w2, w3 = stage_channels if stage_channels is not None else (hidden_channels,) * 3 - - if input_skip and hidden_channels % in_channels != 0: - msg = ( - f"hidden_channels ({hidden_channels}) must be divisible by in_channels ({in_channels}) for input_skip" - ) - raise ValueError(msg) - - # Projection layers - self.input_proj = build_stem(in_channels, w1, dims, bare=bare_stem, temporal_padding=temporal_padding) - self.mid_output_head = build_output_head(w2, in_channels, dims, zq_dim=None, temporal_padding=temporal_padding) - self.output_proj = build_output_head( - w3, - in_channels, - dims, - zq_dim=zq_dim if modulated_output_proj else None, - temporal_padding=temporal_padding, - ) - - # Two 2x upsample stages - upsample_spec = UpsampleSpec( - dims=dims, - mode=upsample_mode, - temporal_mix=temporal_mix, - icnr=icnr, - temporal_padding=temporal_padding, - padding_mode=upsample_padding_mode, - ) - self.upsample_1 = build_stage_upsample(w1, upsample_spec) - self.upsample_2 = build_stage_upsample(w2, upsample_spec) - - # Residual stacks with linear stochastic depth across all stages; the - # first block of a stage narrows from the previous stage width. - total_blocks = num_pre_blocks + num_mid_blocks + num_post_blocks - - def stage_spec(channels: int, first_in: int | None) -> ResidualStackSpec: - return ResidualStackSpec( - channels=channels, - mid_channels=bottleneck_channels if bottleneck_channels is not None else channels * expand_ratio, - total_blocks=total_blocks, - stochastic_depth_rate=stochastic_depth_rate, - dims=dims, - layer_scale_init=layer_scale_init, - grn=grn, - depthwise=depthwise, - kernel_size=kernel_size, - zq_dim=zq_dim, - in_channels=first_in, - temporal_padding=temporal_padding, - ) - - self.pre_blocks = build_residual_stack(num_pre_blocks, 0, stage_spec(w1, None)) - self.mid_blocks = build_residual_stack( - num_mid_blocks, num_pre_blocks, stage_spec(w2, w1 if w1 != w2 else None) - ) - self.post_blocks = build_residual_stack( - num_post_blocks, num_pre_blocks + num_mid_blocks, stage_spec(w3, w2 if w2 != w3 else None) - ) - - # Build x2-exclusive modules after the complete x4 path so enabling an x2 - # entry does not perturb the base-path RNG initialization. - self.mid_input_proj: nn.Sequential | None = None - self.x2_branch: X2Branch | None = None - if enable_x2_entry: - self.mid_input_proj = build_stem( - in_channels, hidden_channels, dims, bare=bare_stem, temporal_padding=temporal_padding - ) - adapter_spec = replace( - stage_spec(w1, None), - total_blocks=max(x2_adapter_blocks, 1), - stochastic_depth_rate=0.0, - ) - adapter = build_residual_stack(x2_adapter_blocks, 0, adapter_spec) - for block in adapter: - block.zero_init_residual() # type: ignore[attr-defined] - finisher = X2Finisher(self.upsample_1, x2_finisher) if x2_finisher != "none" else None - self.x2_branch = X2Branch( - adapter, - x2_tail_mode, - self.shared_second_stage(), - dims=dims, - finisher=finisher, - ) - - if global_skip: - # Zero heads so step 0 reproduces the nearest base exactly at both scales. - for head in (self.mid_output_head, self.output_proj): - final_conv = head[-1] - nn.init.zeros_(final_conv.weight) # type: ignore[arg-type] - nn.init.zeros_(final_conv.bias) # type: ignore[arg-type] - - def shared_second_stage(self) -> SharedSecondStage: - """Return a non-registering view of the e115-compatible second stage.""" - return SharedSecondStage( - upsample=self.upsample_2, - blocks=self.post_blocks, - output_proj=self.output_proj, - mid_blocks=self.mid_blocks, - ) - - def _project_with_skip(self, x: Tensor, proj: nn.Module) -> Tensor: - """Apply a projection and optionally add the channel-repeated input skip.""" - projected = proj(x) - if self.input_skip: - repeats = projected.shape[1] // x.shape[1] - projected = projected + x.repeat_interleave(repeats, dim=1) - return projected - - def _head_x4(self, x: Tensor, zq: Tensor | None) -> Tensor: - """x4 entry: ``input_proj -> input_skip -> pre_blocks -> upsample_1`` (-> 2x grid).""" - x = self._project_with_skip(x, self.input_proj) - x = apply_block_sequence(x, zq, BlockSequence(self.pre_blocks)) - return self.upsample_1(x) - - def _head_x2(self, x: Tensor, zq: Tensor | None) -> Tensor: - """Project a genuine x2 latent and apply the scale-specific adapter.""" - if self.mid_input_proj is None or self.x2_branch is None: - msg = "x2 head requires enable_x2_entry=True" - raise ValueError(msg) - x = self._project_with_skip(x, self.mid_input_proj) - return self.x2_branch.apply_adapter(x, zq) - - def _tail( - self, - x: Tensor, - zq: Tensor | None, - b: int, - t: int, - *, - return_intermediates: bool, - ) -> dict[str, torch.Tensor] | torch.Tensor: - """x4 continuation: ``mid_blocks -> [2x head] -> second stage``.""" - x = apply_block_sequence(x, zq, BlockSequence(self.mid_blocks)) - - # 2x supervision branch (only when the caller will use it) - z_2x = None - if return_intermediates: - z_2x_inner = self.mid_output_head(x) - z_2x = split_batch_time(z_2x_inner, b, t) if self.dims == 2 else z_2x_inner - - z_4x = self._second_stage(x, zq, b, t) - if return_intermediates: - return {"2x": z_2x, "4x": z_4x} # type: ignore[return-value] - return z_4x - - def _second_stage(self, x: Tensor, zq: Tensor | None, b: int, t: int) -> Tensor: - """Shared second stage: ``upsample_2 -> post_blocks -> output_proj`` — 2x of the current grid.""" - x = self.upsample_2(x) - x = apply_block_sequence(x, zq, BlockSequence(self.post_blocks)) - x = apply_output_head(self.output_proj, x, zq) - return split_batch_time(x, b, t) if self.dims == 2 else x - - def forward( - self, - z: torch.Tensor, - *, - entry: Literal["x4", "x2"] = "x4", - return_intermediates: bool | None = None, - ) -> dict[str, torch.Tensor] | torch.Tensor: - """Upsample through the x4 cascade or the configured x2 entry.""" - if entry not in ("x4", "x2"): - msg = f"entry must be 'x4' or 'x2', got {entry!r}" - raise ValueError(msg) - if entry == "x2" and not self.enable_x2_entry: - msg = "forward(entry='x2') requires the model to be built with enable_x2_entry=True" - raise ValueError(msg) - if entry == "x2" and return_intermediates: - msg = "return_intermediates=True is not supported for entry='x2' — the x2 path has no 2x sub-target" - raise ValueError(msg) - if return_intermediates is None: - return_intermediates = self.training - - b, _, t, _, _ = z.shape - x = merge_batch_time(z) if self.dims == 2 else z - zq = x if self.zq_dim is not None else None - - if entry == "x2": - if self.x2_branch is None: - msg = "x2 branch requires enable_x2_entry=True" - raise ValueError(msg) - x = self._head_x2(x, zq) - return self.x2_branch.forward_second_stage(x, zq, b, t, self.shared_second_stage()) - x = self._head_x4(x, zq) - out = self._tail(x, zq, b, t, return_intermediates=return_intermediates) - if not self.global_skip: - return out - return self._add_nearest_base(out, z) - - def _add_nearest_base( - self, - out: dict[str, torch.Tensor] | torch.Tensor, - z: torch.Tensor, - ) -> dict[str, torch.Tensor] | torch.Tensor: - """Add the parameter-free nearest base of the input latent to each output.""" - if isinstance(out, dict): - return { - "2x": out["2x"] + spatial_nearest_2x(z, DIMS_3, STAGE_FACTOR), - "4x": out["4x"] + spatial_nearest_2x(z, DIMS_3, self.upscale_factor), - } - return out + spatial_nearest_2x(z, DIMS_3, self.upscale_factor) - - -# Factory function for building an upsampler model. - - -def build_upsampler(config: Mapping[str, Any]) -> MultiScaleUpsampler: - """Create a `MultiScaleUpsampler` from a serialized model config. - - Args: - config: Multi-scale model configuration, as stored under `latent_upscaler/config.json`'s - `models[i].model`. Fields absent from older published configs fall back to - `MultiScaleUpsampler`'s own defaults. - - Returns: - Initialized upsampler module. - - Raises: - ValueError: If `config["architecture"]` is not `"multi_scale"`. - """ - if config["architecture"] != "multi_scale": - raise ValueError(f"unsupported latent-upscaler architecture={config['architecture']!r}") - stage_channels = config["stage_channels"] - x2_adapter_sources = config.get("x2_adapter_sources") - return MultiScaleUpsampler( - in_channels=config["in_channels"], - hidden_channels=config["hidden_channels"], - bottleneck_channels=config["bottleneck_channels"], - num_pre_blocks=config["num_pre_blocks"], - num_mid_blocks=config["num_mid_blocks"], - num_post_blocks=config["num_post_blocks"], - input_skip=config["input_skip"], - upsample_mode=config["upsample_mode"], - temporal_mix=config["temporal_mix"], - icnr=config.get("icnr", False), - dims=config["dims"], - expand_ratio=config["expand_ratio"], - kernel_size=config["kernel_size"], - layer_scale_init=config["layer_scale_init"], - grn=config["grn"], - stochastic_depth_rate=config["stochastic_depth_rate"], - depthwise=config["depthwise"], - zq_dim=config["in_channels"] if config["modulated_norm"] else None, - bare_stem=config["bare_stem"], - modulated_output_proj=config["modulated_output_proj"], - global_skip=config.get("global_skip", False), - enable_x2_entry=config["enable_x2_entry"], - x2_adapter_blocks=config.get("x2_adapter_blocks", 0), - x2_tail_mode=config.get("x2_tail_mode", "shared"), - x2_adapter_sources=tuple(x2_adapter_sources) if x2_adapter_sources is not None else None, - x2_finisher=config.get("x2_finisher", "none"), - stage_channels=tuple(stage_channels) if stage_channels is not None else None, - temporal_padding=config["temporal_padding"], - upsample_padding_mode=config["upsample_padding_mode"], - ) - - -class Kandinsky6SRLatentUpscalerBank(ModelMixin, ConfigMixin): - """Diffusers wrapper around the self-contained x2/x4 latent-upscaler bank. - - Args: - models (`list[dict]`): Serialized upscaler definitions. Each definition - contains a ``target_scale`` and a validated ``model`` configuration. - scaling_factor (`float`, *optional*, defaults to 1.0): Factor applied to - latents before they are passed to an upscaler. - """ - - _no_split_modules = ["MultiScaleUpsampler"] - - @register_to_config - def __init__( - self, - models: Sequence[Mapping[str, Any]], - scaling_factor: float = 1.0, - ) -> None: - super().__init__() - self._models = nn.ModuleDict() - for spec in models: - target_scale = str(spec["target_scale"]) - if not target_scale.endswith("x"): - target_scale = f"{target_scale}x" - if target_scale in self._models: - raise ValueError(f"duplicate latent upscaler target scale: {target_scale}") - upscaler = build_upsampler(spec["model"]) - upscaler.target_scale = target_scale - upscaler.scaling_factor = float(scaling_factor) - self._models[target_scale] = upscaler - - def for_scale(self, scale: int) -> nn.Module | None: - """Return the bank entry matching ``scale``.""" - key = f"{int(scale)}x" - return self._models[key] if key in self._models else None - - @property - def scales(self) -> tuple[int, ...]: - """Return all serialized upscale factors.""" - return tuple(sorted(int(key.removesuffix("x")) for key in self._models)) - - def forward( - self, - latent: Tensor, - *, - entry: str | None = None, - return_intermediates: bool | None = None, - ) -> Any: - """Forward one bank entry, primarily for direct component use. - - Args: - latent (`torch.Tensor`): Latent tensor to upscale. - entry (`str`, *optional*): Entry name such as ``"x2"`` or ``"2x"``. - Required when the bank contains multiple scales. - return_intermediates (`bool`, *optional*): Whether to return - intermediate upscaler outputs when supported by the entry. - - Returns: - `torch.Tensor` or `dict`: Upscaled latent, or the selected entry's - intermediate-output structure. - """ - if entry is None: - if len(self._models) != 1: - raise ValueError("entry is required when the latent-upscaler bank has multiple scales") - entry = next(iter(self._models)) - key = entry.removeprefix("x") - if not key.endswith("x"): - key = f"{key}x" - model = self._models[key] - kwargs: dict[str, Any] = {"entry": f"x{key.removesuffix('x')}"} - if return_intermediates is not None: - kwargs["return_intermediates"] = return_intermediates - return model(latent, **kwargs) - - -__all__ = ["Kandinsky6SRLatentUpscalerBank"] diff --git a/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py b/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py new file mode 100644 index 000000000000..f5bbcf10fc66 --- /dev/null +++ b/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py @@ -0,0 +1,296 @@ +# 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. + +"""Latent upscalers used by the Kandinsky 6 video super-resolution pipeline.""" + +from __future__ import annotations + +import torch +from torch import Tensor, nn +from torch.nn import functional + +from ...configuration_utils import ConfigMixin, register_to_config +from ...models.autoencoders.vae import DecoderOutput +from ...models.modeling_utils import ModelMixin + + +class Kandinsky6SRLatentUpscalerConv3d(nn.Conv3d): + """`Conv3d` that replicates the edge frame along time and zero-pads height and width, matching the K-VAE + latents this model operates on: repeating the boundary frame avoids a zero "hole" next to frame 0, which already + encodes a single pixel frame while every later latent frame aggregates several. + """ + + def __init__(self, in_channels: int, out_channels: int, kernel_size: int) -> None: + super().__init__(in_channels, out_channels, kernel_size, padding=(0, kernel_size // 2, kernel_size // 2)) + self.temporal_pad = kernel_size // 2 + + def forward(self, hidden_states: Tensor) -> Tensor: + if self.temporal_pad: + hidden_states = functional.pad( + hidden_states, (0, 0, 0, 0, self.temporal_pad, self.temporal_pad), mode="replicate" + ) + return super().forward(hidden_states) + + +class Kandinsky6SRLatentUpscalerRMSNorm(nn.Module): + """Channel-first RMS normalization with a learnable gain.""" + + def __init__(self, num_channels: int) -> None: + super().__init__() + self.scale = num_channels**0.5 + self.gamma = nn.Parameter(torch.ones(num_channels, 1, 1, 1)) + + def forward(self, hidden_states: Tensor) -> Tensor: + return functional.normalize(hidden_states, dim=1) * self.scale * self.gamma + + +class Kandinsky6SRLatentUpscalerModulatedNorm(nn.Module): + """RMS norm followed by a FiLM modulation computed from the input latent `zq`, nearest-upsampled to the feature + grid so the same conditioning serves every resolution of the cascade.""" + + def __init__(self, num_channels: int, zq_channels: int) -> None: + super().__init__() + self.norm = Kandinsky6SRLatentUpscalerRMSNorm(num_channels) + self.conv_y = nn.Conv3d(zq_channels, num_channels, kernel_size=1) + self.conv_b = nn.Conv3d(zq_channels, num_channels, kernel_size=1) + + def forward(self, hidden_states: Tensor, zq: Tensor) -> Tensor: + if zq.shape[2:] != hidden_states.shape[2:]: + zq = functional.interpolate(zq, size=hidden_states.shape[2:], mode="nearest") + return self.norm(hidden_states) * self.conv_y(zq) + self.conv_b(zq) + + +class Kandinsky6SRLatentUpscalerResidualBlock(nn.Module): + """Pre-activation residual block: `norm -> SiLU -> conv3x3x3 -> norm -> SiLU -> conv3x3x3`, with a 1x1 shortcut + when the width changes.""" + + def __init__(self, in_channels: int, out_channels: int, zq_channels: int) -> None: + super().__init__() + self.norm1 = Kandinsky6SRLatentUpscalerModulatedNorm(in_channels, zq_channels) + self.conv1 = Kandinsky6SRLatentUpscalerConv3d(in_channels, out_channels, kernel_size=3) + self.norm2 = Kandinsky6SRLatentUpscalerModulatedNorm(out_channels, zq_channels) + self.conv2 = Kandinsky6SRLatentUpscalerConv3d(out_channels, out_channels, kernel_size=3) + self.shortcut = ( + nn.Identity() if in_channels == out_channels else nn.Conv3d(in_channels, out_channels, kernel_size=1) + ) + + def forward(self, hidden_states: Tensor, zq: Tensor) -> Tensor: + residual = self.conv1(functional.silu(self.norm1(hidden_states, zq))) + residual = self.conv2(functional.silu(self.norm2(residual, zq))) + return self.shortcut(hidden_states) + residual + + +class Kandinsky6SRLatentUpscalerUpsample(nn.Module): + """Spatial 2x: nearest upsample plus a per-frame convolutional residual, mixed by a pointwise convolution.""" + + def __init__(self, channels: int) -> None: + super().__init__() + self.spatial_conv = nn.Conv3d(channels, channels, kernel_size=(1, 3, 3), padding=(0, 1, 1)) + self.linear = nn.Conv3d(channels, channels, kernel_size=1) + + def forward(self, hidden_states: Tensor) -> Tensor: + hidden_states = functional.interpolate(hidden_states, scale_factor=(1, 2, 2), mode="nearest") + return self.linear(hidden_states + self.spatial_conv(hidden_states)) + + +class Kandinsky6SRLatentUpscalerOutputHead(nn.Sequential): + """`modulated norm -> SiLU -> conv3x3x3` projection back to the latent channels, threading `zq` into the norm.""" + + def __init__(self, in_channels: int, out_channels: int, zq_channels: int) -> None: + super().__init__( + Kandinsky6SRLatentUpscalerModulatedNorm(in_channels, zq_channels), + nn.SiLU(), + Kandinsky6SRLatentUpscalerConv3d(in_channels, out_channels, kernel_size=3), + ) + + def forward(self, hidden_states: Tensor, zq: Tensor) -> Tensor: + norm, activation, conv = self + return conv(activation(norm(hidden_states, zq))) + + +def _residual_stack(num_blocks: int, in_channels: int, out_channels: int, zq_channels: int) -> nn.ModuleList: + """`num_blocks` residual blocks; the first one narrows `in_channels` to `out_channels`.""" + widths = [in_channels] + [out_channels] * num_blocks + return nn.ModuleList( + [ + Kandinsky6SRLatentUpscalerResidualBlock(widths[index], widths[index + 1], zq_channels) + for index in range(num_blocks) + ] + ) + + +class Kandinsky6SRLatentUpscalerX2Branch(nn.Module): + """The x2 upscaler tail: adapter blocks and a finisher on the input grid, then one 2x stage.""" + + def __init__( + self, + in_channels: int, + stage_channels: tuple[int, int, int], + num_adapter_blocks: int, + num_mid_blocks: int, + num_post_blocks: int, + ) -> None: + super().__init__() + width_1, width_2, width_3 = stage_channels + self.adapter = _residual_stack(num_adapter_blocks, width_1, width_1, in_channels) + self.finisher = nn.Module() + self.finisher.spatial_conv = nn.Conv3d(width_1, width_1, kernel_size=(1, 3, 3), padding=(0, 1, 1)) + self.finisher.linear = nn.Conv3d(width_1, width_1, kernel_size=1) + self.private_mid_blocks = _residual_stack(num_mid_blocks, width_1, width_2, in_channels) + self.private_upsample = Kandinsky6SRLatentUpscalerUpsample(width_2) + self.private_blocks = _residual_stack(num_post_blocks, width_2, width_3, in_channels) + self.private_output_proj = Kandinsky6SRLatentUpscalerOutputHead(width_3, in_channels, in_channels) + + def forward(self, hidden_states: Tensor, zq: Tensor) -> Tensor: + for block in self.adapter: + hidden_states = block(hidden_states, zq) + hidden_states = self.finisher.linear(hidden_states + self.finisher.spatial_conv(hidden_states)) + for block in self.private_mid_blocks: + hidden_states = block(hidden_states, zq) + hidden_states = self.private_upsample(hidden_states) + for block in self.private_blocks: + hidden_states = block(hidden_states, zq) + return self.private_output_proj(hidden_states, zq) + + +class Kandinsky6SRLatentUpscaler(nn.Module): + """One latent upscaler of the bank: a two-stage 2x+2x cascade for `scale=4`, or its single-stage x2 variant. + + Both share the same widths. The x4 model runs `input_proj -> pre_blocks -> upsample_1 -> mid_blocks -> upsample_2 + -> post_blocks -> output_proj`; the x2 model runs its own stem and the [`Kandinsky6SRLatentUpscalerX2Branch`] tail. + """ + + def __init__( + self, + scale: int, + in_channels: int, + stage_channels: tuple[int, int, int], + num_pre_blocks: int, + num_mid_blocks: int, + num_post_blocks: int, + num_x2_adapter_blocks: int, + ) -> None: + super().__init__() + if scale not in (2, 4): + raise ValueError(f"`scale` must be 2 or 4, got {scale}") + self.scale = scale + width_1, width_2, width_3 = stage_channels + + if scale == 4: + self.input_proj = nn.Sequential(Kandinsky6SRLatentUpscalerConv3d(in_channels, width_1, kernel_size=3)) + self.pre_blocks = _residual_stack(num_pre_blocks, width_1, width_1, in_channels) + self.upsample_1 = Kandinsky6SRLatentUpscalerUpsample(width_1) + self.mid_blocks = _residual_stack(num_mid_blocks, width_1, width_2, in_channels) + self.upsample_2 = Kandinsky6SRLatentUpscalerUpsample(width_2) + self.post_blocks = _residual_stack(num_post_blocks, width_2, width_3, in_channels) + self.output_proj = Kandinsky6SRLatentUpscalerOutputHead(width_3, in_channels, in_channels) + else: + self.mid_input_proj = nn.Sequential(Kandinsky6SRLatentUpscalerConv3d(in_channels, width_1, kernel_size=3)) + self.x2_branch = Kandinsky6SRLatentUpscalerX2Branch( + in_channels, stage_channels, num_x2_adapter_blocks, num_mid_blocks, num_post_blocks + ) + + def forward(self, latents: Tensor) -> Tensor: + zq = latents + if self.scale == 2: + return self.x2_branch(self.mid_input_proj(latents), zq) + + hidden_states = self.input_proj(latents) + for block in self.pre_blocks: + hidden_states = block(hidden_states, zq) + hidden_states = self.upsample_1(hidden_states) + for block in self.mid_blocks: + hidden_states = block(hidden_states, zq) + hidden_states = self.upsample_2(hidden_states) + for block in self.post_blocks: + hidden_states = block(hidden_states, zq) + return self.output_proj(hidden_states, zq) + + +class Kandinsky6SRLatentUpscalerBank(ModelMixin, ConfigMixin): + r""" + Bank of latent upscalers used by [`Kandinsky6SRPipeline`]: one [`Kandinsky6SRLatentUpscaler`] per supported spatial + scale, operating on K-VAE latents. + + Args: + in_channels (`int`, defaults to `64`): + Number of latent channels. + stage_channels (`tuple[int, int, int]`, defaults to `(2048, 1024, 512)`): + Feature widths of the three stages of the cascade. + num_pre_blocks (`int`, defaults to `5`): + Residual blocks before the first upsample of the x4 model. + num_mid_blocks (`int`, defaults to `3`): + Residual blocks between the two upsamples. + num_post_blocks (`int`, defaults to `3`): + Residual blocks after the last upsample. + num_x2_adapter_blocks (`int`, defaults to `2`): + Residual blocks of the x2 model's adapter. + scales (`tuple[int, ...]`, defaults to `(2, 4)`): + Spatial scales the bank provides an upscaler for. + scaling_factor (`float`, defaults to `0.910344`): + Scale the input latents are expected to carry (the K-VAE `scaling_factor`). + """ + + _no_split_modules = ["Kandinsky6SRLatentUpscaler"] + + @register_to_config + def __init__( + self, + in_channels: int = 64, + stage_channels: tuple[int, int, int] = (2048, 1024, 512), + num_pre_blocks: int = 5, + num_mid_blocks: int = 3, + num_post_blocks: int = 3, + num_x2_adapter_blocks: int = 2, + scales: tuple[int, ...] = (2, 4), + scaling_factor: float = 0.910344004631042, + ) -> None: + super().__init__() + self._models = nn.ModuleDict( + { + f"{scale}x": Kandinsky6SRLatentUpscaler( + scale=scale, + in_channels=in_channels, + stage_channels=stage_channels, + num_pre_blocks=num_pre_blocks, + num_mid_blocks=num_mid_blocks, + num_post_blocks=num_post_blocks, + num_x2_adapter_blocks=num_x2_adapter_blocks, + ) + for scale in scales + } + ) + + def forward(self, latents: Tensor, scale: int, return_dict: bool = True) -> DecoderOutput | tuple[Tensor]: + r""" + Args: + latents (`torch.Tensor` of shape `(batch_size, in_channels, num_frames, height, width)`): + K-VAE latents scaled by `scaling_factor`. + scale (`int`): + Spatial upscale factor; one of `scales`. + return_dict (`bool`, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.DecoderOutput`] instead of a plain tuple. + + Returns: + The upscaled latents of shape `(batch_size, in_channels, num_frames, height * scale, width * scale)`. + """ + if scale not in self.config.scales: + raise ValueError(f"No latent upscaler for scale {scale}; available scales: {list(self.config.scales)}") + upscaled = self._models[f"{scale}x"](latents) + if not return_dict: + return (upscaled,) + return DecoderOutput(sample=upscaled) + + +__all__ = ["Kandinsky6SRLatentUpscalerBank"] diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py index 3178dc564974..3579b254a29b 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py @@ -12,1512 +12,378 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Kandinsky 6 super-resolution Diffusers pipeline.""" - +import inspect import math -from contextlib import contextmanager -from dataclasses import dataclass -from typing import Any, Literal, NamedTuple import numpy as np +import PIL.Image import torch -from torch import Tensor, nn -from torch.nn import functional +import torch.nn.functional as F -from ...utils import logging, replace_example_docstring +from ...models import Kandinsky6SRTransformer3DModel, Kandinsky6SRVAE +from ...schedulers import FlowMatchEulerDiscreteScheduler, PiflowScheduler +from ...utils import replace_example_docstring from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline +from .modeling_latent_upscaler import Kandinsky6SRLatentUpscalerBank from .pipeline_output import Kandinsky6SRPipelineOutput -logger = logging.get_logger(__name__) - - EXAMPLE_DOC_STRING = """ Examples: - ```python >>> import torch >>> from diffusers import Kandinsky6SRPipeline, Kandinsky6TI2VAPipeline + >>> from diffusers.utils import export_to_video - >>> t2va = Kandinsky6TI2VAPipeline.from_pretrained( - ... "kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers", torch_dtype=torch.bfloat16 - ... ) - >>> t2va.enable_model_cpu_offload() - >>> generated = t2va(prompt="A cat and a dog baking a cake together in a kitchen.", height=480, width=864) - - >>> pipe = Kandinsky6SRPipeline.from_pretrained( - ... "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers", torch_dtype=torch.bfloat16 + >>> pipe = Kandinsky6TI2VAPipeline.from_pretrained( + ... "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers", torch_dtype=torch.bfloat16 ... ) >>> pipe.enable_model_cpu_offload() - - >>> output = pipe(video=generated.frames, resolution_scale=2.25) - >>> output.frames.shape - torch.Size([1, 3, 121, 1080, 1944]) + >>> video = pipe( + ... prompt="A cat and a dog baking a cake together in a kitchen.", + ... height=480, + ... width=864, + ... num_inference_steps=16, + ... guidance_scale=1.0, + ... sample_audio=False, + ... ).frames[0] + + >>> sr_pipe = Kandinsky6SRPipeline.from_pretrained( + ... "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", torch_dtype=torch.bfloat16 + ... ) + >>> sr_pipe.enable_model_cpu_offload() + >>> output = sr_pipe(video=video, resolution_scale=2.25, num_inference_steps=2) + >>> export_to_video(output.frames[0], "output_sr.mp4", fps=24) ``` """ -# These values are the SR data contract and must stay local to this pipeline. -VAE_SPATIAL_FACTOR = 16 -RESOLUTIONS: dict[int, list[tuple[int, int]]] = {512: [(512, 512), (512, 768), (768, 512)]} -MAX_NUM_FRAMES = 121 -ResolutionScale = Literal[2, 4] - - -# Shared tiling utilities for spatial video tiling with Hanning-window blending. - - -class TileGrid(NamedTuple): - """Spatial tile grid parameters. - - ``tops`` and ``lefts`` are the authoritative per-axis tile start positions. - They cover ``[0, length - tile]``, evenly spread so every gap keeps at - least the configured overlap (see :func:`axis_positions`). - """ - - tile_h: int - tile_w: int - tops: tuple[int, ...] - lefts: tuple[int, ...] - - @property - def n_h(self) -> int: - """Number of tile rows.""" - return len(self.tops) - - @property - def n_w(self) -> int: - """Number of tile columns.""" - return len(self.lefts) - - @property - def total_tiles(self) -> int: - """Total number of tiles (``n_h * n_w``).""" - return self.n_h * self.n_w - - -def axis_positions(length: int, tile: int, min_overlap: float, snap: int) -> tuple[int, ...]: - """Return evenly distributed, ``snap``-aligned tile start positions. - - Uses the minimal position count whose uniform stride keeps the tile - overlap at or above ``min_overlap``, then spreads the positions evenly - over ``[0, length - tile]`` in integer ``snap`` units. When ``length`` or - ``tile`` is not ``snap``-aligned the same layout is computed at pixel - precision (``snap=1``) — the latent-grid guard downstream still enforces - alignment where it actually matters (the LU path). - - Args: - length: Axis length in pixels. - tile: Tile size on this axis. - min_overlap: Overlap floor as a fraction of ``tile`` in ``[0, 1)``. - snap: Position alignment unit (the VAE spatial factor). - - Returns: - Strictly increasing positions covering ``[0, length - tile]``. - - Raises: - ValueError: On non-positive ``tile``/``snap`` or ``min_overlap`` - outside ``[0, 1)``. - """ - if tile <= 0 or snap <= 0 or not 0 <= min_overlap < 1: - msg = f"invalid axis spec: tile={tile}, snap={snap}, min_overlap={min_overlap}" - raise ValueError(msg) - if tile >= length: - return (0,) - unit = snap if length % snap == 0 and tile % snap == 0 else 1 - span_units = (length - tile) // unit - max_stride_units = max(1, math.floor(tile * (1.0 - min_overlap) / unit)) - count = math.ceil(span_units / max_stride_units) + 1 - return tuple(round(i * span_units / (count - 1)) * unit for i in range(count)) - - -def compute_tile_grid(h: int, w: int, tile_hw: tuple[int, int], min_overlap: float, snap: int) -> TileGrid: - """Build a :class:`TileGrid` covering ``(h, w)`` with the given tile size. - - Args: - h: Video height in pixels. - w: Video width in pixels. - tile_hw: Explicit ``(tile_h, tile_w)``. - min_overlap: Overlap floor as a fraction of the tile size. - snap: Position alignment unit (the VAE spatial factor). - - Returns: - ``TileGrid`` with evenly distributed, aligned tile positions. - """ - tile_h, tile_w = tile_hw - return TileGrid( - tile_h=tile_h, - tile_w=tile_w, - tops=axis_positions(h, tile_h, min_overlap, snap), - lefts=axis_positions(w, tile_w, min_overlap, snap), - ) - - -def extract_all_tiles(video: torch.Tensor, grid: TileGrid) -> list[torch.Tensor]: - """Extract all spatial tiles from a video tensor. - - Args: - video: ``[C, T, H, W]`` tensor (only the last two dims are read). - grid: Tile grid parameters from ``compute_tile_grid``. - - Returns: - List of ``[C, T, tile_h, tile_w]`` tensors in row-major order - (``tops`` x ``lefts``). - """ - tiles: list[torch.Tensor] = [] - for top in grid.tops: - for left in grid.lefts: - tile = video[:, :, top : top + grid.tile_h, left : left + grid.tile_w] - tiles.append(tile) - return tiles - - -def hanning_window_2d(h: int, w: int, device: torch.device) -> torch.Tensor: - """Create a 2D Hanning window with non-zero endpoints. - - Uses ``hann_window(n + 2)[1:-1]`` to avoid exact zeros at boundaries, - ensuring non-zero weight where only one tile contributes. - - Args: - h: Window height. - w: Window width. - device: Target device. - - Returns: - ``[h, w]`` float tensor with values in ``(0, 1]``. - """ - wy = torch.hann_window(h + 2, device=device)[1:-1] - wx = torch.hann_window(w + 2, device=device)[1:-1] - return wy[:, None] * wx[None, :] - - -def stitch_tiles_hanning( - tiles: list[torch.Tensor], - grid: TileGrid, - original_h: int, - original_w: int, - scale: int = 1, -) -> torch.Tensor: - """Stitch tiles into a full frame using Hanning-window weighted blending. - - Each tile is multiplied by a 2D Hanning window and accumulated into - the output canvas. The final result is normalised by the accumulated - weights so that overlapping regions blend smoothly. The same scheme - works for non-uniform overlap at the right/bottom edge: in the clamped - region both ``pred_acc`` and ``weight_acc`` receive more contributions, - and the per-pixel division cancels it out. - - When ``scale > 1``, tiles are assumed to be at HR resolution - (i.e. each tile covers ``tile_h * scale x tile_w * scale`` pixels) - and the output canvas is ``original_h * scale x original_w * scale``. - Grid positions are scaled accordingly. - - Args: - tiles: List of ``[C, T, th, tw]`` float tensors in row-major order - (same order as ``extract_all_tiles``). When ``scale == 1``, - ``th == grid.tile_h``; when ``scale > 1``, ``th == grid.tile_h * scale``. - grid: Tile grid parameters (at LQ / original resolution). - original_h: LQ frame height. - original_w: LQ frame width. - scale: Upscale factor. Output resolution is - ``(original_h * scale, original_w * scale)``. - - Returns: - ``[C, T, original_h * scale, original_w * scale]`` float tensor. - """ - first = tiles[0] - c, t = first.shape[0], first.shape[1] - hr_tile_h, hr_tile_w = first.shape[2], first.shape[3] - device = first.device - - window = hanning_window_2d(hr_tile_h, hr_tile_w, device) - window = window.unsqueeze(0).unsqueeze(0) # [1, 1, hr_tile_h, hr_tile_w] - - out_h = original_h * scale - out_w = original_w * scale - pred_acc = torch.zeros(c, t, out_h, out_w, device=device) - weight_acc = torch.zeros(1, 1, out_h, out_w, device=device) - - tile_idx = 0 - for top in grid.tops: - for left in grid.lefts: - y = top * scale - x = left * scale - tile = tiles[tile_idx] - pred_acc[:, :, y : y + hr_tile_h, x : x + hr_tile_w] += tile * window - weight_acc[:, :, y : y + hr_tile_h, x : x + hr_tile_w] += window - tile_idx += 1 - - return pred_acc / weight_acc.clamp(min=1e-6) - - -def resolve_scale_request(scale: float) -> tuple[int, float]: - """Map a requested ``resolution_scale`` to ``(tiling_scale, pre_upscale)``. - - Args: - scale: The requested total upscale — ``2``, ``4`` or ``2.25``. - - Returns: - ``(tiling_scale, pre_upscale)``; ``pre_upscale`` is ``1.0`` for the - integer scales. - """ - requested = float(scale) - if requested == 2.25: - return 2, 1.125 - if requested in (2.0, 4.0): - return int(requested), 1.0 - raise ValueError("SR supports total scales 2, 4, and 2.25") - - -def pre_upscale_video(video: torch.Tensor, factor: float, spatial_multiple: int) -> torch.Tensor: - """Bilinear-upscale a ``[C, T, H, W]`` uint8 video by ``factor`` in pixel space. - - Target dims are rounded to the nearest multiple of ``spatial_multiple`` - (the VAE spatial factor) so the whole-video encode and the latent tile - grid stay integer-aligned (``latent_tile_grid_from_pixel_grid`` rejects - unaligned grids). For sources whose scaled dims already land on the - factor (512x768 x1.125 -> 576x864) the rounding is a no-op and the total - scale is exact. - - Args: - video: ``[C, T, H, W]`` uint8 source video (only the last two dims are read). - factor: Pixel upscale factor (> 1). - spatial_multiple: VAE spatial factor to align the target dims to. - - Returns: - ``[C, T, H', W']`` uint8 video with ``H' ~= H * factor`` aligned. - """ - if video.ndim != 4: - raise ValueError(f"video must have rank 4 [C,T,H,W], got {tuple(video.shape)}") - if factor <= 0 or spatial_multiple <= 0: - raise ValueError("factor and spatial_multiple must be positive") - height, width = video.shape[-2:] - target_h = max(spatial_multiple, round(height * factor / spatial_multiple) * spatial_multiple) - target_w = max(spatial_multiple, round(width * factor / spatial_multiple) * spatial_multiple) - resized = functional.interpolate( - video.float(), - size=(target_h, target_w), - mode="bilinear", - align_corners=False, - ) - return resized.round_().clamp_(0, 255).to(torch.uint8) - - -def get_visual_size(x: torch.Tensor, spatial_factor: int) -> int: - """Return the resolution key matching a visual latent tensor's spatial dims. - - Args: - x: Visual latent tensor of shape ``(T, H, W, C)``. - spatial_factor: The VAE's spatial compression factor. - - Returns: - The matching resolution key. - - Raises: - ValueError: If tensor dimensions do not match any known resolution. - """ - actual_size = (x.shape[1] * spatial_factor, x.shape[2] * spatial_factor) - for key, value in RESOLUTIONS.items(): - if actual_size in value: - return key - valid_sizes = {size for sizes in RESOLUTIONS.values() for size in sizes} - msg = ( - f"Visual tensor spatial dimensions {actual_size} do not match any known resolution. " - f"Tensor shape: {x.shape}. Valid resolutions: {valid_sizes}" - ) - raise ValueError(msg) - - -def degrade_lq_latent( - lq_latent: torch.Tensor, - noise_scale: float = 0.7, - noise_type: Literal["linear", "ddpm"] = "linear", - generator: torch.Generator | None = None, -) -> torch.Tensor: - """Mix random Gaussian noise into an LQ latent for data augmentation. - - Args: - lq_latent: LQ latent tensor of arbitrary shape. - noise_scale: Noise fraction ``s`` in ``(0, 1)``. When ``0``, the tensor - is returned unchanged. - noise_type: ``"linear"`` for ``(1-s)*lq + s*eps`` or ``"ddpm"`` for - ``sqrt(1-s²)*lq + s*eps`` (variance-preserving). - generator: Optional RNG generator for deterministic noise sampling. - - Returns: - Degraded LQ latent with the same shape and dtype as the input. - """ - if noise_scale <= 0: - return lq_latent - eps = randn_tensor(lq_latent.shape, generator=generator, device=lq_latent.device, dtype=lq_latent.dtype) - if noise_type == "ddpm": - return (1 - noise_scale**2) ** 0.5 * lq_latent + noise_scale * eps - return (1 - noise_scale) * lq_latent + noise_scale * eps - - -def cast_to_module_dtype(module: torch.nn.Module, value: torch.Tensor) -> torch.Tensor: - """Cast floating-point inputs to the module's parameter dtype.""" - try: - dtype = next(module.parameters()).dtype - except StopIteration: - return value - if value.is_floating_point() and value.dtype != dtype: - return value.to(dtype=dtype) - return value - - -def encode_pixels_to_latent(vae: torch.nn.Module, pixels: torch.Tensor) -> torch.Tensor: - """Encode ``(B, C, T, H, W)`` pixels in ``[0, 255]`` with the KVAE. - - Normalizes with the KVAE's own convention, then returns the latent - (``(latent, split_list)[0]`` — the regularizer mode; the KVAE always - returns this mode, there is no sampling toggle). - - Args: - vae: Causal video KVAE. - pixels: ``(B, C, T, H, W)`` tensor in ``[0, 255]``. - - Returns: - Latent ``(B, C, T', H', W')`` — not yet scaled by ``scaling_factor``. - """ - x = vae.normalize_data(pixels) - x = cast_to_module_dtype(vae, x) - result = vae.encode(x) - return result[0] - - -def decode_latent_to_uint8(vae: torch.nn.Module, latent: torch.Tensor) -> torch.Tensor: - """Decode a ``(B, C, T, H, W)`` latent to ``uint8`` pixels in ``[0, 255]``. - - The latent must already be unscaled (divided by ``scaling_factor``); - denormalization uses the KVAE's ``denormalize_data``. - - Args: - vae: Causal video KVAE. - latent: ``(B, C, T, H, W)`` unscaled latent. - - Returns: - ``(B, C, T_pixel, H, W)`` ``uint8`` pixels in ``[0, 255]``. - """ - latent = cast_to_module_dtype(vae, latent) - return denormalize_to_uint8(vae, vae.decode(latent).sample) - - -def denormalize_to_uint8(vae: torch.nn.Module, decoded: torch.Tensor) -> torch.Tensor: - """Convert an already-decoded tensor in ``[-1, 1]`` to ``uint8`` ``[0, 255]``. - - Uses the KVAE ``denormalize_data`` convention ((x + 1) * 128). For decode - outputs that did not pass through :func:`decode_latent_to_uint8` (e.g. - debug-clip dumps of already-decoded pixels). - - Args: - vae: Causal video KVAE. - decoded: Decode output in ``[-1, 1]`` (any shape). - - Returns: - ``uint8`` pixels in ``[0, 255]`` with the same shape. - """ - float_pixels = (vae.denormalize_data(decoded) / 255.0).clamp(0.0, 1.0) - return (float_pixels * 255.0).to(torch.uint8) - - -def encode_lq_video_to_lr_latent(lq_video: torch.Tensor, vae: torch.nn.Module, device: str | torch.device) -> Tensor: - """Encode one whole ``[C, T, H, W]`` pixel video into a raw ``[C, T', H', W']`` latent. - - Unlike :func:`Kandinsky6SRPipeline._encode_lq_videos` (a batch of spatial - tiles, channels-last, pre-scaled for the latent upscaler), this encodes a - single un-tiled video and leaves the result unscaled — the caller tiles it - in latent space and applies ``scaling_factor`` per tile (see - ``_upscale_lr_latent_tile``). +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. Args: - lq_video: ``[C, T, H, W]`` uint8 video in ``[0, 255]``. - vae: Causal video KVAE. - device: Target device. + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. Returns: - Unscaled latent ``[C, T', H', W']``. - - Raises: - ValueError: If ``lq_video`` is not rank 4. + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. """ - if lq_video.ndim != 4: - raise ValueError(f"lq_video must have rank 4 [C,T,H,W], got {tuple(lq_video.shape)}") - pixel = lq_video.unsqueeze(0).to(device=device) - latent = encode_pixels_to_latent(vae, pixel) - return latent.squeeze(0).float() - - -@dataclass -class RunConfig: - """Parameters for one tiled SR run.""" - - device: str - num_steps: int = 5 - seed: int = 42 - min_overlap: float = 0.20 - tiles_batch_size: int = 1 - resolution_scale: ResolutionScale = 4 - - def __post_init__(self) -> None: - if self.num_steps < 2: - raise ValueError("num_steps must be at least 2") - if not 0.0 <= self.min_overlap < 1.0: - raise ValueError("min_overlap must satisfy 0 <= min_overlap < 1") - if self.tiles_batch_size <= 0: - raise ValueError("tiles_batch_size must be positive") - if self.resolution_scale not in (2, 4): - raise ValueError("resolution_scale must be 2 or 4") - - -def _closest_base_resolution(h: int, w: int, visual_size: int) -> tuple[int, int]: - if visual_size not in RESOLUTIONS: - raise ValueError(f"Unsupported SR visual_size={visual_size}; known sizes: {sorted(RESOLUTIONS)}") - ratio = w / h if h else 1.0 - return min(RESOLUTIONS[visual_size], key=lambda hw: abs(hw[1] / hw[0] - ratio)) - - -def _tile_geometry( - h: int, - w: int, - visual_size: int, - scale: int, - min_overlap: float, - spatial_factor: int = VAE_SPATIAL_FACTOR, -) -> tuple[tuple[int, int], tuple[int, int], TileGrid]: - base_h, base_w = _closest_base_resolution(h, w, visual_size) - if base_h % scale or base_w % scale: - raise ValueError(f"resolution_scale={scale} does not divide the base resolution {base_h}x{base_w} exactly") - tile_hw = (base_h // scale, base_w // scale) - grid = compute_tile_grid(h, w, tile_hw, min_overlap, spatial_factor) - return (base_h, base_w), tile_hw, grid - - -def latent_tile_grid_from_pixel_grid(pixel_grid: TileGrid, spatial_factor: int) -> TileGrid: - """Convert an aligned pixel grid to the corresponding latent grid.""" - values = (pixel_grid.tile_h, pixel_grid.tile_w, *pixel_grid.tops, *pixel_grid.lefts) - if any(value % spatial_factor for value in values): - raise ValueError( - f"Pixel tile grid (tile {pixel_grid.tile_h}x{pixel_grid.tile_w}, " - f"tops={pixel_grid.tops}, lefts={pixel_grid.lefts}) not aligned to VAE spatial factor " - f"{spatial_factor}." - ) - return TileGrid( - pixel_grid.tile_h // spatial_factor, - pixel_grid.tile_w // spatial_factor, - tuple(top // spatial_factor for top in pixel_grid.tops), - tuple(left // spatial_factor for left in pixel_grid.lefts), - ) - - -def _upsample_tiles_to_base(raw_tiles: list[torch.Tensor], base_h: int, base_w: int) -> list[torch.Tensor]: - """Bilinear-upscale ``[C,T,h,w]`` pixel tiles to ``[T,base_h,base_w,C]`` (channels-last, for ``_encode_lq_videos``).""" - return [ - functional.interpolate(tile.float(), size=(base_h, base_w), mode="bilinear", align_corners=False).permute( - 1, 2, 3, 0 - ) - for tile in raw_tiles - ] - - -def _stitch_batch( - outputs: list[Tensor], - grid: TileGrid, - batch_size: int, - height: int, - width: int, - scale: float, -) -> Tensor: - """Stitch tile-major output back into ``[B,C,T,H,W]`` samples.""" - tile_count = grid.total_tiles - expected = batch_size * tile_count - if len(outputs) != expected: - raise RuntimeError(f"expected {expected} SR tile outputs, got {len(outputs)}") - return torch.stack( - [ - stitch_tiles_hanning( - [outputs[tile_index * batch_size + sample_index] for tile_index in range(tile_count)], - grid, - height, - width, - scale=scale, + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." ) - .clamp(0, 255) - .to(torch.uint8) - for sample_index in range(batch_size) - ] - ) - - -def _prepare_batch(value: Tensor, name: str, channels: int, *, dtype: torch.dtype | None = None) -> Tensor: - """Add a batch dim to a single ``(C,T,H,W)`` sample; validate a ``(B,C,T,H,W)`` batch.""" - if value.ndim == 4: - value = value.unsqueeze(0) - if value.ndim != 5: - raise ValueError(f"`{name}` must have shape (C,T,H,W) or (B,C,T,H,W), got {tuple(value.shape)}") - if value.shape[1] != channels: - raise ValueError(f"`{name}` must have {channels} channels at dim 1, got shape {tuple(value.shape)}") - if dtype is not None and value.dtype != dtype: - raise ValueError(f"`{name}` must be {dtype}, got {value.dtype}") - return value + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps -@contextmanager -def _execution_device_context(device: torch.device | str): - """Scope SR model execution to Diffusers' resolved CUDA device.""" - resolved_device = torch.device(device) - if resolved_device.type == "cuda": - with torch.cuda.device(resolved_device): - yield - else: - yield +def _tile_positions(length: int, tile: int, min_overlap: float, snap: int) -> list[int]: + """Evenly spread, `snap`-aligned tile start positions covering `[0, length - tile]` with at least `min_overlap` + (a fraction of `tile`) shared between neighbours.""" + if tile >= length: + return [0] + span = (length - tile) // snap + max_stride = max(1, math.floor(tile * (1.0 - min_overlap) / snap)) + count = math.ceil(span / max_stride) + 1 + return [round(index * span / (count - 1)) * snap for index in range(count)] -def _decode_source_latent_video(raw_latents: Tensor, source_vae: nn.Module, device: str | torch.device) -> Tensor: - """Decode a base-VAE latent to ``[C,T,H,W]`` pixels for the KVAE bridge.""" - if raw_latents.ndim != 4: - raise ValueError(f"raw_latents must have rank 4 [C,T,H,W], got {tuple(raw_latents.shape)}") - latent_5d = raw_latents.unsqueeze(0).to(device=device) - latent_5d = cast_to_module_dtype(source_vae, latent_5d) - # `source_vae.decode` is decorated with `@apply_forward_hook` (the standard convention on every - # AutoencoderKL-family `encode`/`decode` in this repo), so a Diffusers CPU-offload hook on - # `source_vae` already fires correctly here without any manual hook dispatch. - device_type = torch.device(device).type - with torch.autocast(device_type, dtype=torch.bfloat16, enabled=device_type == "cuda"): - decoded = source_vae.decode(latent_5d).sample - return ((decoded.squeeze(0).float().clamp(-1, 1) + 1.0) * 127.5).round().clamp(0, 255).to(torch.uint8).cpu() +def _hann_window_2d(height: int, width: int, device: torch.device) -> torch.Tensor: + """2D Hann window with non-zero borders, so a region covered by a single tile keeps a positive weight.""" + window_y = torch.hann_window(height + 2, device=device)[1:-1] + window_x = torch.hann_window(width + 2, device=device)[1:-1] + return window_y[:, None] * window_x[None, :] class Kandinsky6SRPipeline(DiffusionPipeline): - r"""Standalone Diffusers pipeline for Kandinsky 6 video super-resolution. + r""" + Pipeline for video super-resolution with Kandinsky 6. - Components are supplied to the constructor or loaded by Diffusers through - ``from_pretrained``. This pipeline deliberately has no config-path - factory: model construction belongs to the Diffusers component package, - while this class owns only SR orchestration. + The video is split into overlapping spatial tiles, every tile is refined by the SR transformer at one of the tile + sizes the model was trained on (`transformer.config.tile_sizes`), and the refined tiles are blended back with Hann + windows. When the pipeline has a `latent_upscaler`, the tiles are cut from the K-VAE latents of the whole video and + upscaled in latent space; otherwise the pixel tiles are bilinearly upscaled and encoded. - ``resolution_scale=2.25`` is the production route. It performs a 1.125x - pixel pre-upscale followed by the x2 latent-upscaler path. The input - components must already be on a device; ``.to(device)`` can be used after - construction. + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). Args: - transformer: SR transformer used to denoise latent tiles. - vae: Video VAE used for pixel-input encoding. - scheduler: Scheduler used for SR denoising. - latent_upscaler: Optional x2/x4 latent upscaler bank. - source_vae: Optional source VAE used only by the KVAE latent bridge. + transformer ([`Kandinsky6SRTransformer3DModel`]): + Transformer that refines the latent tiles. + vae ([`Kandinsky6SRVAE`]): + Causal video K-VAE used to encode the input video and decode the refined tiles. + scheduler ([`FlowMatchEulerDiscreteScheduler`] or [`PiflowScheduler`]): + Scheduler used with `transformer` to denoise the tiles. Distilled checkpoints ship with a + [`PiflowScheduler`]. + latent_upscaler ([`Kandinsky6SRLatentUpscalerBank`], *optional*): + Latent upscalers for the supported scales. """ - model_cpu_offload_seq = "source_vae->latent_upscaler->transformer->vae" - _optional_components = ["latent_upscaler", "source_vae"] + model_cpu_offload_seq = "latent_upscaler->transformer->vae" + _optional_components = ["latent_upscaler"] def __init__( self, - transformer: nn.Module, - vae: nn.Module, - scheduler: Any, - latent_upscaler: nn.Module | None = None, - source_vae: nn.Module | None = None, + transformer: Kandinsky6SRTransformer3DModel, + vae: Kandinsky6SRVAE, + scheduler: FlowMatchEulerDiscreteScheduler | PiflowScheduler, + latent_upscaler: Kandinsky6SRLatentUpscalerBank | None = None, ) -> None: super().__init__() - self.register_modules( - transformer=transformer, - vae=vae, - scheduler=scheduler, - latent_upscaler=latent_upscaler, - source_vae=source_vae, + self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler, latent_upscaler=latent_upscaler) + + self.vae_scale_factor_spatial = ( + self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 16 ) - if not getattr(transformer, "sr_params", None): - raise ValueError("SR transformer config must contain 'sr_params'") - vae_config = getattr(vae, "config", None) - self.spatial_factor = int( - getattr(vae, "spatial_factor", None) or getattr(vae_config, "spatial_factor", None) or VAE_SPATIAL_FACTOR + self.vae_scale_factor_temporal = ( + self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 4 ) + self.transformer_tile_sizes = ( + tuple(tuple(size) for size in self.transformer.config.tile_sizes) + if getattr(self, "transformer", None) is not None + else ((512, 512), (512, 768), (768, 512)) + ) + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) def check_inputs( - self, - video: Tensor | None = None, - latents: Tensor | None = None, - *, - resolution_scale: float, - num_inference_steps: int, - min_overlap: float, - tiles_batch_size: int, - kvae_bridge: bool = False, - output_type: str = "pt", - ) -> None: - """Validate SR input, scale, tiling, and output arguments. - - Args: - video: Pixel video batch, already normalized to rank 5. - latents: Latent batch, already normalized to rank 5. - resolution_scale: Total output scale, including the 2.25x - fractional route. - num_inference_steps: Number of denoising steps. - min_overlap: Minimum tile overlap fraction in ``[0, 1)``. - tiles_batch_size: Number of tiles processed together. - kvae_bridge: Whether ``latents`` come from a different VAE and - must be bridged through ``self.source_vae``. - output_type: ``pt``/``torch`` or ``np``/``numpy``. - - Raises: - ValueError: If inputs conflict or a value is outside the SR - pipeline's supported range. - """ - if (video is None) == (latents is None): - raise ValueError("pass exactly one of `video` or `latents`") - if output_type not in ("pt", "torch", "np", "numpy"): - raise ValueError(f"unsupported output_type={output_type!r}") - if video is not None: - num_frames = video.shape[2] - if (num_frames - 1) % 8 != 0: - raise ValueError( - f"`video` must have 1 + 8k frames (1, 9, 17, ..., {MAX_NUM_FRAMES}, ...), got {num_frames}" - ) - if num_frames > MAX_NUM_FRAMES: - raise ValueError(f"`video` must have at most {MAX_NUM_FRAMES} frames (~5s at 24fps), got {num_frames}") - - resolve_scale_request(float(resolution_scale)) + self, video, resolution_scale, num_inference_steps, lq_noise_scale, min_overlap, tiles_batch_size + ): + if resolution_scale not in (2, 2.25, 4): + raise ValueError(f"`resolution_scale` must be 2, 2.25 or 4 but is {resolution_scale}.") if num_inference_steps < 1: - raise ValueError("num_inference_steps must be positive") + raise ValueError(f"`num_inference_steps` has to be positive but is {num_inference_steps}.") + if not 0.0 < lq_noise_scale <= 1.0: + raise ValueError(f"`lq_noise_scale` has to be in (0, 1] but is {lq_noise_scale}.") if not 0.0 <= min_overlap < 1.0: - raise ValueError("min_overlap must be in [0, 1)") + raise ValueError(f"`min_overlap` has to be in [0, 1) but is {min_overlap}.") if tiles_batch_size < 1: - raise ValueError("tiles_batch_size must be positive") - if latents is not None: - if float(resolution_scale) == 2.25 and not kvae_bridge: - raise ValueError( - "2.25x latent SR requires kvae_bridge=True and source_vae; " - "ordinary latent input supports only x2/x4" - ) - if kvae_bridge and self.source_vae is None: - raise ValueError("KVAE bridge latent SR requires source_vae") - - @staticmethod - def _format_frames(frames: Tensor, output_type: str) -> Tensor | np.ndarray: - if output_type in ("pt", "torch"): - return frames - if output_type in ("np", "numpy"): - return frames.detach().cpu().numpy() - raise ValueError(f"Unsupported output_type={output_type!r}; use 'pt'/'torch' or 'np'/'numpy'") - - def _sparse_params(self, visual: torch.Tensor, visual_cu_seqlens: torch.Tensor) -> dict[str, Any] | None: - """Build NABLA sparse-attention parameters for ``self.transformer`` and the given visual tensor. - - Returns ``None`` for dense (``"flash"``) attention. The target ``P`` is used directly (the - training-time linear warmup schedule does not apply at inference). - """ - model = self.transformer - if model.patch_size[0] != 1: - raise ValueError(f"Expected transformer.patch_size[0] == 1, got {model.patch_size[0]}") - t, h, w, _ = visual.shape - t, h, w = (t // model.patch_size[0], h // model.patch_size[1], w // model.patch_size[2]) - visual_size = get_visual_size(visual, self.spatial_factor) - attention_configs = model.attention_params - try: - attention_params = attention_configs[visual_size] - except KeyError: - # JSON object keys are always strings, while the native YAML config - # uses integer resolution keys. - attention_params = attention_configs[str(visual_size)] - - if attention_params.get("type") != "nabla": - return None - return { - "to_fractal": True, - "P": attention_params.get("P"), - "wT": attention_params.get("wT"), - "wW": attention_params.get("wW"), - "wH": attention_params.get("wH"), - "add_sta": attention_params.get("add_sta"), - "visual_shape": (t, h, w), - "visual_seqlens": visual_cu_seqlens, - } - - @torch.no_grad() - def _encode_lq_videos(self, lq_videos: list[torch.Tensor], device: str | torch.device) -> torch.Tensor: - """VAE-encode a batch of ``[T,H,W,C]`` spatial tiles into a stacked ``[sum(T'),H',W',C]`` latent. - - Each tile is encoded independently (rather than as one batched VAE - call) to mirror the reference implementation exactly; ``self.vae``'s - ``scaling_factor`` is applied here, since these latents feed the - denoising loop directly (unlike ``encode_lq_video_to_lr_latent``, - whose unscaled output is scaled later by ``_upscale_lr_latent_tile``). - """ - vae = self.vae - scaling_factor = float(vae.config.scaling_factor) - latents = [] - for lq in lq_videos: - pixel = lq.permute(3, 0, 1, 2).unsqueeze(0).to(device=device) - latent = encode_pixels_to_latent(vae, pixel) - latents.append(latent.squeeze(0).permute(1, 2, 3, 0).float() * scaling_factor) - return torch.cat(latents, dim=0) - - @torch.no_grad() - def _upscale_lr_latent_tile( - self, lr_latent_tile: torch.Tensor, scale: int, device: str | torch.device - ) -> torch.Tensor: - """Upscale one raw ``[C,T,H,W]`` latent tile via ``self.latent_upscaler``'s ``scale``x entry.""" - scaling_factor = float(getattr(getattr(self.vae, "config", None), "scaling_factor", 1.0)) - tile = lr_latent_tile.unsqueeze(0).to(device=device, dtype=torch.float32) - tile = cast_to_module_dtype(self.latent_upscaler, tile) - device_type = torch.device(device).type - with torch.autocast(device_type, dtype=torch.bfloat16, enabled=device_type == "cuda"): - upscaled = self.latent_upscaler(tile * scaling_factor, entry=f"{scale}x", return_intermediates=False) - return upscaled.squeeze(0).permute(1, 2, 3, 0).float() - - def _build_initial_latent( - self, - lq_latent: torch.Tensor, - bs: int, - duration: int, - height: int, - width: int, - device: str | torch.device, - generator: torch.Generator | None, - lq_noise_scale: float, - lq_noise_type: Literal["linear", "ddpm"], - lq_channel_noise_scale: float, - ) -> torch.Tensor: - """Build the initial latent tensor for the SR denoising loop, per ``self.transformer.instruct_type``. - - Constructs ``[starting | cond | mask]`` depending on ``instruct_type``: - - - ``"noise"``: degraded LQ as starting point, zero LQ cond + zero mask. - - ``"channel"``: random Gaussian noise as starting point, LQ cond + ones mask. - - ``"hybrid"``: degraded LQ as starting point, LQ cond + ones mask. - - ``"hybrid_anchor"``: degraded LQ as starting point, zeroed (anchor-free) - HR anchor + zeroed anchor_mask — the only mode Kandinsky 6 SR ships with. - - Returns: - Initial latent tensor ready for the denoising loop (bs*T, H, W, 33). - - Raises: - ValueError: If ``instruct_type`` is not supported. - """ - dit = self.transformer - lq_latent = lq_latent.to(device) - lq_latent = cast_to_module_dtype(dit, lq_latent) - if dit.instruct_type == "noise": - degraded_lq = degrade_lq_latent(lq_latent, lq_noise_scale, lq_noise_type, generator=generator) - if not getattr(dit, "visual_cond", False): - return degraded_lq - zero_cond = torch.zeros_like(degraded_lq) - zero_mask = torch.zeros([*degraded_lq.shape[:-1], 1], dtype=degraded_lq.dtype, device=device) - return torch.cat([degraded_lq, zero_cond, zero_mask], dim=-1) - if dit.instruct_type in ("channel", "hybrid"): - if dit.instruct_type == "hybrid": - starting = degrade_lq_latent(lq_latent.clone(), lq_noise_scale, lq_noise_type, generator=generator) - else: - starting = randn_tensor( - (bs * duration, height, width, dit.in_visual_dim), generator=generator, device=device - ) - channel_lq = degrade_lq_latent(lq_latent, lq_channel_noise_scale, noise_type="linear") - mask = torch.ones_like(lq_latent[..., :1]) - return torch.cat([starting, channel_lq, mask], dim=-1) - if dit.instruct_type == "hybrid_anchor": - # Kandinsky 6 SR ships anchor-free: the public API has no way to supply a real HR - # anchor, so the anchor and its mask are always zeroed. - anchor_latent = torch.zeros_like(lq_latent) - anchor_mask = torch.zeros_like(lq_latent[..., :1]) - starting = degrade_lq_latent(lq_latent.clone(), lq_noise_scale, lq_noise_type, generator=generator) - anchor = degrade_lq_latent(anchor_latent, lq_channel_noise_scale, noise_type="linear") - return torch.cat([starting, anchor, anchor_mask], dim=-1) - raise ValueError(f"Kandinsky6SRPipeline does not support instruct_type={dit.instruct_type!r}") - - def _prepare_sr_latents( - self, - *, - lq_videos: list[torch.Tensor] | None, - lq_latents: torch.Tensor | None, - n_samples: int | None, - device: str | torch.device, - generator: torch.Generator | None, - lq_noise_scale: float, - lq_noise_type: str, - lq_channel_noise_scale: float, - ) -> tuple[torch.Tensor, torch.Tensor, int, int, int, int]: - """Encode LQ input and build the initial latent for one SR tile batch. - - Returns: - ``(lq_latent, image, batch_size, duration, height, width)``. - """ - if lq_latents is not None: - lq_latent = lq_latents.to(device) - if n_samples is None: - raise ValueError("n_samples is required when lq_latents is provided") - batch_size = n_samples - elif lq_videos is not None: - batch_size = len(lq_videos) - lq_latent = self._encode_lq_videos(lq_videos, device) - else: - raise ValueError("Either lq_videos or lq_latents must be provided") - duration = lq_latent.shape[0] // batch_size - height, width = (lq_latent.shape[1], lq_latent.shape[2]) - image = self._build_initial_latent( - lq_latent, - bs=batch_size, - duration=duration, - height=height, - width=width, - device=device, - generator=generator, - lq_noise_scale=lq_noise_scale, - lq_noise_type=lq_noise_type, - lq_channel_noise_scale=lq_channel_noise_scale, - ) - return lq_latent, image, batch_size, duration, height, width - - def _encode_sr_text( - self, - batch_size: int, - device: str | torch.device, - cached_text_embeds: dict[str, torch.Tensor] | None, - ) -> tuple[dict[str, torch.Tensor], torch.Tensor]: - """Build the SR text state: empty placeholders, or cached empty-caption embeddings. - - Kandinsky 6 SR ships with a text-free transformer (``use_text=False``) and has no - text encoder registered on the pipeline, so this never runs a real text encoder; - it only replicates precomputed ``cached_text_embeds`` to the batch size, or returns - placeholders when the transformer ignores text entirely. + raise ValueError(f"`tiles_batch_size` has to be positive but is {tiles_batch_size}.") - Raises: - ValueError: If the transformer needs text but no ``cached_text_embeds`` were given. - """ - if not getattr(self.transformer, "use_text", True): - # A use_text=False DiT ignores every text input (encoder, cross-attention, and pooled-text - # conditioning are all absent), so these zero-size tensors only satisfy the indexing and - # cu_seqlens bookkeeping in the denoising loop without ever loading a cached empty-caption file. - embeds = { - "text_embeds": torch.zeros(0, 1, device=device), - "pooled_embed": torch.zeros(batch_size, 1, device=device), - } - cu_seqlens = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) - return embeds, cu_seqlens - if cached_text_embeds is None: - raise ValueError("this SR transformer requires text; pass cached_text_embeds") - single_text = cached_text_embeds["text_embeds"] - single_pooled = cached_text_embeds["pooled_embed"] - seq_len = int(cached_text_embeds["cu_seqlens"][-1].item()) - text_embeds = { - "text_embeds": single_text.repeat(batch_size, 1).to(device), - "pooled_embed": single_pooled.repeat(batch_size, 1).to(device), - } - cu_seqlens = seq_len * torch.arange(batch_size + 1, dtype=torch.int32, device=device) - return text_embeds, cu_seqlens - - def _text_positions( - self, - image: torch.Tensor, - batch_size: int, - duration: int, - height: int, - width: int, - text_cu_seqlens: torch.Tensor, - ) -> tuple[torch.Tensor, list[torch.Tensor], torch.Tensor]: - """Build the packed visual/text RoPE positions ``self.transformer`` needs for one SR call.""" - device = image.device - visual_cu_seqlens = duration * torch.arange(batch_size + 1, dtype=torch.int32, device=device) - visual_rope_pos = [ - torch.cat([torch.arange(int(end), device=device) for end in torch.diff(visual_cu_seqlens).cpu()]), - torch.arange(height // self.transformer.patch_size[1], device=device), - torch.arange(width // self.transformer.patch_size[2], device=device), - ] - text_cu_seqlens = text_cu_seqlens.to(device) - text_rope_pos = torch.cat([torch.arange(int(end), device=device) for end in torch.diff(text_cu_seqlens).cpu()]) - return visual_cu_seqlens, visual_rope_pos, text_rope_pos - - def _predict( - self, - x: torch.Tensor, - t: torch.Tensor, - text_embeds: dict[str, torch.Tensor], - visual_cu_seqlens: torch.Tensor, - text_cu_seqlens: torch.Tensor, - visual_rope_pos: list[torch.Tensor], - text_rope_pos: torch.Tensor, - scale_factor: tuple[float, ...], - sparse_params: dict | None, - ) -> torch.Tensor: - """Run one ``self.transformer`` forward. SR has no guidance: a single conditional pass.""" - model_time = (t * 1000).to(dtype=x.dtype) - motion_score = torch.full((1,), 900.0, device=x.device, dtype=x.dtype) - return self.transformer( - x, - text_embeds["text_embeds"], - text_embeds["pooled_embed"], - model_time, - visual_cu_seqlens, - text_cu_seqlens, - visual_rope_pos, - text_rope_pos, - scale_factor=scale_factor, - sparse_params=sparse_params, - motion_score=motion_score, - ) - - @torch.no_grad() - def _denoise_native( - self, - img: torch.Tensor, - device: str | torch.device, - num_steps: int, - text_embeds: dict[str, torch.Tensor], - visual_cu_seqlens: torch.Tensor, - text_cu_seqlens: torch.Tensor, - visual_rope_pos: list[torch.Tensor], - text_rope_pos: torch.Tensor, - scale_factor: tuple[float, ...], - scheduler_scale: float, - start_timestep: float = 1.0, - progress_bar: Any = None, - ) -> torch.Tensor: - """Run the native Euler loop from ``start_timestep`` to 0 for the non-piflow SR transformer. - - This intentionally does not go through ``self.scheduler``: the SR transformer predicts - velocity directly and this loop applies the same shift formula - ``FlowMatchEulerDiscreteScheduler`` uses (``scheduler_scale`` mirrors its ``shift``), matching - the reference implementation bit-for-bit. - """ - img = img.to(device) - visual_cu_seqlens = visual_cu_seqlens.to(device) - text_cu_seqlens = text_cu_seqlens.to(device) - visual_rope_pos = [position.to(device) for position in visual_rope_pos] - text_rope_pos = text_rope_pos.to(device) - text_embeds = {key: value.to(device) for key, value in text_embeds.items()} - sparse_params = self._sparse_params(img, visual_cu_seqlens) - timesteps = torch.linspace(start_timestep, 0, num_steps, device=device) - timesteps = scheduler_scale * timesteps / (1 + (scheduler_scale - 1) * timesteps) - out_channels = self.transformer.in_visual_dim - for timestep, timestep_diff in zip(timesteps[:-1], torch.diff(timesteps), strict=False): - time = timestep.unsqueeze(0).repeat(visual_cu_seqlens.shape[0] - 1) - velocity = self._predict( - img, - time, - text_embeds, - visual_cu_seqlens, - text_cu_seqlens, - visual_rope_pos, - text_rope_pos, - scale_factor, - sparse_params, + num_frames = video.shape[2] + if (num_frames - 1) % self.vae_scale_factor_temporal != 0: + raise ValueError( + f"`video` must have `1 + k * {self.vae_scale_factor_temporal}` frames but has {num_frames}." ) - out_channels = velocity.shape[-1] - img[..., :out_channels] += timestep_diff * velocity - if progress_bar is not None: - progress_bar.update() - progress_bar.refresh() - return img[..., :out_channels] - - @torch.no_grad() - def _denoise_piflow( - self, - img: torch.Tensor, - text_embeds: dict[str, torch.Tensor], - visual_cu_seqlens: torch.Tensor, - text_cu_seqlens: torch.Tensor, - visual_rope_pos: list[torch.Tensor], - text_rope_pos: torch.Tensor, - scale_factor: tuple[float, ...], - *, - nfe: int, - out_dim: int, - device: str | torch.device | None = None, - progress_bar: Any = None, - ) -> torch.Tensor: - """Few-step π-Flow denoise: ``nfe`` transformer calls, via ``self.scheduler``. - ``self.scheduler`` (a ``PiflowScheduler``) owns the segment schedule and the - network-free policy integration between transformer calls entirely from its own config. + def encode_video(self, video: torch.Tensor) -> torch.Tensor: + r""" + Encodes a `(batch_size, channels, num_frames, height, width)` video in `[-1, 1]` into K-VAE latents scaled by + the VAE `scaling_factor`. """ - device = device if device is not None else img.device - img = img.to(device) - visual_cu_seqlens = visual_cu_seqlens.to(device) - text_cu_seqlens = text_cu_seqlens.to(device) - visual_rope_pos = [position.to(device) for position in visual_rope_pos] - text_rope_pos = text_rope_pos.to(device) - text_embeds = {key: value.to(device) for key, value in text_embeds.items()} - sparse_params = self._sparse_params(img, visual_cu_seqlens) - if nfe < 1: - raise ValueError(f"piflow denoising requires nfe >= 1, got {nfe}") - self.scheduler.set_timesteps(nfe, device=device) - if self.scheduler.timesteps.shape[0] != nfe: - raise ValueError(f"Piflow scheduler returned {self.scheduler.timesteps.shape[0]} steps, expected {nfe}") - x = img - for step_index in range(nfe): - timestep = self.scheduler.timesteps[step_index] - n_objects = visual_cu_seqlens.shape[0] - 1 - t_src = torch.full((n_objects,), float(timestep.item()), device=device, dtype=x.dtype) - v0 = self._predict( - x, - t_src, - text_embeds, - visual_cu_seqlens, - text_cu_seqlens, - visual_rope_pos, - text_rope_pos, - scale_factor, - sparse_params, - ) - x_pred = self.scheduler.step(v0, timestep, x[..., :out_dim], return_dict=False)[0] - x = torch.cat([x_pred, x[..., out_dim:]], dim=-1) - if progress_bar is not None: - progress_bar.update() - progress_bar.refresh() - return x[..., :out_dim] + # The K-VAE was trained on `pixels / 128 - 1` rather than the `pixels / 127.5 - 1` of `VideoProcessor`. + video = (video.to(self.vae.dtype) + 1) * (127.5 / 128) - 1 + latents = self.vae.encode(video, return_dict=False)[0].mode() + return latents * self.vae.config.scaling_factor - def _scheduler_nfe(self) -> int | None: - """Return the configured π-Flow step count, or ``None`` for a regular scheduler.""" - if not bool(getattr(self.scheduler, "is_piflow", False)): - return None - nfe = getattr(self.scheduler.config, "nfe", None) - if nfe is None: - raise ValueError("PiflowScheduler used for SR must define nfe in scheduler_config.json") - return int(nfe) - - def _denoise( - self, - image: torch.Tensor, - text_embeds: dict[str, torch.Tensor], - text_cu_seqlens: torch.Tensor, - batch_size: int, - duration: int, - height: int, - width: int, - *, - scale_factor: tuple[float, ...], - num_steps: int, - scheduler_scale: float, - lq_noise_scale: float, - cap_noise_timestep: bool, - device: str | torch.device, - progress_bar: Any = None, - ) -> torch.Tensor: - """Run either the Euler or π-Flow SR denoising stage on ``self.transformer``/``self.scheduler``.""" - nfe = self._scheduler_nfe() - visual_cu_seqlens, visual_rope_pos, text_rope_pos = self._text_positions( - image, batch_size, duration, height, width, text_cu_seqlens - ) - if nfe is not None: - if cap_noise_timestep and self.transformer.instruct_type in ("noise", "hybrid"): - raise NotImplementedError( - "piflow sampler does not support cap_noise_timestep for noise/hybrid instruct" - ) - out_dim = int(getattr(self.transformer, "base_out_visual_dim", self.transformer.in_visual_dim)) - return self._denoise_piflow( - image, - text_embeds, - visual_cu_seqlens, - text_cu_seqlens, - visual_rope_pos, - text_rope_pos, - scale_factor, - nfe=nfe, - out_dim=out_dim, - device=device, - progress_bar=progress_bar, - ) - start_timestep = ( - lq_noise_scale - if cap_noise_timestep and self.transformer.instruct_type in ("noise", "hybrid", "hybrid_anchor") - else 1.0 - ) - return self._denoise_native( - image, - device, - num_steps, - text_embeds, - visual_cu_seqlens, - text_cu_seqlens, - visual_rope_pos, - text_rope_pos, - scale_factor, - scheduler_scale, - start_timestep=start_timestep, - progress_bar=progress_bar, - ) - - @torch.no_grad() - def _vae_decode( - self, - latent_visual: torch.Tensor, - batch_size: int, - duration: int, - height: int, - width: int, - ) -> torch.Tensor: - """Decode denoised SR latents into ``[batch, 3, T, H, W]`` uint8 frames with ``self.vae``.""" - device_type = latent_visual.device.type - with torch.autocast(device_type, dtype=torch.bfloat16, enabled=device_type == "cuda"): - all_latents = latent_visual.reshape(batch_size, duration, height, width, -1) - all_latents = (all_latents / self.vae.config.scaling_factor).permute(0, 4, 1, 2, 3) - decoded = [decode_latent_to_uint8(self.vae, all_latents[i : i + 1]) for i in range(batch_size)] - return torch.cat(decoded, dim=0) - - @torch.no_grad() - def _run_stages( - self, - *, - scale_factor: tuple[float, ...], - num_steps: int, - scheduler_scale: float, - device: str | torch.device, - generator: torch.Generator | None, - lq_noise_scale: float, - lq_noise_type: str, - lq_channel_noise_scale: float, - cap_noise_timestep: bool, - cached_text_embeds: dict[str, torch.Tensor] | None, - lq_videos: list[torch.Tensor] | None = None, - lq_latents: torch.Tensor | None = None, - n_samples: int | None = None, - progress_bar: Any = None, - ) -> torch.Tensor: - """Run the four SR stages — prepare latents, encode text, denoise, decode — for one tile batch.""" - lq_latent, image, batch_size, duration, height, width = self._prepare_sr_latents( - lq_videos=lq_videos, - lq_latents=lq_latents, - n_samples=n_samples, - device=device, - generator=generator, - lq_noise_scale=lq_noise_scale, - lq_noise_type=lq_noise_type, - lq_channel_noise_scale=lq_channel_noise_scale, - ) - text_embeds, text_cu_seqlens = self._encode_sr_text(batch_size, device, cached_text_embeds) - latent_visual = self._denoise( - image, - text_embeds, - text_cu_seqlens, - batch_size, - duration, - height, - width, - scale_factor=scale_factor, - num_steps=num_steps, - scheduler_scale=scheduler_scale, - lq_noise_scale=lq_noise_scale, - cap_noise_timestep=cap_noise_timestep, - device=device, - progress_bar=progress_bar, - ) - return self._vae_decode(latent_visual, batch_size, duration, height, width) - - def _tile_batch_kwargs( - self, - run_config: RunConfig, - scale_factor: tuple[float, ...], - cached_text_embeds: dict[str, torch.Tensor] | None, - seed: int, - **inputs: Any, - ) -> dict[str, Any]: - # A fresh per-tile Generator, seeded deterministically from the run's base seed plus this - # tile batch's offset, so a tile's noise never depends on how many other tiles ran before it. - generator = torch.Generator(device=run_config.device) - generator.manual_seed(seed) - transformer = self.transformer - return { - **inputs, - "scale_factor": scale_factor, - "num_steps": run_config.num_steps, - "scheduler_scale": transformer.scheduler_scale, - "generator": generator, - "device": run_config.device, - "lq_noise_scale": transformer.lq_noise_scale, - "lq_noise_type": transformer.lq_noise_type, - "lq_channel_noise_scale": transformer.lq_channel_noise_scale, - "cap_noise_timestep": transformer.cap_noise_timestep, - "cached_text_embeds": cached_text_embeds, - } - - def _loaded_lu_scales(self) -> tuple[int, ...]: - if self.latent_upscaler is None: - return () - return tuple(self.latent_upscaler.scales) - - def _run_tile_batches( - self, - tile_inputs: list[torch.Tensor], - run_config: RunConfig, - scale_factor: tuple[float, ...], - *, - cached_text_embeds: dict[str, torch.Tensor] | None, - prepare_chunk: Any | None = None, - progress_bar: Any = None, - ) -> list[torch.Tensor]: - """Run the batched sampler loop for latent or pixel tile inputs. - - The transformer and SR VAE stay resident for the complete tile loop. Only the - optional latent-upscaler is streamed per chunk. - """ - outputs: list[torch.Tensor] = [] - for start in range(0, len(tile_inputs), run_config.tiles_batch_size): - raw_chunk = tile_inputs[start : start + run_config.tiles_batch_size] - chunk = prepare_chunk(raw_chunk) if prepare_chunk is not None else raw_chunk - input_kwargs: dict[str, Any] = ( - {"lq_videos": chunk} - if prepare_chunk is None - else {"lq_latents": torch.cat(chunk, dim=0), "n_samples": len(raw_chunk)} - ) - kwargs = self._tile_batch_kwargs( - run_config, - scale_factor, - cached_text_embeds, - run_config.seed + start, - **input_kwargs, - ) - sr_batch = self._run_stages(**kwargs, progress_bar=progress_bar) - if sr_batch.ndim != 5: - raise ValueError(f"SR sampler must return [batch,C,T,H,W], got {tuple(sr_batch.shape)}") - outputs.extend(sample.float().cpu() for sample in sr_batch) - return outputs - - def _run_batched_tiles( - self, - samples: list[Tensor], - run_config: RunConfig, - *, - latent_input: bool, - cached_text_embeds: dict[str, torch.Tensor] | None, - progress_bar: Any, - ) -> Tensor: - """Run equal-shape samples tile-major and restore the sample batch.""" - visual_size = self.transformer.visual_size - batch_size = len(samples) - scale_factor = self.transformer.scale_factor[visual_size] - if latent_input: - spatial_factor = self.spatial_factor - height = samples[0].shape[-2] * spatial_factor - width = samples[0].shape[-1] * spatial_factor - _, _, pixel_grid = _tile_geometry( - height, - width, - visual_size, - run_config.resolution_scale, - run_config.min_overlap, - spatial_factor, - ) - tile_grid = latent_tile_grid_from_pixel_grid(pixel_grid, spatial_factor) - input_grid = tile_grid - stitch_grid = pixel_grid - tile_sets = [extract_all_tiles(sample, tile_grid) for sample in samples] - if run_config.resolution_scale not in self._loaded_lu_scales(): - raise ValueError( - f"The latent-upscaler path has no {run_config.resolution_scale}x model " - f"(loaded LU scales: {self._loaded_lu_scales()})" - ) - - def prepare_chunk(chunk: list[Tensor]) -> list[Tensor]: - return [ - self._upscale_lr_latent_tile(tile, run_config.resolution_scale, run_config.device) - for tile in chunk - ] - else: - height, width = samples[0].shape[-2:] - spatial_factor = self.spatial_factor - base, _tile_hw, grid = _tile_geometry( - height, - width, - visual_size, - run_config.resolution_scale, - run_config.min_overlap, - spatial_factor, - ) - tile_sets = [ - _upsample_tiles_to_base(extract_all_tiles(sample, grid), base[0], base[1]) for sample in samples - ] - input_grid = grid - stitch_grid = grid - prepare_chunk = None - - tile_inputs = [ - tile_sets[sample_index][tile_index] - for tile_index in range(input_grid.total_tiles) - for sample_index in range(batch_size) - ] - outputs = self._run_tile_batches( - tile_inputs, - run_config, - scale_factor, - cached_text_embeds=cached_text_embeds, - prepare_chunk=prepare_chunk, - progress_bar=progress_bar, - ) - return _stitch_batch(outputs, stitch_grid, batch_size, height, width, run_config.resolution_scale) - - def _sr_progress_plan( - self, - value: Tensor, - run_config: RunConfig, - *, - latent_input: bool, - batch_size: int = 1, - ) -> tuple[int, int, int]: - """Return total updates, denoising steps, and tile batches.""" - visual_size = self.transformer.visual_size - spatial_factor = self.spatial_factor - height, width = value.shape[-2:] - if latent_input: - height *= spatial_factor - width *= spatial_factor - _, _, grid = _tile_geometry( - height, width, visual_size, run_config.resolution_scale, run_config.min_overlap, spatial_factor - ) - tile_count = grid.total_tiles * batch_size - chunk_count = max(1, math.ceil(tile_count / run_config.tiles_batch_size)) - nfe = self._scheduler_nfe() - denoise_steps = nfe if nfe is not None else run_config.num_steps - 1 - denoise_steps = max(1, denoise_steps) - return chunk_count * denoise_steps, denoise_steps, chunk_count - - @staticmethod - def _seed_from_generator(generator: torch.Generator | None, device: torch.device) -> int: - if generator is None: - return int(torch.randint(0, 2**31, (1,), device=device).item()) - generator_device = torch.device(getattr(generator, "device", "cpu")) - return int(torch.randint(0, 2**31, (1,), generator=generator, device=generator_device).item()) + def decode_latents(self, latents: torch.Tensor) -> torch.Tensor: + r"""Decodes scaled K-VAE latents into a `(batch_size, channels, num_frames, height, width)` video in `[-1, 1]`.""" + video = self.vae.decode(latents.to(self.vae.dtype) / self.vae.config.scaling_factor, return_dict=False)[0] + return ((video.float() + 1) * (128 / 127.5) - 1).clamp(-1, 1) @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, - video: Tensor | None = None, - latents: Tensor | None = None, + video: list[PIL.Image.Image] | list[list[PIL.Image.Image]] | np.ndarray | torch.Tensor, resolution_scale: float = 2.25, - num_inference_steps: int = 5, - generator: torch.Generator | None = None, - min_overlap: float = 0.20, + num_inference_steps: int = 4, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + lq_noise_scale: float = 0.7, + min_overlap: float = 0.2, tiles_batch_size: int = 1, - kvae_bridge: bool = False, - cached_text_embeds: dict[str, Tensor] | None = None, - output_type: str = "pt", + generator: torch.Generator | list[torch.Generator] | None = None, + output_type: str = "pil", return_dict: bool = True, - ) -> Kandinsky6SRPipelineOutput | tuple[Tensor]: - r"""Super-resolve one or more equal-size pixel videos or latent inputs. + ): + r""" + The call function to the pipeline for super-resolution. Args: - video: LQ pixel video(s) to super-resolve, ``uint8`` in ``[0, 255]``, - shape ``(3, T, H, W)`` or a batch ``(B, 3, T, H, W)`` — the same - layout [`Kandinsky6TI2VAPipeline`] returns with - ``output_type="pt"``. ``T`` must be ``1 + 8k`` frames (align or - truncate longer/misaligned clips yourself) and at most - ``MAX_NUM_FRAMES`` (121, ~5s at 24fps). Mutually exclusive with - ``latents``. - latents: Pre-encoded SR-VAE latent(s) already at the target - (post-upscale) resolution, shape ``(C, T, H, W)`` or a batch - ``(B, C, T, H, W)``, **unscaled** (not yet multiplied by - ``self.vae.config.scaling_factor``). Mutually exclusive with - ``video``. Pass ``kvae_bridge=True`` when these latents - instead come from a different (base) VAE at the *source* - resolution — they are then decoded through ``self.source_vae`` - and re-enter the normal ``video`` path (pre-upscale, encode, - latent-upscale) rather than being used directly. - resolution_scale: Total spatial upscale factor. Supported values - are ``2``, ``4``, and ``2.25``. - num_inference_steps: Number of denoising steps for the SR model. - generator: Optional random generator used to derive the sampling seed. - min_overlap: Minimum fraction of overlap between adjacent tiles. - tiles_batch_size: Number of tiles processed in one sampler batch. - kvae_bridge: Whether latent inputs come from a base VAE and need to - be decoded before SR. - cached_text_embeds: Optional cached empty-caption embeddings. - output_type: ``pt``/``torch`` or ``np``/``numpy``. - return_dict: Whether to return ``Kandinsky6SRPipelineOutput``. + video (`list[PIL.Image.Image]`, `np.ndarray` or `torch.Tensor`): + The low-resolution video(s), in any format [`~video_processor.VideoProcessor.preprocess_video`] + accepts, with `1 + k * 4` frames. Sizes are rounded down to a multiple of the VAE spatial factor. + resolution_scale (`float`, defaults to `2.25`): + Total spatial upscale: `2`, `4`, or `2.25` (a 1.125x bilinear pre-upscale followed by the 2x path). + num_inference_steps (`int`, defaults to `4`): + The number of denoising steps per tile. Use `2` with the distilled checkpoints. + timesteps (`list[int]`, *optional*): + Custom timesteps for schedulers that support them. + sigmas (`list[float]`, *optional*): + Custom sigmas for schedulers that support them. + lq_noise_scale (`float`, defaults to `0.7`): + Amount of Gaussian noise mixed into the low-resolution latents (variance preserving) before denoising. + min_overlap (`float`, defaults to `0.2`): + Minimum overlap between neighbouring tiles as a fraction of the tile size. + tiles_batch_size (`int`, defaults to `1`): + Number of tiles denoised per transformer call. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Generator(s) used for the noise mixed into the tiles. + output_type (`str`, defaults to `"pil"`): + The output format of the generated video: `"pil"`, `"np"` or `"pt"`. + return_dict (`bool`, defaults to `True`): + Whether or not to return a [`Kandinsky6SRPipelineOutput`] instead of a plain tuple. Examples: Returns: - ``Kandinsky6SRPipelineOutput`` or its tuple representation when - ``return_dict=False``. + [`Kandinsky6SRPipelineOutput`] or `tuple`: + The super-resolved video; a one-element tuple when `return_dict=False`. """ - if (video is None) == (latents is None): - raise ValueError("pass exactly one of `video` or `latents`") - if video is not None: - video = _prepare_batch(video, "video", 3, dtype=torch.uint8) - if latents is not None: - latents = _prepare_batch(latents, "latents", int(self.transformer.in_visual_dim)) - - self.check_inputs( - video=video, - latents=latents, - resolution_scale=resolution_scale, - num_inference_steps=num_inference_steps, - min_overlap=min_overlap, - tiles_batch_size=tiles_batch_size, - kvae_bridge=kvae_bridge, - output_type=output_type, - ) - tiling_scale, pre_upscale = resolve_scale_request(float(resolution_scale)) - batch_size = latents.shape[0] if latents is not None else video.shape[0] - sr_video = list(video.unbind(0)) if video is not None else None - sr_latents = list(latents.unbind(0)) if latents is not None else None + device = self._execution_device + dtype = self.transformer.dtype + + # 1. Preprocess the video and check inputs + video = self.video_processor.preprocess_video(video).to(device) + self.check_inputs(video, resolution_scale, num_inference_steps, lq_noise_scale, min_overlap, tiles_batch_size) + + # 2. The 2.25x route bilinearly pre-upscales the pixels by 1.125x before the 2x path + tiling_scale, pre_upscale = (2, 1.125) if resolution_scale == 2.25 else (int(resolution_scale), 1.0) + if pre_upscale != 1.0: + batch_size, channels, num_frames, height, width = video.shape + snap = self.vae_scale_factor_spatial + target = (round(height * pre_upscale / snap) * snap, round(width * pre_upscale / snap) * snap) + frames = video.permute(0, 2, 1, 3, 4).flatten(0, 1) + frames = F.interpolate(frames, size=target, mode="bilinear", align_corners=False) + video = frames.unflatten(0, (batch_size, num_frames)).permute(0, 2, 1, 3, 4) + batch_size, _, num_frames, height, width = video.shape + + # 3. Tile grid at the input resolution; every tile is refined at the closest trained tile resolution + base_height, base_width = min(self.transformer_tile_sizes, key=lambda hw: abs(hw[1] / hw[0] - width / height)) + snap = self.vae_scale_factor_spatial + if base_height % (snap * tiling_scale) or base_width % (snap * tiling_scale): + raise ValueError( + f"The tile size {(base_height, base_width)} must be divisible by the VAE spatial factor times the" + f" upscale factor ({snap * tiling_scale}) so that tiles align with the latent grid." + ) + tile_height, tile_width = base_height // tiling_scale, base_width // tiling_scale + if height < tile_height or width < tile_width: + raise ValueError( + f"`video` must be at least {tile_height}x{tile_width} pixels for `resolution_scale={resolution_scale}`," + f" got {height}x{width}." + ) + tops = _tile_positions(height, tile_height, min_overlap, snap) + lefts = _tile_positions(width, tile_width, min_overlap, snap) + tile_grid = [(top, left) for top in tops for left in lefts] + + # 4. Low-resolution latent tiles at the trained tile resolution + use_latent_upscaler = self.latent_upscaler is not None and tiling_scale in self.latent_upscaler.config.scales + if use_latent_upscaler: + latents = self.encode_video(video) + lq_tiles = [ + self.latent_upscaler( + latents[ + :, :, :, top // snap : (top + tile_height) // snap, left // snap : (left + tile_width) // snap + ], + scale=tiling_scale, + return_dict=False, + )[0] + for top, left in tile_grid + ] + else: + lq_tiles = [] + for top, left in tile_grid: + pixel_tile = video[:, :, :, top : top + tile_height, left : left + tile_width] + frames = pixel_tile.permute(0, 2, 1, 3, 4).flatten(0, 1) + frames = F.interpolate(frames, size=(base_height, base_width), mode="bilinear", align_corners=False) + pixel_tile = frames.unflatten(0, (batch_size, num_frames)).permute(0, 2, 1, 3, 4) + lq_tiles.append(self.encode_video(pixel_tile)) + lq_tiles = torch.stack(lq_tiles, dim=1).to(dtype) # (batch_size, num_tiles, channels, frames, height, width) + + # 5. Denoise the tiles and blend the decoded tiles into the output canvas + output_height, output_width = height * tiling_scale, width * tiling_scale + video_acc = torch.zeros((batch_size, 3, num_frames, output_height, output_width), device=device) + window = _hann_window_2d(base_height, base_width, device) + weight_acc = torch.zeros((1, 1, 1, output_height, output_width), device=device) + for top, left in tile_grid: + top, left = top * tiling_scale, left * tiling_scale + weight_acc[..., top : top + base_height, left : left + base_width] += window + num_chunks = math.ceil(len(tile_grid) / tiles_batch_size) + + with self.progress_bar(total=batch_size * num_chunks * num_inference_steps) as progress_bar: + for sample_index in range(batch_size): + sample_generator = generator[sample_index] if isinstance(generator, list) else generator + for start in range(0, len(tile_grid), tiles_batch_size): + tile_indices = range(start, min(start + tiles_batch_size, len(tile_grid))) + # The transformer works on `(batch, frames, height, width, channels)` tiles + lq_latents = lq_tiles[sample_index, list(tile_indices)].permute(0, 2, 3, 4, 1) + + # Variance-preserving noise mixed into the low-resolution latents is the starting point + noise = randn_tensor(lq_latents.shape, generator=sample_generator, device=device, dtype=dtype) + latents = math.sqrt(1 - lq_noise_scale**2) * lq_latents + lq_noise_scale * noise + + timesteps_, _ = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps, sigmas) + for t in timesteps_: + # The SR transformer takes `[latent | anchor latent | anchor mask]`; the released checkpoints + # are anchor-free, so the anchor channels are zeros. + anchor = torch.zeros_like(latents) + anchor_mask = torch.zeros((*latents.shape[:-1], 1), dtype=dtype, device=device) + latent_model_input = torch.cat([latents, anchor, anchor_mask], dim=-1) + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=t.expand(latents.shape[0]), + return_dict=False, + )[0] + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + progress_bar.update() + + decoded = self.decode_latents(latents.permute(0, 4, 1, 2, 3)) + for tile, tile_index in zip(decoded, tile_indices): + top, left = tile_grid[tile_index] + top, left = top * tiling_scale, left * tiling_scale + video_acc[sample_index, :, :, top : top + base_height, left : left + base_width] += ( + tile * window + ) - with _execution_device_context(self._execution_device): - try: - if sr_latents is not None and kvae_bridge: - sr_video = [ - _decode_source_latent_video(latent, self.source_vae, self._execution_device) - for latent in sr_latents - ] - sr_latents = None - if sr_video is not None and pre_upscale != 1.0: - spatial_factor = self.spatial_factor - sr_video = [pre_upscale_video(item, pre_upscale, spatial_factor) for item in sr_video] + video = (video_acc / weight_acc).clamp(-1, 1) + video = self.video_processor.postprocess_video(video, output_type=output_type) - run_config = RunConfig( - device=str(self._execution_device), - num_steps=num_inference_steps, - seed=self._seed_from_generator(generator, self._execution_device), - min_overlap=min_overlap, - tiles_batch_size=tiles_batch_size * batch_size, - resolution_scale=tiling_scale, - ) - progress_input = sr_latents[0] if sr_latents is not None else sr_video[0] - progress_total, progress_steps, progress_tiles = self._sr_progress_plan( - progress_input, - run_config, - latent_input=sr_latents is not None, - batch_size=batch_size, - ) - with self.progress_bar(total=progress_total) as progress_bar: - progress_bar.set_description(f"SR [{progress_steps} steps x {progress_tiles} tiles]") - has_latent_upscaler = run_config.resolution_scale in self._loaded_lu_scales() - if sr_latents is not None: - frames = self._run_batched_tiles( - sr_latents, - run_config, - latent_input=True, - cached_text_embeds=cached_text_embeds, - progress_bar=progress_bar, - ) - elif has_latent_upscaler: - lr_latents = [ - encode_lq_video_to_lr_latent(item, self.vae, self._execution_device) for item in sr_video - ] - frames = self._run_batched_tiles( - lr_latents, - run_config, - latent_input=True, - cached_text_embeds=cached_text_embeds, - progress_bar=progress_bar, - ) - else: - frames = self._run_batched_tiles( - sr_video, - run_config, - latent_input=False, - cached_text_embeds=cached_text_embeds, - progress_bar=progress_bar, - ) - finally: - # Diffusers expects custom pipelines to restore the offloaded - # modules after a call. This also keeps the next invocation from - # observing a partially resident model chain. - self.maybe_free_model_hooks() + # Offload all models + self.maybe_free_model_hooks() - output = Kandinsky6SRPipelineOutput(frames=self._format_frames(frames, output_type)) if not return_dict: - return (output.frames,) - return output + return (video,) + return Kandinsky6SRPipelineOutput(frames=video) diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index 30a1fe03a977..13a48c5c579f 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -12,35 +12,35 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Kandinsky 6 TI2VA Diffusers pipeline.""" - import copy +import inspect import math -import warnings from collections.abc import Callable -from pathlib import Path -from typing import Any -import numpy as np +import PIL.Image import torch -from torch import Tensor, nn +from transformers import CLIPTextModel, CLIPTokenizer, Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor +from ...image_processor import PipelineImageInput +from ...models import AutoencoderKLHunyuanVideo, Kandinsky6Transformer3DModel, MMAudioVAE +from ...schedulers import FlowMatchEulerDiscreteScheduler, PiflowScheduler from ...utils import replace_example_docstring from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import Kandinsky6TI2VAPipelineOutput EXAMPLE_DOC_STRING = """ Examples: - ```python >>> import torch >>> from diffusers import Kandinsky6TI2VAPipeline >>> from diffusers.utils import encode_video - >>> model_id = "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers" - >>> pipe = Kandinsky6TI2VAPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16) + >>> pipe = Kandinsky6TI2VAPipeline.from_pretrained( + ... "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers", torch_dtype=torch.bfloat16 + ... ) >>> pipe.enable_model_cpu_offload() >>> output = pipe( @@ -50,16 +50,12 @@ ... num_frames=121, ... num_inference_steps=16, ... guidance_scale=1.0, - ... sample_audio=True, ... ) - - >>> video = output.frames[0].permute(1, 2, 3, 0) - >>> audio = torch.as_tensor(output.audio[0])[:, None].repeat(1, 2) >>> encode_video( - ... video, + ... output.frames[0], ... fps=24, ... output_path="output.mp4", - ... audio=audio, + ... audio=output.audio[0][None], ... audio_sample_rate=pipe.audio_sample_rate, ... ) ``` @@ -78,6 +74,7 @@ "<|im_start|>user\n{}<|im_end|>", ] ) +# Number of template tokens preceding the user prompt in the Qwen sequence. _QWEN_CROP_START = 129 _CLIP_MAX_LENGTH = 77 _T2VA_EXPANSION_INSTRUCTION = ( @@ -102,35 +99,113 @@ ) +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +): + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`list[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`list[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + class Kandinsky6TI2VAPipeline(DiffusionPipeline): - r"""Pipeline for text/image-to-video-and-audio generation with Kandinsky 6. + r""" + Pipeline for text/image-to-video-and-audio generation with Kandinsky 6. + + Video and audio latents are denoised together by a single multimodal transformer, conditioned on Qwen2.5-VL text + tokens and a CLIP pooled embedding. An optional reference image conditions the first frame. - This pipeline inherits the generic loading, device placement, and progress handling provided by - [`DiffusionPipeline`]. The transformer must expose the K6 multimodal forward contract, while the VAE adapters - must implement the portable video/audio post-processing contracts. + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). Args: transformer ([`Kandinsky6Transformer3DModel`]): - Multimodal K6 transformer used to denoise video and audio latents. - vae (`nn.Module`): - Video VAE used to encode and decode video latents. - text_encoder: - Qwen2.5-VL text encoder for token-level embeddings. - audio_vae (`nn.Module`, *optional*): - Audio VAE used to decode generated audio latents. May be `None` only when `sample_audio=False`. - scheduler: - Diffusion scheduler used by the denoising loop. - tokenizer: - Qwen2.5-VL processor. - text_encoder_2: - CLIP text encoder for pooled embeddings. - tokenizer_2: - CLIP tokenizer. + Multimodal transformer that denoises the video and audio latents. + vae ([`AutoencoderKLHunyuanVideo`]): + Video VAE used to encode the reference image and decode the generated video. + text_encoder ([`~transformers.Qwen2_5_VLForConditionalGeneration`]): + Qwen2.5-VL model providing the token-level text embeddings and, optionally, prompt expansion. + tokenizer ([`~transformers.Qwen2_5_VLProcessor`]): + Processor of `text_encoder`. + text_encoder_2 ([`~transformers.CLIPTextModel`]): + CLIP text encoder providing the pooled text embedding. + tokenizer_2 ([`~transformers.CLIPTokenizer`]): + Tokenizer of `text_encoder_2`. + scheduler ([`FlowMatchEulerDiscreteScheduler`] or [`PiflowScheduler`]): + Scheduler used with `transformer` to denoise the latents. Distilled checkpoints ship with a + [`PiflowScheduler`] and must be run with `guidance_scale=1.0`. + audio_vae ([`MMAudioVAE`], *optional*): + Audio VAE used to decode the generated audio latents. Only needed when `sample_audio=True`. """ model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae->audio_vae" _optional_components = ["audio_vae"] - _callback_tensor_inputs = ["latents"] + _callback_tensor_inputs = ["latents", "audio_latents", "prompt_embeds", "negative_prompt_embeds"] _DEFAULT_NEGATIVE_PROMPT = ( "Static, 2D cartoon, cartoon, 2d animation, paintings, images, " "worst quality, low quality, ugly, deformed, walking backwards" @@ -138,209 +213,55 @@ class Kandinsky6TI2VAPipeline(DiffusionPipeline): def __init__( self, - transformer: nn.Module, - vae: nn.Module, - text_encoder: Any, - audio_vae: nn.Module | None, - scheduler: Any, - tokenizer: Any, - text_encoder_2: Any, - tokenizer_2: Any, + transformer: Kandinsky6Transformer3DModel, + vae: AutoencoderKLHunyuanVideo, + text_encoder: Qwen2_5_VLForConditionalGeneration, + tokenizer: Qwen2_5_VLProcessor, + text_encoder_2: CLIPTextModel, + tokenizer_2: CLIPTokenizer, + scheduler: FlowMatchEulerDiscreteScheduler | PiflowScheduler, + audio_vae: MMAudioVAE | None = None, ) -> None: super().__init__() self.register_modules( transformer=transformer, vae=vae, text_encoder=text_encoder, - audio_vae=audio_vae, tokenizer=tokenizer, text_encoder_2=text_encoder_2, tokenizer_2=tokenizer_2, scheduler=scheduler, + audio_vae=audio_vae, ) - scale_factor = ( - self.transformer.config.get("scale_factor", (1.0, 2.0, 2.0)) - if getattr(self, "transformer", None) is not None - else (1.0, 2.0, 2.0) - ) - self.scale_factor = tuple(float(value) for value in scale_factor) - - audio_config = self.audio_vae.config if getattr(self, "audio_vae", None) is not None else {} - self.audio_sample_rate = int(audio_config.get("sample_rate", 44_100)) - self.audio_downsample_factor = int(audio_config.get("downsample_factor", 1_024)) - - @staticmethod - def _value_batch_size(value: Any) -> int: - if value is None: - return 1 - if isinstance(value, list): - return len(value) - if isinstance(value, Tensor): - return int(value.shape[0]) if value.ndim in (3, 5) else 1 - return 1 - - @classmethod - def _batch_size( - cls, - prompt: str | list[str] | None, - negative_prompt: str | list[str] | None, - image: str | object | list[str | object] | None, - latents: Tensor | None, - audio_latents: Tensor | None, - prompt_embeds: Tensor | None, - negative_prompt_embeds: Tensor | None, - ) -> int: - sizes = [ - cls._value_batch_size(value) - for value in ( - prompt, - negative_prompt, - image, - latents, - audio_latents, - prompt_embeds, - negative_prompt_embeds, - ) - ] - batch_size = max(sizes) - if batch_size < 1 or any(size not in (1, batch_size) for size in sizes): - raise ValueError(f"all batched inputs must have the same batch size, got {sizes}") - return batch_size - - @staticmethod - def _as_batch(value: Any, batch_size: int, name: str) -> list[Any] | None: - if value is None: - return None - values = value if isinstance(value, list) else [value] - if len(values) == 1: - return values * batch_size - if len(values) != batch_size: - raise ValueError(f"{name} must contain one item or {batch_size} items, got {len(values)}") - return values - - def check_inputs( - self, - prompt: str | list[str] | None, - negative_prompt: str | list[str] | None, - height: int, - width: int, - num_frames: int, - num_inference_steps: int, - max_sequence_length: int, - output_type: str, - prompt_embeds: Tensor | None = None, - pooled_prompt_embeds: Tensor | None = None, - negative_prompt_embeds: Tensor | None = None, - negative_pooled_prompt_embeds: Tensor | None = None, - callback_on_step_end_tensor_inputs: list[str] | None = None, - sample_audio: bool = True, - image: str | object | None = None, - visual_cond_scheme: str = "pretrain", - ) -> None: - r"""Validate arguments shared by text-to-video and image-to-video calls. - - Args: - prompt (`str` or `list[str]`, *optional*): Text prompt or a batch of prompts. - negative_prompt (`str` or `list[str]`, *optional*): Negative text prompt or a batch of prompts. - height (`int`): Requested output height in pixels. - width (`int`): Requested output width in pixels. - num_frames (`int`): Requested output frame count. - num_inference_steps (`int`): Number of denoising steps. - max_sequence_length (`int`): Maximum Qwen prompt length after the template. - output_type (`str`): One of `"pt"`, `"torch"`, `"np"`, `"numpy"`, or `"latent"`. - prompt_embeds (`torch.Tensor`, *optional*): Precomputed Qwen prompt embeddings. - pooled_prompt_embeds (`torch.Tensor`, *optional*): Precomputed CLIP pooled prompt embeddings. - negative_prompt_embeds (`torch.Tensor`, *optional*): Precomputed negative Qwen embeddings. - negative_pooled_prompt_embeds (`torch.Tensor`, *optional*): Precomputed negative CLIP pooled embeddings. - callback_on_step_end_tensor_inputs (`list[str]`, *optional*): Tensor names exposed to the step callback. - sample_audio (`bool`, *optional*, defaults to `True`): Whether to generate and decode audio. - image (`str`, `PIL.Image.Image`, or a list thereof, *optional*): Reference image or batch of images. - visual_cond_scheme (`str`, *optional*, defaults to `"pretrain"`): Image-conditioning scheme when `image` - is set. - - Raises: - ValueError: If an input combination is unsupported. - """ - for name, value in (("prompt", prompt), ("negative_prompt", negative_prompt)): - if value is not None and not isinstance(value, (str, list)): - raise ValueError(f"`{name}` has to be a `str` or a `list`, but is {type(value)}") - if isinstance(value, list) and (not value or not all(isinstance(item, str) for item in value)): - raise ValueError(f"`{name}` must be a non-empty list of strings") - - if prompt is None and prompt_embeds is None: - raise ValueError("provide either `prompt` or `prompt_embeds`") - if prompt is not None and prompt_embeds is not None: - raise ValueError("provide either `prompt` or `prompt_embeds`, not both") - if (prompt_embeds is None) != (pooled_prompt_embeds is None): - raise ValueError("`prompt_embeds` and `pooled_prompt_embeds` must be provided together") - if negative_prompt is not None and negative_prompt_embeds is not None: - raise ValueError("provide either `negative_prompt` or `negative_prompt_embeds`, not both") - if (negative_prompt_embeds is None) != (negative_pooled_prompt_embeds is None): - raise ValueError("`negative_prompt_embeds` and `negative_pooled_prompt_embeds` must be provided together") - if num_frames < 1: - raise ValueError("num_frames must be positive") - if num_inference_steps < 1: - raise ValueError("num_inference_steps must be positive") - if height < 1 or width < 1: - raise ValueError("height and width must be positive") - if height % 8 != 0 or width % 8 != 0: - raise ValueError("height and width must be divisible by 8") - if output_type not in ("pt", "torch", "np", "numpy", "latent"): - raise ValueError("output_type must be 'pt', 'torch', 'np', 'numpy', or 'latent'") - if max_sequence_length < 1 or max_sequence_length > 1024: - raise ValueError("max_sequence_length must be between 1 and 1024") - if sample_audio and self.audio_vae is None: - raise ValueError("sample_audio=True requires an audio_vae") - - patch_size = tuple(int(value) for value in getattr(self.transformer, "patch_size", (1, 2, 2))) - if len(patch_size) < 3 or any(value <= 0 for value in patch_size): - raise ValueError(f"transformer.patch_size must contain three positive values, got {patch_size}") - if (height // 8) % patch_size[1] != 0 or (width // 8) % patch_size[2] != 0: - raise ValueError( - "height and width must produce whole DiT spatial patches; " - f"got latent size {(height // 8, width // 8)} and patch size {patch_size}" - ) - if callback_on_step_end_tensor_inputs is not None and not all( - name in self._callback_tensor_inputs for name in callback_on_step_end_tensor_inputs - ): - raise ValueError( - f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, " - f"but found {callback_on_step_end_tensor_inputs}" - ) - if visual_cond_scheme not in ("pretrain", "i2v", "tail_cond_first_frame"): - raise ValueError(f"unknown visual_cond_scheme={visual_cond_scheme!r}") - if image is not None: - if not getattr(self.transformer, "visual_cond", False): - raise ValueError("image conditioning requires transformer.visual_cond=True") - if ( - visual_cond_scheme == "tail_cond_first_frame" - and getattr(self.transformer, "visual_token_type_num_embeddings", 0) < 2 - ): - raise ValueError("tail_cond_first_frame requires transformer.visual_token_type_num_embeddings >= 2") - - def _encode_single_prompt( - self, text: str | list[str], max_sequence_length: int, device: torch.device - ) -> tuple[Tensor, Tensor, Tensor | None]: - r"""Encode one prompt (or batch of prompts) with Qwen2.5-VL and CLIP. + self.vae_scale_factor_spatial = ( + self.vae.config.spatial_compression_ratio if getattr(self, "vae", None) is not None else 8 + ) + self.vae_scale_factor_temporal = ( + self.vae.config.temporal_compression_ratio if getattr(self, "vae", None) is not None else 4 + ) + self.transformer_patch_size = ( + tuple(self.transformer.config.patch_size) if getattr(self, "transformer", None) is not None else (1, 2, 2) + ) + self.audio_sample_rate = ( + self.audio_vae.config.sample_rate if getattr(self, "audio_vae", None) is not None else 44_100 + ) + self.audio_latent_hop_length = ( + self.audio_vae.latent_hop_length if getattr(self, "audio_vae", None) is not None else 1024 + ) + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) - Args: - text (`str` or `list[str]`): Prompt or prompts to encode. - max_sequence_length (`int`): Maximum number of Qwen text tokens. - device (`torch.device`): Execution device to move tokenized inputs to. This must be the pipeline's - `_execution_device`, not a component's own parameter device: under CPU/group offloading a - component's resting parameter device can be `meta` or `cpu` even though it computes on the - accelerator once its forward hook runs. + def _get_prompt_embeds( + self, prompt: list[str], max_sequence_length: int, device: torch.device + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Encode prompts with Qwen2.5-VL (token embeddings) and CLIP (pooled embedding). - Returns: - `tuple[torch.Tensor, torch.Tensor, torch.Tensor or None]`: `(prompt_embeds, pooled_prompt_embeds, - attention_mask)`. `attention_mask` is `None` when every prompt in the batch fills - `max_sequence_length` (an all-`True` mask carries no extra information). + Returns the token embeddings, the pooled embeddings and a boolean padding mask, which is `None` when no prompt + in the batch is padded (a mask without padding carries no information, and dropping it keeps every attention + backend available). """ - texts = [text] if isinstance(text, str) else text - full_texts = [_PROMPT_TEMPLATE.format(item) for item in texts] inputs = self.tokenizer( - text=full_texts, + text=[_PROMPT_TEMPLATE.format(item) for item in prompt], images=None, videos=None, max_length=max_sequence_length + _QWEN_CROP_START, @@ -355,15 +276,12 @@ def _encode_single_prompt( output_hidden_states=True, ) prompt_embeds = qwen_output["hidden_states"][-1][:, _QWEN_CROP_START:] - attention_mask = inputs["attention_mask"][:, _QWEN_CROP_START:].to(dtype=torch.bool) - if attention_mask.all(): - # An all-True mask carries no information beyond "attend to everything", which is what the - # transformer already does when no mask is passed (see transformer_kandinsky6.py's - # _normalize_attn_mask). Drop it so the dense mask isn't threaded through the denoise loop. - attention_mask = None + prompt_attention_mask = inputs["attention_mask"][:, _QWEN_CROP_START:].to(dtype=torch.bool) + if prompt_attention_mask.all(): + prompt_attention_mask = None clip_inputs = self.tokenizer_2( - texts, + prompt, max_length=_CLIP_MAX_LENGTH, truncation=True, add_special_tokens=True, @@ -371,56 +289,88 @@ def _encode_single_prompt( return_tensors="pt", ).to(device) pooled_prompt_embeds = self.text_encoder_2(**clip_inputs)["pooler_output"] - return prompt_embeds, pooled_prompt_embeds, attention_mask + return prompt_embeds, pooled_prompt_embeds, prompt_attention_mask def encode_prompt( self, prompt: str | list[str], negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, - prompt_embeds: Tensor | None = None, - pooled_prompt_embeds: Tensor | None = None, - prompt_attention_mask: Tensor | None = None, - negative_prompt_embeds: Tensor | None = None, - negative_pooled_prompt_embeds: Tensor | None = None, - negative_prompt_attention_mask: Tensor | None = None, + num_videos_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, device: torch.device | None = None, - ) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None, Tensor | None, Tensor | None]: - r"""Encode the positive and (optionally) negative prompt. + dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, ...]: + r""" + Encodes the prompt into text encoder hidden states. Args: - prompt (`str` or `list[str]`): Prompt or prompts to encode. - negative_prompt (`str` or `list[str]`, *optional*): Negative prompt or prompts, used when - `do_classifier_free_guidance=True` and `negative_prompt_embeds` is not already provided. - do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): Whether to also encode - `negative_prompt`. - prompt_embeds (`torch.Tensor`, *optional*): Precomputed Qwen prompt embeddings, skipping encoding of - `prompt`. - pooled_prompt_embeds (`torch.Tensor`, *optional*): Precomputed CLIP pooled prompt embeddings. - prompt_attention_mask (`torch.Tensor`, *optional*): Qwen attention mask paired with `prompt_embeds`. - negative_prompt_embeds (`torch.Tensor`, *optional*): Precomputed negative Qwen embeddings. - negative_pooled_prompt_embeds (`torch.Tensor`, *optional*): Precomputed negative CLIP pooled embeddings. - negative_prompt_attention_mask (`torch.Tensor`, *optional*): Qwen attention mask paired with - `negative_prompt_embeds`. - max_sequence_length (`int`, *optional*, defaults to 1024): Maximum Qwen prompt length after the - template. - device (`torch.device`, *optional*): Execution device for newly-encoded prompts. Defaults to - `self._execution_device`. - - Returns: - `tuple[torch.Tensor, ...]`: `(prompt_embeds, pooled_prompt_embeds, prompt_attention_mask, - negative_prompt_embeds, negative_pooled_prompt_embeds, negative_prompt_attention_mask)`. + prompt (`str` or `list[str]`): + Prompt to be encoded. + negative_prompt (`str` or `list[str]`, *optional*): + The prompt not to guide the generation. Ignored when `do_classifier_free_guidance` is `False`. + do_classifier_free_guidance (`bool`, defaults to `True`): + Whether to also encode the negative prompt. + num_videos_per_prompt (`int`, defaults to `1`): + Number of videos generated per prompt; the embeddings are repeated accordingly. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated Qwen2.5-VL text embeddings. Skips encoding `prompt`. + pooled_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated CLIP pooled text embeddings. Must be given together with `prompt_embeds`. + prompt_attention_mask (`torch.Tensor`, *optional*): + Boolean padding mask of `prompt_embeds`. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative Qwen2.5-VL text embeddings. + negative_pooled_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative CLIP pooled text embeddings. + negative_prompt_attention_mask (`torch.Tensor`, *optional*): + Boolean padding mask of `negative_prompt_embeds`. + max_sequence_length (`int`, defaults to `1024`): + Maximum number of prompt tokens after the chat template. + device (`torch.device`, *optional*): + Device to run the text encoders on. + dtype (`torch.dtype`, *optional*): + Dtype of the returned embeddings. """ device = device or self._execution_device + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt_embeds is None: - prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self._encode_single_prompt( + prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self._get_prompt_embeds( prompt, max_sequence_length, device ) if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or self._DEFAULT_NEGATIVE_PROMPT + negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt = negative_prompt * len(prompt) if len(negative_prompt) == 1 else negative_prompt negative_prompt_embeds, negative_pooled_prompt_embeds, negative_prompt_attention_mask = ( - self._encode_single_prompt(negative_prompt, max_sequence_length, device) + self._get_prompt_embeds(negative_prompt, max_sequence_length, device) ) + + prompt_embeds = prompt_embeds.to(device=device, dtype=dtype).repeat_interleave(num_videos_per_prompt, dim=0) + pooled_prompt_embeds = pooled_prompt_embeds.to(device=device, dtype=dtype).repeat_interleave( + num_videos_per_prompt, dim=0 + ) + if prompt_attention_mask is not None: + prompt_attention_mask = prompt_attention_mask.to(device).repeat_interleave(num_videos_per_prompt, dim=0) + if do_classifier_free_guidance: + negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype).repeat_interleave( + num_videos_per_prompt, dim=0 + ) + negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.to( + device=device, dtype=dtype + ).repeat_interleave(num_videos_per_prompt, dim=0) + if negative_prompt_attention_mask is not None: + negative_prompt_attention_mask = negative_prompt_attention_mask.to(device).repeat_interleave( + num_videos_per_prompt, dim=0 + ) + return ( prompt_embeds, pooled_prompt_embeds, @@ -433,312 +383,207 @@ def encode_prompt( def expand_prompts( self, prompt: str | list[str], - image: Any | list[Any] | None = None, + image: PIL.Image.Image | list[PIL.Image.Image] | None = None, max_sequence_length: int = 1024, - mode: str | None = None, generator: torch.Generator | None = None, ) -> str | list[str]: - """Expand T2VA prompts, or image-grounded I2VA prompts when an image is supplied. + r""" + Rewrites short prompts into detailed video+audio prompts with the Qwen2.5-VL text encoder, grounding them on + the reference image when one is given. Args: - prompt: Prompt or a batch of prompts to expand. - image: Optional reference image or batch of reference images, for `mode="i2va"`. - max_sequence_length: Maximum number of new tokens to generate per prompt. - mode: `"t2va"` or `"i2va"`. Defaults to `"i2va"` when `image` is given, else `"t2va"`. - generator: Optional generator whose seed makes the (sampled) expansion reproducible. + prompt (`str` or `list[str]`): + Prompt or prompts to expand. + image (`PIL.Image.Image` or `list[PIL.Image.Image]`, *optional*): + Reference image(s) of an image-to-video call. + max_sequence_length (`int`, defaults to `1024`): + Maximum number of generated tokens per prompt. + generator (`torch.Generator`, *optional*): + Seeds the sampled expansion. `generate` draws from the global RNG, so the global RNG is seeded from + this generator's seed; later `randn_tensor` calls keep using `generator` directly. Returns: - The expanded prompt, or a list of expanded prompts when `prompt` is a list. + `str` or `list[str]`: The expanded prompt(s). """ if isinstance(prompt, list): - images = image if isinstance(image, (list, tuple)) else [image] * len(prompt) + images = image if isinstance(image, list) else [image] * len(prompt) return [ self.expand_prompts( - item, image=item_image, max_sequence_length=max_sequence_length, mode=mode, generator=generator + item, image=item_image, max_sequence_length=max_sequence_length, generator=generator ) for item, item_image in zip(prompt, images, strict=True) ] - mode = mode or ("i2va" if image is not None else "t2va") - if mode not in ("t2va", "i2va"): - raise ValueError(f"unknown prompt expansion mode={mode!r}") - image = self._prepare_expansion_image(image) if image is not None else None - instruction = (_I2VA_EXPANSION_INSTRUCTION if mode == "i2va" else _T2VA_EXPANSION_INSTRUCTION).format(prompt) - content = [] - if image is not None: - content.append({"type": "image", "image": image}) + if image is not None and not isinstance(image, PIL.Image.Image): + raise ValueError("`expand_prompts` expects `image` as a `PIL.Image.Image`") + + instruction = (_I2VA_EXPANSION_INSTRUCTION if image is not None else _T2VA_EXPANSION_INSTRUCTION).format( + prompt + ) + content = [{"type": "image", "image": image}] if image is not None else [] content.append({"type": "text", "text": instruction}) - messages = [{"role": "user", "content": content}] - text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) - device = self._execution_device + text = self.tokenizer.apply_chat_template( + [{"role": "user", "content": content}], tokenize=False, add_generation_prompt=True + ) inputs = self.tokenizer( text=[text], images=[image] if image is not None else None, videos=None, padding=True, return_tensors="pt", - ).to(device) + ).to(self._execution_device) if generator is not None: - # `generate()` samples from the global RNG rather than accepting a `torch.Generator`, so seed - # the global RNG from the pipeline's own generator to make the expansion reproducible. This - # doesn't affect later `randn_tensor(..., generator=generator)` calls, which use `generator` - # directly rather than the global RNG state. torch.manual_seed(generator.initial_seed()) generated = self.text_encoder.generate(**inputs, max_new_tokens=max_sequence_length) - qwen_crop_start = inputs["input_ids"].shape[1] - trimmed = [output[qwen_crop_start:] for output in generated] - return self.tokenizer.batch_decode(trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - - @staticmethod - def _prepare_expansion_image(image: Any) -> Any: - if isinstance(image, (str, Path)): - from PIL import Image - - with Image.open(image) as pil_image: - return pil_image.convert("RGB").copy() - return image - - @staticmethod - def _load_pil_rgb(image: str | object): - from PIL import Image - - if isinstance(image, str): - try: - pil_image = Image.open(image).convert("RGB") - pil_image.load() - except Exception as exc: - raise ValueError(f"Cannot decode i2va input image {image!r}: {exc}") from exc - elif isinstance(image, Image.Image): - pil_image = image.convert("RGB") - else: - raise TypeError(f"i2va image must be a path or PIL image, got {type(image).__name__}") - return pil_image - - def _encode_i2va_first_frame(self, image: str | object, device: torch.device, height: int, width: int) -> Tensor: - """Resize, center-crop, and VAE-encode one reference image into the packed K6 first-frame latent.""" - try: - import torchvision.transforms.functional as TF - except (ImportError, OSError) as exc: - raise RuntimeError( - "I2VA image processing requires torchvision. Install it with `pip install torchvision`." - ) from exc - - pil_image = self._load_pil_rgb(image) - tensor = TF.pil_to_tensor(pil_image).unsqueeze(0) - src_h, src_w = tensor.shape[-2:] - scale = min(src_h / height, src_w / width) - tensor = TF.resize(tensor, (int(src_h / scale), int(src_w / scale))) - cur_h, cur_w = tensor.shape[-2:] - tensor = TF.crop(tensor, (cur_h - height) // 2, (cur_w - width) // 2, height, width) - - tensor = tensor.to(device=device, dtype=next(self.vae.parameters()).dtype) / 127.5 - 1.0 - tensor = tensor.transpose(0, 1).unsqueeze(0) - encoded = self.vae.encode(tensor) - posterior = encoded.latent_dist if hasattr(encoded, "latent_dist") else encoded[0] - latent = posterior.sample() - latent = latent.squeeze(0).permute(1, 2, 3, 0) - return latent * self.vae.config.scaling_factor + generated = generated[:, inputs["input_ids"].shape[1] :] + return self.tokenizer.batch_decode(generated, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + + def check_inputs( + self, + prompt, + negative_prompt, + height, + width, + num_frames, + image, + visual_cond_scheme, + sample_audio, + prompt_embeds=None, + pooled_prompt_embeds=None, + negative_prompt_embeds=None, + negative_pooled_prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + ): + spatial_multiple = self.vae_scale_factor_spatial * max(self.transformer_patch_size[1:]) + if height % spatial_multiple != 0 or width % spatial_multiple != 0: + raise ValueError( + f"`height` and `width` have to be divisible by {spatial_multiple} but are {height} and {width}." + ) + if num_frames < 1: + raise ValueError(f"`num_frames` has to be positive but is {num_frames}.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Please make sure to only forward one of the two." + ) + if (prompt_embeds is None) != (pooled_prompt_embeds is None): + raise ValueError("`prompt_embeds` and `pooled_prompt_embeds` must be provided together.") + if (negative_prompt_embeds is None) != (negative_pooled_prompt_embeds is None): + raise ValueError("`negative_prompt_embeds` and `negative_pooled_prompt_embeds` must be provided together.") + + if sample_audio and getattr(self, "audio_vae", None) is None: + raise ValueError("`sample_audio=True` requires an `audio_vae`.") + if visual_cond_scheme not in ("pretrain", "i2v", "tail_cond_first_frame"): + raise ValueError( + f"`visual_cond_scheme` must be one of 'pretrain', 'i2v' or 'tail_cond_first_frame' but is {visual_cond_scheme!r}." + ) + if image is not None: + if not self.transformer.config.visual_cond: + raise ValueError("Image conditioning requires a transformer with `visual_cond=True`.") + if ( + visual_cond_scheme == "tail_cond_first_frame" + and self.transformer.config.visual_token_type_num_embeddings < 2 + ): + raise ValueError( + "`tail_cond_first_frame` requires a transformer with `visual_token_type_num_embeddings >= 2`." + ) + + def encode_image( + self, + image: PipelineImageInput, + height: int, + width: int, + device: torch.device, + dtype: torch.dtype, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + r""" + Encodes the reference image(s) into first-frame latents of shape `(batch_size, latent_height, latent_width, + latent_channels)`, scaled by the VAE `scaling_factor`. PIL images are resized and center-cropped to `height x + width`; tensors and arrays must already have that size. + """ + is_pil = isinstance(image, PIL.Image.Image) or ( + isinstance(image, list) and isinstance(image[0], PIL.Image.Image) + ) + image = self.video_processor.preprocess( + image, height=height, width=width, resize_mode="crop" if is_pil else "default" + ) + image = image.to(device=device, dtype=self.vae.dtype).unsqueeze(2) + latents = retrieve_latents(self.vae.encode(image), generator=generator) * self.vae.config.scaling_factor + return latents[:, :, 0].permute(0, 2, 3, 1).to(dtype) def prepare_latents( self, batch_size: int, num_channels_latents: int, - num_latent_frames: int, height: int, width: int, + num_frames: int, dtype: torch.dtype, device: torch.device, - generator: torch.Generator | None, - latents: Tensor | None = None, - ) -> Tensor: - shape = (batch_size, num_latent_frames, height // 8, width // 8, num_channels_latents) + generator: torch.Generator | list[torch.Generator] | None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + r"""Returns video latents in the transformer's `(batch_size, num_frames, height, width, channels)` layout. + A user-provided `latents` tensor is expected in the `(batch_size, channels, num_frames, height, width)` layout. + """ + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + latent_height = height // self.vae_scale_factor_spatial + latent_width = width // self.vae_scale_factor_spatial if latents is not None: - latents = latents.to(device=device, dtype=dtype) - if latents.ndim == 4: - latents = latents.unsqueeze(0) - bchtw = (batch_size, num_channels_latents, num_latent_frames, height // 8, width // 8) - if tuple(latents.shape) == bchtw: - latents = latents.permute(0, 2, 3, 4, 1) - if tuple(latents.shape) != shape: - raise ValueError( - f"`latents` must have shape {shape} (or the BCHTW equivalent), got {tuple(latents.shape)}" - ) - return latents + return latents.to(device=device, dtype=dtype).permute(0, 2, 3, 4, 1) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + shape = (batch_size, num_latent_frames, latent_height, latent_width, num_channels_latents) return randn_tensor(shape, generator=generator, device=device, dtype=dtype) def prepare_audio_latents( self, batch_size: int, num_channels_latents: int, - audio_duration: int, + audio_length: int, dtype: torch.dtype, device: torch.device, - generator: torch.Generator | None, - audio_latents: Tensor | None = None, - ) -> Tensor: - shape = (batch_size, audio_duration, num_channels_latents) + generator: torch.Generator | list[torch.Generator] | None, + audio_latents: torch.Tensor | None = None, + ) -> torch.Tensor: + r"""Returns audio latents in the transformer's `(batch_size, audio_length, channels)` layout. A user-provided + `audio_latents` tensor is expected in the `(batch_size, channels, audio_length)` layout.""" if audio_latents is not None: - audio_latents = audio_latents.to(device=device, dtype=dtype) - if audio_latents.ndim == 2: - audio_latents = audio_latents.unsqueeze(0) - if tuple(audio_latents.shape) != shape: - raise ValueError(f"`audio_latents` must have shape {shape}, got {tuple(audio_latents.shape)}") - return audio_latents + return audio_latents.to(device=device, dtype=dtype).permute(0, 2, 1) + shape = (batch_size, audio_length, num_channels_latents) return randn_tensor(shape, generator=generator, device=device, dtype=dtype) - @staticmethod - def _audio_latent_duration( - num_video_latent_frames: int, - *, - fps: float = 24.0, - audio_fps: int = 44100, - downsample_factor: int = 1024, - ) -> int: - """Audio latent length matching K5 T2VA: ceil(sample_frames/fps * audio_fps / downsample).""" - sample_frames = (num_video_latent_frames - 1) * 4 + 1 - return int(math.ceil(sample_frames / fps * audio_fps / downsample_factor)) - - @staticmethod - def _append_i2va_tail_condition( - latents: Tensor, - first_frames: Tensor, - batch_size: int, - num_latent_frames: int, - ) -> tuple[Tensor, Tensor, Tensor]: - """Append one clean reference frame to batched `(B, T, H, W, C)` latents. - - Returns the extended latents, the per-frame token-type ids (`0` = generated, `1` = reference), and a - boolean mask selecting the generated frames. - """ - if latents.shape[0] != batch_size or latents.shape[1] != num_latent_frames: - raise ValueError( - "generated visual latent shape mismatch: expected " - f"({batch_size}, {num_latent_frames}, H, W, C), got {tuple(latents.shape)}" - ) - _, _, height, width, channels = latents.shape - first_frames = first_frames.to(device=latents.device, dtype=latents.dtype) - expected = (batch_size, height, width, channels) - if tuple(first_frames.shape) != expected: - raise ValueError( - f"first-frame latent shape mismatch: expected {expected}, got {tuple(first_frames.shape)}" - ) - - latents = torch.cat([latents, first_frames[:, None]], dim=1) - token_types = torch.cat( - [ - torch.zeros((batch_size, num_latent_frames), dtype=torch.long, device=latents.device), - torch.ones((batch_size, 1), dtype=torch.long, device=latents.device), - ], - dim=1, - ) - return latents, token_types, token_types == 0 - - @staticmethod - def _apply_visual_conditioning(latents: Tensor, first_frames: Tensor | None, visual_cond_scheme: str) -> Tensor: - """Append visual-conditioning channels for batched `(B, T, H, W, C)` video latents.""" - cond = torch.zeros_like(latents) - mask = torch.zeros((*latents.shape[:-1], 1), dtype=latents.dtype, device=latents.device) - if first_frames is None: - return torch.cat([latents, cond, mask], dim=-1) - - first_frames = first_frames.to(device=latents.device, dtype=latents.dtype) - tail_cond = visual_cond_scheme == "tail_cond_first_frame" - if visual_cond_scheme == "i2v": - latents[:, 0] = first_frames - elif tail_cond: - latents[:, -1] = first_frames - elif visual_cond_scheme == "pretrain": - cond[:, 0] = first_frames - else: - raise ValueError(f"unknown visual_cond_scheme={visual_cond_scheme!r}") - mask[:, -1 if tail_cond else 0] = 1 - return torch.cat([latents, cond, mask], dim=-1) - - @staticmethod - def _compute_rope1d(rope: nn.Module, length: int, device: torch.device | None = None) -> Tensor: - """Build 1-D RoPE and move non-persistent tables for Diffusers offload.""" - if length < 1: - raise ValueError(f"rope length must be positive, got {length}") - if device is None: - device = next(rope.buffers()).device - else: - device = torch.device(device) - # Accelerate's model CPU-offload hook only moves the transformer when its forward starts. RoPE is - # materialized before that forward, so explicitly move its non-persistent lookup table first. - rope.to(device) - pos = torch.arange(length, device=device) - return rope(pos) - - @staticmethod - def _compute_visual_rope( - rope: nn.Module, - shape: tuple[int, int, int], - scale_factor: tuple[float, float, float], - device: torch.device | None = None, - ) -> Tensor: - """Build 3-D RoPE and move non-persistent tables for Diffusers offload.""" - t, h, w = (int(shape[0]), int(shape[1]), int(shape[2])) - if t < 1 or h < 1 or w < 1: - raise ValueError(f"visual rope shape must be positive, got {shape}") - if device is None: - device = next(rope.buffers()).device - else: - device = torch.device(device) - rope.to(device) - pos = [torch.arange(t, device=device), torch.arange(h, device=device), torch.arange(w, device=device)] - return rope((t, h, w), pos, tuple(float(value) for value in scale_factor)) - - def _prepare_text_ropes( - self, - text_length: int, - negative_text_length: int, - audio_length: int | None, - device: torch.device, - ) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor | None]: - """Build the video/audio text-branch RoPE for the cond/uncond passes, plus the audio latent RoPE.""" - video_text_rope = self._compute_rope1d(self.transformer.video_text_rope_embeddings, text_length, device=device) - audio_text_rope = self._compute_rope1d(self.transformer.audio_text_rope_embeddings, text_length, device=device) - negative_video_text_rope = self._compute_rope1d( - self.transformer.video_text_rope_embeddings, negative_text_length, device=device - ) - negative_audio_text_rope = self._compute_rope1d( - self.transformer.audio_text_rope_embeddings, negative_text_length, device=device - ) - audio_rope = ( - self._compute_rope1d(self.transformer.audio_rope_embeddings, audio_length, device=device) - if audio_length is not None - else None - ) - return video_text_rope, audio_text_rope, negative_video_text_rope, negative_audio_text_rope, audio_rope - - @staticmethod - def _apply_guidance(cond: Tensor, uncond: Tensor, guidance_scale: float) -> Tensor: - return uncond + guidance_scale * (cond - uncond) - - def _postprocess_video(self, latents: Tensor) -> Tensor: - """Decode video latents `(B, T, H, W, C)` into `(B, 3, T, H, W)` `uint8` frames in `[0, 255]`.""" - frames = (latents / self.vae.config.scaling_factor).permute(0, 4, 1, 2, 3) - # Hunyuan VAE loads as fp16; DiT latents are bf16 — match weight dtype. - vae_dtype = next(self.vae.parameters()).dtype - frames = self.vae.decode(frames.to(dtype=vae_dtype)).sample - return ((frames.clamp(-1.0, 1.0) + 1.0) * 127.5).to(torch.uint8) - - def _postprocess_audio(self, audio_latents: Tensor | None) -> list[np.ndarray] | None: - """Decode audio latents `(B, A, D)` into a list of int16 `(samples,)` waveforms, or `None` in T2V mode.""" - if audio_latents is None: - return None - audio = audio_latents / getattr(self.audio_vae, "scaling_factor", 1.0) - audio = audio + getattr(self.audio_vae, "mean_value", 0.0) - decoded = self.audio_vae.wrapped_decode(audio.transpose(1, 2)) - if decoded.ndim == 1: - decoded = decoded.unsqueeze(0) - elif decoded.ndim == 3 and decoded.shape[1] == 1: - decoded = decoded[:, 0] - return [(np.clip(waveform.cpu().float().numpy(), -1.0, 1.0) * 32767).astype(np.int16) for waveform in decoded] - @property def guidance_scale(self): return self._guidance_scale + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + @property def num_timesteps(self): return self._num_timesteps @@ -756,139 +601,141 @@ def interrupt(self): def __call__( self, prompt: str | list[str] | None = None, - image: str | object | list[str | object] | None = None, + image: PipelineImageInput | None = None, negative_prompt: str | list[str] | None = None, height: int = 512, width: int = 768, num_frames: int = 121, - sample_fps: float = 24.0, + frame_rate: float = 24.0, num_inference_steps: int = 50, - max_sequence_length: int = 1024, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, guidance_scale: float = 5.0, - num_images_per_prompt: int = 1, - generator: torch.Generator | None = None, - latents: Tensor | None = None, - audio_latents: Tensor | None = None, - prompt_embeds: Tensor | None = None, - pooled_prompt_embeds: Tensor | None = None, - negative_prompt_embeds: Tensor | None = None, - negative_pooled_prompt_embeds: Tensor | None = None, - sample_audio: bool = True, + num_videos_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + audio_latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_pooled_prompt_embeds: torch.Tensor | None = None, + sample_audio: bool | None = None, expand_prompts: bool = False, - output_type: str = "pt", - return_dict: bool = True, - callback_on_step_end: Callable[..., Any] | None = None, - callback_on_step_end_tensor_inputs: list[str] | None = None, visual_cond_scheme: str | None = None, - ) -> Kandinsky6TI2VAPipelineOutput | tuple[Tensor | np.ndarray, list[np.ndarray] | None]: - r"""Generate synchronized video and, optionally, audio from text or an image. + max_sequence_length: int = 1024, + output_type: str = "pil", + return_dict: bool = True, + callback_on_step_end: Callable[[int, int, dict], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + ): + r""" + The call function to the pipeline for generation. Args: - prompt (`str` or `list[str]`, *optional*): Text prompt or a batch of prompts. - image (`str`, `PIL.Image.Image`, or a list thereof, *optional*): Reference image or batch of images. - negative_prompt (`str` or `list[str]`, *optional*): Negative prompt. - height (`int`, *optional*, defaults to 512): Output video height in pixels; must be divisible by 8. - width (`int`, *optional*, defaults to 768): Output video width in pixels; must be divisible by 8. - num_frames (`int`, *optional*, defaults to 121): Number of decoded video frames. - sample_fps (`float`, *optional*, defaults to 24.0): Output video frame rate. Kandinsky 6 is trained - for 24.0 fps; other values may affect audio/video alignment. - num_inference_steps (`int`, *optional*, defaults to 50): Number of denoising steps. - max_sequence_length (`int`, *optional*, defaults to 1024): Maximum Qwen prompt length after the - template. - guidance_scale (`float`, *optional*, defaults to 5.0): Classifier-free guidance weight. - num_images_per_prompt (`int`, *optional*, defaults to 1): Number of videos to generate per prompt. - generator (`torch.Generator`, *optional*): Random generator used for latent initialization and, when - `expand_prompts=True`, to seed prompt expansion. - latents (`torch.Tensor`, *optional*): Precomputed video latents. - audio_latents (`torch.Tensor`, *optional*): Precomputed audio latents. - prompt_embeds (`torch.Tensor`, *optional*): Precomputed Qwen prompt embeddings. - pooled_prompt_embeds (`torch.Tensor`, *optional*): Precomputed CLIP pooled prompt embeddings. - negative_prompt_embeds (`torch.Tensor`, *optional*): Precomputed negative Qwen embeddings. - negative_pooled_prompt_embeds (`torch.Tensor`, *optional*): Precomputed negative CLIP pooled embeddings. - sample_audio (`bool`, *optional*, defaults to `True`): Whether to generate synchronized audio. - expand_prompts (`bool`, *optional*, defaults to `False`): Whether to use the built-in Qwen video+audio - prompt expander before encoding. - output_type (`str`, *optional*, defaults to `"pt"`): `"pt"`/`"torch"` for tensors, `"np"`/`"numpy"` for - NumPy arrays, or `"latent"` for undecoded video latents. - return_dict (`bool`, *optional*, defaults to `True`): Whether to return `Kandinsky6TI2VAPipelineOutput`. - callback_on_step_end (`Callable`, *optional*): Callback invoked after each step. - callback_on_step_end_tensor_inputs (`list[str]`, *optional*): Names passed to the callback. - visual_cond_scheme (`str`, *optional*): Image-conditioning scheme. Defaults to `"pretrain"` without - `image` and `"tail_cond_first_frame"` with `image`. + prompt (`str` or `list[str]`, *optional*): + The prompt or prompts to guide the generation. Required unless `prompt_embeds` is given. + image (`PipelineImageInput`, *optional*): + Reference image(s) conditioning the first frame (image-to-video-and-audio). + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the generation. Defaults to the Kandinsky 6 negative prompt. + height (`int`, defaults to `512`): + Height of the generated video in pixels. + width (`int`, defaults to `768`): + Width of the generated video in pixels. + num_frames (`int`, defaults to `121`): + Number of generated frames. + frame_rate (`float`, defaults to `24.0`): + Frame rate the video is generated at; sets the length of the synchronized audio. + num_inference_steps (`int`, defaults to `50`): + The number of denoising steps. Use `16` with the distilled checkpoints. + timesteps (`list[int]`, *optional*): + Custom timesteps for schedulers that support them. + sigmas (`list[float]`, *optional*): + Custom sigmas for schedulers that support them. + guidance_scale (`float`, defaults to `5.0`): + Classifier-free guidance scale. Must be `1.0` with a [`PiflowScheduler`]. + num_videos_per_prompt (`int`, defaults to `1`): + The number of videos to generate per prompt. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Generator(s) used for the initial noise and the reference image encoding. + latents (`torch.Tensor`, *optional*): + Pre-generated video latents of shape `(batch_size, channels, num_latent_frames, latent_height, + latent_width)`. + audio_latents (`torch.Tensor`, *optional*): + Pre-generated audio latents of shape `(batch_size, channels, audio_length)`. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated Qwen2.5-VL text embeddings. + pooled_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated CLIP pooled text embeddings. + negative_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative Qwen2.5-VL text embeddings. + negative_pooled_prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated negative CLIP pooled text embeddings. + sample_audio (`bool`, *optional*): + Whether to generate synchronized audio. Defaults to `True` when the pipeline has an `audio_vae` and + `False` otherwise; `True` requires an `audio_vae`. + expand_prompts (`bool`, defaults to `False`): + Whether to rewrite the prompts with [`~Kandinsky6TI2VAPipeline.expand_prompts`] before encoding. + visual_cond_scheme (`str`, *optional*): + How the reference image conditions the video: `"tail_cond_first_frame"` (default with `image`), + `"i2v"`, or `"pretrain"` (default without `image`). + max_sequence_length (`int`, defaults to `1024`): + Maximum number of prompt tokens after the chat template. + output_type (`str`, defaults to `"pil"`): + The output format of the generated video: `"pil"`, `"np"`, `"pt"` or `"latent"`. + return_dict (`bool`, defaults to `True`): + Whether or not to return a [`Kandinsky6TI2VAPipelineOutput`] instead of a plain tuple. + callback_on_step_end (`Callable`, *optional*): + A function called at the end of each denoising step with `callback_on_step_end(self, step, timestep, + callback_kwargs)`. It may return a dict overriding the listed tensors. + callback_on_step_end_tensor_inputs (`list[str]`, defaults to `["latents"]`): + Tensor inputs passed to `callback_on_step_end`; a subset of `_callback_tensor_inputs`. Examples: Returns: - [`Kandinsky6TI2VAPipelineOutput`] or `tuple`: [`Kandinsky6TI2VAPipelineOutput`] containing video - frames and int16 NumPy audio waveforms if `return_dict=True`, otherwise a `(frames, audio)` tuple. + [`Kandinsky6TI2VAPipelineOutput`] or `tuple`: + The generated video and audio; a `(frames, audio)` tuple when `return_dict=False`. """ - if sample_fps != 24.0: - warnings.warn( - f"Kandinsky 6 was trained for 24.0 fps; received sample_fps={sample_fps}.", - UserWarning, - stacklevel=2, - ) - - # 1. Define call parameters - batch_size = self._batch_size( - prompt, negative_prompt, image, latents, audio_latents, prompt_embeds, negative_prompt_embeds - ) - prompt_batch = self._as_batch(prompt, batch_size, "prompt") - negative_prompt_batch = self._as_batch(negative_prompt, batch_size, "negative_prompt") - image_batch = self._as_batch(image, batch_size, "image") + # 1. Check inputs. Raise error if not correct + if sample_audio is None: + sample_audio = getattr(self, "audio_vae", None) is not None if visual_cond_scheme is None: visual_cond_scheme = "tail_cond_first_frame" if image is not None else "pretrain" - - # 2. Check inputs. Raise error if not correct self.check_inputs( - prompt=prompt_batch, - negative_prompt=negative_prompt_batch, + prompt=prompt, + negative_prompt=negative_prompt, height=height, width=width, num_frames=num_frames, - num_inference_steps=num_inference_steps, - max_sequence_length=max_sequence_length, - output_type=output_type, + image=image, + visual_cond_scheme=visual_cond_scheme, + sample_audio=sample_audio, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - sample_audio=sample_audio, - image=image_batch, - visual_cond_scheme=visual_cond_scheme, ) - callback_on_step_end_tensor_inputs = callback_on_step_end_tensor_inputs or self._callback_tensor_inputs + if isinstance(self.scheduler, PiflowScheduler) and guidance_scale != 1.0: + raise ValueError("`PiflowScheduler` runs the distilled checkpoints and requires `guidance_scale=1.0`.") self._guidance_scale = guidance_scale self._current_timestep = None self._interrupt = False + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + prompt = [prompt] + batch_size = len(prompt) if prompt is not None else prompt_embeds.shape[0] device = self._execution_device - dtype = ( - self.transformer.dtype - if isinstance(getattr(self.transformer, "dtype", None), torch.dtype) - else torch.bfloat16 - ) + dtype = self.transformer.dtype - # 3. Encode the reference image, if any - first_frames = None - if image_batch is not None: - first_frames = torch.cat( - [self._encode_i2va_first_frame(item, device, height, width) for item in image_batch], dim=0 + # 3. Encode input prompt + if expand_prompts and prompt is not None: + prompt = self.expand_prompts( + prompt, image=image, max_sequence_length=max_sequence_length, generator=generator ) - - # 4. Encode input prompt - if expand_prompts and prompt_batch is not None and prompt_embeds is None: - prompt_batch = self.expand_prompts( - prompt_batch, - image=image_batch, - max_sequence_length=max_sequence_length, - mode="i2va" if image_batch is not None else "t2va", - generator=generator, - ) - negative_prompt_batch = negative_prompt_batch or [self._DEFAULT_NEGATIVE_PROMPT] * batch_size - ( prompt_embeds, pooled_prompt_embeds, @@ -897,171 +744,140 @@ def __call__( negative_pooled_prompt_embeds, negative_prompt_attention_mask, ) = self.encode_prompt( - prompt=prompt_batch or [""] * batch_size, - negative_prompt=negative_prompt_batch, - do_classifier_free_guidance=True, - device=device, + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=num_videos_per_prompt, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, ) - prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) - pooled_prompt_embeds = pooled_prompt_embeds.to(device=device, dtype=dtype) - negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.to(device=device, dtype=dtype) - if prompt_attention_mask is not None: - prompt_attention_mask = prompt_attention_mask.to(device=device) - if negative_prompt_attention_mask is not None: - negative_prompt_attention_mask = negative_prompt_attention_mask.to(device=device) - - prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - pooled_prompt_embeds = pooled_prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - negative_prompt_embeds = negative_prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) - if prompt_attention_mask is not None: - prompt_attention_mask = prompt_attention_mask.repeat_interleave(num_images_per_prompt, dim=0) - if negative_prompt_attention_mask is not None: - negative_prompt_attention_mask = negative_prompt_attention_mask.repeat_interleave( - num_images_per_prompt, dim=0 - ) - if first_frames is not None: - first_frames = first_frames.repeat_interleave(num_images_per_prompt, dim=0) - batch_size = batch_size * num_images_per_prompt - - # 5. Prepare latent variables - patch_size = tuple(int(value) for value in self.transformer.patch_size) - num_latent_frames = (num_frames - 1) // 4 + 1 - num_channels_latents = int(self.transformer.in_visual_dim) + + # 4. Encode the reference image + first_frame_latents = None + if image is not None: + first_frame_latents = self.encode_image(image, height, width, device, dtype, generator) + first_frame_latents = first_frame_latents.repeat_interleave(num_videos_per_prompt, dim=0) + + # 5. Prepare timesteps. Audio uses a second scheduler instance so that both modalities keep their own step + # counter. + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, timesteps, sigmas + ) + audio_scheduler = copy.deepcopy(self.scheduler) if sample_audio else None + self._num_timesteps = len(timesteps) + + # 6. Prepare latent variables + batch_size = batch_size * num_videos_per_prompt latents = self.prepare_latents( - batch_size, num_channels_latents, num_latent_frames, height, width, dtype, device, generator, latents + batch_size, + self.transformer.config.in_visual_dim, + height, + width, + num_frames, + dtype, + device, + generator, + latents, ) + num_latent_frames = latents.shape[1] if sample_audio: - audio_duration = self._audio_latent_duration( - num_latent_frames, - fps=sample_fps, - audio_fps=self.audio_sample_rate, - downsample_factor=self.audio_downsample_factor, + audio_length = math.ceil( + ((num_latent_frames - 1) * self.vae_scale_factor_temporal + 1) + / frame_rate + * self.audio_sample_rate + / self.audio_latent_hop_length ) audio_latents = self.prepare_audio_latents( - batch_size, - int(self.transformer.in_audio_dim), - audio_duration, - dtype, - device, - generator, - audio_latents, + batch_size, self.transformer.config.in_audio_dim, audio_length, dtype, device, generator, audio_latents ) - elif audio_latents is not None: - raise ValueError("`audio_latents` cannot be provided when `sample_audio=False`") else: audio_latents = None + # The `tail_cond_first_frame` scheme appends the clean reference frame as an extra, masked frame that reuses + # the first frame's rotary position and carries token type `1`. + tail_cond = image is not None and visual_cond_scheme == "tail_cond_first_frame" visual_token_type_ids = None - generated_visual_mask = None - if image_batch is not None and visual_cond_scheme == "tail_cond_first_frame": - latents, visual_token_type_ids, generated_visual_mask = self._append_i2va_tail_condition( - latents, first_frames, batch_size, num_latent_frames - ) - - # 6. Prepare rotary position embeddings - visual_shape = ( - num_latent_frames // patch_size[0], - (height // 8) // patch_size[1], - (width // 8) // patch_size[2], - ) - if visual_shape[0] < 1: - raise ValueError(f"invalid visual latent shape {visual_shape}") - visual_rope = self._compute_visual_rope( - self.transformer.visual_rope_embeddings, visual_shape, self.scale_factor, device=device - ) - if generated_visual_mask is not None: - # The appended reference frame reuses the first frame's rotary position rather than a real T+1 one. - visual_rope = torch.cat([visual_rope, visual_rope[:1]], dim=0) - video_text_rope, audio_text_rope, negative_video_text_rope, negative_audio_text_rope, audio_rope = ( - self._prepare_text_ropes( - text_length=prompt_embeds.shape[1], - negative_text_length=negative_prompt_embeds.shape[1], - audio_length=audio_latents.shape[1] if audio_latents is not None else None, - device=device, + visual_rope_pos = None + if tail_cond: + latents = torch.cat([latents, first_frame_latents[:, None]], dim=1) + visual_token_type_ids = torch.zeros((batch_size, num_latent_frames + 1), dtype=torch.long, device=device) + visual_token_type_ids[:, -1] = 1 + patch_t, patch_h, patch_w = self.transformer_patch_size + visual_rope_pos = ( + torch.cat( + [ + torch.arange(num_latent_frames // patch_t, device=device), + torch.zeros(1, dtype=torch.long, device=device), + ] + ), + torch.arange(latents.shape[2] // patch_h, device=device), + torch.arange(latents.shape[3] // patch_w, device=device), ) - ) # 7. Denoising loop - is_piflow = bool(getattr(self.scheduler, "is_piflow", False)) - if is_piflow and guidance_scale != 1.0: - raise ValueError("PiflowScheduler requires guidance_scale=1.0") - do_classifier_free_guidance = abs(guidance_scale - 1.0) > 1e-6 - tail_cond = visual_cond_scheme == "tail_cond_first_frame" - - self.scheduler.set_timesteps(num_inference_steps, device=device) - timesteps = self.scheduler.timesteps - # A second scheduler instance keeps audio's step index independent, since the scheduler advances its own - # internal step index on every `.step()` call. - audio_scheduler = copy.deepcopy(self.scheduler) if audio_latents is not None else None - if audio_scheduler is not None: - audio_scheduler.set_timesteps(num_inference_steps, device=device) - self._num_timesteps = len(timesteps) - with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): if self.interrupt: continue self._current_timestep = t - timestep = t.unsqueeze(0).expand(batch_size) - - video_model_input = ( - self._apply_visual_conditioning(latents, first_frames, visual_cond_scheme) - if self.transformer.visual_cond - else latents - ) + timestep = t.expand(batch_size) + + # Visual conditioning channels: [latent | conditioning latent | mask]. The reference frame either + # fills the conditioning channels (`pretrain`) or overwrites a latent frame (`i2v`, tail). + latent_model_input = latents + if self.transformer.config.visual_cond: + cond_latents = torch.zeros_like(latents) + cond_mask = torch.zeros((*latents.shape[:-1], 1), dtype=latents.dtype, device=device) + if first_frame_latents is not None: + if visual_cond_scheme == "pretrain": + cond_latents[:, 0] = first_frame_latents + cond_mask[:, 0] = 1 + elif visual_cond_scheme == "i2v": + latents[:, 0] = first_frame_latents + cond_mask[:, 0] = 1 + else: + latents[:, -1] = first_frame_latents + cond_mask[:, -1] = 1 + latent_model_input = torch.cat([latents, cond_latents, cond_mask], dim=-1) with self.transformer.cache_context("cond"): - model_pred = self.transformer( - x_video=video_model_input, - x_audio=audio_latents, - text_embed=prompt_embeds, - pooled_text_embed=pooled_prompt_embeds, - time=timestep, - visual_rope=visual_rope, - audio_rope=audio_rope, - video_text_rope=video_text_rope, - audio_text_rope=audio_text_rope, - attention_mask=prompt_attention_mask, + noise_pred = self.transformer( + hidden_states=latent_model_input, + audio_hidden_states=audio_latents, + encoder_hidden_states=prompt_embeds, + pooled_projections=pooled_prompt_embeds, + timestep=timestep, + visual_rope_pos=visual_rope_pos, + encoder_attention_mask=prompt_attention_mask, visual_token_type_ids=visual_token_type_ids, + return_dict=False, ) - - if do_classifier_free_guidance: + if self.do_classifier_free_guidance: with self.transformer.cache_context("uncond"): - model_pred_uncond = self.transformer( - x_video=video_model_input, - x_audio=audio_latents, - text_embed=negative_prompt_embeds, - pooled_text_embed=negative_pooled_prompt_embeds, - time=timestep, - visual_rope=visual_rope, - audio_rope=audio_rope, - video_text_rope=negative_video_text_rope, - audio_text_rope=negative_audio_text_rope, - attention_mask=negative_prompt_attention_mask, + noise_pred_uncond = self.transformer( + hidden_states=latent_model_input, + audio_hidden_states=audio_latents, + encoder_hidden_states=negative_prompt_embeds, + pooled_projections=negative_pooled_prompt_embeds, + timestep=timestep, + visual_rope_pos=visual_rope_pos, + encoder_attention_mask=negative_prompt_attention_mask, visual_token_type_ids=visual_token_type_ids, + return_dict=False, ) - if audio_latents is not None: - video_pred = self._apply_guidance(model_pred[0], model_pred_uncond[0], self.guidance_scale) - audio_pred = self._apply_guidance(model_pred[1], model_pred_uncond[1], self.guidance_scale) - else: - video_pred = self._apply_guidance(model_pred, model_pred_uncond, self.guidance_scale) - elif audio_latents is not None: - video_pred, audio_pred = model_pred - else: - video_pred = model_pred - - latents = self.scheduler.step(video_pred, t, latents, return_dict=False)[0] - if audio_latents is not None: - audio_latents = audio_scheduler.step(audio_pred, t, audio_latents, return_dict=False)[0] - if tail_cond and first_frames is not None: - latents[:, -1] = first_frames.to(device=device, dtype=latents.dtype) + noise_pred = tuple( + uncond + self.guidance_scale * (cond - uncond) + for cond, uncond in zip(noise_pred, noise_pred_uncond) + ) + + latents = self.scheduler.step(noise_pred[0], t, latents, return_dict=False)[0] + if sample_audio: + audio_latents = audio_scheduler.step(noise_pred[1], t, audio_latents, return_dict=False)[0] if callback_on_step_end is not None: callback_kwargs = {} @@ -1069,32 +885,38 @@ def __call__( callback_kwargs[k] = locals()[k] callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) latents = callback_outputs.pop("latents", latents) + audio_latents = callback_outputs.pop("audio_latents", audio_latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) progress_bar.update() self._current_timestep = None - # 8. Restore the injected reference frame and drop it from the output - if first_frames is not None: - ff = first_frames.to(device=device, dtype=latents.dtype) - if visual_cond_scheme == "i2v": - latents[:, 0] = ff - elif tail_cond: - latents[:, -1] = ff - if generated_visual_mask is not None: - latents = latents[:, generated_visual_mask[0]] - - # 9. Decode outputs + # 8. Keep the injected reference frame clean and drop the appended tail frame + if first_frame_latents is not None and visual_cond_scheme == "i2v": + latents[:, 0] = first_frame_latents + if tail_cond: + latents = latents[:, :-1] + + # 9. Decode + latents = latents.permute(0, 4, 1, 2, 3) + audio_latents = audio_latents.permute(0, 2, 1) if sample_audio else None if output_type == "latent": - frames: Tensor | np.ndarray = latents.permute(0, 4, 1, 2, 3) + video, audio = latents, audio_latents else: - frames = self._postprocess_video(latents) - if output_type not in ("pt", "torch"): - frames = frames.permute(0, 2, 3, 4, 1).cpu().numpy() - audio = self._postprocess_audio(audio_latents) if sample_audio else None - + video = self.vae.decode(latents.to(self.vae.dtype) / self.vae.config.scaling_factor, return_dict=False)[0] + video = self.video_processor.postprocess_video(video, output_type=output_type) + audio = None + if sample_audio: + audio_latents = audio_latents.to(self.audio_vae.dtype) / self.audio_vae.config.scaling_factor + audio = self.audio_vae.decode(audio_latents, return_dict=False)[0][:, 0].float() + if output_type == "np": + audio = audio.cpu().numpy() + + # Offload all models self.maybe_free_model_hooks() if not return_dict: - return frames, audio - return Kandinsky6TI2VAPipelineOutput(frames=frames, audio=audio) + return (video, audio) + return Kandinsky6TI2VAPipelineOutput(frames=video, audio=audio) diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_output.py b/src/diffusers/pipelines/kandinsky6/pipeline_output.py index 1f47e4d5818a..e626995ed257 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_output.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_output.py @@ -12,34 +12,46 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Output containers for the Kandinsky 6 Diffusers pipelines.""" - -from __future__ import annotations - from dataclasses import dataclass -from typing import Any import numpy as np +import PIL.Image +import torch from ...utils import BaseOutput @dataclass class Kandinsky6TI2VAPipelineOutput(BaseOutput): - """Output of the K6 video-and-audio pipeline.""" + r""" + Output class for [`Kandinsky6TI2VAPipeline`]. - frames: Any - audio: list[np.ndarray] | None + Args: + frames (`torch.Tensor`, `np.ndarray`, or `list[list[PIL.Image.Image]]`): + The generated video. A nested list of length `batch_size` holding `num_frames` PIL images each, or a NumPy + array or torch tensor of shape `(batch_size, num_frames, height, width, channels)` / `(batch_size, + num_frames, channels, height, width)`. With `output_type="latent"`, the video latents of shape + `(batch_size, channels, num_latent_frames, latent_height, latent_width)`. + audio (`torch.Tensor` or `np.ndarray`, *optional*): + The generated waveforms of shape `(batch_size, num_samples)` in `[-1, 1]` at the audio VAE's sample rate, + or `None` when audio was not sampled. With `output_type="latent"`, the audio latents of shape `(batch_size, + channels, audio_length)`. + """ + + frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]] + audio: torch.Tensor | np.ndarray | None = None @dataclass class Kandinsky6SRPipelineOutput(BaseOutput): - """Output of the K6 super-resolution pipeline.""" - - frames: Any + r""" + Output class for [`Kandinsky6SRPipeline`]. + Args: + frames (`torch.Tensor`, `np.ndarray`, or `list[list[PIL.Image.Image]]`): + The super-resolved video. A nested list of length `batch_size` holding `num_frames` PIL images each, or a + NumPy array or torch tensor of shape `(batch_size, num_frames, height, width, channels)` / `(batch_size, + num_frames, channels, height, width)`. + """ -__all__ = [ - "Kandinsky6SRPipelineOutput", - "Kandinsky6TI2VAPipelineOutput", -] + frames: torch.Tensor | np.ndarray | list[list[PIL.Image.Image]] diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index baad89b768e9..3bcb1698f148 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -146,7 +146,6 @@ from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler from .scheduling_euler_discrete import EulerDiscreteScheduler from .scheduling_flow_map_euler_discrete import FlowMapEulerDiscreteScheduler - from .scheduling_piflow import PiflowScheduler from .scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler from .scheduling_flow_match_heun_discrete import FlowMatchHeunDiscreteScheduler from .scheduling_flow_match_lcm import FlowMatchLCMScheduler @@ -159,6 +158,7 @@ from .scheduling_lcm import LCMScheduler from .scheduling_ltx_euler_ancestral_rf import LTXEulerAncestralRFScheduler from .scheduling_minimax_h3 import MiniMaxH3Scheduler + from .scheduling_piflow import PiflowScheduler from .scheduling_pndm import PNDMScheduler from .scheduling_repaint import RePaintScheduler from .scheduling_sasolver import SASolverScheduler diff --git a/src/diffusers/schedulers/scheduling_piflow.py b/src/diffusers/schedulers/scheduling_piflow.py index 4e0ac7ad90f6..181ffcc35e0a 100644 --- a/src/diffusers/schedulers/scheduling_piflow.py +++ b/src/diffusers/schedulers/scheduling_piflow.py @@ -141,10 +141,9 @@ def policy_rollout_fm( class PiflowScheduler(SchedulerMixin, ConfigMixin): """Few-step PiFlow scheduler for widened-output diffusion transformers. - PiFlow evaluates the denoising model at a small number of grid points and integrates a network-free policy - between those evaluations. The scheduler is intended for distilled Kandinsky 6 checkpoints, including the main - video/audio model and the video super-resolution model. Their model output contains `n_grid` predictions per - sample channel. + PiFlow evaluates the denoising model at a small number of grid points and integrates a network-free policy between + those evaluations. The scheduler is intended for distilled Kandinsky 6 checkpoints, including the main video/audio + model and the video super-resolution model. Their model output contains `n_grid` predictions per sample channel. This scheduler inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic methods implemented for all schedulers (loading, saving, etc.). @@ -155,7 +154,6 @@ class PiflowScheduler(SchedulerMixin, ConfigMixin): shift (`float`, *optional*, defaults to 5.0): Flow-matching timestep shift. n_grid (`int`, *optional*, defaults to 10): Number of predictions in the widened model output. - nfe (`int`, *optional*): Number of model evaluations used at inference. eps (`float`, *optional*, defaults to 1e-6): Minimum timestep and policy denominator. final_step_size_scale (`float`, *optional*, defaults to 0.5): Relative size of the final raw-timestep segment. @@ -165,7 +163,6 @@ class PiflowScheduler(SchedulerMixin, ConfigMixin): _compatibles = [] order = 1 - is_piflow = True @register_to_config def __init__( @@ -173,7 +170,6 @@ def __init__( num_train_timesteps: int = 1000, shift: float = 5.0, n_grid: int = 10, - nfe: int | None = None, eps: float = 1e-6, final_step_size_scale: float = 0.5, num_policy_substeps: int = 128, diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 4c5400229ef3..33219521ba98 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1575,52 +1575,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class MMAudioVAE(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 Kandinsky6SRLatentUpscalerBank(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 Kandinsky6SRPipeline(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 Kandinsky6SRPipelineOutput(metaclass=DummyObject): +class Kandinsky5Transformer3DModel(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -1665,36 +1620,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class Kandinsky6TI2VAPipeline(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 Kandinsky6TI2VAPipelineOutput(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 Kandinsky6Transformer3DModel(metaclass=DummyObject): _backends = ["torch"] @@ -1710,21 +1635,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class Kandinsky5Transformer3DModel(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"] @@ -1950,6 +1860,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class MMAudioVAE(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 MochiTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] @@ -3557,7 +3482,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class PiflowScheduler(metaclass=DummyObject): +class FlowMatchEulerDiscreteScheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -3572,7 +3497,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class FlowMatchEulerDiscreteScheduler(metaclass=DummyObject): +class FlowMatchHeunDiscreteScheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -3587,7 +3512,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class FlowMatchHeunDiscreteScheduler(metaclass=DummyObject): +class FlowMatchLCMScheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -3602,7 +3527,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class FlowMatchLCMScheduler(metaclass=DummyObject): +class HeliosDMDScheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -3617,7 +3542,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class HeliosDMDScheduler(metaclass=DummyObject): +class HeliosScheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -3632,7 +3557,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class HeliosScheduler(metaclass=DummyObject): +class HeunDiscreteScheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -3647,7 +3572,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class HeunDiscreteScheduler(metaclass=DummyObject): +class IPNDMScheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -3662,7 +3587,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class IPNDMScheduler(metaclass=DummyObject): +class KarrasVeScheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -3677,7 +3602,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class KarrasVeScheduler(metaclass=DummyObject): +class KDPM2AncestralDiscreteScheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -3692,7 +3617,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class KDPM2AncestralDiscreteScheduler(metaclass=DummyObject): +class KDPM2DiscreteScheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -3707,7 +3632,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class KDPM2DiscreteScheduler(metaclass=DummyObject): +class LCMScheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -3722,7 +3647,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class LCMScheduler(metaclass=DummyObject): +class LTXEulerAncestralRFScheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -3737,7 +3662,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class LTXEulerAncestralRFScheduler(metaclass=DummyObject): +class MiniMaxH3Scheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -3752,7 +3677,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class MiniMaxH3Scheduler(metaclass=DummyObject): +class PiflowScheduler(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 376596d632ea..a1f889a6dae6 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -2672,6 +2672,66 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class Kandinsky6SRPipeline(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 Kandinsky6SRPipelineOutput(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 Kandinsky6TI2VAPipeline(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 Kandinsky6TI2VAPipelineOutput(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 KandinskyCombinedPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/autoencoders/test_models_autoencoder_kandinsky6_sr.py b/tests/models/autoencoders/test_models_autoencoder_kandinsky6_sr.py new file mode 100644 index 000000000000..f7a27ce31e78 --- /dev/null +++ b/tests/models/autoencoders/test_models_autoencoder_kandinsky6_sr.py @@ -0,0 +1,105 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from diffusers import Kandinsky6SRVAE +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device +from ..testing_utils import ( + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, +) + + +enable_full_determinism() + + +class Kandinsky6SRVAETesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return Kandinsky6SRVAE + + @property + def pretrained_model_name_or_path(self): + return "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "vae"} + + @property + def main_input_name(self) -> str: + return "sample" + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict: + return { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 4, + "encoder_block_out_channels": (4, 8, 8), + "decoder_block_out_channels": (4, 8, 8), + "layers_per_block": 1, + "temporal_compression_ratio": 4, + "temporal_compression_start_level": 0, + } + + def get_dummy_inputs(self) -> dict: + return {"sample": randn_tensor((2, 3, 9, 16, 16), generator=self.generator, device=torch_device)} + + @property + def input_shape(self) -> tuple[int, ...]: + return (3, 9, 16, 16) + + @property + def output_shape(self) -> tuple[int, ...]: + return (3, 9, 16, 16) + + +class TestKandinsky6SRVAEModel(Kandinsky6SRVAETesterConfig, ModelTesterMixin): + def test_compression_ratios(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + with torch.no_grad(): + latents = model.encode(self.get_dummy_inputs()["sample"]).latent_dist.mode() + assert model.spatial_compression_ratio == 4 + assert model.temporal_compression_ratio == 4 + assert latents.shape == (2, 4, 3, 4, 4) + + def test_segmented_processing_matches_single_pass(self): + # 41 frames span three causal segments; the carried padding must reproduce a single-segment result. + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + sample = randn_tensor((1, 3, 41, 16, 16), generator=self.generator, device=torch_device) + with torch.no_grad(): + latents = model.encode(sample).latent_dist.mode() + latents_prefix = model.encode(sample[:, :, :17]).latent_dist.mode() + decoded = model.decode(latents).sample + decoded_prefix = model.decode(latents[:, :, :5]).sample + assert_tensors_close(latents_prefix, latents[:, :, :5], atol=1e-5, rtol=0) + assert_tensors_close(decoded_prefix, decoded[:, :, :17], atol=1e-4, rtol=0) + + +class TestKandinsky6SRVAEMemory(Kandinsky6SRVAETesterConfig, MemoryTesterMixin): + pass + + +class TestKandinsky6SRVAETorchCompile(Kandinsky6SRVAETesterConfig, TorchCompileTesterMixin): + pass diff --git a/tests/models/autoencoders/test_models_autoencoder_mmaudio.py b/tests/models/autoencoders/test_models_autoencoder_mmaudio.py new file mode 100644 index 000000000000..4747cb2196b2 --- /dev/null +++ b/tests/models/autoencoders/test_models_autoencoder_mmaudio.py @@ -0,0 +1,99 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from diffusers import MMAudioVAE +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, +) + + +enable_full_determinism() + + +class MMAudioVAETesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return MMAudioVAE + + @property + def pretrained_model_name_or_path(self): + return "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "audio_vae"} + + @property + def main_input_name(self) -> str: + return "sample" + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict: + return { + "mel_bins": 8, + "latent_channels": 4, + "hidden_channels": 8, + "channel_multipliers": (1, 2), + "layers_per_block": 1, + "sample_rate": 64, + "n_fft": 16, + "hop_length": 4, + "vocoder_upsample_initial_channel": 8, + "vocoder_upsample_rates": (2, 2), + "vocoder_upsample_kernel_sizes": (4, 4), + "vocoder_resblock_kernel_sizes": (3,), + "vocoder_resblock_dilation_sizes": ((1, 3),), + } + + def get_dummy_inputs(self) -> dict: + waveform = randn_tensor((2, 256), generator=self.generator, device=torch_device).clamp(-1, 1) + return {"sample": waveform} + + @property + def input_shape(self) -> tuple[int, ...]: + return (256,) + + @property + def output_shape(self) -> tuple[int, ...]: + return (1, 256) + + +class TestMMAudioVAEModel(MMAudioVAETesterConfig, ModelTesterMixin): + def test_latent_shape(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + with torch.no_grad(): + latents = model.encode(self.get_dummy_inputs()["sample"]).latent_dist.mode() + # 256 samples -> 64 mel frames (hop 4) -> 32 latent frames (one 2x downsample) + assert model.latent_hop_length == 8 + assert latents.shape == (2, 4, 32) + + +class TestMMAudioVAEMemory(MMAudioVAETesterConfig, MemoryTesterMixin): + pass + + +class TestMMAudioVAETorchCompile(MMAudioVAETesterConfig, TorchCompileTesterMixin): + pass diff --git a/tests/models/transformers/test_models_transformer_kandinsky6.py b/tests/models/transformers/test_models_transformer_kandinsky6.py index 93d1a150531c..88be69c550b7 100644 --- a/tests/models/transformers/test_models_transformer_kandinsky6.py +++ b/tests/models/transformers/test_models_transformer_kandinsky6.py @@ -17,12 +17,12 @@ import torch from diffusers import Kandinsky6Transformer3DModel -from diffusers.models.transformers.transformer_kandinsky6 import Kandinsky6RoPE1D, Kandinsky6RoPE3D from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import ( AttentionTesterMixin, + BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TorchCompileTesterMixin, @@ -33,14 +33,14 @@ enable_full_determinism() -class Kandinsky6TransformerTesterConfig: +class Kandinsky6TransformerTesterConfig(BaseModelTesterConfig): @property def model_class(self): return Kandinsky6Transformer3DModel @property def pretrained_model_name_or_path(self): - return "" # TODO: Set Hub repository ID + return "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers" @property def pretrained_model_kwargs(self): @@ -48,13 +48,13 @@ def pretrained_model_kwargs(self): @property def main_input_name(self) -> str: - return "x_video" + return "hidden_states" @property def generator(self): return torch.Generator("cpu").manual_seed(0) - def get_init_dict(self) -> dict[str, int | list[int]]: + def get_init_dict(self) -> dict: return { "in_visual_dim": 4, "out_visual_dim": 4, @@ -70,87 +70,70 @@ def get_init_dict(self) -> dict[str, int | list[int]]: "visual_cond": True, "in_audio_dim": 4, "out_audio_dim": 4, + "visual_token_type_num_embeddings": 2, } - def _build_dummy_inputs( - self, batch_size: int, num_frames: int, height: int, width: int, device - ) -> dict[str, torch.Tensor]: + def _build_dummy_inputs(self, batch_size: int, num_frames: int, height: int, width: int) -> dict: init_dict = self.get_init_dict() - head_dim = sum(init_dict["axes_dims"]) - patch_t, patch_h, patch_w = init_dict["patch_size"] - vis_in_dim = 2 * init_dict["in_visual_dim"] + 1 # visual_cond=True prepends cond + mask channels - text_length = 6 - - x_video = randn_tensor( - (batch_size, num_frames, height, width, vis_in_dim), generator=self.generator, device=device - ) - text_embed = randn_tensor( - (batch_size, text_length, init_dict["in_text_dim"]), generator=self.generator, device=device - ) - pooled_text_embed = randn_tensor( - (batch_size, init_dict["in_text_dim2"]), generator=self.generator, device=device - ) - # `self.generator` is always a CPU generator (see below); draw on CPU and move, like `randn_tensor` does - # internally, so this also works with an accelerator (e.g. MPS) `device`. - time = torch.randint(0, 1000, (batch_size,), generator=self.generator).float().to(device) - - visual_shape = (num_frames // patch_t, height // patch_h, width // patch_w) - visual_rope = Kandinsky6RoPE3D(init_dict["axes_dims"]).to(device)( - visual_shape, - [torch.arange(size, device=device) for size in visual_shape], - ) - video_text_rope = Kandinsky6RoPE1D(head_dim).to(device)(torch.arange(text_length, device=device)) - audio_text_rope = Kandinsky6RoPE1D(head_dim).to(device)(torch.arange(text_length, device=device)) - + # `visual_cond=True` appends conditioning latents and a mask to the input channels + num_input_channels = 2 * init_dict["in_visual_dim"] + 1 return { - "x_video": x_video, - "text_embed": text_embed, - "pooled_text_embed": pooled_text_embed, - "time": time, - "visual_rope": visual_rope, - "video_text_rope": video_text_rope, - "audio_text_rope": audio_text_rope, + "hidden_states": randn_tensor( + (batch_size, num_frames, height, width, num_input_channels), + generator=self.generator, + device=torch_device, + ), + "audio_hidden_states": randn_tensor( + (batch_size, 5, init_dict["in_audio_dim"]), generator=self.generator, device=torch_device + ), + "encoder_hidden_states": randn_tensor( + (batch_size, 6, init_dict["in_text_dim"]), generator=self.generator, device=torch_device + ), + "pooled_projections": randn_tensor( + (batch_size, init_dict["in_text_dim2"]), generator=self.generator, device=torch_device + ), + "timestep": torch.randint(0, 1000, (batch_size,), generator=self.generator).float().to(torch_device), } - def get_dummy_inputs(self, batch_size: int = 1, device=torch_device) -> dict[str, torch.Tensor]: - # Only `x_video` is passed (no `x_audio`): the fused block no-ops its audio branch when `aud is None`, - # so this still runs the real dual-stream block weights, while `forward` returns a single tensor - # (matching what the generic mixins below expect) instead of a `(video, audio)` tuple. - return self._build_dummy_inputs(batch_size, num_frames=2, height=4, width=4, device=device) + def get_dummy_inputs(self) -> dict: + return self._build_dummy_inputs(batch_size=1, num_frames=2, height=4, width=4) @property def input_shape(self) -> tuple[int, ...]: - # (num_frames, height, width, vis_in_dim), batch dimension excluded. return (2, 4, 4, 2 * 4 + 1) @property def output_shape(self) -> tuple[int, ...]: - # (num_frames, height, width, out_visual_dim), batch dimension excluded. return (2, 4, 4, 4) class TestKandinsky6TransformerModel(Kandinsky6TransformerTesterConfig, ModelTesterMixin): - pass + def test_video_only_forward(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + inputs.pop("audio_hidden_states") + with torch.no_grad(): + output = model(**inputs) + assert output.sample.shape == (1, *self.output_shape) + assert output.audio_sample is None + + def test_tail_conditioning_inputs(self): + # An appended reference frame reuses temporal rotary position 0 and carries token type 1. + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self._build_dummy_inputs(batch_size=1, num_frames=3, height=4, width=4) + inputs["visual_rope_pos"] = ( + torch.tensor([0, 1, 0], device=torch_device), + torch.arange(2, device=torch_device), + torch.arange(2, device=torch_device), + ) + inputs["visual_token_type_ids"] = torch.tensor([[0, 0, 1]], device=torch_device) + with torch.no_grad(): + output = model(**inputs) + assert output.sample.shape == (1, 3, 4, 4, 4) class TestKandinsky6TransformerMemory(Kandinsky6TransformerTesterConfig, MemoryTesterMixin): - # `Kandinsky6TimeEmbeddings`/`Kandinsky6Modulation` bypass leaf-level offload hooks by design (see the - # `_supports_group_offloading = False` comment on the model); the classic accelerate cpu/disk offload - # hooks hit the exact same gap and have no model-level opt-out flag, so skip explicitly here instead. - _OFFLOAD_SKIP_REASON = ( - "Kandinsky6TimeEmbeddings/Kandinsky6Modulation read their nn.Linear weight/bias directly for a " - "fp32-upcast functional.linear call rather than calling the submodule, so leaf-level offload hooks " - "never see (or move) those weights." - ) - - def test_cpu_offload(self, *args, **kwargs): - pytest.skip(self._OFFLOAD_SKIP_REASON) - - def test_disk_offload_without_safetensors(self, *args, **kwargs): - pytest.skip(self._OFFLOAD_SKIP_REASON) - - def test_disk_offload_with_safetensors(self, *args, **kwargs): - pytest.skip(self._OFFLOAD_SKIP_REASON) + pass class TestKandinsky6TransformerTorchCompile(Kandinsky6TransformerTesterConfig, TorchCompileTesterMixin): @@ -158,8 +141,8 @@ class TestKandinsky6TransformerTorchCompile(Kandinsky6TransformerTesterConfig, T def different_shapes_for_compilation(self): return [(4, 4), (4, 8), (8, 8)] - def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]: - return self._build_dummy_inputs(batch_size=1, num_frames=2, height=height, width=width, device=torch_device) + def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict: + return self._build_dummy_inputs(batch_size=1, num_frames=2, height=height, width=width) class TestKandinsky6TransformerTraining(Kandinsky6TransformerTesterConfig, TrainingTesterMixin): @@ -170,6 +153,6 @@ class TestKandinsky6TransformerAttention(Kandinsky6TransformerTesterConfig, Atte def test_fuse_unfuse_qkv_projections(self, *args, **kwargs): pytest.skip( "Kandinsky6Attention names its projections to_query/to_key/to_value/out_layer (matching the " - "native K6 checkpoint layout) rather than to_q/to_k/to_v/to_out, which " + "Kandinsky 5 layout) rather than to_q/to_k/to_v/to_out, which " "AttentionModuleMixin.fuse_projections hardcodes." ) diff --git a/tests/models/transformers/test_models_transformer_kandinsky6_sr.py b/tests/models/transformers/test_models_transformer_kandinsky6_sr.py new file mode 100644 index 000000000000..1cbe29a6a52f --- /dev/null +++ b/tests/models/transformers/test_models_transformer_kandinsky6_sr.py @@ -0,0 +1,122 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers import Kandinsky6SRTransformer3DModel +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class Kandinsky6SRTransformerTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return Kandinsky6SRTransformer3DModel + + @property + def pretrained_model_name_or_path(self): + return "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} + + @property + def main_input_name(self) -> str: + return "hidden_states" + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict: + return { + "in_visual_dim": 4, + "out_visual_dim": 8, + "time_dim": 16, + "patch_size": (1, 2, 2), + "model_dim": 24, + "ff_dim": 32, + "num_visual_blocks": 2, + "axes_dims": (4, 4, 4), + # NABLA needs 8x8 token blocks, far larger than these tiny grids; run dense attention. + "nabla_threshold": None, + } + + def _build_dummy_inputs(self, batch_size: int, num_frames: int, height: int, width: int) -> dict: + # The input concatenates the noisy latent, the anchor latent and the anchor mask + num_input_channels = 2 * self.get_init_dict()["in_visual_dim"] + 1 + return { + "hidden_states": randn_tensor( + (batch_size, num_frames, height, width, num_input_channels), + generator=self.generator, + device=torch_device, + ), + "timestep": torch.randint(0, 1000, (batch_size,), generator=self.generator).float().to(torch_device), + } + + def get_dummy_inputs(self) -> dict: + return self._build_dummy_inputs(batch_size=1, num_frames=2, height=4, width=4) + + @property + def input_shape(self) -> tuple[int, ...]: + return (2, 4, 4, 2 * 4 + 1) + + @property + def output_shape(self) -> tuple[int, ...]: + return (2, 4, 4, 8) + + +class TestKandinsky6SRTransformerModel(Kandinsky6SRTransformerTesterConfig, ModelTesterMixin): + pass + + +class TestKandinsky6SRTransformerMemory(Kandinsky6SRTransformerTesterConfig, MemoryTesterMixin): + pass + + +class TestKandinsky6SRTransformerTorchCompile(Kandinsky6SRTransformerTesterConfig, TorchCompileTesterMixin): + @property + def different_shapes_for_compilation(self): + return [(4, 4), (4, 8), (8, 8)] + + def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict: + return self._build_dummy_inputs(batch_size=1, num_frames=2, height=height, width=width) + + +class TestKandinsky6SRTransformerTraining(Kandinsky6SRTransformerTesterConfig, TrainingTesterMixin): + pass + + +class TestKandinsky6SRTransformerAttention(Kandinsky6SRTransformerTesterConfig, AttentionTesterMixin): + def test_fuse_unfuse_qkv_projections(self, *args, **kwargs): + pytest.skip( + "Kandinsky6SRAttention names its projections to_query/to_key/to_value/out_layer (matching the " + "Kandinsky 5 layout) rather than to_q/to_k/to_v/to_out, which " + "AttentionModuleMixin.fuse_projections hardcodes." + ) diff --git a/tests/pipelines/kandinsky6/test_models_latent_upscaler.py b/tests/pipelines/kandinsky6/test_models_latent_upscaler.py new file mode 100644 index 000000000000..2ba43be3a69e --- /dev/null +++ b/tests/pipelines/kandinsky6/test_models_latent_upscaler.py @@ -0,0 +1,96 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from diffusers.pipelines.kandinsky6 import Kandinsky6SRLatentUpscalerBank +from diffusers.utils.torch_utils import randn_tensor + +from ...models.testing_utils import ( + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, +) +from ...testing_utils import enable_full_determinism, torch_device + + +enable_full_determinism() + + +class Kandinsky6SRLatentUpscalerBankTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return Kandinsky6SRLatentUpscalerBank + + @property + def pretrained_model_name_or_path(self): + return "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "latent_upscaler"} + + @property + def main_input_name(self) -> str: + return "latents" + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict: + return { + "in_channels": 4, + "stage_channels": (8, 8, 4), + "num_pre_blocks": 1, + "num_mid_blocks": 1, + "num_post_blocks": 1, + "num_x2_adapter_blocks": 1, + "scales": (2, 4), + } + + def get_dummy_inputs(self) -> dict: + return { + "latents": randn_tensor((2, 4, 3, 4, 4), generator=self.generator, device=torch_device), + "scale": 2, + } + + @property + def input_shape(self) -> tuple[int, ...]: + return (4, 3, 4, 4) + + @property + def output_shape(self) -> tuple[int, ...]: + return (4, 3, 8, 8) + + +class TestKandinsky6SRLatentUpscalerBankModel(Kandinsky6SRLatentUpscalerBankTesterConfig, ModelTesterMixin): + def test_x4_scale(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + with torch.no_grad(): + output = model(inputs["latents"], scale=4).sample + assert output.shape == (2, 4, 3, 16, 16) + + +class TestKandinsky6SRLatentUpscalerBankMemory(Kandinsky6SRLatentUpscalerBankTesterConfig, MemoryTesterMixin): + pass + + +class TestKandinsky6SRLatentUpscalerBankTorchCompile( + Kandinsky6SRLatentUpscalerBankTesterConfig, TorchCompileTesterMixin +): + pass diff --git a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_sr.py b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_sr.py new file mode 100644 index 000000000000..bc8ab010c3bf --- /dev/null +++ b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_sr.py @@ -0,0 +1,133 @@ +# Copyright 2025 The Kandinsky Team and The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from diffusers import ( + FlowMatchEulerDiscreteScheduler, + Kandinsky6SRPipeline, + Kandinsky6SRTransformer3DModel, + Kandinsky6SRVAE, +) +from diffusers.pipelines.kandinsky6 import Kandinsky6SRLatentUpscalerBank + +from ...testing_utils import torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin + + +class Kandinsky6SRPipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = Kandinsky6SRPipeline + required_input_params_in_call_signature = frozenset(["video", "resolution_scale", "num_inference_steps"]) + batch_input_params = frozenset(["video"]) + optional_input_params = frozenset(["num_inference_steps", "generator", "output_type", "return_dict"]) + # (num_frames, channels, height, width) for output_type="pt": the 5x16x32 input video upscaled 2x. + output_shape = (5, 3, 32, 64) + + def get_dummy_components(self): + torch.manual_seed(0) + transformer = Kandinsky6SRTransformer3DModel( + in_visual_dim=4, + out_visual_dim=4, + time_dim=16, + patch_size=(1, 1, 1), + model_dim=24, + ff_dim=32, + num_visual_blocks=2, + axes_dims=(4, 4, 4), + # NABLA needs 8x8 token blocks; the tiny tiles run dense attention. + nabla_threshold=None, + # Tiles the size of the 16x32 test video (and its transposed / square variants); every tile size must be + # divisible by the VAE spatial factor (4) times the largest scale (4). + tile_sizes=((16, 16), (16, 32), (32, 16)), + ) + + torch.manual_seed(0) + # Three levels give a 4x spatial factor; both non-final levels compress time for a 4x temporal factor. + vae = Kandinsky6SRVAE( + latent_channels=4, + encoder_block_out_channels=(4, 8, 8), + decoder_block_out_channels=(4, 8, 8), + layers_per_block=1, + temporal_compression_ratio=4, + temporal_compression_start_level=0, + ) + + torch.manual_seed(0) + latent_upscaler = Kandinsky6SRLatentUpscalerBank( + in_channels=4, + stage_channels=(8, 8, 4), + num_pre_blocks=1, + num_mid_blocks=1, + num_post_blocks=1, + num_x2_adapter_blocks=1, + scales=(2, 4), + ) + + scheduler = FlowMatchEulerDiscreteScheduler(shift=3.5) + + return { + "transformer": transformer, + "vae": vae, + "scheduler": scheduler, + "latent_upscaler": latent_upscaler, + } + + def get_dummy_inputs(self): + # A `(num_frames, channels, height, width)` video in [0, 1] with 1 + 4k frames. + video = torch.rand((5, 3, 16, 32), generator=self.get_generator(0)) + return { + "video": video, + "resolution_scale": 2, + "num_inference_steps": 2, + "generator": self.get_generator(0), + "min_overlap": 0.2, + "tiles_batch_size": 4, + "output_type": "pt", + } + + +class TestKandinsky6SRPipeline(Kandinsky6SRPipelineTesterConfig, PipelineTesterMixin): + def test_kandinsky6_sr_scales(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + # The 2.25x route rounds the 1.125x pre-upscaled size to the VAE spatial factor (4 here): 16 -> 16, 32 -> 36. + for resolution_scale, expected in ((2, (32, 64)), (2.25, (32, 72)), (4, (64, 128))): + inputs = self.get_dummy_inputs() + inputs["resolution_scale"] = resolution_scale + output = pipe(**inputs) + assert output.frames.shape == (1, 5, 3, *expected), f"unexpected shape for scale {resolution_scale}" + assert not torch.isnan(output.frames).any() + + def test_kandinsky6_sr_without_latent_upscaler(self): + # Without the latent upscaler the pixel tiles are bilinearly upscaled and encoded instead. + components = self.get_dummy_components() + components["latent_upscaler"] = None + pipe = self.pipeline_class(**components).to(torch_device) + output = pipe(**self.get_dummy_inputs()) + assert output.frames.shape == (1, *self.output_shape) + assert not torch.isnan(output.frames).any() + + def test_kandinsky6_sr_rejects_wrong_frame_count(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs() + inputs["video"] = inputs["video"][:4] + try: + pipe(**inputs) + except ValueError as error: + assert "frames" in str(error) + else: + raise AssertionError("expected a ValueError for a video without 1 + 4k frames") + + +class TestKandinsky6SRPipelineMemory(Kandinsky6SRPipelineTesterConfig, MemoryTesterMixin): + pass diff --git a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py index a571bbc11254..834c8fe699c0 100644 --- a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py +++ b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py @@ -14,7 +14,6 @@ import numpy as np import PIL.Image -import pytest import torch from transformers import ( AutoProcessor, @@ -30,6 +29,7 @@ FlowMatchEulerDiscreteScheduler, Kandinsky6TI2VAPipeline, Kandinsky6Transformer3DModel, + MMAudioVAE, ) from ...testing_utils import torch_device @@ -42,15 +42,17 @@ class Kandinsky6TI2VAPipelineTesterConfig(BasePipelineTesterConfig): ["prompt", "height", "width", "num_frames", "num_inference_steps", "guidance_scale"] ) batch_input_params = frozenset(["prompt", "negative_prompt"]) - # (channels, num_frames, height, width) for output_type="pt", matching `get_dummy_inputs()`'s - # (num_frames=5, height=16, width=16) at the tiny VAE's real 8x spatial / 4x temporal compression. - output_shape = (3, 5, 16, 16) + optional_input_params = frozenset( + ["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"] + ) + # (num_frames, channels, height, width) for output_type="pt", matching `get_dummy_inputs()`'s + # (num_frames=5, height=16, width=16) at the tiny VAE's 8x spatial / 4x temporal compression. + output_shape = (5, 3, 16, 16) def get_dummy_components(self): torch.manual_seed(0) - # 3 down/up levels so the tiny VAE's *realized* spatial/temporal compression actually matches the - # declared spatial_compression_ratio=8 / temporal_compression_ratio=4 (2 levels only realizes 4x/2x) — - # the pipeline itself hardcodes `// 8` / `// 4` when sizing latents rather than reading these configs. + # 3 down/up levels so the tiny VAE's realized compression matches the declared + # spatial_compression_ratio=8 / temporal_compression_ratio=4 the pipeline reads. vae = AutoencoderKLHunyuanVideo( act_fn="silu", block_out_channels=[8, 8, 8], @@ -75,11 +77,28 @@ def get_dummy_components(self): ], ) + torch.manual_seed(0) + # A tiny sample rate keeps the audio latent sequence short (2 latent frames for 5 video frames). + audio_vae = MMAudioVAE( + mel_bins=8, + latent_channels=4, + hidden_channels=8, + channel_multipliers=(1, 2), + layers_per_block=1, + sample_rate=64, + n_fft=16, + hop_length=4, + vocoder_upsample_initial_channel=8, + vocoder_upsample_rates=(2, 2), + vocoder_upsample_kernel_sizes=(4, 4), + vocoder_resblock_kernel_sizes=(3,), + vocoder_resblock_dilation_sizes=((1, 3),), + ) + scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0) - # mrope_section must sum to (hidden_size / num_attention_heads) / 2 — matches the known-good - # hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration-compatible config used by Kandinsky5's - # own test suite (tests/pipelines/kandinsky5/test_kandinsky5.py). + # mrope_section must sum to (hidden_size / num_attention_heads) / 2, matching the + # hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration processor used below. qwen_hidden_size = 32 torch.manual_seed(0) qwen_config = Qwen2_5_VLConfig( @@ -152,11 +171,11 @@ def get_dummy_components(self): "transformer": transformer, "vae": vae, "text_encoder": text_encoder, - "audio_vae": None, - "scheduler": scheduler, "tokenizer": tokenizer, "text_encoder_2": text_encoder_2, "tokenizer_2": tokenizer_2, + "scheduler": scheduler, + "audio_vae": audio_vae, } def get_dummy_inputs(self): @@ -170,12 +189,29 @@ def get_dummy_inputs(self): "width": 16, "num_frames": 5, "max_sequence_length": 16, - "sample_audio": False, "output_type": "pt", } class TestKandinsky6TI2VAPipeline(Kandinsky6TI2VAPipelineTesterConfig, PipelineTesterMixin): + def test_kandinsky6_ti2va_audio_output(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + output = pipe(**self.get_dummy_inputs()) + + # 5 frames at 24 fps and 64 Hz -> 2 audio latent frames -> 4 mel frames -> 16 samples + assert output.frames.shape == (1, *self.output_shape) + assert output.audio.shape == (1, 16) + assert not torch.isnan(output.audio).any() + + def test_kandinsky6_ti2va_video_only(self): + pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs() + inputs["sample_audio"] = False + output = pipe(**inputs) + + assert output.frames.shape == (1, *self.output_shape) + assert output.audio is None + def test_kandinsky6_ti2va_i2va(self): pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) inputs = self.get_dummy_inputs() @@ -187,15 +223,9 @@ def test_kandinsky6_ti2va_i2va(self): assert not torch.isnan(output.frames.float()).any() def test_kandinsky6_ti2va_different_images(self): - # A freshly initialized (untrained) transformer zero-initializes every AdaLN modulation gate - # (`Kandinsky6Modulation.out_layer`), so at init every attention/feed-forward residual is gated to - # exactly zero: neither the text branch nor (with the default `tail_cond_first_frame` scheme) a - # reference image reaches the *kept* output frames yet, since that path is also self-attention - # (equally gated). `visual_cond_scheme="i2v"` is different: it overwrites frame 0's latent with the - # image encoding directly (`_apply_visual_conditioning`), and that frame is kept (not appended and - # dropped like `tail_cond_first_frame`), so it passes through two real (non-zero-init) linears — - # `visual_embeddings` and `out_layer` — around the identity (zero-gated) blocks. That's the one - # conditioning signal guaranteed to already produce different output on a fresh model. + # A freshly initialized transformer zero-initializes every AdaLN gate, so the text and reference-image paths + # through attention are identity at init. The `i2v` scheme instead writes the encoded image into frame 0, + # which passes through the real `visual_embeddings` and `out_layer` projections and must change the output. pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) inputs = self.get_dummy_inputs() inputs["visual_cond_scheme"] = "i2v" @@ -211,30 +241,16 @@ def test_kandinsky6_ti2va_different_images(self): max_diff = (output_black_image.float() - output_white_image.float()).abs().max() assert max_diff > 1e-6, "Outputs should be different for different reference images." - def test_kandinsky6_ti2va_num_images_per_prompt(self): + def test_kandinsky6_ti2va_num_videos_per_prompt(self): pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) inputs = self.get_dummy_inputs() - inputs["num_images_per_prompt"] = 2 + inputs["num_videos_per_prompt"] = 2 output = pipe(**inputs) assert output.frames.shape == (2, *self.output_shape) + assert output.audio.shape[0] == 2 class TestKandinsky6TI2VAPipelineMemory(Kandinsky6TI2VAPipelineTesterConfig, MemoryTesterMixin): - # `Kandinsky6TimeEmbeddings`/`Kandinsky6Modulation` (in the transformer) read their `nn.Linear` - # weight/bias directly for a fp32-upcast `functional.linear` call rather than calling the submodule - # (see the `_supports_group_offloading = False` comment on `Kandinsky6Transformer3DModel`), so - # leaf-level offload hooks never see those weights. `enable_sequential_cpu_offload` hits the same gap - # and has no pipeline-level opt-out flag; confirmed this isn't an environment issue — Flux's identical - # tests pass on this same machine/backend. - _OFFLOAD_SKIP_REASON = ( - "Kandinsky6TimeEmbeddings/Kandinsky6Modulation bypass leaf-level offload hooks by design; see " - "Kandinsky6Transformer3DModel._supports_group_offloading." - ) - - def test_sequential_cpu_offload_forward_pass(self, *args, **kwargs): - pytest.skip(self._OFFLOAD_SKIP_REASON) - - def test_sequential_offload_forward_pass_twice(self, *args, **kwargs): - pytest.skip(self._OFFLOAD_SKIP_REASON) + pass diff --git a/tests/schedulers/test_scheduler_piflow.py b/tests/schedulers/test_scheduler_piflow.py index 6a21ae8b5594..53b7d5d2c043 100644 --- a/tests/schedulers/test_scheduler_piflow.py +++ b/tests/schedulers/test_scheduler_piflow.py @@ -52,7 +52,6 @@ def test_instantiation_with_defaults(self): scheduler = self.scheduler_class(**self.get_default_config()) self.assertEqual(scheduler.config.num_train_timesteps, 1000) self.assertEqual(scheduler.config.n_grid, 4) - self.assertTrue(scheduler.is_piflow) def test_invalid_n_grid_raises(self): with self.assertRaises(ValueError): From 65b00fd22ce4c36e72e5537916b4aa8d3fc9144d Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 28 Sep 2026 17:55:05 +0300 Subject: [PATCH 20/62] untrack notebooks directory Co-Authored-By: Claude Sonnet 5 --- .gitignore | 2 +- notebooks/kandinsky6_sr.ipynb | 96 ----------- notebooks/kandinsky6_sr.py | 105 ------------ notebooks/kandinsky6_ti2va.ipynb | 240 ---------------------------- notebooks/kandinsky6_ti2va.py | 182 --------------------- notebooks/kandinsky6_ti2va_sr.ipynb | 130 --------------- notebooks/kandinsky6_ti2va_sr.py | 113 ------------- 7 files changed, 1 insertion(+), 867 deletions(-) delete mode 100644 notebooks/kandinsky6_sr.ipynb delete mode 100644 notebooks/kandinsky6_sr.py delete mode 100644 notebooks/kandinsky6_ti2va.ipynb delete mode 100644 notebooks/kandinsky6_ti2va.py delete mode 100644 notebooks/kandinsky6_ti2va_sr.ipynb delete mode 100644 notebooks/kandinsky6_ti2va_sr.py diff --git a/.gitignore b/.gitignore index 7b156e460abf..44f8e92185bd 100644 --- a/.gitignore +++ b/.gitignore @@ -185,4 +185,4 @@ wandb # AI agent generated symlinks /.agents/skills -/.claude/skills \ No newline at end of file +/.claude/skillsnotebooks/ diff --git a/notebooks/kandinsky6_sr.ipynb b/notebooks/kandinsky6_sr.ipynb deleted file mode 100644 index b5f51cc2fe13..000000000000 --- a/notebooks/kandinsky6_sr.ipynb +++ /dev/null @@ -1,96 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "from pathlib import Path\n", - "\n", - "import numpy as np\n", - "import torch\n", - "from diffusers import Kandinsky6SRPipeline\n", - "from diffusers.utils import encode_video, load_video\n", - "\n", - "REPO = Path.cwd()\n", - "while REPO != REPO.parent and not (REPO / \"dev\" / \"exports\" / \"diffusers\").exists():\n", - " REPO = REPO.parent\n", - "if not (REPO / \"dev\" / \"exports\" / \"diffusers\").exists():\n", - " raise RuntimeError(\"Run this notebook from the K6 repository or one of its subdirectories.\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "MODEL_REPO = \"kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers\"\n", - "GPU_ID = 0\n", - "\n", - "pipe = Kandinsky6SRPipeline.from_pretrained(MODEL_REPO, dtype=torch.bfloat16)\n", - "pipe.enable_model_cpu_offload(gpu_id=GPU_ID)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "BS = 2\n", - "INPUT_VIDEO = os.path.join(REPO, \"assets/vsr_input.mp4\")\n", - "OUTPUT_VIDEOS = [\n", - " os.path.join(REPO, f\"outputs/diffusers/k6_vsr/bs{BS}/vsr_x2.25_batchid{i}.mp4\")\n", - " for i in range(BS)\n", - "]\n", - "SEED = 1137\n", - "GENERATOR = torch.Generator(device=f\"cuda:{GPU_ID}\").manual_seed(SEED)\n", - "\n", - "# The pipeline only takes/returns tensors. Decode the video once, align it to\n", - "# 1 + 8k frames (the SR model's temporal chunk size), and repeat it BS times.\n", - "frames, fps = load_video(INPUT_VIDEO, return_fps=True)\n", - "video = torch.from_numpy(np.stack([np.array(frame) for frame in frames])).permute(3, 0, 1, 2)\n", - "aligned = 1 + 8 * ((video.shape[1] - 1) // 8)\n", - "video_batch = video[:, :aligned].unsqueeze(0).repeat(BS, 1, 1, 1, 1)\n", - "\n", - "# 2.25x means x1.125 pixel pre-upscale followed by the x2 KVAE path.\n", - "result = pipe(\n", - " video=video_batch,\n", - " resolution_scale=2.25,\n", - " output_type=\"pt\",\n", - " generator=GENERATOR,\n", - ")\n", - "print(f\"frames: {tuple(result.frames.shape)}\")\n", - "\n", - "for sr_frames, output_path in zip(result.frames, OUTPUT_VIDEOS, strict=True):\n", - " os.makedirs(os.path.dirname(output_path), exist_ok=True)\n", - " encode_video(sr_frames.permute(1, 2, 3, 0), fps=round(fps), output_path=output_path)\n", - " print(f\"saved: {output_path}\")" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "kandinsky-6-video-pro", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.13.9" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/kandinsky6_sr.py b/notebooks/kandinsky6_sr.py deleted file mode 100644 index 0c2e3f20f5a6..000000000000 --- a/notebooks/kandinsky6_sr.py +++ /dev/null @@ -1,105 +0,0 @@ -#!/usr/bin/env python3 -"""Kandinsky 6 video super-resolution. - -Script version of kandinsky6_sr.ipynb, for running on a headless server. - -The pipeline itself only accepts/returns tensors (no file I/O): this script -does the mp4 <-> tensor conversion with the shared `diffusers.utils` -video helpers, the same way a user integrating the pipeline into their own -code would. - -Usage: - python kandinsky6_sr.py --input-video assets/vsr_input.mp4 - python kandinsky6_sr.py --input-video a.mp4 b.mp4 --gpu-id 1 --resolution-scale 2.25 - -Authenticate with the Hub before running if the model repo is gated: - huggingface-cli login - or: - export HF_TOKEN=hf_... -""" - -from __future__ import annotations - -import argparse -from os import pipe -from pathlib import Path - -import numpy as np -import torch - -from diffusers import Kandinsky6SRPipeline -from diffusers.utils import encode_video, load_video - -DEFAULT_MODEL = "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers" - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--model", default=DEFAULT_MODEL) - parser.add_argument("--gpu-id", type=int, default=0) - parser.add_argument( - "--input-video", - type=Path, - nargs="+", - required=True, - help="One or more input video paths. A single path is replicated --batch-size times.", - ) - parser.add_argument("--batch-size", type=int, default=1, help="Only used when --input-video has one path.") - parser.add_argument("--resolution-scale", type=float, default=2.25) - parser.add_argument("--seed", type=int, default=1137) - parser.add_argument("--output-dir", type=Path, default=Path("outputs/diffusers/k6_vsr")) - return parser.parse_args() - - -def load_video_tensor(path: Path) -> tuple[torch.Tensor, float]: - """Decode ``path`` into a ``[3, T, H, W]`` uint8 tensor, aligned to ``1 + 8k`` frames.""" - frames, fps = load_video(str(path), return_fps=True) - video = torch.from_numpy(np.stack([np.array(frame) for frame in frames])).permute(3, 0, 1, 2) - aligned = 1 + 8 * ((video.shape[1] - 1) // 8) - if aligned == 0: - raise SystemExit(f"{path}: no readable frames") - return video[:, :aligned], fps - - -def main() -> None: - args = parse_args() - - input_videos = args.input_video - for path in input_videos: - if not path.exists(): - raise SystemExit(f"input video not found: {path}") - if len(input_videos) == 1 and args.batch_size > 1: - input_videos = input_videos * args.batch_size - batch_size = len(input_videos) - - loaded = [load_video_tensor(path) for path in input_videos] - videos, fps_values = [item[0] for item in loaded], [item[1] for item in loaded] - if any(video.shape != videos[0].shape for video in videos[1:]): - raise SystemExit("all --input-video clips must share resolution and frame count") - video_batch = torch.stack(videos) - - pipe = Kandinsky6SRPipeline.from_pretrained(args.model, dtype=torch.bfloat16) - pipe.enable_model_cpu_offload(gpu_id=args.gpu_id) - - output_dir = args.output_dir / f"bs{batch_size}" - output_dir.mkdir(parents=True, exist_ok=True) - output_paths = [output_dir / f"vsr_x{args.resolution_scale}_batchid{i}.mp4" for i in range(batch_size)] - - generator = torch.Generator(device=f"cuda:{args.gpu_id}").manual_seed(args.seed) - - # resolution_scale=2.25 means x1.125 pixel pre-upscale followed by the x2 KVAE path - # (the production route). - result = pipe( - video=video_batch, - resolution_scale=args.resolution_scale, - output_type="pt", - generator=generator, - ) - print(f"frames: {tuple(result.frames.shape)}") - for frames, output_path, fps in zip(result.frames, output_paths, fps_values, strict=True): - encode_video(frames.permute(1, 2, 3, 0), fps=round(fps), output_path=str(output_path)) - print(f"saved: {output_path}") - - -if __name__ == "__main__": - main() diff --git a/notebooks/kandinsky6_ti2va.ipynb b/notebooks/kandinsky6_ti2va.ipynb deleted file mode 100644 index 7ca2eca7fc7d..000000000000 --- a/notebooks/kandinsky6_ti2va.ipynb +++ /dev/null @@ -1,240 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "7dead615", - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "from pathlib import Path\n", - "\n", - "os.environ[\"HF_TOKEN\"] = \"TOKEN\"\n", - "\n", - "import torch\n", - "from diffusers import Kandinsky6TI2VAPipeline\n", - "from diffusers.hooks import MagCacheConfig\n", - "from diffusers.models.transformers.transformer_kandinsky6 import (\n", - " Kandinsky6FusedTransformerDecoderBlock,\n", - " Kandinsky6TransformerDecoderBlock,\n", - " Kandinsky6TransformerEncoderBlock,\n", - ")\n", - "from diffusers.utils import encode_video\n", - "\n", - "REPO = Path.cwd()\n", - "while REPO != REPO.parent and not (REPO / \"dev\" / \"exports\" / \"diffusers\").exists():\n", - " REPO = REPO.parent\n", - "if not (REPO / \"dev\" / \"exports\" / \"diffusers\").exists():\n", - " raise RuntimeError(\"Run this notebook from the K6 repository or one of its subdirectories.\")" - ] - }, - { - "cell_type": "markdown", - "id": "b6a61ec1", - "metadata": {}, - "source": [ - "### Initialize pipeline" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "821a098a", - "metadata": {}, - "outputs": [], - "source": [ - "# Keep generation and compile settings in one K6 YAML config. MagCache defaults are exported with the transformer.\n", - "GPU_ID = 2\n", - "\n", - "MODELS = {\n", - " \"k6pro\": \"kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers\",\n", - " \"k6pro_distill\": \"kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers\",\n", - "}\n", - "\n", - "model_name = \"k6pro_distill\"\n", - "model_path = MODELS[model_name]\n", - "expand_prompts = True\n", - "use_magcache = bool(\"distill\" not in model_name)\n", - "use_compile = False\n", - "max_sequence_length = 1024\n", - "sample_fps = 24.0\n", - "seed = 1137\n", - "\n", - "backend = \"native\" if use_compile else \"_flash_3\"\n", - "\n", - "pipe = Kandinsky6TI2VAPipeline.from_pretrained(model_path, dtype=torch.bfloat16)\n", - "pipe.transformer.set_attention_backend(backend)\n", - "\n", - "# Compile after loading the weights and before CPU offload. This applies the\n", - "# configured K6 block-level torch.compile strategy to the exported transformer.\n", - "if use_compile:\n", - " torch.set_float32_matmul_precision('high')\n", - " for block in pipe.transformer.modules():\n", - " if isinstance(block, Kandinsky6FusedTransformerDecoderBlock):\n", - " block.forward = torch.compile(block.forward)\n", - " elif isinstance(block, (Kandinsky6TransformerEncoderBlock, Kandinsky6TransformerDecoderBlock)):\n", - " block.forward = torch.compile(\n", - " block.forward, mode=\"max-autotune-no-cudagraphs\", dynamic=True\n", - " )\n", - "\n", - "is_piflow = pipe.scheduler.__class__.__name__ == \"PiflowScheduler\"\n", - "num_inference_steps = 16 if is_piflow else 50\n", - "guidance_scale = 1.0 if is_piflow else 5.0\n", - "print(f\"Scheduler: {pipe.scheduler.__class__.__name__}; steps: {num_inference_steps}; guidance: {guidance_scale}\")" - ] - }, - { - "cell_type": "markdown", - "id": "4ed6df0c", - "metadata": {}, - "source": [ - "### Run generation" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "0c31961d", - "metadata": {}, - "outputs": [], - "source": [ - "# NOTE: uses different mag_ratios\n", - "gen_mode = \"t2va\" # \"t2va\"\n", - "\n", - "output_suffix = \"\"\n", - "output_suffix += \"_piflow\" if is_piflow else \"\"\n", - "output_suffix += \"_expand\" if expand_prompts else \"\"\n", - "output_suffix += \"_compile\" if use_compile else \"\"\n", - "output_suffix += \"_magcache\" if use_magcache else \"\"\n", - "output_path = REPO / \"outputs\" / \"diffusers\" / model_name\n", - "output_path.mkdir(parents=True, exist_ok=True)\n", - "\n", - "# Use Diffusers' official MagCacheConfig and model enable_cache API. K6\n", - "# stores conditional/unconditional coefficients interleaved, so the\n", - "# configuration counts both CFG forwards.\n", - "magcache = pipe.transformer.config.get(\"magcache\")\n", - "if use_magcache and magcache is not None:\n", - " mag_ratios = magcache[\"mag_ratios\"]\n", - " if isinstance(mag_ratios, dict):\n", - " mag_ratios = mag_ratios[gen_mode]\n", - " diffusers_magcache = MagCacheConfig(\n", - " threshold=magcache[\"threshold\"],\n", - " max_skip_steps=magcache[\"max_skip_steps\"],\n", - " retention_ratio=magcache[\"retention_ratio\"],\n", - " num_inference_steps=2 * num_inference_steps,\n", - " mag_ratios=list(mag_ratios),\n", - " )\n", - " if pipe.transformer.is_cache_enabled:\n", - " pipe.transformer.disable_cache()\n", - " pipe.transformer.enable_cache(diffusers_magcache)\n", - " print(f\"Diffusers MagCache enabled: {pipe.transformer.is_cache_enabled}\")\n", - " print(f\"MagCache coefficients ({gen_mode}): {len(mag_ratios)}; threshold: {magcache['threshold']}\")\n", - "elif getattr(pipe.transformer, \"is_cache_enabled\", False):\n", - " pipe.transformer.disable_cache()\n", - "\n", - "# Model CPU offload keeps one active component on GPU at a time.\n", - "pipe.enable_model_cpu_offload(gpu_id=GPU_ID)\n", - "\n", - "print(f\"Qwen context length: {max_sequence_length}\")\n", - "print(f\"Audio sample rate: {pipe.audio_sample_rate}; video FPS: {sample_fps}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "cfc1e973", - "metadata": {}, - "outputs": [], - "source": [ - "batch_size = 2\n", - "prompt = (\n", - " \"A news anchor delivers the evening news in a clear, standard, and neutral tone, saying: \"\n", - " \"Good evening, ladies and gentlemen, and welcome to the evening news.\"\n", - ")\n", - "negative_prompt = \"blurry, distorted, low quality, noisy audio\"\n", - "\n", - "prompts = [prompt] * batch_size\n", - "negative_prompts = [negative_prompt] * batch_size\n", - "images = None\n", - "\n", - "# NOTE: use example image for I2VA mode\n", - "if gen_mode == \"i2va\":\n", - " image = REPO / \"assets\" / \"i2va_input.png\"\n", - " print(f\"reference image: {image}\")\n", - " images = [str(image)] * batch_size\n", - "\n", - "\n", - "result = pipe(\n", - " prompt=prompts,\n", - " image=images,\n", - " negative_prompt=negative_prompts,\n", - " height=480,\n", - " width=864,\n", - " num_frames=121,\n", - " sample_fps=sample_fps,\n", - " num_inference_steps=num_inference_steps,\n", - " max_sequence_length=max_sequence_length,\n", - " guidance_scale=guidance_scale,\n", - " sample_audio=True,\n", - " expand_prompts=True,\n", - " generator=torch.Generator().manual_seed(seed),\n", - " output_type=\"pt\",\n", - ")\n", - "frames = result.frames\n", - "audios = result.audio\n", - "print(f\"decoded frames: {tuple(frames.shape)}\")\n", - "print(f\"audio streams: {len(audios)}, samples in first stream: {audios[0].shape[0]}\")\n", - "\n", - "if use_magcache and magcache is not None:\n", - " print(\"MagCache run complete\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "f0879c75", - "metadata": {}, - "outputs": [], - "source": [ - "final_output_dir = output_path / f\"bs{batch_size}\"\n", - "final_output_dir.mkdir(parents=True, exist_ok=True)\n", - "\n", - "for i, (video, audio) in enumerate(zip(frames, audios)):\n", - " batch_suffix = f\"_batchid{i}\" if batch_size > 1 else \"\"\n", - " video_path = final_output_dir / f\"kandinsky6_{gen_mode}_{num_inference_steps}steps{output_suffix}{batch_suffix}.mp4\"\n", - " video = video.permute(1, 2, 3, 0)\n", - " audio_samples = torch.as_tensor(audio)[:, None].repeat(1, 2)\n", - " encode_video(\n", - " video,\n", - " fps=int(sample_fps),\n", - " output_path=str(video_path),\n", - " audio=audio_samples,\n", - " audio_sample_rate=int(pipe.audio_sample_rate),\n", - " )\n", - " print(f\"saved: {video_path}\")" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "kandinsky-6-video-pro", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.13.9" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/kandinsky6_ti2va.py b/notebooks/kandinsky6_ti2va.py deleted file mode 100644 index 6b423ad85abd..000000000000 --- a/notebooks/kandinsky6_ti2va.py +++ /dev/null @@ -1,182 +0,0 @@ -#!/usr/bin/env python3 -"""Kandinsky 6 text/image-to-video-and-audio generation. - -Script version of kandinsky6_ti2va.ipynb, for running on a headless server. - -Usage: - python kandinsky6_ti2va.py --gpu-id 0 - python kandinsky6_ti2va.py --model k6pro --batch-size 1 --no-magcache - python kandinsky6_ti2va.py --gen-mode i2va --image assets/i2va_input.png - -Authenticate with the Hub before running if the model repo is gated: - huggingface-cli login - or: - export HF_TOKEN=hf_... -""" - -from __future__ import annotations - -import argparse -from pathlib import Path - -import torch - -from diffusers import Kandinsky6TI2VAPipeline -from diffusers.hooks import MagCacheConfig -from diffusers.models.transformers.transformer_kandinsky6 import ( - Kandinsky6FusedTransformerDecoderBlock, - Kandinsky6TransformerDecoderBlock, - Kandinsky6TransformerEncoderBlock, -) -from diffusers.utils import encode_video - -MODELS = { - "k6pro": "kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers", - "k6pro_distill": "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers", -} - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--model", choices=sorted(MODELS), default="k6pro_distill") - parser.add_argument("--gpu-id", type=int, default=0) - parser.add_argument("--batch-size", type=int, default=2) - parser.add_argument("--gen-mode", choices=["t2va", "i2va"], default="t2va") - parser.add_argument("--image", type=Path, default=None, help="Reference image, required for --gen-mode i2va.") - parser.add_argument( - "--prompt", - default=( - "A news anchor delivers the evening news in a clear, standard, and neutral tone, saying: " - "Good evening, ladies and gentlemen, and welcome to the evening news." - ), - ) - parser.add_argument("--negative-prompt", default="blurry, distorted, low quality, noisy audio") - parser.add_argument("--height", type=int, default=480) - parser.add_argument("--width", type=int, default=864) - parser.add_argument("--num-frames", type=int, default=121) - parser.add_argument("--sample-fps", type=float, default=24.0) - parser.add_argument("--max-sequence-length", type=int, default=1024) - parser.add_argument("--seed", type=int, default=1137) - parser.add_argument("--expand-prompts", action=argparse.BooleanOptionalAction, default=True) - parser.add_argument( - "--magcache", - action=argparse.BooleanOptionalAction, - default=None, - help="Enable MagCache. Defaults to on for non-distilled models, off for distilled ones.", - ) - parser.add_argument("--use-compile", action="store_true") - parser.add_argument("--output-dir", type=Path, default=Path("outputs/diffusers")) - return parser.parse_args() - - -def main() -> None: - args = parse_args() - if args.gen_mode == "i2va" and args.image is None: - raise SystemExit("--gen-mode i2va requires --image ") - - model_path = MODELS[args.model] - use_magcache = bool("distill" not in args.model) if args.magcache is None else args.magcache - backend = "native" if args.use_compile else "_flash_3" - - pipe = Kandinsky6TI2VAPipeline.from_pretrained(model_path, dtype=torch.bfloat16) - pipe.transformer.set_attention_backend(backend) - - # Compile after loading the weights and before CPU offload. - if args.use_compile: - torch.set_float32_matmul_precision("high") - for block in pipe.transformer.modules(): - if isinstance(block, Kandinsky6FusedTransformerDecoderBlock): - block.forward = torch.compile(block.forward) - elif isinstance(block, (Kandinsky6TransformerEncoderBlock, Kandinsky6TransformerDecoderBlock)): - block.forward = torch.compile(block.forward, mode="max-autotune-no-cudagraphs", dynamic=True) - - is_piflow = pipe.scheduler.__class__.__name__ == "PiflowScheduler" - num_inference_steps = 16 if is_piflow else 50 - guidance_scale = 1.0 if is_piflow else 5.0 - print(f"Scheduler: {pipe.scheduler.__class__.__name__}; steps: {num_inference_steps}; guidance: {guidance_scale}") - - output_suffix = "" - output_suffix += "_piflow" if is_piflow else "" - output_suffix += "_expand" if args.expand_prompts else "" - output_suffix += "_compile" if args.use_compile else "" - output_suffix += "_magcache" if use_magcache else "" - output_path = args.output_dir / args.model - output_path.mkdir(parents=True, exist_ok=True) - - # Diffusers' MagCacheConfig, sourced from the metadata the checkpoint carries in its - # own config. K6 stores conditional/unconditional coefficients interleaved, so the - # configuration counts both CFG forwards. - magcache_meta = pipe.transformer.config.get("magcache") - if use_magcache and magcache_meta is not None: - mag_ratios = magcache_meta["mag_ratios"] - if isinstance(mag_ratios, dict): - mag_ratios = mag_ratios[args.gen_mode] - diffusers_magcache = MagCacheConfig( - threshold=magcache_meta["threshold"], - max_skip_steps=magcache_meta["max_skip_steps"], - retention_ratio=magcache_meta["retention_ratio"], - num_inference_steps=2 * num_inference_steps, - mag_ratios=list(mag_ratios), - ) - if pipe.transformer.is_cache_enabled: - pipe.transformer.disable_cache() - pipe.transformer.enable_cache(diffusers_magcache) - print(f"MagCache enabled: {pipe.transformer.is_cache_enabled}") - print(f"MagCache coefficients ({args.gen_mode}): {len(mag_ratios)}; threshold: {magcache_meta['threshold']}") - else: - use_magcache = False - if getattr(pipe.transformer, "is_cache_enabled", False): - pipe.transformer.disable_cache() - - # Model CPU offload keeps one active component on GPU at a time. - pipe.enable_model_cpu_offload(gpu_id=args.gpu_id) - - print(f"Qwen context length: {args.max_sequence_length}") - print(f"Audio sample rate: {pipe.audio_sample_rate}; video FPS: {args.sample_fps}") - - prompts = [args.prompt] * args.batch_size - negative_prompts = [args.negative_prompt] * args.batch_size - images = [str(args.image)] * args.batch_size if args.gen_mode == "i2va" else None - - result = pipe( - prompt=prompts, - image=images, - negative_prompt=negative_prompts, - height=args.height, - width=args.width, - num_frames=args.num_frames, - sample_fps=args.sample_fps, - num_inference_steps=num_inference_steps, - max_sequence_length=args.max_sequence_length, - guidance_scale=guidance_scale, - sample_audio=True, - expand_prompts=args.expand_prompts, - generator=torch.Generator().manual_seed(args.seed), - output_type="pt", - ) - frames = result.frames - audios = result.audio - print(f"decoded frames: {tuple(frames.shape)}") - print(f"audio streams: {len(audios)}, samples in first stream: {audios[0].shape[0]}") - - final_output_dir = output_path / f"bs{args.batch_size}" - final_output_dir.mkdir(parents=True, exist_ok=True) - for i, (video, audio) in enumerate(zip(frames, audios)): - batch_suffix = f"_batchid{i}" if args.batch_size > 1 else "" - video_path = ( - final_output_dir / f"kandinsky6_{args.gen_mode}_{num_inference_steps}steps{output_suffix}{batch_suffix}.mp4" - ) - video = video.permute(1, 2, 3, 0) - audio_samples = torch.as_tensor(audio)[:, None].repeat(1, 2) - encode_video( - video, - fps=int(args.sample_fps), - output_path=str(video_path), - audio=audio_samples, - audio_sample_rate=int(pipe.audio_sample_rate), - ) - print(f"saved: {video_path}") - - -if __name__ == "__main__": - main() diff --git a/notebooks/kandinsky6_ti2va_sr.ipynb b/notebooks/kandinsky6_ti2va_sr.ipynb deleted file mode 100644 index 2f4a42e5daa4..000000000000 --- a/notebooks/kandinsky6_ti2va_sr.ipynb +++ /dev/null @@ -1,130 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Kandinsky 6 T2VA \u2192 SR \u2014 stacked Diffusers example\n", - "\n", - "This notebook follows the two-stage pattern used by LTX-2.5: load the T2VA and SR artifacts independently, generate video and audio with T2VA, then pass the decoded video and generated audio to SR. The pipelines remain separate and can be loaded, offloaded, or replaced independently.\n", - "\n", - "The SR artifact is created with `convert_sr_checkpoint.py` and is a separately packaged Diffusers bundle containing `Kandinsky6SRPipeline`. The production SR route is `resolution_scale=2.25` (1.125x pixel pre-upscale followed by the x2 KVAE path)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from pathlib import Path\n", - "\n", - "import torch\n", - "from diffusers import Kandinsky6SRPipeline, Kandinsky6TI2VAPipeline\n", - "from diffusers.utils import encode_video\n", - "\n", - "REPO = Path.cwd()\n", - "while REPO != REPO.parent and not (REPO / \"dev\" / \"exports\" / \"diffusers\").exists():\n", - " REPO = REPO.parent\n", - "if not (REPO / \"dev\" / \"exports\" / \"diffusers\").exists():\n", - " raise RuntimeError(\"Run this notebook from the K6 repository or one of its subdirectories.\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Stage 1 and stage 2 are independent Diffusers artifacts.\n", - "T2VA_MODEL_PATH = REPO / \"outputs\" / \"diffusers_bundle\" / \"kandinsky6_diffusers\"\n", - "SR_MODEL_PATH = REPO / \"outputs\" / \"diffusers_bundle\" / \"kandinsky6_sr_diffusers\"\n", - "GPU_ID = 0 # Change this to the CUDA device used for inference.\n", - "\n", - "t2va_pipe = Kandinsky6TI2VAPipeline.from_pretrained(\n", - " T2VA_MODEL_PATH,\n", - " torch_dtype=torch.bfloat16,\n", - ")\n", - "t2va_pipe.enable_model_cpu_offload(gpu_id=GPU_ID)\n", - "\n", - "sr_pipe = Kandinsky6SRPipeline.from_pretrained(\n", - " SR_MODEL_PATH,\n", - " torch_dtype=torch.bfloat16,\n", - ")\n", - "sr_pipe.enable_model_cpu_offload(gpu_id=GPU_ID)\n", - "\n", - "print(f\"T2VA artifact: {T2VA_MODEL_PATH}\")\n", - "print(f\"SR artifact: {SR_MODEL_PATH}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "prompt = \"A cinematic shot of a small red sailboat crossing a calm blue lake at sunset, with gentle waves and natural ambient sound.\"\n", - "negative_prompt = \"blurry, distorted, low quality, noisy audio\"\n", - "sample_fps = 24.0\n", - "\n", - "t2va_result = t2va_pipe(\n", - " prompt=prompt,\n", - " negative_prompt=negative_prompt,\n", - " height=480,\n", - " width=864,\n", - " num_frames=121, # 1 + 8k: Kandinsky6SRPipeline rejects other values\n", - " sample_fps=sample_fps,\n", - " num_inference_steps=50,\n", - " guidance_scale=5.0,\n", - " sample_audio=True,\n", - " output_type=\"pt\",\n", - ")\n", - "if t2va_result.audio is None:\n", - " raise RuntimeError(\"The T2VA pipeline returned no audio\")\n", - "print(f\"T2VA frames: {tuple(t2va_result.frames.shape)}\")\n", - "print(f\"T2VA audio streams: {len(t2va_result.audio)}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Stage 2 consumes the stage-1 tensors directly; no intermediate video encode/decode is needed.\n", - "output_path = REPO / \"outputs\" / \"diffusers\" / \"kandinsky6_t2va_sr.mp4\"\n", - "output_path.parent.mkdir(parents=True, exist_ok=True)\n", - "\n", - "sr_result = sr_pipe(\n", - " video=t2va_result.frames,\n", - " resolution_scale=2.25,\n", - " output_type=\"pt\",\n", - ")\n", - "print(f\"SR frames: {tuple(sr_result.frames.shape)}\")\n", - "\n", - "audio = torch.as_tensor(t2va_result.audio[0])[:, None].repeat(1, 2)\n", - "encode_video(\n", - " sr_result.frames[0].permute(1, 2, 3, 0),\n", - " fps=int(sample_fps),\n", - " output_path=str(output_path),\n", - " audio=audio,\n", - " audio_sample_rate=int(t2va_pipe.audio_sample_rate),\n", - ")\n", - "print(f\"saved: {output_path}\")" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "name": "python", - "version": "3.11" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/kandinsky6_ti2va_sr.py b/notebooks/kandinsky6_ti2va_sr.py deleted file mode 100644 index 2870c3da64c9..000000000000 --- a/notebooks/kandinsky6_ti2va_sr.py +++ /dev/null @@ -1,113 +0,0 @@ -#!/usr/bin/env python3 -"""Kandinsky 6 T2VA -> SR, stacked. - -Script version of kandinsky6_ti2va_sr.ipynb, for running on a headless server. - -Follows the two-stage pattern used by LTX-2.5: load the T2VA and SR artifacts -independently, generate video and audio with T2VA, then pass the decoded video -and generated audio to SR directly (no intermediate video encode/decode). The -pipelines remain separate and can be loaded, offloaded, or replaced independently. - -The production SR route is --resolution-scale 2.25 (1.125x pixel pre-upscale -followed by the x2 KVAE path). - -Usage: - python kandinsky6_ti2va_sr.py - python kandinsky6_ti2va_sr.py --t2va-model /path/to/local/bundle --sr-model /path/to/local/sr_bundle - -Authenticate with the Hub before running if a model repo is gated: - huggingface-cli login - or: - export HF_TOKEN=hf_... -""" - -from __future__ import annotations - -import argparse -from pathlib import Path - -import torch - -from diffusers import Kandinsky6SRPipeline, Kandinsky6TI2VAPipeline -from diffusers.utils import encode_video - -DEFAULT_T2VA_MODEL = "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers" -DEFAULT_SR_MODEL = "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers" - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument("--t2va-model", default=DEFAULT_T2VA_MODEL, help="Hub repo id or local bundle path.") - parser.add_argument("--sr-model", default=DEFAULT_SR_MODEL, help="Hub repo id or local bundle path.") - parser.add_argument("--gpu-id", type=int, default=0) - parser.add_argument( - "--prompt", - default=( - "A cinematic shot of a small red sailboat crossing a calm blue lake at sunset, " - "with gentle waves and natural ambient sound." - ), - ) - parser.add_argument("--negative-prompt", default="blurry, distorted, low quality, noisy audio") - parser.add_argument("--height", type=int, default=480) - parser.add_argument("--width", type=int, default=864) - parser.add_argument( - "--num-frames", type=int, default=121, help="Must be 1 + 8k: Kandinsky6SRPipeline rejects other values." - ) - parser.add_argument("--sample-fps", type=float, default=24.0) - parser.add_argument("--num-inference-steps", type=int, default=50) - parser.add_argument("--guidance-scale", type=float, default=5.0) - parser.add_argument("--resolution-scale", type=float, default=2.25) - parser.add_argument("--output", type=Path, default=Path("outputs/diffusers/kandinsky6_t2va_sr.mp4")) - return parser.parse_args() - - -def main() -> None: - args = parse_args() - - t2va_pipe = Kandinsky6TI2VAPipeline.from_pretrained(args.t2va_model, torch_dtype=torch.bfloat16) - t2va_pipe.enable_model_cpu_offload(gpu_id=args.gpu_id) - - sr_pipe = Kandinsky6SRPipeline.from_pretrained(args.sr_model, torch_dtype=torch.bfloat16) - sr_pipe.enable_model_cpu_offload(gpu_id=args.gpu_id) - - print(f"T2VA artifact: {args.t2va_model}") - print(f"SR artifact: {args.sr_model}") - - t2va_result = t2va_pipe( - prompt=args.prompt, - negative_prompt=args.negative_prompt, - height=args.height, - width=args.width, - num_frames=args.num_frames, - sample_fps=args.sample_fps, - num_inference_steps=args.num_inference_steps, - guidance_scale=args.guidance_scale, - sample_audio=True, - output_type="pt", - ) - if t2va_result.audio is None: - raise RuntimeError("The T2VA pipeline returned no audio") - print(f"T2VA frames: {tuple(t2va_result.frames.shape)}") - print(f"T2VA audio streams: {len(t2va_result.audio)}") - - sr_result = sr_pipe( - video=t2va_result.frames, - resolution_scale=args.resolution_scale, - output_type="pt", - ) - print(f"SR frames: {tuple(sr_result.frames.shape)}") - - args.output.parent.mkdir(parents=True, exist_ok=True) - audio = torch.as_tensor(t2va_result.audio[0])[:, None].repeat(1, 2) - encode_video( - sr_result.frames[0].permute(1, 2, 3, 0), - fps=int(args.sample_fps), - output_path=str(args.output), - audio=audio, - audio_sample_rate=int(t2va_pipe.audio_sample_rate), - ) - print(f"saved: {args.output}") - - -if __name__ == "__main__": - main() From 66a0d44df290e19fae64c2f798c8952fd09db19b Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 28 Sep 2026 18:23:52 +0300 Subject: [PATCH 21/62] patch mmaudio VAE --- src/diffusers/models/autoencoders/autoencoder_mmaudio.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index d51200034d49..21df0ca3f35b 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -132,8 +132,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: batch_size, channels, length = x.shape qkv = self.qkv(x).reshape(batch_size, self.num_heads, -1, 3, length) query, key, value = normalize(qkv, dim=2).unbind(3) - # `(B, heads, D, T)` -> `(B, T, heads, D)` for the attention dispatcher. - query, key, value = (t.permute(0, 3, 1, 2) for t in (query, key, value)) + # `(B, heads, D, T)` -> `(B, T, heads, D)` for the attention dispatcher. `D` is not contiguous after the + # permute (q/k/v are interleaved along the channel dim), which some attention backends (e.g. FlashAttention-3) + # require, so make the tensors contiguous here. + query, key, value = (t.permute(0, 3, 1, 2).contiguous() for t in (query, key, value)) hidden_states = dispatch_attention_fn(query, key, value) hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, channels, length) return mp_sum(x, self.proj_out(hidden_states), t=0.3) From 0ec1bd1ac8578c76db5d747adb5173a575897826 Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 28 Sep 2026 18:41:02 +0300 Subject: [PATCH 22/62] patch mmaudio VAE --- src/diffusers/models/autoencoders/autoencoder_mmaudio.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index 21df0ca3f35b..175005f662f6 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -24,7 +24,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...utils.accelerate_utils import apply_forward_hook -from ..attention_dispatch import dispatch_attention_fn +from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin from .vae import DecoderOutput, DiagonalGaussianDistribution @@ -136,7 +136,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # permute (q/k/v are interleaved along the channel dim), which some attention backends (e.g. FlashAttention-3) # require, so make the tensors contiguous here. query, key, value = (t.permute(0, 3, 1, 2).contiguous() for t in (query, key, value)) - hidden_states = dispatch_attention_fn(query, key, value) + # This block always attends with a single head over the full channel width (like the SD-VAE `AttnBlock`), + # so `head_dim` can be in the thousands. FlashAttention-family backends cap `head_dim` at 256, so force + # the native SDPA backend here regardless of whatever backend is globally active for the rest of the + # pipeline (e.g. via `transformer.set_attention_backend(...)`). + hidden_states = dispatch_attention_fn(query, key, value, backend=AttentionBackendName.NATIVE) hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, channels, length) return mp_sum(x, self.proj_out(hidden_states), t=0.3) From c17252c486812b6dff805ff6b7e8f53fa5a88cd7 Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 28 Sep 2026 22:43:16 +0300 Subject: [PATCH 23/62] fix inits --- src/diffusers/__init__.py | 2 ++ src/diffusers/pipelines/__init__.py | 2 ++ .../utils/dummy_torch_and_transformers_objects.py | 15 +++++++++++++++ 3 files changed, 19 insertions(+) diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 1b038bbd2c98..1f42cd05c885 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -699,6 +699,7 @@ "Kandinsky5I2VPipeline", "Kandinsky5T2IPipeline", "Kandinsky5T2VPipeline", + "Kandinsky6SRLatentUpscalerBank", "Kandinsky6SRPipeline", "Kandinsky6SRPipelineOutput", "Kandinsky6TI2VAPipeline", @@ -1554,6 +1555,7 @@ Kandinsky5I2VPipeline, Kandinsky5T2IPipeline, Kandinsky5T2VPipeline, + Kandinsky6SRLatentUpscalerBank, Kandinsky6SRPipeline, Kandinsky6SRPipelineOutput, Kandinsky6TI2VAPipeline, diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index fc415647492d..4f0433b0ca02 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -449,6 +449,7 @@ "Kandinsky5I2IPipeline", ] _import_structure["kandinsky6"] = [ + "Kandinsky6SRLatentUpscalerBank", "Kandinsky6SRPipeline", "Kandinsky6SRPipelineOutput", "Kandinsky6TI2VAPipeline", @@ -784,6 +785,7 @@ Kandinsky5T2VPipeline, ) from .kandinsky6 import ( + Kandinsky6SRLatentUpscalerBank, Kandinsky6SRPipeline, Kandinsky6SRPipelineOutput, Kandinsky6TI2VAPipeline, diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index a1f889a6dae6..042f0b3d6cdf 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -2672,6 +2672,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class Kandinsky6SRLatentUpscalerBank(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 Kandinsky6SRPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] From d6e1f65e5b7ee858a49cf57e540c053393481259 Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 28 Sep 2026 23:39:00 +0300 Subject: [PATCH 24/62] fix flex attn --- .../transformers/transformer_kandinsky6_sr.py | 41 ++++++++++--------- 1 file changed, 22 insertions(+), 19 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index cbc7095fe2de..ab484947c099 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -26,7 +26,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin from ..attention import AttentionMixin, AttentionModuleMixin -from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn +from ..attention_dispatch import dispatch_attention_fn from ..embeddings import get_timestep_embedding from ..modeling_outputs import Transformer2DModelOutput from ..modeling_utils import ModelMixin, get_parameter_dtype @@ -146,7 +146,15 @@ def sliding_tile_mask( class Kandinsky6SRAttnProcessor: - """Self-attention processor of the SR transformer: dense by default, NABLA sparse when `sparse_params` is set.""" + """Self-attention processor of the SR transformer: dense by default, NABLA sparse when `sparse_params` is set. + + Like [`Kandinsky5AttnProcessor`], the attention backend is chosen through the usual `_attention_backend` + (`set_attention_backend`), even in the NABLA sparse branch. Its sparsity is expressed as a `BlockMask`, which + only the `flex` backend can consume, so sparse checkpoints require `set_attention_backend("flex")`. That backend + also needs to run under `torch.compile` (e.g. `transformer.compile_repeated_blocks()`): uncompiled, PyTorch's + flex attention falls back to an eager implementation that materializes the full attention matrix, which does not + fit in memory at video resolutions. + """ _attention_backend = None _parallel_config = None @@ -167,30 +175,25 @@ def __call__( key = apply_rotary(key, rotary_emb) if sparse_params is None: - hidden_states = dispatch_attention_fn( - query, - key, - value, - backend=self._attention_backend, - parallel_config=self._parallel_config, - ) + attn_mask = None else: # The block statistics are computed from the `(B, heads, S, D)` layout the mask builder expects; the - # resulting `BlockMask` only runs on the flex backend. - block_mask = nabla_block_mask( + # resulting `BlockMask` only runs on the flex backend, so sparse checkpoints need + # `set_attention_backend("flex")` (see the class docstring). + attn_mask = nabla_block_mask( query.transpose(1, 2), key.transpose(1, 2), sparse_params["sta_mask"], thr=sparse_params["threshold"], ) - hidden_states = dispatch_attention_fn( - query, - key, - value, - attn_mask=block_mask, - backend=AttentionBackendName.FLEX, - parallel_config=self._parallel_config, - ) + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attn_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) return attn.out_layer(hidden_states.flatten(2, 3)) From b558b74cd7e35ac0713025b096362c20c0f82a07 Mon Sep 17 00:00:00 2001 From: leffff Date: Tue, 29 Sep 2026 10:46:32 +0300 Subject: [PATCH 25/62] refactoring --- src/diffusers/models/attention_dispatch.py | 10 ++++++++++ .../models/autoencoders/autoencoder_mmaudio.py | 6 +++--- .../models/transformers/transformer_kandinsky6.py | 14 ++++++++------ .../kandinsky6/modeling_latent_upscaler.py | 5 +++-- .../pipelines/kandinsky6/pipeline_kandinsky6_sr.py | 6 ++++++ .../kandinsky6/pipeline_kandinsky6_ti2va.py | 11 ++++++++++- .../kandinsky6/test_pipeline_kandinsky6_ti2va.py | 6 ++++++ tests/pipelines/testing_utils/memory.py | 2 +- 8 files changed, 47 insertions(+), 13 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index e7cc20f580d4..9e08525042cd 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -3618,6 +3618,15 @@ def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): else: raise ValueError("Attention mask must be either None, a BlockMask, or a 2D/4D tensor.") + # A `BlockMask`'s own block size is the sparsity granularity the caller built it around (e.g. NABLA's 64-token + # fractal blocks). The forward kernel's tile sizes must divide that granularity evenly; torch's autotuned + # defaults aren't guaranteed to (they're picked independent of the mask), so pin them to the mask's own block + # size, which always divides itself. + kernel_options = None + if block_mask is not None: + q_block_size, kv_block_size = block_mask.BLOCK_SIZE + kernel_options = {"BLOCK_M": q_block_size, "BLOCK_N": kv_block_size} + query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) out = flex_attention.flex_attention( query=query, @@ -3628,6 +3637,7 @@ def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): scale=scale, enable_gqa=enable_gqa, return_lse=return_lse, + kernel_options=kernel_options, ) out = out.permute(0, 2, 1, 3) return out diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index 175005f662f6..9a8d348e23c8 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -55,9 +55,9 @@ def mel_to_hz(mel: torch.Tensor) -> torch.Tensor: mel >= min_log_mel, min_log_hz * torch.exp(log_step * (mel - min_log_mel)), linear_step * mel ) - fft_freqs = torch.linspace(0, sample_rate / 2, 1 + n_fft // 2, dtype=torch.float64) - mel_limits = hz_to_mel(torch.tensor([f_min, f_max], dtype=torch.float64)) - mel_freqs = mel_to_hz(torch.linspace(mel_limits[0], mel_limits[1], num_mels + 2, dtype=torch.float64)) + fft_freqs = torch.linspace(0, sample_rate / 2, 1 + n_fft // 2, dtype=torch.float32) + mel_limits = hz_to_mel(torch.tensor([f_min, f_max], dtype=torch.float32)) + mel_freqs = mel_to_hz(torch.linspace(mel_limits[0], mel_limits[1], num_mels + 2, dtype=torch.float32)) freq_diff = torch.diff(mel_freqs) ramps = mel_freqs[:, None] - fft_freqs[None, :] lower = -ramps[:-2] / freq_diff[:-1, None] diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index fb3383cc5296..ee523dc1287c 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -28,7 +28,7 @@ from ...hooks import MagCacheConfig from ...hooks.hooks import HookRegistry, ModelHook, StateManager from ...hooks.mag_cache import MagCacheState -from ...loaders import FromOriginalModelMixin, PeftAdapterMixin +from ...loaders import PeftAdapterMixin from ...utils import BaseOutput from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn @@ -435,7 +435,13 @@ def disable_cache(self) -> None: class Kandinsky6AttnProcessor: - """Diffusers attention processor used by the TI2VA transformer.""" + """Diffusers attention processor used by the TI2VA transformer. + + The NABLA sparse branch always dispatches on the `flex` backend, regardless of `_attention_backend`: its + `BlockMask` only the `flex` backend can consume, and `_attention_backend` can be transparently remapped to a + masked-attention variant (see the setter above) when `text_token_padding` is enabled, which would otherwise be + given a `BlockMask` it can't handle. + """ _attention_backend = None _parallel_config = None @@ -637,13 +643,10 @@ def forward( sparse_params: dict[str, Any] | None = None, rope_q: Tensor | None = None, rope_kv: Tensor | None = None, - **kwargs: Any, ) -> Tensor: # Native K6 blocks pass ``(hidden, rope, mask_or_sparse)`` for # self-attention. Keep that call shape while exposing Diffusers' # encoder_hidden_states/rotary_emb keyword boundary. - if attn_mask is None: - attn_mask = kwargs.get("attention_mask") rotary_emb = rope_q if rope_q is not None else rotary_emb rotary_emb_kv = rope_kv if rope_kv is not None else rotary_emb return self.processor( @@ -945,7 +948,6 @@ class Kandinsky6Transformer3DModel( ModelMixin, ConfigMixin, PeftAdapterMixin, - FromOriginalModelMixin, CacheMixin, AttentionMixin, ): diff --git a/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py b/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py index f5bbcf10fc66..104ecb2fd27c 100644 --- a/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py +++ b/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py @@ -44,7 +44,7 @@ def forward(self, hidden_states: Tensor) -> Tensor: class Kandinsky6SRLatentUpscalerRMSNorm(nn.Module): - """Channel-first RMS normalization with a learnable gain.""" + """Channel-first RMS normalization with a learnable gain, computed in float32.""" def __init__(self, num_channels: int) -> None: super().__init__() @@ -52,7 +52,8 @@ def __init__(self, num_channels: int) -> None: self.gamma = nn.Parameter(torch.ones(num_channels, 1, 1, 1)) def forward(self, hidden_states: Tensor) -> Tensor: - return functional.normalize(hidden_states, dim=1) * self.scale * self.gamma + normalized = functional.normalize(hidden_states.float(), dim=1).to(hidden_states.dtype) + return normalized * self.scale * self.gamma class Kandinsky6SRLatentUpscalerModulatedNorm(nn.Module): diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py index 3579b254a29b..ae3529aaffe9 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py @@ -291,6 +291,12 @@ def __call__( video = frames.unflatten(0, (batch_size, num_frames)).permute(0, 2, 1, 3, 4) batch_size, _, num_frames, height, width = video.shape + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + # 3. Tile grid at the input resolution; every tile is refined at the closest trained tile resolution base_height, base_width = min(self.transformer_tile_sizes, key=lambda hw: abs(hw[1] / hw[0] - width / height)) snap = self.vae_scale_factor_spatial diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index 13a48c5c579f..a24c5e21dccf 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -24,13 +24,15 @@ from ...image_processor import PipelineImageInput from ...models import AutoencoderKLHunyuanVideo, Kandinsky6Transformer3DModel, MMAudioVAE from ...schedulers import FlowMatchEulerDiscreteScheduler, PiflowScheduler -from ...utils import replace_example_docstring +from ...utils import logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import Kandinsky6TI2VAPipelineOutput +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + EXAMPLE_DOC_STRING = """ Examples: ```python @@ -717,6 +719,13 @@ def __call__( negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, ) + if num_frames % self.vae_scale_factor_temporal != 1: + logger.warning( + f"`num_frames - 1` has to be divisible by {self.vae_scale_factor_temporal}. Rounding to the nearest number." + ) + num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 + num_frames = max(num_frames, 1) + if isinstance(self.scheduler, PiflowScheduler) and guidance_scale != 1.0: raise ValueError("`PiflowScheduler` runs the distilled checkpoints and requires `guidance_scale=1.0`.") diff --git a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py index 834c8fe699c0..5ebaa1d4613d 100644 --- a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py +++ b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py @@ -49,6 +49,12 @@ class Kandinsky6TI2VAPipelineTesterConfig(BasePipelineTesterConfig): # (num_frames=5, height=16, width=16) at the tiny VAE's 8x spatial / 4x temporal compression. output_shape = (5, 3, 16, 16) + # `audio_vae` (`MMAudioVAE`) reads its own `data_std`/`data_mean` buffers directly in + # `MMAudioAutoencoder.decode` rather than through one of its leaf submodules, so leaf-level onload hooks on + # its children never onload them and decoding runs on a mix of onload/offload devices. Every other component + # offloads fine at leaf level, so exclude just this one rather than skipping the test. + group_offloading_leaf_level_exclude_modules = ["audio_vae"] + def get_dummy_components(self): torch.manual_seed(0) # 3 down/up levels so the tiny VAE's realized compression matches the declared diff --git a/tests/pipelines/testing_utils/memory.py b/tests/pipelines/testing_utils/memory.py index e74d50537c8d..4639016f640a 100644 --- a/tests/pipelines/testing_utils/memory.py +++ b/tests/pipelines/testing_utils/memory.py @@ -312,7 +312,7 @@ def enable_group_offload_on_component(pipe, group_offloading_kwargs): for module in component.modules() if hasattr(module, "_diffusers_hook") ) - for component_name in ["vae", "vqvae", "image_encoder"]: + for component_name in ["vae", "vqvae", "image_encoder", "audio_vae", "latent_upscaler"]: component = getattr(pipe, component_name, None) if isinstance(component, torch.nn.Module): component.to(torch_device) From 9debf90ee2494e680af1453a13f1e32b7329f07c Mon Sep 17 00:00:00 2001 From: leffff Date: Tue, 29 Sep 2026 11:17:55 +0300 Subject: [PATCH 26/62] remove __all__ --- .../models/autoencoders/autoencoder_kandinsky6_sr.py | 3 --- .../models/autoencoders/autoencoder_mmaudio.py | 3 --- .../models/transformers/transformer_kandinsky6.py | 12 ------------ 3 files changed, 18 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index c2165cea25aa..5ab7438821fd 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -592,6 +592,3 @@ def forward( if not return_dict: return (decoded,) return DecoderOutput(sample=decoded) - - -__all__ = ["Kandinsky6SRVAE"] diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index 9a8d348e23c8..8d6f6c5b6cfd 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -617,6 +617,3 @@ def forward( if not return_dict: return (waveform,) return DecoderOutput(sample=waveform) - - -__all__ = ["MMAudioVAE"] diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index ee523dc1287c..6ce3b696749f 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -451,18 +451,6 @@ def __init__(self, attention_backend=None, parallel_config=None): self._attention_backend = attention_backend self._parallel_config = parallel_config - @property - def _attention_backend(self): - return self.__attention_backend - - @_attention_backend.setter - def _attention_backend(self, backend): - if self._masked and backend is not None: - name = getattr(backend, "value", backend).lower() - backend = _MASKED_ATTENTION_BACKENDS.get(name, name) - backend = AttentionBackendName(backend) - self.__attention_backend = backend - def __call__( self, attn: Any, From 1cf55a81e0d8747ee7b9f666083f4bf4200d7159 Mon Sep 17 00:00:00 2001 From: leffff Date: Tue, 29 Sep 2026 12:34:35 +0300 Subject: [PATCH 27/62] separate MMAudio VAE and Vocoder --- docs/source/en/api/models/kandinsky6_vae.md | 7 +- src/diffusers/__init__.py | 2 + .../autoencoders/autoencoder_mmaudio.py | 226 ++--------------- src/diffusers/pipelines/__init__.py | 2 + .../pipelines/kandinsky6/__init__.py | 2 + .../pipelines/kandinsky6/modeling_vocoder.py | 228 ++++++++++++++++++ .../kandinsky6/pipeline_kandinsky6_ti2va.py | 26 +- .../dummy_torch_and_transformers_objects.py | 15 ++ .../test_models_autoencoder_mmaudio.py | 9 +- .../kandinsky6/test_models_vocoder.py | 86 +++++++ .../test_pipeline_kandinsky6_ti2va.py | 18 +- 11 files changed, 395 insertions(+), 226 deletions(-) create mode 100644 src/diffusers/pipelines/kandinsky6/modeling_vocoder.py create mode 100644 tests/pipelines/kandinsky6/test_models_vocoder.py diff --git a/docs/source/en/api/models/kandinsky6_vae.md b/docs/source/en/api/models/kandinsky6_vae.md index ba550b3384a7..2dd4c31a0689 100644 --- a/docs/source/en/api/models/kandinsky6_vae.md +++ b/docs/source/en/api/models/kandinsky6_vae.md @@ -13,7 +13,7 @@ specific language governing permissions and limitations under the License. # Kandinsky 6 VAEs Kandinsky 6 uses a causal 3D K-VAE for video super-resolution and the MMAudio mel-spectrogram VAE, paired with a -BigVGAN vocoder, for synchronized audio generation. +separate BigVGAN [`~pipelines.kandinsky6.MMAudioVocoder`], for synchronized audio generation. ## Kandinsky6SRVAE @@ -28,3 +28,8 @@ BigVGAN vocoder, for synchronized audio generation. - encode - decode - all + +## MMAudioVocoder + +[[autodoc]] pipelines.kandinsky6.MMAudioVocoder + - forward diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 1f42cd05c885..5aedf783a997 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -750,6 +750,7 @@ "Lumina2Text2ImgPipeline", "LuminaPipeline", "LuminaText2ImgPipeline", + "MMAudioVocoder", "MarigoldDepthPipeline", "MarigoldIntrinsicsPipeline", "MarigoldNormalsPipeline", @@ -1609,6 +1610,7 @@ MarigoldDepthPipeline, MarigoldIntrinsicsPipeline, MarigoldNormalsPipeline, + MMAudioVocoder, MochiPipeline, MotifVideoImage2VideoPipeline, MotifVideoPipeline, diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index 8d6f6c5b6cfd..ff644daea5c0 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -12,7 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""MMAudio mel-spectrogram VAE and BigVGAN vocoder used by the Kandinsky 6 TI2VA pipeline.""" +"""MMAudio mel-spectrogram VAE used by the Kandinsky 6 TI2VA pipeline. + +The BigVGAN vocoder that turns this VAE's decoded mel spectrograms into waveforms is a separate, pipeline-local +component, [`~pipelines.kandinsky6.MMAudioVocoder`]. +""" from __future__ import annotations @@ -295,171 +299,6 @@ def decode(self, z: torch.Tensor) -> torch.Tensor: return self.decoder(z) * self.data_std + self.data_mean -# BigVGAN vocoder (https://github.com/NVIDIA/BigVGAN, MIT license) with the anti-aliased Snake activations of -# https://github.com/junjun3518/alias-free-torch (Apache 2.0). - - -def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: - """Kaiser-windowed sinc low-pass filter of shape `(1, 1, kernel_size)` normalized to unit sum.""" - even = kernel_size % 2 == 0 - half_size = kernel_size // 2 - - attenuation = 2.285 * (half_size - 1) * math.pi * 4 * half_width + 7.95 - if attenuation > 50.0: - beta = 0.1102 * (attenuation - 8.7) - elif attenuation >= 21.0: - beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0) - else: - beta = 0.0 - window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) - - time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size - filter_ = 2 * cutoff * window * torch.sinc(2 * cutoff * time) - filter_ = filter_ / filter_.sum() - return filter_.view(1, 1, kernel_size) - - -class MMAudioSnakeBeta(nn.Module): - """`x + 1/b * sin^2(a * x)` with per-channel log-scale frequency `a` and magnitude `b`.""" - - def __init__(self, channels: int) -> None: - super().__init__() - self.alpha = nn.Parameter(torch.zeros(channels)) - self.beta = nn.Parameter(torch.zeros(channels)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - alpha = torch.exp(self.alpha)[None, :, None] - beta = torch.exp(self.beta)[None, :, None] - return x + (1.0 / (beta + 1e-9)) * torch.sin(x * alpha).pow(2) - - -class MMAudioLowPassFilter1d(nn.Module): - def __init__(self, cutoff: float, half_width: float, stride: int, kernel_size: int) -> None: - super().__init__() - even = kernel_size % 2 == 0 - self.pad_left = kernel_size // 2 - int(even) - self.pad_right = kernel_size // 2 - self.stride = stride - self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - channels = x.shape[1] - x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate") - return F.conv1d(x, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels) - - -class MMAudioUpSample1d(nn.Module): - def __init__(self, ratio: int, kernel_size: int) -> None: - super().__init__() - self.ratio = ratio - self.stride = ratio - self.pad = kernel_size // ratio - 1 - self.pad_left = self.pad * self.stride + (kernel_size - self.stride) // 2 - self.pad_right = self.pad * self.stride + (kernel_size - self.stride + 1) // 2 - self.register_buffer("filter", kaiser_sinc_filter1d(0.5 / ratio, 0.6 / ratio, kernel_size)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - channels = x.shape[1] - x = F.pad(x, (self.pad, self.pad), mode="replicate") - x = self.ratio * F.conv_transpose1d( - x, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels - ) - return x[..., self.pad_left : -self.pad_right] - - -class MMAudioDownSample1d(nn.Module): - def __init__(self, ratio: int, kernel_size: int) -> None: - super().__init__() - self.lowpass = MMAudioLowPassFilter1d(0.5 / ratio, 0.6 / ratio, stride=ratio, kernel_size=kernel_size) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.lowpass(x) - - -class MMAudioActivation1d(nn.Module): - """Anti-aliased activation: 2x upsample, Snake-beta, 2x downsample.""" - - def __init__(self, channels: int, ratio: int = 2, kernel_size: int = 12) -> None: - super().__init__() - self.act = MMAudioSnakeBeta(channels) - self.upsample = MMAudioUpSample1d(ratio, kernel_size) - self.downsample = MMAudioDownSample1d(ratio, kernel_size) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.downsample(self.act(self.upsample(x))) - - -class MMAudioAMPBlock(nn.Module): - """Anti-aliased multi-periodicity block: dilated convolutions each followed by a dilation-1 convolution.""" - - def __init__(self, channels: int, kernel_size: int, dilations: tuple[int, ...]) -> None: - super().__init__() - self.convs1 = nn.ModuleList( - [ - nn.Conv1d( - channels, - channels, - kernel_size, - dilation=dilation, - padding=(kernel_size * dilation - dilation) // 2, - ) - for dilation in dilations - ] - ) - self.convs2 = nn.ModuleList( - [nn.Conv1d(channels, channels, kernel_size, padding=(kernel_size - 1) // 2) for _ in dilations] - ) - self.activations = nn.ModuleList([MMAudioActivation1d(channels) for _ in range(2 * len(dilations))]) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - activations_1, activations_2 = self.activations[::2], self.activations[1::2] - for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, activations_1, activations_2): - x = conv2(act2(conv1(act1(x)))) + x - return x - - -class MMAudioBigVGAN(nn.Module): - """BigVGAN-v2 vocoder turning a mel spectrogram into a waveform.""" - - def __init__( - self, - num_mels: int, - upsample_initial_channel: int, - upsample_rates: tuple[int, ...], - upsample_kernel_sizes: tuple[int, ...], - resblock_kernel_sizes: tuple[int, ...], - resblock_dilation_sizes: tuple[tuple[int, ...], ...], - ) -> None: - super().__init__() - self.num_kernels = len(resblock_kernel_sizes) - self.conv_pre = nn.Conv1d(num_mels, upsample_initial_channel, 7, padding=3) - - self.ups = nn.ModuleList() - self.resblocks = nn.ModuleList() - channels = upsample_initial_channel - for rate, kernel_size in zip(upsample_rates, upsample_kernel_sizes): - self.ups.append( - nn.ModuleList( - [nn.ConvTranspose1d(channels, channels // 2, kernel_size, rate, padding=(kernel_size - rate) // 2)] - ) - ) - channels //= 2 - for block_kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes): - self.resblocks.append(MMAudioAMPBlock(channels, block_kernel_size, tuple(dilations))) - - self.activation_post = MMAudioActivation1d(channels) - self.conv_post = nn.Conv1d(channels, 1, 7, padding=3, bias=False) - - def forward(self, mel: torch.Tensor) -> torch.Tensor: - hidden_states = self.conv_pre(mel) - for stage, up in enumerate(self.ups): - hidden_states = up[0](hidden_states) - blocks = self.resblocks[stage * self.num_kernels : (stage + 1) * self.num_kernels] - hidden_states = sum(block(hidden_states) for block in blocks) / self.num_kernels - hidden_states = self.conv_post(self.activation_post(hidden_states)) - return torch.clamp(hidden_states, min=-1.0, max=1.0) - - class MMAudioMelSpectrogram(nn.Module): """Log-mel front end of the encoder. The filterbank and window are buffers so they follow the model's device.""" @@ -493,10 +332,11 @@ def forward(self, waveform: torch.Tensor) -> torch.Tensor: class MMAudioVAE(ModelMixin, ConfigMixin): r""" Audio VAE of [`Kandinsky6TI2VAPipeline`]: a magnitude-preserving autoencoder over log-mel spectrograms (MMAudio, - https://arxiv.org/abs/2412.15322) followed by a BigVGAN-v2 vocoder. + https://arxiv.org/abs/2412.15322). - `encode` turns a waveform into a latent distribution; `decode` turns latents into a waveform through the mel - decoder and the vocoder. One latent frame covers `hop_length * 2` samples. + `encode` turns a waveform into a latent distribution; `decode` turns latents back into a mel spectrogram, which + the pipeline-local [`~pipelines.kandinsky6.MMAudioVocoder`] then turns into a waveform. One latent frame covers + `hop_length * 2` samples. Args: mel_bins (`int`, defaults to `128`): @@ -514,22 +354,13 @@ class MMAudioVAE(ModelMixin, ConfigMixin): n_fft (`int`, defaults to `2048`): FFT size of the mel front end. hop_length (`int`, defaults to `512`): - Hop length of the mel front end, also the vocoder's total upsampling factor. - vocoder_upsample_initial_channel (`int`, defaults to `1536`): - Width of the vocoder's first layer. - vocoder_upsample_rates (`tuple[int, ...]`, defaults to `(8, 4, 2, 2, 2, 2)`): - Upsampling factors of the vocoder stages. - vocoder_upsample_kernel_sizes (`tuple[int, ...]`, defaults to `(16, 8, 4, 4, 4, 4)`): - Transposed-convolution kernel sizes of the vocoder stages. - vocoder_resblock_kernel_sizes (`tuple[int, ...]`, defaults to `(3, 7, 11)`): - Kernel sizes of the vocoder's residual blocks. - vocoder_resblock_dilation_sizes (`tuple[tuple[int, ...], ...]`, defaults to `((1, 3, 5), (1, 3, 5), (1, 3, 5))`): - Dilations of the vocoder's residual blocks. + Hop length of the mel front end. Must match the total upsampling factor of the + [`~pipelines.kandinsky6.MMAudioVocoder`] this VAE is paired with. scaling_factor (`float`, defaults to `0.417`): Scale applied to the latents before they enter the diffusion transformer. """ - _no_split_modules = ["MMAudioResnetBlock1D", "MMAudioAttnBlock1D", "MMAudioAMPBlock"] + _no_split_modules = ["MMAudioResnetBlock1D", "MMAudioAttnBlock1D"] @register_to_config def __init__( @@ -542,29 +373,13 @@ def __init__( sample_rate: int = 44_100, n_fft: int = 2048, hop_length: int = 512, - vocoder_upsample_initial_channel: int = 1536, - vocoder_upsample_rates: tuple[int, ...] = (8, 4, 2, 2, 2, 2), - vocoder_upsample_kernel_sizes: tuple[int, ...] = (16, 8, 4, 4, 4, 4), - vocoder_resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11), - vocoder_resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)), scaling_factor: float = 0.417, ) -> None: super().__init__() - if math.prod(vocoder_upsample_rates) != hop_length: - raise ValueError("the vocoder upsampling rates must multiply to `hop_length`") - self.mel_converter = MMAudioMelSpectrogram(sample_rate, n_fft, mel_bins, hop_length) self.vae = MMAudioAutoencoder( mel_bins, latent_channels, hidden_channels, channel_multipliers, layers_per_block ) - self.vocoder = MMAudioBigVGAN( - num_mels=mel_bins, - upsample_initial_channel=vocoder_upsample_initial_channel, - upsample_rates=vocoder_upsample_rates, - upsample_kernel_sizes=vocoder_upsample_kernel_sizes, - resblock_kernel_sizes=vocoder_resblock_kernel_sizes, - resblock_dilation_sizes=vocoder_resblock_dilation_sizes, - ) # The encoder downsamples the mel frames once by 2. self.latent_hop_length = hop_length * 2 @@ -588,7 +403,7 @@ def encode(self, audio: torch.Tensor, return_dict: bool = True) -> AutoencoderKL @apply_forward_hook def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple: r""" - Decode latents into a waveform. + Decode latents into a mel spectrogram. Args: z (`torch.Tensor` of shape `(batch_size, latent_channels, num_latent_frames)`): @@ -597,12 +412,13 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | t Whether to return a [`~models.autoencoder_kl.DecoderOutput`] instead of a plain tuple. Returns: - The waveform of shape `(batch_size, 1, num_samples)` in `[-1, 1]`. + The mel spectrogram of shape `(batch_size, mel_bins, num_mel_frames)`, ready for + [`~pipelines.kandinsky6.MMAudioVocoder`]. """ - waveform = self.vocoder(self.vae.decode(z)) + mel = self.vae.decode(z) if not return_dict: - return (waveform,) - return DecoderOutput(sample=waveform) + return (mel,) + return DecoderOutput(sample=mel) def forward( self, @@ -613,7 +429,7 @@ def forward( ) -> DecoderOutput | tuple: posterior = self.encode(sample).latent_dist z = posterior.sample(generator=generator) if sample_posterior else posterior.mode() - waveform = self.decode(z).sample + mel = self.decode(z).sample if not return_dict: - return (waveform,) - return DecoderOutput(sample=waveform) + return (mel,) + return DecoderOutput(sample=mel) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 4f0433b0ca02..2fd426918859 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -454,6 +454,7 @@ "Kandinsky6SRPipelineOutput", "Kandinsky6TI2VAPipeline", "Kandinsky6TI2VAPipelineOutput", + "MMAudioVocoder", ] _import_structure["z_image"] = [ "ZImageControlNetInpaintPipeline", @@ -790,6 +791,7 @@ Kandinsky6SRPipelineOutput, Kandinsky6TI2VAPipeline, Kandinsky6TI2VAPipelineOutput, + MMAudioVocoder, ) from .krea2 import Krea2Pipeline from .latent_consistency_models import ( diff --git a/src/diffusers/pipelines/kandinsky6/__init__.py b/src/diffusers/pipelines/kandinsky6/__init__.py index 7e18b45cf44e..3ee290995ba8 100644 --- a/src/diffusers/pipelines/kandinsky6/__init__.py +++ b/src/diffusers/pipelines/kandinsky6/__init__.py @@ -22,6 +22,7 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: _import_structure["modeling_latent_upscaler"] = ["Kandinsky6SRLatentUpscalerBank"] + _import_structure["modeling_vocoder"] = ["MMAudioVocoder"] _import_structure["pipeline_kandinsky6_sr"] = ["Kandinsky6SRPipeline"] _import_structure["pipeline_kandinsky6_ti2va"] = ["Kandinsky6TI2VAPipeline"] _import_structure["pipeline_output"] = ["Kandinsky6SRPipelineOutput", "Kandinsky6TI2VAPipelineOutput"] @@ -34,6 +35,7 @@ from ...utils.dummy_torch_and_transformers_objects import * else: from .modeling_latent_upscaler import Kandinsky6SRLatentUpscalerBank + from .modeling_vocoder import MMAudioVocoder from .pipeline_kandinsky6_sr import Kandinsky6SRPipeline from .pipeline_kandinsky6_ti2va import Kandinsky6TI2VAPipeline from .pipeline_output import Kandinsky6SRPipelineOutput, Kandinsky6TI2VAPipelineOutput diff --git a/src/diffusers/pipelines/kandinsky6/modeling_vocoder.py b/src/diffusers/pipelines/kandinsky6/modeling_vocoder.py new file mode 100644 index 000000000000..f5470ba860e2 --- /dev/null +++ b/src/diffusers/pipelines/kandinsky6/modeling_vocoder.py @@ -0,0 +1,228 @@ +# 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. + +"""BigVGAN vocoder that turns the mel spectrograms decoded by [`MMAudioVAE`] into waveforms.""" + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...models.autoencoders.vae import DecoderOutput +from ...models.modeling_utils import ModelMixin + + +def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: + """Kaiser-windowed sinc low-pass filter of shape `(1, 1, kernel_size)` normalized to unit sum.""" + even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + + attenuation = 2.285 * (half_size - 1) * math.pi * 4 * half_width + 7.95 + if attenuation > 50.0: + beta = 0.1102 * (attenuation - 8.7) + elif attenuation >= 21.0: + beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + + time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size + filter_ = 2 * cutoff * window * torch.sinc(2 * cutoff * time) + filter_ = filter_ / filter_.sum() + return filter_.view(1, 1, kernel_size) + + +class MMAudioSnakeBeta(nn.Module): + """`x + 1/b * sin^2(a * x)` with per-channel log-scale frequency `a` and magnitude `b`.""" + + def __init__(self, channels: int) -> None: + super().__init__() + self.alpha = nn.Parameter(torch.zeros(channels)) + self.beta = nn.Parameter(torch.zeros(channels)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + alpha = torch.exp(self.alpha)[None, :, None] + beta = torch.exp(self.beta)[None, :, None] + return x + (1.0 / (beta + 1e-9)) * torch.sin(x * alpha).pow(2) + + +class MMAudioLowPassFilter1d(nn.Module): + def __init__(self, cutoff: float, half_width: float, stride: int, kernel_size: int) -> None: + super().__init__() + even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + channels = x.shape[1] + x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate") + return F.conv1d(x, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels) + + +class MMAudioUpSample1d(nn.Module): + def __init__(self, ratio: int, kernel_size: int) -> None: + super().__init__() + self.ratio = ratio + self.stride = ratio + self.pad = kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (kernel_size - self.stride + 1) // 2 + self.register_buffer("filter", kaiser_sinc_filter1d(0.5 / ratio, 0.6 / ratio, kernel_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + channels = x.shape[1] + x = F.pad(x, (self.pad, self.pad), mode="replicate") + x = self.ratio * F.conv_transpose1d( + x, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels + ) + return x[..., self.pad_left : -self.pad_right] + + +class MMAudioDownSample1d(nn.Module): + def __init__(self, ratio: int, kernel_size: int) -> None: + super().__init__() + self.lowpass = MMAudioLowPassFilter1d(0.5 / ratio, 0.6 / ratio, stride=ratio, kernel_size=kernel_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lowpass(x) + + +class MMAudioActivation1d(nn.Module): + """Anti-aliased activation: 2x upsample, Snake-beta, 2x downsample.""" + + def __init__(self, channels: int, ratio: int = 2, kernel_size: int = 12) -> None: + super().__init__() + self.act = MMAudioSnakeBeta(channels) + self.upsample = MMAudioUpSample1d(ratio, kernel_size) + self.downsample = MMAudioDownSample1d(ratio, kernel_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.downsample(self.act(self.upsample(x))) + + +class MMAudioAMPBlock(nn.Module): + """Anti-aliased multi-periodicity block: dilated convolutions each followed by a dilation-1 convolution.""" + + def __init__(self, channels: int, kernel_size: int, dilations: tuple[int, ...]) -> None: + super().__init__() + self.convs1 = nn.ModuleList( + [ + nn.Conv1d( + channels, + channels, + kernel_size, + dilation=dilation, + padding=(kernel_size * dilation - dilation) // 2, + ) + for dilation in dilations + ] + ) + self.convs2 = nn.ModuleList( + [nn.Conv1d(channels, channels, kernel_size, padding=(kernel_size - 1) // 2) for _ in dilations] + ) + self.activations = nn.ModuleList([MMAudioActivation1d(channels) for _ in range(2 * len(dilations))]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + activations_1, activations_2 = self.activations[::2], self.activations[1::2] + for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, activations_1, activations_2): + x = conv2(act2(conv1(act1(x)))) + x + return x + + +class MMAudioVocoder(ModelMixin, ConfigMixin): + r""" + BigVGAN-v2 vocoder (https://github.com/NVIDIA/BigVGAN, MIT license) with the anti-aliased Snake activations of + https://github.com/junjun3518/alias-free-torch (Apache 2.0), turning the mel spectrograms [`MMAudioVAE`] decodes + into waveforms for [`Kandinsky6TI2VAPipeline`]. + + Args: + num_mels (`int`, defaults to `128`): + Number of mel bins of the input spectrogram. Must match the paired [`MMAudioVAE`]'s `mel_bins`. + upsample_initial_channel (`int`, defaults to `1536`): + Width of the first layer. + upsample_rates (`tuple[int, ...]`, defaults to `(8, 4, 2, 2, 2, 2)`): + Upsampling factors of the vocoder stages. Their product is the total upsampling factor and must match the + paired [`MMAudioVAE`]'s `hop_length`. + upsample_kernel_sizes (`tuple[int, ...]`, defaults to `(16, 8, 4, 4, 4, 4)`): + Transposed-convolution kernel sizes of the vocoder stages. + resblock_kernel_sizes (`tuple[int, ...]`, defaults to `(3, 7, 11)`): + Kernel sizes of the residual blocks. + resblock_dilation_sizes (`tuple[tuple[int, ...], ...]`, defaults to `((1, 3, 5), (1, 3, 5), (1, 3, 5))`): + Dilations of the residual blocks. + """ + + _no_split_modules = ["MMAudioAMPBlock"] + + @register_to_config + def __init__( + self, + num_mels: int = 128, + upsample_initial_channel: int = 1536, + upsample_rates: tuple[int, ...] = (8, 4, 2, 2, 2, 2), + upsample_kernel_sizes: tuple[int, ...] = (16, 8, 4, 4, 4, 4), + resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11), + resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)), + ) -> None: + super().__init__() + if len(upsample_rates) != len(upsample_kernel_sizes): + raise ValueError("`upsample_rates` and `upsample_kernel_sizes` must have the same length") + if len(resblock_kernel_sizes) != len(resblock_dilation_sizes): + raise ValueError("`resblock_kernel_sizes` and `resblock_dilation_sizes` must have the same length") + + self.num_kernels = len(resblock_kernel_sizes) + self.conv_pre = nn.Conv1d(num_mels, upsample_initial_channel, 7, padding=3) + + self.ups = nn.ModuleList() + self.resblocks = nn.ModuleList() + channels = upsample_initial_channel + for rate, kernel_size in zip(upsample_rates, upsample_kernel_sizes): + self.ups.append( + nn.ModuleList( + [nn.ConvTranspose1d(channels, channels // 2, kernel_size, rate, padding=(kernel_size - rate) // 2)] + ) + ) + channels //= 2 + for block_kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes): + self.resblocks.append(MMAudioAMPBlock(channels, block_kernel_size, tuple(dilations))) + + self.activation_post = MMAudioActivation1d(channels) + self.conv_post = nn.Conv1d(channels, 1, 7, padding=3, bias=False) + + def forward(self, mel: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple: + r""" + Args: + mel (`torch.Tensor` of shape `(batch_size, num_mels, num_mel_frames)`): + Mel spectrogram, as decoded by [`MMAudioVAE`]. + return_dict (`bool`, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.DecoderOutput`] instead of a plain tuple. + + Returns: + The waveform of shape `(batch_size, 1, num_samples)` in `[-1, 1]`. + """ + hidden_states = self.conv_pre(mel) + for stage, up in enumerate(self.ups): + hidden_states = up[0](hidden_states) + blocks = self.resblocks[stage * self.num_kernels : (stage + 1) * self.num_kernels] + hidden_states = sum(block(hidden_states) for block in blocks) / self.num_kernels + hidden_states = self.conv_post(self.activation_post(hidden_states)) + waveform = torch.clamp(hidden_states, min=-1.0, max=1.0) + if not return_dict: + return (waveform,) + return DecoderOutput(sample=waveform) diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index a24c5e21dccf..f8c8eb48859f 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -28,6 +28,7 @@ from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline +from .modeling_vocoder import MMAudioVocoder from .pipeline_output import Kandinsky6TI2VAPipelineOutput @@ -202,11 +203,15 @@ class Kandinsky6TI2VAPipeline(DiffusionPipeline): Scheduler used with `transformer` to denoise the latents. Distilled checkpoints ship with a [`PiflowScheduler`] and must be run with `guidance_scale=1.0`. audio_vae ([`MMAudioVAE`], *optional*): - Audio VAE used to decode the generated audio latents. Only needed when `sample_audio=True`. + Audio VAE used to decode the generated audio latents into a mel spectrogram. Only needed when + `sample_audio=True`. + vocoder ([`MMAudioVocoder`], *optional*): + Vocoder used to turn the mel spectrogram `audio_vae` decodes into a waveform. Only needed when + `sample_audio=True`. """ - model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae->audio_vae" - _optional_components = ["audio_vae"] + model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae->audio_vae->vocoder" + _optional_components = ["audio_vae", "vocoder"] _callback_tensor_inputs = ["latents", "audio_latents", "prompt_embeds", "negative_prompt_embeds"] _DEFAULT_NEGATIVE_PROMPT = ( "Static, 2D cartoon, cartoon, 2d animation, paintings, images, " @@ -223,6 +228,7 @@ def __init__( tokenizer_2: CLIPTokenizer, scheduler: FlowMatchEulerDiscreteScheduler | PiflowScheduler, audio_vae: MMAudioVAE | None = None, + vocoder: MMAudioVocoder | None = None, ) -> None: super().__init__() self.register_modules( @@ -234,6 +240,7 @@ def __init__( tokenizer_2=tokenizer_2, scheduler=scheduler, audio_vae=audio_vae, + vocoder=vocoder, ) self.vae_scale_factor_spatial = ( @@ -491,8 +498,8 @@ def check_inputs( if (negative_prompt_embeds is None) != (negative_pooled_prompt_embeds is None): raise ValueError("`negative_prompt_embeds` and `negative_pooled_prompt_embeds` must be provided together.") - if sample_audio and getattr(self, "audio_vae", None) is None: - raise ValueError("`sample_audio=True` requires an `audio_vae`.") + if sample_audio and (getattr(self, "audio_vae", None) is None or getattr(self, "vocoder", None) is None): + raise ValueError("`sample_audio=True` requires an `audio_vae` and a `vocoder`.") if visual_cond_scheme not in ("pretrain", "i2v", "tail_cond_first_frame"): raise ValueError( f"`visual_cond_scheme` must be one of 'pretrain', 'i2v' or 'tail_cond_first_frame' but is {visual_cond_scheme!r}." @@ -674,8 +681,8 @@ def __call__( negative_pooled_prompt_embeds (`torch.Tensor`, *optional*): Pre-generated negative CLIP pooled text embeddings. sample_audio (`bool`, *optional*): - Whether to generate synchronized audio. Defaults to `True` when the pipeline has an `audio_vae` and - `False` otherwise; `True` requires an `audio_vae`. + Whether to generate synchronized audio. Defaults to `True` when the pipeline has an `audio_vae` and a + `vocoder`, `False` otherwise; `True` requires both. expand_prompts (`bool`, defaults to `False`): Whether to rewrite the prompts with [`~Kandinsky6TI2VAPipeline.expand_prompts`] before encoding. visual_cond_scheme (`str`, *optional*): @@ -701,7 +708,7 @@ def __call__( """ # 1. Check inputs. Raise error if not correct if sample_audio is None: - sample_audio = getattr(self, "audio_vae", None) is not None + sample_audio = getattr(self, "audio_vae", None) is not None and getattr(self, "vocoder", None) is not None if visual_cond_scheme is None: visual_cond_scheme = "tail_cond_first_frame" if image is not None else "pretrain" self.check_inputs( @@ -919,7 +926,8 @@ def __call__( audio = None if sample_audio: audio_latents = audio_latents.to(self.audio_vae.dtype) / self.audio_vae.config.scaling_factor - audio = self.audio_vae.decode(audio_latents, return_dict=False)[0][:, 0].float() + mel = self.audio_vae.decode(audio_latents, return_dict=False)[0] + audio = self.vocoder(mel.to(self.vocoder.dtype), return_dict=False)[0][:, 0].float() if output_type == "np": audio = audio.cpu().numpy() diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 042f0b3d6cdf..7a173941256a 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -3482,6 +3482,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class MMAudioVocoder(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 MochiPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/autoencoders/test_models_autoencoder_mmaudio.py b/tests/models/autoencoders/test_models_autoencoder_mmaudio.py index 4747cb2196b2..13481f2b9407 100644 --- a/tests/models/autoencoders/test_models_autoencoder_mmaudio.py +++ b/tests/models/autoencoders/test_models_autoencoder_mmaudio.py @@ -61,11 +61,6 @@ def get_init_dict(self) -> dict: "sample_rate": 64, "n_fft": 16, "hop_length": 4, - "vocoder_upsample_initial_channel": 8, - "vocoder_upsample_rates": (2, 2), - "vocoder_upsample_kernel_sizes": (4, 4), - "vocoder_resblock_kernel_sizes": (3,), - "vocoder_resblock_dilation_sizes": ((1, 3),), } def get_dummy_inputs(self) -> dict: @@ -78,7 +73,9 @@ def input_shape(self) -> tuple[int, ...]: @property def output_shape(self) -> tuple[int, ...]: - return (1, 256) + # `decode`/`forward` return a mel spectrogram (`mel_bins`, num_mel_frames); the waveform is produced by the + # separate, pipeline-local `MMAudioVocoder`. + return (8, 64) class TestMMAudioVAEModel(MMAudioVAETesterConfig, ModelTesterMixin): diff --git a/tests/pipelines/kandinsky6/test_models_vocoder.py b/tests/pipelines/kandinsky6/test_models_vocoder.py new file mode 100644 index 000000000000..d54575bae7a2 --- /dev/null +++ b/tests/pipelines/kandinsky6/test_models_vocoder.py @@ -0,0 +1,86 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from diffusers.pipelines.kandinsky6 import MMAudioVocoder +from diffusers.utils.torch_utils import randn_tensor + +from ...models.testing_utils import ( + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, +) +from ...testing_utils import enable_full_determinism, torch_device + + +enable_full_determinism() + + +class MMAudioVocoderTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return MMAudioVocoder + + @property + def pretrained_model_name_or_path(self): + return "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "vocoder"} + + @property + def main_input_name(self) -> str: + return "mel" + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict: + return { + "num_mels": 8, + "upsample_initial_channel": 8, + "upsample_rates": (2, 2), + "upsample_kernel_sizes": (4, 4), + "resblock_kernel_sizes": (3,), + "resblock_dilation_sizes": ((1, 3),), + } + + def get_dummy_inputs(self) -> dict: + return {"mel": randn_tensor((2, 8, 16), generator=self.generator, device=torch_device)} + + @property + def input_shape(self) -> tuple[int, ...]: + return (8, 16) + + @property + def output_shape(self) -> tuple[int, ...]: + # `upsample_rates` multiply to 4, so 16 mel frames decode to 64 samples. + return (1, 64) + + +class TestMMAudioVocoderModel(MMAudioVocoderTesterConfig, ModelTesterMixin): + pass + + +class TestMMAudioVocoderMemory(MMAudioVocoderTesterConfig, MemoryTesterMixin): + pass + + +class TestMMAudioVocoderTorchCompile(MMAudioVocoderTesterConfig, TorchCompileTesterMixin): + pass diff --git a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py index 5ebaa1d4613d..316c36fa7755 100644 --- a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py +++ b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py @@ -30,6 +30,7 @@ Kandinsky6TI2VAPipeline, Kandinsky6Transformer3DModel, MMAudioVAE, + MMAudioVocoder, ) from ...testing_utils import torch_device @@ -94,11 +95,17 @@ def get_dummy_components(self): sample_rate=64, n_fft=16, hop_length=4, - vocoder_upsample_initial_channel=8, - vocoder_upsample_rates=(2, 2), - vocoder_upsample_kernel_sizes=(4, 4), - vocoder_resblock_kernel_sizes=(3,), - vocoder_resblock_dilation_sizes=((1, 3),), + ) + + torch.manual_seed(0) + # `upsample_rates` must multiply to `audio_vae.config.hop_length` (4 = 2 * 2). + vocoder = MMAudioVocoder( + num_mels=8, + upsample_initial_channel=8, + upsample_rates=(2, 2), + upsample_kernel_sizes=(4, 4), + resblock_kernel_sizes=(3,), + resblock_dilation_sizes=((1, 3),), ) scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0) @@ -182,6 +189,7 @@ def get_dummy_components(self): "tokenizer_2": tokenizer_2, "scheduler": scheduler, "audio_vae": audio_vae, + "vocoder": vocoder, } def get_dummy_inputs(self): From 525281ee795ead451a37fc1bcf6fd5adf2939222 Mon Sep 17 00:00:00 2001 From: leffff Date: Tue, 29 Sep 2026 22:49:03 +0300 Subject: [PATCH 28/62] fix naming --- .../models/transformers/transformer_kandinsky6.py | 12 ++---------- .../models/transformers/transformer_kandinsky6_sr.py | 2 +- 2 files changed, 3 insertions(+), 11 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 6ce3b696749f..2567eb2b6ce9 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -136,7 +136,7 @@ def nabla_block_mask( ) -class RoPE1D(nn.Module): +class Kandinsky6RoPE1D(nn.Module): """1-D Rotary Position Embedding — used for text and audio sequences.""" def __init__( @@ -161,7 +161,7 @@ def forward(self, pos: Tensor) -> Tensor: return rope.view(*rope.shape[:-1], 2, 2).unsqueeze(-4) -class RoPE3D(nn.Module): +class Kandinsky6RoPE3D(nn.Module): """3-D Rotary Position Embedding — used for video spatial-temporal tokens (T, H, W).""" def __init__( @@ -923,14 +923,6 @@ def forward( return vis, aud -class Kandinsky6RoPE1D(RoPE1D): - """Diffusers-exported name for the K6 one-dimensional RoPE.""" - - -class Kandinsky6RoPE3D(RoPE3D): - """Diffusers-exported name for the K6 three-dimensional RoPE.""" - - class Kandinsky6Transformer3DModel( Kandinsky6MagCacheMixin, ModelMixin, diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index ab484947c099..5c5bdd9997f7 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -331,7 +331,7 @@ def forward(self, visual_embed: Tensor, time_embed: Tensor) -> Tensor: ) -# Copied from diffusers.models.transformers.transformer_kandinsky6.RoPE3D with RoPE3D->Kandinsky6SRRoPE3D +# Copied from diffusers.models.transformers.transformer_kandinsky6.Kandinsky6RoPE3D with Kandinsky6RoPE3D->Kandinsky6SRRoPE3D class Kandinsky6SRRoPE3D(nn.Module): """3-D Rotary Position Embedding — used for video spatial-temporal tokens (T, H, W).""" From 695caecf646c5dbe6d7028d0a0101e28e6747d13 Mon Sep 17 00:00:00 2001 From: leffff Date: Wed, 30 Sep 2026 16:42:23 +0300 Subject: [PATCH 29/62] refactor checkpointing, magcache, transformer, sr vae for K6 --- docs/source/en/api/pipelines/kandinsky6.md | 17 + .../autoencoders/autoencoder_kandinsky6_sr.py | 192 ++++++++--- .../autoencoders/autoencoder_mmaudio.py | 6 +- .../transformers/transformer_kandinsky6.py | 312 +----------------- .../transformers/transformer_kandinsky6_sr.py | 10 +- 5 files changed, 187 insertions(+), 350 deletions(-) diff --git a/docs/source/en/api/pipelines/kandinsky6.md b/docs/source/en/api/pipelines/kandinsky6.md index aff11abc55ef..6a3caa06321f 100644 --- a/docs/source/en/api/pipelines/kandinsky6.md +++ b/docs/source/en/api/pipelines/kandinsky6.md @@ -77,6 +77,23 @@ sr_pipe.enable_model_cpu_offload() upscaled = sr_pipe(video=output.frames[0], resolution_scale=2.25, num_inference_steps=2).frames[0] ``` +## Memory optimization + +Refer to the [Reduce memory usage](../../optimization/memory) guide for the general set of techniques. Both +[`Kandinsky6TI2VAPipeline`] and [`Kandinsky6SRPipeline`] support [model offloading](../../optimization/memory#model-offloading) +(used above) and, for a smaller footprint at the cost of speed, [sequential CPU offloading](../../optimization/memory#cpu-offloading): + +```python +pipe.enable_sequential_cpu_offload() +``` + +[`Kandinsky6TI2VAPipeline`]'s video VAE also supports [tiled decoding](../../optimization/memory#vae-tiling) for high +resolutions or long videos: + +```python +pipe.vae.enable_tiling() +``` + ## Kandinsky6TI2VAPipeline [[autodoc]] Kandinsky6TI2VAPipeline diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index 5ab7438821fd..f933a91ddef4 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -36,7 +36,7 @@ CONV_CHUNK_ELEMENTS = 2 * 10**9 -class Kandinsky6SRConv3d(nn.Conv3d): +class Kandinsky6SRSafeConv3d(nn.Conv3d): """`Conv3d` that splits very large inputs along the time axis and convolves them chunk by chunk. Each chunk is padded with the last `kernel_size - 1` frames of the previous chunk before convolving, so the result @@ -73,6 +73,12 @@ class Kandinsky6SRCausalConv3d(nn.Module): Height and width are zero-padded symmetrically. Along time the first segment is padded by repeating its first frame `kernel_size - 1` times; later segments are padded with the frames the previous segment left behind in `cache["padding"]`, so a video processed segment by segment matches a single pass. + + This caching is not an optional performance knob: it is what lets `Kandinsky6SRVAE.encode`/`decode` process an + arbitrarily long video in bounded-memory segments (see `SEGMENT_FRAMES`) while reproducing the exact output of a + single non-causal pass. Without it, each segment would need the raw frames the previous segment already consumed in + order to rebuild correct padding, which means keeping the whole video in memory anyway and defeats the reason + segmentation exists. """ def __init__( @@ -96,13 +102,13 @@ def __init__( self.time_pad = time_kernel_size - 1 self.time_kernel_size = time_kernel_size self.time_stride = stride[0] - self.conv = Kandinsky6SRConv3d(in_channels, out_channels, kernel_size, stride=stride) + self.conv = Kandinsky6SRSafeConv3d(in_channels, out_channels, kernel_size, stride=stride) @staticmethod def make_cache() -> dict: return {"padding": None} - def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, cache: dict) -> tuple[torch.Tensor, dict]: batch_size, _, num_frames, height, width = hidden_states.shape hidden_states = F.pad(hidden_states, (self.width_pad, self.width_pad, self.height_pad, self.height_pad)) @@ -129,19 +135,28 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: if offset_out < output_frames: output[:, :, offset_out:] = self.conv(hidden_states[:, :, offset_in:]) - # Carry the frames the next segment's first window still needs. + # The frames the next segment's first window still needs. Returned rather than written into `cache` in + # place, so this stays correct under `torch.utils.checkpoint`'s backward-time recompute (see + # `Kandinsky6SRResnetBlock3D`): the recompute reads the same, never-mutated `cache` argument as the + # original call. pad_offset = ( offset_in + stride * math.trunc((num_frames - offset_in - self.time_kernel_size) / stride) + stride ) if pad_offset < 0: - cache["padding"] = torch.cat([padding[:, :, pad_offset:], hidden_states], dim=2) + new_padding = torch.cat([padding[:, :, pad_offset:], hidden_states], dim=2) else: - cache["padding"] = hidden_states[:, :, pad_offset:].clone() - return output + new_padding = hidden_states[:, :, pad_offset:].clone() + return output, {"padding": new_padding} class Kandinsky6SRRMSNorm(nn.Module): - """RMS normalization over the channel axis of a `(B, C, T, H, W)` tensor, computed in float32.""" + """RMS normalization over the channel axis of a `(B, C, T, H, W)` tensor, computed in float32. + + Same idea as `WanRMS_norm` (`autoencoder_kl_wan.py`) / `QwenImageRMS_norm` (`autoencoder_kl_qwenimage.py`), but not + a verbatim copy of either, so it is not marked `# Copied from`: it is hardcoded to the channel-first 5D layout used + throughout this VAE instead of taking `channel_first`/`images` flags, always upcasts to float32 rather than only + for fp16/bf16/fp8 inputs, and has no learnable bias term. + """ def __init__(self, num_channels: int) -> None: super().__init__() @@ -154,8 +169,14 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class Kandinsky6SRSpatialNorm3D(nn.Module): - """RMS norm modulated by the latent `zq`, the same conditioning other video VAEs in the library use to inject the - latent back into the decoder, adapted to causal segment-by-segment decoding. + """Spatial normalization conditioned on the latent `zq` (https://huggingface.co/papers/2209.09002), used here so + the decoder can inject the latent back in at every resnet block. Structurally this plays the same role as + `SpatialNorm` (`attention_processor.py`) / `CogVideoXSpatialNorm3D` (`autoencoder_kl_cogvideox.py`) — normalize the + features, then scale and shift by convolutions of the upsampled `zq`, threading a `cache` argument through + `forward` functionally (read, never written in place) the same way `CogVideoXSpatialNorm3D` threads its + `conv_cache` — but it is not a `# Copied from` of either: it normalizes with `Kandinsky6SRRMSNorm` instead of + `GroupNorm` (matching the plain-RMSNorm blocks elsewhere in this VAE), and interpolates `zq` in channel chunks to + bound peak memory. `zq` is nearest-upsampled to the feature grid. In the first segment the first frame is upsampled separately, because the temporal upsampler turns `T + 1` latent frames into `2T + 1` pixel frames. @@ -164,14 +185,14 @@ class Kandinsky6SRSpatialNorm3D(nn.Module): def __init__(self, num_channels: int, zq_channels: int) -> None: super().__init__() self.norm_layer = Kandinsky6SRRMSNorm(num_channels) - self.conv_y = Kandinsky6SRConv3d(zq_channels, num_channels, kernel_size=1) - self.conv_b = Kandinsky6SRConv3d(zq_channels, num_channels, kernel_size=1) + self.conv_y = Kandinsky6SRSafeConv3d(zq_channels, num_channels, kernel_size=1) + self.conv_b = Kandinsky6SRSafeConv3d(zq_channels, num_channels, kernel_size=1) @staticmethod def make_cache() -> dict: return {"is_first_segment": True} - def forward(self, hidden_states: torch.Tensor, zq: torch.Tensor, cache: dict) -> torch.Tensor: + def forward(self, hidden_states: torch.Tensor, zq: torch.Tensor, cache: dict) -> tuple[torch.Tensor, dict]: if cache["is_first_segment"]: zq_first = F.interpolate(zq[:, :, :1], size=hidden_states[:, :, :1].shape[-3:], mode="nearest") if zq.size(2) > 1: @@ -186,7 +207,6 @@ def forward(self, hidden_states: torch.Tensor, zq: torch.Tensor, cache: dict) -> zq = torch.cat([zq_first, zq_rest], dim=2) else: zq = zq_first - cache["is_first_segment"] = False else: zq = torch.cat( [ @@ -195,11 +215,20 @@ def forward(self, hidden_states: torch.Tensor, zq: torch.Tensor, cache: dict) -> ], dim=1, ) - return self.norm_layer(hidden_states) * self.conv_y(zq) + self.conv_b(zq) + output = self.norm_layer(hidden_states) * self.conv_y(zq) + self.conv_b(zq) + return output, {"is_first_segment": False} class Kandinsky6SRResnetBlock3D(nn.Module): - """Causal residual block; decoder blocks modulate their norms with the latent `zq`.""" + """Causal residual block; decoder blocks modulate their norms with the latent `zq`. + + `forward` never mutates the `cache` it is given: `conv1`/`conv2` (`Kandinsky6SRCausalConv3d`) and, for decoder + blocks, `norm1`/`norm2` (`Kandinsky6SRSpatialNorm3D`) each return a new cache value instead of writing into the one + they were passed, the same functional-cache convention `autoencoder_kl_cogvideox.py` uses. That is what makes it + safe to wrap this block in `torch.utils.checkpoint` (see `Kandinsky6SREncoder3D`/`Kandinsky6SRDecoder3D`): the + backward-time recompute reads the same, never-mutated `cache` argument the original call did, so it retraces the + same branch and produces the same activations and gradients. + """ def __init__(self, in_channels: int, out_channels: int, zq_channels: int | None = None) -> None: super().__init__() @@ -214,7 +243,7 @@ def __init__(self, in_channels: int, out_channels: int, zq_channels: int | None self.conv1 = Kandinsky6SRCausalConv3d(in_channels, out_channels, kernel_size=3) self.conv2 = Kandinsky6SRCausalConv3d(out_channels, out_channels, kernel_size=3) if in_channels != out_channels: - self.nin_shortcut = Kandinsky6SRConv3d(in_channels, out_channels, kernel_size=1) + self.nin_shortcut = Kandinsky6SRSafeConv3d(in_channels, out_channels, kernel_size=1) def make_cache(self) -> dict: cache = {"conv1": self.conv1.make_cache(), "conv2": self.conv2.make_cache()} @@ -223,26 +252,43 @@ def make_cache(self) -> dict: cache["norm2"] = self.norm2.make_cache() return cache - def forward(self, hidden_states: torch.Tensor, cache: dict, zq: torch.Tensor | None = None) -> torch.Tensor: + def forward( + self, hidden_states: torch.Tensor, cache: dict, zq: torch.Tensor | None = None + ) -> tuple[torch.Tensor, dict]: residual = hidden_states - hidden_states = self.norm1(hidden_states) if zq is None else self.norm1(hidden_states, zq, cache["norm1"]) - hidden_states = self.conv1(F.silu(hidden_states), cache["conv1"]) - hidden_states = self.norm2(hidden_states) if zq is None else self.norm2(hidden_states, zq, cache["norm2"]) - hidden_states = self.conv2(F.silu(hidden_states), cache["conv2"]) + new_cache = {} + if zq is None: + hidden_states = self.norm1(hidden_states) + else: + hidden_states, new_cache["norm1"] = self.norm1(hidden_states, zq, cache["norm1"]) + hidden_states, new_cache["conv1"] = self.conv1(F.silu(hidden_states), cache["conv1"]) + if zq is None: + hidden_states = self.norm2(hidden_states) + else: + hidden_states, new_cache["norm2"] = self.norm2(hidden_states, zq, cache["norm2"]) + hidden_states, new_cache["conv2"] = self.conv2(F.silu(hidden_states), cache["conv2"]) if self.in_channels != self.out_channels: residual = self.nin_shortcut(residual) - return residual + hidden_states + return residual + hidden_states, new_cache class Kandinsky6SRDownsample(nn.Module): - """Spatial 2x downsample (strided conv plus pixel-unshuffle average) with an optional causal temporal 2x.""" + """Spatial 2x downsample (strided conv plus pixel-unshuffle average) with an optional causal temporal 2x. + + Implemented from scratch rather than reusing `KVAECachedPXSDownsample` (`autoencoder_kl_kvae_video.py`), the + closest existing causal, stateful down/upsample block in the library: that class is a private helper of a different + Kandinsky VAE (not a shared, importable utility) and uses its own cache convention — a `list` of per-conv dicts + mutated in place — rather than this file's dict-keyed cache that `forward` returns functionally instead of mutating + (see `Kandinsky6SRCausalConv3d`), so reusing it as-is would mean threading two incompatible cache conventions + through the same model. See `Kandinsky6SRUpsample` for the upsampling counterpart. + """ def __init__(self, in_channels: int, compress_time: bool) -> None: super().__init__() out_channels = 2 * in_channels self.compress_time = compress_time self.unshuffle = nn.PixelUnshuffle(2) - self.spatial_conv = Kandinsky6SRConv3d( + self.spatial_conv = Kandinsky6SRSafeConv3d( in_channels, out_channels, kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1) ) if compress_time: @@ -252,7 +298,7 @@ def __init__(self, in_channels: int, compress_time: bool) -> None: Kandinsky6SRCausalConv3d(out_channels, out_channels, kernel_size=(2, 1, 1), stride=(2, 1, 1)), ] ) - self.linear = Kandinsky6SRConv3d(out_channels, out_channels, kernel_size=1) + self.linear = Kandinsky6SRSafeConv3d(out_channels, out_channels, kernel_size=1) def make_cache(self) -> dict: if not self.compress_time: @@ -284,21 +330,26 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: pooled = F.avg_pool1d(sequence, kernel_size=2, stride=2) pooled = pooled.reshape(batch_size, height, width, channels, -1).permute(0, 3, 4, 1, 2) - conv_out = self.temporal_conv[0](hidden_states, cache["temporal_conv"][0]) - conv_out = self.temporal_conv[1](conv_out, cache["temporal_conv"][1]) + conv_out, cache["temporal_conv"][0] = self.temporal_conv[0](hidden_states, cache["temporal_conv"][0]) + conv_out, cache["temporal_conv"][1] = self.temporal_conv[1](conv_out, cache["temporal_conv"][1]) return self.linear(conv_out + pooled) class Kandinsky6SRUpsample(nn.Module): - """Spatial 2x nearest upsample with a convolutional residual, preceded by an optional causal temporal 2x.""" + """Spatial 2x nearest upsample with a convolutional residual, preceded by an optional causal temporal 2x. + + Bespoke for the same reason as `Kandinsky6SRDownsample`: the closest existing precedent, `KVAECachedPXSUpsample` + (`autoencoder_kl_kvae_video.py`), is a private helper of a different Kandinsky VAE with its own, incompatible cache + convention, not a shared utility this file can import and reuse. + """ def __init__(self, channels: int, compress_time: bool) -> None: super().__init__() self.compress_time = compress_time - self.spatial_conv = Kandinsky6SRConv3d(channels, channels, kernel_size=(1, 3, 3), padding=(0, 1, 1)) + self.spatial_conv = Kandinsky6SRSafeConv3d(channels, channels, kernel_size=(1, 3, 3), padding=(0, 1, 1)) if compress_time: self.temporal_conv = Kandinsky6SRCausalConv3d(channels, channels, kernel_size=(3, 1, 1)) - self.linear = Kandinsky6SRConv3d(channels, channels, kernel_size=1) + self.linear = Kandinsky6SRSafeConv3d(channels, channels, kernel_size=1) def make_cache(self) -> dict: return {"temporal_conv": self.temporal_conv.make_cache()} if self.compress_time else {} @@ -310,7 +361,8 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: repeated = hidden_states.repeat_interleave(2, dim=2) if cache["temporal_conv"]["padding"] is None: repeated = repeated[:, :, 1:] - hidden_states = self.temporal_conv(repeated, cache["temporal_conv"]) + repeated + conv_out, cache["temporal_conv"] = self.temporal_conv(repeated, cache["temporal_conv"]) + hidden_states = conv_out + repeated hidden_states = F.interpolate(hidden_states, scale_factor=(1, 2, 2), mode="nearest") hidden_states = hidden_states + self.spatial_conv(hidden_states) @@ -318,6 +370,8 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: class Kandinsky6SREncoder3D(nn.Module): + _supports_gradient_checkpointing = True + def __init__( self, in_channels: int, @@ -354,6 +408,8 @@ def __init__( self.norm_out = Kandinsky6SRRMSNorm(block_in) self.conv_out = Kandinsky6SRCausalConv3d(block_in, 2 * latent_channels, kernel_size=3) + self.gradient_checkpointing = False + def make_cache(self) -> dict: return { "conv_in": self.conv_in.make_cache(), @@ -370,19 +426,38 @@ def make_cache(self) -> dict: } def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: - hidden_states = self.conv_in(hidden_states, cache["conv_in"]) + hidden_states, cache["conv_in"] = self.conv_in(hidden_states, cache["conv_in"]) for down, level_cache in zip(self.down, cache["down"]): - for block, block_cache in zip(down.block, level_cache["block"]): - hidden_states = block(hidden_states, block_cache) + for i, block in enumerate(down.block): + block_cache = level_cache["block"][i] + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states, level_cache["block"][i] = self._gradient_checkpointing_func( + block, hidden_states, block_cache + ) + else: + hidden_states, level_cache["block"][i] = block(hidden_states, block_cache) if hasattr(down, "downsample"): hidden_states = down.downsample(hidden_states, level_cache["downsample"]) - hidden_states = self.mid.block_1(hidden_states, cache["mid_1"]) - hidden_states = self.mid.block_2(hidden_states, cache["mid_2"]) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states, cache["mid_1"] = self._gradient_checkpointing_func( + self.mid.block_1, hidden_states, cache["mid_1"] + ) + hidden_states, cache["mid_2"] = self._gradient_checkpointing_func( + self.mid.block_2, hidden_states, cache["mid_2"] + ) + else: + hidden_states, cache["mid_1"] = self.mid.block_1(hidden_states, cache["mid_1"]) + hidden_states, cache["mid_2"] = self.mid.block_2(hidden_states, cache["mid_2"]) + hidden_states = F.silu(self.norm_out(hidden_states)) - return self.conv_out(hidden_states, cache["conv_out"]) + hidden_states, cache["conv_out"] = self.conv_out(hidden_states, cache["conv_out"]) + return hidden_states class Kandinsky6SRDecoder3D(nn.Module): + _supports_gradient_checkpointing = True + def __init__( self, out_channels: int, @@ -423,6 +498,8 @@ def __init__( self.norm_out = Kandinsky6SRSpatialNorm3D(block_in, latent_channels) self.conv_out = Kandinsky6SRCausalConv3d(block_in, out_channels, kernel_size=3) + self.gradient_checkpointing = False + def make_cache(self) -> dict: return { "conv_in": self.conv_in.make_cache(), @@ -440,17 +517,36 @@ def make_cache(self) -> dict: } def forward(self, latents: torch.Tensor, cache: dict) -> torch.Tensor: - hidden_states = self.conv_in(latents, cache["conv_in"]) - hidden_states = self.mid.block_1(hidden_states, cache["mid_1"], zq=latents) - hidden_states = self.mid.block_2(hidden_states, cache["mid_2"], zq=latents) + hidden_states, cache["conv_in"] = self.conv_in(latents, cache["conv_in"]) + + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states, cache["mid_1"] = self._gradient_checkpointing_func( + self.mid.block_1, hidden_states, cache["mid_1"], latents + ) + hidden_states, cache["mid_2"] = self._gradient_checkpointing_func( + self.mid.block_2, hidden_states, cache["mid_2"], latents + ) + else: + hidden_states, cache["mid_1"] = self.mid.block_1(hidden_states, cache["mid_1"], zq=latents) + hidden_states, cache["mid_2"] = self.mid.block_2(hidden_states, cache["mid_2"], zq=latents) + for level in reversed(range(len(self.up))): up, level_cache = self.up[level], cache["up"][level] - for block, block_cache in zip(up.block, level_cache["block"]): - hidden_states = block(hidden_states, block_cache, zq=latents) + for i, block in enumerate(up.block): + block_cache = level_cache["block"][i] + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states, level_cache["block"][i] = self._gradient_checkpointing_func( + block, hidden_states, block_cache, latents + ) + else: + hidden_states, level_cache["block"][i] = block(hidden_states, block_cache, zq=latents) if hasattr(up, "upsample"): hidden_states = up.upsample(hidden_states, level_cache["upsample"]) - hidden_states = F.silu(self.norm_out(hidden_states, latents, cache["norm_out"])) - return self.conv_out(hidden_states, cache["conv_out"]) + + hidden_states, cache["norm_out"] = self.norm_out(hidden_states, latents, cache["norm_out"]) + hidden_states = F.silu(hidden_states) + hidden_states, cache["conv_out"] = self.conv_out(hidden_states, cache["conv_out"]) + return hidden_states class Kandinsky6SRVAE(ModelMixin, ConfigMixin): @@ -461,6 +557,11 @@ class Kandinsky6SRVAE(ModelMixin, ConfigMixin): carry their padding state between segments, so the segmentation only bounds peak memory and does not change the result. + Supports gradient checkpointing (see [`~ModelMixin.enable_gradient_checkpointing`]): every + [`Kandinsky6SRResnetBlock3D`] in the encoder and decoder is checkpointed, since its `cache` argument is threaded + functionally (read, never mutated in place — see [`Kandinsky6SRCausalConv3d`]), so recomputing it during backward + reproduces the exact same forward pass. + Args: in_channels (`int`, defaults to `3`): Number of pixel channels. @@ -484,6 +585,7 @@ class Kandinsky6SRVAE(ModelMixin, ConfigMixin): """ _no_split_modules = ["Kandinsky6SREncoder3D", "Kandinsky6SRDecoder3D"] + _supports_gradient_checkpointing = True @register_to_config def __init__( diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index ff644daea5c0..14419b2b1818 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -87,6 +87,8 @@ def normalize(x: torch.Tensor, dim: list[int] | None = None, eps: float = 1e-4) def mp_silu(x: torch.Tensor) -> torch.Tensor: """SiLU rescaled so that a unit-variance input stays unit-variance.""" + # 0.596 is the (empirically measured) standard deviation of SiLU(z) for z ~ N(0, 1), i.e. the EDM2-style + # magnitude-preserving rescaling constant for this nonlinearity (see Karras et al. reference above). return F.silu(x) / 0.596 @@ -334,8 +336,8 @@ class MMAudioVAE(ModelMixin, ConfigMixin): Audio VAE of [`Kandinsky6TI2VAPipeline`]: a magnitude-preserving autoencoder over log-mel spectrograms (MMAudio, https://arxiv.org/abs/2412.15322). - `encode` turns a waveform into a latent distribution; `decode` turns latents back into a mel spectrogram, which - the pipeline-local [`~pipelines.kandinsky6.MMAudioVocoder`] then turns into a waveform. One latent frame covers + `encode` turns a waveform into a latent distribution; `decode` turns latents back into a mel spectrogram, which the + pipeline-local [`~pipelines.kandinsky6.MMAudioVocoder`] then turns into a waveform. One latent frame covers `hop_length * 2` samples. Args: diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 2567eb2b6ce9..d61903b3415f 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -25,9 +25,6 @@ from torch.nn.attention.flex_attention import BlockMask from ...configuration_utils import ConfigMixin, register_to_config -from ...hooks import MagCacheConfig -from ...hooks.hooks import HookRegistry, ModelHook, StateManager -from ...hooks.mag_cache import MagCacheState from ...loaders import PeftAdapterMixin from ...utils import BaseOutput from ..attention import AttentionMixin, AttentionModuleMixin @@ -203,221 +200,6 @@ def forward( return rope.unsqueeze(-4) # (T, H, W, 1, total_dim, 2, 2) -# Diffusers MagCache adapter for the K6 multimodal transformer. -# -# Diffusers provides the public :class:`MagCacheConfig` and the stateful hook -# infrastructure. K6's fused visual blocks carry video and audio streams -# together, so the stock hook needs a small adapter to preserve both tensors -# when a block is skipped. The adapter keeps the Diffusers ``enable_cache`` / -# ``disable_cache`` API and uses the Diffusers MagCache state and configuration. - - -_HEAD_HOOK = "kandinsky6_mag_cache_head" -_BLOCK_HOOK = "kandinsky6_mag_cache_block" - - -def _streams_from_args(args: tuple[Any, ...], kwargs: dict[str, Any]) -> tuple[Tensor | None, Tensor | None]: - video = kwargs.get("vis", args[0] if args else None) - audio = kwargs.get("aud", args[1] if len(args) > 1 else None) - return video, audio - - -def _add_residual( - video: Tensor | None, - audio: Tensor | None, - residual: Tensor | tuple[Tensor | None, Tensor | None], -) -> Tensor | tuple[Tensor | None, Tensor | None]: - if audio is None: - if not isinstance(residual, Tensor): - raise RuntimeError("K6 MagCache residual does not match a video-only block") - return video + residual if video is not None else video - if not isinstance(residual, tuple): - raise RuntimeError("K6 MagCache residual does not contain an audio stream") - video_residual, audio_residual = residual - return ( - video + video_residual if video is not None and video_residual is not None else video, - audio + audio_residual if audio_residual is not None else audio, - ) - - -def _residual( - output: Tensor | tuple[Tensor | None, Tensor | None], - input_video: Tensor | None, - input_audio: Tensor | None, -) -> Tensor | tuple[Tensor | None, Tensor | None]: - output_video, output_audio = _streams_from_args((output,), {}) if isinstance(output, Tensor) else output - if input_audio is None: - if output_video is None or input_video is None: - return output_video - return output_video - input_video - return ( - output_video - input_video if output_video is not None and input_video is not None else None, - output_audio - input_audio if output_audio is not None else None, - ) - - -def _should_compute(state: MagCacheState, config: MagCacheConfig, *, lane: int, num_steps: int) -> bool: - if config.calibrate: - return True - - ratio_index = state.step_index * 2 + lane - if config.mag_ratios is None or ratio_index >= len(config.mag_ratios): - current_scale = 1.0 - else: - current_scale = float(config.mag_ratios[ratio_index]) - - retention_step = int(config.retention_ratio * num_steps + 0.5) - if state.step_index < retention_step: - return True - - state.accumulated_ratio *= current_scale - state.accumulated_steps += 1 - state.accumulated_err += abs(1.0 - state.accumulated_ratio) - if ( - state.previous_residual is not None - and state.accumulated_err <= config.threshold - and state.accumulated_steps <= config.max_skip_steps - ): - return False - - state.accumulated_ratio = 1.0 - state.accumulated_steps = 0 - state.accumulated_err = 0.0 - return True - - -def _advance(state: MagCacheState, config: MagCacheConfig, num_steps: int) -> None: - state.step_index += 1 - if state.step_index < num_steps: - return - state.step_index = 0 - state.accumulated_ratio = 1.0 - state.accumulated_steps = 0 - state.accumulated_err = 0.0 - state.previous_residual = None - state.head_block_input = None - state.should_compute = True - state.calibration_ratios = [] - - -class _Kandinsky6MagCacheHeadHook(ModelHook): - _is_stateful = True - - def __init__(self, state_manager: StateManager, config: MagCacheConfig, num_steps: int): - super().__init__() - self.state_manager = state_manager - self.config = config - self.num_steps = num_steps - - @torch.compiler.disable - def new_forward(self, module: nn.Module, *args, **kwargs): - if self.state_manager._current_context is None: - self.state_manager.set_context("inference") - state: MagCacheState = self.state_manager.get_state() - video, audio = _streams_from_args(args, kwargs) - state.head_block_input = (video, audio) - lane = 1 if self.state_manager._current_context in {"uncond", "negative"} else 0 - state.should_compute = _should_compute(state, self.config, lane=lane, num_steps=self.num_steps) - - if not state.should_compute: - if state.previous_residual is None: - raise RuntimeError("K6 MagCache requested a skip before a residual was computed") - return _add_residual(video, audio, state.previous_residual) - return self.fn_ref.original_forward(*args, **kwargs) - - def reset_state(self, module: nn.Module): - self.state_manager.reset() - return module - - -class _Kandinsky6MagCacheBlockHook(ModelHook): - def __init__(self, state_manager: StateManager, config: MagCacheConfig, num_steps: int, is_tail: bool): - super().__init__() - self.state_manager = state_manager - self.config = config - self.num_steps = num_steps - self.is_tail = is_tail - - @torch.compiler.disable - def new_forward(self, module: nn.Module, *args, **kwargs): - if self.state_manager._current_context is None: - self.state_manager.set_context("inference") - state: MagCacheState = self.state_manager.get_state() - video, audio = _streams_from_args(args, kwargs) - - if not state.should_compute: - if self.is_tail: - _advance(state, self.config, self.num_steps) - return video if audio is None else (video, audio) - - output = self.fn_ref.original_forward(*args, **kwargs) - if self.is_tail: - input_video, input_audio = state.head_block_input - state.previous_residual = _residual(output, input_video, input_audio) - _advance(state, self.config, self.num_steps) - return output - - def reset_state(self, module: nn.Module): - self.state_manager.reset() - return module - - -def _apply_kandinsky6_mag_cache(module: nn.Module, config: MagCacheConfig) -> None: - blocks = getattr(module, "visual_transformer_blocks", None) - if not isinstance(blocks, nn.ModuleList) or len(blocks) == 0: - raise ValueError("K6 MagCache requires a non-empty transformer.visual_transformer_blocks ModuleList") - - registry = HookRegistry.check_if_exists_or_initialize(module) - registry.remove_hook(_HEAD_HOOK, recurse=True) - registry.remove_hook(_BLOCK_HOOK, recurse=True) - state_manager = StateManager(MagCacheState, (), {}) - # K6 stores conditional and unconditional coefficients interleaved. The - # Diffusers pipeline sets separate cache contexts for the two CFG passes. - num_steps = config.num_inference_steps // 2 - - head_registry = HookRegistry.check_if_exists_or_initialize(blocks[0]) - head_registry.register_hook( - _Kandinsky6MagCacheHeadHook(state_manager, config, num_steps), - _HEAD_HOOK, - ) - for block in blocks[1:-1]: - block_registry = HookRegistry.check_if_exists_or_initialize(block) - block_registry.register_hook( - _Kandinsky6MagCacheBlockHook(state_manager, config, num_steps, is_tail=False), - _BLOCK_HOOK, - ) - if len(blocks) > 1: - tail_registry = HookRegistry.check_if_exists_or_initialize(blocks[-1]) - tail_registry.register_hook( - _Kandinsky6MagCacheBlockHook(state_manager, config, num_steps, is_tail=True), - _BLOCK_HOOK, - ) - - -class Kandinsky6MagCacheMixin(CacheMixin): - """Expose Diffusers' cache API for K6's dual-stream visual blocks.""" - - def enable_cache(self, config) -> None: - if not isinstance(config, MagCacheConfig): - return super().enable_cache(config) - if self.is_cache_enabled: - raise ValueError(f"Caching has already been enabled with {type(self._cache_config)}") - if config.num_inference_steps % 2 != 0: - raise ValueError( - "K6 MagCacheConfig.num_inference_steps must count interleaved CFG forwards (an even number)" - ) - _apply_kandinsky6_mag_cache(self, config) - self._cache_config = config - - def disable_cache(self) -> None: - if not isinstance(self._cache_config, MagCacheConfig): - return super().disable_cache() - registry = HookRegistry.check_if_exists_or_initialize(self) - registry.remove_hook(_HEAD_HOOK, recurse=True) - registry.remove_hook(_BLOCK_HOOK, recurse=True) - self._cache_config = None - - # Diffusers-style K6 transformer components. # # The classes use the same inner-module vocabulary as the Diffusers Kandinsky5 @@ -437,10 +219,8 @@ def disable_cache(self) -> None: class Kandinsky6AttnProcessor: """Diffusers attention processor used by the TI2VA transformer. - The NABLA sparse branch always dispatches on the `flex` backend, regardless of `_attention_backend`: its - `BlockMask` only the `flex` backend can consume, and `_attention_backend` can be transparently remapped to a - masked-attention variant (see the setter above) when `text_token_padding` is enabled, which would otherwise be - given a `BlockMask` it can't handle. + The NABLA sparse branch always dispatches on the `flex` backend, regardless of `_attention_backend`, since its + `BlockMask` only the `flex` backend can consume. """ _attention_backend = None @@ -602,7 +382,6 @@ def __init__( num_channels: int, head_dim: int, kv_dim: int | None = None, - text_token_padding: bool = False, processor: Kandinsky6AttnProcessor | None = None, ): super().__init__() @@ -616,11 +395,7 @@ def __init__( self.query_norm = nn.RMSNorm(head_dim) self.key_norm = nn.RMSNorm(head_dim) self.out_layer = nn.Linear(num_channels, num_channels) - self.text_token_padding = text_token_padding self.set_processor(processor or self._default_processor_cls()) - if self.text_token_padding: - self.processor._masked = True - self.processor._attention_backend = AttentionBackendName.NATIVE def forward( self, @@ -707,12 +482,11 @@ def __init__( time_dim: int, ff_dim: int, head_dim: int, - text_token_padding: bool = False, ): super().__init__() self.text_modulation = Kandinsky6Modulation(time_dim, model_dim, 6) self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) - self.self_attention = Kandinsky6Attention(model_dim, head_dim, text_token_padding=text_token_padding) + self.self_attention = Kandinsky6Attention(model_dim, head_dim) self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.feed_forward = Kandinsky6FeedForward(model_dim, ff_dim) @@ -748,7 +522,6 @@ def __init__( time_dim: int, ff_dim: int, head_dim: int, - text_token_padding: bool = False, ): super().__init__() self.visual_modulation = Kandinsky6Modulation(time_dim, model_dim, 9) @@ -759,7 +532,6 @@ def __init__( model_dim, head_dim, kv_dim=model_dim, - text_token_padding=text_token_padding, ) self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.feed_forward = Kandinsky6FeedForward(model_dim, ff_dim) @@ -778,16 +550,10 @@ def __init__( time_dim_a: int, ff_dim_a: int, head_dim_a: int, - text_token_padding: bool = False, - ca_rope: bool = False, - cross_gates: bool = False, - fix_modulation: bool = False, ): super().__init__() - self.videoT = Kandinsky6TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim, text_token_padding) - self.audioT = Kandinsky6TransformerDecoderBlock( - model_dim_a, time_dim_a, ff_dim_a, head_dim_a, text_token_padding - ) + self.videoT = Kandinsky6TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim) + self.audioT = Kandinsky6TransformerDecoderBlock(model_dim_a, time_dim_a, ff_dim_a, head_dim_a) self.va_cross_attention = Kandinsky6Attention( model_dim, head_dim, @@ -798,23 +564,10 @@ def __init__( head_dim_a, kv_dim=model_dim, ) - self.va_modulation = Kandinsky6Modulation( - time_dim, - model_dim if not cross_gates else model_dim * 2 + model_dim_a, - 1 if cross_gates else 3, - ) - self.av_modulation = Kandinsky6Modulation( - time_dim_a, - model_dim_a if not cross_gates else model_dim_a * 2 + model_dim, - 1 if cross_gates else 3, - ) + self.va_modulation = Kandinsky6Modulation(time_dim, model_dim, 3) + self.av_modulation = Kandinsky6Modulation(time_dim_a, model_dim_a, 3) self.va_normalization = nn.LayerNorm(model_dim, elementwise_affine=False) self.av_normalization = nn.LayerNorm(model_dim_a, elementwise_affine=False) - self.ca_rope = ca_rope - self.cross_gates = cross_gates - self.fix_modulation = fix_modulation - self.model_dim = model_dim - self.model_dim_a = model_dim_a def forward( self, @@ -870,39 +623,23 @@ def forward( aud = apply_gate_sum(aud, aud_out_t, gate_a).type_as(aud) if vis is not None: - t_va_mod = t_a if not self.fix_modulation else t_v - t_av_mod = t_v if not self.fix_modulation else t_a - va_params = self.va_modulation(t_va_mod) - av_params = self.av_modulation(t_av_mod) - if self.cross_gates: - va_shift, va_scale, va_gate = torch.split( - va_params, [self.model_dim, self.model_dim, self.model_dim_a], dim=-1 - ) - av_shift, av_scale, av_gate = torch.split( - av_params, [self.model_dim_a, self.model_dim_a, self.model_dim], dim=-1 - ) - else: - va_shift, va_scale, va_gate = torch.chunk(va_params, 3, dim=-1) - av_shift, av_scale, av_gate = torch.chunk(av_params, 3, dim=-1) + va_params = self.va_modulation(t_a) + av_params = self.av_modulation(t_v) + va_shift, va_scale, va_gate = torch.chunk(va_params, 3, dim=-1) + av_shift, av_scale, av_gate = torch.chunk(av_params, 3, dim=-1) vis = apply_gate_sum(vis, vis_out_t, gate_v).type_as(vis) vis_for_va = apply_scale_shift_norm(self.va_normalization, vis, va_scale, va_shift) aud_for_av = apply_scale_shift_norm(self.av_normalization, aud, av_scale, av_shift) - rq_v = vis_rope if self.ca_rope else None - rk_a = aud_rope if self.ca_rope else None vis_from_aud = self.va_cross_attention( vis_for_va, encoder_hidden_states=aud_pre_ca, - rope_q=rq_v, - rope_kv=rk_a, ) aud_from_vis = self.av_cross_attention( aud_for_av, encoder_hidden_states=vis_pre_ca, - rope_q=rk_a, - rope_kv=rq_v, ) - vis = apply_gate_sum(vis, vis_from_aud, va_gate if not self.cross_gates else av_gate).type_as(vis) - aud = apply_gate_sum(aud, aud_from_vis, av_gate if not self.cross_gates else va_gate).type_as(aud) + vis = apply_gate_sum(vis, vis_from_aud, va_gate).type_as(vis) + aud = apply_gate_sum(aud, aud_from_vis, av_gate).type_as(aud) elif vis is not None: vis = apply_gate_sum(vis, vis_out_t, gate_v).type_as(vis) @@ -924,7 +661,6 @@ def forward( class Kandinsky6Transformer3DModel( - Kandinsky6MagCacheMixin, ModelMixin, ConfigMixin, PeftAdapterMixin, @@ -958,14 +694,9 @@ class Kandinsky6Transformer3DModel( ff_dim_a (`int`, *optional*): Audio feed-forward hidden dimension. Defaults to `ff_dim`. axes_dims_a (`tuple[int, int, int]`, *optional*): Audio RoPE dimensions. Defaults to `axes_dims`. audio_freqs_scaling (`float`, *optional*, defaults to 1.0): Audio RoPE frequency scaling. - text_token_padding (`bool`, *optional*, defaults to False): Whether text sequences are padded. scale_factor (`tuple[float, float, float]`, *optional*, defaults to `(1.0, 2.0, 2.0)`): Per-axis `(t, h, w)` RoPE frequency scaling applied to the video positions. - ca_rope (`bool`, *optional*, defaults to False): Whether to use cross-modal audio RoPE. - cross_gates (`bool`, *optional*, defaults to False): Whether to use cross-modal residual gates. - fix_modulation (`bool`, *optional*, defaults to False): Whether to use the fixed modulation variant. visual_token_type_num_embeddings (`int`, *optional*, defaults to 0): Number of visual token type embeddings. - magcache (`dict`, *optional*): MagCache configuration. """ _repeated_blocks = [ @@ -1001,20 +732,14 @@ def __init__( ff_dim_a: int | None = None, axes_dims_a: tuple | None = None, audio_freqs_scaling: float = 1.0, - text_token_padding: bool = False, scale_factor: tuple | list[float] = (1.0, 2.0, 2.0), - ca_rope: bool = False, - cross_gates: bool = False, - fix_modulation: bool = False, visual_token_type_num_embeddings: int = 0, - magcache: dict[str, Any] | None = None, ) -> None: super().__init__() self.patch_size = patch_size self.visual_cond = visual_cond self.in_visual_dim = in_visual_dim self.in_audio_dim = in_audio_dim - self.text_token_padding = text_token_padding self.scale_factor = tuple(float(value) for value in scale_factor) self.visual_token_type_num_embeddings = int(visual_token_type_num_embeddings or 0) head_dim = sum(axes_dims) @@ -1045,12 +770,7 @@ def __init__( setattr( self, f"{prefix}_text_transformer_blocks", - nn.ModuleList( - [ - Kandinsky6TransformerEncoderBlock(md, td, fd, hd, text_token_padding) - for _ in range(num_text_blocks) - ] - ), + nn.ModuleList([Kandinsky6TransformerEncoderBlock(md, td, fd, hd) for _ in range(num_text_blocks)]), ) self.visual_transformer_blocks = nn.ModuleList( [ @@ -1063,10 +783,6 @@ def __init__( time_dim_a, ff_dim_a, head_dim_a, - text_token_padding, - ca_rope, - cross_gates, - fix_modulation, ) for _ in range(num_visual_blocks) ] diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index 5c5bdd9997f7..6c3acfaf6f3b 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -149,11 +149,11 @@ class Kandinsky6SRAttnProcessor: """Self-attention processor of the SR transformer: dense by default, NABLA sparse when `sparse_params` is set. Like [`Kandinsky5AttnProcessor`], the attention backend is chosen through the usual `_attention_backend` - (`set_attention_backend`), even in the NABLA sparse branch. Its sparsity is expressed as a `BlockMask`, which - only the `flex` backend can consume, so sparse checkpoints require `set_attention_backend("flex")`. That backend - also needs to run under `torch.compile` (e.g. `transformer.compile_repeated_blocks()`): uncompiled, PyTorch's - flex attention falls back to an eager implementation that materializes the full attention matrix, which does not - fit in memory at video resolutions. + (`set_attention_backend`), even in the NABLA sparse branch. Its sparsity is expressed as a `BlockMask`, which only + the `flex` backend can consume, so sparse checkpoints require `set_attention_backend("flex")`. That backend also + needs to run under `torch.compile` (e.g. `transformer.compile_repeated_blocks()`): uncompiled, PyTorch's flex + attention falls back to an eager implementation that materializes the full attention matrix, which does not fit in + memory at video resolutions. """ _attention_backend = None From 3d9c95cb84b09e3079a77dff9e594156454755be Mon Sep 17 00:00:00 2001 From: leffff Date: Wed, 30 Sep 2026 18:06:10 +0300 Subject: [PATCH 30/62] refactor --- .../autoencoders/autoencoder_kandinsky6_sr.py | 4 +- .../autoencoders/autoencoder_mmaudio.py | 5 +- .../transformers/transformer_kandinsky6.py | 89 ++++++++++++++++--- 3 files changed, 82 insertions(+), 16 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index f933a91ddef4..47731d00e008 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -77,8 +77,8 @@ class Kandinsky6SRCausalConv3d(nn.Module): This caching is not an optional performance knob: it is what lets `Kandinsky6SRVAE.encode`/`decode` process an arbitrarily long video in bounded-memory segments (see `SEGMENT_FRAMES`) while reproducing the exact output of a single non-causal pass. Without it, each segment would need the raw frames the previous segment already consumed in - order to rebuild correct padding, which means keeping the whole video in memory anyway and defeats the reason - segmentation exists. + order to rebuild correct padding — a lookback window that grows with network depth, rather than the bounded, + constant-size cache this class carries instead — which defeats the point of processing the video in segments. """ def __init__( diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index 14419b2b1818..147681fe56cc 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -87,8 +87,9 @@ def normalize(x: torch.Tensor, dim: list[int] | None = None, eps: float = 1e-4) def mp_silu(x: torch.Tensor) -> torch.Tensor: """SiLU rescaled so that a unit-variance input stays unit-variance.""" - # 0.596 is the (empirically measured) standard deviation of SiLU(z) for z ~ N(0, 1), i.e. the EDM2-style - # magnitude-preserving rescaling constant for this nonlinearity (see Karras et al. reference above). + # 0.596 is the (empirically measured) RMS of SiLU(z) for z ~ N(0, 1), i.e. sqrt(E[SiLU(z)^2]) rather than its + # standard deviation (SiLU(z) has nonzero mean), matching the EDM2 magnitude-preserving convention of keeping + # E[x^2] = 1 rather than mean-subtracted variance = 1 (see Karras et al. reference above). return F.silu(x) / 0.596 diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index d61903b3415f..a4d01e7c7995 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -382,6 +382,7 @@ def __init__( num_channels: int, head_dim: int, kv_dim: int | None = None, + text_token_padding: bool = False, processor: Kandinsky6AttnProcessor | None = None, ): super().__init__() @@ -395,7 +396,11 @@ def __init__( self.query_norm = nn.RMSNorm(head_dim) self.key_norm = nn.RMSNorm(head_dim) self.out_layer = nn.Linear(num_channels, num_channels) + self.text_token_padding = text_token_padding self.set_processor(processor or self._default_processor_cls()) + if self.text_token_padding: + self.processor._masked = True + self.processor._attention_backend = AttentionBackendName.NATIVE def forward( self, @@ -482,11 +487,12 @@ def __init__( time_dim: int, ff_dim: int, head_dim: int, + text_token_padding: bool = False, ): super().__init__() self.text_modulation = Kandinsky6Modulation(time_dim, model_dim, 6) self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) - self.self_attention = Kandinsky6Attention(model_dim, head_dim) + self.self_attention = Kandinsky6Attention(model_dim, head_dim, text_token_padding=text_token_padding) self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.feed_forward = Kandinsky6FeedForward(model_dim, ff_dim) @@ -522,6 +528,7 @@ def __init__( time_dim: int, ff_dim: int, head_dim: int, + text_token_padding: bool = False, ): super().__init__() self.visual_modulation = Kandinsky6Modulation(time_dim, model_dim, 9) @@ -532,6 +539,7 @@ def __init__( model_dim, head_dim, kv_dim=model_dim, + text_token_padding=text_token_padding, ) self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.feed_forward = Kandinsky6FeedForward(model_dim, ff_dim) @@ -550,10 +558,16 @@ def __init__( time_dim_a: int, ff_dim_a: int, head_dim_a: int, + text_token_padding: bool = False, + ca_rope: bool = False, + cross_gates: bool = False, + fix_modulation: bool = False, ): super().__init__() - self.videoT = Kandinsky6TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim) - self.audioT = Kandinsky6TransformerDecoderBlock(model_dim_a, time_dim_a, ff_dim_a, head_dim_a) + self.videoT = Kandinsky6TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim, text_token_padding) + self.audioT = Kandinsky6TransformerDecoderBlock( + model_dim_a, time_dim_a, ff_dim_a, head_dim_a, text_token_padding + ) self.va_cross_attention = Kandinsky6Attention( model_dim, head_dim, @@ -564,10 +578,23 @@ def __init__( head_dim_a, kv_dim=model_dim, ) - self.va_modulation = Kandinsky6Modulation(time_dim, model_dim, 3) - self.av_modulation = Kandinsky6Modulation(time_dim_a, model_dim_a, 3) + self.va_modulation = Kandinsky6Modulation( + time_dim, + model_dim if not cross_gates else model_dim * 2 + model_dim_a, + 1 if cross_gates else 3, + ) + self.av_modulation = Kandinsky6Modulation( + time_dim_a, + model_dim_a if not cross_gates else model_dim_a * 2 + model_dim, + 1 if cross_gates else 3, + ) self.va_normalization = nn.LayerNorm(model_dim, elementwise_affine=False) self.av_normalization = nn.LayerNorm(model_dim_a, elementwise_affine=False) + self.ca_rope = ca_rope + self.cross_gates = cross_gates + self.fix_modulation = fix_modulation + self.model_dim = model_dim + self.model_dim_a = model_dim_a def forward( self, @@ -623,23 +650,39 @@ def forward( aud = apply_gate_sum(aud, aud_out_t, gate_a).type_as(aud) if vis is not None: - va_params = self.va_modulation(t_a) - av_params = self.av_modulation(t_v) - va_shift, va_scale, va_gate = torch.chunk(va_params, 3, dim=-1) - av_shift, av_scale, av_gate = torch.chunk(av_params, 3, dim=-1) + t_va_mod = t_a if not self.fix_modulation else t_v + t_av_mod = t_v if not self.fix_modulation else t_a + va_params = self.va_modulation(t_va_mod) + av_params = self.av_modulation(t_av_mod) + if self.cross_gates: + va_shift, va_scale, va_gate = torch.split( + va_params, [self.model_dim, self.model_dim, self.model_dim_a], dim=-1 + ) + av_shift, av_scale, av_gate = torch.split( + av_params, [self.model_dim_a, self.model_dim_a, self.model_dim], dim=-1 + ) + else: + va_shift, va_scale, va_gate = torch.chunk(va_params, 3, dim=-1) + av_shift, av_scale, av_gate = torch.chunk(av_params, 3, dim=-1) vis = apply_gate_sum(vis, vis_out_t, gate_v).type_as(vis) vis_for_va = apply_scale_shift_norm(self.va_normalization, vis, va_scale, va_shift) aud_for_av = apply_scale_shift_norm(self.av_normalization, aud, av_scale, av_shift) + rq_v = vis_rope if self.ca_rope else None + rk_a = aud_rope if self.ca_rope else None vis_from_aud = self.va_cross_attention( vis_for_va, encoder_hidden_states=aud_pre_ca, + rope_q=rq_v, + rope_kv=rk_a, ) aud_from_vis = self.av_cross_attention( aud_for_av, encoder_hidden_states=vis_pre_ca, + rope_q=rk_a, + rope_kv=rq_v, ) - vis = apply_gate_sum(vis, vis_from_aud, va_gate).type_as(vis) - aud = apply_gate_sum(aud, aud_from_vis, av_gate).type_as(aud) + vis = apply_gate_sum(vis, vis_from_aud, va_gate if not self.cross_gates else av_gate).type_as(vis) + aud = apply_gate_sum(aud, aud_from_vis, av_gate if not self.cross_gates else va_gate).type_as(aud) elif vis is not None: vis = apply_gate_sum(vis, vis_out_t, gate_v).type_as(vis) @@ -696,7 +739,15 @@ class Kandinsky6Transformer3DModel( audio_freqs_scaling (`float`, *optional*, defaults to 1.0): Audio RoPE frequency scaling. scale_factor (`tuple[float, float, float]`, *optional*, defaults to `(1.0, 2.0, 2.0)`): Per-axis `(t, h, w)` RoPE frequency scaling applied to the video positions. + text_token_padding (`bool`, *optional*, defaults to False): Whether text sequences are padded. + ca_rope (`bool`, *optional*, defaults to False): Whether to use cross-modal audio RoPE. + cross_gates (`bool`, *optional*, defaults to False): Whether to use cross-modal residual gates. + fix_modulation (`bool`, *optional*, defaults to False): Whether to use the fixed modulation variant. visual_token_type_num_embeddings (`int`, *optional*, defaults to 0): Number of visual token type embeddings. + + Released checkpoints set `text_token_padding`, `ca_rope`, `cross_gates`, and `fix_modulation` to `True` (see each + checkpoint's `transformer/config.json`); the `False` defaults only describe an architecture variant this repo + does not ship weights for. """ _repeated_blocks = [ @@ -733,6 +784,10 @@ def __init__( axes_dims_a: tuple | None = None, audio_freqs_scaling: float = 1.0, scale_factor: tuple | list[float] = (1.0, 2.0, 2.0), + text_token_padding: bool = False, + ca_rope: bool = False, + cross_gates: bool = False, + fix_modulation: bool = False, visual_token_type_num_embeddings: int = 0, ) -> None: super().__init__() @@ -740,6 +795,7 @@ def __init__( self.visual_cond = visual_cond self.in_visual_dim = in_visual_dim self.in_audio_dim = in_audio_dim + self.text_token_padding = text_token_padding self.scale_factor = tuple(float(value) for value in scale_factor) self.visual_token_type_num_embeddings = int(visual_token_type_num_embeddings or 0) head_dim = sum(axes_dims) @@ -770,7 +826,12 @@ def __init__( setattr( self, f"{prefix}_text_transformer_blocks", - nn.ModuleList([Kandinsky6TransformerEncoderBlock(md, td, fd, hd) for _ in range(num_text_blocks)]), + nn.ModuleList( + [ + Kandinsky6TransformerEncoderBlock(md, td, fd, hd, text_token_padding) + for _ in range(num_text_blocks) + ] + ), ) self.visual_transformer_blocks = nn.ModuleList( [ @@ -783,6 +844,10 @@ def __init__( time_dim_a, ff_dim_a, head_dim_a, + text_token_padding, + ca_rope, + cross_gates, + fix_modulation, ) for _ in range(num_visual_blocks) ] From 002f4fb73401fc46cdb213f3643d98d1e228d993 Mon Sep 17 00:00:00 2001 From: leffff Date: Wed, 30 Sep 2026 19:24:14 +0300 Subject: [PATCH 31/62] address apply_scale_shift_norm problem --- .../transformers/transformer_kandinsky6.py | 40 +++++++++++-------- .../transformers/transformer_kandinsky6_sr.py | 22 ++++++---- 2 files changed, 38 insertions(+), 24 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index a4d01e7c7995..6a597613fb5a 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -64,12 +64,16 @@ def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) -def apply_scale_shift_norm(norm, x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: - """AdaLN-style affine in fp32, cast back to ``x.dtype``.""" +def apply_scale_shift(normed: Tensor, x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: + """Apply an AdaLN-style scale/shift affine to an already-normalized tensor, in fp32, cast back to ``x.dtype``. + + Callers compute ``normed`` themselves (e.g. ``self.some_norm(x.float())``) so the norm-layer call stays visible + in `forward` instead of being hidden inside this helper. + """ if x.ndim > 2 and scale.ndim == 2: shape = (scale.shape[0],) + (1,) * (x.ndim - 2) + (scale.shape[-1],) scale, shift = scale.reshape(shape), shift.reshape(shape) - return (norm(x.float()) * (scale.float() + 1.0) + shift.float()).to(dtype=x.dtype) + return (normed * (scale.float() + 1.0) + shift.float()).to(dtype=x.dtype) def apply_gate_sum(x: Tensor, out: Tensor, gate: Tensor) -> Tensor: @@ -441,8 +445,8 @@ def __init__(self, model_dim: int, time_dim: int, visual_dim: int, patch_size: t def forward(self, visual_embed: Tensor, time_embed: Tensor) -> Tensor: shift, scale = torch.chunk(self.modulation(time_embed), 2, dim=-1) condition_shape = (scale.shape[0],) + (1,) * (visual_embed.ndim - 2) + (scale.shape[-1],) - x = apply_scale_shift_norm( - self.norm, + x = apply_scale_shift( + self.norm(visual_embed.float()), visual_embed, scale.reshape(condition_shape), shift.reshape(condition_shape), @@ -471,7 +475,7 @@ def __init__(self, model_dim: int, time_dim: int, audio_dim: int): def forward(self, audio_embed: Tensor, time_embed: Tensor) -> Tensor: shift, scale = torch.chunk(self.modulation(time_embed), 2, dim=-1) - x = apply_scale_shift_norm(self.norm, audio_embed, scale, shift) + x = apply_scale_shift(self.norm(audio_embed.float()), audio_embed, scale, shift) # The reference audio head normalizes a second time after the AdaLN affine; the released weights were # trained this way, so the double norm is kept for parity even though the video head has none. x = self.norm(x) @@ -502,7 +506,7 @@ def forward(self, x: Tensor, time_embed: Tensor, rope: Tensor, attn_mask: Tensor x = apply_gate_sum( x, self.self_attention( - apply_scale_shift_norm(self.self_attention_norm, x, scale, shift), + apply_scale_shift(self.self_attention_norm(x.float()), x, scale, shift), rotary_emb=rope, attn_mask=attn_mask, ), @@ -510,7 +514,7 @@ def forward(self, x: Tensor, time_embed: Tensor, rope: Tensor, attn_mask: Tensor ) shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) return apply_gate_sum( - x, self.feed_forward(apply_scale_shift_norm(self.feed_forward_norm, x, scale, shift)), gate + x, self.feed_forward(apply_scale_shift(self.feed_forward_norm(x.float()), x, scale, shift)), gate ) @@ -615,14 +619,14 @@ def forward( vis = apply_gate_sum( vis, self.videoT.self_attention( - apply_scale_shift_norm(self.videoT.self_attention_norm, vis, scale, shift), + apply_scale_shift(self.videoT.self_attention_norm(vis.float()), vis, scale, shift), rotary_emb=vis_rope, sparse_params=sparse_params, ), gate, ).type_as(vis) shift, scale, gate_v = torch.chunk(ca_p, 3, dim=-1) - vis_pre_ca = apply_scale_shift_norm(self.videoT.cross_attention_norm, vis, scale, shift) + vis_pre_ca = apply_scale_shift(self.videoT.cross_attention_norm(vis.float()), vis, scale, shift) vis_out_t = self.videoT.cross_attention( vis_pre_ca, encoder_hidden_states=text_v, @@ -635,13 +639,13 @@ def forward( aud = apply_gate_sum( aud, self.audioT.self_attention( - apply_scale_shift_norm(self.audioT.self_attention_norm, aud, scale, shift), + apply_scale_shift(self.audioT.self_attention_norm(aud.float()), aud, scale, shift), rotary_emb=aud_rope, ), gate, ).type_as(aud) shift, scale, gate_a = torch.chunk(ca_p, 3, dim=-1) - aud_pre_ca = apply_scale_shift_norm(self.audioT.cross_attention_norm, aud, scale, shift) + aud_pre_ca = apply_scale_shift(self.audioT.cross_attention_norm(aud.float()), aud, scale, shift) aud_out_t = self.audioT.cross_attention( aud_pre_ca, encoder_hidden_states=text_a, @@ -665,8 +669,8 @@ def forward( va_shift, va_scale, va_gate = torch.chunk(va_params, 3, dim=-1) av_shift, av_scale, av_gate = torch.chunk(av_params, 3, dim=-1) vis = apply_gate_sum(vis, vis_out_t, gate_v).type_as(vis) - vis_for_va = apply_scale_shift_norm(self.va_normalization, vis, va_scale, va_shift) - aud_for_av = apply_scale_shift_norm(self.av_normalization, aud, av_scale, av_shift) + vis_for_va = apply_scale_shift(self.va_normalization(vis.float()), vis, va_scale, va_shift) + aud_for_av = apply_scale_shift(self.av_normalization(aud.float()), aud, av_scale, av_shift) rq_v = vis_rope if self.ca_rope else None rk_a = aud_rope if self.ca_rope else None vis_from_aud = self.va_cross_attention( @@ -690,14 +694,18 @@ def forward( shift, scale, gate = torch.chunk(ff_p, 3, dim=-1) vis = apply_gate_sum( vis, - self.videoT.feed_forward(apply_scale_shift_norm(self.videoT.feed_forward_norm, vis, scale, shift)), + self.videoT.feed_forward( + apply_scale_shift(self.videoT.feed_forward_norm(vis.float()), vis, scale, shift) + ), gate, ).type_as(vis) if aud is not None: shift, scale, gate = torch.chunk(ff_p_a, 3, dim=-1) aud = apply_gate_sum( aud, - self.audioT.feed_forward(apply_scale_shift_norm(self.audioT.feed_forward_norm, aud, scale, shift)), + self.audioT.feed_forward( + apply_scale_shift(self.audioT.feed_forward_norm(aud.float()), aud, scale, shift) + ), gate, ).type_as(aud) return vis, aud diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index 6c3acfaf6f3b..cb43ddf771bf 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -51,13 +51,17 @@ def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) -# Copied from diffusers.models.transformers.transformer_kandinsky6.apply_scale_shift_norm -def apply_scale_shift_norm(norm, x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: - """AdaLN-style affine in fp32, cast back to ``x.dtype``.""" +# Copied from diffusers.models.transformers.transformer_kandinsky6.apply_scale_shift +def apply_scale_shift(normed: Tensor, x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: + """Apply an AdaLN-style scale/shift affine to an already-normalized tensor, in fp32, cast back to ``x.dtype``. + + Callers compute ``normed`` themselves (e.g. ``self.some_norm(x.float())``) so the norm-layer call stays visible + in `forward` instead of being hidden inside this helper. + """ if x.ndim > 2 and scale.ndim == 2: shape = (scale.shape[0],) + (1,) * (x.ndim - 2) + (scale.shape[-1],) scale, shift = scale.reshape(shape), shift.reshape(shape) - return (norm(x.float()) * (scale.float() + 1.0) + shift.float()).to(dtype=x.dtype) + return (normed * (scale.float() + 1.0) + shift.float()).to(dtype=x.dtype) # Copied from diffusers.models.transformers.transformer_kandinsky6.apply_gate_sum @@ -312,8 +316,8 @@ def __init__(self, model_dim: int, time_dim: int, visual_dim: int, patch_size: t def forward(self, visual_embed: Tensor, time_embed: Tensor) -> Tensor: shift, scale = torch.chunk(self.modulation(time_embed), 2, dim=-1) condition_shape = (scale.shape[0],) + (1,) * (visual_embed.ndim - 2) + (scale.shape[-1],) - x = apply_scale_shift_norm( - self.norm, + x = apply_scale_shift( + self.norm(visual_embed.float()), visual_embed, scale.reshape(condition_shape), shift.reshape(condition_shape), @@ -397,7 +401,7 @@ def forward( hidden_states = apply_gate_sum( hidden_states, self.self_attention( - apply_scale_shift_norm(self.self_attention_norm, hidden_states, scale, shift), + apply_scale_shift(self.self_attention_norm(hidden_states.float()), hidden_states, scale, shift), rotary_emb, sparse_params, ), @@ -406,7 +410,9 @@ def forward( shift, scale, gate = torch.chunk(feed_forward_params, 3, dim=-1) return apply_gate_sum( hidden_states, - self.feed_forward(apply_scale_shift_norm(self.feed_forward_norm, hidden_states, scale, shift)), + self.feed_forward( + apply_scale_shift(self.feed_forward_norm(hidden_states.float()), hidden_states, scale, shift) + ), gate, ) From 0e5c59971682d1b5f696dd9b92e1e71feb155c7a Mon Sep 17 00:00:00 2001 From: leffff Date: Wed, 30 Sep 2026 19:35:55 +0300 Subject: [PATCH 32/62] update docs --- .../en/api/models/kandinsky6_transformer.md | 22 ++++++++++++++++ docs/source/en/api/models/kandinsky6_vae.md | 24 ++++++++++++++++++ docs/source/en/api/pipelines/kandinsky6.md | 25 ++++++++++++++++++- 3 files changed, 70 insertions(+), 1 deletion(-) diff --git a/docs/source/en/api/models/kandinsky6_transformer.md b/docs/source/en/api/models/kandinsky6_transformer.md index b712daf88cfa..7684a54a670c 100644 --- a/docs/source/en/api/models/kandinsky6_transformer.md +++ b/docs/source/en/api/models/kandinsky6_transformer.md @@ -17,6 +17,17 @@ text/image-to-video-and-audio generation, and a text-free diffusion transformer ## Kandinsky6Transformer3DModel +The multimodal transformer used by [`Kandinsky6TI2VAPipeline`]. + +```python +import torch +from diffusers import Kandinsky6Transformer3DModel + +transformer = Kandinsky6Transformer3DModel.from_pretrained( + "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16 +) +``` + [[autodoc]] Kandinsky6Transformer3DModel - all - forward @@ -27,6 +38,17 @@ text/image-to-video-and-audio generation, and a text-free diffusion transformer ## Kandinsky6SRTransformer3DModel +The text-free transformer used by [`Kandinsky6SRPipeline`] to refine one tile of the upscaled video at a time. + +```python +import torch +from diffusers import Kandinsky6SRTransformer3DModel + +transformer = Kandinsky6SRTransformer3DModel.from_pretrained( + "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16 +) +``` + [[autodoc]] Kandinsky6SRTransformer3DModel - all - forward diff --git a/docs/source/en/api/models/kandinsky6_vae.md b/docs/source/en/api/models/kandinsky6_vae.md index 2dd4c31a0689..b7c3924ff9dd 100644 --- a/docs/source/en/api/models/kandinsky6_vae.md +++ b/docs/source/en/api/models/kandinsky6_vae.md @@ -17,6 +17,18 @@ separate BigVGAN [`~pipelines.kandinsky6.MMAudioVocoder`], for synchronized audi ## Kandinsky6SRVAE +The causal 3D K-VAE used by [`Kandinsky6SRPipeline`]. It processes arbitrarily long videos in bounded-memory +segments while reproducing the exact output of a single, non-segmented pass. + +```python +import torch +from diffusers import Kandinsky6SRVAE + +vae = Kandinsky6SRVAE.from_pretrained( + "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers", subfolder="vae", torch_dtype=torch.bfloat16 +) +``` + [[autodoc]] Kandinsky6SRVAE - encode - decode @@ -24,6 +36,18 @@ separate BigVGAN [`~pipelines.kandinsky6.MMAudioVocoder`], for synchronized audi ## MMAudioVAE +The mel-spectrogram VAE used by [`Kandinsky6TI2VAPipeline`] when `sample_audio=True`. Its `decode` output is a mel +spectrogram; pass it through [`~pipelines.kandinsky6.MMAudioVocoder`] to get a waveform. + +```python +import torch +from diffusers import MMAudioVAE + +audio_vae = MMAudioVAE.from_pretrained( + "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers", subfolder="audio_vae", torch_dtype=torch.bfloat16 +) +``` + [[autodoc]] MMAudioVAE - encode - decode diff --git a/docs/source/en/api/pipelines/kandinsky6.md b/docs/source/en/api/pipelines/kandinsky6.md index 6a3caa06321f..1c16067293b5 100644 --- a/docs/source/en/api/pipelines/kandinsky6.md +++ b/docs/source/en/api/pipelines/kandinsky6.md @@ -18,7 +18,11 @@ transformer: video and audio latents are denoised together through fused blocks modalities, each conditioned on its own Qwen2.5-VL text branch and a CLIP pooled embedding. A separate super-resolution model upscales the generated video tile by tile in the latent space of a causal 3D K-VAE. -The distilled checkpoints use the few-step [`PiflowScheduler`] and must be run with `guidance_scale=1.0`. +> [!TIP] +> Check out the [Kandinsky Lab](https://huggingface.co/kandinskylab) organization on the Hub for the full set of +> official checkpoints, including flow-matching and distilled variants of both the base and super-resolution models. +> +> Distilled checkpoints ship with the few-step [`PiflowScheduler`] and must be run with `guidance_scale=1.0`. ## Available models @@ -94,6 +98,25 @@ resolutions or long videos: pipe.vae.enable_tiling() ``` +## Notes + +- `height` and `width` must be divisible by the video VAE's spatial compression ratio times the transformer's patch + size — `16` with the default [`Kandinsky6TI2VAPipeline`] configuration (`AutoencoderKLHunyuanVideo` at a + compression ratio of `8`, `patch_size=(1, 2, 2)`). `480x864`, used in the example above, satisfies this. +- [`Kandinsky6SRPipeline`]'s input `video` must have `1 + k * vae_scale_factor_temporal` frames for some integer `k` + (a temporal compression ratio of `4` with the default K-VAE configuration, so `121` frames works but `120` doesn't) + — trim or pad a video that doesn't already satisfy this before upscaling it. +- `visual_cond_scheme` is inferred automatically from whether `image` is passed (`"tail_cond_first_frame"` when it + is, `"pretrain"` otherwise). Only set it explicitly together with a matching `image`: passing + `visual_cond_scheme="i2v"` or `"tail_cond_first_frame"` without `image` doesn't raise an error, it silently falls + back to text-only generation. +- `expand_prompts=True` reuses the already-loaded Qwen2.5-VL text encoder for an extra generation pass before + denoising, so it adds latency but no extra model weights. +- Compile the repeated transformer blocks for faster repeated inference: + ```python + pipe.transformer.compile_repeated_blocks(fullgraph=True) + ``` + ## Kandinsky6TI2VAPipeline [[autodoc]] Kandinsky6TI2VAPipeline From 40fe659cb30c0af4b3925a71d4ea9de4971469e1 Mon Sep 17 00:00:00 2001 From: leffff Date: Wed, 30 Sep 2026 20:05:44 +0300 Subject: [PATCH 33/62] apply style chanes --- docs/source/en/api/models/kandinsky6_vae.md | 7 +++++++ docs/source/en/api/schedulers/piflow.md | 2 ++ src/diffusers/__init__.py | 6 +++--- src/diffusers/models/__init__.py | 4 ++-- src/diffusers/models/autoencoders/autoencoder_mmaudio.py | 5 ++++- .../models/transformers/transformer_kandinsky6.py | 8 ++++---- .../models/transformers/transformer_kandinsky6_sr.py | 4 ++-- src/diffusers/pipelines/kandinsky6/modeling_vocoder.py | 7 ++++++- src/diffusers/schedulers/__init__.py | 2 +- src/diffusers/schedulers/scheduling_piflow.py | 3 +++ 10 files changed, 34 insertions(+), 14 deletions(-) diff --git a/docs/source/en/api/models/kandinsky6_vae.md b/docs/source/en/api/models/kandinsky6_vae.md index b7c3924ff9dd..e72646c8bf2b 100644 --- a/docs/source/en/api/models/kandinsky6_vae.md +++ b/docs/source/en/api/models/kandinsky6_vae.md @@ -39,6 +39,9 @@ vae = Kandinsky6SRVAE.from_pretrained( The mel-spectrogram VAE used by [`Kandinsky6TI2VAPipeline`] when `sample_audio=True`. Its `decode` output is a mel spectrogram; pass it through [`~pipelines.kandinsky6.MMAudioVocoder`] to get a waveform. +The reference implementation can be found at [hkchengrex/MMAudio](https://github.com/hkchengrex/MMAudio) (MIT +license). + ```python import torch from diffusers import MMAudioVAE @@ -55,5 +58,9 @@ audio_vae = MMAudioVAE.from_pretrained( ## MMAudioVocoder +Adapted from the BigVGAN-v2 vocoder MMAudio bundles, itself from +[NVIDIA/BigVGAN](https://github.com/NVIDIA/BigVGAN) (MIT license), with the anti-aliased Snake activations of +[alias-free-torch](https://github.com/junjun3518/alias-free-torch) (Apache License 2.0). + [[autodoc]] pipelines.kandinsky6.MMAudioVocoder - forward diff --git a/docs/source/en/api/schedulers/piflow.md b/docs/source/en/api/schedulers/piflow.md index 03ff161a2eba..1845c05aa5cd 100644 --- a/docs/source/en/api/schedulers/piflow.md +++ b/docs/source/en/api/schedulers/piflow.md @@ -17,6 +17,8 @@ text/image-to-video-and-audio model and the video super-resolution model. It imp [π-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. +The reference implementation can be found at [Lakonik/LakonLab](https://github.com/Lakonik/LakonLab). + ## PiflowScheduler [[autodoc]] PiflowScheduler diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 5aedf783a997..b5be02858f75 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -316,12 +316,12 @@ "LTXVideoTransformer3DModel", "Lumina2Transformer2DModel", "LuminaNextDiT2DModel", - "MMAudioVAE", "MiniMaxH3Transformer3DModel", "MiniMaxMusic3ConditionEncoder", "MiniMaxMusic3RVQDepthDecoder", "MiniMaxMusic3Transformer1DModel", "MiniMaxMusic3Vocoder", + "MMAudioVAE", "MochiTransformer3DModel", "ModelMixin", "MotifVideoTransformer3DModel", @@ -460,8 +460,8 @@ "LCMScheduler", "LTXEulerAncestralRFScheduler", "MiniMaxH3Scheduler", - "PNDMScheduler", "PiflowScheduler", + "PNDMScheduler", "RePaintScheduler", "SASolverScheduler", "SchedulerMixin", @@ -750,10 +750,10 @@ "Lumina2Text2ImgPipeline", "LuminaPipeline", "LuminaText2ImgPipeline", - "MMAudioVocoder", "MarigoldDepthPipeline", "MarigoldIntrinsicsPipeline", "MarigoldNormalsPipeline", + "MMAudioVocoder", "MochiPipeline", "MotifVideoImage2VideoPipeline", "MotifVideoPipeline", diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 2d56ee1cdeb2..66fbec1bdb4b 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -31,12 +31,12 @@ _import_structure["autoencoders.autoencoder_asym_kl"] = ["AsymmetricAutoencoderKL"] _import_structure["autoencoders.autoencoder_cosmos3_audio"] = ["Cosmos3AVAEAudioTokenizer"] _import_structure["autoencoders.autoencoder_dc"] = ["AutoencoderDC"] + _import_structure["autoencoders.autoencoder_kandinsky6_sr"] = ["Kandinsky6SRVAE"] _import_structure["autoencoders.autoencoder_kl"] = ["AutoencoderKL"] _import_structure["autoencoders.autoencoder_kl_allegro"] = ["AutoencoderKLAllegro"] _import_structure["autoencoders.autoencoder_kl_cogvideox"] = ["AutoencoderKLCogVideoX"] _import_structure["autoencoders.autoencoder_kl_cosmos"] = ["AutoencoderKLCosmos"] _import_structure["autoencoders.autoencoder_kl_flux2"] = ["AutoencoderKLFlux2"] - _import_structure["autoencoders.autoencoder_kandinsky6_sr"] = ["Kandinsky6SRVAE"] _import_structure["autoencoders.autoencoder_kl_hunyuan_video"] = ["AutoencoderKLHunyuanVideo"] _import_structure["autoencoders.autoencoder_kl_hunyuanimage"] = ["AutoencoderKLHunyuanImage"] _import_structure["autoencoders.autoencoder_kl_hunyuanimage_refiner"] = ["AutoencoderKLHunyuanImageRefiner"] @@ -53,8 +53,8 @@ _import_structure["autoencoders.autoencoder_kl_qwenimage"] = ["AutoencoderKLQwenImage"] _import_structure["autoencoders.autoencoder_kl_temporal_decoder"] = ["AutoencoderKLTemporalDecoder"] _import_structure["autoencoders.autoencoder_kl_wan"] = ["AutoencoderKLWan"] - _import_structure["autoencoders.autoencoder_mmaudio"] = ["MMAudioVAE"] _import_structure["autoencoders.autoencoder_longcat_audio_dit"] = ["LongCatAudioDiTVae"] + _import_structure["autoencoders.autoencoder_mmaudio"] = ["MMAudioVAE"] _import_structure["autoencoders.autoencoder_oobleck"] = ["AutoencoderOobleck"] _import_structure["autoencoders.autoencoder_rae"] = ["AutoencoderRAE"] _import_structure["autoencoders.autoencoder_same"] = ["AutoencoderSAME"] diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index 147681fe56cc..32388325479f 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -1,4 +1,4 @@ -# Copyright 2025 The Kandinsky Team and The HuggingFace Team. All rights reserved. +# 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. @@ -12,6 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +# Adapted from the MMAudio VAE reference implementation (MIT license, Copyright (c) 2024 Sony Research Inc.): +# https://github.com/hkchengrex/MMAudio/tree/main/mmaudio/ext/autoencoder + """MMAudio mel-spectrogram VAE used by the Kandinsky 6 TI2VA pipeline. The BigVGAN vocoder that turns this VAE's decoded mel spectrograms into waveforms is a separate, pipeline-local diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 6a597613fb5a..0bd57c6c07e1 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -67,8 +67,8 @@ def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: def apply_scale_shift(normed: Tensor, x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: """Apply an AdaLN-style scale/shift affine to an already-normalized tensor, in fp32, cast back to ``x.dtype``. - Callers compute ``normed`` themselves (e.g. ``self.some_norm(x.float())``) so the norm-layer call stays visible - in `forward` instead of being hidden inside this helper. + Callers compute ``normed`` themselves (e.g. ``self.some_norm(x.float())``) so the norm-layer call stays visible in + `forward` instead of being hidden inside this helper. """ if x.ndim > 2 and scale.ndim == 2: shape = (scale.shape[0],) + (1,) * (x.ndim - 2) + (scale.shape[-1],) @@ -754,8 +754,8 @@ class Kandinsky6Transformer3DModel( visual_token_type_num_embeddings (`int`, *optional*, defaults to 0): Number of visual token type embeddings. Released checkpoints set `text_token_padding`, `ca_rope`, `cross_gates`, and `fix_modulation` to `True` (see each - checkpoint's `transformer/config.json`); the `False` defaults only describe an architecture variant this repo - does not ship weights for. + checkpoint's `transformer/config.json`); the `False` defaults only describe an architecture variant this repo does + not ship weights for. """ _repeated_blocks = [ diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index cb43ddf771bf..58cc67ab755a 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -55,8 +55,8 @@ def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: def apply_scale_shift(normed: Tensor, x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: """Apply an AdaLN-style scale/shift affine to an already-normalized tensor, in fp32, cast back to ``x.dtype``. - Callers compute ``normed`` themselves (e.g. ``self.some_norm(x.float())``) so the norm-layer call stays visible - in `forward` instead of being hidden inside this helper. + Callers compute ``normed`` themselves (e.g. ``self.some_norm(x.float())``) so the norm-layer call stays visible in + `forward` instead of being hidden inside this helper. """ if x.ndim > 2 and scale.ndim == 2: shape = (scale.shape[0],) + (1,) * (x.ndim - 2) + (scale.shape[-1],) diff --git a/src/diffusers/pipelines/kandinsky6/modeling_vocoder.py b/src/diffusers/pipelines/kandinsky6/modeling_vocoder.py index f5470ba860e2..42132ba0b0d1 100644 --- a/src/diffusers/pipelines/kandinsky6/modeling_vocoder.py +++ b/src/diffusers/pipelines/kandinsky6/modeling_vocoder.py @@ -1,4 +1,4 @@ -# Copyright 2025 The Kandinsky Team and The HuggingFace Team. All rights reserved. +# 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. @@ -12,6 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +# Adapted from the BigVGAN-v2 vocoder MMAudio bundles at +# https://github.com/hkchengrex/MMAudio/tree/main/mmaudio/ext/bigvgan_v2, itself adapted from +# https://github.com/NVIDIA/BigVGAN (MIT license), with the anti-aliased Snake activations of +# https://github.com/junjun3518/alias-free-torch (Apache License 2.0). + """BigVGAN vocoder that turns the mel spectrograms decoded by [`MMAudioVAE`] into waveforms.""" from __future__ import annotations diff --git a/src/diffusers/schedulers/__init__.py b/src/diffusers/schedulers/__init__.py index 3bcb1698f148..47f4a0857090 100644 --- a/src/diffusers/schedulers/__init__.py +++ b/src/diffusers/schedulers/__init__.py @@ -61,7 +61,6 @@ _import_structure["scheduling_euler_ancestral_discrete"] = ["EulerAncestralDiscreteScheduler"] _import_structure["scheduling_euler_discrete"] = ["EulerDiscreteScheduler"] _import_structure["scheduling_flow_map_euler_discrete"] = ["FlowMapEulerDiscreteScheduler"] - _import_structure["scheduling_piflow"] = ["PiflowScheduler"] _import_structure["scheduling_flow_match_euler_discrete"] = ["FlowMatchEulerDiscreteScheduler"] _import_structure["scheduling_flow_match_heun_discrete"] = ["FlowMatchHeunDiscreteScheduler"] _import_structure["scheduling_flow_match_lcm"] = ["FlowMatchLCMScheduler"] @@ -74,6 +73,7 @@ _import_structure["scheduling_lcm"] = ["LCMScheduler"] _import_structure["scheduling_ltx_euler_ancestral_rf"] = ["LTXEulerAncestralRFScheduler"] _import_structure["scheduling_minimax_h3"] = ["MiniMaxH3Scheduler"] + _import_structure["scheduling_piflow"] = ["PiflowScheduler"] _import_structure["scheduling_pndm"] = ["PNDMScheduler"] _import_structure["scheduling_repaint"] = ["RePaintScheduler"] _import_structure["scheduling_sasolver"] = ["SASolverScheduler"] diff --git a/src/diffusers/schedulers/scheduling_piflow.py b/src/diffusers/schedulers/scheduling_piflow.py index 181ffcc35e0a..d858b425a9fd 100644 --- a/src/diffusers/schedulers/scheduling_piflow.py +++ b/src/diffusers/schedulers/scheduling_piflow.py @@ -12,6 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +# DISCLAIMER: This file is strongly influenced by the π-Flow reference implementation at +# https://github.com/Lakonik/LakonLab (https://huggingface.co/papers/2510.14974) + """Diffusers scheduler for distilled Kandinsky 6 PiFlow checkpoints.""" from dataclasses import dataclass From 28c5326da8cb2e6fd262da2207c1d09ec12277e2 Mon Sep 17 00:00:00 2001 From: leffff Date: Thu, 1 Oct 2026 14:03:58 +0300 Subject: [PATCH 34/62] address review feedback on transformer, SR VAE, and flex attention Co-Authored-By: Claude Sonnet 5 --- .../en/api/models/kandinsky6_transformer.md | 4 - src/diffusers/models/attention_dispatch.py | 14 +- .../autoencoders/autoencoder_kandinsky6_sr.py | 15 +- .../transformers/transformer_kandinsky6.py | 143 +++++++----------- .../transformers/transformer_kandinsky6_sr.py | 42 +++-- 5 files changed, 85 insertions(+), 133 deletions(-) diff --git a/docs/source/en/api/models/kandinsky6_transformer.md b/docs/source/en/api/models/kandinsky6_transformer.md index 7684a54a670c..1425c01fd07c 100644 --- a/docs/source/en/api/models/kandinsky6_transformer.md +++ b/docs/source/en/api/models/kandinsky6_transformer.md @@ -32,10 +32,6 @@ transformer = Kandinsky6Transformer3DModel.from_pretrained( - all - forward -## Kandinsky6Transformer3DModelOutput - -[[autodoc]] models.transformers.transformer_kandinsky6.Kandinsky6Transformer3DModelOutput - ## Kandinsky6SRTransformer3DModel The text-free transformer used by [`Kandinsky6SRPipeline`] to refine one tile of the upscaled video at a time. diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 9e08525042cd..74aef0d8536e 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -3620,12 +3620,20 @@ def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): # A `BlockMask`'s own block size is the sparsity granularity the caller built it around (e.g. NABLA's 64-token # fractal blocks). The forward kernel's tile sizes must divide that granularity evenly; torch's autotuned - # defaults aren't guaranteed to (they're picked independent of the mask), so pin them to the mask's own block - # size, which always divides itself. + # defaults are only overridden when they don't already divide the mask's block size, to avoid unintended side + # effects from pinning tile sizes that otherwise would have been fine (e.g. the default 128x128 mask size would + # otherwise always force `BLOCK_N` to 128, even though 64 is a valid size). kernel_options = None if block_mask is not None: q_block_size, kv_block_size = block_mask.BLOCK_SIZE - kernel_options = {"BLOCK_M": q_block_size, "BLOCK_N": kv_block_size} + default_block_m = 64 if query.dtype == torch.float32 else 128 + default_block_n = 64 + kernel_options = {} + if q_block_size % default_block_m != 0: + kernel_options["BLOCK_M"] = q_block_size + if kv_block_size % default_block_n != 0: + kernel_options["BLOCK_N"] = kv_block_size + kernel_options = kernel_options or None query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) out = flex_attention.flex_attention( diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index 47731d00e008..6f737a7980a4 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -275,12 +275,7 @@ def forward( class Kandinsky6SRDownsample(nn.Module): """Spatial 2x downsample (strided conv plus pixel-unshuffle average) with an optional causal temporal 2x. - Implemented from scratch rather than reusing `KVAECachedPXSDownsample` (`autoencoder_kl_kvae_video.py`), the - closest existing causal, stateful down/upsample block in the library: that class is a private helper of a different - Kandinsky VAE (not a shared, importable utility) and uses its own cache convention — a `list` of per-conv dicts - mutated in place — rather than this file's dict-keyed cache that `forward` returns functionally instead of mutating - (see `Kandinsky6SRCausalConv3d`), so reusing it as-is would mean threading two incompatible cache conventions - through the same model. See `Kandinsky6SRUpsample` for the upsampling counterpart. + See `Kandinsky6SRUpsample` for the upsampling counterpart. """ def __init__(self, in_channels: int, compress_time: bool) -> None: @@ -338,9 +333,7 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: class Kandinsky6SRUpsample(nn.Module): """Spatial 2x nearest upsample with a convolutional residual, preceded by an optional causal temporal 2x. - Bespoke for the same reason as `Kandinsky6SRDownsample`: the closest existing precedent, `KVAECachedPXSUpsample` - (`autoencoder_kl_kvae_video.py`), is a private helper of a different Kandinsky VAE with its own, incompatible cache - convention, not a shared utility this file can import and reuse. + See `Kandinsky6SRDownsample` for the downsampling counterpart. """ def __init__(self, channels: int, compress_time: bool) -> None: @@ -370,8 +363,6 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: class Kandinsky6SREncoder3D(nn.Module): - _supports_gradient_checkpointing = True - def __init__( self, in_channels: int, @@ -456,8 +447,6 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: class Kandinsky6SRDecoder3D(nn.Module): - _supports_gradient_checkpointing = True - def __init__( self, out_channels: int, diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 0bd57c6c07e1..b1060f1b16fe 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -17,7 +17,6 @@ from __future__ import annotations import math -from dataclasses import dataclass from typing import Any import torch @@ -26,28 +25,12 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin -from ...utils import BaseOutput from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn from ..cache_utils import CacheMixin -from ..embeddings import get_timestep_embedding +from ..embeddings import TimestepEmbedding, Timesteps from ..modeling_utils import ModelMixin, get_parameter_dtype - - -@dataclass -class Kandinsky6Transformer3DModelOutput(BaseOutput): - r""" - Output of [`Kandinsky6Transformer3DModel`]. - - Args: - sample (`torch.Tensor` of shape `(batch_size, num_frames, height, width, out_visual_dim)`): - The predicted video velocity, in the unpatchified `(B, T, H, W, C)` latent layout. - audio_sample (`torch.Tensor` of shape `(batch_size, audio_length, out_audio_dim)`, *optional*): - The predicted audio velocity. `None` when no `audio_hidden_states` were passed. - """ - - sample: torch.Tensor - audio_sample: torch.Tensor | None = None +from .transformer_ltx2 import AudioVisualModelOutput def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: @@ -114,7 +97,8 @@ def nabla_block_mask( thr: float = 0.9, block_size: int = 64, ) -> BlockMask: - """Build a dynamic NABLA block mask from query/key statistics and an STA prior.""" + """Build a dynamic NABLA block mask from query/key statistics and an STA (Sliding-Tile Attention, see + https://huggingface.co/papers/2502.04507) prior.""" B, h, S, D = q.shape s1 = S // block_size qa = q.reshape(B, h, s1, block_size, D).mean(-2) @@ -153,12 +137,12 @@ def __init__( self.max_period = max_period self.freqs_scaling = freqs_scaling freq = get_freqs(dim // 2, max_period) * freqs_scaling - self.register_buffer("args", torch.outer(torch.arange(max_pos, dtype=freq.dtype), freq), persistent=False) + self.register_buffer("angles", torch.outer(torch.arange(max_pos, dtype=freq.dtype), freq), persistent=False) def forward(self, pos: Tensor) -> Tensor: # RoPE tables are fp32; keep trig in fp32. - args = self.args[pos] # (seq_len, dim//2) - rope = torch.stack([torch.cos(args), -torch.sin(args), torch.sin(args), torch.cos(args)], dim=-1) + angles = self.angles[pos] # (seq_len, dim//2) + rope = torch.stack([torch.cos(angles), -torch.sin(angles), torch.sin(angles), torch.cos(angles)], dim=-1) return rope.view(*rope.shape[:-1], 2, 2).unsqueeze(-4) @@ -177,7 +161,9 @@ def __init__( self.max_period = max_period for i, (d, mp) in enumerate(zip(axes_dims, max_pos)): freq = get_freqs(d // 2, max_period) - self.register_buffer(f"args_{i}", torch.outer(torch.arange(mp, dtype=freq.dtype), freq), persistent=False) + self.register_buffer( + f"angles_{i}", torch.outer(torch.arange(mp, dtype=freq.dtype), freq), persistent=False + ) def forward( self, @@ -186,40 +172,24 @@ def forward( ) -> Tensor: # `pos` holds one index tensor per axis; the grid size is implied by their lengths. num_frames, height, width = (int(axis_pos.shape[0]) for axis_pos in pos) - args_t = self.args_0[pos[0]] / scale_factor[0] # (T, d//2) - args_h = self.args_1[pos[1]] / scale_factor[1] # (H, d//2) - args_w = self.args_2[pos[2]] / scale_factor[2] # (W, d//2) + angles_t = self.angles_0[pos[0]] / scale_factor[0] # (T, d//2) + angles_h = self.angles_1[pos[1]] / scale_factor[1] # (H, d//2) + angles_w = self.angles_2[pos[2]] / scale_factor[2] # (W, d//2) - args = torch.cat( + angles = torch.cat( [ - args_t.view(num_frames, 1, 1, -1).expand(num_frames, height, width, -1), - args_h.view(1, height, 1, -1).expand(num_frames, height, width, -1), - args_w.view(1, 1, width, -1).expand(num_frames, height, width, -1), + angles_t.view(num_frames, 1, 1, -1).expand(num_frames, height, width, -1), + angles_h.view(1, height, 1, -1).expand(num_frames, height, width, -1), + angles_w.view(1, 1, width, -1).expand(num_frames, height, width, -1), ], dim=-1, ) - cos, sin = torch.cos(args), torch.sin(args) + cos, sin = torch.cos(angles), torch.sin(angles) rope = torch.stack([cos, -sin, sin, cos], dim=-1) # (T, H, W, total_dim, 4) rope = rope.view(*rope.shape[:-1], 2, 2) # (T, H, W, total_dim, 2, 2) return rope.unsqueeze(-4) # (T, H, W, 1, total_dim, 2, 2) -# Diffusers-style K6 transformer components. -# -# The classes use the same inner-module vocabulary as the Diffusers Kandinsky5 -# transformer (for example ``in_layer``, ``modulation``, ``self_attention`` and -# ``feed_forward``). Checkpoint conversion maps the native K6 names to this -# public Diffusers layout. - - -_MASKED_ATTENTION_BACKENDS = { - "flash": "flash_varlen", - "_flash_3": "_flash_varlen_3", - "sage": "sage_varlen", - "native": "native", -} - - class Kandinsky6AttnProcessor: """Diffusers attention processor used by the TI2VA transformer. @@ -231,7 +201,6 @@ class Kandinsky6AttnProcessor: _parallel_config = None def __init__(self, attention_backend=None, parallel_config=None): - self._masked = False self._attention_backend = attention_backend self._parallel_config = parallel_config @@ -297,24 +266,19 @@ def __call__( class Kandinsky6TimeEmbeddings(nn.Module): """Sinusoidal timestep embedding with a K6-compatible parameter layout.""" - def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): + def __init__(self, model_dim: int, time_dim: int): super().__init__() if model_dim % 2: raise ValueError("model_dim must be even") - self.model_dim = model_dim - self.max_period = max_period - self.in_layer = nn.Linear(model_dim, time_dim) - self.activation = nn.SiLU() - self.out_layer = nn.Linear(time_dim, time_dim) + self.time_proj = Timesteps(num_channels=model_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=model_dim, time_embed_dim=time_dim, act_fn="silu") def forward(self, timestep: Tensor) -> Tensor: # The sinusoidal embedding is float32; `_keep_in_fp32_modules` keeps these layers float32 under # `from_pretrained(torch_dtype=...)`, and the cast aligns the input with whatever dtype they hold. - embed = get_timestep_embedding( - timestep, self.model_dim, flip_sin_to_cos=True, downscale_freq_shift=0, max_period=self.max_period - ) - embed = embed.to(get_parameter_dtype(self.in_layer)) - return self.out_layer(self.activation(self.in_layer(embed))) + embed = self.time_proj(timestep) + embed = embed.to(get_parameter_dtype(self.timestep_embedder)) + return self.timestep_embedder(embed) class Kandinsky6TextEmbeddings(nn.Module): @@ -403,7 +367,6 @@ def __init__( self.text_token_padding = text_token_padding self.set_processor(processor or self._default_processor_cls()) if self.text_token_padding: - self.processor._masked = True self.processor._attention_backend = AttentionBackendName.NATIVE def forward( @@ -568,8 +531,10 @@ def __init__( fix_modulation: bool = False, ): super().__init__() - self.videoT = Kandinsky6TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim, text_token_padding) - self.audioT = Kandinsky6TransformerDecoderBlock( + self.video_dec_block = Kandinsky6TransformerDecoderBlock( + model_dim, time_dim, ff_dim, head_dim, text_token_padding + ) + self.audio_dec_block = Kandinsky6TransformerDecoderBlock( model_dim_a, time_dim_a, ff_dim_a, head_dim_a, text_token_padding ) self.va_cross_attention = Kandinsky6Attention( @@ -614,39 +579,39 @@ def forward( ) -> tuple[Tensor | None, Tensor | None]: t_v, t_a = time_embed if vis is not None: - sa_p, ca_p, ff_p = torch.chunk(self.videoT.visual_modulation(t_v), 3, dim=-1) + sa_p, ca_p, ff_p = torch.chunk(self.video_dec_block.visual_modulation(t_v), 3, dim=-1) shift, scale, gate = torch.chunk(sa_p, 3, dim=-1) vis = apply_gate_sum( vis, - self.videoT.self_attention( - apply_scale_shift(self.videoT.self_attention_norm(vis.float()), vis, scale, shift), + self.video_dec_block.self_attention( + apply_scale_shift(self.video_dec_block.self_attention_norm(vis.float()), vis, scale, shift), rotary_emb=vis_rope, sparse_params=sparse_params, ), gate, ).type_as(vis) shift, scale, gate_v = torch.chunk(ca_p, 3, dim=-1) - vis_pre_ca = apply_scale_shift(self.videoT.cross_attention_norm(vis.float()), vis, scale, shift) - vis_out_t = self.videoT.cross_attention( + vis_pre_ca = apply_scale_shift(self.video_dec_block.cross_attention_norm(vis.float()), vis, scale, shift) + vis_out_t = self.video_dec_block.cross_attention( vis_pre_ca, encoder_hidden_states=text_v, attn_mask=attn_mask, ) if aud is not None: - sa_p, ca_p, ff_p_a = torch.chunk(self.audioT.visual_modulation(t_a), 3, dim=-1) + sa_p, ca_p, ff_p_a = torch.chunk(self.audio_dec_block.visual_modulation(t_a), 3, dim=-1) shift, scale, gate = torch.chunk(sa_p, 3, dim=-1) aud = apply_gate_sum( aud, - self.audioT.self_attention( - apply_scale_shift(self.audioT.self_attention_norm(aud.float()), aud, scale, shift), + self.audio_dec_block.self_attention( + apply_scale_shift(self.audio_dec_block.self_attention_norm(aud.float()), aud, scale, shift), rotary_emb=aud_rope, ), gate, ).type_as(aud) shift, scale, gate_a = torch.chunk(ca_p, 3, dim=-1) - aud_pre_ca = apply_scale_shift(self.audioT.cross_attention_norm(aud.float()), aud, scale, shift) - aud_out_t = self.audioT.cross_attention( + aud_pre_ca = apply_scale_shift(self.audio_dec_block.cross_attention_norm(aud.float()), aud, scale, shift) + aud_out_t = self.audio_dec_block.cross_attention( aud_pre_ca, encoder_hidden_states=text_a, attn_mask=attn_mask, @@ -694,8 +659,8 @@ def forward( shift, scale, gate = torch.chunk(ff_p, 3, dim=-1) vis = apply_gate_sum( vis, - self.videoT.feed_forward( - apply_scale_shift(self.videoT.feed_forward_norm(vis.float()), vis, scale, shift) + self.video_dec_block.feed_forward( + apply_scale_shift(self.video_dec_block.feed_forward_norm(vis.float()), vis, scale, shift) ), gate, ).type_as(vis) @@ -703,8 +668,8 @@ def forward( shift, scale, gate = torch.chunk(ff_p_a, 3, dim=-1) aud = apply_gate_sum( aud, - self.audioT.feed_forward( - apply_scale_shift(self.audioT.feed_forward_norm(aud.float()), aud, scale, shift) + self.audio_dec_block.feed_forward( + apply_scale_shift(self.audio_dec_block.feed_forward_norm(aud.float()), aud, scale, shift) ), gate, ).type_as(aud) @@ -875,7 +840,7 @@ def forward( visual_token_type_ids: Tensor | None = None, sparse_params: dict | None = None, return_dict: bool = True, - ) -> Kandinsky6Transformer3DModelOutput | tuple[Tensor, ...]: + ) -> AudioVisualModelOutput | tuple[Tensor, ...]: r""" Args: hidden_states (`torch.Tensor` of shape `(batch_size, num_frames, height, width, in_channels)`): @@ -899,12 +864,17 @@ def forward( Per-frame token type ids, embedded through `visual_token_type_embeddings`. Requires `visual_token_type_num_embeddings > 0`. sparse_params (`dict`, *optional*): - NABLA sparse-attention configuration for the video self-attention. + NABLA sparse-attention configuration for the video self-attention, dispatched through the `flex` + attention backend regardless of `transformer.set_attention_backend(...)`. This trades the regular + SDPA backend's dense attention for a block-sparse pattern that scales better to longer sequences, at + the cost of FlexAttention's one-time kernel-autotuning overhead on the first call. None of the + currently released checkpoints were trained with NABLA; this is forward-looking support for a future, + longer-duration (e.g. 10s) model. return_dict (`bool`, defaults to `True`): - Whether to return a [`Kandinsky6Transformer3DModelOutput`] instead of a plain tuple. + Whether to return an [`AudioVisualModelOutput`] instead of a plain tuple. Returns: - [`Kandinsky6Transformer3DModelOutput`] or `tuple`: + [`AudioVisualModelOutput`] or `tuple`: The predicted video velocity in the `(B, T, H, W, out_visual_dim)` layout and, when `audio_hidden_states` was given, the predicted audio velocity of shape `(B, audio_length, out_audio_dim)`. @@ -985,13 +955,4 @@ def forward( if not return_dict: return (video_out,) if audio_out is None else (video_out, audio_out) - return Kandinsky6Transformer3DModelOutput(sample=video_out, audio_sample=audio_out) - - -__all__ = [ - "Kandinsky6Transformer3DModel", - "Kandinsky6Transformer3DModelOutput", - "Kandinsky6TransformerEncoderBlock", - "Kandinsky6TransformerDecoderBlock", - "Kandinsky6FusedTransformerDecoderBlock", -] + return AudioVisualModelOutput(sample=video_out, audio_sample=audio_out) diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index 58cc67ab755a..7a535fa44d32 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -27,7 +27,7 @@ from ...loaders import PeftAdapterMixin from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn -from ..embeddings import get_timestep_embedding +from ..embeddings import TimestepEmbedding, Timesteps from ..modeling_outputs import Transformer2DModelOutput from ..modeling_utils import ModelMixin, get_parameter_dtype @@ -107,7 +107,8 @@ def nabla_block_mask( thr: float = 0.9, block_size: int = 64, ) -> BlockMask: - """Build a dynamic NABLA block mask from query/key statistics and an STA prior.""" + """Build a dynamic NABLA block mask from query/key statistics and an STA (Sliding-Tile Attention, see + https://huggingface.co/papers/2502.04507) prior.""" B, h, S, D = q.shape s1 = S // block_size qa = q.reshape(B, h, s1, block_size, D).mean(-2) @@ -233,24 +234,19 @@ def forward( class Kandinsky6SRTimeEmbeddings(nn.Module): """Sinusoidal timestep embedding with a K6-compatible parameter layout.""" - def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): + def __init__(self, model_dim: int, time_dim: int): super().__init__() if model_dim % 2: raise ValueError("model_dim must be even") - self.model_dim = model_dim - self.max_period = max_period - self.in_layer = nn.Linear(model_dim, time_dim) - self.activation = nn.SiLU() - self.out_layer = nn.Linear(time_dim, time_dim) + self.time_proj = Timesteps(num_channels=model_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(in_channels=model_dim, time_embed_dim=time_dim, act_fn="silu") def forward(self, timestep: Tensor) -> Tensor: # The sinusoidal embedding is float32; `_keep_in_fp32_modules` keeps these layers float32 under # `from_pretrained(torch_dtype=...)`, and the cast aligns the input with whatever dtype they hold. - embed = get_timestep_embedding( - timestep, self.model_dim, flip_sin_to_cos=True, downscale_freq_shift=0, max_period=self.max_period - ) - embed = embed.to(get_parameter_dtype(self.in_layer)) - return self.out_layer(self.activation(self.in_layer(embed))) + embed = self.time_proj(timestep) + embed = embed.to(get_parameter_dtype(self.timestep_embedder)) + return self.timestep_embedder(embed) # Copied from diffusers.models.transformers.transformer_kandinsky6.Kandinsky6VisualEmbeddings with Kandinsky6->Kandinsky6SR @@ -351,7 +347,9 @@ def __init__( self.max_period = max_period for i, (d, mp) in enumerate(zip(axes_dims, max_pos)): freq = get_freqs(d // 2, max_period) - self.register_buffer(f"args_{i}", torch.outer(torch.arange(mp, dtype=freq.dtype), freq), persistent=False) + self.register_buffer( + f"angles_{i}", torch.outer(torch.arange(mp, dtype=freq.dtype), freq), persistent=False + ) def forward( self, @@ -360,19 +358,19 @@ def forward( ) -> Tensor: # `pos` holds one index tensor per axis; the grid size is implied by their lengths. num_frames, height, width = (int(axis_pos.shape[0]) for axis_pos in pos) - args_t = self.args_0[pos[0]] / scale_factor[0] # (T, d//2) - args_h = self.args_1[pos[1]] / scale_factor[1] # (H, d//2) - args_w = self.args_2[pos[2]] / scale_factor[2] # (W, d//2) + angles_t = self.angles_0[pos[0]] / scale_factor[0] # (T, d//2) + angles_h = self.angles_1[pos[1]] / scale_factor[1] # (H, d//2) + angles_w = self.angles_2[pos[2]] / scale_factor[2] # (W, d//2) - args = torch.cat( + angles = torch.cat( [ - args_t.view(num_frames, 1, 1, -1).expand(num_frames, height, width, -1), - args_h.view(1, height, 1, -1).expand(num_frames, height, width, -1), - args_w.view(1, 1, width, -1).expand(num_frames, height, width, -1), + angles_t.view(num_frames, 1, 1, -1).expand(num_frames, height, width, -1), + angles_h.view(1, height, 1, -1).expand(num_frames, height, width, -1), + angles_w.view(1, 1, width, -1).expand(num_frames, height, width, -1), ], dim=-1, ) - cos, sin = torch.cos(args), torch.sin(args) + cos, sin = torch.cos(angles), torch.sin(angles) rope = torch.stack([cos, -sin, sin, cos], dim=-1) # (T, H, W, total_dim, 4) rope = rope.view(*rope.shape[:-1], 2, 2) # (T, H, W, total_dim, 2, 2) return rope.unsqueeze(-4) # (T, H, W, 1, total_dim, 2, 2) From 372f29088c4e7c1db26ad69802d73c378e20461b Mon Sep 17 00:00:00 2001 From: leffff Date: Thu, 1 Oct 2026 18:48:59 +0300 Subject: [PATCH 35/62] address comments --- .../autoencoders/autoencoder_kandinsky6_sr.py | 19 +++++++++ .../autoencoders/autoencoder_mmaudio.py | 30 +++++++++----- .../transformers/transformer_kandinsky6_sr.py | 3 -- .../kandinsky6/modeling_latent_upscaler.py | 40 ++++++++----------- .../pipelines/kandinsky6/modeling_vocoder.py | 6 +-- .../kandinsky6/pipeline_kandinsky6_ti2va.py | 12 ++++-- 6 files changed, 68 insertions(+), 42 deletions(-) diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index 6f737a7980a4..c17da9b26fc9 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -363,6 +363,16 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: class Kandinsky6SREncoder3D(nn.Module): + """Causal encoder: `conv_in`, downsampling resnet levels, a resnet bottleneck, then `norm_out`/`conv_out`. + + `forward` only returns `hidden_states`, not `cache`: each submodule still returns its own updated cache value + rather than mutating the one it was given (see `Kandinsky6SRCausalConv3d`), but this method writes every one of + those return values back into the matching key of the `cache` dict it received, e.g. + `hidden_states, cache["conv_out"] = self.conv_out(hidden_states, cache["conv_out"])`. Because the caller + (`Kandinsky6SRVAE.encode`) passes the same `cache` dict to every segment, that key assignment is what carries + the padding state from one segment to the next — `cache` does not need to be returned for this to work. + """ + def __init__( self, in_channels: int, @@ -447,6 +457,15 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: class Kandinsky6SRDecoder3D(nn.Module): + """Causal decoder: `conv_in`, a `zq`-conditioned resnet bottleneck, upsampling resnet levels, then + `norm_out`/`conv_out`. + + Like `Kandinsky6SREncoder3D.forward`, this only returns the decoded tensor, not `cache`: writing each + submodule's returned cache value back into the matching key of the `cache` dict it received (e.g. + `hidden_states, cache["conv_out"] = self.conv_out(hidden_states, cache["conv_out"])`) is enough, since + `Kandinsky6SRVAE.decode` passes that same dict to every segment. + """ + def __init__( self, out_channels: int, diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index 32388325479f..566479ce8aa3 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -31,6 +31,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...utils.accelerate_utils import apply_forward_hook +from ..attention import AttentionModuleMixin from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin @@ -131,16 +132,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return mp_sum(x, hidden_states, t=0.3) -class MMAudioAttnBlock1D(nn.Module): - def __init__(self, channels: int, num_heads: int = 1) -> None: - super().__init__() - self.num_heads = num_heads - self.qkv = MMAudioMPConv1d(channels, channels * 3, kernel_size=1) - self.proj_out = MMAudioMPConv1d(channels, channels, kernel_size=1) +class MMAudioAttnProcessor: + """Attention processor used by [`MMAudioAttnBlock1D`].""" - def forward(self, x: torch.Tensor) -> torch.Tensor: + def __call__(self, attn: "MMAudioAttnBlock1D", x: torch.Tensor) -> torch.Tensor: batch_size, channels, length = x.shape - qkv = self.qkv(x).reshape(batch_size, self.num_heads, -1, 3, length) + qkv = attn.qkv(x).reshape(batch_size, attn.num_heads, -1, 3, length) query, key, value = normalize(qkv, dim=2).unbind(3) # `(B, heads, D, T)` -> `(B, T, heads, D)` for the attention dispatcher. `D` is not contiguous after the # permute (q/k/v are interleaved along the channel dim), which some attention backends (e.g. FlashAttention-3) @@ -152,7 +149,22 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # pipeline (e.g. via `transformer.set_attention_backend(...)`). hidden_states = dispatch_attention_fn(query, key, value, backend=AttentionBackendName.NATIVE) hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_size, channels, length) - return mp_sum(x, self.proj_out(hidden_states), t=0.3) + return attn.proj_out(hidden_states) + + +class MMAudioAttnBlock1D(nn.Module, AttentionModuleMixin): + _default_processor_cls = MMAudioAttnProcessor + _available_processors = [MMAudioAttnProcessor] + + def __init__(self, channels: int, num_heads: int = 1, processor: MMAudioAttnProcessor | None = None) -> None: + super().__init__() + self.num_heads = num_heads + self.qkv = MMAudioMPConv1d(channels, channels * 3, kernel_size=1) + self.proj_out = MMAudioMPConv1d(channels, channels, kernel_size=1) + self.set_processor(processor or self._default_processor_cls()) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return mp_sum(x, self.processor(self, x), t=0.3) class MMAudioUpsample1D(nn.Module): diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index 7a535fa44d32..e677327010ce 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -565,6 +565,3 @@ def forward( if not return_dict: return (output,) return Transformer2DModelOutput(sample=output) - - -__all__ = ["Kandinsky6SRTransformer3DModel"] diff --git a/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py b/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py index 104ecb2fd27c..64186ed13a5b 100644 --- a/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py +++ b/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py @@ -36,10 +36,9 @@ def __init__(self, in_channels: int, out_channels: int, kernel_size: int) -> Non self.temporal_pad = kernel_size // 2 def forward(self, hidden_states: Tensor) -> Tensor: - if self.temporal_pad: - hidden_states = functional.pad( - hidden_states, (0, 0, 0, 0, self.temporal_pad, self.temporal_pad), mode="replicate" - ) + hidden_states = functional.pad( + hidden_states, (0, 0, 0, 0, self.temporal_pad, self.temporal_pad), mode="replicate" + ) return super().forward(hidden_states) @@ -105,19 +104,17 @@ def forward(self, hidden_states: Tensor) -> Tensor: return self.linear(hidden_states + self.spatial_conv(hidden_states)) -class Kandinsky6SRLatentUpscalerOutputHead(nn.Sequential): +class Kandinsky6SRLatentUpscalerOutputHead(nn.Module): """`modulated norm -> SiLU -> conv3x3x3` projection back to the latent channels, threading `zq` into the norm.""" def __init__(self, in_channels: int, out_channels: int, zq_channels: int) -> None: - super().__init__( - Kandinsky6SRLatentUpscalerModulatedNorm(in_channels, zq_channels), - nn.SiLU(), - Kandinsky6SRLatentUpscalerConv3d(in_channels, out_channels, kernel_size=3), - ) + super().__init__() + self.norm = Kandinsky6SRLatentUpscalerModulatedNorm(in_channels, zq_channels) + self.activation = nn.SiLU() + self.conv = Kandinsky6SRLatentUpscalerConv3d(in_channels, out_channels, kernel_size=3) def forward(self, hidden_states: Tensor, zq: Tensor) -> Tensor: - norm, activation, conv = self - return conv(activation(norm(hidden_states, zq))) + return self.conv(self.activation(self.norm(hidden_states, zq))) def _residual_stack(num_blocks: int, in_channels: int, out_channels: int, zq_channels: int) -> nn.ModuleList: @@ -148,21 +145,21 @@ def __init__( self.finisher = nn.Module() self.finisher.spatial_conv = nn.Conv3d(width_1, width_1, kernel_size=(1, 3, 3), padding=(0, 1, 1)) self.finisher.linear = nn.Conv3d(width_1, width_1, kernel_size=1) - self.private_mid_blocks = _residual_stack(num_mid_blocks, width_1, width_2, in_channels) - self.private_upsample = Kandinsky6SRLatentUpscalerUpsample(width_2) - self.private_blocks = _residual_stack(num_post_blocks, width_2, width_3, in_channels) - self.private_output_proj = Kandinsky6SRLatentUpscalerOutputHead(width_3, in_channels, in_channels) + self.mid_blocks = _residual_stack(num_mid_blocks, width_1, width_2, in_channels) + self.upsample = Kandinsky6SRLatentUpscalerUpsample(width_2) + self.blocks = _residual_stack(num_post_blocks, width_2, width_3, in_channels) + self.output_proj = Kandinsky6SRLatentUpscalerOutputHead(width_3, in_channels, in_channels) def forward(self, hidden_states: Tensor, zq: Tensor) -> Tensor: for block in self.adapter: hidden_states = block(hidden_states, zq) hidden_states = self.finisher.linear(hidden_states + self.finisher.spatial_conv(hidden_states)) - for block in self.private_mid_blocks: + for block in self.mid_blocks: hidden_states = block(hidden_states, zq) - hidden_states = self.private_upsample(hidden_states) - for block in self.private_blocks: + hidden_states = self.upsample(hidden_states) + for block in self.blocks: hidden_states = block(hidden_states, zq) - return self.private_output_proj(hidden_states, zq) + return self.output_proj(hidden_states, zq) class Kandinsky6SRLatentUpscaler(nn.Module): @@ -292,6 +289,3 @@ def forward(self, latents: Tensor, scale: int, return_dict: bool = True) -> Deco if not return_dict: return (upscaled,) return DecoderOutput(sample=upscaled) - - -__all__ = ["Kandinsky6SRLatentUpscalerBank"] diff --git a/src/diffusers/pipelines/kandinsky6/modeling_vocoder.py b/src/diffusers/pipelines/kandinsky6/modeling_vocoder.py index 42132ba0b0d1..0121c7547a92 100644 --- a/src/diffusers/pipelines/kandinsky6/modeling_vocoder.py +++ b/src/diffusers/pipelines/kandinsky6/modeling_vocoder.py @@ -47,9 +47,9 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size - filter_ = 2 * cutoff * window * torch.sinc(2 * cutoff * time) - filter_ = filter_ / filter_.sum() - return filter_.view(1, 1, kernel_size) + filter = 2 * cutoff * window * torch.sinc(2 * cutoff * time) + filter = filter / filter.sum() + return filter.view(1, 1, kernel_size) class MMAudioSnakeBeta(nn.Module): diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index f8c8eb48859f..9c344b47cf1f 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -522,12 +522,14 @@ def encode_image( width: int, device: torch.device, dtype: torch.dtype, + num_videos_per_prompt: int = 1, generator: torch.Generator | None = None, ) -> torch.Tensor: r""" Encodes the reference image(s) into first-frame latents of shape `(batch_size, latent_height, latent_width, latent_channels)`, scaled by the VAE `scaling_factor`. PIL images are resized and center-cropped to `height x - width`; tensors and arrays must already have that size. + width`; tensors and arrays must already have that size. The latents are repeated `num_videos_per_prompt` + times along the batch dimension. """ is_pil = isinstance(image, PIL.Image.Image) or ( isinstance(image, list) and isinstance(image[0], PIL.Image.Image) @@ -537,7 +539,8 @@ def encode_image( ) image = image.to(device=device, dtype=self.vae.dtype).unsqueeze(2) latents = retrieve_latents(self.vae.encode(image), generator=generator) * self.vae.config.scaling_factor - return latents[:, :, 0].permute(0, 2, 3, 1).to(dtype) + latents = latents[:, :, 0].permute(0, 2, 3, 1).to(dtype) + return latents.repeat_interleave(num_videos_per_prompt, dim=0) def prepare_latents( self, @@ -776,8 +779,9 @@ def __call__( # 4. Encode the reference image first_frame_latents = None if image is not None: - first_frame_latents = self.encode_image(image, height, width, device, dtype, generator) - first_frame_latents = first_frame_latents.repeat_interleave(num_videos_per_prompt, dim=0) + first_frame_latents = self.encode_image( + image, height, width, device, dtype, num_videos_per_prompt, generator + ) # 5. Prepare timesteps. Audio uses a second scheduler instance so that both modalities keep their own step # counter. From cc703db2a16409bc5ea856509a688b86e72d77c4 Mon Sep 17 00:00:00 2001 From: leffff Date: Thu, 1 Oct 2026 23:56:42 +0300 Subject: [PATCH 36/62] fix critical bugs --- .../en/api/models/kandinsky6_transformer.md | 3 ++ docs/source/en/api/pipelines/kandinsky6.md | 4 +++ .../transformers/transformer_kandinsky6.py | 30 +++++-------------- .../kandinsky6/pipeline_kandinsky6_sr.py | 4 +++ .../kandinsky6/pipeline_kandinsky6_ti2va.py | 14 +++++---- tests/pipelines/testing_utils/memory.py | 2 +- 6 files changed, 27 insertions(+), 30 deletions(-) diff --git a/docs/source/en/api/models/kandinsky6_transformer.md b/docs/source/en/api/models/kandinsky6_transformer.md index 1425c01fd07c..984aa7e3f3d2 100644 --- a/docs/source/en/api/models/kandinsky6_transformer.md +++ b/docs/source/en/api/models/kandinsky6_transformer.md @@ -43,6 +43,9 @@ from diffusers import Kandinsky6SRTransformer3DModel transformer = Kandinsky6SRTransformer3DModel.from_pretrained( "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16 ) +# The default config (`nabla_threshold=0.8`) runs NABLA sparse attention, like Kandinsky 5's `nabla` +# checkpoints — only the `flex` backend can execute it. +transformer.set_attention_backend("flex") ``` [[autodoc]] Kandinsky6SRTransformer3DModel diff --git a/docs/source/en/api/pipelines/kandinsky6.md b/docs/source/en/api/pipelines/kandinsky6.md index 1c16067293b5..6e3d0353d047 100644 --- a/docs/source/en/api/pipelines/kandinsky6.md +++ b/docs/source/en/api/pipelines/kandinsky6.md @@ -76,6 +76,10 @@ windows. sr_pipe = Kandinsky6SRPipeline.from_pretrained( "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", torch_dtype=torch.bfloat16 ) +# The SR transformer's default config runs NABLA sparse attention (like Kandinsky 5's `nabla` +# checkpoints), which only the `flex` backend can execute. +sr_pipe.transformer.set_attention_backend("flex") +sr_pipe.transformer.compile_repeated_blocks(fullgraph=True) sr_pipe.enable_model_cpu_offload() upscaled = sr_pipe(video=output.frames[0], resolution_scale=2.25, num_inference_steps=2).frames[0] diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index b1060f1b16fe..b243dafab98e 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -350,7 +350,6 @@ def __init__( num_channels: int, head_dim: int, kv_dim: int | None = None, - text_token_padding: bool = False, processor: Kandinsky6AttnProcessor | None = None, ): super().__init__() @@ -364,10 +363,7 @@ def __init__( self.query_norm = nn.RMSNorm(head_dim) self.key_norm = nn.RMSNorm(head_dim) self.out_layer = nn.Linear(num_channels, num_channels) - self.text_token_padding = text_token_padding self.set_processor(processor or self._default_processor_cls()) - if self.text_token_padding: - self.processor._attention_backend = AttentionBackendName.NATIVE def forward( self, @@ -454,12 +450,11 @@ def __init__( time_dim: int, ff_dim: int, head_dim: int, - text_token_padding: bool = False, ): super().__init__() self.text_modulation = Kandinsky6Modulation(time_dim, model_dim, 6) self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) - self.self_attention = Kandinsky6Attention(model_dim, head_dim, text_token_padding=text_token_padding) + self.self_attention = Kandinsky6Attention(model_dim, head_dim) self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.feed_forward = Kandinsky6FeedForward(model_dim, ff_dim) @@ -495,7 +490,6 @@ def __init__( time_dim: int, ff_dim: int, head_dim: int, - text_token_padding: bool = False, ): super().__init__() self.visual_modulation = Kandinsky6Modulation(time_dim, model_dim, 9) @@ -506,7 +500,6 @@ def __init__( model_dim, head_dim, kv_dim=model_dim, - text_token_padding=text_token_padding, ) self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.feed_forward = Kandinsky6FeedForward(model_dim, ff_dim) @@ -525,18 +518,13 @@ def __init__( time_dim_a: int, ff_dim_a: int, head_dim_a: int, - text_token_padding: bool = False, ca_rope: bool = False, cross_gates: bool = False, fix_modulation: bool = False, ): super().__init__() - self.video_dec_block = Kandinsky6TransformerDecoderBlock( - model_dim, time_dim, ff_dim, head_dim, text_token_padding - ) - self.audio_dec_block = Kandinsky6TransformerDecoderBlock( - model_dim_a, time_dim_a, ff_dim_a, head_dim_a, text_token_padding - ) + self.video_dec_block = Kandinsky6TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim) + self.audio_dec_block = Kandinsky6TransformerDecoderBlock(model_dim_a, time_dim_a, ff_dim_a, head_dim_a) self.va_cross_attention = Kandinsky6Attention( model_dim, head_dim, @@ -712,7 +700,9 @@ class Kandinsky6Transformer3DModel( audio_freqs_scaling (`float`, *optional*, defaults to 1.0): Audio RoPE frequency scaling. scale_factor (`tuple[float, float, float]`, *optional*, defaults to `(1.0, 2.0, 2.0)`): Per-axis `(t, h, w)` RoPE frequency scaling applied to the video positions. - text_token_padding (`bool`, *optional*, defaults to False): Whether text sequences are padded. + text_token_padding (`bool`, *optional*, defaults to False): Checkpoint metadata recording whether the + reference model's text sequences are padded. `forward` always accepts an optional `encoder_attention_mask` + regardless of this flag; whether one is actually passed is entirely up to the caller. ca_rope (`bool`, *optional*, defaults to False): Whether to use cross-modal audio RoPE. cross_gates (`bool`, *optional*, defaults to False): Whether to use cross-modal residual gates. fix_modulation (`bool`, *optional*, defaults to False): Whether to use the fixed modulation variant. @@ -799,12 +789,7 @@ def __init__( setattr( self, f"{prefix}_text_transformer_blocks", - nn.ModuleList( - [ - Kandinsky6TransformerEncoderBlock(md, td, fd, hd, text_token_padding) - for _ in range(num_text_blocks) - ] - ), + nn.ModuleList([Kandinsky6TransformerEncoderBlock(md, td, fd, hd) for _ in range(num_text_blocks)]), ) self.visual_transformer_blocks = nn.ModuleList( [ @@ -817,7 +802,6 @@ def __init__( time_dim_a, ff_dim_a, head_dim_a, - text_token_padding, ca_rope, cross_gates, fix_modulation, diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py index ae3529aaffe9..b83eb8c36582 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py @@ -53,6 +53,10 @@ >>> sr_pipe = Kandinsky6SRPipeline.from_pretrained( ... "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", torch_dtype=torch.bfloat16 ... ) + >>> # The transformer's default config runs NABLA sparse attention, which only the `flex` backend can + >>> # execute; compiling avoids the eager fallback's much higher memory use at video resolutions. + >>> sr_pipe.transformer.set_attention_backend("flex") + >>> sr_pipe.transformer.compile_repeated_blocks(fullgraph=True) >>> sr_pipe.enable_model_cpu_offload() >>> output = sr_pipe(video=video, resolution_scale=2.25, num_inference_steps=2) >>> export_to_video(output.frames[0], "output_sr.mp4", fps=24) diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index 9c344b47cf1f..f8973983335d 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -394,7 +394,7 @@ def expand_prompts( prompt: str | list[str], image: PIL.Image.Image | list[PIL.Image.Image] | None = None, max_sequence_length: int = 1024, - generator: torch.Generator | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, ) -> str | list[str]: r""" Rewrites short prompts into detailed video+audio prompts with the Qwen2.5-VL text encoder, grounding them on @@ -407,20 +407,22 @@ def expand_prompts( Reference image(s) of an image-to-video call. max_sequence_length (`int`, defaults to `1024`): Maximum number of generated tokens per prompt. - generator (`torch.Generator`, *optional*): - Seeds the sampled expansion. `generate` draws from the global RNG, so the global RNG is seeded from - this generator's seed; later `randn_tensor` calls keep using `generator` directly. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Seeds the sampled expansion; a list must match `prompt`'s length, one generator per item. `generate` + draws from the global RNG, so the global RNG is seeded from this generator's seed; later + `randn_tensor` calls keep using `generator` directly. Returns: `str` or `list[str]`: The expanded prompt(s). """ if isinstance(prompt, list): images = image if isinstance(image, list) else [image] * len(prompt) + generators = generator if isinstance(generator, list) else [generator] * len(prompt) return [ self.expand_prompts( - item, image=item_image, max_sequence_length=max_sequence_length, generator=generator + item, image=item_image, max_sequence_length=max_sequence_length, generator=item_generator ) - for item, item_image in zip(prompt, images, strict=True) + for item, item_image, item_generator in zip(prompt, images, generators, strict=True) ] if image is not None and not isinstance(image, PIL.Image.Image): raise ValueError("`expand_prompts` expects `image` as a `PIL.Image.Image`") diff --git a/tests/pipelines/testing_utils/memory.py b/tests/pipelines/testing_utils/memory.py index 4639016f640a..ba18f86fd0a8 100644 --- a/tests/pipelines/testing_utils/memory.py +++ b/tests/pipelines/testing_utils/memory.py @@ -312,7 +312,7 @@ def enable_group_offload_on_component(pipe, group_offloading_kwargs): for module in component.modules() if hasattr(module, "_diffusers_hook") ) - for component_name in ["vae", "vqvae", "image_encoder", "audio_vae", "latent_upscaler"]: + for component_name in ["vae", "vqvae", "image_encoder", "audio_vae", "vocoder", "latent_upscaler"]: component = getattr(pipe, component_name, None) if isinstance(component, torch.nn.Module): component.to(torch_device) From 9ae71f0d04eddfa3f3aa5e0a189b3d68345ec968 Mon Sep 17 00:00:00 2001 From: leffff Date: Fri, 2 Oct 2026 00:38:07 +0300 Subject: [PATCH 37/62] Revert checkpoint-breaking module renames from 28c5326da Commit 28c5326da (today, 14:03 +0300) renamed Kandinsky6FusedTransformerDecoderBlock's videoT/audioT submodules to video_dec_block/audio_dec_block, and replaced both Kandinsky6TimeEmbeddings and Kandinsky6SRTimeEmbeddings's hand-rolled in_layer/out_layer with diffusers' Timesteps/TimestepEmbedding, purely for naming clarity. No checkpoint conversion was re-run to match, so every published K6 checkpoint (unchanged since Sep 27, confirmed via identical blob SHA256 across recent Hub commits) still ships the old names. Verified against kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers and kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers: this was silently dropping ~65% of the main transformer's weights (every visual_transformer_blocks.*.{video,audio}T.* parameter) and the SR transformer's time embeddings, then crashing with "Cannot copy out of meta tensor" on enable_model_cpu_offload(). After this revert, all checkpoint keys match exactly (4471/4471 for the main transformer, 459/459 for the SR transformer). Co-Authored-By: Claude Sonnet 5 --- .../transformers/transformer_kandinsky6.py | 51 ++++++++++--------- .../transformers/transformer_kandinsky6_sr.py | 19 ++++--- 2 files changed, 40 insertions(+), 30 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index b243dafab98e..ea6d40340b5a 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -28,7 +28,7 @@ from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn from ..cache_utils import CacheMixin -from ..embeddings import TimestepEmbedding, Timesteps +from ..embeddings import get_timestep_embedding from ..modeling_utils import ModelMixin, get_parameter_dtype from .transformer_ltx2 import AudioVisualModelOutput @@ -266,19 +266,24 @@ def __call__( class Kandinsky6TimeEmbeddings(nn.Module): """Sinusoidal timestep embedding with a K6-compatible parameter layout.""" - def __init__(self, model_dim: int, time_dim: int): + def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): super().__init__() if model_dim % 2: raise ValueError("model_dim must be even") - self.time_proj = Timesteps(num_channels=model_dim, flip_sin_to_cos=True, downscale_freq_shift=0) - self.timestep_embedder = TimestepEmbedding(in_channels=model_dim, time_embed_dim=time_dim, act_fn="silu") + self.model_dim = model_dim + self.max_period = max_period + self.in_layer = nn.Linear(model_dim, time_dim) + self.activation = nn.SiLU() + self.out_layer = nn.Linear(time_dim, time_dim) def forward(self, timestep: Tensor) -> Tensor: # The sinusoidal embedding is float32; `_keep_in_fp32_modules` keeps these layers float32 under # `from_pretrained(torch_dtype=...)`, and the cast aligns the input with whatever dtype they hold. - embed = self.time_proj(timestep) - embed = embed.to(get_parameter_dtype(self.timestep_embedder)) - return self.timestep_embedder(embed) + embed = get_timestep_embedding( + timestep, self.model_dim, flip_sin_to_cos=True, downscale_freq_shift=0, max_period=self.max_period + ) + embed = embed.to(get_parameter_dtype(self.in_layer)) + return self.out_layer(self.activation(self.in_layer(embed))) class Kandinsky6TextEmbeddings(nn.Module): @@ -523,8 +528,8 @@ def __init__( fix_modulation: bool = False, ): super().__init__() - self.video_dec_block = Kandinsky6TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim) - self.audio_dec_block = Kandinsky6TransformerDecoderBlock(model_dim_a, time_dim_a, ff_dim_a, head_dim_a) + self.videoT = Kandinsky6TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim) + self.audioT = Kandinsky6TransformerDecoderBlock(model_dim_a, time_dim_a, ff_dim_a, head_dim_a) self.va_cross_attention = Kandinsky6Attention( model_dim, head_dim, @@ -567,39 +572,39 @@ def forward( ) -> tuple[Tensor | None, Tensor | None]: t_v, t_a = time_embed if vis is not None: - sa_p, ca_p, ff_p = torch.chunk(self.video_dec_block.visual_modulation(t_v), 3, dim=-1) + sa_p, ca_p, ff_p = torch.chunk(self.videoT.visual_modulation(t_v), 3, dim=-1) shift, scale, gate = torch.chunk(sa_p, 3, dim=-1) vis = apply_gate_sum( vis, - self.video_dec_block.self_attention( - apply_scale_shift(self.video_dec_block.self_attention_norm(vis.float()), vis, scale, shift), + self.videoT.self_attention( + apply_scale_shift(self.videoT.self_attention_norm(vis.float()), vis, scale, shift), rotary_emb=vis_rope, sparse_params=sparse_params, ), gate, ).type_as(vis) shift, scale, gate_v = torch.chunk(ca_p, 3, dim=-1) - vis_pre_ca = apply_scale_shift(self.video_dec_block.cross_attention_norm(vis.float()), vis, scale, shift) - vis_out_t = self.video_dec_block.cross_attention( + vis_pre_ca = apply_scale_shift(self.videoT.cross_attention_norm(vis.float()), vis, scale, shift) + vis_out_t = self.videoT.cross_attention( vis_pre_ca, encoder_hidden_states=text_v, attn_mask=attn_mask, ) if aud is not None: - sa_p, ca_p, ff_p_a = torch.chunk(self.audio_dec_block.visual_modulation(t_a), 3, dim=-1) + sa_p, ca_p, ff_p_a = torch.chunk(self.audioT.visual_modulation(t_a), 3, dim=-1) shift, scale, gate = torch.chunk(sa_p, 3, dim=-1) aud = apply_gate_sum( aud, - self.audio_dec_block.self_attention( - apply_scale_shift(self.audio_dec_block.self_attention_norm(aud.float()), aud, scale, shift), + self.audioT.self_attention( + apply_scale_shift(self.audioT.self_attention_norm(aud.float()), aud, scale, shift), rotary_emb=aud_rope, ), gate, ).type_as(aud) shift, scale, gate_a = torch.chunk(ca_p, 3, dim=-1) - aud_pre_ca = apply_scale_shift(self.audio_dec_block.cross_attention_norm(aud.float()), aud, scale, shift) - aud_out_t = self.audio_dec_block.cross_attention( + aud_pre_ca = apply_scale_shift(self.audioT.cross_attention_norm(aud.float()), aud, scale, shift) + aud_out_t = self.audioT.cross_attention( aud_pre_ca, encoder_hidden_states=text_a, attn_mask=attn_mask, @@ -647,8 +652,8 @@ def forward( shift, scale, gate = torch.chunk(ff_p, 3, dim=-1) vis = apply_gate_sum( vis, - self.video_dec_block.feed_forward( - apply_scale_shift(self.video_dec_block.feed_forward_norm(vis.float()), vis, scale, shift) + self.videoT.feed_forward( + apply_scale_shift(self.videoT.feed_forward_norm(vis.float()), vis, scale, shift) ), gate, ).type_as(vis) @@ -656,8 +661,8 @@ def forward( shift, scale, gate = torch.chunk(ff_p_a, 3, dim=-1) aud = apply_gate_sum( aud, - self.audio_dec_block.feed_forward( - apply_scale_shift(self.audio_dec_block.feed_forward_norm(aud.float()), aud, scale, shift) + self.audioT.feed_forward( + apply_scale_shift(self.audioT.feed_forward_norm(aud.float()), aud, scale, shift) ), gate, ).type_as(aud) diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index e677327010ce..d31307c997de 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -27,7 +27,7 @@ from ...loaders import PeftAdapterMixin from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn -from ..embeddings import TimestepEmbedding, Timesteps +from ..embeddings import get_timestep_embedding from ..modeling_outputs import Transformer2DModelOutput from ..modeling_utils import ModelMixin, get_parameter_dtype @@ -234,19 +234,24 @@ def forward( class Kandinsky6SRTimeEmbeddings(nn.Module): """Sinusoidal timestep embedding with a K6-compatible parameter layout.""" - def __init__(self, model_dim: int, time_dim: int): + def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): super().__init__() if model_dim % 2: raise ValueError("model_dim must be even") - self.time_proj = Timesteps(num_channels=model_dim, flip_sin_to_cos=True, downscale_freq_shift=0) - self.timestep_embedder = TimestepEmbedding(in_channels=model_dim, time_embed_dim=time_dim, act_fn="silu") + self.model_dim = model_dim + self.max_period = max_period + self.in_layer = nn.Linear(model_dim, time_dim) + self.activation = nn.SiLU() + self.out_layer = nn.Linear(time_dim, time_dim) def forward(self, timestep: Tensor) -> Tensor: # The sinusoidal embedding is float32; `_keep_in_fp32_modules` keeps these layers float32 under # `from_pretrained(torch_dtype=...)`, and the cast aligns the input with whatever dtype they hold. - embed = self.time_proj(timestep) - embed = embed.to(get_parameter_dtype(self.timestep_embedder)) - return self.timestep_embedder(embed) + embed = get_timestep_embedding( + timestep, self.model_dim, flip_sin_to_cos=True, downscale_freq_shift=0, max_period=self.max_period + ) + embed = embed.to(get_parameter_dtype(self.in_layer)) + return self.out_layer(self.activation(self.in_layer(embed))) # Copied from diffusers.models.transformers.transformer_kandinsky6.Kandinsky6VisualEmbeddings with Kandinsky6->Kandinsky6SR From a7e596cb9d15721cf4b9c0ec110feba6b4f7dedc Mon Sep 17 00:00:00 2001 From: leffff Date: Fri, 2 Oct 2026 00:45:37 +0300 Subject: [PATCH 38/62] Revert checkpoint-breaking latent-upscaler renames from 372f29088 Same class of regression as 9ae71f0d0, this time in Kandinsky6SRLatentUpscalerX2Branch and Kandinsky6SRLatentUpscalerOutputHead: today's "address comments" commit (18:48 +0300) dropped the private_ prefix from private_mid_blocks/private_upsample/private_blocks/ private_output_proj, and converted OutputHead from nn.Sequential to a plain nn.Module with named norm/activation/conv submodules, again without a matching checkpoint re-conversion. Verified against kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers's latent_upscaler: this left every Kandinsky6SRLatentUpscalerX2Branch parameter (and the OutputHead-based output_proj used standalone by the x2/x4 branches) unfilled on the meta device, crashing enable_model_cpu_offload() the same way as the main transformer did before 9ae71f0d0. Co-Authored-By: Claude Sonnet 5 --- .../kandinsky6/modeling_latent_upscaler.py | 30 ++++++++++--------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py b/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py index 64186ed13a5b..ec10ec4231f4 100644 --- a/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py +++ b/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py @@ -104,17 +104,19 @@ def forward(self, hidden_states: Tensor) -> Tensor: return self.linear(hidden_states + self.spatial_conv(hidden_states)) -class Kandinsky6SRLatentUpscalerOutputHead(nn.Module): +class Kandinsky6SRLatentUpscalerOutputHead(nn.Sequential): """`modulated norm -> SiLU -> conv3x3x3` projection back to the latent channels, threading `zq` into the norm.""" def __init__(self, in_channels: int, out_channels: int, zq_channels: int) -> None: - super().__init__() - self.norm = Kandinsky6SRLatentUpscalerModulatedNorm(in_channels, zq_channels) - self.activation = nn.SiLU() - self.conv = Kandinsky6SRLatentUpscalerConv3d(in_channels, out_channels, kernel_size=3) + super().__init__( + Kandinsky6SRLatentUpscalerModulatedNorm(in_channels, zq_channels), + nn.SiLU(), + Kandinsky6SRLatentUpscalerConv3d(in_channels, out_channels, kernel_size=3), + ) def forward(self, hidden_states: Tensor, zq: Tensor) -> Tensor: - return self.conv(self.activation(self.norm(hidden_states, zq))) + norm, activation, conv = self + return conv(activation(norm(hidden_states, zq))) def _residual_stack(num_blocks: int, in_channels: int, out_channels: int, zq_channels: int) -> nn.ModuleList: @@ -145,21 +147,21 @@ def __init__( self.finisher = nn.Module() self.finisher.spatial_conv = nn.Conv3d(width_1, width_1, kernel_size=(1, 3, 3), padding=(0, 1, 1)) self.finisher.linear = nn.Conv3d(width_1, width_1, kernel_size=1) - self.mid_blocks = _residual_stack(num_mid_blocks, width_1, width_2, in_channels) - self.upsample = Kandinsky6SRLatentUpscalerUpsample(width_2) - self.blocks = _residual_stack(num_post_blocks, width_2, width_3, in_channels) - self.output_proj = Kandinsky6SRLatentUpscalerOutputHead(width_3, in_channels, in_channels) + self.private_mid_blocks = _residual_stack(num_mid_blocks, width_1, width_2, in_channels) + self.private_upsample = Kandinsky6SRLatentUpscalerUpsample(width_2) + self.private_blocks = _residual_stack(num_post_blocks, width_2, width_3, in_channels) + self.private_output_proj = Kandinsky6SRLatentUpscalerOutputHead(width_3, in_channels, in_channels) def forward(self, hidden_states: Tensor, zq: Tensor) -> Tensor: for block in self.adapter: hidden_states = block(hidden_states, zq) hidden_states = self.finisher.linear(hidden_states + self.finisher.spatial_conv(hidden_states)) - for block in self.mid_blocks: + for block in self.private_mid_blocks: hidden_states = block(hidden_states, zq) - hidden_states = self.upsample(hidden_states) - for block in self.blocks: + hidden_states = self.private_upsample(hidden_states) + for block in self.private_blocks: hidden_states = block(hidden_states, zq) - return self.output_proj(hidden_states, zq) + return self.private_output_proj(hidden_states, zq) class Kandinsky6SRLatentUpscaler(nn.Module): From 8fadd97f2405a11f68f2931e056a097a6e236067 Mon Sep 17 00:00:00 2001 From: leffff Date: Fri, 2 Oct 2026 20:14:28 +0300 Subject: [PATCH 39/62] move vocoder and LU --- docs/source/en/api/models/kandinsky6_vae.md | 6 +- docs/source/en/api/pipelines/kandinsky6.md | 2 +- src/diffusers/__init__.py | 8 +- src/diffusers/models/__init__.py | 4 + src/diffusers/models/autoencoders/__init__.py | 1 + .../autoencoders/autoencoder_kandinsky6_sr.py | 134 ++++++------------ .../autoencoders/autoencoder_mmaudio.py | 16 +-- .../autoencoders/mmaudio_vocoder.py} | 4 +- .../models/latent_upscaler/__init__.py | 5 + .../latent_upscaler_kandinsky6_sr.py} | 4 +- .../transformers/transformer_kandinsky6.py | 8 +- .../pipelines/kandinsky6/__init__.py | 4 - .../kandinsky6/pipeline_kandinsky6_sr.py | 3 +- .../kandinsky6/pipeline_kandinsky6_ti2va.py | 11 +- src/diffusers/utils/dummy_pt_objects.py | 30 ++++ .../dummy_torch_and_transformers_objects.py | 30 ---- .../test_models_autoencoder_mmaudio.py | 2 +- .../autoencoders}/test_models_vocoder.py | 6 +- tests/models/latent_upscaler/__init__.py | 0 .../test_models_latent_upscaler.py | 6 +- .../kandinsky6/test_pipeline_kandinsky6_sr.py | 2 +- 21 files changed, 119 insertions(+), 167 deletions(-) rename src/diffusers/{pipelines/kandinsky6/modeling_vocoder.py => models/autoencoders/mmaudio_vocoder.py} (99%) create mode 100644 src/diffusers/models/latent_upscaler/__init__.py rename src/diffusers/{pipelines/kandinsky6/modeling_latent_upscaler.py => models/latent_upscaler/latent_upscaler_kandinsky6_sr.py} (99%) rename tests/{pipelines/kandinsky6 => models/autoencoders}/test_models_vocoder.py (96%) create mode 100644 tests/models/latent_upscaler/__init__.py rename tests/{pipelines/kandinsky6 => models/latent_upscaler}/test_models_latent_upscaler.py (96%) diff --git a/docs/source/en/api/models/kandinsky6_vae.md b/docs/source/en/api/models/kandinsky6_vae.md index e72646c8bf2b..b4e42a23f65b 100644 --- a/docs/source/en/api/models/kandinsky6_vae.md +++ b/docs/source/en/api/models/kandinsky6_vae.md @@ -13,7 +13,7 @@ specific language governing permissions and limitations under the License. # Kandinsky 6 VAEs Kandinsky 6 uses a causal 3D K-VAE for video super-resolution and the MMAudio mel-spectrogram VAE, paired with a -separate BigVGAN [`~pipelines.kandinsky6.MMAudioVocoder`], for synchronized audio generation. +separate BigVGAN [`MMAudioVocoder`], for synchronized audio generation. ## Kandinsky6SRVAE @@ -37,7 +37,7 @@ vae = Kandinsky6SRVAE.from_pretrained( ## MMAudioVAE The mel-spectrogram VAE used by [`Kandinsky6TI2VAPipeline`] when `sample_audio=True`. Its `decode` output is a mel -spectrogram; pass it through [`~pipelines.kandinsky6.MMAudioVocoder`] to get a waveform. +spectrogram; pass it through [`MMAudioVocoder`] to get a waveform. The reference implementation can be found at [hkchengrex/MMAudio](https://github.com/hkchengrex/MMAudio) (MIT license). @@ -62,5 +62,5 @@ Adapted from the BigVGAN-v2 vocoder MMAudio bundles, itself from [NVIDIA/BigVGAN](https://github.com/NVIDIA/BigVGAN) (MIT license), with the anti-aliased Snake activations of [alias-free-torch](https://github.com/junjun3518/alias-free-torch) (Apache License 2.0). -[[autodoc]] pipelines.kandinsky6.MMAudioVocoder +[[autodoc]] MMAudioVocoder - forward diff --git a/docs/source/en/api/pipelines/kandinsky6.md b/docs/source/en/api/pipelines/kandinsky6.md index 6e3d0353d047..34583cc07b9a 100644 --- a/docs/source/en/api/pipelines/kandinsky6.md +++ b/docs/source/en/api/pipelines/kandinsky6.md @@ -135,7 +135,7 @@ pipe.vae.enable_tiling() ## Kandinsky6SRLatentUpscalerBank -[[autodoc]] pipelines.kandinsky6.modeling_latent_upscaler.Kandinsky6SRLatentUpscalerBank +[[autodoc]] Kandinsky6SRLatentUpscalerBank - forward ## Kandinsky6TI2VAPipelineOutput diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index b5be02858f75..58dcf342e314 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -303,6 +303,7 @@ "JoyImageEditTransformer3DModel", "Kandinsky3UNet", "Kandinsky5Transformer3DModel", + "Kandinsky6SRLatentUpscalerBank", "Kandinsky6SRTransformer3DModel", "Kandinsky6SRVAE", "Kandinsky6Transformer3DModel", @@ -322,6 +323,7 @@ "MiniMaxMusic3Transformer1DModel", "MiniMaxMusic3Vocoder", "MMAudioVAE", + "MMAudioVocoder", "MochiTransformer3DModel", "ModelMixin", "MotifVideoTransformer3DModel", @@ -699,7 +701,6 @@ "Kandinsky5I2VPipeline", "Kandinsky5T2IPipeline", "Kandinsky5T2VPipeline", - "Kandinsky6SRLatentUpscalerBank", "Kandinsky6SRPipeline", "Kandinsky6SRPipelineOutput", "Kandinsky6TI2VAPipeline", @@ -753,7 +754,6 @@ "MarigoldDepthPipeline", "MarigoldIntrinsicsPipeline", "MarigoldNormalsPipeline", - "MMAudioVocoder", "MochiPipeline", "MotifVideoImage2VideoPipeline", "MotifVideoPipeline", @@ -1185,6 +1185,7 @@ JoyImageEditTransformer3DModel, Kandinsky3UNet, Kandinsky5Transformer3DModel, + Kandinsky6SRLatentUpscalerBank, Kandinsky6SRTransformer3DModel, Kandinsky6SRVAE, Kandinsky6Transformer3DModel, @@ -1204,6 +1205,7 @@ MiniMaxMusic3Transformer1DModel, MiniMaxMusic3Vocoder, MMAudioVAE, + MMAudioVocoder, MochiTransformer3DModel, ModelMixin, MotifVideoTransformer3DModel, @@ -1556,7 +1558,6 @@ Kandinsky5I2VPipeline, Kandinsky5T2IPipeline, Kandinsky5T2VPipeline, - Kandinsky6SRLatentUpscalerBank, Kandinsky6SRPipeline, Kandinsky6SRPipelineOutput, Kandinsky6TI2VAPipeline, @@ -1610,7 +1611,6 @@ MarigoldDepthPipeline, MarigoldIntrinsicsPipeline, MarigoldNormalsPipeline, - MMAudioVocoder, MochiPipeline, MotifVideoImage2VideoPipeline, MotifVideoPipeline, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 66fbec1bdb4b..a395b493cdf2 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -63,6 +63,7 @@ _import_structure["autoencoders.consistency_decoder_vae"] = ["ConsistencyDecoderVAE"] _import_structure["autoencoders.ltx2_diffusion_decoder"] = ["LTX2VideoDiffusionDecoderModel"] _import_structure["autoencoders.minimax_music3_vocoder"] = ["MiniMaxMusic3Vocoder"] + _import_structure["autoencoders.mmaudio_vocoder"] = ["MMAudioVocoder"] _import_structure["autoencoders.vq_model"] = ["VQModel"] _import_structure["cache_utils"] = ["CacheMixin"] _import_structure["condition_embedders.condition_embedder_anima"] = ["AnimaTextConditioner"] @@ -87,6 +88,7 @@ _import_structure["controlnets.multicontrolnet"] = ["MultiControlNetModel"] _import_structure["controlnets.multicontrolnet_union"] = ["MultiControlNetUnionModel"] _import_structure["embeddings"] = ["ImageProjection"] + _import_structure["latent_upscaler.latent_upscaler_kandinsky6_sr"] = ["Kandinsky6SRLatentUpscalerBank"] _import_structure["modeling_utils"] = ["ModelMixin"] _import_structure["transformers.ace_step_transformer"] = ["AceStepTransformer1DModel"] _import_structure["transformers.auraflow_transformer_2d"] = ["AuraFlowTransformer2DModel"] @@ -213,6 +215,7 @@ LTX2VideoDiffusionDecoderModel, MiniMaxMusic3Vocoder, MMAudioVAE, + MMAudioVocoder, VQModel, ) from .cache_utils import CacheMixin @@ -238,6 +241,7 @@ ZImageControlNetModel, ) from .embeddings import ImageProjection + from .latent_upscaler import Kandinsky6SRLatentUpscalerBank from .modeling_utils import ModelMixin from .transformers import ( AceStepTransformer1DModel, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index 22c2416d71b2..4997e1459d1a 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/src/diffusers/models/autoencoders/__init__.py @@ -33,4 +33,5 @@ from .consistency_decoder_vae import ConsistencyDecoderVAE from .ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel from .minimax_music3_vocoder import MiniMaxMusic3Vocoder +from .mmaudio_vocoder import MMAudioVocoder from .vq_model import VQModel diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index c17da9b26fc9..a8dcf117173b 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -79,6 +79,9 @@ class Kandinsky6SRCausalConv3d(nn.Module): single non-causal pass. Without it, each segment would need the raw frames the previous segment already consumed in order to rebuild correct padding — a lookback window that grows with network depth, rather than the bounded, constant-size cache this class carries instead — which defeats the point of processing the video in segments. + + `cache` is mutated in place: the caller passes the same dict to every segment, and this class writes the padding it + leaves behind directly into `cache["padding"]` before returning. """ def __init__( @@ -108,7 +111,7 @@ def __init__( def make_cache() -> dict: return {"padding": None} - def forward(self, hidden_states: torch.Tensor, cache: dict) -> tuple[torch.Tensor, dict]: + def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: batch_size, _, num_frames, height, width = hidden_states.shape hidden_states = F.pad(hidden_states, (self.width_pad, self.width_pad, self.height_pad, self.height_pad)) @@ -135,18 +138,15 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> tuple[torch.Tenso if offset_out < output_frames: output[:, :, offset_out:] = self.conv(hidden_states[:, :, offset_in:]) - # The frames the next segment's first window still needs. Returned rather than written into `cache` in - # place, so this stays correct under `torch.utils.checkpoint`'s backward-time recompute (see - # `Kandinsky6SRResnetBlock3D`): the recompute reads the same, never-mutated `cache` argument as the - # original call. + # The frames the next segment's first window still needs. pad_offset = ( offset_in + stride * math.trunc((num_frames - offset_in - self.time_kernel_size) / stride) + stride ) if pad_offset < 0: - new_padding = torch.cat([padding[:, :, pad_offset:], hidden_states], dim=2) + cache["padding"] = torch.cat([padding[:, :, pad_offset:], hidden_states], dim=2) else: - new_padding = hidden_states[:, :, pad_offset:].clone() - return output, {"padding": new_padding} + cache["padding"] = hidden_states[:, :, pad_offset:].clone() + return output class Kandinsky6SRRMSNorm(nn.Module): @@ -172,11 +172,10 @@ class Kandinsky6SRSpatialNorm3D(nn.Module): """Spatial normalization conditioned on the latent `zq` (https://huggingface.co/papers/2209.09002), used here so the decoder can inject the latent back in at every resnet block. Structurally this plays the same role as `SpatialNorm` (`attention_processor.py`) / `CogVideoXSpatialNorm3D` (`autoencoder_kl_cogvideox.py`) — normalize the - features, then scale and shift by convolutions of the upsampled `zq`, threading a `cache` argument through - `forward` functionally (read, never written in place) the same way `CogVideoXSpatialNorm3D` threads its - `conv_cache` — but it is not a `# Copied from` of either: it normalizes with `Kandinsky6SRRMSNorm` instead of - `GroupNorm` (matching the plain-RMSNorm blocks elsewhere in this VAE), and interpolates `zq` in channel chunks to - bound peak memory. + features, then scale and shift by convolutions of the upsampled `zq`, carrying a `cache` dict across segments the + same way `CogVideoXSpatialNorm3D` threads its `conv_cache` — but it is not a `# Copied from` of either: it + normalizes with `Kandinsky6SRRMSNorm` instead of `GroupNorm` (matching the plain-RMSNorm blocks elsewhere in this + VAE), and interpolates `zq` in channel chunks to bound peak memory. `zq` is nearest-upsampled to the feature grid. In the first segment the first frame is upsampled separately, because the temporal upsampler turns `T + 1` latent frames into `2T + 1` pixel frames. @@ -192,7 +191,7 @@ def __init__(self, num_channels: int, zq_channels: int) -> None: def make_cache() -> dict: return {"is_first_segment": True} - def forward(self, hidden_states: torch.Tensor, zq: torch.Tensor, cache: dict) -> tuple[torch.Tensor, dict]: + def forward(self, hidden_states: torch.Tensor, zq: torch.Tensor, cache: dict) -> torch.Tensor: if cache["is_first_segment"]: zq_first = F.interpolate(zq[:, :, :1], size=hidden_states[:, :, :1].shape[-3:], mode="nearest") if zq.size(2) > 1: @@ -216,18 +215,15 @@ def forward(self, hidden_states: torch.Tensor, zq: torch.Tensor, cache: dict) -> dim=1, ) output = self.norm_layer(hidden_states) * self.conv_y(zq) + self.conv_b(zq) - return output, {"is_first_segment": False} + cache["is_first_segment"] = False + return output class Kandinsky6SRResnetBlock3D(nn.Module): """Causal residual block; decoder blocks modulate their norms with the latent `zq`. - `forward` never mutates the `cache` it is given: `conv1`/`conv2` (`Kandinsky6SRCausalConv3d`) and, for decoder - blocks, `norm1`/`norm2` (`Kandinsky6SRSpatialNorm3D`) each return a new cache value instead of writing into the one - they were passed, the same functional-cache convention `autoencoder_kl_cogvideox.py` uses. That is what makes it - safe to wrap this block in `torch.utils.checkpoint` (see `Kandinsky6SREncoder3D`/`Kandinsky6SRDecoder3D`): the - backward-time recompute reads the same, never-mutated `cache` argument the original call did, so it retraces the - same branch and produces the same activations and gradients. + `conv1`/`conv2` (`Kandinsky6SRCausalConv3d`) and, for decoder blocks, `norm1`/`norm2` (`Kandinsky6SRSpatialNorm3D`) + each mutate the matching key of the `cache` dict this block is given in place. """ def __init__(self, in_channels: int, out_channels: int, zq_channels: int | None = None) -> None: @@ -252,24 +248,21 @@ def make_cache(self) -> dict: cache["norm2"] = self.norm2.make_cache() return cache - def forward( - self, hidden_states: torch.Tensor, cache: dict, zq: torch.Tensor | None = None - ) -> tuple[torch.Tensor, dict]: + def forward(self, hidden_states: torch.Tensor, cache: dict, zq: torch.Tensor | None = None) -> torch.Tensor: residual = hidden_states - new_cache = {} if zq is None: hidden_states = self.norm1(hidden_states) else: - hidden_states, new_cache["norm1"] = self.norm1(hidden_states, zq, cache["norm1"]) - hidden_states, new_cache["conv1"] = self.conv1(F.silu(hidden_states), cache["conv1"]) + hidden_states = self.norm1(hidden_states, zq, cache["norm1"]) + hidden_states = self.conv1(F.silu(hidden_states), cache["conv1"]) if zq is None: hidden_states = self.norm2(hidden_states) else: - hidden_states, new_cache["norm2"] = self.norm2(hidden_states, zq, cache["norm2"]) - hidden_states, new_cache["conv2"] = self.conv2(F.silu(hidden_states), cache["conv2"]) + hidden_states = self.norm2(hidden_states, zq, cache["norm2"]) + hidden_states = self.conv2(F.silu(hidden_states), cache["conv2"]) if self.in_channels != self.out_channels: residual = self.nin_shortcut(residual) - return residual + hidden_states, new_cache + return residual + hidden_states class Kandinsky6SRDownsample(nn.Module): @@ -325,8 +318,8 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: pooled = F.avg_pool1d(sequence, kernel_size=2, stride=2) pooled = pooled.reshape(batch_size, height, width, channels, -1).permute(0, 3, 4, 1, 2) - conv_out, cache["temporal_conv"][0] = self.temporal_conv[0](hidden_states, cache["temporal_conv"][0]) - conv_out, cache["temporal_conv"][1] = self.temporal_conv[1](conv_out, cache["temporal_conv"][1]) + conv_out = self.temporal_conv[0](hidden_states, cache["temporal_conv"][0]) + conv_out = self.temporal_conv[1](conv_out, cache["temporal_conv"][1]) return self.linear(conv_out + pooled) @@ -354,7 +347,7 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: repeated = hidden_states.repeat_interleave(2, dim=2) if cache["temporal_conv"]["padding"] is None: repeated = repeated[:, :, 1:] - conv_out, cache["temporal_conv"] = self.temporal_conv(repeated, cache["temporal_conv"]) + conv_out = self.temporal_conv(repeated, cache["temporal_conv"]) hidden_states = conv_out + repeated hidden_states = F.interpolate(hidden_states, scale_factor=(1, 2, 2), mode="nearest") @@ -365,12 +358,9 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: class Kandinsky6SREncoder3D(nn.Module): """Causal encoder: `conv_in`, downsampling resnet levels, a resnet bottleneck, then `norm_out`/`conv_out`. - `forward` only returns `hidden_states`, not `cache`: each submodule still returns its own updated cache value - rather than mutating the one it was given (see `Kandinsky6SRCausalConv3d`), but this method writes every one of - those return values back into the matching key of the `cache` dict it received, e.g. - `hidden_states, cache["conv_out"] = self.conv_out(hidden_states, cache["conv_out"])`. Because the caller - (`Kandinsky6SRVAE.encode`) passes the same `cache` dict to every segment, that key assignment is what carries - the padding state from one segment to the next — `cache` does not need to be returned for this to work. + Each submodule mutates the matching key of the `cache` dict it is given in place. Because the caller + (`Kandinsky6SRVAE.encode`) passes the same `cache` dict to every segment, that in-place mutation is what carries + the padding state from one segment to the next. """ def __init__( @@ -409,8 +399,6 @@ def __init__( self.norm_out = Kandinsky6SRRMSNorm(block_in) self.conv_out = Kandinsky6SRCausalConv3d(block_in, 2 * latent_channels, kernel_size=3) - self.gradient_checkpointing = False - def make_cache(self) -> dict: return { "conv_in": self.conv_in.make_cache(), @@ -427,32 +415,18 @@ def make_cache(self) -> dict: } def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: - hidden_states, cache["conv_in"] = self.conv_in(hidden_states, cache["conv_in"]) + hidden_states = self.conv_in(hidden_states, cache["conv_in"]) for down, level_cache in zip(self.down, cache["down"]): for i, block in enumerate(down.block): - block_cache = level_cache["block"][i] - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, level_cache["block"][i] = self._gradient_checkpointing_func( - block, hidden_states, block_cache - ) - else: - hidden_states, level_cache["block"][i] = block(hidden_states, block_cache) + hidden_states = block(hidden_states, level_cache["block"][i]) if hasattr(down, "downsample"): hidden_states = down.downsample(hidden_states, level_cache["downsample"]) - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, cache["mid_1"] = self._gradient_checkpointing_func( - self.mid.block_1, hidden_states, cache["mid_1"] - ) - hidden_states, cache["mid_2"] = self._gradient_checkpointing_func( - self.mid.block_2, hidden_states, cache["mid_2"] - ) - else: - hidden_states, cache["mid_1"] = self.mid.block_1(hidden_states, cache["mid_1"]) - hidden_states, cache["mid_2"] = self.mid.block_2(hidden_states, cache["mid_2"]) + hidden_states = self.mid.block_1(hidden_states, cache["mid_1"]) + hidden_states = self.mid.block_2(hidden_states, cache["mid_2"]) hidden_states = F.silu(self.norm_out(hidden_states)) - hidden_states, cache["conv_out"] = self.conv_out(hidden_states, cache["conv_out"]) + hidden_states = self.conv_out(hidden_states, cache["conv_out"]) return hidden_states @@ -460,10 +434,8 @@ class Kandinsky6SRDecoder3D(nn.Module): """Causal decoder: `conv_in`, a `zq`-conditioned resnet bottleneck, upsampling resnet levels, then `norm_out`/`conv_out`. - Like `Kandinsky6SREncoder3D.forward`, this only returns the decoded tensor, not `cache`: writing each - submodule's returned cache value back into the matching key of the `cache` dict it received (e.g. - `hidden_states, cache["conv_out"] = self.conv_out(hidden_states, cache["conv_out"])`) is enough, since - `Kandinsky6SRVAE.decode` passes that same dict to every segment. + Like `Kandinsky6SREncoder3D`, each submodule mutates the matching key of the `cache` dict it is given in place, + which is what carries the padding state across the segments `Kandinsky6SRVAE.decode` passes it. """ def __init__( @@ -506,8 +478,6 @@ def __init__( self.norm_out = Kandinsky6SRSpatialNorm3D(block_in, latent_channels) self.conv_out = Kandinsky6SRCausalConv3d(block_in, out_channels, kernel_size=3) - self.gradient_checkpointing = False - def make_cache(self) -> dict: return { "conv_in": self.conv_in.make_cache(), @@ -525,35 +495,21 @@ def make_cache(self) -> dict: } def forward(self, latents: torch.Tensor, cache: dict) -> torch.Tensor: - hidden_states, cache["conv_in"] = self.conv_in(latents, cache["conv_in"]) + hidden_states = self.conv_in(latents, cache["conv_in"]) - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, cache["mid_1"] = self._gradient_checkpointing_func( - self.mid.block_1, hidden_states, cache["mid_1"], latents - ) - hidden_states, cache["mid_2"] = self._gradient_checkpointing_func( - self.mid.block_2, hidden_states, cache["mid_2"], latents - ) - else: - hidden_states, cache["mid_1"] = self.mid.block_1(hidden_states, cache["mid_1"], zq=latents) - hidden_states, cache["mid_2"] = self.mid.block_2(hidden_states, cache["mid_2"], zq=latents) + hidden_states = self.mid.block_1(hidden_states, cache["mid_1"], zq=latents) + hidden_states = self.mid.block_2(hidden_states, cache["mid_2"], zq=latents) for level in reversed(range(len(self.up))): up, level_cache = self.up[level], cache["up"][level] for i, block in enumerate(up.block): - block_cache = level_cache["block"][i] - if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, level_cache["block"][i] = self._gradient_checkpointing_func( - block, hidden_states, block_cache, latents - ) - else: - hidden_states, level_cache["block"][i] = block(hidden_states, block_cache, zq=latents) + hidden_states = block(hidden_states, level_cache["block"][i], zq=latents) if hasattr(up, "upsample"): hidden_states = up.upsample(hidden_states, level_cache["upsample"]) - hidden_states, cache["norm_out"] = self.norm_out(hidden_states, latents, cache["norm_out"]) + hidden_states = self.norm_out(hidden_states, latents, cache["norm_out"]) hidden_states = F.silu(hidden_states) - hidden_states, cache["conv_out"] = self.conv_out(hidden_states, cache["conv_out"]) + hidden_states = self.conv_out(hidden_states, cache["conv_out"]) return hidden_states @@ -565,11 +521,6 @@ class Kandinsky6SRVAE(ModelMixin, ConfigMixin): carry their padding state between segments, so the segmentation only bounds peak memory and does not change the result. - Supports gradient checkpointing (see [`~ModelMixin.enable_gradient_checkpointing`]): every - [`Kandinsky6SRResnetBlock3D`] in the encoder and decoder is checkpointed, since its `cache` argument is threaded - functionally (read, never mutated in place — see [`Kandinsky6SRCausalConv3d`]), so recomputing it during backward - reproduces the exact same forward pass. - Args: in_channels (`int`, defaults to `3`): Number of pixel channels. @@ -593,7 +544,6 @@ class Kandinsky6SRVAE(ModelMixin, ConfigMixin): """ _no_split_modules = ["Kandinsky6SREncoder3D", "Kandinsky6SRDecoder3D"] - _supports_gradient_checkpointing = True @register_to_config def __init__( diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index 566479ce8aa3..e9d73417cd3f 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -17,8 +17,8 @@ """MMAudio mel-spectrogram VAE used by the Kandinsky 6 TI2VA pipeline. -The BigVGAN vocoder that turns this VAE's decoded mel spectrograms into waveforms is a separate, pipeline-local -component, [`~pipelines.kandinsky6.MMAudioVocoder`]. +The BigVGAN vocoder that turns this VAE's decoded mel spectrograms into waveforms is a separate component, +[`MMAudioVocoder`]. """ from __future__ import annotations @@ -352,9 +352,8 @@ class MMAudioVAE(ModelMixin, ConfigMixin): Audio VAE of [`Kandinsky6TI2VAPipeline`]: a magnitude-preserving autoencoder over log-mel spectrograms (MMAudio, https://arxiv.org/abs/2412.15322). - `encode` turns a waveform into a latent distribution; `decode` turns latents back into a mel spectrogram, which the - pipeline-local [`~pipelines.kandinsky6.MMAudioVocoder`] then turns into a waveform. One latent frame covers - `hop_length * 2` samples. + `encode` turns a waveform into a latent distribution; `decode` turns latents back into a mel spectrogram, which + [`MMAudioVocoder`] then turns into a waveform. One latent frame covers `hop_length * 2` samples. Args: mel_bins (`int`, defaults to `128`): @@ -372,8 +371,8 @@ class MMAudioVAE(ModelMixin, ConfigMixin): n_fft (`int`, defaults to `2048`): FFT size of the mel front end. hop_length (`int`, defaults to `512`): - Hop length of the mel front end. Must match the total upsampling factor of the - [`~pipelines.kandinsky6.MMAudioVocoder`] this VAE is paired with. + Hop length of the mel front end. Must match the total upsampling factor of the [`MMAudioVocoder`] this VAE + is paired with. scaling_factor (`float`, defaults to `0.417`): Scale applied to the latents before they enter the diffusion transformer. """ @@ -430,8 +429,7 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | t Whether to return a [`~models.autoencoder_kl.DecoderOutput`] instead of a plain tuple. Returns: - The mel spectrogram of shape `(batch_size, mel_bins, num_mel_frames)`, ready for - [`~pipelines.kandinsky6.MMAudioVocoder`]. + The mel spectrogram of shape `(batch_size, mel_bins, num_mel_frames)`, ready for [`MMAudioVocoder`]. """ mel = self.vae.decode(z) if not return_dict: diff --git a/src/diffusers/pipelines/kandinsky6/modeling_vocoder.py b/src/diffusers/models/autoencoders/mmaudio_vocoder.py similarity index 99% rename from src/diffusers/pipelines/kandinsky6/modeling_vocoder.py rename to src/diffusers/models/autoencoders/mmaudio_vocoder.py index 0121c7547a92..d1f12fb5ef62 100644 --- a/src/diffusers/pipelines/kandinsky6/modeling_vocoder.py +++ b/src/diffusers/models/autoencoders/mmaudio_vocoder.py @@ -28,8 +28,8 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config -from ...models.autoencoders.vae import DecoderOutput -from ...models.modeling_utils import ModelMixin +from ..modeling_utils import ModelMixin +from .vae import DecoderOutput def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: diff --git a/src/diffusers/models/latent_upscaler/__init__.py b/src/diffusers/models/latent_upscaler/__init__.py new file mode 100644 index 000000000000..1c3f83192f85 --- /dev/null +++ b/src/diffusers/models/latent_upscaler/__init__.py @@ -0,0 +1,5 @@ +from ...utils import is_torch_available + + +if is_torch_available(): + from .latent_upscaler_kandinsky6_sr import Kandinsky6SRLatentUpscalerBank diff --git a/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py b/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py similarity index 99% rename from src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py rename to src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py index ec10ec4231f4..b3e10416a7d8 100644 --- a/src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py +++ b/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py @@ -21,8 +21,8 @@ from torch.nn import functional from ...configuration_utils import ConfigMixin, register_to_config -from ...models.autoencoders.vae import DecoderOutput -from ...models.modeling_utils import ModelMixin +from ..autoencoders.vae import DecoderOutput +from ..modeling_utils import ModelMixin class Kandinsky6SRLatentUpscalerConv3d(nn.Conv3d): diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index ea6d40340b5a..620c5e3ee9c7 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -854,10 +854,10 @@ def forward( `visual_token_type_num_embeddings > 0`. sparse_params (`dict`, *optional*): NABLA sparse-attention configuration for the video self-attention, dispatched through the `flex` - attention backend regardless of `transformer.set_attention_backend(...)`. This trades the regular - SDPA backend's dense attention for a block-sparse pattern that scales better to longer sequences, at - the cost of FlexAttention's one-time kernel-autotuning overhead on the first call. None of the - currently released checkpoints were trained with NABLA; this is forward-looking support for a future, + attention backend regardless of `transformer.set_attention_backend(...)`. This trades the regular SDPA + backend's dense attention for a block-sparse pattern that scales better to longer sequences, at the + cost of FlexAttention's one-time kernel-autotuning overhead on the first call. None of the currently + released checkpoints were trained with NABLA; this is forward-looking support for a future, longer-duration (e.g. 10s) model. return_dict (`bool`, defaults to `True`): Whether to return an [`AudioVisualModelOutput`] instead of a plain tuple. diff --git a/src/diffusers/pipelines/kandinsky6/__init__.py b/src/diffusers/pipelines/kandinsky6/__init__.py index 3ee290995ba8..23a2c473afd0 100644 --- a/src/diffusers/pipelines/kandinsky6/__init__.py +++ b/src/diffusers/pipelines/kandinsky6/__init__.py @@ -21,8 +21,6 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: - _import_structure["modeling_latent_upscaler"] = ["Kandinsky6SRLatentUpscalerBank"] - _import_structure["modeling_vocoder"] = ["MMAudioVocoder"] _import_structure["pipeline_kandinsky6_sr"] = ["Kandinsky6SRPipeline"] _import_structure["pipeline_kandinsky6_ti2va"] = ["Kandinsky6TI2VAPipeline"] _import_structure["pipeline_output"] = ["Kandinsky6SRPipelineOutput", "Kandinsky6TI2VAPipelineOutput"] @@ -34,8 +32,6 @@ except OptionalDependencyNotAvailable: from ...utils.dummy_torch_and_transformers_objects import * else: - from .modeling_latent_upscaler import Kandinsky6SRLatentUpscalerBank - from .modeling_vocoder import MMAudioVocoder from .pipeline_kandinsky6_sr import Kandinsky6SRPipeline from .pipeline_kandinsky6_ti2va import Kandinsky6TI2VAPipeline from .pipeline_output import Kandinsky6SRPipelineOutput, Kandinsky6TI2VAPipelineOutput diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py index b83eb8c36582..8c8aad02cc3f 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py @@ -20,13 +20,12 @@ import torch import torch.nn.functional as F -from ...models import Kandinsky6SRTransformer3DModel, Kandinsky6SRVAE +from ...models import Kandinsky6SRLatentUpscalerBank, Kandinsky6SRTransformer3DModel, Kandinsky6SRVAE from ...schedulers import FlowMatchEulerDiscreteScheduler, PiflowScheduler from ...utils import replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline -from .modeling_latent_upscaler import Kandinsky6SRLatentUpscalerBank from .pipeline_output import Kandinsky6SRPipelineOutput diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index f8973983335d..1ec53e071ee0 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -22,13 +22,12 @@ from transformers import CLIPTextModel, CLIPTokenizer, Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor from ...image_processor import PipelineImageInput -from ...models import AutoencoderKLHunyuanVideo, Kandinsky6Transformer3DModel, MMAudioVAE +from ...models import AutoencoderKLHunyuanVideo, Kandinsky6Transformer3DModel, MMAudioVAE, MMAudioVocoder from ...schedulers import FlowMatchEulerDiscreteScheduler, PiflowScheduler from ...utils import logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline -from .modeling_vocoder import MMAudioVocoder from .pipeline_output import Kandinsky6TI2VAPipelineOutput @@ -409,8 +408,8 @@ def expand_prompts( Maximum number of generated tokens per prompt. generator (`torch.Generator` or `list[torch.Generator]`, *optional*): Seeds the sampled expansion; a list must match `prompt`'s length, one generator per item. `generate` - draws from the global RNG, so the global RNG is seeded from this generator's seed; later - `randn_tensor` calls keep using `generator` directly. + draws from the global RNG, so the global RNG is seeded from this generator's seed; later `randn_tensor` + calls keep using `generator` directly. Returns: `str` or `list[str]`: The expanded prompt(s). @@ -530,8 +529,8 @@ def encode_image( r""" Encodes the reference image(s) into first-frame latents of shape `(batch_size, latent_height, latent_width, latent_channels)`, scaled by the VAE `scaling_factor`. PIL images are resized and center-cropped to `height x - width`; tensors and arrays must already have that size. The latents are repeated `num_videos_per_prompt` - times along the batch dimension. + width`; tensors and arrays must already have that size. The latents are repeated `num_videos_per_prompt` times + along the batch dimension. """ is_pil = isinstance(image, PIL.Image.Image) or ( isinstance(image, list) and isinstance(image[0], PIL.Image.Image) diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 33219521ba98..b105c80e6436 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -1590,6 +1590,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class Kandinsky6SRLatentUpscalerBank(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 Kandinsky6SRTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] @@ -1875,6 +1890,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class MMAudioVocoder(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 MochiTransformer3DModel(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 7a173941256a..a1f889a6dae6 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -2672,21 +2672,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class Kandinsky6SRLatentUpscalerBank(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 Kandinsky6SRPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] @@ -3482,21 +3467,6 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class MMAudioVocoder(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 MochiPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/autoencoders/test_models_autoencoder_mmaudio.py b/tests/models/autoencoders/test_models_autoencoder_mmaudio.py index 13481f2b9407..a12e01ee4049 100644 --- a/tests/models/autoencoders/test_models_autoencoder_mmaudio.py +++ b/tests/models/autoencoders/test_models_autoencoder_mmaudio.py @@ -74,7 +74,7 @@ def input_shape(self) -> tuple[int, ...]: @property def output_shape(self) -> tuple[int, ...]: # `decode`/`forward` return a mel spectrogram (`mel_bins`, num_mel_frames); the waveform is produced by the - # separate, pipeline-local `MMAudioVocoder`. + # separate `MMAudioVocoder`. return (8, 64) diff --git a/tests/pipelines/kandinsky6/test_models_vocoder.py b/tests/models/autoencoders/test_models_vocoder.py similarity index 96% rename from tests/pipelines/kandinsky6/test_models_vocoder.py rename to tests/models/autoencoders/test_models_vocoder.py index d54575bae7a2..225cfdbf4386 100644 --- a/tests/pipelines/kandinsky6/test_models_vocoder.py +++ b/tests/models/autoencoders/test_models_vocoder.py @@ -15,16 +15,16 @@ import torch -from diffusers.pipelines.kandinsky6 import MMAudioVocoder +from diffusers import MMAudioVocoder from diffusers.utils.torch_utils import randn_tensor -from ...models.testing_utils import ( +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TorchCompileTesterMixin, ) -from ...testing_utils import enable_full_determinism, torch_device enable_full_determinism() diff --git a/tests/models/latent_upscaler/__init__.py b/tests/models/latent_upscaler/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/pipelines/kandinsky6/test_models_latent_upscaler.py b/tests/models/latent_upscaler/test_models_latent_upscaler.py similarity index 96% rename from tests/pipelines/kandinsky6/test_models_latent_upscaler.py rename to tests/models/latent_upscaler/test_models_latent_upscaler.py index 2ba43be3a69e..059041af075d 100644 --- a/tests/pipelines/kandinsky6/test_models_latent_upscaler.py +++ b/tests/models/latent_upscaler/test_models_latent_upscaler.py @@ -15,16 +15,16 @@ import torch -from diffusers.pipelines.kandinsky6 import Kandinsky6SRLatentUpscalerBank +from diffusers import Kandinsky6SRLatentUpscalerBank from diffusers.utils.torch_utils import randn_tensor -from ...models.testing_utils import ( +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TorchCompileTesterMixin, ) -from ...testing_utils import enable_full_determinism, torch_device enable_full_determinism() diff --git a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_sr.py b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_sr.py index bc8ab010c3bf..3a31c1e77c5b 100644 --- a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_sr.py +++ b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_sr.py @@ -16,11 +16,11 @@ from diffusers import ( FlowMatchEulerDiscreteScheduler, + Kandinsky6SRLatentUpscalerBank, Kandinsky6SRPipeline, Kandinsky6SRTransformer3DModel, Kandinsky6SRVAE, ) -from diffusers.pipelines.kandinsky6 import Kandinsky6SRLatentUpscalerBank from ...testing_utils import torch_device from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin From 0b2ca95aa897bdc6f50e906248d627bae624cee2 Mon Sep 17 00:00:00 2001 From: leffff Date: Fri, 2 Oct 2026 21:40:47 +0300 Subject: [PATCH 40/62] add coppied from resolve i2v cond mode --- docs/source/en/api/pipelines/kandinsky6.md | 4 -- .../autoencoders/autoencoder_kandinsky6_sr.py | 12 ++--- .../transformers/transformer_kandinsky6.py | 11 ++--- .../kandinsky6/pipeline_kandinsky6_ti2va.py | 48 +++++-------------- .../test_pipeline_kandinsky6_ti2va.py | 6 --- 5 files changed, 21 insertions(+), 60 deletions(-) diff --git a/docs/source/en/api/pipelines/kandinsky6.md b/docs/source/en/api/pipelines/kandinsky6.md index 34583cc07b9a..bef76ac4f8bf 100644 --- a/docs/source/en/api/pipelines/kandinsky6.md +++ b/docs/source/en/api/pipelines/kandinsky6.md @@ -110,10 +110,6 @@ pipe.vae.enable_tiling() - [`Kandinsky6SRPipeline`]'s input `video` must have `1 + k * vae_scale_factor_temporal` frames for some integer `k` (a temporal compression ratio of `4` with the default K-VAE configuration, so `121` frames works but `120` doesn't) — trim or pad a video that doesn't already satisfy this before upscaling it. -- `visual_cond_scheme` is inferred automatically from whether `image` is passed (`"tail_cond_first_frame"` when it - is, `"pretrain"` otherwise). Only set it explicitly together with a matching `image`: passing - `visual_cond_scheme="i2v"` or `"tail_cond_first_frame"` without `image` doesn't raise an error, it silently falls - back to text-only generation. - `expand_prompts=True` reuses the already-loaded Qwen2.5-VL text encoder for an extra generation pass before denoising, so it adds latency but no extra model weights. - Compile the repeated transformer blocks for faster repeated inference: diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index a8dcf117173b..c92d8833db58 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -149,7 +149,7 @@ def forward(self, hidden_states: torch.Tensor, cache: dict) -> torch.Tensor: return output -class Kandinsky6SRRMSNorm(nn.Module): +class Kandinsky6VAERMSNorm(nn.Module): """RMS normalization over the channel axis of a `(B, C, T, H, W)` tensor, computed in float32. Same idea as `WanRMS_norm` (`autoencoder_kl_wan.py`) / `QwenImageRMS_norm` (`autoencoder_kl_qwenimage.py`), but not @@ -174,7 +174,7 @@ class Kandinsky6SRSpatialNorm3D(nn.Module): `SpatialNorm` (`attention_processor.py`) / `CogVideoXSpatialNorm3D` (`autoencoder_kl_cogvideox.py`) — normalize the features, then scale and shift by convolutions of the upsampled `zq`, carrying a `cache` dict across segments the same way `CogVideoXSpatialNorm3D` threads its `conv_cache` — but it is not a `# Copied from` of either: it - normalizes with `Kandinsky6SRRMSNorm` instead of `GroupNorm` (matching the plain-RMSNorm blocks elsewhere in this + normalizes with `Kandinsky6VAERMSNorm` instead of `GroupNorm` (matching the plain-RMSNorm blocks elsewhere in this VAE), and interpolates `zq` in channel chunks to bound peak memory. `zq` is nearest-upsampled to the feature grid. In the first segment the first frame is upsampled separately, @@ -183,7 +183,7 @@ class Kandinsky6SRSpatialNorm3D(nn.Module): def __init__(self, num_channels: int, zq_channels: int) -> None: super().__init__() - self.norm_layer = Kandinsky6SRRMSNorm(num_channels) + self.norm_layer = Kandinsky6VAERMSNorm(num_channels) self.conv_y = Kandinsky6SRSafeConv3d(zq_channels, num_channels, kernel_size=1) self.conv_b = Kandinsky6SRSafeConv3d(zq_channels, num_channels, kernel_size=1) @@ -231,8 +231,8 @@ def __init__(self, in_channels: int, out_channels: int, zq_channels: int | None self.in_channels = in_channels self.out_channels = out_channels if zq_channels is None: - self.norm1 = Kandinsky6SRRMSNorm(in_channels) - self.norm2 = Kandinsky6SRRMSNorm(out_channels) + self.norm1 = Kandinsky6VAERMSNorm(in_channels) + self.norm2 = Kandinsky6VAERMSNorm(out_channels) else: self.norm1 = Kandinsky6SRSpatialNorm3D(in_channels, zq_channels) self.norm2 = Kandinsky6SRSpatialNorm3D(out_channels, zq_channels) @@ -396,7 +396,7 @@ def __init__( self.mid = nn.Module() self.mid.block_1 = Kandinsky6SRResnetBlock3D(block_in, block_in) self.mid.block_2 = Kandinsky6SRResnetBlock3D(block_in, block_in) - self.norm_out = Kandinsky6SRRMSNorm(block_in) + self.norm_out = Kandinsky6VAERMSNorm(block_in) self.conv_out = Kandinsky6SRCausalConv3d(block_in, 2 * latent_channels, kernel_size=3) def make_cache(self) -> dict: diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 620c5e3ee9c7..e0080932fb87 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -28,7 +28,7 @@ from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn from ..cache_utils import CacheMixin -from ..embeddings import get_timestep_embedding +from ..embeddings import Timesteps from ..modeling_utils import ModelMixin, get_parameter_dtype from .transformer_ltx2 import AudioVisualModelOutput @@ -266,12 +266,11 @@ def __call__( class Kandinsky6TimeEmbeddings(nn.Module): """Sinusoidal timestep embedding with a K6-compatible parameter layout.""" - def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): + def __init__(self, model_dim: int, time_dim: int): super().__init__() if model_dim % 2: raise ValueError("model_dim must be even") - self.model_dim = model_dim - self.max_period = max_period + self.time_proj = Timesteps(model_dim, flip_sin_to_cos=True, downscale_freq_shift=0) self.in_layer = nn.Linear(model_dim, time_dim) self.activation = nn.SiLU() self.out_layer = nn.Linear(time_dim, time_dim) @@ -279,9 +278,7 @@ def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): def forward(self, timestep: Tensor) -> Tensor: # The sinusoidal embedding is float32; `_keep_in_fp32_modules` keeps these layers float32 under # `from_pretrained(torch_dtype=...)`, and the cast aligns the input with whatever dtype they hold. - embed = get_timestep_embedding( - timestep, self.model_dim, flip_sin_to_cos=True, downscale_freq_shift=0, max_period=self.max_period - ) + embed = self.time_proj(timestep) embed = embed.to(get_parameter_dtype(self.in_layer)) return self.out_layer(self.activation(self.in_layer(embed))) diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index 1ec53e071ee0..17ff490e9b22 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -455,7 +455,6 @@ def check_inputs( width, num_frames, image, - visual_cond_scheme, sample_audio, prompt_embeds=None, pooled_prompt_embeds=None, @@ -501,19 +500,12 @@ def check_inputs( if sample_audio and (getattr(self, "audio_vae", None) is None or getattr(self, "vocoder", None) is None): raise ValueError("`sample_audio=True` requires an `audio_vae` and a `vocoder`.") - if visual_cond_scheme not in ("pretrain", "i2v", "tail_cond_first_frame"): - raise ValueError( - f"`visual_cond_scheme` must be one of 'pretrain', 'i2v' or 'tail_cond_first_frame' but is {visual_cond_scheme!r}." - ) if image is not None: if not self.transformer.config.visual_cond: raise ValueError("Image conditioning requires a transformer with `visual_cond=True`.") - if ( - visual_cond_scheme == "tail_cond_first_frame" - and self.transformer.config.visual_token_type_num_embeddings < 2 - ): + if self.transformer.config.visual_token_type_num_embeddings < 2: raise ValueError( - "`tail_cond_first_frame` requires a transformer with `visual_token_type_num_embeddings >= 2`." + "Image conditioning requires a transformer with `visual_token_type_num_embeddings >= 2`." ) def encode_image( @@ -634,7 +626,6 @@ def __call__( negative_pooled_prompt_embeds: torch.Tensor | None = None, sample_audio: bool | None = None, expand_prompts: bool = False, - visual_cond_scheme: str | None = None, max_sequence_length: int = 1024, output_type: str = "pil", return_dict: bool = True, @@ -689,9 +680,6 @@ def __call__( `vocoder`, `False` otherwise; `True` requires both. expand_prompts (`bool`, defaults to `False`): Whether to rewrite the prompts with [`~Kandinsky6TI2VAPipeline.expand_prompts`] before encoding. - visual_cond_scheme (`str`, *optional*): - How the reference image conditions the video: `"tail_cond_first_frame"` (default with `image`), - `"i2v"`, or `"pretrain"` (default without `image`). max_sequence_length (`int`, defaults to `1024`): Maximum number of prompt tokens after the chat template. output_type (`str`, defaults to `"pil"`): @@ -713,8 +701,6 @@ def __call__( # 1. Check inputs. Raise error if not correct if sample_audio is None: sample_audio = getattr(self, "audio_vae", None) is not None and getattr(self, "vocoder", None) is not None - if visual_cond_scheme is None: - visual_cond_scheme = "tail_cond_first_frame" if image is not None else "pretrain" self.check_inputs( prompt=prompt, negative_prompt=negative_prompt, @@ -722,7 +708,6 @@ def __call__( width=width, num_frames=num_frames, image=image, - visual_cond_scheme=visual_cond_scheme, sample_audio=sample_audio, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, @@ -737,9 +722,6 @@ def __call__( num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1 num_frames = max(num_frames, 1) - if isinstance(self.scheduler, PiflowScheduler) and guidance_scale != 1.0: - raise ValueError("`PiflowScheduler` runs the distilled checkpoints and requires `guidance_scale=1.0`.") - self._guidance_scale = guidance_scale self._current_timestep = None self._interrupt = False @@ -819,9 +801,9 @@ def __call__( else: audio_latents = None - # The `tail_cond_first_frame` scheme appends the clean reference frame as an extra, masked frame that reuses - # the first frame's rotary position and carries token type `1`. - tail_cond = image is not None and visual_cond_scheme == "tail_cond_first_frame" + # Image conditioning appends the clean reference frame as an extra, masked frame that reuses the first + # frame's rotary position and carries token type `1`. + tail_cond = image is not None visual_token_type_ids = None visual_rope_pos = None if tail_cond: @@ -848,22 +830,16 @@ def __call__( self._current_timestep = t timestep = t.expand(batch_size) - # Visual conditioning channels: [latent | conditioning latent | mask]. The reference frame either - # fills the conditioning channels (`pretrain`) or overwrites a latent frame (`i2v`, tail). + # Visual conditioning channels: [latent | conditioning latent | mask]. The conditioning-latent + # channel is unused (always zero) and only kept to match the transformer's fixed input width; the + # reference frame itself is appended as an extra, masked tail frame (see `tail_cond` above). latent_model_input = latents if self.transformer.config.visual_cond: cond_latents = torch.zeros_like(latents) cond_mask = torch.zeros((*latents.shape[:-1], 1), dtype=latents.dtype, device=device) if first_frame_latents is not None: - if visual_cond_scheme == "pretrain": - cond_latents[:, 0] = first_frame_latents - cond_mask[:, 0] = 1 - elif visual_cond_scheme == "i2v": - latents[:, 0] = first_frame_latents - cond_mask[:, 0] = 1 - else: - latents[:, -1] = first_frame_latents - cond_mask[:, -1] = 1 + latents[:, -1] = first_frame_latents + cond_mask[:, -1] = 1 latent_model_input = torch.cat([latents, cond_latents, cond_mask], dim=-1) with self.transformer.cache_context("cond"): @@ -914,9 +890,7 @@ def __call__( self._current_timestep = None - # 8. Keep the injected reference frame clean and drop the appended tail frame - if first_frame_latents is not None and visual_cond_scheme == "i2v": - latents[:, 0] = first_frame_latents + # 8. Drop the appended tail frame used for reference-image conditioning if tail_cond: latents = latents[:, :-1] diff --git a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py index 316c36fa7755..ac68deeddbf1 100644 --- a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py +++ b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py @@ -237,18 +237,12 @@ def test_kandinsky6_ti2va_i2va(self): assert not torch.isnan(output.frames.float()).any() def test_kandinsky6_ti2va_different_images(self): - # A freshly initialized transformer zero-initializes every AdaLN gate, so the text and reference-image paths - # through attention are identity at init. The `i2v` scheme instead writes the encoded image into frame 0, - # which passes through the real `visual_embeddings` and `out_layer` projections and must change the output. pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) inputs = self.get_dummy_inputs() - inputs["visual_cond_scheme"] = "i2v" - inputs["image"] = PIL.Image.fromarray(np.zeros((16, 16, 3), dtype=np.uint8)) output_black_image = pipe(**inputs).frames inputs = self.get_dummy_inputs() - inputs["visual_cond_scheme"] = "i2v" inputs["image"] = PIL.Image.fromarray(np.full((16, 16, 3), 255, dtype=np.uint8)) output_white_image = pipe(**inputs).frames From 6f035c44c6de67976e69534dce93979e8bd19af8 Mon Sep 17 00:00:00 2001 From: leffff Date: Sat, 3 Oct 2026 11:11:20 +0300 Subject: [PATCH 41/62] rename variables --- .../latent_upscaler_kandinsky6_sr.py | 40 +++++------ .../transformers/transformer_kandinsky6.py | 69 ++++++++----------- .../transformers/transformer_kandinsky6_sr.py | 37 +++------- 3 files changed, 55 insertions(+), 91 deletions(-) diff --git a/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py b/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py index b3e10416a7d8..b97d605c8839 100644 --- a/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py +++ b/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py @@ -104,19 +104,17 @@ def forward(self, hidden_states: Tensor) -> Tensor: return self.linear(hidden_states + self.spatial_conv(hidden_states)) -class Kandinsky6SRLatentUpscalerOutputHead(nn.Sequential): +class Kandinsky6SRLatentUpscalerOutputHead(nn.Module): """`modulated norm -> SiLU -> conv3x3x3` projection back to the latent channels, threading `zq` into the norm.""" def __init__(self, in_channels: int, out_channels: int, zq_channels: int) -> None: - super().__init__( - Kandinsky6SRLatentUpscalerModulatedNorm(in_channels, zq_channels), - nn.SiLU(), - Kandinsky6SRLatentUpscalerConv3d(in_channels, out_channels, kernel_size=3), - ) + super().__init__() + self.norm = Kandinsky6SRLatentUpscalerModulatedNorm(in_channels, zq_channels) + self.activation = nn.SiLU() + self.conv = Kandinsky6SRLatentUpscalerConv3d(in_channels, out_channels, kernel_size=3) def forward(self, hidden_states: Tensor, zq: Tensor) -> Tensor: - norm, activation, conv = self - return conv(activation(norm(hidden_states, zq))) + return self.conv(self.activation(self.norm(hidden_states, zq))) def _residual_stack(num_blocks: int, in_channels: int, out_channels: int, zq_channels: int) -> nn.ModuleList: @@ -147,21 +145,21 @@ def __init__( self.finisher = nn.Module() self.finisher.spatial_conv = nn.Conv3d(width_1, width_1, kernel_size=(1, 3, 3), padding=(0, 1, 1)) self.finisher.linear = nn.Conv3d(width_1, width_1, kernel_size=1) - self.private_mid_blocks = _residual_stack(num_mid_blocks, width_1, width_2, in_channels) - self.private_upsample = Kandinsky6SRLatentUpscalerUpsample(width_2) - self.private_blocks = _residual_stack(num_post_blocks, width_2, width_3, in_channels) - self.private_output_proj = Kandinsky6SRLatentUpscalerOutputHead(width_3, in_channels, in_channels) + self.mid_blocks = _residual_stack(num_mid_blocks, width_1, width_2, in_channels) + self.upsample = Kandinsky6SRLatentUpscalerUpsample(width_2) + self.blocks = _residual_stack(num_post_blocks, width_2, width_3, in_channels) + self.output_proj = Kandinsky6SRLatentUpscalerOutputHead(width_3, in_channels, in_channels) def forward(self, hidden_states: Tensor, zq: Tensor) -> Tensor: for block in self.adapter: hidden_states = block(hidden_states, zq) hidden_states = self.finisher.linear(hidden_states + self.finisher.spatial_conv(hidden_states)) - for block in self.private_mid_blocks: + for block in self.mid_blocks: hidden_states = block(hidden_states, zq) - hidden_states = self.private_upsample(hidden_states) - for block in self.private_blocks: + hidden_states = self.upsample(hidden_states) + for block in self.blocks: hidden_states = block(hidden_states, zq) - return self.private_output_proj(hidden_states, zq) + return self.output_proj(hidden_states, zq) class Kandinsky6SRLatentUpscaler(nn.Module): @@ -257,9 +255,9 @@ def __init__( scaling_factor: float = 0.910344004631042, ) -> None: super().__init__() - self._models = nn.ModuleDict( - { - f"{scale}x": Kandinsky6SRLatentUpscaler( + self._models = nn.ModuleList( + [ + Kandinsky6SRLatentUpscaler( scale=scale, in_channels=in_channels, stage_channels=stage_channels, @@ -269,7 +267,7 @@ def __init__( num_x2_adapter_blocks=num_x2_adapter_blocks, ) for scale in scales - } + ] ) def forward(self, latents: Tensor, scale: int, return_dict: bool = True) -> DecoderOutput | tuple[Tensor]: @@ -287,7 +285,7 @@ def forward(self, latents: Tensor, scale: int, return_dict: bool = True) -> Deco """ if scale not in self.config.scales: raise ValueError(f"No latent upscaler for scale {scale}; available scales: {list(self.config.scales)}") - upscaled = self._models[f"{scale}x"](latents) + upscaled = self._models[self.config.scales.index(scale)](latents) if not return_dict: return (upscaled,) return DecoderOutput(sample=upscaled) diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index e0080932fb87..9bfddfd40587 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -25,10 +25,10 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin -from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn from ..cache_utils import CacheMixin -from ..embeddings import Timesteps +from ..embeddings import TimestepEmbedding, Timesteps from ..modeling_utils import ModelMixin, get_parameter_dtype from .transformer_ltx2 import AudioVisualModelOutput @@ -271,16 +271,14 @@ def __init__(self, model_dim: int, time_dim: int): if model_dim % 2: raise ValueError("model_dim must be even") self.time_proj = Timesteps(model_dim, flip_sin_to_cos=True, downscale_freq_shift=0) - self.in_layer = nn.Linear(model_dim, time_dim) - self.activation = nn.SiLU() - self.out_layer = nn.Linear(time_dim, time_dim) + self.timestep_embedder = TimestepEmbedding(model_dim, time_dim, act_fn="silu") def forward(self, timestep: Tensor) -> Tensor: # The sinusoidal embedding is float32; `_keep_in_fp32_modules` keeps these layers float32 under # `from_pretrained(torch_dtype=...)`, and the cast aligns the input with whatever dtype they hold. embed = self.time_proj(timestep) - embed = embed.to(get_parameter_dtype(self.in_layer)) - return self.out_layer(self.activation(self.in_layer(embed))) + embed = embed.to(get_parameter_dtype(self.timestep_embedder)) + return self.timestep_embedder(embed) class Kandinsky6TextEmbeddings(nn.Module): @@ -328,19 +326,6 @@ def forward(self, x: Tensor) -> Tensor: return self.out_layer(self.activation(x.to(get_parameter_dtype(self.out_layer)))) -class Kandinsky6FeedForward(nn.Module): - """K6 bias-free GELU feed-forward network.""" - - def __init__(self, dim: int, ff_dim: int): - 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: Tensor) -> Tensor: - return self.out_layer(self.activation(self.in_layer(x))) - - class Kandinsky6Attention(nn.Module, AttentionModuleMixin): """K6 attention with the Diffusers ``set_processor`` contract.""" @@ -455,18 +440,18 @@ def __init__( ): super().__init__() self.text_modulation = Kandinsky6Modulation(time_dim, model_dim, 6) - self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) - self.self_attention = Kandinsky6Attention(model_dim, head_dim) + self.attn_norm = nn.LayerNorm(model_dim, elementwise_affine=False) + self.attn = Kandinsky6Attention(model_dim, head_dim) self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) - self.feed_forward = Kandinsky6FeedForward(model_dim, ff_dim) + self.feed_forward = FeedForward(model_dim, inner_dim=ff_dim, activation_fn="gelu", bias=False) def forward(self, x: Tensor, time_embed: Tensor, rope: Tensor, attn_mask: Tensor | None = None) -> Tensor: sa_params, ff_params = torch.chunk(self.text_modulation(time_embed), 2, dim=-1) shift, scale, gate = torch.chunk(sa_params, 3, dim=-1) x = apply_gate_sum( x, - self.self_attention( - apply_scale_shift(self.self_attention_norm(x.float()), x, scale, shift), + self.attn( + apply_scale_shift(self.attn_norm(x.float()), x, scale, shift), rotary_emb=rope, attn_mask=attn_mask, ), @@ -504,7 +489,7 @@ def __init__( kv_dim=model_dim, ) self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) - self.feed_forward = Kandinsky6FeedForward(model_dim, ff_dim) + self.feed_forward = FeedForward(model_dim, inner_dim=ff_dim, activation_fn="gelu", bias=False) class Kandinsky6FusedTransformerDecoderBlock(nn.Module): @@ -525,8 +510,8 @@ def __init__( fix_modulation: bool = False, ): super().__init__() - self.videoT = Kandinsky6TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim) - self.audioT = Kandinsky6TransformerDecoderBlock(model_dim_a, time_dim_a, ff_dim_a, head_dim_a) + self.video_dec_block = Kandinsky6TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim) + self.audio_dec_block = Kandinsky6TransformerDecoderBlock(model_dim_a, time_dim_a, ff_dim_a, head_dim_a) self.va_cross_attention = Kandinsky6Attention( model_dim, head_dim, @@ -569,39 +554,39 @@ def forward( ) -> tuple[Tensor | None, Tensor | None]: t_v, t_a = time_embed if vis is not None: - sa_p, ca_p, ff_p = torch.chunk(self.videoT.visual_modulation(t_v), 3, dim=-1) + sa_p, ca_p, ff_p = torch.chunk(self.video_dec_block.visual_modulation(t_v), 3, dim=-1) shift, scale, gate = torch.chunk(sa_p, 3, dim=-1) vis = apply_gate_sum( vis, - self.videoT.self_attention( - apply_scale_shift(self.videoT.self_attention_norm(vis.float()), vis, scale, shift), + self.video_dec_block.self_attention( + apply_scale_shift(self.video_dec_block.self_attention_norm(vis.float()), vis, scale, shift), rotary_emb=vis_rope, sparse_params=sparse_params, ), gate, ).type_as(vis) shift, scale, gate_v = torch.chunk(ca_p, 3, dim=-1) - vis_pre_ca = apply_scale_shift(self.videoT.cross_attention_norm(vis.float()), vis, scale, shift) - vis_out_t = self.videoT.cross_attention( + vis_pre_ca = apply_scale_shift(self.video_dec_block.cross_attention_norm(vis.float()), vis, scale, shift) + vis_out_t = self.video_dec_block.cross_attention( vis_pre_ca, encoder_hidden_states=text_v, attn_mask=attn_mask, ) if aud is not None: - sa_p, ca_p, ff_p_a = torch.chunk(self.audioT.visual_modulation(t_a), 3, dim=-1) + sa_p, ca_p, ff_p_a = torch.chunk(self.audio_dec_block.visual_modulation(t_a), 3, dim=-1) shift, scale, gate = torch.chunk(sa_p, 3, dim=-1) aud = apply_gate_sum( aud, - self.audioT.self_attention( - apply_scale_shift(self.audioT.self_attention_norm(aud.float()), aud, scale, shift), + self.audio_dec_block.self_attention( + apply_scale_shift(self.audio_dec_block.self_attention_norm(aud.float()), aud, scale, shift), rotary_emb=aud_rope, ), gate, ).type_as(aud) shift, scale, gate_a = torch.chunk(ca_p, 3, dim=-1) - aud_pre_ca = apply_scale_shift(self.audioT.cross_attention_norm(aud.float()), aud, scale, shift) - aud_out_t = self.audioT.cross_attention( + aud_pre_ca = apply_scale_shift(self.audio_dec_block.cross_attention_norm(aud.float()), aud, scale, shift) + aud_out_t = self.audio_dec_block.cross_attention( aud_pre_ca, encoder_hidden_states=text_a, attn_mask=attn_mask, @@ -649,8 +634,8 @@ def forward( shift, scale, gate = torch.chunk(ff_p, 3, dim=-1) vis = apply_gate_sum( vis, - self.videoT.feed_forward( - apply_scale_shift(self.videoT.feed_forward_norm(vis.float()), vis, scale, shift) + self.video_dec_block.feed_forward( + apply_scale_shift(self.video_dec_block.feed_forward_norm(vis.float()), vis, scale, shift) ), gate, ).type_as(vis) @@ -658,8 +643,8 @@ def forward( shift, scale, gate = torch.chunk(ff_p_a, 3, dim=-1) aud = apply_gate_sum( aud, - self.audioT.feed_forward( - apply_scale_shift(self.audioT.feed_forward_norm(aud.float()), aud, scale, shift) + self.audio_dec_block.feed_forward( + apply_scale_shift(self.audio_dec_block.feed_forward_norm(aud.float()), aud, scale, shift) ), gate, ).type_as(aud) diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index d31307c997de..d796c954a18f 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -25,9 +25,9 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin -from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn -from ..embeddings import get_timestep_embedding +from ..embeddings import TimestepEmbedding, Timesteps from ..modeling_outputs import Transformer2DModelOutput from ..modeling_utils import ModelMixin, get_parameter_dtype @@ -234,24 +234,19 @@ def forward( class Kandinsky6SRTimeEmbeddings(nn.Module): """Sinusoidal timestep embedding with a K6-compatible parameter layout.""" - def __init__(self, model_dim: int, time_dim: int, max_period: float = 10000.0): + def __init__(self, model_dim: int, time_dim: int): super().__init__() if model_dim % 2: raise ValueError("model_dim must be even") - self.model_dim = model_dim - self.max_period = max_period - self.in_layer = nn.Linear(model_dim, time_dim) - self.activation = nn.SiLU() - self.out_layer = nn.Linear(time_dim, time_dim) + self.time_proj = Timesteps(model_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding(model_dim, time_dim, act_fn="silu") def forward(self, timestep: Tensor) -> Tensor: # The sinusoidal embedding is float32; `_keep_in_fp32_modules` keeps these layers float32 under # `from_pretrained(torch_dtype=...)`, and the cast aligns the input with whatever dtype they hold. - embed = get_timestep_embedding( - timestep, self.model_dim, flip_sin_to_cos=True, downscale_freq_shift=0, max_period=self.max_period - ) - embed = embed.to(get_parameter_dtype(self.in_layer)) - return self.out_layer(self.activation(self.in_layer(embed))) + embed = self.time_proj(timestep) + embed = embed.to(get_parameter_dtype(self.timestep_embedder)) + return self.timestep_embedder(embed) # Copied from diffusers.models.transformers.transformer_kandinsky6.Kandinsky6VisualEmbeddings with Kandinsky6->Kandinsky6SR @@ -289,20 +284,6 @@ def forward(self, x: Tensor) -> Tensor: return self.out_layer(self.activation(x.to(get_parameter_dtype(self.out_layer)))) -# Copied from diffusers.models.transformers.transformer_kandinsky6.Kandinsky6FeedForward with Kandinsky6->Kandinsky6SR -class Kandinsky6SRFeedForward(nn.Module): - """K6 bias-free GELU feed-forward network.""" - - def __init__(self, dim: int, ff_dim: int): - 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: Tensor) -> Tensor: - return self.out_layer(self.activation(self.in_layer(x))) - - # Copied from diffusers.models.transformers.transformer_kandinsky6.Kandinsky6OutLayer with Kandinsky6->Kandinsky6SR class Kandinsky6SROutLayer(nn.Module): """Projects visual hidden states back to packed latent patches.""" @@ -390,7 +371,7 @@ def __init__(self, model_dim: int, time_dim: int, ff_dim: int, head_dim: int): self.self_attention_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.self_attention = Kandinsky6SRAttention(model_dim, head_dim) self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) - self.feed_forward = Kandinsky6SRFeedForward(model_dim, ff_dim) + self.feed_forward = FeedForward(model_dim, inner_dim=ff_dim, activation_fn="gelu", bias=False) def forward( self, From 1cfe98c248ef786b5a0cfaa531707942a56db965 Mon Sep 17 00:00:00 2001 From: leffff Date: Sat, 3 Oct 2026 11:33:49 +0300 Subject: [PATCH 42/62] fix lu rename bug --- .../latent_upscaler_kandinsky6_sr.py | 37 +++++++++++++------ 1 file changed, 26 insertions(+), 11 deletions(-) diff --git a/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py b/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py index b97d605c8839..79117ae810c6 100644 --- a/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py +++ b/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py @@ -165,8 +165,18 @@ def forward(self, hidden_states: Tensor, zq: Tensor) -> Tensor: class Kandinsky6SRLatentUpscaler(nn.Module): """One latent upscaler of the bank: a two-stage 2x+2x cascade for `scale=4`, or its single-stage x2 variant. - Both share the same widths. The x4 model runs `input_proj -> pre_blocks -> upsample_1 -> mid_blocks -> upsample_2 - -> post_blocks -> output_proj`; the x2 model runs its own stem and the [`Kandinsky6SRLatentUpscalerX2Branch`] tail. + Both share the same widths and the same `input_proj -> pre_blocks -> upsample_1 -> mid_blocks -> upsample_2 -> + post_blocks -> output_proj` backbone (plus an auxiliary `mid_output_head`, a deep-supervision head the + checkpoint trains with zero loss weight -- see below). The x2 model additionally runs the + [`Kandinsky6SRLatentUpscalerX2Branch`] tail; the x4 model's `forward` is just its backbone. + + Released checkpoints train the x2 entry's backbone jointly with its [`Kandinsky6SRLatentUpscalerX2Branch`] + tail (the `x2_adapter_sources` config field names which backbone activations the tail's `adapter` reads from), + but `forward` here only runs the tail, matching this class's behavior before the backbone was known to exist: + the backbone and `mid_output_head` are declared so their trained weights load from the checkpoint instead of + raising "unused weights"/leaving parameters meta/randomly-initialized, but neither is wired into `forward`, so + numerically this is unchanged from before. Wiring the backbone into the x2 tail's forward pass needs the + original training code to confirm the exact tap points first -- guessing would risk silently wrong output. """ def __init__( @@ -185,15 +195,20 @@ def __init__( self.scale = scale width_1, width_2, width_3 = stage_channels - if scale == 4: - self.input_proj = nn.Sequential(Kandinsky6SRLatentUpscalerConv3d(in_channels, width_1, kernel_size=3)) - self.pre_blocks = _residual_stack(num_pre_blocks, width_1, width_1, in_channels) - self.upsample_1 = Kandinsky6SRLatentUpscalerUpsample(width_1) - self.mid_blocks = _residual_stack(num_mid_blocks, width_1, width_2, in_channels) - self.upsample_2 = Kandinsky6SRLatentUpscalerUpsample(width_2) - self.post_blocks = _residual_stack(num_post_blocks, width_2, width_3, in_channels) - self.output_proj = Kandinsky6SRLatentUpscalerOutputHead(width_3, in_channels, in_channels) - else: + self.input_proj = nn.Sequential(Kandinsky6SRLatentUpscalerConv3d(in_channels, width_1, kernel_size=3)) + self.pre_blocks = _residual_stack(num_pre_blocks, width_1, width_1, in_channels) + self.upsample_1 = Kandinsky6SRLatentUpscalerUpsample(width_1) + self.mid_blocks = _residual_stack(num_mid_blocks, width_1, width_2, in_channels) + self.mid_output_head = nn.Sequential( + Kandinsky6SRLatentUpscalerRMSNorm(width_2), + nn.SiLU(), + Kandinsky6SRLatentUpscalerConv3d(width_2, in_channels, kernel_size=3), + ) + self.upsample_2 = Kandinsky6SRLatentUpscalerUpsample(width_2) + self.post_blocks = _residual_stack(num_post_blocks, width_2, width_3, in_channels) + self.output_proj = Kandinsky6SRLatentUpscalerOutputHead(width_3, in_channels, in_channels) + + if scale == 2: self.mid_input_proj = nn.Sequential(Kandinsky6SRLatentUpscalerConv3d(in_channels, width_1, kernel_size=3)) self.x2_branch = Kandinsky6SRLatentUpscalerX2Branch( in_channels, stage_channels, num_x2_adapter_blocks, num_mid_blocks, num_post_blocks From f16156a36185bc51850c568e9fa8afc5f55d32f9 Mon Sep 17 00:00:00 2001 From: Lev Novitskiy <57654885+leffff@users.noreply.github.com> Date: Sun, 4 Oct 2026 00:07:01 +0300 Subject: [PATCH 43/62] Update src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py Co-authored-by: YiYi Xu --- src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index 17ff490e9b22..384b3ac09cba 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -734,7 +734,7 @@ def __call__( dtype = self.transformer.dtype # 3. Encode input prompt - if expand_prompts and prompt is not None: + if expand_prompts: prompt = self.expand_prompts( prompt, image=image, max_sequence_length=max_sequence_length, generator=generator ) From 348c34149ac2a92ef1efcd2739d13ff6a7c21e7f Mon Sep 17 00:00:00 2001 From: Lev Novitskiy <57654885+leffff@users.noreply.github.com> Date: Sun, 4 Oct 2026 00:07:19 +0300 Subject: [PATCH 44/62] Update src/diffusers/models/transformers/transformer_kandinsky6.py Co-authored-by: YiYi Xu --- src/diffusers/models/transformers/transformer_kandinsky6.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 9bfddfd40587..2188b1fe2e7c 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -319,8 +319,6 @@ 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) - nn.init.zeros_(self.out_layer.weight) - nn.init.zeros_(self.out_layer.bias) def forward(self, x: Tensor) -> Tensor: return self.out_layer(self.activation(x.to(get_parameter_dtype(self.out_layer)))) From 0ef584f7062438ef09b531e3927b745f3a7e5fab Mon Sep 17 00:00:00 2001 From: leffff Date: Sun, 4 Oct 2026 00:25:30 +0300 Subject: [PATCH 45/62] address yiyixuxu's comments --- .../latent_upscaler_kandinsky6_sr.py | 69 +++++---- .../transformers/transformer_kandinsky6.py | 134 +++--------------- .../transformers/transformer_kandinsky6_sr.py | 26 ++-- .../kandinsky6/pipeline_kandinsky6_sr.py | 5 +- .../kandinsky6/pipeline_kandinsky6_ti2va.py | 92 ++++++++---- .../test_pipeline_kandinsky6_ti2va.py | 41 +++++- 6 files changed, 177 insertions(+), 190 deletions(-) diff --git a/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py b/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py index 79117ae810c6..630c6edd488b 100644 --- a/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py +++ b/src/diffusers/models/latent_upscaler/latent_upscaler_kandinsky6_sr.py @@ -117,17 +117,6 @@ def forward(self, hidden_states: Tensor, zq: Tensor) -> Tensor: return self.conv(self.activation(self.norm(hidden_states, zq))) -def _residual_stack(num_blocks: int, in_channels: int, out_channels: int, zq_channels: int) -> nn.ModuleList: - """`num_blocks` residual blocks; the first one narrows `in_channels` to `out_channels`.""" - widths = [in_channels] + [out_channels] * num_blocks - return nn.ModuleList( - [ - Kandinsky6SRLatentUpscalerResidualBlock(widths[index], widths[index + 1], zq_channels) - for index in range(num_blocks) - ] - ) - - class Kandinsky6SRLatentUpscalerX2Branch(nn.Module): """The x2 upscaler tail: adapter blocks and a finisher on the input grid, then one 2x stage.""" @@ -141,13 +130,27 @@ def __init__( ) -> None: super().__init__() width_1, width_2, width_3 = stage_channels - self.adapter = _residual_stack(num_adapter_blocks, width_1, width_1, in_channels) + self.adapter = nn.ModuleList( + [Kandinsky6SRLatentUpscalerResidualBlock(width_1, width_1, in_channels) for _ in range(num_adapter_blocks)] + ) self.finisher = nn.Module() self.finisher.spatial_conv = nn.Conv3d(width_1, width_1, kernel_size=(1, 3, 3), padding=(0, 1, 1)) self.finisher.linear = nn.Conv3d(width_1, width_1, kernel_size=1) - self.mid_blocks = _residual_stack(num_mid_blocks, width_1, width_2, in_channels) + mid_widths = [width_1] + [width_2] * num_mid_blocks + self.mid_blocks = nn.ModuleList( + [ + Kandinsky6SRLatentUpscalerResidualBlock(mid_widths[index], mid_widths[index + 1], in_channels) + for index in range(num_mid_blocks) + ] + ) self.upsample = Kandinsky6SRLatentUpscalerUpsample(width_2) - self.blocks = _residual_stack(num_post_blocks, width_2, width_3, in_channels) + post_widths = [width_2] + [width_3] * num_post_blocks + self.blocks = nn.ModuleList( + [ + Kandinsky6SRLatentUpscalerResidualBlock(post_widths[index], post_widths[index + 1], in_channels) + for index in range(num_post_blocks) + ] + ) self.output_proj = Kandinsky6SRLatentUpscalerOutputHead(width_3, in_channels, in_channels) def forward(self, hidden_states: Tensor, zq: Tensor) -> Tensor: @@ -166,17 +169,17 @@ class Kandinsky6SRLatentUpscaler(nn.Module): """One latent upscaler of the bank: a two-stage 2x+2x cascade for `scale=4`, or its single-stage x2 variant. Both share the same widths and the same `input_proj -> pre_blocks -> upsample_1 -> mid_blocks -> upsample_2 -> - post_blocks -> output_proj` backbone (plus an auxiliary `mid_output_head`, a deep-supervision head the - checkpoint trains with zero loss weight -- see below). The x2 model additionally runs the + post_blocks -> output_proj` backbone (plus an auxiliary `mid_output_head`, a deep-supervision head the checkpoint + trains with zero loss weight -- see below). The x2 model additionally runs the [`Kandinsky6SRLatentUpscalerX2Branch`] tail; the x4 model's `forward` is just its backbone. - Released checkpoints train the x2 entry's backbone jointly with its [`Kandinsky6SRLatentUpscalerX2Branch`] - tail (the `x2_adapter_sources` config field names which backbone activations the tail's `adapter` reads from), - but `forward` here only runs the tail, matching this class's behavior before the backbone was known to exist: - the backbone and `mid_output_head` are declared so their trained weights load from the checkpoint instead of - raising "unused weights"/leaving parameters meta/randomly-initialized, but neither is wired into `forward`, so - numerically this is unchanged from before. Wiring the backbone into the x2 tail's forward pass needs the - original training code to confirm the exact tap points first -- guessing would risk silently wrong output. + Released checkpoints train the x2 entry's backbone jointly with its [`Kandinsky6SRLatentUpscalerX2Branch`] tail + (the `x2_adapter_sources` config field names which backbone activations the tail's `adapter` reads from), but + `forward` here only runs the tail, matching this class's behavior before the backbone was known to exist: the + backbone and `mid_output_head` are declared so their trained weights load from the checkpoint instead of raising + "unused weights"/leaving parameters meta/randomly-initialized, but neither is wired into `forward`, so numerically + this is unchanged from before. Wiring the backbone into the x2 tail's forward pass needs the original training code + to confirm the exact tap points first -- guessing would risk silently wrong output. """ def __init__( @@ -196,16 +199,30 @@ def __init__( width_1, width_2, width_3 = stage_channels self.input_proj = nn.Sequential(Kandinsky6SRLatentUpscalerConv3d(in_channels, width_1, kernel_size=3)) - self.pre_blocks = _residual_stack(num_pre_blocks, width_1, width_1, in_channels) + self.pre_blocks = nn.ModuleList( + [Kandinsky6SRLatentUpscalerResidualBlock(width_1, width_1, in_channels) for _ in range(num_pre_blocks)] + ) self.upsample_1 = Kandinsky6SRLatentUpscalerUpsample(width_1) - self.mid_blocks = _residual_stack(num_mid_blocks, width_1, width_2, in_channels) + mid_widths = [width_1] + [width_2] * num_mid_blocks + self.mid_blocks = nn.ModuleList( + [ + Kandinsky6SRLatentUpscalerResidualBlock(mid_widths[index], mid_widths[index + 1], in_channels) + for index in range(num_mid_blocks) + ] + ) self.mid_output_head = nn.Sequential( Kandinsky6SRLatentUpscalerRMSNorm(width_2), nn.SiLU(), Kandinsky6SRLatentUpscalerConv3d(width_2, in_channels, kernel_size=3), ) self.upsample_2 = Kandinsky6SRLatentUpscalerUpsample(width_2) - self.post_blocks = _residual_stack(num_post_blocks, width_2, width_3, in_channels) + post_widths = [width_2] + [width_3] * num_post_blocks + self.post_blocks = nn.ModuleList( + [ + Kandinsky6SRLatentUpscalerResidualBlock(post_widths[index], post_widths[index + 1], in_channels) + for index in range(num_post_blocks) + ] + ) self.output_proj = Kandinsky6SRLatentUpscalerOutputHead(width_3, in_channels, in_channels) if scale == 2: diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 9bfddfd40587..8ba18265a67c 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -21,12 +21,11 @@ import torch from torch import Tensor, nn -from torch.nn.attention.flex_attention import BlockMask from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward -from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn +from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin from ..embeddings import TimestepEmbedding, Timesteps from ..modeling_utils import ModelMixin, get_parameter_dtype @@ -72,55 +71,6 @@ def apply_rotary(x: Tensor, rope: Tensor) -> Tensor: return (rope.float() * x_).sum(dim=-1).reshape(*x.shape).to(dtype=x.dtype) -def _local_patch(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Tensor: - T, H, W = shape - g1, g2, g3 = group_size - x = x.reshape(*x.shape[:dim], T // g1, g1, H // g2, g2, W // g3, g3, *x.shape[dim + 3 :]) - d = len(x.shape[:dim]) - x = x.permute(*range(d), d, d + 2, d + 4, d + 1, d + 3, d + 5, *range(d + 6, len(x.shape))) - return x.flatten(dim, dim + 2).flatten(dim + 1, dim + 3) - - -def _local_merge(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Tensor: - T, H, W = shape - g1, g2, g3 = group_size - x = x.reshape(*x.shape[:dim], T // g1, H // g2, W // g3, g1, g2, g3, *x.shape[dim + 2 :]) - d = len(x.shape[:dim]) - x = x.permute(*range(d), d, d + 3, d + 1, d + 4, d + 2, d + 5, *range(d + 6, len(x.shape))) - return x.flatten(dim, dim + 1).flatten(dim + 1, dim + 2).flatten(dim + 2, dim + 3) - - -def nabla_block_mask( - q: Tensor, - k: Tensor, - sta: Tensor, - thr: float = 0.9, - block_size: int = 64, -) -> BlockMask: - """Build a dynamic NABLA block mask from query/key statistics and an STA (Sliding-Tile Attention, see - https://huggingface.co/papers/2502.04507) prior.""" - B, h, S, D = q.shape - s1 = S // block_size - qa = q.reshape(B, h, s1, block_size, D).mean(-2) - ka = k.reshape(B, h, s1, block_size, D).mean(-2).transpose(-2, -1) - attn_map = torch.softmax((qa @ ka) / math.sqrt(D), dim=-1) - - vals, inds = attn_map.sort(-1) - mask = (vals.cumsum_(-1) >= 1 - thr).int().gather(-1, inds.argsort(-1)) - mask = torch.logical_or(mask, sta) - - kv_nb = mask.sum(-1).to(torch.int32) - kv_inds = mask.argsort(dim=-1, descending=True).to(torch.int32) - return BlockMask.from_kv_blocks( - torch.zeros_like(kv_nb), - kv_inds, - kv_nb, - kv_inds, - BLOCK_SIZE=block_size, - mask_mod=None, - ) - - class Kandinsky6RoPE1D(nn.Module): """1-D Rotary Position Embedding — used for text and audio sequences.""" @@ -191,11 +141,7 @@ def forward( class Kandinsky6AttnProcessor: - """Diffusers attention processor used by the TI2VA transformer. - - The NABLA sparse branch always dispatches on the `flex` backend, regardless of `_attention_backend`, since its - `BlockMask` only the `flex` backend can consume. - """ + """Diffusers attention processor used by the TI2VA transformer.""" _attention_backend = None _parallel_config = None @@ -211,7 +157,6 @@ def __call__( encoder_hidden_states: Tensor | None = None, rotary_emb: Tensor | None = None, rotary_emb_kv: Tensor | None = None, - sparse_params: dict[str, Any] | None = None, attn_mask: Tensor | None = None, ) -> Tensor: query = attn.to_query(hidden_states) @@ -233,32 +178,14 @@ def __call__( if rotary_emb_kv is not None: key = apply_rotary(key, rotary_emb_kv).to(dtype=key.dtype) - if sparse_params is not None: - # The block statistics are computed from the `(B, heads, S, D)` layout the mask builder expects; the - # resulting `BlockMask` only runs on the flex backend. - block_mask = nabla_block_mask( - query.transpose(1, 2), - key.transpose(1, 2), - sparse_params["sta_mask"], - thr=sparse_params["P"], - ) - output = dispatch_attention_fn( - query, - key, - value, - attn_mask=block_mask, - backend=AttentionBackendName.FLEX, - parallel_config=self._parallel_config, - ) - else: - output = dispatch_attention_fn( - query, - key, - value, - attn_mask=attn_mask, - backend=self._attention_backend, - parallel_config=self._parallel_config, - ) + output = dispatch_attention_fn( + query, + key, + value, + attn_mask=attn_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) return attn.out_layer(output.flatten(-2, -1)) @@ -313,14 +240,12 @@ def forward(self, x: Tensor) -> Tensor: class Kandinsky6Modulation(nn.Module): - """Zero-initialized AdaLN modulation projection.""" + """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) - nn.init.zeros_(self.out_layer.weight) - nn.init.zeros_(self.out_layer.bias) def forward(self, x: Tensor) -> Tensor: return self.out_layer(self.activation(x.to(get_parameter_dtype(self.out_layer)))) @@ -358,13 +283,11 @@ def forward( encoder_hidden_states: Tensor | None = None, attn_mask: Tensor | None = None, rotary_emb: Tensor | None = None, - sparse_params: dict[str, Any] | None = None, rope_q: Tensor | None = None, rope_kv: Tensor | None = None, ) -> Tensor: - # Native K6 blocks pass ``(hidden, rope, mask_or_sparse)`` for - # self-attention. Keep that call shape while exposing Diffusers' - # encoder_hidden_states/rotary_emb keyword boundary. + # Native K6 blocks pass ``(hidden, rope, mask)`` for self-attention. Keep that call shape while exposing + # Diffusers' encoder_hidden_states/rotary_emb keyword boundary. rotary_emb = rope_q if rope_q is not None else rotary_emb rotary_emb_kv = rope_kv if rope_kv is not None else rotary_emb return self.processor( @@ -373,7 +296,6 @@ def forward( encoder_hidden_states=encoder_hidden_states, rotary_emb=rotary_emb, rotary_emb_kv=rotary_emb_kv, - sparse_params=sparse_params, attn_mask=attn_mask, ) @@ -549,7 +471,6 @@ def forward( time_embed: tuple[Tensor, Tensor], vis_rope: Tensor | None, aud_rope: Tensor | None, - sparse_params: dict | None, attn_mask=None, ) -> tuple[Tensor | None, Tensor | None]: t_v, t_a = time_embed @@ -561,7 +482,6 @@ def forward( self.video_dec_block.self_attention( apply_scale_shift(self.video_dec_block.self_attention_norm(vis.float()), vis, scale, shift), rotary_emb=vis_rope, - sparse_params=sparse_params, ), gate, ).type_as(vis) @@ -809,7 +729,6 @@ def forward( visual_rope_pos: tuple[Tensor, Tensor, Tensor] | None = None, encoder_attention_mask: Tensor | None = None, visual_token_type_ids: Tensor | None = None, - sparse_params: dict | None = None, return_dict: bool = True, ) -> AudioVisualModelOutput | tuple[Tensor, ...]: r""" @@ -834,13 +753,6 @@ def forward( visual_token_type_ids (`torch.Tensor` of shape `(batch_size, num_frames)`, *optional*): Per-frame token type ids, embedded through `visual_token_type_embeddings`. Requires `visual_token_type_num_embeddings > 0`. - sparse_params (`dict`, *optional*): - NABLA sparse-attention configuration for the video self-attention, dispatched through the `flex` - attention backend regardless of `transformer.set_attention_backend(...)`. This trades the regular SDPA - backend's dense attention for a block-sparse pattern that scales better to longer sequences, at the - cost of FlexAttention's one-time kernel-autotuning overhead on the first call. None of the currently - released checkpoints were trained with NABLA; this is forward-looking support for a future, - longer-duration (e.g. 10s) model. return_dict (`bool`, defaults to `True`): Whether to return an [`AudioVisualModelOutput`] instead of a plain tuple. @@ -882,13 +794,8 @@ def forward( if visual_rope_pos is None: visual_rope_pos = tuple(torch.arange(size, device=device) for size in visual_shape) visual_rope = self.visual_rope_embeddings(visual_rope_pos, self.scale_factor) - to_fractal = bool(sparse_params and sparse_params.get("to_fractal")) - if to_fractal: - visual_embed = _local_patch(visual_embed, visual_shape, (1, 8, 8), dim=1).flatten(1, 2) - visual_rope = _local_patch(visual_rope, visual_shape, (1, 8, 8), dim=0).flatten(0, 1) - else: - visual_embed = visual_embed.flatten(1, 3) - visual_rope = visual_rope.flatten(0, 2) + visual_embed = visual_embed.flatten(1, 3) + visual_rope = visual_rope.flatten(0, 2) # 4. Embed the audio latents and build their rotary embeddings audio_embed = audio_rope = None @@ -906,21 +813,12 @@ def forward( (video_temb, audio_temb), visual_rope, audio_rope, - sparse_params, encoder_attention_mask, ) visual_embed, audio_embed = self._gradient_checkpointing_func(block, *args) if checkpoint else block(*args) # 6. Project back to the latent space - if to_fractal: - visual_embed = _local_merge( - visual_embed.reshape(visual_embed.shape[0], -1, 64, visual_embed.shape[-1]), - visual_shape, - (1, 8, 8), - dim=1, - ) - else: - visual_embed = visual_embed.reshape(-1, *visual_shape, visual_embed.shape[-1]) + visual_embed = visual_embed.reshape(-1, *visual_shape, visual_embed.shape[-1]) video_out = self.out_layer(visual_embed, video_temb) audio_out = self.audio_out_layer(audio_embed, audio_temb) if audio_embed is not None else None diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index d796c954a18f..aafefe5988bc 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -26,7 +26,7 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward -from ..attention_dispatch import dispatch_attention_fn +from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn from ..embeddings import TimestepEmbedding, Timesteps from ..modeling_outputs import Transformer2DModelOutput from ..modeling_utils import ModelMixin, get_parameter_dtype @@ -79,7 +79,6 @@ def apply_rotary(x: Tensor, rope: Tensor) -> Tensor: return (rope.float() * x_).sum(dim=-1).reshape(*x.shape).to(dtype=x.dtype) -# Copied from diffusers.models.transformers.transformer_kandinsky6._local_patch def _local_patch(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Tensor: T, H, W = shape g1, g2, g3 = group_size @@ -89,7 +88,6 @@ def _local_patch(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Te return x.flatten(dim, dim + 2).flatten(dim + 1, dim + 3) -# Copied from diffusers.models.transformers.transformer_kandinsky6._local_merge def _local_merge(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Tensor: T, H, W = shape g1, g2, g3 = group_size @@ -99,7 +97,6 @@ def _local_merge(x: Tensor, shape: tuple, group_size: tuple, dim: int = 0) -> Te return x.flatten(dim, dim + 1).flatten(dim + 1, dim + 2).flatten(dim + 2, dim + 3) -# Copied from diffusers.models.transformers.transformer_kandinsky6.nabla_block_mask def nabla_block_mask( q: Tensor, k: Tensor, @@ -153,15 +150,12 @@ def sliding_tile_mask( class Kandinsky6SRAttnProcessor: """Self-attention processor of the SR transformer: dense by default, NABLA sparse when `sparse_params` is set. - Like [`Kandinsky5AttnProcessor`], the attention backend is chosen through the usual `_attention_backend` - (`set_attention_backend`), even in the NABLA sparse branch. Its sparsity is expressed as a `BlockMask`, which only - the `flex` backend can consume, so sparse checkpoints require `set_attention_backend("flex")`. That backend also - needs to run under `torch.compile` (e.g. `transformer.compile_repeated_blocks()`): uncompiled, PyTorch's flex - attention falls back to an eager implementation that materializes the full attention matrix, which does not fit in - memory at video resolutions. + Always dispatches on the `flex` backend: NABLA's sparsity is expressed as a `BlockMask`, which only `flex` can + consume. That backend also needs to run under `torch.compile` (e.g. `transformer.compile_repeated_blocks()`): + uncompiled, PyTorch's flex attention falls back to an eager implementation that materializes the full attention + matrix, which does not fit in memory at video resolutions. """ - _attention_backend = None _parallel_config = None def __call__( @@ -182,9 +176,7 @@ def __call__( if sparse_params is None: attn_mask = None else: - # The block statistics are computed from the `(B, heads, S, D)` layout the mask builder expects; the - # resulting `BlockMask` only runs on the flex backend, so sparse checkpoints need - # `set_attention_backend("flex")` (see the class docstring). + # The block statistics are computed from the `(B, heads, S, D)` layout the mask builder expects. attn_mask = nabla_block_mask( query.transpose(1, 2), key.transpose(1, 2), @@ -196,7 +188,7 @@ def __call__( key, value, attn_mask=attn_mask, - backend=self._attention_backend, + backend=AttentionBackendName.FLEX, parallel_config=self._parallel_config, ) return attn.out_layer(hidden_states.flatten(2, 3)) @@ -271,14 +263,12 @@ def forward(self, x: Tensor) -> Tensor: # Copied from diffusers.models.transformers.transformer_kandinsky6.Kandinsky6Modulation with Kandinsky6->Kandinsky6SR class Kandinsky6SRModulation(nn.Module): - """Zero-initialized AdaLN modulation projection.""" + """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) - nn.init.zeros_(self.out_layer.weight) - nn.init.zeros_(self.out_layer.bias) def forward(self, x: Tensor) -> Tensor: return self.out_layer(self.activation(x.to(get_parameter_dtype(self.out_layer)))) diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py index 8c8aad02cc3f..5da4578e1465 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py @@ -52,9 +52,8 @@ >>> sr_pipe = Kandinsky6SRPipeline.from_pretrained( ... "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", torch_dtype=torch.bfloat16 ... ) - >>> # The transformer's default config runs NABLA sparse attention, which only the `flex` backend can - >>> # execute; compiling avoids the eager fallback's much higher memory use at video resolutions. - >>> sr_pipe.transformer.set_attention_backend("flex") + >>> # The transformer always runs attention through the `flex` backend; compiling avoids the eager + >>> # fallback's much higher memory use at video resolutions. >>> sr_pipe.transformer.compile_repeated_blocks(fullgraph=True) >>> sr_pipe.enable_model_cpu_offload() >>> output = sr_pipe(video=video, resolution_scale=2.25, num_inference_steps=2) diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index 17ff490e9b22..3c26416407d1 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -259,8 +259,15 @@ def __init__( ) self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + @staticmethod def _get_prompt_embeds( - self, prompt: list[str], max_sequence_length: int, device: torch.device + prompt: list[str], + tokenizer, + text_encoder, + tokenizer_2, + text_encoder_2, + max_sequence_length: int, + device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: """Encode prompts with Qwen2.5-VL (token embeddings) and CLIP (pooled embedding). @@ -268,7 +275,7 @@ def _get_prompt_embeds( in the batch is padded (a mask without padding carries no information, and dropping it keeps every attention backend available). """ - inputs = self.tokenizer( + inputs = tokenizer( text=[_PROMPT_TEMPLATE.format(item) for item in prompt], images=None, videos=None, @@ -277,7 +284,7 @@ def _get_prompt_embeds( return_tensors="pt", padding="max_length", ).to(device) - qwen_output = self.text_encoder( + qwen_output = text_encoder( input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], return_dict=True, @@ -288,7 +295,7 @@ def _get_prompt_embeds( if prompt_attention_mask.all(): prompt_attention_mask = None - clip_inputs = self.tokenizer_2( + clip_inputs = tokenizer_2( prompt, max_length=_CLIP_MAX_LENGTH, truncation=True, @@ -296,7 +303,7 @@ def _get_prompt_embeds( padding="max_length", return_tensors="pt", ).to(device) - pooled_prompt_embeds = self.text_encoder_2(**clip_inputs)["pooler_output"] + pooled_prompt_embeds = text_encoder_2(**clip_inputs)["pooler_output"] return prompt_embeds, pooled_prompt_embeds, prompt_attention_mask def encode_prompt( @@ -351,14 +358,28 @@ def encode_prompt( if prompt_embeds is None: prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = self._get_prompt_embeds( - prompt, max_sequence_length, device + prompt, + self.tokenizer, + self.text_encoder, + self.tokenizer_2, + self.text_encoder_2, + max_sequence_length, + device, ) if do_classifier_free_guidance and negative_prompt_embeds is None: negative_prompt = negative_prompt or self._DEFAULT_NEGATIVE_PROMPT negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt negative_prompt = negative_prompt * len(prompt) if len(negative_prompt) == 1 else negative_prompt negative_prompt_embeds, negative_pooled_prompt_embeds, negative_prompt_attention_mask = ( - self._get_prompt_embeds(negative_prompt, max_sequence_length, device) + self._get_prompt_embeds( + negative_prompt, + self.tokenizer, + self.text_encoder, + self.tokenizer_2, + self.text_encoder_2, + max_sequence_length, + device, + ) ) prompt_embeds = prompt_embeds.to(device=device, dtype=dtype).repeat_interleave(num_videos_per_prompt, dim=0) @@ -388,20 +409,30 @@ def encode_prompt( negative_prompt_attention_mask, ) + @staticmethod def expand_prompts( - self, prompt: str | list[str], + tokenizer, + text_encoder, + device: torch.device, image: PIL.Image.Image | list[PIL.Image.Image] | None = None, max_sequence_length: int = 1024, generator: torch.Generator | list[torch.Generator] | None = None, ) -> str | list[str]: r""" Rewrites short prompts into detailed video+audio prompts with the Qwen2.5-VL text encoder, grounding them on - the reference image when one is given. + the reference image when one is given. A `staticmethod` so it can be used standalone, before running the + pipeline. Args: prompt (`str` or `list[str]`): Prompt or prompts to expand. + tokenizer: + The Qwen2.5-VL processor, e.g. `pipe.tokenizer`. + text_encoder: + The Qwen2.5-VL model, e.g. `pipe.text_encoder`. + device (`torch.device`): + Device to run the text encoder on. image (`PIL.Image.Image` or `list[PIL.Image.Image]`, *optional*): Reference image(s) of an image-to-video call. max_sequence_length (`int`, defaults to `1024`): @@ -418,8 +449,14 @@ def expand_prompts( images = image if isinstance(image, list) else [image] * len(prompt) generators = generator if isinstance(generator, list) else [generator] * len(prompt) return [ - self.expand_prompts( - item, image=item_image, max_sequence_length=max_sequence_length, generator=item_generator + Kandinsky6TI2VAPipeline.expand_prompts( + item, + tokenizer, + text_encoder, + device, + image=item_image, + max_sequence_length=max_sequence_length, + generator=item_generator, ) for item, item_image, item_generator in zip(prompt, images, generators, strict=True) ] @@ -431,21 +468,21 @@ def expand_prompts( ) content = [{"type": "image", "image": image}] if image is not None else [] content.append({"type": "text", "text": instruction}) - text = self.tokenizer.apply_chat_template( + text = tokenizer.apply_chat_template( [{"role": "user", "content": content}], tokenize=False, add_generation_prompt=True ) - inputs = self.tokenizer( + inputs = tokenizer( text=[text], images=[image] if image is not None else None, videos=None, padding=True, return_tensors="pt", - ).to(self._execution_device) + ).to(device) if generator is not None: torch.manual_seed(generator.initial_seed()) - generated = self.text_encoder.generate(**inputs, max_new_tokens=max_sequence_length) + generated = text_encoder.generate(**inputs, max_new_tokens=max_sequence_length) generated = generated[:, inputs["input_ids"].shape[1] :] - return self.tokenizer.batch_decode(generated, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + return tokenizer.batch_decode(generated, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] def check_inputs( self, @@ -456,6 +493,7 @@ def check_inputs( num_frames, image, sample_audio, + expand_prompts, prompt_embeds=None, pooled_prompt_embeds=None, negative_prompt_embeds=None, @@ -488,6 +526,8 @@ def check_inputs( ) elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + if expand_prompts and prompt_embeds is not None: + raise ValueError("`expand_prompts=True` requires `prompt`; it cannot be used with `prompt_embeds`.") if negative_prompt is not None and negative_prompt_embeds is not None: raise ValueError( f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" @@ -624,7 +664,7 @@ def __call__( pooled_prompt_embeds: torch.Tensor | None = None, negative_prompt_embeds: torch.Tensor | None = None, negative_pooled_prompt_embeds: torch.Tensor | None = None, - sample_audio: bool | None = None, + sample_audio: bool = True, expand_prompts: bool = False, max_sequence_length: int = 1024, output_type: str = "pil", @@ -675,9 +715,8 @@ def __call__( Pre-generated negative Qwen2.5-VL text embeddings. negative_pooled_prompt_embeds (`torch.Tensor`, *optional*): Pre-generated negative CLIP pooled text embeddings. - sample_audio (`bool`, *optional*): - Whether to generate synchronized audio. Defaults to `True` when the pipeline has an `audio_vae` and a - `vocoder`, `False` otherwise; `True` requires both. + sample_audio (`bool`, defaults to `True`): + Whether to generate synchronized audio. Requires the pipeline to have an `audio_vae` and a `vocoder`. expand_prompts (`bool`, defaults to `False`): Whether to rewrite the prompts with [`~Kandinsky6TI2VAPipeline.expand_prompts`] before encoding. max_sequence_length (`int`, defaults to `1024`): @@ -699,8 +738,6 @@ def __call__( The generated video and audio; a `(frames, audio)` tuple when `return_dict=False`. """ # 1. Check inputs. Raise error if not correct - if sample_audio is None: - sample_audio = getattr(self, "audio_vae", None) is not None and getattr(self, "vocoder", None) is not None self.check_inputs( prompt=prompt, negative_prompt=negative_prompt, @@ -709,6 +746,7 @@ def __call__( num_frames=num_frames, image=image, sample_audio=sample_audio, + expand_prompts=expand_prompts, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, @@ -734,9 +772,15 @@ def __call__( dtype = self.transformer.dtype # 3. Encode input prompt - if expand_prompts and prompt is not None: + if expand_prompts: prompt = self.expand_prompts( - prompt, image=image, max_sequence_length=max_sequence_length, generator=generator + prompt, + self.tokenizer, + self.text_encoder, + device, + image=image, + max_sequence_length=max_sequence_length, + generator=generator, ) ( prompt_embeds, diff --git a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py index ac68deeddbf1..235d32f6b5c4 100644 --- a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py +++ b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_ti2va.py @@ -14,6 +14,7 @@ import numpy as np import PIL.Image +import pytest import torch from transformers import ( AutoProcessor, @@ -33,7 +34,7 @@ MMAudioVocoder, ) -from ...testing_utils import torch_device +from ...testing_utils import assert_tensors_close, torch_device from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin @@ -208,6 +209,44 @@ def get_dummy_inputs(self): class TestKandinsky6TI2VAPipeline(Kandinsky6TI2VAPipelineTesterConfig, PipelineTesterMixin): + def test_save_load_optional_components(self, tmp_path, expected_max_difference=1e-4): + # Dropping the optional `audio_vae`/`vocoder` components also drops the pipeline's ability to sample + # audio, so `sample_audio` must be turned off explicitly instead of relying on its default. + if not getattr(self.pipeline_class, "_optional_components", None): + pytest.skip(f"Skipping test because {self.pipeline_class} has no `_optional_components`.") + + pipe = self.get_pipeline().to(torch_device) + + for optional_component in pipe._optional_components: + setattr(pipe, optional_component, None) + + inputs = self.get_dummy_inputs() + inputs["sample_audio"] = False + torch.manual_seed(0) + output = pipe(**inputs)[0] + + pipe.save_pretrained(tmp_path, safe_serialization=False) + pipe_loaded = self.pipeline_class.from_pretrained(tmp_path) + pipe_loaded.to(torch_device) + pipe_loaded.set_progress_bar_config(disable=None) + + for optional_component in pipe._optional_components: + assert getattr(pipe_loaded, optional_component) is None, ( + f"`{optional_component}` did not stay set to None after loading." + ) + + inputs = self.get_dummy_inputs() + inputs["sample_audio"] = False + torch.manual_seed(0) + output_loaded = pipe_loaded(**inputs)[0] + + assert_tensors_close( + output_loaded, + output, + atol=expected_max_difference, + msg="Output changed after dropping optional components.", + ) + def test_kandinsky6_ti2va_audio_output(self): pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) output = pipe(**self.get_dummy_inputs()) From 4c0ad1243e2b744891473397b5822547b9415f82 Mon Sep 17 00:00:00 2001 From: leffff Date: Sun, 4 Oct 2026 15:14:18 +0300 Subject: [PATCH 46/62] fix audio channels bug --- src/diffusers/utils/export_utils.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/diffusers/utils/export_utils.py b/src/diffusers/utils/export_utils.py index f06f172e79b2..841b34d6c116 100644 --- a/src/diffusers/utils/export_utils.py +++ b/src/diffusers/utils/export_utils.py @@ -277,23 +277,25 @@ def _write_audio( import torch if samples.ndim == 1: - samples = samples[:, None] + samples = samples[None, :] - if samples.shape[1] != 2 and samples.shape[0] == 2: + # `samples` is documented as [audio_channels, samples]; accept the transposed form too. + if samples.shape[0] not in (1, 2) and samples.shape[1] in (1, 2): samples = samples.T - if samples.shape[1] != 2: - raise ValueError(f"Expected samples with 2 channels; got shape {samples.shape}.") + if samples.shape[0] not in (1, 2): + raise ValueError(f"Expected samples with 1 or 2 channels; got shape {samples.shape}.") # Convert to int16 packed for ingestion; resampler converts to the encoder format. if samples.dtype != torch.int16: samples = torch.clip(samples, -1.0, 1.0) samples = (samples * 32767.0).to(torch.int16) + num_channels = samples.shape[0] frame_in = av_module.AudioFrame.from_ndarray( - samples.contiguous().reshape(1, -1).cpu().numpy(), + samples.T.contiguous().reshape(1, -1).cpu().numpy(), format="s16", - layout="stereo", + layout="mono" if num_channels == 1 else "stereo", ) frame_in.sample_rate = audio_sample_rate From 1f9b471826d72d4e4e66b4093eb9713e82674f7c Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 5 Oct 2026 12:48:00 +0300 Subject: [PATCH 47/62] fix tests --- .../api/models/anyflow_far_transformer3d.md | 2 +- .../en/api/models/anyflow_transformer3d.md | 2 +- .../en/api/models/kandinsky6_transformer.md | 4 +- docs/source/en/api/pipelines/kandinsky6.md | 4 +- docs/source/en/api/pipelines/pag.md | 2 +- .../en/modular_diffusers/modular_pipeline.md | 2 +- .../en/optimization/attention_backends.md | 2 +- docs/source/en/optimization/memory.md | 2 +- .../en/training/distributed_inference.md | 2 +- docs/source/en/using-diffusers/callback.md | 2 +- docs/source/en/using-diffusers/cli.md | 4 +- .../autoencoders/autoencoder_kandinsky6_sr.py | 17 +++ .../autoencoders/autoencoder_mmaudio.py | 18 +++ .../transformers/transformer_kandinsky6_sr.py | 110 ++++++++++-------- .../models/transformers/transformer_ltx2.py | 2 +- .../test_models_transformer_kandinsky6_sr.py | 13 +-- .../kandinsky6/test_pipeline_kandinsky6_sr.py | 8 +- 17 files changed, 122 insertions(+), 74 deletions(-) diff --git a/docs/source/en/api/models/anyflow_far_transformer3d.md b/docs/source/en/api/models/anyflow_far_transformer3d.md index 7f818c44ef20..b10bc866c16b 100644 --- a/docs/source/en/api/models/anyflow_far_transformer3d.md +++ b/docs/source/en/api/models/anyflow_far_transformer3d.md @@ -12,7 +12,7 @@ specific language governing permissions and limitations under the License. # AnyFlowFARTransformer3DModel -The causal (FAR) 3D Transformer used by [`AnyFlowFARPipeline`](../pipelines/anyflow#anyflowfarpipeline) — +The causal (FAR) 3D Transformer used by [`AnyFlowFARPipeline`](../pipelines/anyflow) — the FAR variant of [AnyFlow](https://huggingface.co/papers/2605.13724). See the [`AnyFlowFARPipeline`](../pipelines/anyflow) page for paper, authors, and released checkpoints. It extends the v0.35.1 Wan2.1 backbone with three additions: diff --git a/docs/source/en/api/models/anyflow_transformer3d.md b/docs/source/en/api/models/anyflow_transformer3d.md index d37f7fba62fb..a62ccf1ab6b5 100644 --- a/docs/source/en/api/models/anyflow_transformer3d.md +++ b/docs/source/en/api/models/anyflow_transformer3d.md @@ -12,7 +12,7 @@ specific language governing permissions and limitations under the License. # AnyFlowTransformer3DModel -The bidirectional 3D Transformer used by [`AnyFlowPipeline`](../pipelines/anyflow#anyflowpipeline). It is the +The bidirectional 3D Transformer used by [`AnyFlowPipeline`](../pipelines/anyflow). It is the v0.35.1 Wan2.1 backbone with one structural change: the timestep embedder is replaced by ``AnyFlowDualTimestepTextImageEmbedding``, so every forward call conditions on both the source timestep ``t`` and the target timestep ``r``. This is the embedding required to learn the flow map diff --git a/docs/source/en/api/models/kandinsky6_transformer.md b/docs/source/en/api/models/kandinsky6_transformer.md index 984aa7e3f3d2..e48aac25884d 100644 --- a/docs/source/en/api/models/kandinsky6_transformer.md +++ b/docs/source/en/api/models/kandinsky6_transformer.md @@ -43,8 +43,8 @@ from diffusers import Kandinsky6SRTransformer3DModel transformer = Kandinsky6SRTransformer3DModel.from_pretrained( "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16 ) -# The default config (`nabla_threshold=0.8`) runs NABLA sparse attention, like Kandinsky 5's `nabla` -# checkpoints — only the `flex` backend can execute it. +# The transformer always runs NABLA sparse attention (`nabla_threshold`, 0.8 by default), like Kandinsky 5's `nabla` +# checkpoints. Only the `flex` backend can execute it. transformer.set_attention_backend("flex") ``` diff --git a/docs/source/en/api/pipelines/kandinsky6.md b/docs/source/en/api/pipelines/kandinsky6.md index bef76ac4f8bf..7f9ba2e16e73 100644 --- a/docs/source/en/api/pipelines/kandinsky6.md +++ b/docs/source/en/api/pipelines/kandinsky6.md @@ -76,8 +76,8 @@ windows. sr_pipe = Kandinsky6SRPipeline.from_pretrained( "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", torch_dtype=torch.bfloat16 ) -# The SR transformer's default config runs NABLA sparse attention (like Kandinsky 5's `nabla` -# checkpoints), which only the `flex` backend can execute. +# The SR transformer always runs NABLA sparse attention (like Kandinsky 5's `nabla` checkpoints), which only the +# `flex` backend can execute. sr_pipe.transformer.set_attention_backend("flex") sr_pipe.transformer.compile_repeated_blocks(fullgraph=True) sr_pipe.enable_model_cpu_offload() diff --git a/docs/source/en/api/pipelines/pag.md b/docs/source/en/api/pipelines/pag.md index 055feee8cebe..4b2bf2b24c0c 100644 --- a/docs/source/en/api/pipelines/pag.md +++ b/docs/source/en/api/pipelines/pag.md @@ -333,7 +333,7 @@ PAG reduces artifacts and improves the overall compposition. ## Configure parameters -### pag_applied_layers +### pag_applied_layers[[pag_applied_layers]] The `pag_applied_layers` argument allows you to specify which layers PAG is applied to. By default, it applies only to the mid blocks. Changing this setting will significantly impact the output. You can use the `set_pag_applied_layers` method to adjust the PAG layers after the pipeline is created, helping you find the optimal layers for your model. diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index 68177c7e7fb8..1a5ebda47979 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/docs/source/en/modular_diffusers/modular_pipeline.md @@ -195,7 +195,7 @@ ModularPipeline { If a component in your block doesn't exist in the repository, it remains `null` and is skipped during [`~ModularPipeline.load_components`]. -### from_pretrained +### from_pretrained[[from_pretrained]] [`~ModularPipeline.from_pretrained`] is a convenient way to create a [`ModularPipeline`] without defining blocks yourself. diff --git a/docs/source/en/optimization/attention_backends.md b/docs/source/en/optimization/attention_backends.md index 08df0e75a952..651fefce8677 100644 --- a/docs/source/en/optimization/attention_backends.md +++ b/docs/source/en/optimization/attention_backends.md @@ -28,7 +28,7 @@ Refer to the table below for an overview of the available attention families and This guide will show you how to set and use the different attention backends. -## set_attention_backend +## set_attention_backend[[set_attention_backend]] The [`~ModelMixin.set_attention_backend`] method iterates through all the modules in the model and sets the appropriate attention backend to use. The attention backend setting persists until [`~ModelMixin.reset_attention_backend`] is called. diff --git a/docs/source/en/optimization/memory.md b/docs/source/en/optimization/memory.md index 950d05b81e69..21f3d980034e 100644 --- a/docs/source/en/optimization/memory.md +++ b/docs/source/en/optimization/memory.md @@ -468,7 +468,7 @@ apply_layerwise_casting( ) ``` -## torch.channels_last +## torch.channels_last[[torchchannels_last]] [torch.channels_last](https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html) flips how tensors are stored from `(batch size, channels, height, width)` to `(batch size, height, width, channels)`. This aligns the tensors with how the hardware sequentially accesses the tensors stored in memory and avoids skipping around in memory to access the pixel values. diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md index 856572c2ff08..3b764621d7c1 100644 --- a/docs/source/en/training/distributed_inference.md +++ b/docs/source/en/training/distributed_inference.md @@ -114,7 +114,7 @@ Call `torchrun` to run the inference script and use the `--nproc_per_node` argum torchrun --nproc_per_node=2 run_distributed.py ``` -## device_map +## device_map[[device_map]] The `device_map` argument enables distributed inference by automatically placing model components on separate GPUs. This is especially useful when a model doesn't fit on a single GPU. You can use `device_map` to selectively load and unload the required model components at a given stage as shown in the example below (assumes two GPUs are available). diff --git a/docs/source/en/using-diffusers/callback.md b/docs/source/en/using-diffusers/callback.md index 3847b376cc49..443bdc86fd78 100644 --- a/docs/source/en/using-diffusers/callback.md +++ b/docs/source/en/using-diffusers/callback.md @@ -14,7 +14,7 @@ specific language governing permissions and limitations under the License. A callback is a function that modifies [`DiffusionPipeline`] behavior and it is executed at the end of a denoising step. The changes are propagated to subsequent steps in the denoising process. It is useful for adjusting pipeline attributes or tensor variables to support new features without rewriting the underlying pipeline code. -Diffusers provides several callbacks in the pipeline [overview](../api/pipelines/overview#callbacks). +Diffusers provides several callbacks in the pipeline [overview](../api/pipelines/overview). To enable a callback, configure when the callback is executed after a certain number of denoising steps with one of the following arguments. diff --git a/docs/source/en/using-diffusers/cli.md b/docs/source/en/using-diffusers/cli.md index c16653c9ee7c..060555cc2cb9 100644 --- a/docs/source/en/using-diffusers/cli.md +++ b/docs/source/en/using-diffusers/cli.md @@ -287,7 +287,7 @@ diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 \ hf sandbox kill ``` -## `custom_blocks` +## `custom_blocks`[[custom_blocks]] Package a local `ModularPipelineBlocks` subclass for upload to the Hub. Reads a Python file, AST-scans it for subclasses of `ModularPipelineBlocks`, instantiates the chosen one, and calls `save_pretrained` in the current @@ -304,7 +304,7 @@ diffusers-cli custom_blocks --block_module_name my_block.py --block_class_name M The block class must be instantiable with zero constructor args — hardcode defaults in `__init__` or read config from the pipeline `state` at call time. -## `fp16_safetensors` +## `fp16_safetensors`[[fp16_safetensors]] > [!CAUTION] > This command is now deprecated and will be removed in a future version. diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py index c92d8833db58..25ed70941f82 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py +++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py @@ -646,6 +646,23 @@ def forward( return_dict: bool = True, generator: torch.Generator | None = None, ) -> DecoderOutput | tuple: + r""" + Args: + sample (`torch.Tensor` of shape `(batch_size, channels, num_frames, height, width)`): + Pixel video in `[-1, 1]`. `num_frames` should be `1 + k * temporal_compression_ratio`. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the latent posterior instead of using its mode. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.DecoderOutput`] instead of a plain tuple. + generator (`torch.Generator`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling + deterministic. + + Returns: + [`~models.autoencoder_kl.DecoderOutput`] or `tuple`: + If `return_dict` is True, a [`~models.autoencoder_kl.DecoderOutput`] is returned, otherwise a plain + `tuple` is returned. Its `sample` is the reconstructed video. + """ posterior = self.encode(sample).latent_dist z = posterior.sample(generator=generator) if sample_posterior else posterior.mode() decoded = self.decode(z).sample diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py index e9d73417cd3f..c19ad4fd3619 100644 --- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py +++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py @@ -443,6 +443,24 @@ def forward( return_dict: bool = True, generator: torch.Generator | None = None, ) -> DecoderOutput | tuple: + r""" + Args: + sample (`torch.Tensor` of shape `(batch_size, num_samples)`): + Mono waveform in `[-1, 1]` at `sample_rate` to encode and reconstruct as a mel spectrogram. + sample_posterior (`bool`, *optional*, defaults to `False`): + Whether to sample from the latent posterior instead of using its mode. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~models.autoencoder_kl.DecoderOutput`] instead of a plain tuple. + generator (`torch.Generator`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling + deterministic. + + Returns: + [`~models.autoencoder_kl.DecoderOutput`] or `tuple`: + If `return_dict` is True, a [`~models.autoencoder_kl.DecoderOutput`] is returned, otherwise a plain + `tuple` is returned. Its `sample` is the reconstructed mel spectrogram of shape `(batch_size, mel_bins, + num_mel_frames)`. + """ posterior = self.encode(sample).latent_dist z = posterior.sample(generator=generator) if sample_posterior else posterior.mode() mel = self.decode(z).sample diff --git a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py index aafefe5988bc..b11bba2af7a3 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6_sr.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6_sr.py @@ -16,22 +16,30 @@ from __future__ import annotations +import functools import math -from typing import Any +from typing import TYPE_CHECKING, Any import torch from torch import Tensor, nn -from torch.nn.attention.flex_attention import BlockMask from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin +from ...utils import logging from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward -from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn +from ..attention_dispatch import _CAN_USE_FLEX_ATTN, AttentionBackendName, dispatch_attention_fn from ..embeddings import TimestepEmbedding, Timesteps from ..modeling_outputs import Transformer2DModelOutput from ..modeling_utils import ModelMixin, get_parameter_dtype +if TYPE_CHECKING: + from torch.nn.attention.flex_attention import BlockMask + + +logger = logging.get_logger(__name__) + + # Side of the local token block that NABLA sparse attention groups into one 64-token attention block. FRACTAL_BLOCK_SIZE = 8 @@ -106,6 +114,8 @@ def nabla_block_mask( ) -> BlockMask: """Build a dynamic NABLA block mask from query/key statistics and an STA (Sliding-Tile Attention, see https://huggingface.co/papers/2502.04507) prior.""" + from torch.nn.attention.flex_attention import BlockMask + B, h, S, D = q.shape s1 = S // block_size qa = q.reshape(B, h, s1, block_size, D).mean(-2) @@ -147,8 +157,17 @@ def sliding_tile_mask( return near.transpose(1, 2).reshape(num_frames * height * width, num_frames * height * width) +@functools.lru_cache(maxsize=None) +def _warn_uncompiled_flex_attention() -> None: + logger.warning( + "Kandinsky 6 SR attention runs PyTorch's flex attention eagerly, which materializes the full attention " + "matrix and does not fit in memory at video resolutions. Compile the transformer, e.g. with " + "`transformer.compile_repeated_blocks()`." + ) + + class Kandinsky6SRAttnProcessor: - """Self-attention processor of the SR transformer: dense by default, NABLA sparse when `sparse_params` is set. + """Self-attention processor of the SR transformer: NABLA sparse attention over the `sparse_params` block pattern. Always dispatches on the `flex` backend: NABLA's sparsity is expressed as a `BlockMask`, which only `flex` can consume. That backend also needs to run under `torch.compile` (e.g. `transformer.compile_repeated_blocks()`): @@ -163,7 +182,7 @@ def __call__( attn: "Kandinsky6SRAttention", hidden_states: Tensor, rotary_emb: Tensor, - sparse_params: dict[str, Any] | None = None, + sparse_params: dict[str, Any], ) -> Tensor: query = attn.to_query(hidden_states).unflatten(-1, (attn.num_heads, -1)) key = attn.to_key(hidden_states).unflatten(-1, (attn.num_heads, -1)) @@ -173,16 +192,15 @@ def __call__( query = apply_rotary(query, rotary_emb) key = apply_rotary(key, rotary_emb) - if sparse_params is None: - attn_mask = None - else: - # The block statistics are computed from the `(B, heads, S, D)` layout the mask builder expects. - attn_mask = nabla_block_mask( - query.transpose(1, 2), - key.transpose(1, 2), - sparse_params["sta_mask"], - thr=sparse_params["threshold"], - ) + if not torch.compiler.is_compiling(): + _warn_uncompiled_flex_attention() + # The block statistics are computed from the `(B, heads, S, D)` layout the mask builder expects. + attn_mask = nabla_block_mask( + query.transpose(1, 2), + key.transpose(1, 2), + sparse_params["sta_mask"], + thr=sparse_params["threshold"], + ) hidden_states = dispatch_attention_fn( query, key, @@ -217,7 +235,7 @@ def forward( self, hidden_states: Tensor, rotary_emb: Tensor, - sparse_params: dict[str, Any] | None = None, + sparse_params: dict[str, Any], ) -> Tensor: return self.processor(self, hidden_states, rotary_emb, sparse_params) @@ -368,7 +386,7 @@ def forward( hidden_states: Tensor, temb: Tensor, rotary_emb: Tensor, - sparse_params: dict[str, Any] | None = None, + sparse_params: dict[str, Any], ) -> Tensor: self_attention_params, feed_forward_params = torch.chunk(self.visual_modulation(temb), 2, dim=-1) shift, scale, gate = torch.chunk(self_attention_params, 3, dim=-1) @@ -398,8 +416,9 @@ class Kandinsky6SRTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, The model denoises tiles of the K-VAE latent video. Its input concatenates the noisy latent with the anchor latent and the anchor mask that condition the super-resolution (`2 * in_visual_dim + 1` channels), and its output holds `out_visual_dim` channels: `in_visual_dim` for a plain velocity checkpoint, or a widened `n_grid * in_visual_dim` - for the [`PiflowScheduler`] distilled checkpoints. Video self-attention runs dense, or through the NABLA sparse - block pattern when `nabla_threshold` is set. + for the [`PiflowScheduler`] distilled checkpoints. Video self-attention runs through the NABLA sparse block pattern + on the `flex` attention backend, which needs the token grid (`height` and `width` after `patch_size`) to be + divisible by 8. Args: in_visual_dim (`int`, defaults to `64`): @@ -420,8 +439,8 @@ class Kandinsky6SRTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, RoPE dimensions per `(t, h, w)` axis; their sum is the attention head dimension. scale_factor (`tuple[float, float, float]`, defaults to `(1.0, 2.0, 2.0)`): Per-axis RoPE frequency scaling applied to the token positions. - nabla_threshold (`float`, *optional*, defaults to `0.8`): - Cumulative-attention threshold of the NABLA block selection. `None` runs dense attention. + nabla_threshold (`float`, defaults to `0.8`): + Cumulative-attention threshold of the NABLA block selection. nabla_window (`tuple[int, int, int]`, defaults to `(11, 7, 7)`): Odd `(t, h, w)` extents of the sliding-tile prior that every 8x8 token block always attends to. tile_sizes (`tuple[tuple[int, int], ...]`, defaults to `((512, 512), (512, 768), (768, 512))`): @@ -449,11 +468,16 @@ def __init__( num_visual_blocks: int = 32, axes_dims: tuple[int, int, int] = (16, 24, 24), scale_factor: tuple[float, float, float] = (1.0, 2.0, 2.0), - nabla_threshold: float | None = 0.8, + nabla_threshold: float = 0.8, nabla_window: tuple[int, int, int] = (11, 7, 7), tile_sizes: tuple[tuple[int, int], ...] = ((512, 512), (512, 768), (768, 512)), ) -> None: super().__init__() + if not _CAN_USE_FLEX_ATTN: + raise ImportError( + "Kandinsky6SRTransformer3DModel requires PyTorch>=2.5.0 with `torch.nn.attention.flex_attention`" + " for its NABLA sparse attention." + ) head_dim = sum(axes_dims) self.time_embeddings = Kandinsky6SRTimeEmbeddings(model_dim, time_dim) @@ -501,23 +525,18 @@ def forward( visual_rope = self.visual_rope_embeddings(visual_rope_pos, self.config.scale_factor) # 2. Flatten the grid into a token sequence; NABLA groups 8x8 spatial neighbourhoods into attention blocks - sparse_params = None - if self.config.nabla_threshold is not None: - num_frames, height, width = visual_shape - if height % FRACTAL_BLOCK_SIZE or width % FRACTAL_BLOCK_SIZE: - raise ValueError( - f"NABLA attention needs the token grid {(height, width)} to be divisible by {FRACTAL_BLOCK_SIZE}" - ) - group_size = (1, FRACTAL_BLOCK_SIZE, FRACTAL_BLOCK_SIZE) - visual_embed = _local_patch(visual_embed, visual_shape, group_size, dim=1).flatten(1, 2) - visual_rope = _local_patch(visual_rope, visual_shape, group_size, dim=0).flatten(0, 1) - sta_mask = sliding_tile_mask( - num_frames, height // FRACTAL_BLOCK_SIZE, width // FRACTAL_BLOCK_SIZE, self.config.nabla_window, device + num_frames, height, width = visual_shape + if height % FRACTAL_BLOCK_SIZE or width % FRACTAL_BLOCK_SIZE: + raise ValueError( + f"NABLA attention needs the token grid {(height, width)} to be divisible by {FRACTAL_BLOCK_SIZE}" ) - sparse_params = {"sta_mask": sta_mask, "threshold": self.config.nabla_threshold} - else: - visual_embed = visual_embed.flatten(1, 3) - visual_rope = visual_rope.flatten(0, 2) + group_size = (1, FRACTAL_BLOCK_SIZE, FRACTAL_BLOCK_SIZE) + visual_embed = _local_patch(visual_embed, visual_shape, group_size, dim=1).flatten(1, 2) + visual_rope = _local_patch(visual_rope, visual_shape, group_size, dim=0).flatten(0, 1) + sta_mask = sliding_tile_mask( + num_frames, height // FRACTAL_BLOCK_SIZE, width // FRACTAL_BLOCK_SIZE, self.config.nabla_window, device + ) + sparse_params = {"sta_mask": sta_mask, "threshold": self.config.nabla_threshold} # 3. Transformer blocks for block in self.visual_transformer_blocks: @@ -527,15 +546,12 @@ def forward( visual_embed = block(visual_embed, temb, visual_rope, sparse_params) # 4. Restore the grid and project back to the latent space - if sparse_params is not None: - visual_embed = _local_merge( - visual_embed.reshape(visual_embed.shape[0], -1, FRACTAL_BLOCK_SIZE**2, visual_embed.shape[-1]), - visual_shape, - group_size, - dim=1, - ) - else: - visual_embed = visual_embed.reshape(-1, *visual_shape, visual_embed.shape[-1]) + visual_embed = _local_merge( + visual_embed.reshape(visual_embed.shape[0], -1, FRACTAL_BLOCK_SIZE**2, visual_embed.shape[-1]), + visual_shape, + group_size, + dim=1, + ) output = self.out_layer(visual_embed, temb) if not return_dict: diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py index 27ac7acb89e6..b180735119fd 100644 --- a/src/diffusers/models/transformers/transformer_ltx2.py +++ b/src/diffusers/models/transformers/transformer_ltx2.py @@ -98,7 +98,7 @@ class AudioVisualModelOutput(BaseOutput): """ sample: "torch.Tensor" # noqa: F821 - audio_sample: "torch.Tensor" # noqa: F821 + audio_sample: "torch.Tensor | None" = None # noqa: F821 class LTX2AdaLayerNormSingle(nn.Module): diff --git a/tests/models/transformers/test_models_transformer_kandinsky6_sr.py b/tests/models/transformers/test_models_transformer_kandinsky6_sr.py index 1cbe29a6a52f..27d2a56a4a7a 100644 --- a/tests/models/transformers/test_models_transformer_kandinsky6_sr.py +++ b/tests/models/transformers/test_models_transformer_kandinsky6_sr.py @@ -64,8 +64,6 @@ def get_init_dict(self) -> dict: "ff_dim": 32, "num_visual_blocks": 2, "axes_dims": (4, 4, 4), - # NABLA needs 8x8 token blocks, far larger than these tiny grids; run dense attention. - "nabla_threshold": None, } def _build_dummy_inputs(self, batch_size: int, num_frames: int, height: int, width: int) -> dict: @@ -81,15 +79,15 @@ def _build_dummy_inputs(self, batch_size: int, num_frames: int, height: int, wid } def get_dummy_inputs(self) -> dict: - return self._build_dummy_inputs(batch_size=1, num_frames=2, height=4, width=4) + return self._build_dummy_inputs(batch_size=1, num_frames=2, height=16, width=16) @property def input_shape(self) -> tuple[int, ...]: - return (2, 4, 4, 2 * 4 + 1) + return (2, 16, 16, 2 * 4 + 1) @property def output_shape(self) -> tuple[int, ...]: - return (2, 4, 4, 8) + return (2, 16, 16, 8) class TestKandinsky6SRTransformerModel(Kandinsky6SRTransformerTesterConfig, ModelTesterMixin): @@ -103,12 +101,13 @@ class TestKandinsky6SRTransformerMemory(Kandinsky6SRTransformerTesterConfig, Mem class TestKandinsky6SRTransformerTorchCompile(Kandinsky6SRTransformerTesterConfig, TorchCompileTesterMixin): @property def different_shapes_for_compilation(self): - return [(4, 4), (4, 8), (8, 8)] + return [(16, 16), (16, 32), (32, 32)] - def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict: + def get_dummy_inputs(self, height: int = 16, width: int = 16) -> dict: return self._build_dummy_inputs(batch_size=1, num_frames=2, height=height, width=width) +@pytest.mark.skipif(torch_device == "cpu", reason="FlexAttention does not support backward on CPU.") class TestKandinsky6SRTransformerTraining(Kandinsky6SRTransformerTesterConfig, TrainingTesterMixin): pass diff --git a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_sr.py b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_sr.py index 3a31c1e77c5b..14e1b7376524 100644 --- a/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_sr.py +++ b/tests/pipelines/kandinsky6/test_pipeline_kandinsky6_sr.py @@ -45,11 +45,9 @@ def get_dummy_components(self): ff_dim=32, num_visual_blocks=2, axes_dims=(4, 4, 4), - # NABLA needs 8x8 token blocks; the tiny tiles run dense attention. - nabla_threshold=None, - # Tiles the size of the 16x32 test video (and its transposed / square variants); every tile size must be - # divisible by the VAE spatial factor (4) times the largest scale (4). - tile_sizes=((16, 16), (16, 32), (32, 16)), + # Trained tile resolutions. NABLA needs each tile's latent token grid divisible by 8, so every size is a + # multiple of 32 (VAE spatial factor 4 times 8). The 2x route tiles the 16x32 test video at (16, 32). + tile_sizes=((32, 32), (32, 64), (64, 32)), ) torch.manual_seed(0) From 60cad43415e6032a23f6fe2eae2d008145095009 Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 5 Oct 2026 15:37:33 +0300 Subject: [PATCH 48/62] Fix CI failures on the Kandinsky6 PR - tests/others/test_utils.py: assert the caller file path without assuming the checkout directory is named `diffusers`. - .github/workflows/pr_tests.yml: raise PYTEST_TIMEOUT to 300 s for the example and pipeline CPU steps. Both run several workers with four threads each on an 8-vCPU runner, and the checkpointing example tests run several training subprocesses in sequence. Co-Authored-By: Claude Sonnet 5 --- .github/workflows/pr_tests.yml | 4 ++++ tests/others/test_utils.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/pr_tests.yml b/.github/workflows/pr_tests.yml index 66ce53d576e6..8df8dd0146b3 100644 --- a/.github/workflows/pr_tests.yml +++ b/.github/workflows/pr_tests.yml @@ -136,6 +136,8 @@ jobs: - name: Run fast PyTorch Pipeline CPU tests if: ${{ matrix.config.framework == 'pytorch_pipelines' }} + env: + PYTEST_TIMEOUT: 300 run: | pytest -n 8 --max-worker-restart=0 --dist=loadfile \ -k "not Onnx" \ @@ -152,6 +154,8 @@ jobs: - name: Run example PyTorch CPU tests if: ${{ matrix.config.framework == 'pytorch_examples' }} + env: + PYTEST_TIMEOUT: 300 run: | uv pip install ".[training]" pytest -n 4 --max-worker-restart=0 --dist=loadfile \ diff --git a/tests/others/test_utils.py b/tests/others/test_utils.py index d1a59cec52f1..3741b1ee7012 100755 --- a/tests/others/test_utils.py +++ b/tests/others/test_utils.py @@ -175,7 +175,7 @@ def test_deprecate_stacklevel(self): with pytest.warns(FutureWarning) as warning: deprecate(("deprecated_arg", self.higher_version, "This message is better!!!"), standard_warn=False) assert str(warning[0].message) == "This message is better!!!" - assert "diffusers/tests/others/test_utils.py" in warning[0].filename + assert "tests/others/test_utils.py" in warning[0].filename def test_deprecate_testing_utils_module(self): import diffusers.utils.testing_utils From 35e6cc4efdc2fe55fc8f864560c380d02a124fdc Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 5 Oct 2026 15:53:06 +0300 Subject: [PATCH 49/62] Avoid staging LFS dedup race in the org push test The second push in TestModelPushToHub.test_push_to_hub_in_organization sent the same model bytes as the first push. The staging server can deduplicate those bytes against the first push's LFS object and then reject the commit with "LFS pointer pointed to a file that does not exist". Change the weights before the second push so its bytes differ. Co-Authored-By: Claude Sonnet 5 --- tests/models/test_modeling_common.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 9968add19dd9..9986ef98e9ec 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -368,7 +368,12 @@ def test_push_to_hub_in_organization(self): # Push to hub via save_pretrained to a separate repo. Reusing `self.org_repo_id` after # deleting it makes the staging server's LFS GC reject the next commit with # "LFS pointer pointed to a file that does not exist" when the model bytes are identical. + # Change the weights so the second push's bytes differ from the first; identical bytes are deduplicated + # against the first push's LFS object, which the same GC race can then reject. save_org_repo_id = f"{self.org_repo_id}-saved" + with torch.no_grad(): + for p in model.parameters(): + p.add_(1.0) with tempfile.TemporaryDirectory() as tmp_dir: model.save_pretrained(tmp_dir, push_to_hub=True, token=TOKEN, repo_id=save_org_repo_id) From 6d8a643472996e62cea5508e983668b011f0d008 Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 5 Oct 2026 16:32:14 +0300 Subject: [PATCH 50/62] Revert "Avoid staging LFS dedup race in the org push test" This reverts commit 35e6cc4ef. Co-Authored-By: Claude Sonnet 5 --- tests/models/test_modeling_common.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 9986ef98e9ec..9968add19dd9 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -368,12 +368,7 @@ def test_push_to_hub_in_organization(self): # Push to hub via save_pretrained to a separate repo. Reusing `self.org_repo_id` after # deleting it makes the staging server's LFS GC reject the next commit with # "LFS pointer pointed to a file that does not exist" when the model bytes are identical. - # Change the weights so the second push's bytes differ from the first; identical bytes are deduplicated - # against the first push's LFS object, which the same GC race can then reject. save_org_repo_id = f"{self.org_repo_id}-saved" - with torch.no_grad(): - for p in model.parameters(): - p.add_(1.0) with tempfile.TemporaryDirectory() as tmp_dir: model.save_pretrained(tmp_dir, push_to_hub=True, token=TOKEN, repo_id=save_org_repo_id) From 0b4439be39dd362a07988dce87b13ef725d65537 Mon Sep 17 00:00:00 2001 From: leffff Date: Mon, 5 Oct 2026 16:32:14 +0300 Subject: [PATCH 51/62] Revert "Fix CI failures on the Kandinsky6 PR" This reverts commit 60cad4341. Co-Authored-By: Claude Sonnet 5 --- .github/workflows/pr_tests.yml | 4 ---- tests/others/test_utils.py | 2 +- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/.github/workflows/pr_tests.yml b/.github/workflows/pr_tests.yml index 8df8dd0146b3..66ce53d576e6 100644 --- a/.github/workflows/pr_tests.yml +++ b/.github/workflows/pr_tests.yml @@ -136,8 +136,6 @@ jobs: - name: Run fast PyTorch Pipeline CPU tests if: ${{ matrix.config.framework == 'pytorch_pipelines' }} - env: - PYTEST_TIMEOUT: 300 run: | pytest -n 8 --max-worker-restart=0 --dist=loadfile \ -k "not Onnx" \ @@ -154,8 +152,6 @@ jobs: - name: Run example PyTorch CPU tests if: ${{ matrix.config.framework == 'pytorch_examples' }} - env: - PYTEST_TIMEOUT: 300 run: | uv pip install ".[training]" pytest -n 4 --max-worker-restart=0 --dist=loadfile \ diff --git a/tests/others/test_utils.py b/tests/others/test_utils.py index 3741b1ee7012..d1a59cec52f1 100755 --- a/tests/others/test_utils.py +++ b/tests/others/test_utils.py @@ -175,7 +175,7 @@ def test_deprecate_stacklevel(self): with pytest.warns(FutureWarning) as warning: deprecate(("deprecated_arg", self.higher_version, "This message is better!!!"), standard_warn=False) assert str(warning[0].message) == "This message is better!!!" - assert "tests/others/test_utils.py" in warning[0].filename + assert "diffusers/tests/others/test_utils.py" in warning[0].filename def test_deprecate_testing_utils_module(self): import diffusers.utils.testing_utils From 82f46eff81f56c1285c4e56c627ed030a65f49a3 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Mon, 5 Oct 2026 10:47:42 -1000 Subject: [PATCH 52/62] Apply suggestion from @yiyixuxu --- docs/source/en/api/pipelines/pag.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/api/pipelines/pag.md b/docs/source/en/api/pipelines/pag.md index 18c0cfc26bf6..96026dec06ff 100644 --- a/docs/source/en/api/pipelines/pag.md +++ b/docs/source/en/api/pipelines/pag.md @@ -333,7 +333,7 @@ PAG reduces artifacts and improves the overall compposition. ## Configure parameters -### pag_applied_layers[[pag_applied_layers]] +### pag_applied_layers The `pag_applied_layers` argument allows you to specify which layers PAG is applied to. By default, it applies only to the mid blocks. Changing this setting will significantly impact the output. You can use the `set_pag_applied_layers` method to adjust the PAG layers after the pipeline is created, helping you find the optimal layers for your model. From 3bc67adedbba8abebd06cafb3ec082d429d32ca1 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Mon, 5 Oct 2026 10:48:19 -1000 Subject: [PATCH 53/62] Apply suggestion from @yiyixuxu --- docs/source/en/modular_diffusers/modular_pipeline.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/modular_diffusers/modular_pipeline.md b/docs/source/en/modular_diffusers/modular_pipeline.md index dcc493f4691a..e234a8aa4aae 100644 --- a/docs/source/en/modular_diffusers/modular_pipeline.md +++ b/docs/source/en/modular_diffusers/modular_pipeline.md @@ -195,7 +195,7 @@ ModularPipeline { If a component in your block doesn't exist in the repository, it remains `null` and is skipped during [`~ModularPipeline.load_components`]. -### from_pretrained[[from_pretrained]] +### from_pretrained [`~ModularPipeline.from_pretrained`] is a convenient way to create a [`ModularPipeline`] without defining blocks yourself. From ebea53af0bf4398849977765d682638b40a467be Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Mon, 5 Oct 2026 11:08:21 -1000 Subject: [PATCH 54/62] Apply batched suggestions from code review Co-authored-by: YiYi Xu --- docs/source/en/optimization/attention_backends.md | 2 +- docs/source/en/optimization/memory.md | 2 +- docs/source/en/training/distributed_inference.md | 2 +- docs/source/en/using-diffusers/cli.md | 4 ++-- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/source/en/optimization/attention_backends.md b/docs/source/en/optimization/attention_backends.md index 1318034639d5..a33fbe89815f 100644 --- a/docs/source/en/optimization/attention_backends.md +++ b/docs/source/en/optimization/attention_backends.md @@ -28,7 +28,7 @@ Refer to the table below for an overview of the available attention families and This guide will show you how to set and use the different attention backends. -## set_attention_backend[[set_attention_backend]] +## set_attention_backend The [`~ModelMixin.set_attention_backend`] method iterates through all the modules in the model and sets the appropriate attention backend to use. The attention backend setting persists until [`~ModelMixin.reset_attention_backend`] is called. diff --git a/docs/source/en/optimization/memory.md b/docs/source/en/optimization/memory.md index b79dafb7d634..a68628946015 100644 --- a/docs/source/en/optimization/memory.md +++ b/docs/source/en/optimization/memory.md @@ -468,7 +468,7 @@ apply_layerwise_casting( ) ``` -## torch.channels_last[[torchchannels_last]] +## torch.channels_last [torch.channels_last](https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html) flips how tensors are stored from `(batch size, channels, height, width)` to `(batch size, height, width, channels)`. This aligns the tensors with how the hardware sequentially accesses the tensors stored in memory and avoids skipping around in memory to access the pixel values. diff --git a/docs/source/en/training/distributed_inference.md b/docs/source/en/training/distributed_inference.md index abcbb0a5f286..818488b4cfce 100644 --- a/docs/source/en/training/distributed_inference.md +++ b/docs/source/en/training/distributed_inference.md @@ -114,7 +114,7 @@ Call `torchrun` to run the inference script and use the `--nproc_per_node` argum torchrun --nproc_per_node=2 run_distributed.py ``` -## device_map[[device_map]] +## device_map The `device_map` argument enables distributed inference by automatically placing model components on separate GPUs. This is especially useful when a model doesn't fit on a single GPU. You can use `device_map` to selectively load and unload the required model components at a given stage as shown in the example below (assumes two GPUs are available). diff --git a/docs/source/en/using-diffusers/cli.md b/docs/source/en/using-diffusers/cli.md index 0b75b7df8be4..01339f4628a2 100644 --- a/docs/source/en/using-diffusers/cli.md +++ b/docs/source/en/using-diffusers/cli.md @@ -287,7 +287,7 @@ diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 \ hf sandbox kill ``` -## `custom_blocks`[[custom_blocks]] +## `custom_blocks` Package a local `ModularPipelineBlocks` subclass for upload to the Hub. Reads a Python file, AST-scans it for subclasses of `ModularPipelineBlocks`, instantiates the chosen one, and calls `save_pretrained` in the current @@ -304,7 +304,7 @@ diffusers-cli custom_blocks --block_module_name my_block.py --block_class_name M The block class must be instantiable with zero constructor args and hardcodes defaults in `__init__` or read config from the pipeline `state` at call time. -## `fp16_safetensors`[[fp16_safetensors]] +## `fp16_safetensors` > [!CAUTION] > This command is now deprecated and will be removed in a future version. From 82d3eb59769479e7c78e151b2ed2884bc4418987 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Mon, 5 Oct 2026 23:12:50 +0200 Subject: [PATCH 55/62] remove review-only files from the Kandinsky6 PR Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- K6_DIFFUSERS_LAUNCH_GUIDE.md | 402 ----------------------------------- param_counts.json | 11 - 2 files changed, 413 deletions(-) delete mode 100644 K6_DIFFUSERS_LAUNCH_GUIDE.md delete mode 100644 param_counts.json diff --git a/K6_DIFFUSERS_LAUNCH_GUIDE.md b/K6_DIFFUSERS_LAUNCH_GUIDE.md deleted file mode 100644 index a204b3eb1753..000000000000 --- a/K6_DIFFUSERS_LAUNCH_GUIDE.md +++ /dev/null @@ -1,402 +0,0 @@ -# Launching Kandinsky 6 through Diffusers - -How to run every Kandinsky 6 (K6) model with the Diffusers code in -`diffusers-new-model-addition-kandinskyX` (branch `Kandinsky6`). The launches follow the notebooks in -`diffusers-new-model-addition-kandinskyX/notebooks/`. Code and Hub configs were checked as described in -[Verification status](#verification-status). - -## 1. What you can launch - -| # | Launch | Hub repo | Pipeline | Steps / guidance | Notes | -|---|--------|----------|----------|------------------|-------| -| 1 | Text → video + audio (T2VA), quality | `kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers` | `Kandinsky6TI2VAPipeline` | 50 steps, `guidance_scale=5.0` | `FlowMatchEulerDiscreteScheduler`; MagCache available | -| 2 | Text → video + audio, fast | `kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers` | `Kandinsky6TI2VAPipeline` | 16 steps, `guidance_scale=1.0` (mandatory) | `PiflowScheduler` rejects any other guidance; no MagCache | -| 3 | Image → video + audio (I2VA) | either repo above | `Kandinsky6TI2VAPipeline` | as 1 or 2 | pass `image=` | -| 4 | Video super-resolution, flow matching | `kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers` | `Kandinsky6SRPipeline` | `num_inference_steps=5` (= 4 Euler steps per tile) | x2, x2.25 (default), x4 | -| 5 | Video super-resolution, distilled | `kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers` | `Kandinsky6SRPipeline` | 2 model evaluations per tile, `num_inference_steps` ignored | `PiflowScheduler` | -| 6 | T2VA → SR chain | one of 1/2 + one of 4/5 | both | as above | video tensor handed over directly; audio muxed in afterward | - -Component layout (from each repo's `model_index.json`): - -- **T2VA repos:** `transformer` (`Kandinsky6Transformer3DModel`), `vae` (`AutoencoderKLHunyuanVideo`), `scheduler`, - `text_encoder` (Qwen2.5-VL), `tokenizer`, `text_encoder_2` (CLIP), `tokenizer_2`, `audio_vae` (`Kandinsky6AudioVAE`). - Audio is 44.1 kHz, video is 24 fps. -- **VSR repos:** `transformer` (`Kandinsky6SRTransformer3DModel`, 1.41B), `vae` (`Kandinsky6SRVAE`, 1.74B), - `latent_upscaler` (`Kandinsky6SRLatentUpscalerBank`, x2 + x4, 3.65B), `scheduler`. - -All K6 repos are private: the Hub answers 401 without a login (section 2.3). `Kandinsky-6.0-Lite-Diffusers` appears -only in the FastVideo and SGLang docs, not in the Diffusers code or notebooks, so it is not covered here. - -## 2. Setup - -### 2.1 Install this branch - -Always use the uv `.venv` of the checkout. Do not use conda and do not use `uv run`/`uv sync`, which can swap your -torch build. - -```bash -BASE_PATH="something" # folder that contains the K6 repos -cd "$BASE_PATH/diffusers-new-model-addition-kandinskyX" -uv venv && source .venv/bin/activate -uv pip install -e . -uv pip install av ipykernel -python -c "from diffusers import Kandinsky6TI2VAPipeline, Kandinsky6SRPipeline; print('K6 OK')" -``` - -- Stock Diffusers from PyPI has no `Kandinsky6*` classes, so it fails with an `ImportError`. Use this branch. -- Install a torch build that matches the GPU driver first (torch >= 2.6). Then run `uv pip install -e .`. -- `av` (PyAV) is required by `diffusers.utils.encode_video` to write mp4 files (with or without audio). Both - pipelines take/return tensors only — neither reads nor writes video files itself. -- `natten` and `flash_attn` are not required. The default `native` backend needs neither. -- On a remote GPU box, run the same block from the checkout on that box after `git pull` there. - -### 2.2 GPU memory - -The pipelines are meant to run with `enable_model_cpu_offload()`. That keeps one component on the GPU at a time. The -T2VA order is `text_encoder -> text_encoder_2 -> transformer -> vae -> audio_vae`. The SR order is -`source_vae -> latent_upscaler -> transformer -> vae`. Without offload, the full T2VA pipeline (Qwen2.5-VL plus a -60-block, 4096-wide transformer) needs a large GPU. - -### 2.3 Hugging Face access - -Log in once, interactively, and never paste a token into a script or notebook: - -```bash -hf auth login -hf auth whoami -``` - -For non-interactive shells use `export HF_TOKEN=...` in the shell, never in a file. After the first download the -snapshot is cached in `~/.cache/huggingface/hub`. Set `HF_HUB_OFFLINE=1` to run without network. - -Note: the notebooks contain `os.environ["HF_TOKEN"] = "TOKEN"`. Delete that line rather than putting a real token there. - -## 3. Text/image → video + audio (`Kandinsky6TI2VAPipeline`) - -Defaults of `__call__`: `height=512, width=768, num_frames=121, sample_fps=24.0, num_inference_steps=50, -guidance_scale=5.0, max_sequence_length=1024, sample_audio=True, expand_prompts=False, output_type="pt"`. -`height` and `width` must be divisible by 8, and `max_sequence_length` must be between 1 and 1024. The notebooks use -`480 x 864`, which is the size the `Pro` checkpoints were configured for (`k6_pro_125_480_864_*` in `k6_video`). - -`output_type="pt"` returns `frames` as a **uint8 tensor `[B, 3, T, H, W]` in 0..255**. `audio` is a list with one 1-D -mono waveform per video. Prompts can carry speech as `spoken words`. - -### 3.1 Launch 1: Pro-sft, quality (50 steps, CFG 5.0, MagCache) - -Save as `run_t2va_sft.py` and run it from the venv: - -```python -import torch -from diffusers import Kandinsky6TI2VAPipeline -from diffusers.hooks import MagCacheConfig -from diffusers.utils import encode_video - -MODEL = "kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers" -GPU_ID = 0 -STEPS, GUIDANCE, FPS, SEED = 50, 5.0, 24.0, 1137 -USE_MAGCACHE = True -MODE = "t2va" # "t2va" or "i2va": selects the MagCache coefficients - -pipe = Kandinsky6TI2VAPipeline.from_pretrained(MODEL, dtype=torch.bfloat16) -pipe.transformer.set_attention_backend("native") # "_flash_3" needs FlashAttention 3 (Hopper) - -cfg = pipe.transformer.config.get("magcache") -if USE_MAGCACHE and cfg is not None: - ratios = cfg["mag_ratios"] - ratios = ratios[MODE] if isinstance(ratios, dict) else ratios - pipe.transformer.enable_cache( - MagCacheConfig( - threshold=cfg["threshold"], - max_skip_steps=cfg["max_skip_steps"], - retention_ratio=cfg["retention_ratio"], - num_inference_steps=2 * STEPS, # coefficients interleave cond and uncond, so 2 forwards per step - mag_ratios=list(ratios), - ) - ) - -pipe.enable_model_cpu_offload(gpu_id=GPU_ID) - -result = pipe( - prompt="A news anchor delivers the evening news in a clear, neutral tone, saying: " - "Good evening, ladies and gentlemen, and welcome to the evening news.", - negative_prompt="blurry, distorted, low quality, noisy audio", - height=480, - width=864, - num_frames=121, - sample_fps=FPS, - num_inference_steps=STEPS, - guidance_scale=GUIDANCE, - sample_audio=True, - generator=torch.Generator().manual_seed(SEED), - output_type="pt", -) - -video = result.frames[0].permute(1, 2, 3, 0) # [T,H,W,3] uint8 -audio = torch.as_tensor(result.audio[0])[:, None].repeat(1, 2) # mono -> stereo [samples, 2] -encode_video(video, fps=int(FPS), output_path="t2va_sft.mp4", - audio=audio, audio_sample_rate=int(pipe.audio_sample_rate)) -print("saved t2va_sft.mp4") -``` - -Notes: - -- The MagCache coefficients live in the transformer config under `magcache`, with separate lists for `t2va` and `i2va`. - Set `MODE` to match the generation you run. -- To turn MagCache off, set `USE_MAGCACHE = False`. You can also call `pipe.transformer.disable_cache()`. -- MagCache is used only with the sft checkpoint. The authors' notebook disables it for the distilled one. - -### 3.2 Launch 2: Pro-distill, fast (PiFlow, 16 steps, guidance 1.0) - -Same script with these changes: - -```python -MODEL = "kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers" -STEPS, GUIDANCE = 16, 1.0 -USE_MAGCACHE = False -``` - -The output path can become `t2va_distill.mp4`. `guidance_scale` must stay `1.0` here: the denoise loop raises -`ValueError("PiflowScheduler requires guidance_weight=1.0")` for any other value. With guidance 1.0 there is no -unconditional pass, so no `negative_prompt` is needed. - -You can choose steps and guidance automatically, as the notebook does: - -```python -is_piflow = pipe.scheduler.__class__.__name__ == "PiflowScheduler" -STEPS = 16 if is_piflow else 50 -GUIDANCE = 1.0 if is_piflow else 5.0 -``` - -### 3.3 Launch 3: image → video + audio (I2VA) - -Pass a reference image path (or a PIL image) as `image=`. Set `MODE = "i2va"` for MagCache. For a batch, pass one image -per prompt: - -```python -BASE_PATH = "something" # folder that contains the K6 repos - -result = pipe( - prompt=["A woman turns to the camera and says Hello there!"], - image=[f"{BASE_PATH}/k6_video/assets/i2va_input.png"], - negative_prompt=["blurry, distorted, low quality, noisy audio"], - height=480, width=864, num_frames=121, sample_fps=24.0, - num_inference_steps=STEPS, guidance_scale=GUIDANCE, - sample_audio=True, generator=torch.Generator().manual_seed(1137), output_type="pt", -) -``` - -The repo has a sample image at `k6_video/assets/i2va_input.png`. It is not in the diffusers checkout. - -### 3.4 Optional switches - -| Switch | How | Comment | -|--------|-----|---------| -| Video only | `sample_audio=False` | `result.audio` is then `None`; skip the `audio=` arguments | -| Batching | lists for `prompt`, `negative_prompt`, `image` | the notebook uses batch size 2 | -| Prompt expansion | `expand_prompts=True` | rewrites the prompt with the Qwen2.5-VL model first; slower. Default is `False`. The notebook hard-codes `True` in the call | -| NumPy output | `output_type="np"` | `[B, T, H, W, 3]` | -| Faster attention | `pipe.transformer.set_attention_backend("_flash_3")` | needs FlashAttention 3 (Hopper). Other valid names: `flash`, `native`, `sage`, `xformers`, ... | -| `torch.compile` | see the block below | the notebook uses `native` attention when compiling | - -Notebook-style block compile (only for repeated runs; call it after loading and before `enable_model_cpu_offload`): - -```python -from diffusers.models.transformers.transformer_kandinsky6 import ( - Kandinsky6FusedTransformerDecoderBlock, - Kandinsky6TransformerDecoderBlock, - Kandinsky6TransformerEncoderBlock, -) - -torch.set_float32_matmul_precision("high") -pipe.transformer.set_attention_backend("native") -for block in pipe.transformer.modules(): - if isinstance(block, Kandinsky6FusedTransformerDecoderBlock): - block.forward = torch.compile(block.forward) - elif isinstance(block, (Kandinsky6TransformerEncoderBlock, Kandinsky6TransformerDecoderBlock)): - block.forward = torch.compile(block.forward, mode="max-autotune-no-cudagraphs", dynamic=True) -``` - -## 4. Video super-resolution (`Kandinsky6SRPipeline`) - -Like `Kandinsky6TI2VAPipeline`, this pipeline only takes and returns tensors — no file I/O. Use -`diffusers.utils.load_video` / `encode_video` (or your own decoder) to bridge to mp4 files. - -Inputs and constants: - -- `video`: `uint8` pixels in `[0, 255]`, shape `(3, T, H, W)` or a batch `(B, 3, T, H, W)` — the same layout - `Kandinsky6TI2VAPipeline` returns with `output_type="pt"`. `T` must be `1 + 8k` (the SR model's temporal chunk - size; align/truncate yourself — see `notebooks/kandinsky6_sr.py` for a `load_video`-based helper). A batch requires - every clip to already share resolution and frame count. -- `latents`: a pre-encoded SR-VAE latent, shape `(C, T', H', W')` or a batch `(B, C, T', H', W')`. Mutually exclusive - with `video`. Pass `kvae_bridge=True` when the latents come from a *different* (base T2VA) VAE — the pipeline - decodes them through `source_vae` and re-encodes in its own latent space before tiling. -- `resolution_scale` is `2`, `2.25` (default), or `4`. The `2.25` route is a x1.125 pixel pre-upscale followed by the - x2 model. It is the route used after K6 generation. -- Other knobs are `min_overlap` (minimum tile overlap fraction), `tiles_batch_size`, and `output_type` (`"pt"` or - `"np"`). - -### 4.1 Launch 4: VSR flow-matching - -```python -import numpy as np -import torch -from diffusers import Kandinsky6SRPipeline -from diffusers.utils import encode_video, load_video - -pipe = Kandinsky6SRPipeline.from_pretrained( - "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers", dtype=torch.bfloat16 -) -pipe.enable_model_cpu_offload(gpu_id=0) # or: pipe.to("cuda") - -frames, fps = load_video("input.mp4", return_fps=True) -video = torch.stack([torch.as_tensor(np.array(f)) for f in frames]).permute(3, 0, 1, 2) # [3,T,H,W] uint8 -video = video[:, : 1 + 8 * ((video.shape[1] - 1) // 8)] # align to 1+8k frames - -result = pipe( - video=video, - resolution_scale=2.25, # 2, 2.25 or 4 - num_inference_steps=5, # 4 Euler steps per tile (default) - generator=torch.Generator("cuda:0").manual_seed(1137), -) -print(result.frames.shape) # [1, 3, T, H*scale, W*scale] uint8 -encode_video(result.frames[0].permute(1, 2, 3, 0), fps=round(fps), output_path="output_sr.mp4") -``` - -Batch call (equal-size clips, stacked into one tensor): - -```python -result = pipe( - video=torch.stack([video_a, video_b]), # both [3,T,H,W], same shape - resolution_scale=2.25, - output_type="pt", - generator=torch.Generator("cuda:0").manual_seed(1137), -) -``` - -### 4.2 Launch 5: VSR distilled (2 evaluations per tile) - -The same code with the other repo. Do not pass `num_inference_steps`; it has no effect here. - -```python -pipe = Kandinsky6SRPipeline.from_pretrained( - "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", dtype=torch.bfloat16 -) -pipe.enable_model_cpu_offload(gpu_id=0) -result = pipe(video=video, resolution_scale=2, generator=torch.Generator("cuda:0").manual_seed(1137)) -``` - -## 5. Launch 6: T2VA → SR chain - -Both pipelines stay independent. Stage 2 takes the stage-1 video tensor directly, with no intermediate mp4; the -source audio is kept aside and muxed in only at the very end. This version uses the Hub repos instead of the local -bundles in the stacked notebook. - -```python -import torch -from diffusers import Kandinsky6SRPipeline, Kandinsky6TI2VAPipeline -from diffusers.utils import encode_video - -GPU_ID, FPS = 0, 24.0 - -t2va = Kandinsky6TI2VAPipeline.from_pretrained( - "kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers", dtype=torch.bfloat16) -t2va.enable_model_cpu_offload(gpu_id=GPU_ID) - -sr = Kandinsky6SRPipeline.from_pretrained( - "kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers", dtype=torch.bfloat16) -sr.enable_model_cpu_offload(gpu_id=GPU_ID) - -gen = t2va( - prompt="A cinematic shot of a small red sailboat crossing a calm blue lake at sunset, " - "with gentle waves and natural ambient sound.", - negative_prompt="blurry, distorted, low quality, noisy audio", - height=480, width=864, - num_frames=121, # not 125: SR keeps at most 121 frames - sample_fps=FPS, num_inference_steps=50, guidance_scale=5.0, - sample_audio=True, output_type="pt", -) - -out = sr(video=gen.frames, resolution_scale=2.25, output_type="pt") # uint8 [B,3,T,H,W] in, out -print(tuple(out.frames.shape)) - -audio = torch.as_tensor(gen.audio[0])[:, None].repeat(1, 2) -encode_video( - out.frames[0].permute(1, 2, 3, 0), - fps=int(FPS), - output_path="t2va_sr.mp4", - audio=audio, - audio_sample_rate=int(t2va.audio_sample_rate), -) -``` - -The output is 2.25x the generation size: `480 x 864` becomes about `1080 x 1944`. For the fast variant swap in the two -distilled repos, `STEPS=16`, `guidance_scale=1.0`, and drop `negative_prompt`. - -## 6. Making the notebooks run - -The three notebooks in `diffusers-new-model-addition-kandinskyX/notebooks/` were copied from a `k6_video` layout. -Three edits are needed: - -1. **Repo detection.** Each notebook walks up until `dev/exports/diffusers` exists. That folder is in `k6_video`, not in - the diffusers checkout, so the first cell raises `RuntimeError`. Replace that block with - `from pathlib import Path; REPO = Path.cwd()`. -2. **Inputs.** `assets/vsr_input.mp4` does not exist anywhere in the workspace. Point `INPUT_VIDEO` at a real mp4. - `assets/i2va_input.png` exists only in `k6_video/assets/`. -3. **Stacked notebook.** It loads `outputs/diffusers_bundle/kandinsky6_diffusers` and `.../kandinsky6_sr_diffusers`, - which are local bundles from `k6_video`. Replace them with the Hub ids from section 1. It also generates - `num_frames=125`; `Kandinsky6SRPipeline` requires `1 + 8k` frames and raises `ValueError` on anything else - (it no longer silently truncates), so use 121. - -Also remove the `HF_TOKEN = "TOKEN"` line in `kandinsky6_ti2va.ipynb`, and register the kernel with -`python -m ipykernel install --user --name k6-diffusers` from the activated venv. - -## 7. Building a local Diffusers bundle from a native checkpoint (only if needed) - -Official Hub repos exist for every launch above, so this is not needed for normal use. To make a bundle from a native -checkpoint, the `k6_video` repo (its own `.venv` with `just setup`) provides: - -```bash -BASE_PATH="something" # folder that contains the K6 repos -cd "$BASE_PATH/k6_video" -just convert-diffusers --config src/kandinsky/configs/k6_pro_125_480_864_mCache_mOffload.yaml --output-dir outputs/kandinsky6_diffusers -just convert-diffusers-sr --config src/kandinsky/configs/k6_pro_125_480_864_mCache_mOffload.yaml --output-dir outputs/kandinsky6_sr_diffusers -``` - -Bundles created this way are loaded with the same `from_pretrained("")` call. The SR recipe above needs -private checkpoints and extra flags (`--checkpoint-path`, `--vae-path`, `--latent-upscaler-config`); see the `JUSTFILE` -comments. `just integrate-diffusers ` overlays the generated K6 code onto a Diffusers checkout. It does not -commit or push. - -## 8. Troubleshooting - -| Symptom | Cause / fix | -|---------|-------------| -| `ImportError: cannot import name 'Kandinsky6TI2VAPipeline'` | stock Diffusers is installed; activate the venv where this branch is installed with `-e` | -| `401` / `RepositoryNotFoundError` on a K6 repo | not logged in, or the account has no access; `hf auth login` | -| `ValueError: PiflowScheduler requires guidance_weight=1.0` | distilled repo with `guidance_scale != 1.0` | -| `ValueError: height and width must be divisible by 8` | choose sizes divisible by 8 (480x864 works) | -| `ValueError: Expected samples with 2 channels` in `encode_video` | audio must be stereo: `torch.as_tensor(a)[:, None].repeat(1, 2)` | -| `No module named 'av'` | `uv pip install av` (needed by `encode_video` for mp4 output, with or without audio) | -| Error setting `_flash_3` | FlashAttention 3 not available; use `native` | -| CUDA OOM | make sure `enable_model_cpu_offload()` is on; lower `height`/`width`/`num_frames`; use the distilled model | -| `ValueError: \`video\` must have 1 + 8k frames` / `... at most 121 frames` | align/split the input yourself: 121 frames at 24 fps is about 5 s | - -## Verification status - -Checked on the Mac (CPU only, no GPU): - -- The branch imports; every keyword used above exists in the current signatures of `Kandinsky6TI2VAPipeline.__call__`, - `Kandinsky6SRPipeline.__call__`, `encode_video`, `load_video`, `MagCacheConfig` and - `transformer.set_attention_backend`/`enable_cache`. `dtype=` and `torch_dtype=` are both accepted by `from_pretrained`. -- `_flash_3` and `native` are valid attention-backend names, and `PiflowScheduler` guidance is enforced as stated. -- Defaults, output layouts (`uint8 [B,3,T,H,W]`), the SR tensor layouts, and the 121-frame clip come from the code. -- The three Hub repos `Pro-sft-5s`, `VSR-5s` and `VSR-distilled2steps-5s` (all `-Diffusers`) were read from the local HF - cache. Their `model_index.json`, scheduler and component configs match what is written here. - -**Not verified:** - -- No end-to-end run happened: no model weights were loaded and nothing was generated. Runtime, peak VRAM and output - quality are untested. -- `Kandinsky-6.0-Pro-distill-5s-Diffusers` is not in the local cache and cannot be checked without access. Its id and - the 16-step / guidance-1.0 settings come from your notebook and the class docstring. -- The MagCache path was checked against `MagCacheConfig`'s signature, not run. diff --git a/param_counts.json b/param_counts.json deleted file mode 100644 index 694d3a13df5d..000000000000 --- a/param_counts.json +++ /dev/null @@ -1,11 +0,0 @@ -{ - "pro.transformer": 30136285268, - "pro.vae": 246478803, - "pro.audio_vae(MMAudioVAE)": 427632434, - "vsr.transformer": 1413231168, - "vsr.vae": 1738878243, - "vsr.latent_upscaler": 3654021952, - "vsrd.transformer": 1414263936, - "vsrd.vae": 1738878243, - "vsrd.latent_upscaler": 3654021952 -} \ No newline at end of file From bbab023342f698d2b3739855686a14af0a762fa7 Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Mon, 5 Oct 2026 11:15:46 -1000 Subject: [PATCH 56/62] Apply suggestion from @yiyixuxu --- .gitignore | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 44f8e92185bd..95b27c9084f2 100644 --- a/.gitignore +++ b/.gitignore @@ -185,4 +185,4 @@ wandb # AI agent generated symlinks /.agents/skills -/.claude/skillsnotebooks/ +/.claude/skills From 34f329fcb91989d9939bf350f3203f780c229e3c Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Mon, 5 Oct 2026 23:18:15 +0200 Subject: [PATCH 57/62] restore .gitignore to match main Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- .gitignore | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 95b27c9084f2..7b156e460abf 100644 --- a/.gitignore +++ b/.gitignore @@ -185,4 +185,4 @@ wandb # AI agent generated symlinks /.agents/skills -/.claude/skills +/.claude/skills \ No newline at end of file From 1ea5fb3008a098d85d1a2260a703bcb809a15eca Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Mon, 5 Oct 2026 23:42:21 +0200 Subject: [PATCH 58/62] simplify the mono audio handling in _write_audio Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- src/diffusers/utils/export_utils.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/src/diffusers/utils/export_utils.py b/src/diffusers/utils/export_utils.py index 841b34d6c116..5ceb3c5b3452 100644 --- a/src/diffusers/utils/export_utils.py +++ b/src/diffusers/utils/export_utils.py @@ -276,26 +276,25 @@ def _write_audio( ) -> None: import torch - if samples.ndim == 1: - samples = samples[None, :] + # Mono audio is written as stereo, with the channel duplicated. + if samples.ndim == 1 or 1 in samples.shape: + samples = samples.reshape(-1, 1).repeat(1, 2) - # `samples` is documented as [audio_channels, samples]; accept the transposed form too. - if samples.shape[0] not in (1, 2) and samples.shape[1] in (1, 2): + if samples.shape[1] != 2 and samples.shape[0] == 2: samples = samples.T - if samples.shape[0] not in (1, 2): - raise ValueError(f"Expected samples with 1 or 2 channels; got shape {samples.shape}.") + if samples.shape[1] != 2: + raise ValueError(f"Expected samples with 2 channels; got shape {samples.shape}.") # Convert to int16 packed for ingestion; resampler converts to the encoder format. if samples.dtype != torch.int16: samples = torch.clip(samples, -1.0, 1.0) samples = (samples * 32767.0).to(torch.int16) - num_channels = samples.shape[0] frame_in = av_module.AudioFrame.from_ndarray( - samples.T.contiguous().reshape(1, -1).cpu().numpy(), + samples.contiguous().reshape(1, -1).cpu().numpy(), format="s16", - layout="mono" if num_channels == 1 else "stereo", + layout="stereo", ) frame_in.sample_rate = audio_sample_rate From 871acef5487a67812135e7789bd5759491697b8e Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 6 Oct 2026 00:15:11 +0200 Subject: [PATCH 59/62] remove the no-op set_attention_backend calls from the Kandinsky6 docs Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- docs/source/en/api/models/kandinsky6_transformer.md | 6 +++--- docs/source/en/api/pipelines/kandinsky6.md | 5 ++--- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/docs/source/en/api/models/kandinsky6_transformer.md b/docs/source/en/api/models/kandinsky6_transformer.md index e48aac25884d..a6139b2fb14f 100644 --- a/docs/source/en/api/models/kandinsky6_transformer.md +++ b/docs/source/en/api/models/kandinsky6_transformer.md @@ -43,9 +43,9 @@ from diffusers import Kandinsky6SRTransformer3DModel transformer = Kandinsky6SRTransformer3DModel.from_pretrained( "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", subfolder="transformer", torch_dtype=torch.bfloat16 ) -# The transformer always runs NABLA sparse attention (`nabla_threshold`, 0.8 by default), like Kandinsky 5's `nabla` -# checkpoints. Only the `flex` backend can execute it. -transformer.set_attention_backend("flex") +# The transformer always runs NABLA sparse attention (`nabla_threshold`, 0.8 by default) on the `flex` backend. +# Compile it, otherwise flex falls back to an eager implementation that needs far more memory at video resolutions. +transformer.compile_repeated_blocks(fullgraph=True) ``` [[autodoc]] Kandinsky6SRTransformer3DModel diff --git a/docs/source/en/api/pipelines/kandinsky6.md b/docs/source/en/api/pipelines/kandinsky6.md index 7f9ba2e16e73..f3380c394244 100644 --- a/docs/source/en/api/pipelines/kandinsky6.md +++ b/docs/source/en/api/pipelines/kandinsky6.md @@ -76,9 +76,8 @@ windows. sr_pipe = Kandinsky6SRPipeline.from_pretrained( "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", torch_dtype=torch.bfloat16 ) -# The SR transformer always runs NABLA sparse attention (like Kandinsky 5's `nabla` checkpoints), which only the -# `flex` backend can execute. -sr_pipe.transformer.set_attention_backend("flex") +# The SR transformer always runs NABLA sparse attention on the `flex` backend. Compile it, otherwise flex falls +# back to an eager implementation that needs far more memory at video resolutions. sr_pipe.transformer.compile_repeated_blocks(fullgraph=True) sr_pipe.enable_model_cpu_offload() From 9c780584b63cbc4d252eda0b5e8eecdea25a7679 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 6 Oct 2026 00:25:58 +0200 Subject: [PATCH 60/62] give the Kandinsky6 transformer its own output class instead of reusing LTX-2's Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- .../transformers/transformer_kandinsky6.py | 27 +++++++++++++++---- .../models/transformers/transformer_ltx2.py | 2 +- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_kandinsky6.py b/src/diffusers/models/transformers/transformer_kandinsky6.py index 8ba18265a67c..3acfd8afe484 100644 --- a/src/diffusers/models/transformers/transformer_kandinsky6.py +++ b/src/diffusers/models/transformers/transformer_kandinsky6.py @@ -17,6 +17,7 @@ from __future__ import annotations import math +from dataclasses import dataclass from typing import Any import torch @@ -24,12 +25,28 @@ from ...configuration_utils import ConfigMixin, register_to_config from ...loaders import PeftAdapterMixin +from ...utils import BaseOutput from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn from ..cache_utils import CacheMixin from ..embeddings import TimestepEmbedding, Timesteps from ..modeling_utils import ModelMixin, get_parameter_dtype -from .transformer_ltx2 import AudioVisualModelOutput + + +@dataclass +class Kandinsky6TransformerOutput(BaseOutput): + r""" + The output of [`Kandinsky6Transformer3DModel`]. + + Args: + sample (`torch.Tensor` of shape `(batch_size, num_frames, height, width, out_visual_dim)`): + The predicted video velocity. + audio_sample (`torch.Tensor` of shape `(batch_size, audio_length, out_audio_dim)`, *optional*): + The predicted audio velocity, `None` when the model was called without `audio_hidden_states`. + """ + + sample: torch.Tensor + audio_sample: torch.Tensor | None = None def get_freqs(dim: int, max_period: float = 10000.0) -> Tensor: @@ -730,7 +747,7 @@ def forward( encoder_attention_mask: Tensor | None = None, visual_token_type_ids: Tensor | None = None, return_dict: bool = True, - ) -> AudioVisualModelOutput | tuple[Tensor, ...]: + ) -> Kandinsky6TransformerOutput | tuple[Tensor, ...]: r""" Args: hidden_states (`torch.Tensor` of shape `(batch_size, num_frames, height, width, in_channels)`): @@ -754,10 +771,10 @@ def forward( Per-frame token type ids, embedded through `visual_token_type_embeddings`. Requires `visual_token_type_num_embeddings > 0`. return_dict (`bool`, defaults to `True`): - Whether to return an [`AudioVisualModelOutput`] instead of a plain tuple. + Whether to return a [`Kandinsky6TransformerOutput`] instead of a plain tuple. Returns: - [`AudioVisualModelOutput`] or `tuple`: + [`Kandinsky6TransformerOutput`] or `tuple`: The predicted video velocity in the `(B, T, H, W, out_visual_dim)` layout and, when `audio_hidden_states` was given, the predicted audio velocity of shape `(B, audio_length, out_audio_dim)`. @@ -824,4 +841,4 @@ def forward( if not return_dict: return (video_out,) if audio_out is None else (video_out, audio_out) - return AudioVisualModelOutput(sample=video_out, audio_sample=audio_out) + return Kandinsky6TransformerOutput(sample=video_out, audio_sample=audio_out) diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py index 23f49c51256b..e5462e4b047a 100644 --- a/src/diffusers/models/transformers/transformer_ltx2.py +++ b/src/diffusers/models/transformers/transformer_ltx2.py @@ -98,7 +98,7 @@ class AudioVisualModelOutput(BaseOutput): """ sample: "torch.Tensor" # noqa: F821 - audio_sample: "torch.Tensor | None" = None # noqa: F821 + audio_sample: "torch.Tensor" # noqa: F821 class LTX2AdaLayerNormSingle(nn.Module): From c8a0c1f6e21f63930db635ebe2c5dc9cd4ab349d Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Tue, 6 Oct 2026 06:02:17 +0200 Subject: [PATCH 61/62] fix the stale pipeline exports and add return annotations for the Kandinsky6 pipelines Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS --- src/diffusers/pipelines/__init__.py | 4 ---- src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py | 2 +- .../pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py | 2 +- 3 files changed, 2 insertions(+), 6 deletions(-) diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 886c6b6b749d..b8cb626d8283 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -452,12 +452,10 @@ "Kandinsky5I2IPipeline", ] _import_structure["kandinsky6"] = [ - "Kandinsky6SRLatentUpscalerBank", "Kandinsky6SRPipeline", "Kandinsky6SRPipelineOutput", "Kandinsky6TI2VAPipeline", "Kandinsky6TI2VAPipelineOutput", - "MMAudioVocoder", ] _import_structure["z_image"] = [ "ZImageControlNetInpaintPipeline", @@ -790,12 +788,10 @@ Kandinsky5T2VPipeline, ) from .kandinsky6 import ( - Kandinsky6SRLatentUpscalerBank, Kandinsky6SRPipeline, Kandinsky6SRPipelineOutput, Kandinsky6TI2VAPipeline, Kandinsky6TI2VAPipelineOutput, - MMAudioVocoder, ) from .krea2 import Krea2Pipeline from .latent_consistency_models import ( diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py index 5da4578e1465..12d33d62abb4 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_sr.py @@ -240,7 +240,7 @@ def __call__( generator: torch.Generator | list[torch.Generator] | None = None, output_type: str = "pil", return_dict: bool = True, - ): + ) -> Kandinsky6SRPipelineOutput | tuple: r""" The call function to the pipeline for super-resolution. diff --git a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py index 3c26416407d1..656f2077c47c 100644 --- a/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py +++ b/src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py @@ -671,7 +671,7 @@ def __call__( return_dict: bool = True, callback_on_step_end: Callable[[int, int, dict], None] | None = None, callback_on_step_end_tensor_inputs: list[str] = ["latents"], - ): + ) -> Kandinsky6TI2VAPipelineOutput | tuple: r""" The call function to the pipeline for generation. From 7245998ace13799f8e7c7facde6d288539dd0133 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 6 Oct 2026 16:29:59 +0530 Subject: [PATCH 62/62] address feedback from @leffff --- docs/source/en/api/pipelines/kandinsky6.md | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/docs/source/en/api/pipelines/kandinsky6.md b/docs/source/en/api/pipelines/kandinsky6.md index f3380c394244..db4cdb1944cc 100644 --- a/docs/source/en/api/pipelines/kandinsky6.md +++ b/docs/source/en/api/pipelines/kandinsky6.md @@ -28,8 +28,12 @@ super-resolution model upscales the generated video tile by tile in the latent s | Model | Pipeline | Notes | |---|---|---| -| [`kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Pro-sft-5s-Diffusers) | [`Kandinsky6TI2VAPipeline`] | Flow matching, `guidance_scale=5.0`, 50 steps | -| [`kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers) | [`Kandinsky6TI2VAPipeline`] | Distilled, `guidance_scale=1.0`, 16 steps | +| [`kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Pro-5s-Diffusers) | [`Kandinsky6TI2VAPipeline`] | Flow matching, `guidance_scale=5.0`, 50 steps | +| [`kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Pro-distill-5s-Diffusers) | [`Kandinsky6TI2VAPipeline`] | Distilled, `guidance_scale=1.0`, 10 steps | +| [`kandinskylab/Kandinsky-6.0-Lite-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Lite-5s-Diffusers) | [`Kandinsky6TI2VAPipeline`] | Flow matching, `guidance_scale=5.0`, 50 steps | +| [`kandinskylab/Kandinsky-6.0-Lite-distill-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Lite-distill-5s-Diffusers) | [`Kandinsky6TI2VAPipeline`] | Distilled, `guidance_scale=1.0`, 10 steps | +| [`kandinskylab/Kandinsky-6.0-Pro-pretrain-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Pro-pretrain-5s-Diffusers) | [`Kandinsky6TI2VAPipeline`] | Flow matching, `guidance_scale=5.0`, 50 steps | +| [`kandinskylab/Kandinsky-6.0-Lite-pretrain-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-Lite-pretrain-5s-Diffusers) | [`Kandinsky6TI2VAPipeline`] | Flow matching, `guidance_scale=5.0`, 50 steps | | [`kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-VSR-5s-Diffusers) | [`Kandinsky6SRPipeline`] | Flow matching super-resolution | | [`kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers`](https://huggingface.co/kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers) | [`Kandinsky6SRPipeline`] | Distilled super-resolution, 2 steps | @@ -50,7 +54,7 @@ output = pipe( height=480, width=864, num_frames=121, - num_inference_steps=16, + num_inference_steps=10, guidance_scale=1.0, ) encode_video( @@ -73,13 +77,17 @@ is refined at one of the tile sizes the SR transformer was trained on, and the t windows. ```python +# required: lets inductor pick flex-attention tiles that fit the SR block mask +torch._inductor.config.max_autotune = True + sr_pipe = Kandinsky6SRPipeline.from_pretrained( "kandinskylab/Kandinsky-6.0-VSR-distilled2steps-5s-Diffusers", torch_dtype=torch.bfloat16 ) # The SR transformer always runs NABLA sparse attention on the `flex` backend. Compile it, otherwise flex falls # back to an eager implementation that needs far more memory at video resolutions. -sr_pipe.transformer.compile_repeated_blocks(fullgraph=True) sr_pipe.enable_model_cpu_offload() +sr_pipe.transformer.set_attention_backend("flex") +sr_pipe.transformer.compile_repeated_blocks(fullgraph=True) upscaled = sr_pipe(video=output.frames[0], resolution_scale=2.25, num_inference_steps=2).frames[0] ```