AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-26 08:59:24 +03:00
Code Activity

address comments

Browse Source
This commit is contained in:
leffff
2026-10-01 18:48:59 +03:00
parent 28c5326da8
commit 372f29088c
6 changed files with 68 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py
View File

@@ -363,6 +363,16 @@ class Kandinsky6SRUpsample(nn.Module):
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 @@ class Kandinsky6SREncoder3D(nn.Module):
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,

30
src/diffusers/models/autoencoders/autoencoder_mmaudio.py
View File

@@ -31,6 +31,7 @@ import torch.nn.functional as F
from ...configuration_utils import ConfigMixin, register_to_config
from ...utils.accelerate_utils import apply_forward_hook
from ..attention import AttentionModuleMixin
from ..attention_dispatch import AttentionBackendName, dispatch_attention_fn
from ..modeling_outputs import AutoencoderKLOutput
from ..modeling_utils import ModelMixin
@@ -131,16 +132,12 @@ class MMAudioResnetBlock1D(nn.Module):
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 @@ class MMAudioAttnBlock1D(nn.Module):
# 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):

3
src/diffusers/models/transformers/transformer_kandinsky6_sr.py
View File

@@ -565,6 +565,3 @@ class Kandinsky6SRTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin,
if not return_dict:
return (output,)
return Transformer2DModelOutput(sample=output)
__all__ = ["Kandinsky6SRTransformer3DModel"]

40
src/diffusers/pipelines/kandinsky6/modeling_latent_upscaler.py
View File

@@ -36,10 +36,9 @@ class Kandinsky6SRLatentUpscalerConv3d(nn.Conv3d):
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 @@ class Kandinsky6SRLatentUpscalerUpsample(nn.Module):
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 @@ class Kandinsky6SRLatentUpscalerX2Branch(nn.Module):
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 @@ class Kandinsky6SRLatentUpscalerBank(ModelMixin, ConfigMixin):
if not return_dict:
return (upscaled,)
return DecoderOutput(sample=upscaled)
__all__ = ["Kandinsky6SRLatentUpscalerBank"]

6
src/diffusers/pipelines/kandinsky6/modeling_vocoder.py
View File

@@ -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):

12
src/diffusers/pipelines/kandinsky6/pipeline_kandinsky6_ti2va.py
View File

@@ -522,12 +522,14 @@ class Kandinsky6TI2VAPipeline(DiffusionPipeline):
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 @@ class Kandinsky6TI2VAPipeline(DiffusionPipeline):
)
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 @@ class Kandinsky6TI2VAPipeline(DiffusionPipeline):
# 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.