# Copyright (c) 2026 ByteDance Ltd. and/or its affiliates. # # 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. """ DreamLite UNet model and its constituent 2D blocks. This single file mirrors the structure used by recent diffusers transformer model files: it defines all DreamLite building blocks (Down / Mid / Up) and the top-level :class:`DreamLiteUNetModel` together. Compared to the upstream ``unet_2d_blocks`` Down/Mid/Up cross-attention blocks, the DreamLite variants additionally thread the following knobs: - ``use_sep_conv``: replace standard convs in :class:`ResnetBlock2DDreamLite` with depthwise-separable convs (mobile-friendly). - ``qk_norm``, ``num_kv_heads``, ``ff_mult``: propagated into :class:`DreamLiteTransformer2DModel` / :class:`BasicTransformerBlockDreamLite`. The two "no self-attention" variants hard-code ``use_self_attention=False`` in their :class:`DreamLiteTransformer2DModel` calls. The U-Net itself defaults its attention processors to :class:`DreamLiteAttnProcessor2_0` (GQA-aware SDPA), which is required because the upstream ``AttnProcessor2_0`` does not handle ``kv_heads != heads`` correctly. """ from __future__ import annotations from functools import partial from typing import Any, Optional import torch import torch.nn.functional as F from torch import nn from ...configuration_utils import register_to_config from ..activations import get_activation from ..attention_dispatch import dispatch_attention_fn from ..attention_processor import Attention from ..downsampling import Downsample2D as _CoreDownsample2D from ..downsampling import downsample_2d from ..modeling_utils import ModelMixin from ..normalization import RMSNorm from ..transformers.dual_transformer_2d import DualTransformer2DModel from ..transformers.transformer_2d_dreamlite import DreamLiteTransformer2DModel from ..upsampling import Upsample2D as _CoreUpsample2D from ..upsampling import upsample_2d from .unet_2d_blocks import Downsample2D, Upsample2D, apply_freeu from .unet_2d_condition import UNet2DConditionModel # --------------------------------------------------------------------------- # Building blocks (resnet + attention processor) # --------------------------------------------------------------------------- class DepthwiseSeparableConv(nn.Module): """ Depthwise separable convolution used by DreamLite mobile-friendly ResNet blocks. A depthwise convolution (groups == in_channels) followed by a 1x1 pointwise convolution. The pointwise output channel count is multiplied by `expand_ratio` to support inverted-residual style expansion / contraction inside [`ResnetBlock2DDreamLite`]. """ def __init__( self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, bias: bool = False, expand_ratio: float = 1, ): super().__init__() self.depthwise = nn.Conv2d( in_channels, in_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=in_channels, bias=bias, ) self.pointwise = nn.Conv2d(in_channels, int(out_channels * expand_ratio), kernel_size=1, bias=bias) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.depthwise(hidden_states) hidden_states = self.pointwise(hidden_states) return hidden_states class ResnetBlock2DDreamLite(nn.Module): r""" A ResNet block used by DreamLite. Mirrors [`diffusers.models.resnet.ResnetBlock2D`] with one extra option: use_sep_conv (`bool`, *optional*, defaults to `False`): Replace the two 3x3 convolutions with [`DepthwiseSeparableConv`]. The first conv expands the channel count by 2x; the second conv contracts it back. Used by the mobile-friendly DreamLite checkpoints. All other parameters behave identically to [`diffusers.models.resnet.ResnetBlock2D`]. """ def __init__( self, *, in_channels: int, out_channels: Optional[int] = None, conv_shortcut: bool = False, dropout: float = 0.0, temb_channels: int = 512, groups: int = 32, groups_out: Optional[int] = None, pre_norm: bool = True, eps: float = 1e-6, non_linearity: str = "swish", skip_time_act: bool = False, time_embedding_norm: str = "default", kernel: Optional[torch.Tensor] = None, output_scale_factor: float = 1.0, use_in_shortcut: Optional[bool] = None, up: bool = False, down: bool = False, conv_shortcut_bias: bool = True, conv_2d_out_channels: Optional[int] = None, use_sep_conv: bool = False, ): super().__init__() if time_embedding_norm in ("ada_group", "spatial"): raise ValueError( f"`time_embedding_norm`={time_embedding_norm!r} is not supported by `ResnetBlock2DDreamLite`. " "Use `diffusers.models.resnet.ResnetBlockCondNorm2D` instead." ) self.pre_norm = True 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.up = up self.down = down self.output_scale_factor = output_scale_factor self.time_embedding_norm = time_embedding_norm self.skip_time_act = skip_time_act if groups_out is None: groups_out = groups self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) # Inverted-residual style expansion when `use_sep_conv=True`: conv1 expands channels by 2x, # conv2 contracts them back. For the standard branch this is just a regular 3x3 conv. if use_sep_conv: expand_ratio = 2 self.conv1 = DepthwiseSeparableConv( in_channels, out_channels, kernel_size=3, stride=1, padding=1, expand_ratio=expand_ratio ) out_channels = out_channels * expand_ratio else: expand_ratio = 1 self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) if temb_channels is not None: if self.time_embedding_norm == "default": self.time_emb_proj = nn.Linear(temb_channels, out_channels) elif self.time_embedding_norm == "scale_shift": self.time_emb_proj = nn.Linear(temb_channels, 2 * out_channels) else: raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm}") else: self.time_emb_proj = None self.norm2 = nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True) self.dropout = nn.Dropout(dropout) conv_2d_out_channels = conv_2d_out_channels or out_channels if use_sep_conv: self.conv2 = DepthwiseSeparableConv( out_channels, conv_2d_out_channels, kernel_size=3, stride=1, padding=1, expand_ratio=1 / expand_ratio, ) conv_2d_out_channels = conv_2d_out_channels // expand_ratio else: self.conv2 = nn.Conv2d(out_channels, conv_2d_out_channels, kernel_size=3, stride=1, padding=1) self.nonlinearity = get_activation(non_linearity) self.upsample = self.downsample = None if self.up: if kernel == "fir": fir_kernel = (1, 3, 3, 1) self.upsample = lambda x: upsample_2d(x, kernel=fir_kernel) elif kernel == "sde_vp": self.upsample = partial(F.interpolate, scale_factor=2.0, mode="nearest") else: self.upsample = _CoreUpsample2D(in_channels, use_conv=False) elif self.down: if kernel == "fir": fir_kernel = (1, 3, 3, 1) self.downsample = lambda x: downsample_2d(x, kernel=fir_kernel) elif kernel == "sde_vp": self.downsample = partial(F.avg_pool2d, kernel_size=2, stride=2) else: self.downsample = _CoreDownsample2D(in_channels, use_conv=False, padding=1, name="op") self.use_in_shortcut = self.in_channels != conv_2d_out_channels if use_in_shortcut is None else use_in_shortcut self.conv_shortcut = None if self.use_in_shortcut: self.conv_shortcut = nn.Conv2d( in_channels, conv_2d_out_channels, kernel_size=1, stride=1, padding=0, bias=conv_shortcut_bias, ) def forward(self, input_tensor: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: hidden_states = input_tensor hidden_states = self.norm1(hidden_states) hidden_states = self.nonlinearity(hidden_states) if self.upsample is not None: # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 if hidden_states.shape[0] >= 64: input_tensor = input_tensor.contiguous() hidden_states = hidden_states.contiguous() input_tensor = self.upsample(input_tensor) hidden_states = self.upsample(hidden_states) elif self.downsample is not None: input_tensor = self.downsample(input_tensor) hidden_states = self.downsample(hidden_states) hidden_states = self.conv1(hidden_states) if self.time_emb_proj is not None: if not self.skip_time_act: temb = self.nonlinearity(temb) temb = self.time_emb_proj(temb)[:, :, None, None] if self.time_embedding_norm == "default": if temb is not None: hidden_states = hidden_states + temb hidden_states = self.norm2(hidden_states) elif self.time_embedding_norm == "scale_shift": if temb is None: raise ValueError(f"`temb` should not be None when `time_embedding_norm` is {self.time_embedding_norm}") time_scale, time_shift = torch.chunk(temb, 2, dim=1) hidden_states = self.norm2(hidden_states) hidden_states = hidden_states * (1 + time_scale) + time_shift else: hidden_states = self.norm2(hidden_states) hidden_states = self.nonlinearity(hidden_states) hidden_states = self.dropout(hidden_states) hidden_states = self.conv2(hidden_states) if self.conv_shortcut is not None: # Only call .contiguous() under training, to avoid DDP gradient-stride warnings while keeping # inference fast (especially on CPU). Mirrors the upstream fix from huggingface/diffusers#12975. if self.training: input_tensor = input_tensor.contiguous() input_tensor = self.conv_shortcut(input_tensor) output_tensor = (input_tensor + hidden_states) / self.output_scale_factor return output_tensor class DreamLiteAttnProcessor2_0: r""" Processor for implementing scaled dot-product attention with Grouped Query Attention (GQA / MQA) support. Identical to :class:`AttnProcessor2_0` except the key/value reshape branch correctly handles ``attn.kv_heads != attn.heads`` by reshaping K/V to ``kv_heads`` and then ``repeat_interleave``-ing them up to ``attn.heads``. This is required by the DreamLite UNet, which combines GQA with ``qk_norm`` — a combination the default :class:`AttnProcessor2_0` does not handle. SDPA is delegated to :func:`dispatch_attention_fn` so any of the diffusers attention backends (native PyTorch SDPA, FlashAttention, etc.) can be used. """ _attention_backend = None _parallel_config = None def __call__( self, attn: Attention, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, temb: torch.Tensor | None = None, ) -> torch.Tensor: residual = hidden_states if attn.spatial_norm is not None: hidden_states = attn.spatial_norm(hidden_states, temb) input_ndim = hidden_states.ndim if input_ndim == 4: batch_size, channel, height, width = hidden_states.shape hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) batch_size, sequence_length, _ = ( hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape ) if attention_mask is not None: attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1]) if attn.group_norm is not None: hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) query = attn.to_q(hidden_states) if encoder_hidden_states is None: encoder_hidden_states = hidden_states elif attn.norm_cross: encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) key = attn.to_k(encoder_hidden_states) value = attn.to_v(encoder_hidden_states) # --- GQA-aware reshape (the only real difference vs AttnProcessor2_0) --- # ``dispatch_attention_fn`` expects (batch, seq, heads, head_dim) — keep Q/K/V in that layout # and let the dispatched backend handle the transpose internally. head_dim = query.shape[-1] // attn.heads kv_heads = key.shape[-1] // head_dim query = query.view(batch_size, -1, attn.heads, head_dim) key = key.view(batch_size, -1, kv_heads, head_dim) value = value.view(batch_size, -1, kv_heads, head_dim) if attn.norm_q is not None: query = attn.norm_q(query) if attn.norm_k is not None: key = attn.norm_k(key) if kv_heads != attn.heads: # GQA / MQA: repeat K/V heads up to query heads for SDPA. heads_per_kv_head = attn.heads // kv_heads key = torch.repeat_interleave(key, heads_per_kv_head, dim=2, output_size=key.shape[2] * heads_per_kv_head) value = torch.repeat_interleave( value, heads_per_kv_head, dim=2, output_size=value.shape[2] * heads_per_kv_head ) # ------------------------------------------------------------------------ # the output of sdp = (batch, seq_len, num_heads, head_dim) hidden_states = dispatch_attention_fn( query, key, value, attn_mask=attention_mask, backend=self._attention_backend, parallel_config=self._parallel_config, ) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) # linear proj hidden_states = attn.to_out[0](hidden_states) # dropout hidden_states = attn.to_out[1](hidden_states) if input_ndim == 4: hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) if attn.residual_connection: hidden_states = hidden_states + residual hidden_states = hidden_states / attn.rescale_output_factor return hidden_states # --------------------------------------------------------------------------- # Mid block # --------------------------------------------------------------------------- class DreamLiteUNetMidBlock2DCrossAttn(nn.Module): def __init__( self, in_channels: int, temb_channels: int, out_channels: int | None = None, dropout: float = 0.0, num_layers: int = 1, transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", resnet_groups: int = 32, resnet_groups_out: int | None = None, resnet_pre_norm: bool = True, num_attention_heads: int = 1, output_scale_factor: float = 1.0, cross_attention_dim: int = 1280, dual_cross_attention: bool = False, use_linear_projection: bool = False, upcast_attention: bool = False, attention_type: str = "default", # DreamLite extras qk_norm: str | None = None, use_sep_conv: bool = False, ff_mult: int = 4, num_kv_heads: int | None = None, num_mid_layers: int = 1, ): super().__init__() out_channels = out_channels or in_channels self.in_channels = in_channels self.out_channels = out_channels self.has_cross_attention = True self.num_attention_heads = num_attention_heads resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) if isinstance(transformer_layers_per_block, int): transformer_layers_per_block = [transformer_layers_per_block] * num_layers resnet_groups_out = resnet_groups_out or resnet_groups resnets = [ ResnetBlock2DDreamLite( in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, groups_out=resnet_groups_out, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, use_sep_conv=use_sep_conv, ) ] attentions = [] for i in range(num_layers): if not dual_cross_attention: attentions.append( DreamLiteTransformer2DModel( num_attention_heads, out_channels // num_attention_heads, in_channels=out_channels, num_layers=transformer_layers_per_block[i], cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups_out, use_linear_projection=use_linear_projection, upcast_attention=upcast_attention, attention_type=attention_type, qk_norm=qk_norm, ff_mult=ff_mult, num_kv_heads=num_kv_heads, ) ) else: attentions.append( DualTransformer2DModel( num_attention_heads, out_channels // num_attention_heads, in_channels=out_channels, num_layers=1, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, ) ) resnets.append( ResnetBlock2DDreamLite( in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups_out, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, use_sep_conv=use_sep_conv, ) ) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, cross_attention_kwargs: dict[str, Any] | None = None, encoder_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: hidden_states = self.resnets[0](hidden_states, temb) for attn, resnet in zip(self.attentions, self.resnets[1:]): if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = attn( hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False, )[0] hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) else: hidden_states = attn( hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False, )[0] hidden_states = resnet(hidden_states, temb) return hidden_states # --------------------------------------------------------------------------- # Down blocks # --------------------------------------------------------------------------- class DreamLiteCrossAttnDownBlock2D(nn.Module): """DreamLite down block with both self- and cross-attention in each transformer layer.""" def __init__( self, in_channels: int, out_channels: int, temb_channels: int, dropout: float = 0.0, num_layers: int = 1, transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", resnet_groups: int = 32, resnet_pre_norm: bool = True, num_attention_heads: int = 1, cross_attention_dim: int = 1280, output_scale_factor: float = 1.0, downsample_padding: int = 1, add_downsample: bool = True, dual_cross_attention: bool = False, use_linear_projection: bool = False, only_cross_attention: bool = False, upcast_attention: bool = False, attention_type: str = "default", # DreamLite extras qk_norm: str | None = None, use_sep_conv: bool = False, ff_mult: int = 4, num_kv_heads: int | None = None, ): super().__init__() resnets = [] attentions = [] self.has_cross_attention = True self.num_attention_heads = num_attention_heads if isinstance(transformer_layers_per_block, int): transformer_layers_per_block = [transformer_layers_per_block] * num_layers for i in range(num_layers): in_ch = in_channels if i == 0 else out_channels resnets.append( ResnetBlock2DDreamLite( in_channels=in_ch, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, use_sep_conv=use_sep_conv, ) ) if not dual_cross_attention: attentions.append( DreamLiteTransformer2DModel( num_attention_heads=num_attention_heads, attention_head_dim=out_channels // num_attention_heads, in_channels=out_channels, num_layers=transformer_layers_per_block[i], cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention, upcast_attention=upcast_attention, attention_type=attention_type, qk_norm=qk_norm, ff_mult=ff_mult, num_kv_heads=num_kv_heads, ) ) else: attentions.append( DualTransformer2DModel( num_attention_heads, out_channels // num_attention_heads, in_channels=out_channels, num_layers=1, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, ) ) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) if add_downsample: self.downsamplers = nn.ModuleList( [ Downsample2D( out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op", ) ] ) else: self.downsamplers = None self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, cross_attention_kwargs: dict[str, Any] | None = None, encoder_attention_mask: torch.Tensor | None = None, additional_residuals: torch.Tensor | None = None, ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: output_states: tuple[torch.Tensor, ...] = () blocks = list(zip(self.resnets, self.attentions)) for i, (resnet, attn) in enumerate(blocks): if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) else: hidden_states = resnet(hidden_states, temb) hidden_states = attn( hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False, )[0] if i == len(blocks) - 1 and additional_residuals is not None: hidden_states = hidden_states + additional_residuals output_states = output_states + (hidden_states,) if self.downsamplers is not None: for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states) output_states = output_states + (hidden_states,) return hidden_states, output_states class DreamLiteCrossAttnNoSelfAttnDownBlock2D(nn.Module): """DreamLite down block with cross-attention only (self-attention is removed).""" def __init__( self, in_channels: int, out_channels: int, temb_channels: int, dropout: float = 0.0, num_layers: int = 1, transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", resnet_groups: int = 32, resnet_pre_norm: bool = True, num_attention_heads: int = 1, cross_attention_dim: int = 1280, output_scale_factor: float = 1.0, downsample_padding: int = 1, add_downsample: bool = True, dual_cross_attention: bool = False, use_linear_projection: bool = False, only_cross_attention: bool = False, upcast_attention: bool = False, attention_type: str = "default", # DreamLite extras qk_norm: str | None = None, use_sep_conv: bool = False, ff_mult: int = 4, num_kv_heads: int | None = None, ): super().__init__() resnets = [] attentions = [] self.has_cross_attention = True self.num_attention_heads = num_attention_heads if isinstance(transformer_layers_per_block, int): transformer_layers_per_block = [transformer_layers_per_block] * num_layers for i in range(num_layers): in_ch = in_channels if i == 0 else out_channels resnets.append( ResnetBlock2DDreamLite( in_channels=in_ch, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, use_sep_conv=use_sep_conv, ) ) if not dual_cross_attention: attentions.append( DreamLiteTransformer2DModel( num_attention_heads=num_attention_heads, attention_head_dim=out_channels // num_attention_heads, in_channels=out_channels, num_layers=transformer_layers_per_block[i], cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention, upcast_attention=upcast_attention, attention_type=attention_type, qk_norm=qk_norm, ff_mult=ff_mult, num_kv_heads=num_kv_heads, # DreamLite "remove self-attention" path: use_self_attention=False, ) ) else: attentions.append( DualTransformer2DModel( num_attention_heads, out_channels // num_attention_heads, in_channels=out_channels, num_layers=1, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, ) ) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) if add_downsample: self.downsamplers = nn.ModuleList( [ Downsample2D( out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op", ) ] ) else: self.downsamplers = None self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, cross_attention_kwargs: dict[str, Any] | None = None, encoder_attention_mask: torch.Tensor | None = None, additional_residuals: torch.Tensor | None = None, ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: output_states: tuple[torch.Tensor, ...] = () blocks = list(zip(self.resnets, self.attentions)) for i, (resnet, attn) in enumerate(blocks): if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) else: hidden_states = resnet(hidden_states, temb) hidden_states = attn( hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False, )[0] if i == len(blocks) - 1 and additional_residuals is not None: hidden_states = hidden_states + additional_residuals output_states = output_states + (hidden_states,) if self.downsamplers is not None: for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states) output_states = output_states + (hidden_states,) return hidden_states, output_states class DreamLiteDownBlock2D(nn.Module): """DreamLite plain resnet-only down block (no attention).""" def __init__( self, in_channels: int, out_channels: int, temb_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", resnet_groups: int = 32, resnet_pre_norm: bool = True, output_scale_factor: float = 1.0, add_downsample: bool = True, downsample_padding: int = 1, use_sep_conv: bool = False, ): super().__init__() resnets = [] for i in range(num_layers): in_ch = in_channels if i == 0 else out_channels resnets.append( ResnetBlock2DDreamLite( in_channels=in_ch, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, use_sep_conv=use_sep_conv, ) ) self.resnets = nn.ModuleList(resnets) if add_downsample: self.downsamplers = nn.ModuleList( [ Downsample2D( out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op", ) ] ) else: self.downsamplers = None self.gradient_checkpointing = False def forward( self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None, **kwargs, ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]: output_states: tuple[torch.Tensor, ...] = () for resnet in self.resnets: if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) else: hidden_states = resnet(hidden_states, temb) output_states = output_states + (hidden_states,) if self.downsamplers is not None: for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states) output_states = output_states + (hidden_states,) return hidden_states, output_states # --------------------------------------------------------------------------- # Up blocks # --------------------------------------------------------------------------- class DreamLiteCrossAttnUpBlock2D(nn.Module): """DreamLite up block with both self- and cross-attention in each transformer layer.""" def __init__( self, in_channels: int, out_channels: int, prev_output_channel: int, temb_channels: int, resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", resnet_groups: int = 32, resnet_pre_norm: bool = True, num_attention_heads: int = 1, cross_attention_dim: int = 1280, output_scale_factor: float = 1.0, add_upsample: bool = True, dual_cross_attention: bool = False, use_linear_projection: bool = False, only_cross_attention: bool = False, upcast_attention: bool = False, attention_type: str = "default", # DreamLite extras qk_norm: str | None = None, use_sep_conv: bool = False, ff_mult: int = 4, num_kv_heads: int | None = None, ): super().__init__() resnets = [] attentions = [] self.has_cross_attention = True self.num_attention_heads = num_attention_heads if isinstance(transformer_layers_per_block, int): transformer_layers_per_block = [transformer_layers_per_block] * num_layers for i in range(num_layers): res_skip_channels = in_channels if (i == num_layers - 1) else out_channels resnet_in_channels = prev_output_channel if i == 0 else out_channels resnets.append( ResnetBlock2DDreamLite( in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, use_sep_conv=use_sep_conv, ) ) if not dual_cross_attention: attentions.append( DreamLiteTransformer2DModel( num_attention_heads=num_attention_heads, attention_head_dim=out_channels // num_attention_heads, in_channels=out_channels, num_layers=transformer_layers_per_block[i], cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention, upcast_attention=upcast_attention, attention_type=attention_type, qk_norm=qk_norm, ff_mult=ff_mult, num_kv_heads=num_kv_heads, ) ) else: attentions.append( DualTransformer2DModel( num_attention_heads, out_channels // num_attention_heads, in_channels=out_channels, num_layers=1, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, ) ) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) if add_upsample: self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) else: self.upsamplers = None self.gradient_checkpointing = False self.resolution_idx = resolution_idx def forward( self, hidden_states: torch.Tensor, res_hidden_states_tuple: tuple[torch.Tensor, ...], temb: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, cross_attention_kwargs: dict[str, Any] | None = None, upsample_size: int | None = None, attention_mask: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: is_freeu_enabled = ( getattr(self, "s1", None) and getattr(self, "s2", None) and getattr(self, "b1", None) and getattr(self, "b2", None) ) for resnet, attn in zip(self.resnets, self.attentions): res_hidden_states = res_hidden_states_tuple[-1] res_hidden_states_tuple = res_hidden_states_tuple[:-1] if is_freeu_enabled: hidden_states, res_hidden_states = apply_freeu( self.resolution_idx, hidden_states, res_hidden_states, s1=self.s1, s2=self.s2, b1=self.b1, b2=self.b2, ) hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) else: hidden_states = resnet(hidden_states, temb) hidden_states = attn( hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False, )[0] if self.upsamplers is not None: for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states, upsample_size) return hidden_states class DreamLiteCrossAttnNoSelfAttnUpBlock2D(nn.Module): """DreamLite up block with cross-attention only (self-attention is removed).""" def __init__( self, in_channels: int, out_channels: int, prev_output_channel: int, temb_channels: int, resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, transformer_layers_per_block: int | tuple[int] = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", resnet_groups: int = 32, resnet_pre_norm: bool = True, num_attention_heads: int = 1, cross_attention_dim: int = 1280, output_scale_factor: float = 1.0, add_upsample: bool = True, dual_cross_attention: bool = False, use_linear_projection: bool = False, only_cross_attention: bool = False, upcast_attention: bool = False, attention_type: str = "default", # DreamLite extras qk_norm: str | None = None, use_sep_conv: bool = False, ff_mult: int = 4, num_kv_heads: int | None = None, ): super().__init__() resnets = [] attentions = [] self.has_cross_attention = True self.num_attention_heads = num_attention_heads if isinstance(transformer_layers_per_block, int): transformer_layers_per_block = [transformer_layers_per_block] * num_layers for i in range(num_layers): res_skip_channels = in_channels if (i == num_layers - 1) else out_channels resnet_in_channels = prev_output_channel if i == 0 else out_channels resnets.append( ResnetBlock2DDreamLite( in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, use_sep_conv=use_sep_conv, ) ) if not dual_cross_attention: attentions.append( DreamLiteTransformer2DModel( num_attention_heads=num_attention_heads, attention_head_dim=out_channels // num_attention_heads, in_channels=out_channels, num_layers=transformer_layers_per_block[i], cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention, upcast_attention=upcast_attention, attention_type=attention_type, qk_norm=qk_norm, ff_mult=ff_mult, num_kv_heads=num_kv_heads, # DreamLite "remove self-attention" path: use_self_attention=False, ) ) else: attentions.append( DualTransformer2DModel( num_attention_heads, out_channels // num_attention_heads, in_channels=out_channels, num_layers=1, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, ) ) self.attentions = nn.ModuleList(attentions) self.resnets = nn.ModuleList(resnets) if add_upsample: self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) else: self.upsamplers = None self.gradient_checkpointing = False self.resolution_idx = resolution_idx def forward( self, hidden_states: torch.Tensor, res_hidden_states_tuple: tuple[torch.Tensor, ...], temb: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, cross_attention_kwargs: dict[str, Any] | None = None, upsample_size: int | None = None, attention_mask: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: is_freeu_enabled = ( getattr(self, "s1", None) and getattr(self, "s2", None) and getattr(self, "b1", None) and getattr(self, "b2", None) ) for resnet, attn in zip(self.resnets, self.attentions): res_hidden_states = res_hidden_states_tuple[-1] res_hidden_states_tuple = res_hidden_states_tuple[:-1] if is_freeu_enabled: hidden_states, res_hidden_states = apply_freeu( self.resolution_idx, hidden_states, res_hidden_states, s1=self.s1, s2=self.s2, b1=self.b1, b2=self.b2, ) hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) else: hidden_states = resnet(hidden_states, temb) hidden_states = attn( hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False, )[0] if self.upsamplers is not None: for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states, upsample_size) return hidden_states class DreamLiteUpBlock2D(nn.Module): """DreamLite plain resnet-only up block (no attention).""" def __init__( self, in_channels: int, prev_output_channel: int, out_channels: int, temb_channels: int, resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, resnet_time_scale_shift: str = "default", resnet_act_fn: str = "swish", resnet_groups: int = 32, resnet_pre_norm: bool = True, output_scale_factor: float = 1.0, add_upsample: bool = True, use_sep_conv: bool = False, ): super().__init__() resnets = [] for i in range(num_layers): res_skip_channels = in_channels if (i == num_layers - 1) else out_channels resnet_in_channels = prev_output_channel if i == 0 else out_channels resnets.append( ResnetBlock2DDreamLite( in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, use_sep_conv=use_sep_conv, ) ) self.resnets = nn.ModuleList(resnets) if add_upsample: self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) else: self.upsamplers = None self.gradient_checkpointing = False self.resolution_idx = resolution_idx def forward( self, hidden_states: torch.Tensor, res_hidden_states_tuple: tuple[torch.Tensor, ...], temb: torch.Tensor | None = None, upsample_size: int | None = None, **kwargs, ) -> torch.Tensor: is_freeu_enabled = ( getattr(self, "s1", None) and getattr(self, "s2", None) and getattr(self, "b1", None) and getattr(self, "b2", None) ) for resnet in self.resnets: res_hidden_states = res_hidden_states_tuple[-1] res_hidden_states_tuple = res_hidden_states_tuple[:-1] if is_freeu_enabled: hidden_states, res_hidden_states = apply_freeu( self.resolution_idx, hidden_states, res_hidden_states, s1=self.s1, s2=self.s2, b1=self.b1, b2=self.b2, ) hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func(resnet, hidden_states, temb) else: hidden_states = resnet(hidden_states, temb) if self.upsamplers is not None: for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states, upsample_size) return hidden_states # --------------------------------------------------------------------------- # Local block dispatch (DreamLite-only) # # The string ``down_block_type`` / ``up_block_type`` / ``mid_block_type`` keys # persisted in saved checkpoints' ``config.json`` usually mirror the Python class # names defined above. Some configs use upstream UNet block names instead. # --------------------------------------------------------------------------- _DREAMLITE_DOWN_BLOCK_ALIASES = { "CrossAttnDownRemoveSelfAttnBlock2D": "DreamLiteCrossAttnNoSelfAttnDownBlock2D", "CrossAttnDownBlock2D": "DreamLiteCrossAttnDownBlock2D", "DownBlock2D": "DreamLiteDownBlock2D", } _DREAMLITE_MID_BLOCK_ALIASES = { "UNetMidBlock2DCrossAttn": "DreamLiteUNetMidBlock2DCrossAttn", } _DREAMLITE_UP_BLOCK_ALIASES = { "CrossAttnUpRemoveSelfAttnBlock2D": "DreamLiteCrossAttnNoSelfAttnUpBlock2D", "CrossAttnUpRemoveSelfAttnBlock2DV1": "DreamLiteCrossAttnNoSelfAttnUpBlock2D", "CrossAttnUpBlock2D": "DreamLiteCrossAttnUpBlock2D", "UpBlock2D": "DreamLiteUpBlock2D", } def _get_down_block_dreamlite( down_block_type: str, *, num_layers, transformer_layers_per_block, in_channels, out_channels, temb_channels, add_downsample, resnet_eps, resnet_act_fn, resnet_groups, cross_attention_dim, num_attention_heads, downsample_padding, dual_cross_attention, use_linear_projection, only_cross_attention, upcast_attention, resnet_time_scale_shift, attention_type, dropout, qk_norm, use_sep_conv, ff_mult, num_kv_heads, ): down_block_type = _DREAMLITE_DOWN_BLOCK_ALIASES.get(down_block_type, down_block_type) if down_block_type == "DreamLiteDownBlock2D": return DreamLiteDownBlock2D( num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, dropout=dropout, add_downsample=add_downsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, downsample_padding=downsample_padding, resnet_time_scale_shift=resnet_time_scale_shift, use_sep_conv=use_sep_conv, ) if down_block_type in ("DreamLiteCrossAttnDownBlock2D", "DreamLiteCrossAttnNoSelfAttnDownBlock2D"): if cross_attention_dim is None: raise ValueError(f"cross_attention_dim must be specified for {down_block_type}") cls = ( DreamLiteCrossAttnDownBlock2D if down_block_type == "DreamLiteCrossAttnDownBlock2D" else DreamLiteCrossAttnNoSelfAttnDownBlock2D ) return cls( num_layers=num_layers, transformer_layers_per_block=transformer_layers_per_block, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, dropout=dropout, add_downsample=add_downsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, downsample_padding=downsample_padding, cross_attention_dim=cross_attention_dim, num_attention_heads=num_attention_heads, dual_cross_attention=dual_cross_attention, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention, upcast_attention=upcast_attention, attention_type=attention_type, qk_norm=qk_norm, use_sep_conv=use_sep_conv, ff_mult=ff_mult, num_kv_heads=num_kv_heads, ) raise ValueError(f"DreamLite does not support down_block_type={down_block_type!r}") def _get_mid_block_dreamlite( mid_block_type, *, temb_channels, in_channels, resnet_eps, resnet_act_fn, resnet_groups, output_scale_factor, transformer_layers_per_block, num_attention_heads, cross_attention_dim, dual_cross_attention, use_linear_projection, upcast_attention, resnet_time_scale_shift, attention_type, dropout, qk_norm, use_sep_conv, ff_mult, num_kv_heads, num_mid_layers=1, ): if mid_block_type is None: return None mid_block_type = _DREAMLITE_MID_BLOCK_ALIASES.get(mid_block_type, mid_block_type) if mid_block_type == "DreamLiteUNetMidBlock2DCrossAttn": return DreamLiteUNetMidBlock2DCrossAttn( transformer_layers_per_block=transformer_layers_per_block, in_channels=in_channels, temb_channels=temb_channels, dropout=dropout, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, output_scale_factor=output_scale_factor, resnet_time_scale_shift=resnet_time_scale_shift, cross_attention_dim=cross_attention_dim, num_attention_heads=num_attention_heads, resnet_groups=resnet_groups, dual_cross_attention=dual_cross_attention, use_linear_projection=use_linear_projection, upcast_attention=upcast_attention, attention_type=attention_type, qk_norm=qk_norm, use_sep_conv=use_sep_conv, ff_mult=ff_mult, num_kv_heads=num_kv_heads, num_layers=num_mid_layers, ) raise ValueError(f"DreamLite does not support mid_block_type={mid_block_type!r}") def _get_up_block_dreamlite( up_block_type, *, num_layers, transformer_layers_per_block, in_channels, out_channels, prev_output_channel, temb_channels, add_upsample, resnet_eps, resnet_act_fn, resolution_idx, resnet_groups, cross_attention_dim, num_attention_heads, dual_cross_attention, use_linear_projection, only_cross_attention, upcast_attention, resnet_time_scale_shift, attention_type, dropout, qk_norm, use_sep_conv, ff_mult, num_kv_heads, ): up_block_type = _DREAMLITE_UP_BLOCK_ALIASES.get(up_block_type, up_block_type) if up_block_type == "DreamLiteUpBlock2D": return DreamLiteUpBlock2D( num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, prev_output_channel=prev_output_channel, temb_channels=temb_channels, resolution_idx=resolution_idx, dropout=dropout, add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, resnet_time_scale_shift=resnet_time_scale_shift, use_sep_conv=use_sep_conv, ) if up_block_type in ("DreamLiteCrossAttnUpBlock2D", "DreamLiteCrossAttnNoSelfAttnUpBlock2D"): if cross_attention_dim is None: raise ValueError(f"cross_attention_dim must be specified for {up_block_type}") cls = ( DreamLiteCrossAttnUpBlock2D if up_block_type == "DreamLiteCrossAttnUpBlock2D" else DreamLiteCrossAttnNoSelfAttnUpBlock2D ) return cls( num_layers=num_layers, transformer_layers_per_block=transformer_layers_per_block, in_channels=in_channels, out_channels=out_channels, prev_output_channel=prev_output_channel, temb_channels=temb_channels, resolution_idx=resolution_idx, dropout=dropout, add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, cross_attention_dim=cross_attention_dim, num_attention_heads=num_attention_heads, dual_cross_attention=dual_cross_attention, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention, upcast_attention=upcast_attention, attention_type=attention_type, qk_norm=qk_norm, use_sep_conv=use_sep_conv, ff_mult=ff_mult, num_kv_heads=num_kv_heads, ) raise ValueError(f"DreamLite does not support up_block_type={up_block_type!r}") # --------------------------------------------------------------------------- # Model # --------------------------------------------------------------------------- class DreamLiteUNetModel(UNet2DConditionModel): r""" DreamLite variant of :class:`UNet2DConditionModel`. Differences vs the parent class: * Down / Mid / Up blocks are dispatched to the DreamLite variants defined above, which support depthwise-separable convolutions in resnets and Grouped Query Attention with RMSNorm ``qk_norm`` in attention. * ``default_attn_processor`` returns :class:`DreamLiteAttnProcessor2_0` so SDPA is GQA-aware out of the box. """ _supports_gradient_checkpointing = True _no_split_modules = [ "BasicTransformerBlockDreamLite", "ResnetBlock2DDreamLite", "DreamLiteCrossAttnUpBlock2D", "DreamLiteUpBlock2D", ] _repeated_blocks = ["BasicTransformerBlockDreamLite"] @register_to_config def __init__( self, sample_size: int | tuple[int, int] | None = None, in_channels: int = 4, out_channels: int = 4, center_input_sample: bool = False, flip_sin_to_cos: bool = True, freq_shift: int = 0, down_block_types: tuple[str, ...] = ( "DreamLiteCrossAttnNoSelfAttnDownBlock2D", "DreamLiteCrossAttnNoSelfAttnDownBlock2D", "DreamLiteCrossAttnDownBlock2D", ), mid_block_type: str | None = "DreamLiteUNetMidBlock2DCrossAttn", up_block_types: tuple[str, ...] = ( "DreamLiteCrossAttnUpBlock2D", "DreamLiteCrossAttnNoSelfAttnUpBlock2D", "DreamLiteUpBlock2D", ), only_cross_attention: bool | tuple[bool, ...] = False, block_out_channels: tuple[int, ...] = (320, 640, 1280), layers_per_block: int | tuple[int, ...] = 2, downsample_padding: int = 1, mid_block_scale_factor: float = 1, dropout: float = 0.0, act_fn: str = "silu", norm_num_groups: int | None = 32, norm_eps: float = 1e-5, cross_attention_dim: int | tuple[int, ...] = 2048, transformer_layers_per_block: int | tuple[int, ...] | tuple[tuple, ...] = 1, reverse_transformer_layers_per_block: tuple[tuple[int, ...], ...] | None = None, encoder_hid_dim: int | None = None, encoder_hid_dim_type: str | None = None, attention_head_dim: int | tuple[int, ...] = 64, num_attention_heads: int | tuple[int, ...] | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, class_embed_type: str | None = None, addition_embed_type: str | None = None, addition_time_embed_dim: int | None = None, num_class_embeds: int | None = None, upcast_attention: bool = False, resnet_time_scale_shift: str = "default", resnet_skip_time_act: bool = False, resnet_out_scale_factor: float = 1.0, time_embedding_type: str = "positional", time_embedding_dim: int | None = None, time_embedding_act_fn: str | None = None, timestep_post_act: str | None = None, time_cond_proj_dim: int | None = None, conv_in_kernel: int = 3, conv_out_kernel: int = 3, projection_class_embeddings_input_dim: int | None = None, attention_type: str = "default", class_embeddings_concat: bool = False, mid_block_only_cross_attention: bool | None = None, cross_attention_norm: str | None = None, addition_embed_type_num_heads: int = 64, # ---- DreamLite extras ---- qk_norm: str | None = "rms_norm", use_sep_conv: bool = True, ff_mult: int = 6, num_kv_heads: int | None = 1, num_mid_layers: int = 1, ): # NOTE: deliberately skip UNet2DConditionModel.__init__ because we replicate # the body with DreamLite block dispatch, but call ModelMixin.__init__ so that # mixin state (e.g. _gradient_checkpointing_func) is properly initialised. ModelMixin.__init__(self) self.sample_size = sample_size if num_attention_heads is not None: raise ValueError( "At the moment it is not possible to define the number of attention heads via " "`num_attention_heads` because of a naming issue as described in " "https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. " "Passing `num_attention_heads` will only be supported in diffusers v0.19." ) num_attention_heads = num_attention_heads or attention_head_dim # Reuse parent helpers (they only touch self, no super().__init__ required). self._check_config( down_block_types=down_block_types, up_block_types=up_block_types, only_cross_attention=only_cross_attention, block_out_channels=block_out_channels, layers_per_block=layers_per_block, cross_attention_dim=cross_attention_dim, transformer_layers_per_block=transformer_layers_per_block, reverse_transformer_layers_per_block=reverse_transformer_layers_per_block, attention_head_dim=attention_head_dim, num_attention_heads=num_attention_heads, ) self.projection_class_embeddings_input_dim = projection_class_embeddings_input_dim # input conv_in_padding = (conv_in_kernel - 1) // 2 self.conv_in = nn.Conv2d( in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding ) # time time_embed_dim, timestep_input_dim = self._set_time_proj( time_embedding_type, block_out_channels=block_out_channels, flip_sin_to_cos=flip_sin_to_cos, freq_shift=freq_shift, time_embedding_dim=time_embedding_dim, ) from ..embeddings import TimestepEmbedding # local import to avoid cycle self.time_embedding = TimestepEmbedding( timestep_input_dim, time_embed_dim, act_fn=act_fn, post_act_fn=timestep_post_act, cond_proj_dim=time_cond_proj_dim, ) self._set_encoder_hid_proj( encoder_hid_dim_type, cross_attention_dim=cross_attention_dim, encoder_hid_dim=encoder_hid_dim, ) self._set_class_embedding( class_embed_type, act_fn=act_fn, num_class_embeds=num_class_embeds, projection_class_embeddings_input_dim=projection_class_embeddings_input_dim, time_embed_dim=time_embed_dim, timestep_input_dim=timestep_input_dim, ) self._set_add_embedding( addition_embed_type, addition_embed_type_num_heads=addition_embed_type_num_heads, addition_time_embed_dim=addition_time_embed_dim, cross_attention_dim=cross_attention_dim, encoder_hid_dim=encoder_hid_dim, flip_sin_to_cos=flip_sin_to_cos, freq_shift=freq_shift, projection_class_embeddings_input_dim=projection_class_embeddings_input_dim, time_embed_dim=time_embed_dim, ) self.time_embed_act = None if time_embedding_act_fn is None else get_activation(time_embedding_act_fn) self.down_blocks = nn.ModuleList([]) self.up_blocks = nn.ModuleList([]) # Normalize per-stage args if isinstance(only_cross_attention, bool): if mid_block_only_cross_attention is None: mid_block_only_cross_attention = only_cross_attention only_cross_attention = [only_cross_attention] * len(down_block_types) if mid_block_only_cross_attention is None: mid_block_only_cross_attention = False if isinstance(num_attention_heads, int): num_attention_heads = (num_attention_heads,) * len(down_block_types) if isinstance(attention_head_dim, int): attention_head_dim = (attention_head_dim,) * len(down_block_types) if isinstance(cross_attention_dim, int): cross_attention_dim = (cross_attention_dim,) * len(down_block_types) if isinstance(layers_per_block, int): layers_per_block = [layers_per_block] * len(down_block_types) if isinstance(transformer_layers_per_block, int): transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types) blocks_time_embed_dim = time_embed_dim * 2 if class_embeddings_concat else time_embed_dim # ---- Down ---- output_channel = block_out_channels[0] for i, down_block_type in enumerate(down_block_types): input_channel = output_channel output_channel = block_out_channels[i] is_final_block = i == len(block_out_channels) - 1 self.down_blocks.append( _get_down_block_dreamlite( down_block_type, num_layers=layers_per_block[i], transformer_layers_per_block=transformer_layers_per_block[i], in_channels=input_channel, out_channels=output_channel, temb_channels=blocks_time_embed_dim, add_downsample=not is_final_block, resnet_eps=norm_eps, resnet_act_fn=act_fn, resnet_groups=norm_num_groups, cross_attention_dim=cross_attention_dim[i], num_attention_heads=num_attention_heads[i], downsample_padding=downsample_padding, dual_cross_attention=dual_cross_attention, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention[i], upcast_attention=upcast_attention, resnet_time_scale_shift=resnet_time_scale_shift, attention_type=attention_type, dropout=dropout, qk_norm=qk_norm, use_sep_conv=use_sep_conv, ff_mult=ff_mult, num_kv_heads=num_kv_heads, ) ) # ---- Mid ---- self.mid_block = _get_mid_block_dreamlite( mid_block_type, temb_channels=blocks_time_embed_dim, in_channels=block_out_channels[-1], resnet_eps=norm_eps, resnet_act_fn=act_fn, resnet_groups=norm_num_groups, output_scale_factor=mid_block_scale_factor, transformer_layers_per_block=transformer_layers_per_block[-1], num_attention_heads=num_attention_heads[-1], cross_attention_dim=cross_attention_dim[-1], dual_cross_attention=dual_cross_attention, use_linear_projection=use_linear_projection, upcast_attention=upcast_attention, resnet_time_scale_shift=resnet_time_scale_shift, attention_type=attention_type, dropout=dropout, qk_norm=qk_norm, use_sep_conv=use_sep_conv, ff_mult=ff_mult, num_kv_heads=num_kv_heads, num_mid_layers=num_mid_layers, ) # ---- Up ---- self.num_upsamplers = 0 reversed_block_out_channels = list(reversed(block_out_channels)) reversed_num_attention_heads = list(reversed(num_attention_heads)) reversed_layers_per_block = list(reversed(layers_per_block)) reversed_cross_attention_dim = list(reversed(cross_attention_dim)) reversed_transformer_layers_per_block = ( list(reversed(transformer_layers_per_block)) if reverse_transformer_layers_per_block is None else reverse_transformer_layers_per_block ) only_cross_attention = list(reversed(only_cross_attention)) output_channel = reversed_block_out_channels[0] for i, up_block_type in enumerate(up_block_types): is_final_block = i == len(block_out_channels) - 1 prev_output_channel = output_channel output_channel = reversed_block_out_channels[i] input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)] if not is_final_block: add_upsample = True self.num_upsamplers += 1 else: add_upsample = False self.up_blocks.append( _get_up_block_dreamlite( up_block_type, num_layers=reversed_layers_per_block[i] + 1, transformer_layers_per_block=reversed_transformer_layers_per_block[i], in_channels=input_channel, out_channels=output_channel, prev_output_channel=prev_output_channel, temb_channels=blocks_time_embed_dim, add_upsample=add_upsample, resnet_eps=norm_eps, resnet_act_fn=act_fn, resolution_idx=i, resnet_groups=norm_num_groups, cross_attention_dim=reversed_cross_attention_dim[i], num_attention_heads=reversed_num_attention_heads[i], dual_cross_attention=dual_cross_attention, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention[i], upcast_attention=upcast_attention, resnet_time_scale_shift=resnet_time_scale_shift, attention_type=attention_type, dropout=dropout, qk_norm=qk_norm, use_sep_conv=use_sep_conv, ff_mult=ff_mult, num_kv_heads=num_kv_heads, ) ) # ---- Out ---- if norm_num_groups is not None: self.conv_norm_out = nn.GroupNorm( num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps ) self.conv_act = get_activation(act_fn) else: self.conv_norm_out = None self.conv_act = None conv_out_padding = (conv_out_kernel - 1) // 2 self.conv_out = nn.Conv2d( block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding ) self._set_pos_net_if_use_gligen(attention_type=attention_type, cross_attention_dim=cross_attention_dim) # ---- DreamLite: install GQA-aware processor everywhere ---- for module in self.modules(): if isinstance(module, Attention): module.set_processor(DreamLiteAttnProcessor2_0()) # ----- override default processor so set_attn_processor("default") restores GQA ---- @property def default_attn_processor(self): # type: ignore[override] return DreamLiteAttnProcessor2_0() def set_default_attn_processor(self): # type: ignore[override] """Reinstall :class:`DreamLiteAttnProcessor2_0` everywhere. The parent implementation only knows about the diffusers stock processor sets and would raise for our GQA-aware processor; override so utilities that round-trip through this method (CPU offload, save/load, layerwise casting, ...) keep working unchanged. """ self.set_attn_processor(DreamLiteAttnProcessor2_0()) # ----- DreamLite extension: support `text_proj_rms` encoder_hid_proj ----- def _set_encoder_hid_proj( # type: ignore[override] self, encoder_hid_dim_type, cross_attention_dim, encoder_hid_dim, ): """ Override to support DreamLite's `text_proj_rms` variant (Linear → RMSNorm). All other variants fall back to the parent implementation, preserving full compatibility with upstream configs (`text_proj`, `text_image_proj`, `image_proj`, ...). """ if encoder_hid_dim_type == "text_proj_rms": if encoder_hid_dim is None: raise ValueError( "`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to 'text_proj_rms'." ) self.encoder_hid_proj = nn.Sequential( nn.Linear(encoder_hid_dim, cross_attention_dim), RMSNorm(cross_attention_dim, eps=1e-5, elementwise_affine=True), ) return super()._set_encoder_hid_proj( encoder_hid_dim_type=encoder_hid_dim_type, cross_attention_dim=cross_attention_dim, encoder_hid_dim=encoder_hid_dim, ) # ----- DreamLite extension: dispatch `text_proj_rms` like `text_proj` ----- def process_encoder_hidden_states( # type: ignore[override] self, encoder_hidden_states, added_cond_kwargs ): """ For `text_proj_rms`, the projection is a plain `nn.Sequential` applied to `encoder_hidden_states` (same call signature as `text_proj`). All other variants are delegated to the parent. """ if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj_rms": return self.encoder_hid_proj(encoder_hidden_states) return super().process_encoder_hidden_states( encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs, ) # ----- DreamLite extension: support `addition_embed_type == "time"` ----- def _set_add_embedding( # type: ignore[override] self, addition_embed_type, addition_embed_type_num_heads, addition_time_embed_dim, flip_sin_to_cos, freq_shift, cross_attention_dim, encoder_hid_dim, projection_class_embeddings_input_dim, time_embed_dim, ): """ Override to support DreamLite's `addition_embed_type == "time"` variant (same module layout as `text_time` but `get_aug_embed` does not require `text_embeds`). All other variants delegate to the parent implementation. """ if addition_embed_type == "time": from ..embeddings import TimestepEmbedding, Timesteps self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift) self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim) return super()._set_add_embedding( addition_embed_type=addition_embed_type, addition_embed_type_num_heads=addition_embed_type_num_heads, addition_time_embed_dim=addition_time_embed_dim, flip_sin_to_cos=flip_sin_to_cos, freq_shift=freq_shift, cross_attention_dim=cross_attention_dim, encoder_hid_dim=encoder_hid_dim, projection_class_embeddings_input_dim=projection_class_embeddings_input_dim, time_embed_dim=time_embed_dim, ) # ----- DreamLite extension: dispatch `addition_embed_type == "time"` ----- def get_aug_embed( # type: ignore[override] self, emb, encoder_hidden_states, added_cond_kwargs ): """ For `addition_embed_type == "time"`, build aug_emb from `time_ids` only (no `text_embeds`). All other variants are delegated to the parent. """ if self.config.addition_embed_type == "time": if "time_ids" not in added_cond_kwargs: raise ValueError( f"{self.__class__} has the config param `addition_embed_type` set to 'time' " "which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" ) time_ids = added_cond_kwargs.get("time_ids") time_embeds = self.add_time_proj(time_ids.flatten()) time_embeds = time_embeds.reshape((-1, self.config.projection_class_embeddings_input_dim)) add_embeds = time_embeds.to(emb.dtype) return self.add_embedding(add_embeds) return super().get_aug_embed( emb=emb, encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs, ) __all__ = [ "DreamLiteUNetModel", "DreamLiteUNetMidBlock2DCrossAttn", "DreamLiteCrossAttnDownBlock2D", "DreamLiteCrossAttnNoSelfAttnDownBlock2D", "DreamLiteCrossAttnUpBlock2D", "DreamLiteCrossAttnNoSelfAttnUpBlock2D", "DreamLiteDownBlock2D", "DreamLiteUpBlock2D", ]