modules/NeuralNetwork/unet.py

import logging
import math
from typing import Any, Dict, List, Optional
import torch.nn as nn
import torch as th
import torch

from modules.Utilities import util
from modules.AutoEncoders import ResBlock
from modules.NeuralNetwork import transformer
from modules.cond import cast
from modules.sample import sampling, sampling_util

UNET_MAP_ATTENTIONS = {
    "proj_in.weight",
    "proj_in.bias",
    "proj_out.weight",
    "proj_out.bias",
    "norm.weight",
    "norm.bias",
}

TRANSFORMER_BLOCKS = {
    "norm1.weight",
    "norm1.bias",
    "norm2.weight",
    "norm2.bias",
    "norm3.weight",
    "norm3.bias",
    "attn1.to_q.weight",
    "attn1.to_k.weight",
    "attn1.to_v.weight",
    "attn1.to_out.0.weight",
    "attn1.to_out.0.bias",
    "attn2.to_q.weight",
    "attn2.to_k.weight",
    "attn2.to_v.weight",
    "attn2.to_out.0.weight",
    "attn2.to_out.0.bias",
    "ff.net.0.proj.weight",
    "ff.net.0.proj.bias",
    "ff.net.2.weight",
    "ff.net.2.bias",
}

UNET_MAP_RESNET = {
    "in_layers.2.weight": "conv1.weight",
    "in_layers.2.bias": "conv1.bias",
    "emb_layers.1.weight": "time_emb_proj.weight",
    "emb_layers.1.bias": "time_emb_proj.bias",
    "out_layers.3.weight": "conv2.weight",
    "out_layers.3.bias": "conv2.bias",
    "skip_connection.weight": "conv_shortcut.weight",
    "skip_connection.bias": "conv_shortcut.bias",
    "in_layers.0.weight": "norm1.weight",
    "in_layers.0.bias": "norm1.bias",
    "out_layers.0.weight": "norm2.weight",
    "out_layers.0.bias": "norm2.bias",
}

UNET_MAP_BASIC = {
    ("label_emb.0.0.weight", "class_embedding.linear_1.weight"),
    ("label_emb.0.0.bias", "class_embedding.linear_1.bias"),
    ("label_emb.0.2.weight", "class_embedding.linear_2.weight"),
    ("label_emb.0.2.bias", "class_embedding.linear_2.bias"),
    ("label_emb.0.0.weight", "add_embedding.linear_1.weight"),
    ("label_emb.0.0.bias", "add_embedding.linear_1.bias"),
    ("label_emb.0.2.weight", "add_embedding.linear_2.weight"),
    ("label_emb.0.2.bias", "add_embedding.linear_2.bias"),
    ("input_blocks.0.0.weight", "conv_in.weight"),
    ("input_blocks.0.0.bias", "conv_in.bias"),
    ("out.0.weight", "conv_norm_out.weight"),
    ("out.0.bias", "conv_norm_out.bias"),
    ("out.2.weight", "conv_out.weight"),
    ("out.2.bias", "conv_out.bias"),
    ("time_embed.0.weight", "time_embedding.linear_1.weight"),
    ("time_embed.0.bias", "time_embedding.linear_1.bias"),
    ("time_embed.2.weight", "time_embedding.linear_2.weight"),
    ("time_embed.2.bias", "time_embedding.linear_2.bias"),
}

# taken from https://github.com/TencentARC/T2I-Adapter


def unet_to_diffusers(unet_config: dict) -> dict:
    """#### Convert a UNet configuration to a diffusers configuration.

    #### Args:
        - `unet_config` (dict): The UNet configuration.

    #### Returns:
        - `dict`: The diffusers configuration.
    """
    if "num_res_blocks" not in unet_config:
        return {}
    num_res_blocks = unet_config["num_res_blocks"]
    channel_mult = unet_config["channel_mult"]
    transformer_depth = unet_config["transformer_depth"][:]
    transformer_depth_output = unet_config["transformer_depth_output"][:]
    num_blocks = len(channel_mult)

    transformers_mid = unet_config.get("transformer_depth_middle", None)

    diffusers_unet_map = {}
    for x in range(num_blocks):
        n = 1 + (num_res_blocks[x] + 1) * x
        for i in range(num_res_blocks[x]):
            for b in UNET_MAP_RESNET:
                diffusers_unet_map[
                    "down_blocks.{}.resnets.{}.{}".format(x, i, UNET_MAP_RESNET[b])
                ] = "input_blocks.{}.0.{}".format(n, b)
            num_transformers = transformer_depth.pop(0)
            if num_transformers > 0:
                for b in UNET_MAP_ATTENTIONS:
                    diffusers_unet_map[
                        "down_blocks.{}.attentions.{}.{}".format(x, i, b)
                    ] = "input_blocks.{}.1.{}".format(n, b)
                for t in range(num_transformers):
                    for b in TRANSFORMER_BLOCKS:
                        diffusers_unet_map[
                            "down_blocks.{}.attentions.{}.transformer_blocks.{}.{}".format(
                                x, i, t, b
                            )
                        ] = "input_blocks.{}.1.transformer_blocks.{}.{}".format(n, t, b)
            n += 1
        for k in ["weight", "bias"]:
            diffusers_unet_map["down_blocks.{}.downsamplers.0.conv.{}".format(x, k)] = (
                "input_blocks.{}.0.op.{}".format(n, k)
            )

    i = 0
    for b in UNET_MAP_ATTENTIONS:
        diffusers_unet_map["mid_block.attentions.{}.{}".format(i, b)] = (
            "middle_block.1.{}".format(b)
        )
    for t in range(transformers_mid):
        for b in TRANSFORMER_BLOCKS:
            diffusers_unet_map[
                "mid_block.attentions.{}.transformer_blocks.{}.{}".format(i, t, b)
            ] = "middle_block.1.transformer_blocks.{}.{}".format(t, b)

    for i, n in enumerate([0, 2]):
        for b in UNET_MAP_RESNET:
            diffusers_unet_map[
                "mid_block.resnets.{}.{}".format(i, UNET_MAP_RESNET[b])
            ] = "middle_block.{}.{}".format(n, b)

    num_res_blocks = list(reversed(num_res_blocks))
    for x in range(num_blocks):
        n = (num_res_blocks[x] + 1) * x
        length = num_res_blocks[x] + 1
        for i in range(length):
            c = 0
            for b in UNET_MAP_RESNET:
                diffusers_unet_map[
                    "up_blocks.{}.resnets.{}.{}".format(x, i, UNET_MAP_RESNET[b])
                ] = "output_blocks.{}.0.{}".format(n, b)
            c += 1
            num_transformers = transformer_depth_output.pop()
            if num_transformers > 0:
                c += 1
                for b in UNET_MAP_ATTENTIONS:
                    diffusers_unet_map[
                        "up_blocks.{}.attentions.{}.{}".format(x, i, b)
                    ] = "output_blocks.{}.1.{}".format(n, b)
                for t in range(num_transformers):
                    for b in TRANSFORMER_BLOCKS:
                        diffusers_unet_map[
                            "up_blocks.{}.attentions.{}.transformer_blocks.{}.{}".format(
                                x, i, t, b
                            )
                        ] = "output_blocks.{}.1.transformer_blocks.{}.{}".format(
                            n, t, b
                        )
            if i == length - 1:
                for k in ["weight", "bias"]:
                    diffusers_unet_map[
                        "up_blocks.{}.upsamplers.0.conv.{}".format(x, k)
                    ] = "output_blocks.{}.{}.conv.{}".format(n, c, k)
            n += 1

    for k in UNET_MAP_BASIC:
        diffusers_unet_map[k[1]] = k[0]

    return diffusers_unet_map


def apply_control1(h: th.Tensor, control: any, name: str) -> th.Tensor:
    """#### Apply control to a tensor.

    #### Args:
        - `h` (torch.Tensor): The input tensor.
        - `control` (any): The control to apply.
        - `name` (str): The name of the control.

    #### Returns:
        - `torch.Tensor`: The controlled tensor.
    """
    return h


oai_ops = cast.disable_weight_init


class UNetModel1(nn.Module):
    """#### UNet Model class."""

    def __init__(
        self,
        image_size: int,
        in_channels: int,
        model_channels: int,
        out_channels: int,
        num_res_blocks: list,
        dropout: float = 0,
        channel_mult: tuple = (1, 2, 4, 8),
        conv_resample: bool = True,
        dims: int = 2,
        num_classes: int = None,
        use_checkpoint: bool = False,
        dtype: th.dtype = th.float32,
        num_heads: int = -1,
        num_head_channels: int = -1,
        num_heads_upsample: int = -1,
        use_scale_shift_norm: bool = False,
        resblock_updown: bool = False,
        use_new_attention_order: bool = False,
        use_spatial_transformer: bool = False,  # custom transformer support
        transformer_depth: int = 1,  # custom transformer support
        context_dim: int = None,  # custom transformer support
        n_embed: int = None,  # custom support for prediction of discrete ids into codebook of first stage vq model
        legacy: bool = True,
        disable_self_attentions: list = None,
        num_attention_blocks: list = None,
        disable_middle_self_attn: bool = False,
        use_linear_in_transformer: bool = False,
        adm_in_channels: int = None,
        transformer_depth_middle: int = None,
        transformer_depth_output: list = None,
        use_temporal_resblock: bool = False,
        use_temporal_attention: bool = False,
        time_context_dim: int = None,
        extra_ff_mix_layer: bool = False,
        use_spatial_context: bool = False,
        merge_strategy: any = None,
        merge_factor: float = 0.0,
        video_kernel_size: int = None,
        disable_temporal_crossattention: bool = False,
        max_ddpm_temb_period: int = 10000,
        device: th.device = None,
        operations: any = oai_ops,
    ):
        """#### Initialize the UNetModel1 class.

        #### Args:
            - `image_size` (int): The size of the input image.
            - `in_channels` (int): The number of input channels.
            - `model_channels` (int): The number of model channels.
            - `out_channels` (int): The number of output channels.
            - `num_res_blocks` (list): The number of residual blocks.
            - `dropout` (float, optional): The dropout rate. Defaults to 0.
            - `channel_mult` (tuple, optional): The channel multiplier. Defaults to (1, 2, 4, 8).
            - `conv_resample` (bool, optional): Whether to use convolutional resampling. Defaults to True.
            - `dims` (int, optional): The number of dimensions. Defaults to 2.
            - `num_classes` (int, optional): The number of classes. Defaults to None.
            - `use_checkpoint` (bool, optional): Whether to use checkpointing. Defaults to False.
            - `dtype` (torch.dtype, optional): The data type. Defaults to torch.float32.
            - `num_heads` (int, optional): The number of heads. Defaults to -1.
            - `num_head_channels` (int, optional): The number of head channels. Defaults to -1.
            - `num_heads_upsample` (int, optional): The number of heads for upsampling. Defaults to -1.
            - `use_scale_shift_norm` (bool, optional): Whether to use scale-shift normalization. Defaults to False.
            - `resblock_updown` (bool, optional): Whether to use residual blocks for up/down sampling. Defaults to False.
            - `use_new_attention_order` (bool, optional): Whether to use a new attention order. Defaults to False.
            - `use_spatial_transformer` (bool, optional): Whether to use a spatial transformer. Defaults to False.
            - `transformer_depth` (int, optional): The depth of the transformer. Defaults to 1.
            - `context_dim` (int, optional): The context dimension. Defaults to None.
            - `n_embed` (int, optional): The number of embeddings. Defaults to None.
            - `legacy` (bool, optional): Whether to use legacy mode. Defaults to True.
            - `disable_self_attentions` (list, optional): The list of self-attentions to disable. Defaults to None.
            - `num_attention_blocks` (list, optional): The number of attention blocks. Defaults to None.
            - `disable_middle_self_attn` (bool, optional): Whether to disable middle self-attention. Defaults to False.
            - `use_linear_in_transformer` (bool, optional): Whether to use linear in transformer. Defaults to False.
            - `adm_in_channels` (int, optional): The number of ADM input channels. Defaults to None.
            - `transformer_depth_middle` (int, optional): The depth of the middle transformer. Defaults to None.
            - `transformer_depth_output` (list, optional): The depth of the output transformer. Defaults to None.
            - `use_temporal_resblock` (bool, optional): Whether to use temporal residual blocks. Defaults to False.
            - `use_temporal_attention` (bool, optional): Whether to use temporal attention. Defaults to False.
            - `time_context_dim` (int, optional): The time context dimension. Defaults to None.
            - `extra_ff_mix_layer` (bool, optional): Whether to use an extra feed-forward mix layer. Defaults to False.
            - `use_spatial_context` (bool, optional): Whether to use spatial context. Defaults to False.
            - `merge_strategy` (any, optional): The merge strategy. Defaults to None.
            - `merge_factor` (float, optional): The merge factor. Defaults to 0.0.
            - `video_kernel_size` (int, optional): The video kernel size. Defaults to None.
            - `disable_temporal_crossattention` (bool, optional): Whether to disable temporal cross-attention. Defaults to False.
            - `max_ddpm_temb_period` (int, optional): The maximum DDPM temporal embedding period. Defaults to 10000.
            - `device` (torch.device, optional): The device to use. Defaults to None.
            - `operations` (any, optional): The operations to use. Defaults to oai_ops.
        """
        super().__init__()

        if context_dim is not None:
            self.context_dim = context_dim

        if num_heads_upsample == -1:
            num_heads_upsample = num_heads
        if num_head_channels == -1:
            assert num_heads != -1, (
                "Either num_heads or num_head_channels has to be set"
            )

        self.in_channels = in_channels
        self.model_channels = model_channels
        self.out_channels = out_channels
        self.num_res_blocks = num_res_blocks

        transformer_depth = transformer_depth[:]
        transformer_depth_output = transformer_depth_output[:]

        self.dropout = dropout
        self.channel_mult = channel_mult
        self.conv_resample = conv_resample
        self.num_classes = num_classes
        self.use_checkpoint = use_checkpoint
        self.dtype = dtype
        self.num_heads = num_heads
        self.num_head_channels = num_head_channels
        self.num_heads_upsample = num_heads_upsample
        self.use_temporal_resblocks = use_temporal_resblock
        self.predict_codebook_ids = n_embed is not None

        self.default_num_video_frames = None

        time_embed_dim = model_channels * 4
        self.time_embed = nn.Sequential(
            operations.Linear(
                model_channels, time_embed_dim, dtype=self.dtype, device=device
            ),
            nn.SiLU(),
            operations.Linear(
                time_embed_dim, time_embed_dim, dtype=self.dtype, device=device
            ),
        )

        self.input_blocks = nn.ModuleList(
            [
                sampling.TimestepEmbedSequential1(
                    operations.conv_nd(
                        dims,
                        in_channels,
                        model_channels,
                        3,
                        padding=1,
                        dtype=self.dtype,
                        device=device,
                    )
                )
            ]
        )
        self._feature_size = model_channels
        input_block_chans = [model_channels]
        ch = model_channels
        ds = 1

        def get_attention_layer(
            ch: int,
            num_heads: int,
            dim_head: int,
            depth: int = 1,
            context_dim: int = None,
            use_checkpoint: bool = False,
            disable_self_attn: bool = False,
        ) -> transformer.SpatialTransformer:
            """#### Get an attention layer.

            #### Args:
                - `ch` (int): The number of channels.
                - `num_heads` (int): The number of heads.
                - `dim_head` (int): The dimension of each head.
                - `depth` (int, optional): The depth of the transformer. Defaults to 1.
                - `context_dim` (int, optional): The context dimension. Defaults to None.
                - `use_checkpoint` (bool, optional): Whether to use checkpointing. Defaults to False.
                - `disable_self_attn` (bool, optional): Whether to disable self-attention. Defaults to False.

            #### Returns:
                - `transformer.SpatialTransformer`: The attention layer.
            """
            return transformer.SpatialTransformer(
                ch,
                num_heads,
                dim_head,
                depth=depth,
                context_dim=context_dim,
                disable_self_attn=disable_self_attn,
                use_linear=use_linear_in_transformer,
                use_checkpoint=use_checkpoint,
                dtype=self.dtype,
                device=device,
                operations=operations,
            )

        def get_resblock(
            merge_factor: float,
            merge_strategy: any,
            video_kernel_size: int,
            ch: int,
            time_embed_dim: int,
            dropout: float,
            out_channels: int,
            dims: int,
            use_checkpoint: bool,
            use_scale_shift_norm: bool,
            down: bool = False,
            up: bool = False,
            dtype: th.dtype = None,
            device: th.device = None,
            operations: any = oai_ops,
        ) -> ResBlock.ResBlock1:
            """#### Get a residual block.

            #### Args:
                - `merge_factor` (float): The merge factor.
                - `merge_strategy` (any): The merge strategy.
                - `video_kernel_size` (int): The video kernel size.
                - `ch` (int): The number of channels.
                - `time_embed_dim` (int): The time embedding dimension.
                - `dropout` (float): The dropout rate.
                - `out_channels` (int): The number of output channels.
                - `dims` (int): The number of dimensions.
                - `use_checkpoint` (bool): Whether to use checkpointing.
                - `use_scale_shift_norm` (bool): Whether to use scale-shift normalization.
                - `down` (bool, optional): Whether to use downsampling. Defaults to False.
                - `up` (bool, optional): Whether to use upsampling. Defaults to False.
                - `dtype` (torch.dtype, optional): The data type. Defaults to None.
                - `device` (torch.device, optional): The device. Defaults to None.
                - `operations` (any, optional): The operations to use. Defaults to oai_ops.

            #### Returns:
                - `ResBlock.ResBlock1`: The residual block.
            """
            return ResBlock.ResBlock1(
                channels=ch,
                emb_channels=time_embed_dim,
                dropout=dropout,
                out_channels=out_channels,
                use_checkpoint=use_checkpoint,
                dims=dims,
                use_scale_shift_norm=use_scale_shift_norm,
                down=down,
                up=up,
                dtype=dtype,
                device=device,
                operations=operations,
            )

        self.double_blocks = nn.ModuleList()
        for level, mult in enumerate(channel_mult):
            for nr in range(self.num_res_blocks[level]):
                layers = [
                    get_resblock(
                        merge_factor=merge_factor,
                        merge_strategy=merge_strategy,
                        video_kernel_size=video_kernel_size,
                        ch=ch,
                        time_embed_dim=time_embed_dim,
                        dropout=dropout,
                        out_channels=mult * model_channels,
                        dims=dims,
                        use_checkpoint=use_checkpoint,
                        use_scale_shift_norm=use_scale_shift_norm,
                        dtype=self.dtype,
                        device=device,
                        operations=operations,
                    )
                ]
                ch = mult * model_channels
                num_transformers = transformer_depth.pop(0)
                if num_transformers > 0:
                    dim_head = ch // num_heads
                    disabled_sa = False

                    if (
                        not util.exists(num_attention_blocks)
                        or nr < num_attention_blocks[level]
                    ):
                        layers.append(
                            get_attention_layer(
                                ch,
                                num_heads,
                                dim_head,
                                depth=num_transformers,
                                context_dim=context_dim,
                                disable_self_attn=disabled_sa,
                                use_checkpoint=use_checkpoint,
                            )
                        )
                self.input_blocks.append(sampling.TimestepEmbedSequential1(*layers))
                self._feature_size += ch
                input_block_chans.append(ch)
            if level != len(channel_mult) - 1:
                out_ch = ch
                self.input_blocks.append(
                    sampling.TimestepEmbedSequential1(
                        get_resblock(
                            merge_factor=merge_factor,
                            merge_strategy=merge_strategy,
                            video_kernel_size=video_kernel_size,
                            ch=ch,
                            time_embed_dim=time_embed_dim,
                            dropout=dropout,
                            out_channels=out_ch,
                            dims=dims,
                            use_checkpoint=use_checkpoint,
                            use_scale_shift_norm=use_scale_shift_norm,
                            down=True,
                            dtype=self.dtype,
                            device=device,
                            operations=operations,
                        )
                        if resblock_updown
                        else ResBlock.Downsample1(
                            ch,
                            conv_resample,
                            dims=dims,
                            out_channels=out_ch,
                            dtype=self.dtype,
                            device=device,
                            operations=operations,
                        )
                    )
                )
                ch = out_ch
                input_block_chans.append(ch)
                ds *= 2
                self._feature_size += ch

        dim_head = ch // num_heads
        mid_block = [
            get_resblock(
                merge_factor=merge_factor,
                merge_strategy=merge_strategy,
                video_kernel_size=video_kernel_size,
                ch=ch,
                time_embed_dim=time_embed_dim,
                dropout=dropout,
                out_channels=None,
                dims=dims,
                use_checkpoint=use_checkpoint,
                use_scale_shift_norm=use_scale_shift_norm,
                dtype=self.dtype,
                device=device,
                operations=operations,
            )
        ]

        self.middle_block = None
        if transformer_depth_middle >= -1:
            if transformer_depth_middle >= 0:
                mid_block += [
                    get_attention_layer(  # always uses a self-attn
                        ch,
                        num_heads,
                        dim_head,
                        depth=transformer_depth_middle,
                        context_dim=context_dim,
                        disable_self_attn=disable_middle_self_attn,
                        use_checkpoint=use_checkpoint,
                    ),
                    get_resblock(
                        merge_factor=merge_factor,
                        merge_strategy=merge_strategy,
                        video_kernel_size=video_kernel_size,
                        ch=ch,
                        time_embed_dim=time_embed_dim,
                        dropout=dropout,
                        out_channels=None,
                        dims=dims,
                        use_checkpoint=use_checkpoint,
                        use_scale_shift_norm=use_scale_shift_norm,
                        dtype=self.dtype,
                        device=device,
                        operations=operations,
                    ),
                ]
            self.middle_block = sampling.TimestepEmbedSequential1(*mid_block)
        self._feature_size += ch

        self.output_blocks = nn.ModuleList([])
        for level, mult in list(enumerate(channel_mult))[::-1]:
            for i in range(self.num_res_blocks[level] + 1):
                ich = input_block_chans.pop()
                layers = [
                    get_resblock(
                        merge_factor=merge_factor,
                        merge_strategy=merge_strategy,
                        video_kernel_size=video_kernel_size,
                        ch=ch + ich,
                        time_embed_dim=time_embed_dim,
                        dropout=dropout,
                        out_channels=model_channels * mult,
                        dims=dims,
                        use_checkpoint=use_checkpoint,
                        use_scale_shift_norm=use_scale_shift_norm,
                        dtype=self.dtype,
                        device=device,
                        operations=operations,
                    )
                ]
                ch = model_channels * mult
                num_transformers = transformer_depth_output.pop()
                if num_transformers > 0:
                    dim_head = ch // num_heads
                    disabled_sa = False

                    if (
                        not util.exists(num_attention_blocks)
                        or i < num_attention_blocks[level]
                    ):
                        layers.append(
                            get_attention_layer(
                                ch,
                                num_heads,
                                dim_head,
                                depth=num_transformers,
                                context_dim=context_dim,
                                disable_self_attn=disabled_sa,
                                use_checkpoint=use_checkpoint,
                            )
                        )
                if level and i == self.num_res_blocks[level]:
                    out_ch = ch
                    layers.append(
                        get_resblock(
                            merge_factor=merge_factor,
                            merge_strategy=merge_strategy,
                            video_kernel_size=video_kernel_size,
                            ch=ch,
                            time_embed_dim=time_embed_dim,
                            dropout=dropout,
                            out_channels=out_ch,
                            dims=dims,
                            use_checkpoint=use_checkpoint,
                            use_scale_shift_norm=use_scale_shift_norm,
                            up=True,
                            dtype=self.dtype,
                            device=device,
                            operations=operations,
                        )
                        if resblock_updown
                        else ResBlock.Upsample1(
                            ch,
                            conv_resample,
                            dims=dims,
                            out_channels=out_ch,
                            dtype=self.dtype,
                            device=device,
                            operations=operations,
                        )
                    )
                    ds //= 2
                self.output_blocks.append(sampling.TimestepEmbedSequential1(*layers))
                self._feature_size += ch

        self.out = nn.Sequential(
            operations.GroupNorm(32, ch, dtype=self.dtype, device=device),
            nn.SiLU(),
            util.zero_module(
                operations.conv_nd(
                    dims,
                    model_channels,
                    out_channels,
                    3,
                    padding=1,
                    dtype=self.dtype,
                    device=device,
                )
            ),
        )

    def forward(
        self,
        x: torch.Tensor,
        timesteps: Optional[torch.Tensor] = None,
        context: Optional[torch.Tensor] = None,
        y: Optional[torch.Tensor] = None,
        control: Optional[torch.Tensor] = None,
        transformer_options: Dict[str, Any] = {},
        **kwargs: Any,
    ) -> torch.Tensor:
        """#### Forward pass of the UNet model with optimized calculations."""
        # Setup transformer options (avoid unused variable)
        transformer_options["original_shape"] = list(x.shape)
        transformer_options["transformer_index"] = 0

        # Extract kwargs efficiently
        num_video_frames = kwargs.get("num_video_frames", self.default_num_video_frames)
        image_only_indicator = kwargs.get("image_only_indicator", None)
        time_context = kwargs.get("time_context", None)

        # Validation
        assert (y is not None) == (self.num_classes is not None), (
            "must specify y if and only if the model is class-conditional"
        )

        # Time embedding - optimize by computing with target dtype directly
        t_emb = sampling_util.timestep_embedding(timesteps, self.model_channels).to(
            x.dtype
        )
        emb = self.time_embed(t_emb)

        # Input blocks processing
        hs = []
        h = x
        for id, module in enumerate(self.input_blocks):
            transformer_options["block"] = ("input", id)
            h = ResBlock.forward_timestep_embed1(
                module,
                h,
                emb,
                context,
                transformer_options,
                time_context=time_context,
                num_video_frames=num_video_frames,
                image_only_indicator=image_only_indicator,
            )
            h = apply_control1(h, control, "input")
            hs.append(h)

        # Middle block processing
        transformer_options["block"] = ("middle", 0)
        if self.middle_block is not None:
            h = ResBlock.forward_timestep_embed1(
                self.middle_block,
                h,
                emb,
                context,
                transformer_options,
                time_context=time_context,
                num_video_frames=num_video_frames,
                image_only_indicator=image_only_indicator,
            )
        h = apply_control1(h, control, "middle")

        # Output blocks processing - optimize memory usage
        for id, module in enumerate(self.output_blocks):
            transformer_options["block"] = ("output", id)
            hsp = hs.pop()
            hsp = apply_control1(hsp, control, "output")

            # Concatenate tensors
            h = torch.cat([h, hsp], dim=1)
            del hsp  # Free memory immediately

            # Only calculate output shape when needed
            output_shape = hs[-1].shape if hs else None

            h = ResBlock.forward_timestep_embed1(
                module,
                h,
                emb,
                context,
                transformer_options,
                output_shape,
                time_context=time_context,
                num_video_frames=num_video_frames,
                image_only_indicator=image_only_indicator,
            )

        # Ensure output has correct dtype
        h = h.type(x.dtype)
        return self.out(h)


def detect_unet_config(
    state_dict: Dict[str, torch.Tensor], key_prefix: str
) -> Dict[str, Any]:
    """#### Detect the UNet configuration from a state dictionary.

    #### Args:
        - `state_dict` (Dict[str, torch.Tensor]): The state dictionary.
        - `key_prefix` (str): The key prefix.

    #### Returns:
        - `Dict[str, Any]`: The detected UNet configuration.
    """
    state_dict_keys = list(state_dict.keys())

    if (
        "{}joint_blocks.0.context_block.attn.qkv.weight".format(key_prefix)
        in state_dict_keys
    ):  # mmdit model
        unet_config = {}
        unet_config["in_channels"] = state_dict[
            "{}x_embedder.proj.weight".format(key_prefix)
        ].shape[1]
        patch_size = state_dict["{}x_embedder.proj.weight".format(key_prefix)].shape[2]
        unet_config["patch_size"] = patch_size
        final_layer = "{}final_layer.linear.weight".format(key_prefix)
        if final_layer in state_dict:
            unet_config["out_channels"] = state_dict[final_layer].shape[0] // (
                patch_size * patch_size
            )

        unet_config["depth"] = (
            state_dict["{}x_embedder.proj.weight".format(key_prefix)].shape[0] // 64
        )
        unet_config["input_size"] = None
        y_key = "{}y_embedder.mlp.0.weight".format(key_prefix)
        if y_key in state_dict_keys:
            unet_config["adm_in_channels"] = state_dict[y_key].shape[1]

        context_key = "{}context_embedder.weight".format(key_prefix)
        if context_key in state_dict_keys:
            in_features = state_dict[context_key].shape[1]
            out_features = state_dict[context_key].shape[0]
            unet_config["context_embedder_config"] = {
                "target": "torch.nn.Linear",
                "params": {"in_features": in_features, "out_features": out_features},
            }
        num_patches_key = "{}pos_embed".format(key_prefix)
        if num_patches_key in state_dict_keys:
            num_patches = state_dict[num_patches_key].shape[1]
            unet_config["num_patches"] = num_patches
            unet_config["pos_embed_max_size"] = round(math.sqrt(num_patches))

        rms_qk = "{}joint_blocks.0.context_block.attn.ln_q.weight".format(key_prefix)
        if rms_qk in state_dict_keys:
            unet_config["qk_norm"] = "rms"

        unet_config["pos_embed_scaling_factor"] = None  # unused for inference
        context_processor = "{}context_processor.layers.0.attn.qkv.weight".format(
            key_prefix
        )
        if context_processor in state_dict_keys:
            unet_config["context_processor_layers"] = transformer.count_blocks(
                state_dict_keys,
                "{}context_processor.layers.".format(key_prefix) + "{}.",
            )
        return unet_config

    if "{}clf.1.weight".format(key_prefix) in state_dict_keys:  # stable cascade
        unet_config = {}
        text_mapper_name = "{}clip_txt_mapper.weight".format(key_prefix)
        if text_mapper_name in state_dict_keys:
            unet_config["stable_cascade_stage"] = "c"
            w = state_dict[text_mapper_name]
            if w.shape[0] == 1536:  # stage c lite
                unet_config["c_cond"] = 1536
                unet_config["c_hidden"] = [1536, 1536]
                unet_config["nhead"] = [24, 24]
                unet_config["blocks"] = [[4, 12], [12, 4]]
            elif w.shape[0] == 2048:  # stage c full
                unet_config["c_cond"] = 2048
        elif "{}clip_mapper.weight".format(key_prefix) in state_dict_keys:
            unet_config["stable_cascade_stage"] = "b"
            w = state_dict["{}down_blocks.1.0.channelwise.0.weight".format(key_prefix)]
            if w.shape[-1] == 640:
                unet_config["c_hidden"] = [320, 640, 1280, 1280]
                unet_config["nhead"] = [-1, -1, 20, 20]
                unet_config["blocks"] = [[2, 6, 28, 6], [6, 28, 6, 2]]
                unet_config["block_repeat"] = [[1, 1, 1, 1], [3, 3, 2, 2]]
            elif w.shape[-1] == 576:  # stage b lite
                unet_config["c_hidden"] = [320, 576, 1152, 1152]
                unet_config["nhead"] = [-1, 9, 18, 18]
                unet_config["blocks"] = [[2, 4, 14, 4], [4, 14, 4, 2]]
                unet_config["block_repeat"] = [[1, 1, 1, 1], [2, 2, 2, 2]]
        return unet_config

    if (
        "{}transformer.rotary_pos_emb.inv_freq".format(key_prefix) in state_dict_keys
    ):  # stable audio dit
        unet_config = {}
        unet_config["audio_model"] = "dit1.0"
        return unet_config

    if (
        "{}double_layers.0.attn.w1q.weight".format(key_prefix) in state_dict_keys
    ):  # aura flow dit
        unet_config = {}
        unet_config["max_seq"] = state_dict[
            "{}positional_encoding".format(key_prefix)
        ].shape[1]
        unet_config["cond_seq_dim"] = state_dict[
            "{}cond_seq_linear.weight".format(key_prefix)
        ].shape[1]
        double_layers = transformer.count_blocks(
            state_dict_keys, "{}double_layers.".format(key_prefix) + "{}."
        )
        single_layers = transformer.count_blocks(
            state_dict_keys, "{}single_layers.".format(key_prefix) + "{}."
        )
        unet_config["n_double_layers"] = double_layers
        unet_config["n_layers"] = double_layers + single_layers
        return unet_config

    if "{}mlp_t5.0.weight".format(key_prefix) in state_dict_keys:  # Hunyuan DiT
        unet_config = {}
        unet_config["image_model"] = "hydit"
        unet_config["depth"] = transformer.count_blocks(
            state_dict_keys, "{}blocks.".format(key_prefix) + "{}."
        )
        unet_config["hidden_size"] = state_dict[
            "{}x_embedder.proj.weight".format(key_prefix)
        ].shape[0]
        if unet_config["hidden_size"] == 1408 and unet_config["depth"] == 40:  # DiT-g/2
            unet_config["mlp_ratio"] = 4.3637
        if state_dict["{}extra_embedder.0.weight".format(key_prefix)].shape[1] == 3968:
            unet_config["size_cond"] = True
            unet_config["use_style_cond"] = True
            unet_config["image_model"] = "hydit1"
        return unet_config

    if (
        "{}double_blocks.0.img_attn.norm.key_norm.scale".format(key_prefix)
        in state_dict_keys
    ):  # Flux
        dit_config = {}
        dit_config["image_model"] = "flux"
        dit_config["in_channels"] = 16
        dit_config["vec_in_dim"] = 768
        dit_config["context_in_dim"] = 4096
        dit_config["hidden_size"] = 3072
        dit_config["mlp_ratio"] = 4.0
        dit_config["num_heads"] = 24
        dit_config["depth"] = transformer.count_blocks(
            state_dict_keys, "{}double_blocks.".format(key_prefix) + "{}."
        )
        dit_config["depth_single_blocks"] = transformer.count_blocks(
            state_dict_keys, "{}single_blocks.".format(key_prefix) + "{}."
        )
        dit_config["axes_dim"] = [16, 56, 56]
        dit_config["theta"] = 10000
        dit_config["qkv_bias"] = True
        dit_config["guidance_embed"] = (
            "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys
        )
        return dit_config

    if "{}input_blocks.0.0.weight".format(key_prefix) not in state_dict_keys:
        return None

    unet_config = {
        "use_checkpoint": False,
        "image_size": 32,
        "use_spatial_transformer": True,
        "legacy": False,
    }

    y_input = "{}label_emb.0.0.weight".format(key_prefix)
    if y_input in state_dict_keys:
        unet_config["num_classes"] = "sequential"
        unet_config["adm_in_channels"] = state_dict[y_input].shape[1]
    else:
        unet_config["adm_in_channels"] = None

    model_channels = state_dict["{}input_blocks.0.0.weight".format(key_prefix)].shape[0]
    in_channels = state_dict["{}input_blocks.0.0.weight".format(key_prefix)].shape[1]

    out_key = "{}out.2.weight".format(key_prefix)
    if out_key in state_dict:
        out_channels = state_dict[out_key].shape[0]
    else:
        out_channels = 4

    num_res_blocks = []
    channel_mult = []
    transformer_depth = []
    transformer_depth_output = []
    context_dim = None
    use_linear_in_transformer = False

    video_model = False
    video_model_cross = False

    current_res = 1
    count = 0

    last_res_blocks = 0
    last_channel_mult = 0

    input_block_count = transformer.count_blocks(
        state_dict_keys, "{}input_blocks".format(key_prefix) + ".{}."
    )
    for count in range(input_block_count):
        prefix = "{}input_blocks.{}.".format(key_prefix, count)
        prefix_output = "{}output_blocks.{}.".format(
            key_prefix, input_block_count - count - 1
        )

        block_keys = sorted(
            list(filter(lambda a: a.startswith(prefix), state_dict_keys))
        )
        if len(block_keys) == 0:
            break

        block_keys_output = sorted(
            list(filter(lambda a: a.startswith(prefix_output), state_dict_keys))
        )

        if "{}0.op.weight".format(prefix) in block_keys:  # new layer
            num_res_blocks.append(last_res_blocks)
            channel_mult.append(last_channel_mult)

            current_res *= 2
            last_res_blocks = 0
            last_channel_mult = 0
            out = transformer.calculate_transformer_depth(
                prefix_output, state_dict_keys, state_dict
            )
            if out is not None:
                transformer_depth_output.append(out[0])
            else:
                transformer_depth_output.append(0)
        else:
            res_block_prefix = "{}0.in_layers.0.weight".format(prefix)
            if res_block_prefix in block_keys:
                last_res_blocks += 1
                last_channel_mult = (
                    state_dict["{}0.out_layers.3.weight".format(prefix)].shape[0]
                    // model_channels
                )

                out = transformer.calculate_transformer_depth(
                    prefix, state_dict_keys, state_dict
                )
                if out is not None:
                    transformer_depth.append(out[0])
                    if context_dim is None:
                        context_dim = out[1]
                        use_linear_in_transformer = out[2]
                        out[3]
                else:
                    transformer_depth.append(0)

            res_block_prefix = "{}0.in_layers.0.weight".format(prefix_output)
            if res_block_prefix in block_keys_output:
                out = transformer.calculate_transformer_depth(
                    prefix_output, state_dict_keys, state_dict
                )
                if out is not None:
                    transformer_depth_output.append(out[0])
                else:
                    transformer_depth_output.append(0)

    num_res_blocks.append(last_res_blocks)
    channel_mult.append(last_channel_mult)
    if "{}middle_block.1.proj_in.weight".format(key_prefix) in state_dict_keys:
        transformer_depth_middle = transformer.count_blocks(
            state_dict_keys,
            "{}middle_block.1.transformer_blocks.".format(key_prefix) + "{}",
        )
    elif "{}middle_block.0.in_layers.0.weight".format(key_prefix) in state_dict_keys:
        transformer_depth_middle = -1
    else:
        transformer_depth_middle = -2

    unet_config["in_channels"] = in_channels
    unet_config["out_channels"] = out_channels
    unet_config["model_channels"] = model_channels
    unet_config["num_res_blocks"] = num_res_blocks
    unet_config["transformer_depth"] = transformer_depth
    unet_config["transformer_depth_output"] = transformer_depth_output
    unet_config["channel_mult"] = channel_mult
    unet_config["transformer_depth_middle"] = transformer_depth_middle
    unet_config["use_linear_in_transformer"] = use_linear_in_transformer
    unet_config["context_dim"] = context_dim

    if video_model:
        unet_config["extra_ff_mix_layer"] = True
        unet_config["use_spatial_context"] = True
        unet_config["merge_strategy"] = "learned_with_images"
        unet_config["merge_factor"] = 0.0
        unet_config["video_kernel_size"] = [3, 1, 1]
        unet_config["use_temporal_resblock"] = True
        unet_config["use_temporal_attention"] = True
        unet_config["disable_temporal_crossattention"] = not video_model_cross
    else:
        unet_config["use_temporal_resblock"] = False
        unet_config["use_temporal_attention"] = False

    return unet_config


def model_config_from_unet_config(
    unet_config: Dict[str, Any], state_dict: Optional[Dict[str, torch.Tensor]] = None
) -> Any:
    """#### Get the model configuration from a UNet configuration.

    #### Args:
        - `unet_config` (Dict[str, Any]): The UNet configuration.
        - `state_dict` (Optional[Dict[str, torch.Tensor]], optional): The state dictionary. Defaults to None.

    #### Returns:
        - `Any`: The model configuration.
    """
    from modules.SD15 import SD15

    for model_config in SD15.models:
        if model_config.matches(unet_config, state_dict):
            return model_config(unet_config)

    logging.error("no match {}".format(unet_config))
    return None


def model_config_from_unet(
    state_dict: Dict[str, torch.Tensor],
    unet_key_prefix: str,
    use_base_if_no_match: bool = False,
) -> Any:
    """#### Get the model configuration from a UNet state dictionary.

    #### Args:
        - `state_dict` (Dict[str, torch.Tensor]): The state dictionary.
        - `unet_key_prefix` (str): The UNet key prefix.
        - `use_base_if_no_match` (bool, optional): Whether to use the base configuration if no match is found. Defaults to False.

    #### Returns:
        - `Any`: The model configuration.
    """
    unet_config = detect_unet_config(state_dict, unet_key_prefix)
    if unet_config is None:
        return None
    model_config = model_config_from_unet_config(unet_config, state_dict)
    return model_config


def unet_dtype1(
    device: Optional[torch.device] = None,
    model_params: int = 0,
    supported_dtypes: List[torch.dtype] = [
        torch.float16,
        torch.bfloat16,
        torch.float32,
    ],
) -> torch.dtype:
    """#### Get the dtype for the UNet model.

    #### Args:
        - `device` (Optional[torch.device], optional): The device. Defaults to None.
        - `model_params` (int, optional): The model parameters. Defaults to 0.
        - `supported_dtypes` (List[torch.dtype], optional): The supported dtypes. Defaults to [torch.float16, torch.bfloat16, torch.float32].

    #### Returns:
        - `torch.dtype`: The dtype for the UNet model.
    """
    return torch.float16