From dc64dbfd3fbe8dacc4e88769dfd6316ee8e1a321 Mon Sep 17 00:00:00 2001 From: City <125218114+city96@users.noreply.github.com> Date: Sat, 14 Dec 2024 20:44:10 +0100 Subject: [PATCH] Reduce duplicate code --- comfy/ldm/pixart/blocks.py | 65 ++---------------------------------- comfy/ldm/pixart/pixart.py | 57 +------------------------------ comfy/ldm/pixart/pixartms.py | 51 ++++++++++------------------ 3 files changed, 20 insertions(+), 153 deletions(-) diff --git a/comfy/ldm/pixart/blocks.py b/comfy/ldm/pixart/blocks.py index 4aad7f54a..c2d2f1409 100644 --- a/comfy/ldm/pixart/blocks.py +++ b/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. diff --git a/comfy/ldm/pixart/pixart.py b/comfy/ldm/pixart/pixart.py index d79a59b9b..d2e0f45d0 100644 --- a/comfy/ldm/pixart/pixart.py +++ b/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 diff --git a/comfy/ldm/pixart/pixartms.py b/comfy/ldm/pixart/pixartms.py index d5edb2ab7..64f9a87f6 100644 --- a/comfy/ldm/pixart/pixartms.py +++ b/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 .pixart import PixArt, get_2d_sincos_pos_embed, get_2d_sincos_pos_embed_torch - -class PatchEmbed(nn.Module): - """ - 2D Image to Patch Embedding - """ - def __init__( - self, - patch_size=16, - 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 +from .blocks import ( + t2i_modulate, + CaptionEmbedder, + AttentionKVCompress, + MultiHeadCrossAttention, + T2IFinalLayer, + SizeEmbedder, +) +from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder, PatchEmbed, Mlp +from .pixart import PixArt, get_2d_sincos_pos_embed_torch 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, - dtype=dtype, device=device, operations=operations + patch_size=patch_size, + 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,