# 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,