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