Reduce duplicate code

comfy/ldm/pixart/blocks.py

@@ -8,6 +8,8 @@ import torch.nn.functional as F
 from einops import rearrange
 
 from comfy import model_management
+from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder, Mlp
+
 if model_management.xformers_enabled():
     import xformers.ops
     if int((xformers.__version__).split(".")[2]) >= 28:
@@ -254,48 +256,6 @@ class DecoderLayer(nn.Module):
         return x
 
 
-#################################################################################
-#               Embedding Layers for Timesteps and Class Labels                 #
-#################################################################################
-class TimestepEmbedder(nn.Module):
-    """
-    Embeds scalar timesteps into vector representations.
-    """
-    def __init__(self, hidden_size, frequency_embedding_size=256, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.mlp = nn.Sequential(
-            operations.Linear(frequency_embedding_size, hidden_size, bias=True, dtype=dtype, device=device),
-            nn.SiLU(),
-            operations.Linear(hidden_size, hidden_size, bias=True, dtype=dtype, device=device),
-        )
-        self.frequency_embedding_size = frequency_embedding_size
-
-    @staticmethod
-    def timestep_embedding(t, dim, max_period=10000):
-        """
-        Create sinusoidal timestep embeddings.
-        :param t: a 1-D Tensor of N indices, one per batch element.
-                          These may be fractional.
-        :param dim: the dimension of the output.
-        :param max_period: controls the minimum frequency of the embeddings.
-        :return: an (N, D) Tensor of positional embeddings.
-        """
-        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
-        half = dim // 2
-        freqs = torch.exp(
-            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half)
-        args = t[:, None].float() * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-        return embedding
-
-    def forward(self, t, dtype):
-        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
-        t_emb = self.mlp(t_freq.to(dtype))
-        return t_emb
-
-
 class SizeEmbedder(TimestepEmbedder):
     """
     Embeds scalar timesteps into vector representations.
@@ -355,27 +315,6 @@ class LabelEmbedder(nn.Module):
         return embeddings
 
 
-class Mlp(nn.Module):
-    """
-    Adapted from timm
-    """
-    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, bias=True, drop=None, dtype=None, device=None, operations=None) -> None:
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-
-        self.fc1 = operations.Linear(in_features, hidden_features, bias=bias, dtype=dtype, device=device)
-        self.act = act_layer()
-        self.fc2 = operations.Linear(hidden_features, out_features, bias=bias, dtype=dtype, device=device)
-
-        self.drop1 = nn.Identity()
-        self.drop2 = nn.Identity()
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.act(self.fc1(x))
-        return self.fc2(x)
-
-
 class CaptionEmbedder(nn.Module):
     """
     Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.

comfy/ldm/pixart/pixart.py

@@ -12,9 +12,8 @@ from .blocks import (
     MultiHeadCrossAttention,
     T2IFinalLayer,
     TimestepEmbedder,
-    Mlp
 )
-from comfy.ldm.modules.diffusionmodules.mmdit import PatchEmbed, get_1d_sincos_pos_embed_from_grid_torch
+from comfy.ldm.modules.diffusionmodules.mmdit import PatchEmbed, TimestepEmbedder, Mlp, get_1d_sincos_pos_embed_from_grid_torch
 
 
 class PixArtBlock(nn.Module):
@@ -196,63 +195,9 @@ def get_2d_sincos_pos_embed_torch(embed_dim, w, h, pe_interpolation=1.0, base_si
     grid_h, grid_w = torch.meshgrid(
         torch.arange(h, device=device, dtype=dtype) / (h/base_size) / pe_interpolation,
         torch.arange(w, device=device, dtype=dtype) / (w/base_size) / pe_interpolation,
-        # torch.linspace(-val_h + val_center, val_h + val_center, h, device=device, dtype=dtype),
-        # torch.linspace(-val_w + val_center, val_w + val_center, w, device=device, dtype=dtype),
         indexing='ij'
     )
     emb_h = get_1d_sincos_pos_embed_from_grid_torch(embed_dim // 2, grid_h, device=device, dtype=dtype)
     emb_w = get_1d_sincos_pos_embed_from_grid_torch(embed_dim // 2, grid_w, device=device, dtype=dtype)
     emb = torch.cat([emb_w, emb_h], dim=1)  # (H*W, D)
     return emb
-
-# Unused
-def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, pe_interpolation=1.0, base_size=16):
-    """
-    grid_size: int of the grid height and width
-    return:
-    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
-    """
-    if isinstance(grid_size, int):
-        grid_size = (grid_size, grid_size)
-    grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0]/base_size) / pe_interpolation
-    grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1]/base_size) / pe_interpolation
-    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
-    grid = np.stack(grid, axis=0)
-    grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])
-
-    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
-    if cls_token and extra_tokens > 0:
-        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
-    return pos_embed.astype(np.float32)
-
-
-def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
-    assert embed_dim % 2 == 0
-
-    # use half of dimensions to encode grid_h
-    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
-    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
-
-    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
-    return emb
-
-
-def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
-    """
-    embed_dim: output dimension for each position
-    pos: a list of positions to be encoded: size (M,)
-    out: (M, D)
-    """
-    assert embed_dim % 2 == 0
-    omega = np.arange(embed_dim // 2, dtype=np.float64)
-    omega /= embed_dim / 2.
-    omega = 1. / 10000 ** omega  # (D/2,)
-
-    pos = pos.reshape(-1)  # (M,)
-    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
-
-    emb_sin = np.sin(out)  # (M, D/2)
-    emb_cos = np.cos(out)  # (M, D/2)
-
-    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
-    return emb

comfy/ldm/pixart/pixartms.py

@@ -4,38 +4,16 @@
 import torch
 import torch.nn as nn
 
-from .blocks import t2i_modulate, CaptionEmbedder, AttentionKVCompress, MultiHeadCrossAttention, T2IFinalLayer, TimestepEmbedder, SizeEmbedder, Mlp
+from .blocks import (
-from .pixart import PixArt, get_2d_sincos_pos_embed, get_2d_sincos_pos_embed_torch
+    t2i_modulate,
-
+    CaptionEmbedder,
-class PatchEmbed(nn.Module):
+    AttentionKVCompress,
-    """
+    MultiHeadCrossAttention,
-    2D Image to Patch Embedding
+    T2IFinalLayer,
-    """
+    SizeEmbedder,
-    def __init__(
+)
-            self,
+from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder, PatchEmbed, Mlp
-            patch_size=16,
+from .pixart import PixArt, get_2d_sincos_pos_embed_torch
-            in_chans=3,
-            embed_dim=768,
-            norm_layer=None,
-            flatten=True,
-            bias=True,
-            dtype=None,
-            device=None,
-            operations=None
-    ):
-        super().__init__()
-        patch_size = (patch_size, patch_size)
-        self.patch_size = patch_size
-        self.flatten = flatten
-        self.proj = operations.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias, dtype=dtype, device=device)
-        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
-
-    def forward(self, x):
-        x = self.proj(x)
-        if self.flatten:
-            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
-        x = self.norm(x)
-        return x
 
 
 class PixArtMSBlock(nn.Module):
@@ -123,8 +101,13 @@ class PixArtMS(PixArt):
             operations.Linear(hidden_size, 6 * hidden_size, bias=True, dtype=dtype, device=device)
         )
         self.x_embedder = PatchEmbed(
-            patch_size, in_channels, hidden_size, bias=True,
+            patch_size=patch_size,
-            dtype=dtype, device=device, operations=operations
+            in_chans=in_channels,
+            embed_dim=hidden_size,
+            bias=True,
+            dtype=dtype,
+            device=device,
+            operations=operations
         )
         self.t_embedder = TimestepEmbedder(
             hidden_size, dtype=dtype, device=device, operations=operations,