# Copyright 2025 The HuggingFace Team. All rights reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

from dataclasses import dataclass

from typing import Tuple, Union

import torch

from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config

from ...utils import BaseOutput, logging

from ..attention import AttentionMixin

from ..attention_processor import Attention, AttnProcessor

from ..embeddings import TimestepEmbedding, Timesteps

from ..modeling_utils import ModelMixin

logger = logging.get_logger(__name__) # pylint: disable=invalid-name

@dataclass

class Kandinsky3UNetOutput(BaseOutput):

sample: torch.Tensor = None

class Kandinsky3EncoderProj(nn.Module):

def __init__(self, encoder_hid_dim, cross_attention_dim):

super().__init__()

self.projection_linear = nn.Linear(encoder_hid_dim, cross_attention_dim, bias=False)

self.projection_norm = nn.LayerNorm(cross_attention_dim)

def forward(self, x):

x = self.projection_linear(x)

x = self.projection_norm(x)

return x

class Kandinsky3UNet(ModelMixin, AttentionMixin, ConfigMixin):

@register_to_config

def __init__(

self,

in_channels: int = 4,

time_embedding_dim: int = 1536,

groups: int = 32,

attention_head_dim: int = 64,

layers_per_block: Union[int, Tuple[int]] = 3,

block_out_channels: Tuple[int, ...] = (384, 768, 1536, 3072),

cross_attention_dim: Union[int, Tuple[int]] = 4096,

encoder_hid_dim: int = 4096,

):

super().__init__()

# TODO(Yiyi): Give better name and put into config for the following 4 parameters

expansion_ratio = 4

compression_ratio = 2

add_cross_attention = (False, True, True, True)

add_self_attention = (False, True, True, True)

out_channels = in_channels

init_channels = block_out_channels[0] // 2

self.time_proj = Timesteps(init_channels, flip_sin_to_cos=False, downscale_freq_shift=1)

self.time_embedding = TimestepEmbedding(

init_channels,

time_embedding_dim,

)

self.add_time_condition = Kandinsky3AttentionPooling(

time_embedding_dim, cross_attention_dim, attention_head_dim

)

self.conv_in = nn.Conv2d(in_channels, init_channels, kernel_size=3, padding=1)

self.encoder_hid_proj = Kandinsky3EncoderProj(encoder_hid_dim, cross_attention_dim)

hidden_dims = [init_channels] + list(block_out_channels)

in_out_dims = list(zip(hidden_dims[:-1], hidden_dims[1:]))

text_dims = [cross_attention_dim if is_exist else None for is_exist in add_cross_attention]

num_blocks = len(block_out_channels) * [layers_per_block]

layer_params = [num_blocks, text_dims, add_self_attention]

rev_layer_params = map(reversed, layer_params)

cat_dims = []

self.num_levels = len(in_out_dims)

self.down_blocks = nn.ModuleList([])

for level, ((in_dim, out_dim), res_block_num, text_dim, self_attention) in enumerate(

zip(in_out_dims, *layer_params)

):

down_sample = level != (self.num_levels - 1)

cat_dims.append(out_dim if level != (self.num_levels - 1) else 0)

self.down_blocks.append(

Kandinsky3DownSampleBlock(

in_dim,

out_dim,

time_embedding_dim,

text_dim,

res_block_num,

groups,

attention_head_dim,

expansion_ratio,

compression_ratio,

down_sample,

self_attention,

)

)

self.up_blocks = nn.ModuleList([])

for level, ((out_dim, in_dim), res_block_num, text_dim, self_attention) in enumerate(

zip(reversed(in_out_dims), *rev_layer_params)

):

up_sample = level != 0

self.up_blocks.append(

Kandinsky3UpSampleBlock(

in_dim,

cat_dims.pop(),

out_dim,

time_embedding_dim,

text_dim,

res_block_num,

groups,

attention_head_dim,

expansion_ratio,

compression_ratio,

up_sample,

self_attention,

)

)

self.conv_norm_out = nn.GroupNorm(groups, init_channels)

self.conv_act_out = nn.SiLU()

self.conv_out = nn.Conv2d(init_channels, out_channels, kernel_size=3, padding=1)

def set_default_attn_processor(self):

"""

Disables custom attention processors and sets the default attention implementation.

"""

self.set_attn_processor(AttnProcessor())

def forward(self, sample, timestep, encoder_hidden_states=None, encoder_attention_mask=None, return_dict=True):

if encoder_attention_mask is not None:

encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0

encoder_attention_mask = encoder_attention_mask.unsqueeze(1)

if not torch.is_tensor(timestep):

dtype = torch.float32 if isinstance(timestep, float) else torch.int32

timestep = torch.tensor([timestep], dtype=dtype, device=sample.device)

elif len(timestep.shape) == 0:

timestep = timestep[None].to(sample.device)

# broadcast to batch dimension in a way that's compatible with ONNX/Core ML

timestep = timestep.expand(sample.shape[0])

time_embed_input = self.time_proj(timestep).to(sample.dtype)

time_embed = self.time_embedding(time_embed_input)

encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)

if encoder_hidden_states is not None:

time_embed = self.add_time_condition(time_embed, encoder_hidden_states, encoder_attention_mask)

hidden_states = []

sample = self.conv_in(sample)

for level, down_sample in enumerate(self.down_blocks):

sample = down_sample(sample, time_embed, encoder_hidden_states, encoder_attention_mask)

if level != self.num_levels - 1:

hidden_states.append(sample)

for level, up_sample in enumerate(self.up_blocks):

if level != 0:

sample = torch.cat([sample, hidden_states.pop()], dim=1)

sample = up_sample(sample, time_embed, encoder_hidden_states, encoder_attention_mask)

sample = self.conv_norm_out(sample)

sample = self.conv_act_out(sample)

sample = self.conv_out(sample)

if not return_dict:

return (sample,)

return Kandinsky3UNetOutput(sample=sample)

class Kandinsky3UpSampleBlock(nn.Module):

def __init__(

self,

in_channels,

cat_dim,

out_channels,

time_embed_dim,

context_dim=None,

num_blocks=3,

groups=32,

head_dim=64,

expansion_ratio=4,

compression_ratio=2,

up_sample=True,

self_attention=True,

):

super().__init__()

up_resolutions = [[None, True if up_sample else None, None, None]] + [[None] * 4] * (num_blocks - 1)

hidden_channels = (

[(in_channels + cat_dim, in_channels)]

+ [(in_channels, in_channels)] * (num_blocks - 2)

+ [(in_channels, out_channels)]

)

attentions = []

resnets_in = []

resnets_out = []

self.self_attention = self_attention

self.context_dim = context_dim

if self_attention:

attentions.append(

Kandinsky3AttentionBlock(out_channels, time_embed_dim, None, groups, head_dim, expansion_ratio)

)

else:

attentions.append(nn.Identity())

for (in_channel, out_channel), up_resolution in zip(hidden_channels, up_resolutions):

resnets_in.append(

Kandinsky3ResNetBlock(in_channel, in_channel, time_embed_dim, groups, compression_ratio, up_resolution)

)

if context_dim is not None:

attentions.append(

Kandinsky3AttentionBlock(

in_channel, time_embed_dim, context_dim, groups, head_dim, expansion_ratio

)

)

else:

attentions.append(nn.Identity())

resnets_out.append(

Kandinsky3ResNetBlock(in_channel, out_channel, time_embed_dim, groups, compression_ratio)

)

self.attentions = nn.ModuleList(attentions)

self.resnets_in = nn.ModuleList(resnets_in)

self.resnets_out = nn.ModuleList(resnets_out)

def forward(self, x, time_embed, context=None, context_mask=None, image_mask=None):

for attention, resnet_in, resnet_out in zip(self.attentions[1:], self.resnets_in, self.resnets_out):

x = resnet_in(x, time_embed)

if self.context_dim is not None:

x = attention(x, time_embed, context, context_mask, image_mask)

x = resnet_out(x, time_embed)

if self.self_attention:

x = self.attentions[0](x, time_embed, image_mask=image_mask)

return x

class Kandinsky3DownSampleBlock(nn.Module):

def __init__(

self,

in_channels,

out_channels,

time_embed_dim,

context_dim=None,

num_blocks=3,

groups=32,

head_dim=64,

expansion_ratio=4,

compression_ratio=2,

down_sample=True,

self_attention=True,

):

super().__init__()

attentions = []

resnets_in = []

resnets_out = []

self.self_attention = self_attention

self.context_dim = context_dim

if self_attention:

attentions.append(

Kandinsky3AttentionBlock(in_channels, time_embed_dim, None, groups, head_dim, expansion_ratio)

)

else:

attentions.append(nn.Identity())

up_resolutions = [[None] * 4] * (num_blocks - 1) + [[None, None, False if down_sample else None, None]]

hidden_channels = [(in_channels, out_channels)] + [(out_channels, out_channels)] * (num_blocks - 1)

for (in_channel, out_channel), up_resolution in zip(hidden_channels, up_resolutions):

resnets_in.append(

Kandinsky3ResNetBlock(in_channel, out_channel, time_embed_dim, groups, compression_ratio)

)

if context_dim is not None:

attentions.append(

Kandinsky3AttentionBlock(

out_channel, time_embed_dim, context_dim, groups, head_dim, expansion_ratio

)

)

else:

attentions.append(nn.Identity())

resnets_out.append(

Kandinsky3ResNetBlock(

out_channel, out_channel, time_embed_dim, groups, compression_ratio, up_resolution

)

)

self.attentions = nn.ModuleList(attentions)

self.resnets_in = nn.ModuleList(resnets_in)

self.resnets_out = nn.ModuleList(resnets_out)

def forward(self, x, time_embed, context=None, context_mask=None, image_mask=None):

if self.self_attention:

x = self.attentions[0](x, time_embed, image_mask=image_mask)

for attention, resnet_in, resnet_out in zip(self.attentions[1:], self.resnets_in, self.resnets_out):

x = resnet_in(x, time_embed)

if self.context_dim is not None:

x = attention(x, time_embed, context, context_mask, image_mask)

x = resnet_out(x, time_embed)

return x

class Kandinsky3ConditionalGroupNorm(nn.Module):

def __init__(self, groups, normalized_shape, context_dim):

super().__init__()

self.norm = nn.GroupNorm(groups, normalized_shape, affine=False)

self.context_mlp = nn.Sequential(nn.SiLU(), nn.Linear(context_dim, 2 * normalized_shape))

self.context_mlp[1].weight.data.zero_()

self.context_mlp[1].bias.data.zero_()

def forward(self, x, context):

context = self.context_mlp(context)

for _ in range(len(x.shape[2:])):

context = context.unsqueeze(-1)

scale, shift = context.chunk(2, dim=1)

x = self.norm(x) * (scale + 1.0) + shift

return x

class Kandinsky3Block(nn.Module):

def __init__(self, in_channels, out_channels, time_embed_dim, kernel_size=3, norm_groups=32, up_resolution=None):

super().__init__()

self.group_norm = Kandinsky3ConditionalGroupNorm(norm_groups, in_channels, time_embed_dim)

self.activation = nn.SiLU()

if up_resolution is not None and up_resolution:

self.up_sample = nn.ConvTranspose2d(in_channels, in_channels, kernel_size=2, stride=2)

else:

self.up_sample = nn.Identity()

padding = int(kernel_size > 1)

self.projection = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=padding)

if up_resolution is not None and not up_resolution:

self.down_sample = nn.Conv2d(out_channels, out_channels, kernel_size=2, stride=2)

else:

self.down_sample = nn.Identity()

def forward(self, x, time_embed):

x = self.group_norm(x, time_embed)

x = self.activation(x)

x = self.up_sample(x)

x = self.projection(x)

x = self.down_sample(x)

return x

class Kandinsky3ResNetBlock(nn.Module):

def __init__(

self, in_channels, out_channels, time_embed_dim, norm_groups=32, compression_ratio=2, up_resolutions=4 * [None]

):

super().__init__()

kernel_sizes = [1, 3, 3, 1]

hidden_channel = max(in_channels, out_channels) // compression_ratio

hidden_channels = (

[(in_channels, hidden_channel)] + [(hidden_channel, hidden_channel)] * 2 + [(hidden_channel, out_channels)]

)

self.resnet_blocks = nn.ModuleList(

[

Kandinsky3Block(in_channel, out_channel, time_embed_dim, kernel_size, norm_groups, up_resolution)

for (in_channel, out_channel), kernel_size, up_resolution in zip(

hidden_channels, kernel_sizes, up_resolutions

)

]

)

self.shortcut_up_sample = (

nn.ConvTranspose2d(in_channels, in_channels, kernel_size=2, stride=2)

if True in up_resolutions

else nn.Identity()

)

self.shortcut_projection = (

nn.Conv2d(in_channels, out_channels, kernel_size=1) if in_channels != out_channels else nn.Identity()

)

self.shortcut_down_sample = (

nn.Conv2d(out_channels, out_channels, kernel_size=2, stride=2)

if False in up_resolutions

else nn.Identity()

)

def forward(self, x, time_embed):

out = x

for resnet_block in self.resnet_blocks:

out = resnet_block(out, time_embed)

x = self.shortcut_up_sample(x)

x = self.shortcut_projection(x)

x = self.shortcut_down_sample(x)

x = x + out

return x

class Kandinsky3AttentionPooling(nn.Module):

def __init__(self, num_channels, context_dim, head_dim=64):

super().__init__()

self.attention = Attention(

context_dim,

context_dim,

dim_head=head_dim,

out_dim=num_channels,

out_bias=False,

)

def forward(self, x, context, context_mask=None):

context_mask = context_mask.to(dtype=context.dtype)

context = self.attention(context.mean(dim=1, keepdim=True), context, context_mask)

return x + context.squeeze(1)

class Kandinsky3AttentionBlock(nn.Module):

def __init__(self, num_channels, time_embed_dim, context_dim=None, norm_groups=32, head_dim=64, expansion_ratio=4):

super().__init__()

self.in_norm = Kandinsky3ConditionalGroupNorm(norm_groups, num_channels, time_embed_dim)

self.attention = Attention(

num_channels,

context_dim or num_channels,

dim_head=head_dim,

out_dim=num_channels,

out_bias=False,

)

hidden_channels = expansion_ratio * num_channels

self.out_norm = Kandinsky3ConditionalGroupNorm(norm_groups, num_channels, time_embed_dim)

self.feed_forward = nn.Sequential(

nn.Conv2d(num_channels, hidden_channels, kernel_size=1, bias=False),

nn.SiLU(),

nn.Conv2d(hidden_channels, num_channels, kernel_size=1, bias=False),

)

def forward(self, x, time_embed, context=None, context_mask=None, image_mask=None):

height, width = x.shape[-2:]

out = self.in_norm(x, time_embed)

out = out.reshape(x.shape[0], -1, height * width).permute(0, 2, 1)

context = context if context is not None else out

if context_mask is not None:

context_mask = context_mask.to(dtype=context.dtype)

out = self.attention(out, context, context_mask)

out = out.permute(0, 2, 1).unsqueeze(-1).reshape(out.shape[0], -1, height, width)

x = x + out

out = self.out_norm(x, time_embed)

out = self.feed_forward(out)

x = x + out

return x