import torch from torch import nn import math import comfy.ldm.common_dit from comfy.ldm.modules.attention import optimized_attention from comfy.ldm.flux.math import apply_rope1 from comfy.ldm.flux.layers import EmbedND def attention(q, k, v, heads, transformer_options={}): return optimized_attention( q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), heads=heads, skip_reshape=True, transformer_options=transformer_options ) def apply_scale_shift_norm(norm, x, scale, shift): return torch.addcmul(shift, norm(x), scale + 1.0) def apply_gate_sum(x, out, gate): return torch.addcmul(x, gate, out) def get_shift_scale_gate(params): shift, scale, gate = torch.chunk(params, 3, dim=-1) return tuple(x.unsqueeze(1) for x in (shift, scale, gate)) def get_freqs(dim, max_period=10000.0): return torch.exp(-math.log(max_period) * torch.arange(start=0, end=dim, dtype=torch.float32) / dim) class TimeEmbeddings(nn.Module): def __init__(self, model_dim, time_dim, max_period=10000.0, operation_settings=None): super().__init__() assert model_dim % 2 == 0 self.model_dim = model_dim self.max_period = max_period self.register_buffer("freqs", get_freqs(model_dim // 2, max_period), persistent=False) operations = operation_settings.get("operations") self.in_layer = operations.Linear(model_dim, time_dim, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.activation = nn.SiLU() self.out_layer = operations.Linear(time_dim, time_dim, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) def forward(self, timestep, dtype): args = torch.outer(timestep, self.freqs.to(device=timestep.device)) time_embed = torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(dtype) time_embed = self.out_layer(self.activation(self.in_layer(time_embed))) return time_embed class TextEmbeddings(nn.Module): def __init__(self, text_dim, model_dim, operation_settings=None): super().__init__() operations = operation_settings.get("operations") self.in_layer = operations.Linear(text_dim, model_dim, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.norm = operations.LayerNorm(model_dim, elementwise_affine=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) def forward(self, text_embed): text_embed = self.in_layer(text_embed) return self.norm(text_embed).type_as(text_embed) class VisualEmbeddings(nn.Module): def __init__(self, visual_dim, model_dim, patch_size, operation_settings=None): super().__init__() self.patch_size = patch_size operations = operation_settings.get("operations") self.in_layer = operations.Linear(visual_dim, model_dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) def forward(self, x): x = x.movedim(1, -1) # B C T H W -> B T H W C B, T, H, W, dim = x.shape pt, ph, pw = self.patch_size x = x.view( B, T // pt, pt, H // ph, ph, W // pw, pw, dim, ).permute(0, 1, 3, 5, 2, 4, 6, 7).flatten(4, 7) return self.in_layer(x) class Modulation(nn.Module): def __init__(self, time_dim, model_dim, num_params, operation_settings=None): super().__init__() self.activation = nn.SiLU() self.out_layer = operation_settings.get("operations").Linear(time_dim, num_params * model_dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) def forward(self, x): return self.out_layer(self.activation(x)) class SelfAttention(nn.Module): def __init__(self, num_channels, head_dim, operation_settings=None): super().__init__() assert num_channels % head_dim == 0 self.num_heads = num_channels // head_dim self.head_dim = head_dim operations = operation_settings.get("operations") self.to_query = operations.Linear(num_channels, num_channels, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.to_key = operations.Linear(num_channels, num_channels, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.to_value = operations.Linear(num_channels, num_channels, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.query_norm = operations.RMSNorm(head_dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.key_norm = operations.RMSNorm(head_dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.out_layer = operations.Linear(num_channels, num_channels, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.num_chunks = 2 def _compute_qk(self, x, freqs, proj_fn, norm_fn): result = proj_fn(x).view(*x.shape[:-1], self.num_heads, -1) return apply_rope1(norm_fn(result), freqs) def _forward(self, x, freqs, transformer_options={}): q = self._compute_qk(x, freqs, self.to_query, self.query_norm) k = self._compute_qk(x, freqs, self.to_key, self.key_norm) v = self.to_value(x).view(*x.shape[:-1], self.num_heads, -1) out = attention(q, k, v, self.num_heads, transformer_options=transformer_options) return self.out_layer(out) def _forward_chunked(self, x, freqs, transformer_options={}): def process_chunks(proj_fn, norm_fn): x_chunks = torch.chunk(x, self.num_chunks, dim=1) freqs_chunks = torch.chunk(freqs, self.num_chunks, dim=1) chunks = [] for x_chunk, freqs_chunk in zip(x_chunks, freqs_chunks): chunks.append(self._compute_qk(x_chunk, freqs_chunk, proj_fn, norm_fn)) return torch.cat(chunks, dim=1) q = process_chunks(self.to_query, self.query_norm) k = process_chunks(self.to_key, self.key_norm) v = self.to_value(x).view(*x.shape[:-1], self.num_heads, -1) out = attention(q, k, v, self.num_heads, transformer_options=transformer_options) return self.out_layer(out) def forward(self, x, freqs, transformer_options={}): if x.shape[1] > 8192: return self._forward_chunked(x, freqs, transformer_options=transformer_options) else: return self._forward(x, freqs, transformer_options=transformer_options) class CrossAttention(SelfAttention): def get_qkv(self, x, context): q = self.to_query(x).view(*x.shape[:-1], self.num_heads, -1) k = self.to_key(context).view(*context.shape[:-1], self.num_heads, -1) v = self.to_value(context).view(*context.shape[:-1], self.num_heads, -1) return q, k, v def forward(self, x, context, transformer_options={}): q, k, v = self.get_qkv(x, context) out = attention(self.query_norm(q), self.key_norm(k), v, self.num_heads, transformer_options=transformer_options) return self.out_layer(out) class FeedForward(nn.Module): def __init__(self, dim, ff_dim, operation_settings=None): super().__init__() operations = operation_settings.get("operations") self.in_layer = operations.Linear(dim, ff_dim, bias=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.activation = nn.GELU() self.out_layer = operations.Linear(ff_dim, dim, bias=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.num_chunks = 4 def _forward(self, x): return self.out_layer(self.activation(self.in_layer(x))) def _forward_chunked(self, x): chunks = torch.chunk(x, self.num_chunks, dim=1) output_chunks = [] for chunk in chunks: output_chunks.append(self._forward(chunk)) return torch.cat(output_chunks, dim=1) def forward(self, x): if x.shape[1] > 8192: return self._forward_chunked(x) else: return self._forward(x) class OutLayer(nn.Module): def __init__(self, model_dim, time_dim, visual_dim, patch_size, operation_settings=None): super().__init__() self.patch_size = patch_size self.modulation = Modulation(time_dim, model_dim, 2, operation_settings=operation_settings) operations = operation_settings.get("operations") self.norm = operations.LayerNorm(model_dim, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.out_layer = operations.Linear(model_dim, math.prod(patch_size) * visual_dim, bias=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) def forward(self, visual_embed, time_embed): B, T, H, W, _ = visual_embed.shape shift, scale = torch.chunk(self.modulation(time_embed), 2, dim=-1) scale = scale[:, None, None, None, :] shift = shift[:, None, None, None, :] visual_embed = apply_scale_shift_norm(self.norm, visual_embed, scale, shift) x = self.out_layer(visual_embed) out_dim = x.shape[-1] // (self.patch_size[0] * self.patch_size[1] * self.patch_size[2]) x = x.view( B, T, H, W, out_dim, self.patch_size[0], self.patch_size[1], self.patch_size[2] ) return x.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(2, 3).flatten(3, 4).flatten(4, 5) class TransformerEncoderBlock(nn.Module): def __init__(self, model_dim, time_dim, ff_dim, head_dim, operation_settings=None): super().__init__() self.text_modulation = Modulation(time_dim, model_dim, 6, operation_settings=operation_settings) operations = operation_settings.get("operations") self.self_attention_norm = operations.LayerNorm(model_dim, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.self_attention = SelfAttention(model_dim, head_dim, operation_settings=operation_settings) self.feed_forward_norm = operations.LayerNorm(model_dim, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.feed_forward = FeedForward(model_dim, ff_dim, operation_settings=operation_settings) def forward(self, x, time_embed, freqs, transformer_options={}): self_attn_params, ff_params = torch.chunk(self.text_modulation(time_embed), 2, dim=-1) shift, scale, gate = get_shift_scale_gate(self_attn_params) out = apply_scale_shift_norm(self.self_attention_norm, x, scale, shift) out = self.self_attention(out, freqs, transformer_options=transformer_options) x = apply_gate_sum(x, out, gate) shift, scale, gate = get_shift_scale_gate(ff_params) out = apply_scale_shift_norm(self.feed_forward_norm, x, scale, shift) out = self.feed_forward(out) x = apply_gate_sum(x, out, gate) return x class TransformerDecoderBlock(nn.Module): def __init__(self, model_dim, time_dim, ff_dim, head_dim, operation_settings=None): super().__init__() self.visual_modulation = Modulation(time_dim, model_dim, 9, operation_settings=operation_settings) operations = operation_settings.get("operations") self.self_attention_norm = operations.LayerNorm(model_dim, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.self_attention = SelfAttention(model_dim, head_dim, operation_settings=operation_settings) self.cross_attention_norm = operations.LayerNorm(model_dim, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.cross_attention = CrossAttention(model_dim, head_dim, operation_settings=operation_settings) self.feed_forward_norm = operations.LayerNorm(model_dim, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) self.feed_forward = FeedForward(model_dim, ff_dim, operation_settings=operation_settings) def forward(self, visual_embed, text_embed, time_embed, freqs, transformer_options={}): self_attn_params, cross_attn_params, ff_params = torch.chunk(self.visual_modulation(time_embed), 3, dim=-1) # self attention shift, scale, gate = get_shift_scale_gate(self_attn_params) visual_out = apply_scale_shift_norm(self.self_attention_norm, visual_embed, scale, shift) visual_out = self.self_attention(visual_out, freqs, transformer_options=transformer_options) visual_embed = apply_gate_sum(visual_embed, visual_out, gate) # cross attention shift, scale, gate = get_shift_scale_gate(cross_attn_params) visual_out = apply_scale_shift_norm(self.cross_attention_norm, visual_embed, scale, shift) visual_out = self.cross_attention(visual_out, text_embed, transformer_options=transformer_options) visual_embed = apply_gate_sum(visual_embed, visual_out, gate) # feed forward shift, scale, gate = get_shift_scale_gate(ff_params) visual_out = apply_scale_shift_norm(self.feed_forward_norm, visual_embed, scale, shift) visual_out = self.feed_forward(visual_out) visual_embed = apply_gate_sum(visual_embed, visual_out, gate) return visual_embed class Kandinsky5(nn.Module): def __init__( self, in_visual_dim=16, out_visual_dim=16, in_text_dim=3584, in_text_dim2=768, time_dim=512, model_dim=1792, ff_dim=7168, visual_embed_dim=132, patch_size=(1, 2, 2), num_text_blocks=2, num_visual_blocks=32, axes_dims=(16, 24, 24), rope_scale_factor=(1.0, 2.0, 2.0), dtype=None, device=None, operations=None, **kwargs ): super().__init__() head_dim = sum(axes_dims) self.rope_scale_factor = rope_scale_factor self.in_visual_dim = in_visual_dim self.model_dim = model_dim self.patch_size = patch_size self.visual_embed_dim = visual_embed_dim self.dtype = dtype self.device = device operation_settings = {"operations": operations, "device": device, "dtype": dtype} self.time_embeddings = TimeEmbeddings(model_dim, time_dim, operation_settings=operation_settings) self.text_embeddings = TextEmbeddings(in_text_dim, model_dim, operation_settings=operation_settings) self.pooled_text_embeddings = TextEmbeddings(in_text_dim2, time_dim, operation_settings=operation_settings) self.visual_embeddings = VisualEmbeddings(visual_embed_dim, model_dim, patch_size, operation_settings=operation_settings) self.text_transformer_blocks = nn.ModuleList( [TransformerEncoderBlock(model_dim, time_dim, ff_dim, head_dim, operation_settings=operation_settings) for _ in range(num_text_blocks)] ) self.visual_transformer_blocks = nn.ModuleList( [TransformerDecoderBlock(model_dim, time_dim, ff_dim, head_dim, operation_settings=operation_settings) for _ in range(num_visual_blocks)] ) self.out_layer = OutLayer(model_dim, time_dim, out_visual_dim, patch_size, operation_settings=operation_settings) self.rope_embedder_3d = EmbedND(dim=head_dim, theta=10000.0, axes_dim=axes_dims) self.rope_embedder_1d = EmbedND(dim=head_dim, theta=10000.0, axes_dim=[head_dim]) def rope_encode_1d(self, seq_len, seq_start=0, steps=None, device=None, dtype=None, transformer_options={}): steps = seq_len if steps is None else steps seq_ids = torch.linspace(seq_start, seq_start + (seq_len - 1), steps=steps, device=device, dtype=dtype) seq_ids = seq_ids.reshape(-1, 1).unsqueeze(0) # Shape: (1, steps, 1) freqs = self.rope_embedder_1d(seq_ids).movedim(1, 2) return freqs def rope_encode_3d(self, t, h, w, t_start=0, steps_t=None, steps_h=None, steps_w=None, device=None, dtype=None, transformer_options={}): patch_size = self.patch_size t_len = ((t + (patch_size[0] // 2)) // patch_size[0]) h_len = ((h + (patch_size[1] // 2)) // patch_size[1]) w_len = ((w + (patch_size[2] // 2)) // patch_size[2]) if steps_t is None: steps_t = t_len if steps_h is None: steps_h = h_len if steps_w is None: steps_w = w_len h_start = 0 w_start = 0 rope_options = transformer_options.get("rope_options", None) if rope_options is not None: t_len = (t_len - 1.0) * rope_options.get("scale_t", 1.0) + 1.0 h_len = (h_len - 1.0) * rope_options.get("scale_y", 1.0) + 1.0 w_len = (w_len - 1.0) * rope_options.get("scale_x", 1.0) + 1.0 t_start += rope_options.get("shift_t", 0.0) h_start += rope_options.get("shift_y", 0.0) w_start += rope_options.get("shift_x", 0.0) else: rope_scale_factor = self.rope_scale_factor if self.model_dim == 4096: # pro video model uses different rope scaling at higher resolutions if h * w >= 14080: rope_scale_factor = (1.0, 3.16, 3.16) t_len = (t_len - 1.0) / rope_scale_factor[0] + 1.0 h_len = (h_len - 1.0) / rope_scale_factor[1] + 1.0 w_len = (w_len - 1.0) / rope_scale_factor[2] + 1.0 img_ids = torch.zeros((steps_t, steps_h, steps_w, 3), device=device, dtype=dtype) img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + torch.linspace(t_start, t_start + (t_len - 1), steps=steps_t, device=device, dtype=dtype).reshape(-1, 1, 1) img_ids[:, :, :, 1] = img_ids[:, :, :, 1] + torch.linspace(h_start, h_start + (h_len - 1), steps=steps_h, device=device, dtype=dtype).reshape(1, -1, 1) img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.linspace(w_start, w_start + (w_len - 1), steps=steps_w, device=device, dtype=dtype).reshape(1, 1, -1) img_ids = img_ids.reshape(1, -1, img_ids.shape[-1]) freqs = self.rope_embedder_3d(img_ids).movedim(1, 2) return freqs def forward_orig(self, x, timestep, context, y, freqs, freqs_text, transformer_options={}, **kwargs): patches_replace = transformer_options.get("patches_replace", {}) context = self.text_embeddings(context) time_embed = self.time_embeddings(timestep, x.dtype) + self.pooled_text_embeddings(y) for block in self.text_transformer_blocks: context = block(context, time_embed, freqs_text, transformer_options=transformer_options) visual_embed = self.visual_embeddings(x) visual_shape = visual_embed.shape[:-1] visual_embed = visual_embed.flatten(1, -2) blocks_replace = patches_replace.get("dit", {}) transformer_options["total_blocks"] = len(self.visual_transformer_blocks) transformer_options["block_type"] = "double" for i, block in enumerate(self.visual_transformer_blocks): transformer_options["block_index"] = i if ("double_block", i) in blocks_replace: def block_wrap(args): return block(x=args["x"], context=args["context"], time_embed=args["time_embed"], freqs=args["freqs"], transformer_options=args.get("transformer_options")) visual_embed = blocks_replace[("double_block", i)]({"x": visual_embed, "context": context, "time_embed": time_embed, "freqs": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})["x"] else: visual_embed = block(visual_embed, context, time_embed, freqs=freqs, transformer_options=transformer_options) visual_embed = visual_embed.reshape(*visual_shape, -1) return self.out_layer(visual_embed, time_embed) def _forward(self, x, timestep, context, y, time_dim_replace=None, transformer_options={}, **kwargs): original_dims = x.ndim if original_dims == 4: x = x.unsqueeze(2) bs, c, t_len, h, w = x.shape x = comfy.ldm.common_dit.pad_to_patch_size(x, self.patch_size) if time_dim_replace is not None: time_dim_replace = comfy.ldm.common_dit.pad_to_patch_size(time_dim_replace, self.patch_size) x[:, :time_dim_replace.shape[1], :time_dim_replace.shape[2]] = time_dim_replace freqs = self.rope_encode_3d(t_len, h, w, device=x.device, dtype=x.dtype, transformer_options=transformer_options) freqs_text = self.rope_encode_1d(context.shape[1], device=x.device, dtype=x.dtype, transformer_options=transformer_options) out = self.forward_orig(x, timestep, context, y, freqs, freqs_text, transformer_options=transformer_options, **kwargs) if original_dims == 4: out = out.squeeze(2) return out def forward(self, x, timestep, context, y, time_dim_replace=None, transformer_options={}, **kwargs): return comfy.patcher_extension.WrapperExecutor.new_class_executor( self._forward, self, comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options) ).execute(x, timestep, context, y, time_dim_replace=time_dim_replace, transformer_options=transformer_options, **kwargs)