915 lines
34 KiB
Python
915 lines
34 KiB
Python
# 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] |