| import math |
| from typing import List, Optional, Tuple |
|
|
| import torch |
| import torch.nn as nn |
| from torch import Tensor |
|
|
| from comfy.ldm.modules.diffusionmodules.mmdit import DismantledBlock, PatchEmbed, VectorEmbedder, TimestepEmbedder, get_2d_sincos_pos_embed_torch |
|
|
|
|
| class ControlNetEmbedder(nn.Module): |
|
|
| def __init__( |
| self, |
| img_size: int, |
| patch_size: int, |
| in_chans: int, |
| attention_head_dim: int, |
| num_attention_heads: int, |
| adm_in_channels: int, |
| num_layers: int, |
| main_model_double: int, |
| double_y_emb: bool, |
| device: torch.device, |
| dtype: torch.dtype, |
| pos_embed_max_size: Optional[int] = None, |
| operations = None, |
| ): |
| super().__init__() |
| self.main_model_double = main_model_double |
| self.dtype = dtype |
| self.hidden_size = num_attention_heads * attention_head_dim |
| self.patch_size = patch_size |
| self.x_embedder = PatchEmbed( |
| img_size=img_size, |
| patch_size=patch_size, |
| in_chans=in_chans, |
| embed_dim=self.hidden_size, |
| strict_img_size=pos_embed_max_size is None, |
| device=device, |
| dtype=dtype, |
| operations=operations, |
| ) |
|
|
| self.t_embedder = TimestepEmbedder(self.hidden_size, dtype=dtype, device=device, operations=operations) |
|
|
| self.double_y_emb = double_y_emb |
| if self.double_y_emb: |
| self.orig_y_embedder = VectorEmbedder( |
| adm_in_channels, self.hidden_size, dtype, device, operations=operations |
| ) |
| self.y_embedder = VectorEmbedder( |
| self.hidden_size, self.hidden_size, dtype, device, operations=operations |
| ) |
| else: |
| self.y_embedder = VectorEmbedder( |
| adm_in_channels, self.hidden_size, dtype, device, operations=operations |
| ) |
|
|
| self.transformer_blocks = nn.ModuleList( |
| DismantledBlock( |
| hidden_size=self.hidden_size, num_heads=num_attention_heads, qkv_bias=True, |
| dtype=dtype, device=device, operations=operations |
| ) |
| for _ in range(num_layers) |
| ) |
|
|
| |
| |
| self.use_y_embedder = True |
|
|
| self.controlnet_blocks = nn.ModuleList([]) |
| for _ in range(len(self.transformer_blocks)): |
| controlnet_block = operations.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device) |
| self.controlnet_blocks.append(controlnet_block) |
|
|
| self.pos_embed_input = PatchEmbed( |
| img_size=img_size, |
| patch_size=patch_size, |
| in_chans=in_chans, |
| embed_dim=self.hidden_size, |
| strict_img_size=False, |
| device=device, |
| dtype=dtype, |
| operations=operations, |
| ) |
|
|
| def forward( |
| self, |
| x: torch.Tensor, |
| timesteps: torch.Tensor, |
| y: Optional[torch.Tensor] = None, |
| context: Optional[torch.Tensor] = None, |
| hint = None, |
| ) -> Tuple[Tensor, List[Tensor]]: |
| x_shape = list(x.shape) |
| x = self.x_embedder(x) |
| if not self.double_y_emb: |
| h = (x_shape[-2] + 1) // self.patch_size |
| w = (x_shape[-1] + 1) // self.patch_size |
| x += get_2d_sincos_pos_embed_torch(self.hidden_size, w, h, device=x.device) |
| c = self.t_embedder(timesteps, dtype=x.dtype) |
| if y is not None and self.y_embedder is not None: |
| if self.double_y_emb: |
| y = self.orig_y_embedder(y) |
| y = self.y_embedder(y) |
| c = c + y |
|
|
| x = x + self.pos_embed_input(hint) |
|
|
| block_out = () |
|
|
| repeat = math.ceil(self.main_model_double / len(self.transformer_blocks)) |
| for i in range(len(self.transformer_blocks)): |
| out = self.transformer_blocks[i](x, c) |
| if not self.double_y_emb: |
| x = out |
| block_out += (self.controlnet_blocks[i](out),) * repeat |
|
|
| return {"output": block_out} |
|
|
