From b5fe39211a36f5d5e77ae8e6d3cf08a59c9927ea Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 20 Dec 2024 17:43:50 -0500 Subject: [PATCH 1/4] Remove some useless code. --- comfy/ldm/pixart/pixart.py | 196 ----------------------------------- comfy/ldm/pixart/pixartms.py | 16 ++- 2 files changed, 13 insertions(+), 199 deletions(-) delete mode 100644 comfy/ldm/pixart/pixart.py diff --git a/comfy/ldm/pixart/pixart.py b/comfy/ldm/pixart/pixart.py deleted file mode 100644 index e1e61faf8..000000000 --- a/comfy/ldm/pixart/pixart.py +++ /dev/null @@ -1,196 +0,0 @@ -# Based on: -# https://github.com/PixArt-alpha/PixArt-alpha [Apache 2.0 license] -# https://github.com/PixArt-alpha/PixArt-sigma [Apache 2.0 license] -import torch -import torch.nn as nn - -from .blocks import ( - t2i_modulate, - CaptionEmbedder, - AttentionKVCompress, - MultiHeadCrossAttention, - T2IFinalLayer, -) -from comfy.ldm.modules.diffusionmodules.mmdit import PatchEmbed, TimestepEmbedder, Mlp, get_1d_sincos_pos_embed_from_grid_torch - - -class PixArtBlock(nn.Module): - """ - A PixArt block with adaptive layer norm (adaLN-single) conditioning. - """ - def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, drop_path=0, input_size=None, sampling=None, sr_ratio=1, qk_norm=False, **block_kwargs): - super().__init__() - self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.attn = AttentionKVCompress( - hidden_size, num_heads=num_heads, qkv_bias=True, sampling=sampling, sr_ratio=sr_ratio, - qk_norm=qk_norm, **block_kwargs - ) - self.cross_attn = MultiHeadCrossAttention(hidden_size, num_heads, **block_kwargs) - self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - # to be compatible with lower version pytorch - approx_gelu = lambda: nn.GELU(approximate="tanh") - self.mlp = Mlp(in_features=hidden_size, hidden_features=int(hidden_size * mlp_ratio), act_layer=approx_gelu, drop=0) - self.drop_path = nn.Identity() #DropPath(drop_path) if drop_path > 0. else nn.Identity() - self.scale_shift_table = nn.Parameter(torch.randn(6, hidden_size) / hidden_size ** 0.5) - self.sampling = sampling - self.sr_ratio = sr_ratio - - def forward(self, x, y, t, mask=None, **kwargs): - B, N, C = x.shape - - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.scale_shift_table[None] + t.reshape(B, 6, -1)).chunk(6, dim=1) - x = x + self.drop_path(gate_msa * self.attn(t2i_modulate(self.norm1(x), shift_msa, scale_msa)).reshape(B, N, C)) - x = x + self.cross_attn(x, y, mask) - x = x + self.drop_path(gate_mlp * self.mlp(t2i_modulate(self.norm2(x), shift_mlp, scale_mlp))) - - return x - - -### Core PixArt Model ### -class PixArt(nn.Module): - """ - Diffusion model with a Transformer backbone. - """ - def __init__( - self, - input_size=32, - patch_size=2, - in_channels=4, - hidden_size=1152, - depth=28, - num_heads=16, - mlp_ratio=4.0, - class_dropout_prob=0.1, - pred_sigma=True, - drop_path: float = 0., - caption_channels=4096, - pe_interpolation=1.0, - pe_precision=None, - config=None, - model_max_length=120, - qk_norm=False, - kv_compress_config=None, - **kwargs, - ): - super().__init__() - self.pred_sigma = pred_sigma - self.in_channels = in_channels - self.out_channels = in_channels * 2 if pred_sigma else in_channels - self.patch_size = patch_size - self.num_heads = num_heads - self.pe_interpolation = pe_interpolation - self.pe_precision = pe_precision - self.depth = depth - - self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True) - self.t_embedder = TimestepEmbedder(hidden_size) - num_patches = self.x_embedder.num_patches - self.base_size = input_size // self.patch_size - # Will use fixed sin-cos embedding: - self.register_buffer("pos_embed", torch.zeros(1, num_patches, hidden_size)) - - approx_gelu = lambda: nn.GELU(approximate="tanh") - self.t_block = nn.Sequential( - nn.SiLU(), - nn.Linear(hidden_size, 6 * hidden_size, bias=True) - ) - self.y_embedder = CaptionEmbedder( - in_channels=caption_channels, hidden_size=hidden_size, uncond_prob=class_dropout_prob, - act_layer=approx_gelu, token_num=model_max_length - ) - drop_path = [x.item() for x in torch.linspace(0, drop_path, depth)] # stochastic depth decay rule - self.kv_compress_config = kv_compress_config - if kv_compress_config is None: - self.kv_compress_config = { - 'sampling': None, - 'scale_factor': 1, - 'kv_compress_layer': [], - } - self.blocks = nn.ModuleList([ - PixArtBlock( - hidden_size, num_heads, mlp_ratio=mlp_ratio, drop_path=drop_path[i], - input_size=(input_size // patch_size, input_size // patch_size), - sampling=self.kv_compress_config['sampling'], - sr_ratio=int( - self.kv_compress_config['scale_factor'] - ) if i in self.kv_compress_config['kv_compress_layer'] else 1, - qk_norm=qk_norm, - ) - for i in range(depth) - ]) - self.final_layer = T2IFinalLayer(hidden_size, patch_size, self.out_channels) - - def forward_raw(self, x, t, y, mask=None, data_info=None): - """ - Original forward pass of PixArt. - x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) - t: (N,) tensor of diffusion timesteps - y: (N, 1, 120, C) tensor of class labels - """ - x = self.x_embedder(x) + pos_embed # (N, T, D), where T = H * W / patch_size ** 2 - t = self.t_embedder(timestep) # (N, D) - t0 = self.t_block(t) - y = self.y_embedder(y, self.training) # (N, 1, L, D) - if mask is not None: - if mask.shape[0] != y.shape[0]: - mask = mask.repeat(y.shape[0] // mask.shape[0], 1) - mask = mask.squeeze(1).squeeze(1) - y = y.squeeze(1).masked_select(mask.unsqueeze(-1) != 0).view(1, -1, x.shape[-1]) - y_lens = mask.sum(dim=1).tolist() - else: - y_lens = None - y = y.squeeze(1).view(1, -1, x.shape[-1]) - for block in self.blocks: - x = block(x, y, t0, y_lens) # (N, T, D) - x = self.final_layer(x, t) # (N, T, patch_size ** 2 * out_channels) - x = self.unpatchify(x) # (N, out_channels, H, W) - return x - - def forward(self, x, timesteps, context, y=None, **kwargs): - """ - Forward pass that adapts comfy input to original forward function - x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) - timesteps: (N,) tensor of diffusion timesteps - context: (N, 1, 120, C) conditioning - y: extra conditioning. - """ - ## Still accepts the input w/o that dim but returns garbage - if len(context.shape) == 3: - context = context.unsqueeze(1) - - ## run original forward pass - out = self.forward_raw( - x = x, - t = timesteps, - y = context, - ) - - ## only return EPS - eps, _ = out[:, :self.in_channels], out[:, self.in_channels:] - return eps - - def unpatchify(self, x): - """ - x: (N, T, patch_size**2 * C) - imgs: (N, H, W, C) - """ - c = self.out_channels - p = self.x_embedder.patch_size[0] - h = w = int(x.shape[1] ** 0.5) - assert h * w == x.shape[1] - - x = x.reshape(shape=(x.shape[0], h, w, p, p, c)) - x = torch.einsum('nhwpqc->nchpwq', x) - imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p)) - return imgs - -def get_2d_sincos_pos_embed_torch(embed_dim, w, h, pe_interpolation=1.0, base_size=16, device=None, dtype=torch.float32): - 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, - 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 diff --git a/comfy/ldm/pixart/pixartms.py b/comfy/ldm/pixart/pixartms.py index 8ff0d0a41..50dc58c23 100644 --- a/comfy/ldm/pixart/pixartms.py +++ b/comfy/ldm/pixart/pixartms.py @@ -12,10 +12,20 @@ from .blocks import ( T2IFinalLayer, SizeEmbedder, ) -from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder, PatchEmbed, Mlp -from .pixart import PixArt, get_2d_sincos_pos_embed_torch +from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder, PatchEmbed, Mlp, get_1d_sincos_pos_embed_from_grid_torch +def get_2d_sincos_pos_embed_torch(embed_dim, w, h, pe_interpolation=1.0, base_size=16, device=None, dtype=torch.float32): + 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, + 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 + class PixArtMSBlock(nn.Module): """ A PixArt block with adaptive layer norm zero (adaLN-Zero) conditioning. @@ -53,7 +63,7 @@ class PixArtMSBlock(nn.Module): ### Core PixArt Model ### -class PixArtMS(PixArt): +class PixArtMS(nn.Module): """ Diffusion model with a Transformer backbone. """ From c86cd585731a86515f57e325ddcd5eb918a61479 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 20 Dec 2024 17:50:03 -0500 Subject: [PATCH 2/4] Remove useless code. --- comfy/ldm/pixart/blocks.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/comfy/ldm/pixart/blocks.py b/comfy/ldm/pixart/blocks.py index f60bfd797..48b27008c 100644 --- a/comfy/ldm/pixart/blocks.py +++ b/comfy/ldm/pixart/blocks.py @@ -147,7 +147,6 @@ class AttentionKVCompress(nn.Module): qkv = self.qkv(x).reshape(B, N, 3, C) q, k, v = qkv.unbind(2) - dtype = q.dtype q = self.q_norm(q) k = self.k_norm(k) @@ -209,7 +208,6 @@ class T2IFinalLayer(nn.Module): self.out_channels = out_channels def forward(self, x, t): - dtype = x.dtype shift, scale = (self.scale_shift_table[None].to(dtype=x.dtype, device=x.device) + t[:, None]).chunk(2, dim=1) x = t2i_modulate(self.norm_final(x), shift, scale) x = self.linear(x) From da13b6b827d9291718a51fc9cc4ea54266a43f23 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Fri, 20 Dec 2024 18:02:12 -0500 Subject: [PATCH 3/4] Get rid of meshgrid warning. --- comfy/ldm/lightricks/symmetric_patchifier.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/ldm/lightricks/symmetric_patchifier.py b/comfy/ldm/lightricks/symmetric_patchifier.py index 51ce50589..c58dfb20b 100644 --- a/comfy/ldm/lightricks/symmetric_patchifier.py +++ b/comfy/ldm/lightricks/symmetric_patchifier.py @@ -53,7 +53,7 @@ class Patchifier(ABC): grid_h = torch.arange(h, dtype=torch.float32, device=device) grid_w = torch.arange(w, dtype=torch.float32, device=device) grid_f = torch.arange(f, dtype=torch.float32, device=device) - grid = torch.meshgrid(grid_f, grid_h, grid_w) + grid = torch.meshgrid(grid_f, grid_h, grid_w, indexing='ij') grid = torch.stack(grid, dim=0) grid = grid.unsqueeze(0).repeat(batch_size, 1, 1, 1, 1) From 1419dee9156051aed322c55c7ad8b14fddb7d2d3 Mon Sep 17 00:00:00 2001 From: Qiacheng Li Date: Fri, 20 Dec 2024 15:04:03 -0800 Subject: [PATCH 4/4] Update README.md for Intel GPUs (#6069) --- README.md | 36 ++++++++++++++++++++++++------------ 1 file changed, 24 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 1ec39d06f..844bd2866 100644 --- a/README.md +++ b/README.md @@ -149,6 +149,30 @@ This is the command to install the nightly with ROCm 6.2 which might have some p ```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.2.4``` +### Intel GPUs (Windows and Linux) + +(Option 1) Intel Arc GPU users can install native PyTorch with torch.xpu support using pip (currently available in PyTorch nightly builds). More information can be found [here](https://pytorch.org/docs/main/notes/get_start_xpu.html) + +1. To install PyTorch nightly, use the following command: + +```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/xpu``` + +2. Launch ComfyUI by running `python main.py` + + +(Option 2) Alternatively, Intel GPUs supported by Intel Extension for PyTorch (IPEX) can leverage IPEX for improved performance. + +1. For Intel® Arc™ A-Series Graphics utilizing IPEX, create a conda environment and use the commands below: + +``` +conda install libuv +pip install torch==2.3.1.post0+cxx11.abi torchvision==0.18.1.post0+cxx11.abi torchaudio==2.3.1.post0+cxx11.abi intel-extension-for-pytorch==2.3.110.post0+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/ --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/cn/ +``` + +For other supported Intel GPUs with IPEX, visit [Installation](https://intel.github.io/intel-extension-for-pytorch/index.html#installation?platform=gpu) for more information. + +Additional discussion and help can be found [here](https://github.com/comfyanonymous/ComfyUI/discussions/476). + ### NVIDIA Nvidia users should install stable pytorch using this command: @@ -177,17 +201,6 @@ After this you should have everything installed and can proceed to running Comfy ### Others: -#### Intel GPUs - -Intel GPU support is available for all Intel GPUs supported by Intel's Extension for Pytorch (IPEX) with the support requirements listed in the [Installation](https://intel.github.io/intel-extension-for-pytorch/index.html#installation?platform=gpu) page. Choose your platform and method of install and follow the instructions. The steps are as follows: - -1. Start by installing the drivers or kernel listed or newer in the Installation page of IPEX linked above for Windows and Linux if needed. -1. Follow the instructions to install [Intel's oneAPI Basekit](https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit-download.html) for your platform. -1. Install the packages for IPEX using the instructions provided in the Installation page for your platform. -1. Follow the [ComfyUI manual installation](#manual-install-windows-linux) instructions for Windows and Linux and run ComfyUI normally as described above after everything is installed. - -Additional discussion and help can be found [here](https://github.com/comfyanonymous/ComfyUI/discussions/476). - #### Apple Mac silicon You can install ComfyUI in Apple Mac silicon (M1 or M2) with any recent macOS version. @@ -308,4 +321,3 @@ This will use a snapshot of the legacy frontend preserved in the [ComfyUI Legacy ### Which GPU should I buy for this? [See this page for some recommendations](https://github.com/comfyanonymous/ComfyUI/wiki/Which-GPU-should-I-buy-for-ComfyUI) -