Create min_sdxl.py

module/min_sdxl.py

@@ -0,0 +1,907 @@
+# Modified from minSDXL by Simo Ryu: 
+# https://github.com/cloneofsimo/minSDXL ,
+# which is in turn modified from the original code of:
+# https://github.com/huggingface/diffusers
+# So has APACHE 2.0 license
+
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import inspect
+
+from collections import namedtuple
+
+from torch.fft import fftn, fftshift, ifftn, ifftshift
+
+from diffusers.models.attention_processor import AttnProcessor, AttnProcessor2_0
+
+# Implementation of FreeU for minsdxl
+
+def fourier_filter(x_in: "torch.Tensor", threshold: int, scale: int) -> "torch.Tensor":
+    """Fourier filter as introduced in FreeU (https://arxiv.org/abs/2309.11497).
+    This version of the method comes from here:
+    https://github.com/huggingface/diffusers/pull/5164#issuecomment-1732638706
+    """
+    x = x_in
+    B, C, H, W = x.shape
+
+    # Non-power of 2 images must be float32
+    if (W & (W - 1)) != 0 or (H & (H - 1)) != 0:
+        x = x.to(dtype=torch.float32)
+
+    # FFT
+    x_freq = fftn(x, dim=(-2, -1))
+    x_freq = fftshift(x_freq, dim=(-2, -1))
+
+    B, C, H, W = x_freq.shape
+    mask = torch.ones((B, C, H, W), device=x.device)
+
+    crow, ccol = H // 2, W // 2
+    mask[..., crow - threshold : crow + threshold, ccol - threshold : ccol + threshold] = scale
+    x_freq = x_freq * mask
+
+    # IFFT
+    x_freq = ifftshift(x_freq, dim=(-2, -1))
+    x_filtered = ifftn(x_freq, dim=(-2, -1)).real
+
+    return x_filtered.to(dtype=x_in.dtype)
+
+
+def apply_freeu(
+    resolution_idx: int, hidden_states: "torch.Tensor", res_hidden_states: "torch.Tensor", **freeu_kwargs):
+    """Applies the FreeU mechanism as introduced in https:
+    //arxiv.org/abs/2309.11497. Adapted from the official code repository: https://github.com/ChenyangSi/FreeU.
+    Args:
+        resolution_idx (`int`): Integer denoting the UNet block where FreeU is being applied.
+        hidden_states (`torch.Tensor`): Inputs to the underlying block.
+        res_hidden_states (`torch.Tensor`): Features from the skip block corresponding to the underlying block.
+        s1 (`float`): Scaling factor for stage 1 to attenuate the contributions of the skip features.
+        s2 (`float`): Scaling factor for stage 2 to attenuate the contributions of the skip features.
+        b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features.
+        b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features.
+    """
+    if resolution_idx == 0:
+        num_half_channels = hidden_states.shape[1] // 2
+        hidden_states[:, :num_half_channels] = hidden_states[:, :num_half_channels] * freeu_kwargs["b1"]
+        res_hidden_states = fourier_filter(res_hidden_states, threshold=1, scale=freeu_kwargs["s1"])
+    if resolution_idx == 1:
+        num_half_channels = hidden_states.shape[1] // 2
+        hidden_states[:, :num_half_channels] = hidden_states[:, :num_half_channels] * freeu_kwargs["b2"]
+        res_hidden_states = fourier_filter(res_hidden_states, threshold=1, scale=freeu_kwargs["s2"])
+
+    return hidden_states, res_hidden_states
+
+# Diffusers-style LoRA to keep everything in the min_sdxl.py file
+
+class LoRALinearLayer(nn.Module):
+    r"""
+    A linear layer that is used with LoRA.
+    Parameters:
+        in_features (`int`):
+            Number of input features.
+        out_features (`int`):
+            Number of output features.
+        rank (`int`, `optional`, defaults to 4):
+            The rank of the LoRA layer.
+        network_alpha (`float`, `optional`, defaults to `None`):
+            The value of the network alpha used for stable learning and preventing underflow. This value has the same
+            meaning as the `--network_alpha` option in the kohya-ss trainer script. See
+            https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning
+        device (`torch.device`, `optional`, defaults to `None`):
+            The device to use for the layer's weights.
+        dtype (`torch.dtype`, `optional`, defaults to `None`):
+            The dtype to use for the layer's weights.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        rank: int = 4,
+        network_alpha: Optional[float] = None,
+        device: Optional[Union[torch.device, str]] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        super().__init__()
+
+        self.down = nn.Linear(in_features, rank, bias=False, device=device, dtype=dtype)
+        self.up = nn.Linear(rank, out_features, bias=False, device=device, dtype=dtype)
+        # This value has the same meaning as the `--network_alpha` option in the kohya-ss trainer script.
+        # See https://github.com/darkstorm2150/sd-scripts/blob/main/docs/train_network_README-en.md#execute-learning
+        self.network_alpha = network_alpha
+        self.rank = rank
+        self.out_features = out_features
+        self.in_features = in_features
+
+        nn.init.normal_(self.down.weight, std=1 / rank)
+        nn.init.zeros_(self.up.weight)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        orig_dtype = hidden_states.dtype
+        dtype = self.down.weight.dtype
+
+        down_hidden_states = self.down(hidden_states.to(dtype))
+        up_hidden_states = self.up(down_hidden_states)
+
+        if self.network_alpha is not None:
+            up_hidden_states *= self.network_alpha / self.rank
+
+        return up_hidden_states.to(orig_dtype)
+
+class LoRACompatibleLinear(nn.Linear):
+    """
+    A Linear layer that can be used with LoRA.
+    """
+
+    def __init__(self, *args, lora_layer: Optional[LoRALinearLayer] = None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.lora_layer = lora_layer
+
+    def set_lora_layer(self, lora_layer: Optional[LoRALinearLayer]):
+        self.lora_layer = lora_layer
+
+    def _fuse_lora(self, lora_scale: float = 1.0, safe_fusing: bool = False):
+        if self.lora_layer is None:
+            return
+
+        dtype, device = self.weight.data.dtype, self.weight.data.device
+
+        w_orig = self.weight.data.float()
+        w_up = self.lora_layer.up.weight.data.float()
+        w_down = self.lora_layer.down.weight.data.float()
+
+        if self.lora_layer.network_alpha is not None:
+            w_up = w_up * self.lora_layer.network_alpha / self.lora_layer.rank
+
+        fused_weight = w_orig + (lora_scale * torch.bmm(w_up[None, :], w_down[None, :])[0])
+
+        if safe_fusing and torch.isnan(fused_weight).any().item():
+            raise ValueError(
+                "This LoRA weight seems to be broken. "
+                f"Encountered NaN values when trying to fuse LoRA weights for {self}."
+                "LoRA weights will not be fused."
+            )
+
+        self.weight.data = fused_weight.to(device=device, dtype=dtype)
+
+        # we can drop the lora layer now
+        self.lora_layer = None
+
+        # offload the up and down matrices to CPU to not blow the memory
+        self.w_up = w_up.cpu()
+        self.w_down = w_down.cpu()
+        self._lora_scale = lora_scale
+
+    def _unfuse_lora(self):
+        if not (getattr(self, "w_up", None) is not None and getattr(self, "w_down", None) is not None):
+            return
+
+        fused_weight = self.weight.data
+        dtype, device = fused_weight.dtype, fused_weight.device
+
+        w_up = self.w_up.to(device=device).float()
+        w_down = self.w_down.to(device).float()
+
+        unfused_weight = fused_weight.float() - (self._lora_scale * torch.bmm(w_up[None, :], w_down[None, :])[0])
+        self.weight.data = unfused_weight.to(device=device, dtype=dtype)
+
+        self.w_up = None
+        self.w_down = None
+
+    def forward(self, hidden_states: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
+        if self.lora_layer is None:
+            out = super().forward(hidden_states)
+            return out
+        else:
+            out = super().forward(hidden_states) + (scale * self.lora_layer(hidden_states))
+            return out
+
+class Timesteps(nn.Module):
+    def __init__(self, num_channels: int = 320):
+        super().__init__()
+        self.num_channels = num_channels
+
+    def forward(self, timesteps):
+        half_dim = self.num_channels // 2
+        exponent = -math.log(10000) * torch.arange(
+            half_dim, dtype=torch.float32, device=timesteps.device
+        )
+        exponent = exponent / (half_dim - 0.0)
+
+        emb = torch.exp(exponent)
+        emb = timesteps[:, None].float() * emb[None, :]
+
+        sin_emb = torch.sin(emb)
+        cos_emb = torch.cos(emb)
+        emb = torch.cat([cos_emb, sin_emb], dim=-1)
+
+        return emb
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, in_features, out_features):
+        super(TimestepEmbedding, self).__init__()
+        self.linear_1 = nn.Linear(in_features, out_features, bias=True)
+        self.act = nn.SiLU()
+        self.linear_2 = nn.Linear(out_features, out_features, bias=True)
+
+    def forward(self, sample):
+        sample = self.linear_1(sample)
+        sample = self.act(sample)
+        sample = self.linear_2(sample)
+
+        return sample
+
+
+class ResnetBlock2D(nn.Module):
+    def __init__(self, in_channels, out_channels, conv_shortcut=True):
+        super(ResnetBlock2D, self).__init__()
+        self.norm1 = nn.GroupNorm(32, in_channels, eps=1e-05, affine=True)
+        self.conv1 = nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        self.time_emb_proj = nn.Linear(1280, out_channels, bias=True)
+        self.norm2 = nn.GroupNorm(32, out_channels, eps=1e-05, affine=True)
+        self.dropout = nn.Dropout(p=0.0, inplace=False)
+        self.conv2 = nn.Conv2d(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        self.nonlinearity = nn.SiLU()
+        self.conv_shortcut = None
+        if conv_shortcut:
+            self.conv_shortcut = nn.Conv2d(
+                in_channels, out_channels, kernel_size=1, stride=1
+            )
+
+    def forward(self, input_tensor, temb):
+        hidden_states = input_tensor
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+
+        hidden_states = self.conv1(hidden_states)
+
+        temb = self.nonlinearity(temb)
+        temb = self.time_emb_proj(temb)[:, :, None, None]
+        hidden_states = hidden_states + temb
+        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:
+            input_tensor = self.conv_shortcut(input_tensor)
+
+        output_tensor = input_tensor + hidden_states
+
+        return output_tensor
+
+
+class Attention(nn.Module):
+    def __init__(
+        self, inner_dim, cross_attention_dim=None, num_heads=None, dropout=0.0, processor=None, scale_qk=True
+    ):
+        super(Attention, self).__init__()
+        if num_heads is None:
+            self.head_dim = 64
+            self.num_heads = inner_dim // self.head_dim
+        else:
+            self.num_heads = num_heads
+            self.head_dim = inner_dim // num_heads
+
+        self.scale = self.head_dim**-0.5
+        if cross_attention_dim is None:
+            cross_attention_dim = inner_dim
+        self.to_q = LoRACompatibleLinear(inner_dim, inner_dim, bias=False)
+        self.to_k = LoRACompatibleLinear(cross_attention_dim, inner_dim, bias=False)
+        self.to_v = LoRACompatibleLinear(cross_attention_dim, inner_dim, bias=False)
+
+        self.to_out = nn.ModuleList(
+            [LoRACompatibleLinear(inner_dim, inner_dim), nn.Dropout(dropout, inplace=False)]
+        )
+
+        self.scale_qk = scale_qk
+        if processor is None:
+            processor = (
+                AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+            )
+        self.set_processor(processor)
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        **cross_attention_kwargs,
+    ) -> torch.Tensor:
+        r"""
+        The forward method of the `Attention` class.
+        Args:
+            hidden_states (`torch.Tensor`):
+                The hidden states of the query.
+            encoder_hidden_states (`torch.Tensor`, *optional*):
+                The hidden states of the encoder.
+            attention_mask (`torch.Tensor`, *optional*):
+                The attention mask to use. If `None`, no mask is applied.
+            **cross_attention_kwargs:
+                Additional keyword arguments to pass along to the cross attention.
+        Returns:
+            `torch.Tensor`: The output of the attention layer.
+        """
+        # The `Attention` class can call different attention processors / attention functions
+        # here we simply pass along all tensors to the selected processor class
+        # For standard processors that are defined here, `**cross_attention_kwargs` is empty
+
+        attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
+        unused_kwargs = [k for k, _ in cross_attention_kwargs.items() if k not in attn_parameters]
+        if len(unused_kwargs) > 0:
+            print(
+                f"cross_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored."
+            )
+        cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters}
+
+        return self.processor(
+            self,
+            hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            attention_mask=attention_mask,
+            **cross_attention_kwargs,
+        )
+    
+    def orig_forward(self, hidden_states, encoder_hidden_states=None):
+        q = self.to_q(hidden_states)
+        k = (
+            self.to_k(encoder_hidden_states)
+            if encoder_hidden_states is not None
+            else self.to_k(hidden_states)
+        )
+        v = (
+            self.to_v(encoder_hidden_states)
+            if encoder_hidden_states is not None
+            else self.to_v(hidden_states)
+        )
+        b, t, c = q.size()
+
+        q = q.view(q.size(0), q.size(1), self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(k.size(0), k.size(1), self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(v.size(0), v.size(1), self.num_heads, self.head_dim).transpose(1, 2)
+
+        # scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+        # attn_weights = torch.softmax(scores, dim=-1)
+        # attn_output = torch.matmul(attn_weights, v)
+
+        attn_output = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False, scale=self.scale,
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(b, t, c)
+
+        for layer in self.to_out:
+            attn_output = layer(attn_output)
+
+        return attn_output
+    
+    def set_processor(self, processor) -> None:
+        r"""
+        Set the attention processor to use.
+        Args:
+            processor (`AttnProcessor`):
+                The attention processor to use.
+        """
+        # if current processor is in `self._modules` and if passed `processor` is not, we need to
+        # pop `processor` from `self._modules`
+        if (
+            hasattr(self, "processor")
+            and isinstance(self.processor, torch.nn.Module)
+            and not isinstance(processor, torch.nn.Module)
+        ):
+            print(f"You are removing possibly trained weights of {self.processor} with {processor}")
+            self._modules.pop("processor")
+
+        self.processor = processor
+    
+    def get_processor(self, return_deprecated_lora: bool = False):
+        r"""
+        Get the attention processor in use.
+        Args:
+            return_deprecated_lora (`bool`, *optional*, defaults to `False`):
+                Set to `True` to return the deprecated LoRA attention processor.
+        Returns:
+            "AttentionProcessor": The attention processor in use.
+        """
+        if not return_deprecated_lora:
+            return self.processor
+
+        # TODO(Sayak, Patrick). The rest of the function is needed to ensure backwards compatible
+        # serialization format for LoRA Attention Processors. It should be deleted once the integration
+        # with PEFT is completed.
+        is_lora_activated = {
+            name: module.lora_layer is not None
+            for name, module in self.named_modules()
+            if hasattr(module, "lora_layer")
+        }
+
+        # 1. if no layer has a LoRA activated we can return the processor as usual
+        if not any(is_lora_activated.values()):
+            return self.processor
+
+        # If doesn't apply LoRA do `add_k_proj` or `add_v_proj`
+        is_lora_activated.pop("add_k_proj", None)
+        is_lora_activated.pop("add_v_proj", None)
+        # 2. else it is not possible that only some layers have LoRA activated
+        if not all(is_lora_activated.values()):
+            raise ValueError(
+                f"Make sure that either all layers or no layers have LoRA activated, but have {is_lora_activated}"
+            )
+
+        # 3. And we need to merge the current LoRA layers into the corresponding LoRA attention processor
+        non_lora_processor_cls_name = self.processor.__class__.__name__
+        lora_processor_cls = getattr(import_module(__name__), "LoRA" + non_lora_processor_cls_name)
+
+        hidden_size = self.inner_dim
+
+        # now create a LoRA attention processor from the LoRA layers
+        if lora_processor_cls in [LoRAAttnProcessor, LoRAAttnProcessor2_0, LoRAXFormersAttnProcessor]:
+            kwargs = {
+                "cross_attention_dim": self.cross_attention_dim,
+                "rank": self.to_q.lora_layer.rank,
+                "network_alpha": self.to_q.lora_layer.network_alpha,
+                "q_rank": self.to_q.lora_layer.rank,
+                "q_hidden_size": self.to_q.lora_layer.out_features,
+                "k_rank": self.to_k.lora_layer.rank,
+                "k_hidden_size": self.to_k.lora_layer.out_features,
+                "v_rank": self.to_v.lora_layer.rank,
+                "v_hidden_size": self.to_v.lora_layer.out_features,
+                "out_rank": self.to_out[0].lora_layer.rank,
+                "out_hidden_size": self.to_out[0].lora_layer.out_features,
+            }
+
+            if hasattr(self.processor, "attention_op"):
+                kwargs["attention_op"] = self.processor.attention_op
+
+            lora_processor = lora_processor_cls(hidden_size, **kwargs)
+            lora_processor.to_q_lora.load_state_dict(self.to_q.lora_layer.state_dict())
+            lora_processor.to_k_lora.load_state_dict(self.to_k.lora_layer.state_dict())
+            lora_processor.to_v_lora.load_state_dict(self.to_v.lora_layer.state_dict())
+            lora_processor.to_out_lora.load_state_dict(self.to_out[0].lora_layer.state_dict())
+        elif lora_processor_cls == LoRAAttnAddedKVProcessor:
+            lora_processor = lora_processor_cls(
+                hidden_size,
+                cross_attention_dim=self.add_k_proj.weight.shape[0],
+                rank=self.to_q.lora_layer.rank,
+                network_alpha=self.to_q.lora_layer.network_alpha,
+            )
+            lora_processor.to_q_lora.load_state_dict(self.to_q.lora_layer.state_dict())
+            lora_processor.to_k_lora.load_state_dict(self.to_k.lora_layer.state_dict())
+            lora_processor.to_v_lora.load_state_dict(self.to_v.lora_layer.state_dict())
+            lora_processor.to_out_lora.load_state_dict(self.to_out[0].lora_layer.state_dict())
+
+            # only save if used
+            if self.add_k_proj.lora_layer is not None:
+                lora_processor.add_k_proj_lora.load_state_dict(self.add_k_proj.lora_layer.state_dict())
+                lora_processor.add_v_proj_lora.load_state_dict(self.add_v_proj.lora_layer.state_dict())
+            else:
+                lora_processor.add_k_proj_lora = None
+                lora_processor.add_v_proj_lora = None
+        else:
+            raise ValueError(f"{lora_processor_cls} does not exist.")
+
+        return lora_processor
+    
+class GEGLU(nn.Module):
+    def __init__(self, in_features, out_features):
+        super(GEGLU, self).__init__()
+        self.proj = nn.Linear(in_features, out_features * 2, bias=True)
+
+    def forward(self, x):
+        x_proj = self.proj(x)
+        x1, x2 = x_proj.chunk(2, dim=-1)
+        return x1 * torch.nn.functional.gelu(x2)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, in_features, out_features):
+        super(FeedForward, self).__init__()
+
+        self.net = nn.ModuleList(
+            [
+                GEGLU(in_features, out_features * 4),
+                nn.Dropout(p=0.0, inplace=False),
+                nn.Linear(out_features * 4, out_features, bias=True),
+            ]
+        )
+
+    def forward(self, x):
+        for layer in self.net:
+            x = layer(x)
+        return x
+
+
+class BasicTransformerBlock(nn.Module):
+    def __init__(self, hidden_size):
+        super(BasicTransformerBlock, self).__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, eps=1e-05, elementwise_affine=True)
+        self.attn1 = Attention(hidden_size)
+        self.norm2 = nn.LayerNorm(hidden_size, eps=1e-05, elementwise_affine=True)
+        self.attn2 = Attention(hidden_size, 2048)
+        self.norm3 = nn.LayerNorm(hidden_size, eps=1e-05, elementwise_affine=True)
+        self.ff = FeedForward(hidden_size, hidden_size)
+
+    def forward(self, x, encoder_hidden_states=None):
+        residual = x
+
+        x = self.norm1(x)
+        x = self.attn1(x)
+        x = x + residual
+
+        residual = x
+
+        x = self.norm2(x)
+        if encoder_hidden_states is not None:
+            x = self.attn2(x, encoder_hidden_states)
+        else:
+            x = self.attn2(x)
+        x = x + residual
+
+        residual = x
+
+        x = self.norm3(x)
+        x = self.ff(x)
+        x = x + residual
+        return x
+
+
+class Transformer2DModel(nn.Module):
+    def __init__(self, in_channels, out_channels, n_layers):
+        super(Transformer2DModel, self).__init__()
+        self.norm = nn.GroupNorm(32, in_channels, eps=1e-06, affine=True)
+        self.proj_in = nn.Linear(in_channels, out_channels, bias=True)
+        self.transformer_blocks = nn.ModuleList(
+            [BasicTransformerBlock(out_channels) for _ in range(n_layers)]
+        )
+        self.proj_out = nn.Linear(out_channels, out_channels, bias=True)
+
+    def forward(self, hidden_states, encoder_hidden_states=None):
+        batch, _, height, width = hidden_states.shape
+        res = hidden_states
+        hidden_states = self.norm(hidden_states)
+        inner_dim = hidden_states.shape[1]
+        hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
+            batch, height * width, inner_dim
+        )
+        hidden_states = self.proj_in(hidden_states)
+
+        for block in self.transformer_blocks:
+            hidden_states = block(hidden_states, encoder_hidden_states)
+
+        hidden_states = self.proj_out(hidden_states)
+        hidden_states = (
+            hidden_states.reshape(batch, height, width, inner_dim)
+            .permute(0, 3, 1, 2)
+            .contiguous()
+        )
+
+        return hidden_states + res
+
+
+class Downsample2D(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(Downsample2D, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=2, padding=1
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class Upsample2D(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(Upsample2D, self).__init__()
+        self.conv = nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+
+    def forward(self, x):
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+        return self.conv(x)
+
+
+class DownBlock2D(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(DownBlock2D, self).__init__()
+        self.resnets = nn.ModuleList(
+            [
+                ResnetBlock2D(in_channels, out_channels, conv_shortcut=False),
+                ResnetBlock2D(out_channels, out_channels, conv_shortcut=False),
+            ]
+        )
+        self.downsamplers = nn.ModuleList([Downsample2D(out_channels, out_channels)])
+
+    def forward(self, hidden_states, temb):
+        output_states = []
+        for module in self.resnets:
+            hidden_states = module(hidden_states, temb)
+            output_states.append(hidden_states)
+
+        hidden_states = self.downsamplers[0](hidden_states)
+        output_states.append(hidden_states)
+
+        return hidden_states, output_states
+
+
+class CrossAttnDownBlock2D(nn.Module):
+    def __init__(self, in_channels, out_channels, n_layers, has_downsamplers=True):
+        super(CrossAttnDownBlock2D, self).__init__()
+        self.attentions = nn.ModuleList(
+            [
+                Transformer2DModel(out_channels, out_channels, n_layers),
+                Transformer2DModel(out_channels, out_channels, n_layers),
+            ]
+        )
+        self.resnets = nn.ModuleList(
+            [
+                ResnetBlock2D(in_channels, out_channels),
+                ResnetBlock2D(out_channels, out_channels, conv_shortcut=False),
+            ]
+        )
+        self.downsamplers = None
+        if has_downsamplers:
+            self.downsamplers = nn.ModuleList(
+                [Downsample2D(out_channels, out_channels)]
+            )
+
+    def forward(self, hidden_states, temb, encoder_hidden_states):
+        output_states = []
+        for resnet, attn in zip(self.resnets, self.attentions):
+            hidden_states = resnet(hidden_states, temb)
+            hidden_states = attn(
+                hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+            )
+            output_states.append(hidden_states)
+
+        if self.downsamplers is not None:
+            hidden_states = self.downsamplers[0](hidden_states)
+            output_states.append(hidden_states)
+
+        return hidden_states, output_states
+
+
+class CrossAttnUpBlock2D(nn.Module):
+    def __init__(self, in_channels, out_channels, prev_output_channel, n_layers):
+        super(CrossAttnUpBlock2D, self).__init__()
+        self.attentions = nn.ModuleList(
+            [
+                Transformer2DModel(out_channels, out_channels, n_layers),
+                Transformer2DModel(out_channels, out_channels, n_layers),
+                Transformer2DModel(out_channels, out_channels, n_layers),
+            ]
+        )
+        self.resnets = nn.ModuleList(
+            [
+                ResnetBlock2D(prev_output_channel + out_channels, out_channels),
+                ResnetBlock2D(2 * out_channels, out_channels),
+                ResnetBlock2D(out_channels + in_channels, out_channels),
+            ]
+        )
+        self.upsamplers = nn.ModuleList([Upsample2D(out_channels, out_channels)])
+
+    def forward(
+        self, hidden_states, res_hidden_states_tuple, temb, encoder_hidden_states
+    ):
+        for resnet, attn in zip(self.resnets, self.attentions):
+            # pop res hidden states
+            res_hidden_states = res_hidden_states_tuple[-1]
+            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+            hidden_states = resnet(hidden_states, temb)
+            hidden_states = attn(
+                hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+            )
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states)
+
+        return hidden_states
+
+
+class UpBlock2D(nn.Module):
+    def __init__(self, in_channels, out_channels, prev_output_channel):
+        super(UpBlock2D, self).__init__()
+        self.resnets = nn.ModuleList(
+            [
+                ResnetBlock2D(out_channels + prev_output_channel, out_channels),
+                ResnetBlock2D(out_channels * 2, out_channels),
+                ResnetBlock2D(out_channels + in_channels, out_channels),
+            ]
+        )
+
+    def forward(self, hidden_states, res_hidden_states_tuple, temb=None):
+
+        is_freeu_enabled = (
+            getattr(self, "s1", None)
+            and getattr(self, "s2", None)
+            and getattr(self, "b1", None)
+            and getattr(self, "b2", None)
+            and getattr(self, "resolution_idx", 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)
+            hidden_states = resnet(hidden_states, temb)
+
+        return hidden_states
+
+class UNetMidBlock2DCrossAttn(nn.Module):
+    def __init__(self, in_features):
+        super(UNetMidBlock2DCrossAttn, self).__init__()
+        self.attentions = nn.ModuleList(
+            [Transformer2DModel(in_features, in_features, n_layers=10)]
+        )
+        self.resnets = nn.ModuleList(
+            [
+                ResnetBlock2D(in_features, in_features, conv_shortcut=False),
+                ResnetBlock2D(in_features, in_features, conv_shortcut=False),
+            ]
+        )
+
+    def forward(self, hidden_states, temb=None, encoder_hidden_states=None):
+        hidden_states = self.resnets[0](hidden_states, temb)
+        for attn, resnet in zip(self.attentions, self.resnets[1:]):
+            hidden_states = attn(
+                hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+            )
+            hidden_states = resnet(hidden_states, temb)
+
+        return hidden_states
+
+
+class UNet2DConditionModel(nn.Module):
+    def __init__(self):
+        super(UNet2DConditionModel, self).__init__()
+
+        # This is needed to imitate huggingface config behavior
+        # has nothing to do with the model itself
+        # remove this if you don't use diffuser's pipeline
+        self.config = namedtuple(
+            "config", "in_channels addition_time_embed_dim sample_size"
+        )
+        self.config.in_channels = 4
+        self.config.addition_time_embed_dim = 256
+        self.config.sample_size = 128
+
+        self.conv_in = nn.Conv2d(4, 320, kernel_size=3, stride=1, padding=1)
+        self.time_proj = Timesteps()
+        self.time_embedding = TimestepEmbedding(in_features=320, out_features=1280)
+        self.add_time_proj = Timesteps(256)
+        self.add_embedding = TimestepEmbedding(in_features=2816, out_features=1280)
+        self.down_blocks = nn.ModuleList(
+            [
+                DownBlock2D(in_channels=320, out_channels=320),
+                CrossAttnDownBlock2D(in_channels=320, out_channels=640, n_layers=2),
+                CrossAttnDownBlock2D(
+                    in_channels=640,
+                    out_channels=1280,
+                    n_layers=10,
+                    has_downsamplers=False,
+                ),
+            ]
+        )
+        self.up_blocks = nn.ModuleList(
+            [
+                CrossAttnUpBlock2D(
+                    in_channels=640,
+                    out_channels=1280,
+                    prev_output_channel=1280,
+                    n_layers=10,
+                ),
+                CrossAttnUpBlock2D(
+                    in_channels=320,
+                    out_channels=640,
+                    prev_output_channel=1280,
+                    n_layers=2,
+                ),
+                UpBlock2D(in_channels=320, out_channels=320, prev_output_channel=640),
+            ]
+        )
+        self.mid_block = UNetMidBlock2DCrossAttn(1280)
+        self.conv_norm_out = nn.GroupNorm(32, 320, eps=1e-05, affine=True)
+        self.conv_act = nn.SiLU()
+        self.conv_out = nn.Conv2d(320, 4, kernel_size=3, stride=1, padding=1)
+
+    def forward(
+        self, sample, timesteps, encoder_hidden_states, added_cond_kwargs, **kwargs
+    ):
+        # Implement the forward pass through the model
+        timesteps = timesteps.expand(sample.shape[0])
+        t_emb = self.time_proj(timesteps).to(dtype=sample.dtype)
+
+        emb = self.time_embedding(t_emb)
+
+        text_embeds = added_cond_kwargs.get("text_embeds")
+        time_ids = added_cond_kwargs.get("time_ids")
+
+        time_embeds = self.add_time_proj(time_ids.flatten())
+        time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
+
+        add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
+        add_embeds = add_embeds.to(emb.dtype)
+        aug_emb = self.add_embedding(add_embeds)
+
+        emb = emb + aug_emb
+
+        sample = self.conv_in(sample)
+
+        # 3. down
+        s0 = sample
+        sample, [s1, s2, s3] = self.down_blocks[0](
+            sample,
+            temb=emb,
+        )
+
+        sample, [s4, s5, s6] = self.down_blocks[1](
+            sample,
+            temb=emb,
+            encoder_hidden_states=encoder_hidden_states,
+        )
+
+        sample, [s7, s8] = self.down_blocks[2](
+            sample,
+            temb=emb,
+            encoder_hidden_states=encoder_hidden_states,
+        )
+
+        # 4. mid
+        sample = self.mid_block(
+            sample, emb, encoder_hidden_states=encoder_hidden_states
+        )
+
+        # 5. up
+        sample = self.up_blocks[0](
+            hidden_states=sample,
+            temb=emb,
+            res_hidden_states_tuple=[s6, s7, s8],
+            encoder_hidden_states=encoder_hidden_states,
+        )
+
+        sample = self.up_blocks[1](
+            hidden_states=sample,
+            temb=emb,
+            res_hidden_states_tuple=[s3, s4, s5],
+            encoder_hidden_states=encoder_hidden_states,
+        )
+
+        sample = self.up_blocks[2](
+            hidden_states=sample,
+            temb=emb,
+            res_hidden_states_tuple=[s0, s1, s2],
+        )
+
+        # 6. post-process
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        return [sample]