diff --git a/comfy/ldm/anima/lllite.py b/comfy/ldm/anima/lllite.py new file mode 100644 index 000000000..5c950ec89 --- /dev/null +++ b/comfy/ldm/anima/lllite.py @@ -0,0 +1,278 @@ +import re + +import torch +from torch import nn +import torch.nn.functional as F + +import comfy.ops +import comfy.utils + + +MODULE_PATTERN = re.compile(r"lllite_dit_blocks_(\d+)_(self_attn_[qkv]_proj|cross_attn_q_proj|mlp_layer1)$") + + +def _group_norm(channels, device=None, dtype=None, operations=None): + groups = 8 + while groups > 1 and channels % groups != 0: + groups //= 2 + return operations.GroupNorm(groups, channels, device=device, dtype=dtype) + + +class AnimaLLLiteResBlock(nn.Module): + def __init__(self, channels, device=None, dtype=None, operations=None): + super().__init__() + self.norm1 = _group_norm(channels, device=device, dtype=dtype, operations=operations) + self.conv1 = operations.Conv2d(channels, channels, kernel_size=3, padding=1, device=device, dtype=dtype) + self.norm2 = _group_norm(channels, device=device, dtype=dtype, operations=operations) + self.conv2 = operations.Conv2d(channels, channels, kernel_size=3, padding=1, device=device, dtype=dtype) + + def forward(self, x): + h = self.conv1(F.silu(self.norm1(x))) + h = self.conv2(F.silu(self.norm2(h))) + return x + h + + +class AnimaLLLiteASPP(nn.Module): + def __init__(self, channels, dilations, device=None, dtype=None, operations=None): + super().__init__() + branches = [] + for dilation in dilations: + if dilation == 1: + conv = operations.Conv2d(channels, channels, kernel_size=1, device=device, dtype=dtype) + else: + conv = operations.Conv2d(channels, channels, kernel_size=3, padding=dilation, dilation=dilation, device=device, dtype=dtype) + branches.append(nn.Sequential(conv, _group_norm(channels, device=device, dtype=dtype, operations=operations), nn.SiLU())) + self.branches = nn.ModuleList(branches) + self.global_pool = nn.AdaptiveAvgPool2d(1) + self.global_conv = nn.Sequential( + operations.Conv2d(channels, channels, kernel_size=1, device=device, dtype=dtype), + _group_norm(channels, device=device, dtype=dtype, operations=operations), + nn.SiLU(), + ) + self.proj = nn.Sequential( + operations.Conv2d(channels * (len(dilations) + 1), channels, kernel_size=1, device=device, dtype=dtype), + _group_norm(channels, device=device, dtype=dtype, operations=operations), + nn.SiLU(), + ) + + def forward(self, x): + height, width = x.shape[-2:] + outputs = [branch(x) for branch in self.branches] + pooled = self.global_conv(self.global_pool(x)) + outputs.append(F.interpolate(pooled, size=(height, width), mode="bilinear", align_corners=False)) + return self.proj(torch.cat(outputs, dim=1)) + + +class AnimaLLLiteConditioning(nn.Module): + def __init__(self, cond_in_channels, cond_dim, cond_emb_dim, cond_resblocks, aspp_dilations, device=None, dtype=None, operations=None): + super().__init__() + half_dim = cond_dim // 2 + self.conv1 = operations.Conv2d(cond_in_channels, half_dim, kernel_size=4, stride=4, device=device, dtype=dtype) + self.norm1 = _group_norm(half_dim, device=device, dtype=dtype, operations=operations) + self.conv2 = operations.Conv2d(half_dim, half_dim, kernel_size=3, padding=1, device=device, dtype=dtype) + self.norm2 = _group_norm(half_dim, device=device, dtype=dtype, operations=operations) + self.conv3 = operations.Conv2d(half_dim, cond_dim, kernel_size=4, stride=4, device=device, dtype=dtype) + self.norm3 = _group_norm(cond_dim, device=device, dtype=dtype, operations=operations) + self.resblocks = nn.ModuleList([ + AnimaLLLiteResBlock(cond_dim, device=device, dtype=dtype, operations=operations) + for _ in range(cond_resblocks) + ]) + self.aspp = AnimaLLLiteASPP(cond_dim, aspp_dilations, device=device, dtype=dtype, operations=operations) if aspp_dilations else None + self.proj = operations.Conv2d(cond_dim, cond_emb_dim, kernel_size=1, device=device, dtype=dtype) + self.out_norm = operations.LayerNorm(cond_emb_dim, device=device, dtype=dtype) + + def forward(self, x): + x = F.silu(self.norm1(self.conv1(x))) + x = F.silu(self.norm2(self.conv2(x))) + x = F.silu(self.norm3(self.conv3(x))) + for block in self.resblocks: + x = block(x) + if self.aspp is not None: + x = self.aspp(x) + x = self.proj(x).flatten(2).transpose(1, 2).contiguous() + return self.out_norm(x) + + +class AnimaLLLiteModule(nn.Module): + def __init__(self, in_dim, cond_emb_dim, mlp_dim, device=None, dtype=None, operations=None): + super().__init__() + self.down = operations.Linear(in_dim, mlp_dim, device=device, dtype=dtype) + self.mid = operations.Linear(mlp_dim + cond_emb_dim, mlp_dim, device=device, dtype=dtype) + self.cond_to_film = operations.Linear(cond_emb_dim, 2 * mlp_dim, device=device, dtype=dtype) + self.up = operations.Linear(mlp_dim, in_dim, device=device, dtype=dtype) + self.depth_embed = nn.Parameter(torch.empty(cond_emb_dim, device=device, dtype=dtype), requires_grad=False) + + def forward(self, x, cond_emb, strength): + original_shape = x.shape + if x.ndim == 5: + x = x.flatten(1, 3) + + if x.shape[0] != cond_emb.shape[0]: + if x.shape[0] % cond_emb.shape[0] != 0: + raise ValueError(f"Anima LLLite batch mismatch: model input batch {x.shape[0]}, control batch {cond_emb.shape[0]}") + cond_emb = cond_emb.repeat(x.shape[0] // cond_emb.shape[0], 1, 1) + if x.shape[1] != cond_emb.shape[1]: + raise ValueError(f"Anima LLLite sequence mismatch: model input has {x.shape[1]} tokens, control has {cond_emb.shape[1]}") + + cond_local = cond_emb + comfy.ops.cast_to_input(self.depth_embed, cond_emb) + hidden = F.silu(self.down(x)) + gamma, beta = self.cond_to_film(cond_local).chunk(2, dim=-1) + hidden = self.mid(torch.cat((cond_local, hidden), dim=-1)) + hidden = F.silu(hidden * (1 + gamma) + beta) + x = x + self.up(hidden) * strength + + if len(original_shape) == 5: + x = x.reshape(original_shape) + return x + + +class AnimaLLLite(nn.Module): + def __init__(self, state_dict, metadata, device=None, dtype=None, operations=None): + super().__init__() + metadata = metadata or {} + version = metadata.get("lllite.version", "2") + if version != "2": + raise ValueError(f"Unsupported Anima LLLite version {version!r}; only named-key v2 checkpoints are supported") + + module_names = sorted({key.split(".", 1)[0] for key in state_dict if key.startswith("lllite_dit_blocks_")}) + if not module_names: + raise ValueError("Anima LLLite checkpoint has no lllite_dit_blocks_* modules") + + cond_in_channels = state_dict["lllite_conditioning1.conv1.weight"].shape[1] + cond_dim = state_dict["lllite_conditioning1.conv3.weight"].shape[0] + cond_emb_dim = state_dict["lllite_conditioning1.proj.weight"].shape[0] + resblock_ids = {int(key.split(".")[2]) for key in state_dict if key.startswith("lllite_conditioning1.resblocks.")} + cond_resblocks = max(resblock_ids) + 1 if resblock_ids else 0 + use_aspp = any(key.startswith("lllite_conditioning1.aspp.") for key in state_dict) + dilation_string = metadata.get("lllite.aspp_dilations", "1,2,4,8") + aspp_dilations = tuple(int(value) for value in dilation_string.split(",") if value.strip()) if use_aspp else () + + self.cond_in_channels = cond_in_channels + self.inpaint_masked_input = metadata.get("lllite.inpaint_masked_input", "false").lower() == "true" + self.lllite_conditioning1 = AnimaLLLiteConditioning( + cond_in_channels, cond_dim, cond_emb_dim, cond_resblocks, aspp_dilations, + device=device, dtype=dtype, operations=operations, + ) + + self.module_names = set() + self.block_count = 0 + self.model_dim = None + for name in module_names: + match = MODULE_PATTERN.fullmatch(name) + if match is None: + raise ValueError(f"Unsupported Anima LLLite module name: {name}") + down_shape = state_dict[f"{name}.down.weight"].shape + mlp_dim, in_dim = down_shape + module_cond_dim = state_dict[f"{name}.cond_to_film.weight"].shape[1] + if module_cond_dim != cond_emb_dim: + raise ValueError(f"Anima LLLite conditioning dimension mismatch in {name}: {module_cond_dim} != {cond_emb_dim}") + if self.model_dim is None: + self.model_dim = in_dim + elif self.model_dim != in_dim: + raise ValueError(f"Anima LLLite model dimension mismatch in {name}: {in_dim} != {self.model_dim}") + self.add_module(name, AnimaLLLiteModule(in_dim, cond_emb_dim, mlp_dim, device=device, dtype=dtype, operations=operations)) + self.module_names.add(name) + self.block_count = max(self.block_count, int(match.group(1)) + 1) + + def encode_conditioning(self, image): + return self.lllite_conditioning1(image) + + def apply(self, x, cond_emb, block_index, target, strength): + name = f"lllite_dit_blocks_{block_index}_{target}" + if name not in self.module_names: + return x + return self.get_submodule(name)(x, cond_emb, strength) + + +class AnimaLLLitePatch: + def __init__(self, model_patch, image, mask, strength, sigma_start, sigma_end): + self.model_patch = model_patch + self.image = image + self.mask = mask + self.strength = strength + self.sigma_start = sigma_start + self.sigma_end = sigma_end + + def __call__(self, args): + x = args["x"] + transformer_options = args["transformer_options"] + if self.strength == 0.0: + return args + sigmas = transformer_options.get("sigmas") + if sigmas is not None: + sigma = float(sigmas.max().item()) + if not self.sigma_end <= sigma <= self.sigma_start: + return args + if x.shape[2] != 1: + raise ValueError(f"Anima LLLite only supports T=1, got T={x.shape[2]}") + + target_height = x.shape[-2] * 8 + target_width = x.shape[-1] * 8 + image = comfy.utils.common_upscale( + self.image.movedim(-1, 1), target_width, target_height, "bicubic", crop="center" + ).clamp(0.0, 1.0) + image = image.to(device=x.device, dtype=x.dtype) * 2.0 - 1.0 + + if self.model_patch.model.cond_in_channels == 4: + mask = self.mask + if mask.ndim == 3: + mask = mask.unsqueeze(1) + if mask.ndim != 4 or mask.shape[1] != 1: + raise ValueError(f"Anima LLLite mask must have one channel, got shape {tuple(mask.shape)}") + mask = comfy.utils.common_upscale( + mask.float(), target_width, target_height, "nearest-exact", crop="center" + ) + if mask.shape[0] != image.shape[0]: + if image.shape[0] % mask.shape[0] != 0: + raise ValueError( + f"Anima LLLite mask batch {mask.shape[0]} cannot be broadcast to image batch {image.shape[0]}" + ) + mask = mask.repeat(image.shape[0] // mask.shape[0], 1, 1, 1) + mask = (mask >= 0.5).to(device=x.device, dtype=x.dtype) + if self.model_patch.model.inpaint_masked_input: + image = image * (mask < 0.5).to(image.dtype) + image = torch.cat((image, mask * 2.0 - 1.0), dim=1) + + cond_emb = self.model_patch.model.encode_conditioning(image) + transformer_options["model_patch_data"][self] = cond_emb + return args + + def to(self, device_or_dtype): + return self + + def models(self): + return [self.model_patch] + + +class AnimaLLLiteAttentionPatch: + def __init__(self, patch, targets): + self.patch = patch + self.targets = targets + + def __call__(self, q, k, v, pe=None, attn_mask=None, extra_options=None): + cond_emb = extra_options["model_patch_data"].get(self.patch) + if cond_emb is None: + return {"q": q, "k": k, "v": v, "pe": pe, "attn_mask": attn_mask} + + block_index = extra_options["block_index"] + values = {"q": q, "k": k, "v": v} + for value_name, target in self.targets.items(): + values[value_name] = self.patch.model_patch.model.apply( + values[value_name], cond_emb, block_index, target, self.patch.strength + ) + + return {"q": values["q"], "k": values["k"], "v": values["v"], "pe": pe, "attn_mask": attn_mask} + + +class AnimaLLLiteMLPPatch: + def __init__(self, patch): + self.patch = patch + + def __call__(self, args): + cond_emb = args["transformer_options"]["model_patch_data"].get(self.patch) + if cond_emb is None: + return args + args["x"] = self.patch.model_patch.model.apply( + args["x"], cond_emb, args["transformer_options"]["block_index"], "mlp_layer1", self.patch.strength + ) + return args diff --git a/comfy/ldm/cosmos/predict2.py b/comfy/ldm/cosmos/predict2.py index 371296e21..d391d50b1 100644 --- a/comfy/ldm/cosmos/predict2.py +++ b/comfy/ldm/cosmos/predict2.py @@ -149,11 +149,29 @@ class Attention(nn.Module): x: torch.Tensor, context: Optional[torch.Tensor] = None, rope_emb: Optional[torch.Tensor] = None, + transformer_options: Optional[dict] = {}, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - q = self.q_proj(x) context = x if context is None else context - k = self.k_proj(context) - v = self.v_proj(context) + q_input = x + k_input = context + v_input = context + + transformer_patches = transformer_options.get("patches", {}) + patch_name = "attn1_patch" if self.is_selfattn else "attn2_patch" + if patch_name in transformer_patches: + extra_options = transformer_options.copy() + extra_options["n_heads"] = self.n_heads + extra_options["dim_head"] = self.head_dim + for patch in transformer_patches[patch_name]: + out = patch(q_input, k_input, v_input, pe=rope_emb, attn_mask=None, extra_options=extra_options) + q_input = out.get("q", q_input) + k_input = out.get("k", k_input) + v_input = out.get("v", v_input) + rope_emb = out.get("pe", rope_emb) + + q = self.q_proj(q_input) + k = self.k_proj(k_input) + v = self.v_proj(v_input) q, k, v = map( lambda t: rearrange(t, "b ... (h d) -> b ... h d", h=self.n_heads, d=self.head_dim), (q, k, v), @@ -194,7 +212,7 @@ class Attention(nn.Module): x (Tensor): The query tensor of shape [B, Mq, K] context (Optional[Tensor]): The key tensor of shape [B, Mk, K] or use x as context [self attention] if None """ - q, k, v = self.compute_qkv(x, context, rope_emb=rope_emb) + q, k, v = self.compute_qkv(x, context, rope_emb=rope_emb, transformer_options=transformer_options) return self.compute_attention(q, k, v, transformer_options=transformer_options) @@ -561,8 +579,14 @@ class Block(nn.Module): self.layer_norm_mlp, scale_mlp_B_T_1_1_D, shift_mlp_B_T_1_1_D, - ) - result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D.to(compute_dtype)) + ).to(compute_dtype) + patches = transformer_options.get("patches", {}) + if "mlp_patch" in patches: + args = {"x": normalized_x_B_T_H_W_D, "transformer_options": transformer_options} + for patch in patches["mlp_patch"]: + args = patch(args) + normalized_x_B_T_H_W_D = args["x"] + result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D) x_B_T_H_W_D = torch.addcmul(x_B_T_H_W_D, gate_mlp_B_T_1_1_D.to(residual_dtype), result_B_T_H_W_D.to(residual_dtype)) return x_B_T_H_W_D @@ -869,11 +893,22 @@ class MiniTrainDIT(nn.Module): x_B_T_H_W_D.shape == extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape ), f"{x_B_T_H_W_D.shape} != {extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape}" + transformer_options = kwargs.get("transformer_options", {}) + patches = transformer_options.get("patches", {}) + if "post_input" in patches: + transformer_options = transformer_options.copy() + transformer_options["model_patch_data"] = {} + + if "post_input" in patches: + for patch in patches["post_input"]: + out = patch({"img": x_B_T_H_W_D, "x": x_B_C_T_H_W, "transformer_options": transformer_options}) + x_B_T_H_W_D = out["img"] + block_kwargs = { "rope_emb_L_1_1_D": rope_emb_L_1_1_D.unsqueeze(1).unsqueeze(0), "adaln_lora_B_T_3D": adaln_lora_B_T_3D, "extra_per_block_pos_emb": extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D, - "transformer_options": kwargs.get("transformer_options", {}), + "transformer_options": transformer_options, } # The residual stream for this model has large values. To make fp16 compute_dtype work, we keep the residual stream @@ -883,7 +918,8 @@ class MiniTrainDIT(nn.Module): if x_B_T_H_W_D.dtype == torch.float16: x_B_T_H_W_D = x_B_T_H_W_D.float() - for block in self.blocks: + for block_index, block in enumerate(self.blocks): + transformer_options["block_index"] = block_index x_B_T_H_W_D = block( x_B_T_H_W_D, t_embedding_B_T_D, diff --git a/comfy_extras/nodes_model_patch.py b/comfy_extras/nodes_model_patch.py index 3f785c8b5..0935af09d 100644 --- a/comfy_extras/nodes_model_patch.py +++ b/comfy_extras/nodes_model_patch.py @@ -8,6 +8,7 @@ import comfy.ldm.common_dit import comfy.latent_formats import comfy.ldm.lumina.controlnet import comfy.ldm.supir.supir_modules +import comfy.ldm.anima.lllite from comfy.ldm.wan.model_multitalk import WanMultiTalkAttentionBlock, MultiTalkAudioProjModel from comfy_api.latest import io from comfy.ldm.supir.supir_patch import SUPIRPatch @@ -236,10 +237,12 @@ class ModelPatchLoader: def load_model_patch(self, name): model_patch_path = folder_paths.get_full_path_or_raise("model_patches", name) - sd = comfy.utils.load_torch_file(model_patch_path, safe_load=True) + sd, metadata = comfy.utils.load_torch_file(model_patch_path, safe_load=True, return_metadata=True) dtype = comfy.utils.weight_dtype(sd) - if 'controlnet_blocks.0.y_rms.weight' in sd: + if 'lllite_conditioning1.conv1.weight' in sd: + model = comfy.ldm.anima.lllite.AnimaLLLite(sd, metadata, device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast) + elif 'controlnet_blocks.0.y_rms.weight' in sd: additional_in_dim = sd["img_in.weight"].shape[1] - 64 model = QwenImageBlockWiseControlNet(additional_in_dim=additional_in_dim, device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast) elif 'feature_embedder.mid_layer_norm.bias' in sd: @@ -296,6 +299,50 @@ class ModelPatchLoader: return (model_patcher,) +class AnimaLLLiteApply: + @classmethod + def INPUT_TYPES(s): + return {"required": {"model": ("MODEL",), + "model_patch": ("MODEL_PATCH",), + "image": ("IMAGE",), + "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), + "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}), + "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}), + }, + "optional": {"mask": ("MASK",), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply_patch" + EXPERIMENTAL = True + + CATEGORY = "model_patches/anima" + + def apply_patch(self, model, model_patch, image, strength, start_percent, end_percent, mask=None): + image = image[..., :3] + + if model_patch.model.cond_in_channels == 4 and mask is None: + mask = torch.zeros_like(image[..., 0]) + elif model_patch.model.cond_in_channels != 4: + mask = None + + model_sampling = model.get_model_object("model_sampling") + sigma_start = float(model_sampling.percent_to_sigma(start_percent)) + sigma_end = float(model_sampling.percent_to_sigma(end_percent)) + patch = comfy.ldm.anima.lllite.AnimaLLLitePatch(model_patch, image, mask, strength, sigma_start, sigma_end) + model_patched = model.clone() + model_patched.set_model_post_input_patch(patch) + model_patched.set_model_attn1_patch(comfy.ldm.anima.lllite.AnimaLLLiteAttentionPatch( + patch, + {"q": "self_attn_q_proj", "k": "self_attn_k_proj", "v": "self_attn_v_proj"}, + )) + model_patched.set_model_attn2_patch(comfy.ldm.anima.lllite.AnimaLLLiteAttentionPatch( + patch, + {"q": "cross_attn_q_proj"}, + )) + model_patched.set_model_patch(comfy.ldm.anima.lllite.AnimaLLLiteMLPPatch(patch), "mlp_patch") + return (model_patched,) + + class DiffSynthCnetPatch: def __init__(self, model_patch, vae, image, strength, mask=None): self.model_patch = model_patch @@ -674,6 +721,7 @@ NODE_CLASS_MAPPINGS = { "ZImageFunControlnet": ZImageFunControlnet, "USOStyleReference": USOStyleReference, "SUPIRApply": SUPIRApply, + "AnimaLLLiteApply": AnimaLLLiteApply, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -682,4 +730,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ZImageFunControlnet": "Apply Z-Image Fun ControlNet", "USOStyleReference": "Apply USO Style Reference", "SUPIRApply": "Apply SUPIR Patch", + "AnimaLLLiteApply": "Apply Anima LLLite", }