|
import math |
|
from typing import List, Optional, Tuple |
|
|
|
import numpy as np |
|
import torch |
|
import torch.nn as nn |
|
from einops import rearrange |
|
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} |
|
|
