refactor

Author: a-r-r-o-w

scripts/convert_dcae_to_diffusers.py

@@ -19,7 +19,11 @@ def remap_qkv_(key: str, state_dict: Dict[str, Any]):
     state_dict[f"{parent_module}.to_q.weight"] = q.squeeze()
     state_dict[f"{parent_module}.to_k.weight"] = k.squeeze()
     state_dict[f"{parent_module}.to_v.weight"] = v.squeeze()
-    # state_dict[key.replace("qkv.conv", "to_qkv")] = state_dict.pop(key)
+
+
+def remap_proj_conv_(key: str, state_dict: Dict[str, Any]):
+    parent_module, _, _ = key.rpartition(".proj.conv.weight")
+    state_dict[f"{parent_module}.to_out.weight"] = state_dict.pop(key).squeeze()
 
 
 AE_KEYS_RENAME_DICT = {
@@ -40,7 +44,6 @@ def remap_qkv_(key: str, state_dict: Dict[str, Any]):
     "conv1.conv": "conv1",
     "conv2.conv": "conv2",
     "conv2.norm": "norm",
-    "proj.conv": "proj_out",
     "proj.norm": "norm_out",
     # encoder
     "encoder.project_in.conv": "encoder.conv_in",
@@ -76,6 +79,7 @@ def remap_qkv_(key: str, state_dict: Dict[str, Any]):
 
 AE_SPECIAL_KEYS_REMAP = {
     "qkv.conv.weight": remap_qkv_,
+    "proj.conv.weight": remap_proj_conv_,
 }
 
 

src/diffusers/models/attention_processor.py

@@ -752,7 +752,7 @@ def fuse_projections(self, fuse=True):
         self.fused_projections = fuse
 
 
-class MultiscaleAttentionProjection(nn.Module):
+class SanaMultiscaleAttentionProjection(nn.Module):
     def __init__(
         self,
         in_channels: int,
@@ -761,25 +761,24 @@ def __init__(
     ) -> None:
         super().__init__()
 
+        channels = 3 * in_channels
         self.proj_in = nn.Conv2d(
-            3 * in_channels,
-            3 * in_channels,
+            channels,
+            channels,
             kernel_size,
             padding=kernel_size // 2,
             groups=3 * in_channels,
             bias=False,
         )
-        self.proj_out = nn.Conv2d(
-            3 * in_channels, 3 * in_channels, 1, 1, 0, groups=3 * num_attention_heads, bias=False
-        )
+        self.proj_out = nn.Conv2d(channels, channels, 1, 1, 0, groups=3 * num_attention_heads, bias=False)
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         hidden_states = self.proj_in(hidden_states)
         hidden_states = self.proj_out(hidden_states)
         return hidden_states
 
 
-class MultiscaleLinearAttention(nn.Module):
+class SanaMultiscaleLinearAttention(nn.Module):
     r"""Lightweight multi-scale linear attention"""
 
     def __init__(
@@ -792,6 +791,7 @@ def __init__(
         norm_type: str = "batch_norm",
         kernel_sizes: Tuple[int, ...] = (5,),
         eps: float = 1e-15,
+        residual_connection: bool = False,
     ):
         super().__init__()
 
@@ -801,6 +801,7 @@ def __init__(
         self.eps = eps
         self.attention_head_dim = attention_head_dim
         self.norm_type = norm_type
+        self.residual_connection = residual_connection
 
         num_attention_heads = (
             int(in_channels // attention_head_dim * heads_ratio)
@@ -809,102 +810,32 @@ def __init__(
         )
         inner_dim = num_attention_heads * attention_head_dim
 
-        # self.to_qkv = nn.Conv2d(in_channels, 3 * inner_dim, 1, 1, 0, bias=False)
         self.to_q = nn.Linear(in_channels, inner_dim, bias=False)
         self.to_k = nn.Linear(in_channels, inner_dim, bias=False)
         self.to_v = nn.Linear(in_channels, inner_dim, bias=False)
 
         self.to_qkv_multiscale = nn.ModuleList()
         for kernel_size in kernel_sizes:
-            self.to_qkv_multiscale.append(MultiscaleAttentionProjection(inner_dim, num_attention_heads, kernel_size))
+            self.to_qkv_multiscale.append(
+                SanaMultiscaleAttentionProjection(inner_dim, num_attention_heads, kernel_size)
+            )
 
-        self.kernel_nonlinearity = nn.ReLU()
-        self.proj_out = nn.Conv2d(inner_dim * (1 + len(kernel_sizes)), out_channels, 1, 1, 0, bias=False)
+        self.nonlinearity = nn.ReLU()
+        self.to_out = nn.Linear(inner_dim * (1 + len(kernel_sizes)), out_channels, bias=False)
         self.norm_out = get_normalization(norm_type, num_features=out_channels)
 
-    def linear_attention(self, qkv: torch.Tensor) -> torch.Tensor:
-        batch_size, _, height, width = qkv.shape
-
-        qkv = qkv.float()
-        qkv = torch.reshape(qkv, (batch_size, -1, 3 * self.attention_head_dim, height * width))
-
-        query, key, value = (
-            qkv[:, :, 0 : self.attention_head_dim],
-            qkv[:, :, self.attention_head_dim : 2 * self.attention_head_dim],
-            qkv[:, :, 2 * self.attention_head_dim :],
-        )
-
-        # lightweight linear attention
-        query = self.kernel_nonlinearity(query)
-        key = self.kernel_nonlinearity(key)
-        value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1)
-
-        key_T = key.transpose(-1, -2)
-        scores = torch.matmul(value, key_T)
-        output = torch.matmul(scores, query)
-
-        output = output.float()
-        output = output[:, :, :-1] / (output[:, :, -1:] + self.eps)
-        output = torch.reshape(output, (batch_size, -1, height, width))
-
-        return output
-
-    def quadratic_attention(self, qkv: torch.Tensor) -> torch.Tensor:
-        batch_size, _, height, width = list(qkv.size())
-
-        qkv = torch.reshape(qkv, (batch_size, -1, 3 * self.attention_head_dim, height * width))
-        query, key, value = (
-            qkv[:, :, 0 : self.attention_head_dim],
-            qkv[:, :, self.attention_head_dim : 2 * self.attention_head_dim],
-            qkv[:, :, 2 * self.attention_head_dim :],
-        )
-
-        query = self.kernel_nonlinearity(query)
-        key = self.kernel_nonlinearity(key)
-
-        scores = torch.matmul(key.transpose(-1, -2), query)
-
-        original_dtype = scores.dtype
-        scores = scores.float()
-        scores = scores / (torch.sum(scores, dim=2, keepdim=True) + self.eps)
-        scores = scores.to(original_dtype)
-
-        output = torch.matmul(value, scores)
-        output = torch.reshape(output, (batch_size, -1, height, width))
-
-        return output
+        self.processor = SanaMultiscaleLinearAttnProcessor2_0()
+        self.processor_quadratic = SanaMultiscaleQuadraticAttnProcessor2_0()
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        residual = hidden_states
-
-        # qkv = self.to_qkv(hidden_states)
-        hidden_states = hidden_states.movedim(1, 3)
-        query = self.to_q(hidden_states)
-        key = self.to_k(hidden_states)
-        value = self.to_v(hidden_states)
-        qkv = torch.cat([query, key, value], dim=3)
-        qkv = qkv.movedim(3, 1)
-
-        multi_scale_qkv = [qkv]
-        for block in self.to_qkv_multiscale:
-            multi_scale_qkv.append(block(qkv))
+        height, width = hidden_states.shape[-2:]
 
-        qkv = torch.cat(multi_scale_qkv, dim=1)
-
-        height, width = qkv.shape[-2:]
         if height * width > self.attention_head_dim:
-            hidden_states = self.linear_attention(qkv).to(qkv.dtype)
-        else:
-            hidden_states = self.quadratic_attention(qkv)
-
-        hidden_states = self.proj_out(hidden_states)
-
-        if self.norm_type == "rms_norm":
-            hidden_states = self.norm_out(hidden_states.movedim(1, -1)).movedim(-1, 1)
+            hidden_states = self.processor(self, hidden_states)
         else:
-            hidden_states = self.norm_out(hidden_states)
+            hidden_states = self.processor_quadratic(self, hidden_states)
 
-        return hidden_states + residual
+        return hidden_states
 
 
 class AttnProcessor:
@@ -5160,6 +5091,109 @@ def __call__(
         return hidden_states
 
 
+class SanaMultiscaleLinearAttnProcessor2_0:
+    r"""
+    Processor for implementing multiscale linear attention.
+    """
+
+    def __call__(self, attn: SanaMultiscaleLinearAttention, hidden_states: torch.Tensor) -> torch.Tensor:
+        residual = hidden_states
+
+        batch_size, _, height, width = hidden_states.shape
+        original_dtype = hidden_states.dtype
+
+        hidden_states = hidden_states.movedim(1, -1)
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+        hidden_states = torch.cat([query, key, value], dim=3)
+        hidden_states = hidden_states.movedim(-1, 1)
+
+        multiscale_hidden_states = [hidden_states]
+        for block in attn.to_qkv_multiscale:
+            multiscale_hidden_states.append(block(hidden_states))
+
+        hidden_states = torch.cat(multiscale_hidden_states, dim=1)
+
+        hidden_states = hidden_states.to(dtype=torch.float32)
+        hidden_states = hidden_states.reshape(batch_size, -1, 3 * attn.attention_head_dim, height * width)
+
+        query, key, value = hidden_states.chunk(3, dim=2)
+        query = attn.nonlinearity(query)
+        key = attn.nonlinearity(key)
+        value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1)
+
+        scores = torch.matmul(value, key.transpose(-1, -2))
+        hidden_states = torch.matmul(scores, query)
+
+        hidden_states = hidden_states.to(dtype=torch.float32)
+        hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + attn.eps)
+        hidden_states = hidden_states.to(dtype=original_dtype)
+
+        hidden_states = torch.reshape(hidden_states, (batch_size, -1, height, width))
+        hidden_states = attn.to_out(hidden_states.movedim(1, -1)).movedim(-1, 1)
+
+        if attn.norm_type == "rms_norm":
+            hidden_states = attn.norm_out(hidden_states.movedim(1, -1)).movedim(-1, 1)
+        else:
+            hidden_states = attn.norm_out(hidden_states)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        return hidden_states
+
+
+class SanaMultiscaleQuadraticAttnProcessor2_0:
+    r"""
+    Processor for implementing multiscale quadratic attention.
+    """
+
+    def __call__(self, attn: SanaMultiscaleLinearAttention, hidden_states: torch.Tensor) -> torch.Tensor:
+        residual = hidden_states
+
+        batch_size, _, height, width = list(hidden_states.size())
+        original_dtype = hidden_states.dtype
+
+        hidden_states = hidden_states.movedim(1, -1)
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+        hidden_states = torch.cat([query, key, value], dim=3)
+        hidden_states = hidden_states.movedim(-1, 1)
+
+        multi_scale_qkv = [hidden_states]
+        for block in attn.to_qkv_multiscale:
+            multi_scale_qkv.append(block(hidden_states))
+
+        hidden_states = torch.cat(multi_scale_qkv, dim=1)
+
+        hidden_states = hidden_states.reshape(batch_size, -1, 3 * attn.attention_head_dim, height * width)
+
+        query, key, value = hidden_states.chunk(3, dim=2)
+        query = attn.nonlinearity(query)
+        key = attn.nonlinearity(key)
+
+        scores = torch.matmul(key.transpose(-1, -2), query)
+        scores = scores.to(dtype=torch.float32)
+        scores = scores / (torch.sum(scores, dim=2, keepdim=True) + attn.eps)
+        scores = scores.to(dtype=original_dtype)
+        hidden_states = torch.matmul(value, scores)
+
+        hidden_states = torch.reshape(hidden_states, (batch_size, -1, height, width))
+        hidden_states = attn.to_out(hidden_states.movedim(1, -1)).movedim(-1, 1)
+
+        if attn.norm_type == "rms_norm":
+            hidden_states = attn.norm_out(hidden_states.movedim(1, -1)).movedim(-1, 1)
+        else:
+            hidden_states = attn.norm_out(hidden_states)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        return hidden_states
+
+
 class LoRAAttnProcessor:
     def __init__(self):
         pass

src/diffusers/models/autoencoders/autoencoder_dc.py

@@ -13,15 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple, Union
+from typing import Tuple, Union
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 from ...configuration_utils import ConfigMixin, register_to_config
 from ..activations import get_activation
-from ..attention_processor import MultiscaleLinearAttention
+from ..attention_processor import SanaMultiscaleLinearAttention
 from ..modeling_utils import ModelMixin
 from ..normalization import RMSNorm, get_normalization
 
@@ -82,7 +82,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             hidden_states = self.norm(hidden_states.movedim(1, -1)).movedim(-1, 1)
         else:
             hidden_states = self.norm(hidden_states)
-        
+
         return hidden_states + residual
 
 
@@ -97,13 +97,14 @@ def __init__(
     ) -> None:
         super().__init__()
 
-        self.attn = MultiscaleLinearAttention(
+        self.attn = SanaMultiscaleLinearAttention(
             in_channels=in_channels,
             out_channels=in_channels,
             heads_ratio=heads_ratio,
             attention_head_dim=dim,
             norm_type=norm_type,
             kernel_sizes=qkv_multiscales,
+            residual_connection=True,
         )
 
         self.conv_out = GLUMBConv(

src/diffusers/models/normalization.py

@@ -574,7 +574,13 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         return F.normalize(hidden_states, p=self.p, dim=self.dim, eps=self.eps)
 
 
-def get_normalization(norm_type: str = "batch_norm", num_features: Optional[int] = None, eps: float = 1e-5, elementwise_affine: bool = True, bias: bool = True) -> nn.Module:
+def get_normalization(
+    norm_type: str = "batch_norm",
+    num_features: Optional[int] = None,
+    eps: float = 1e-5,
+    elementwise_affine: bool = True,
+    bias: bool = True,
+) -> nn.Module:
     if norm_type == "rms_norm":
         norm = RMSNorm(num_features, eps=eps, elementwise_affine=elementwise_affine, bias=bias)
     elif norm_type == "layer_norm":