commit a56f0cf887c54bf0019ae4c91f65a31eaae71d36 Author: quinteroac Date: Wed May 27 14:28:18 2026 -0600 Initial AnimaFastTrain nodes diff --git a/README.md b/README.md new file mode 100644 index 0000000..d261cd4 --- /dev/null +++ b/README.md @@ -0,0 +1,72 @@ +# ComfyUI AnimaFastTrain + +Experimental custom nodes for training Anima reference context in ComfyUI with +in-memory reference context tokens. + +## Installation + +Clone this repository into your ComfyUI `custom_nodes` directory: + +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/quinteroac/ComfyUI-AnimaFastTrain.git +``` + +Then restart ComfyUI. The nodes will appear under: + +```text +Anima/AnimaFastTrain +``` + +To update an existing installation: + +```bash +cd ComfyUI/custom_nodes/ComfyUI-AnimaFastTrain +git pull +``` + +Restart ComfyUI after updating. + +## Nodes + +- **AnimaFastTrain - Train Context Tokens** + - Inputs: `MODEL`, `CLIP`, `VAE`, and a reference `IMAGE`. + - Trains per-block context tokens in memory. + - Does not save safetensors or any weight cache to disk. + - Training uses scale `1.0`, matching the experimental Python pipeline. + Runtime strength is controlled by the patch node. + - Default parameters: + - `training_steps=80` + - `learning_rate=0.02` + - internal Anima shape: `num_blocks=28`, `text_dim=1024` + +- **AnimaFastTrain - Patch Model** + - Inputs: `MODEL` and the in-memory context object. + - Outputs a patched `MODEL`. + - Connect the patched model to a normal ComfyUI sampler. + - The patch is applied during the model call and restored immediately after + each call. + - Default runtime strength: + - `reference_context_scale=0.45` + +## Important graph order + +Connect the final model into **Patch Model** after any LoRA/model patch nodes. +For example: + +`Checkpoint Loader -> LoRA Loader -> Patch Model -> KSampler` + +If another model patch node runs after **Patch Model**, it may clone the model +and drop this experimental wrapper. + +## Suggested workflow + +1. Load the Anima model, CLIP, and VAE. +2. Load or create a reference image. A face-dominant crop often improves + facial identity consistency. +3. Run **Train Context Tokens**. +4. Run **Patch Model** with the returned context. +5. Use the patched model with normal conditioning, latent, sampler, and VAE + decode nodes. + +Everything is intentionally memory-only for experimentation. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..efd723c --- /dev/null +++ b/__init__.py @@ -0,0 +1,7 @@ +"""Experimental AnimaFastTrain nodes for ComfyUI.""" + +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +WEB_DIRECTORY = None + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..7652b87 --- /dev/null +++ b/nodes.py @@ -0,0 +1,618 @@ +"""Experimental AnimaFastTrain custom nodes. + +These nodes keep identity context tokens in memory only: + +- Train Reference Context Tokens: optimizes per-block Anima context tokens from + one reference image and returns an in-memory token object. +- Patch Model With Context Tokens: returns a MODEL whose UNet call temporarily + injects those tokens into Anima block cross-attention. + +No safetensors cache is written by these nodes. +""" + +from __future__ import annotations + +import hashlib +import logging +from dataclasses import dataclass +from types import MethodType +from typing import Any, Callable + +import torch + +LOGGER = logging.getLogger("ComfyUI-AnimaFastTrain") + +ANIMA_NUM_BLOCKS = 28 +ANIMA_TEXT_DIM = 1024 + + +@dataclass +class AnimaReferenceContext: + tokens: Any + losses: list[float] + training_steps: int + learning_rate: float + training_seed: int + reference_count: int + num_blocks: int + num_tokens: int + text_dim: int + application_count: int = 0 + + +def _normalize_prompt_text(text: str) -> str: + return " " if text == "" else text + + +def _encode_prompt(clip: Any, text: str) -> Any: + tokens = clip.tokenize(_normalize_prompt_text(text)) + return clip.encode_from_tokens_scheduled(tokens) + + +def _reference_image_seed(reference_image: Any) -> int: + images = reference_image.detach().cpu().clamp(0, 1) + image_bytes = (images * 255).to(torch.uint8).numpy().tobytes() + digest = hashlib.sha256(image_bytes).digest() + return int.from_bytes(digest[:8], byteorder="little", signed=False) + + +def _resize_reference_image(reference_image: Any, size: int) -> Any: + if size <= 0: + return reference_image[:1] + image = reference_image[:1].movedim(-1, 1) + image = torch.nn.functional.interpolate( + image, + size=(size, size), + mode="bicubic", + align_corners=False, + ) + return image.movedim(1, -1).clamp(0, 1) + + +def _prepare_reference_images( + reference_images: list[Any], + *, + size: int, +) -> Any: + resized_images = [] + for image in reference_images: + if image is None: + continue + resized_images.append(_resize_reference_image(image, size)) + if not resized_images: + raise ValueError("at least one reference image is required") + return torch.cat(resized_images, dim=0) + + +def _create_reference_context_tokens( + reference_image: Any, + *, + num_blocks: int, + num_tokens: int, + text_dim: int, + dtype: str, + std: float, +) -> Any: + if num_blocks < 1: + raise ValueError("num_blocks must be >= 1") + if num_tokens < 1: + raise ValueError("num_tokens must be >= 1") + if text_dim < 1: + raise ValueError("text_dim must be >= 1") + if std < 0: + raise ValueError("init_std must be >= 0") + + torch_dtype = getattr(torch, dtype) + generator = torch.Generator(device="cpu") + generator.manual_seed(_reference_image_seed(reference_image)) + return torch.randn( + (num_blocks, num_tokens, text_dim), + generator=generator, + dtype=torch_dtype, + ) * std + + +def _get_anima_diffusion_model(model: Any) -> Any: + if hasattr(model, "model") and hasattr(model.model, "diffusion_model"): + return model.model.diffusion_model + if hasattr(model, "diffusion_model"): + return model.diffusion_model + if hasattr(model, "inner_model"): + return _get_anima_diffusion_model(model.inner_model) + if hasattr(model, "inner_model_wrapper"): + return _get_anima_diffusion_model(model.inner_model_wrapper) + if hasattr(model, "blocks"): + return model + raise TypeError("model must be an Anima diffusion model or ComfyUI ModelPatcher") + + +def _prepare_block_context_tokens(tokens: Any, block_count: int) -> Any: + if len(tokens.shape) == 2: + tokens = tokens.unsqueeze(0).unsqueeze(0) + elif len(tokens.shape) == 3: + if tokens.shape[0] == block_count: + tokens = tokens.unsqueeze(1) + else: + tokens = tokens.unsqueeze(0) + elif len(tokens.shape) != 4: + raise ValueError( + "context tokens must be shaped [tokens, text_dim], " + "[batch, tokens, text_dim], [blocks, tokens, text_dim], " + "or [blocks, batch, tokens, text_dim]" + ) + + if tokens.shape[0] == 1 and block_count != 1: + tokens = tokens.repeat(block_count, 1, 1, 1) + if tokens.shape[0] != block_count: + raise ValueError( + f"context token block count must be 1 or match model blocks ({tokens.shape[0]} != {block_count})" + ) + if not torch.is_floating_point(tokens): + raise ValueError("context tokens must be floating point") + return tokens + + +def _apply_anima_context_injection( + model: Any, + tokens: Any, + *, + scale: float, + mode: str, + debug_state: dict[str, int] | None = None, +) -> Callable[[], None]: + if mode not in {"positive", "all"}: + raise ValueError("mode must be one of: positive, all") + + diffusion_model = _get_anima_diffusion_model(model) + blocks = getattr(diffusion_model, "blocks", None) + if blocks is None: + raise TypeError("Anima diffusion model must expose a blocks attribute") + if debug_state is not None and not debug_state.get("patch_logged"): + debug_state["patch_logged"] = 1 + LOGGER.info( + "Resolved Anima diffusion model for context injection: active_type=%s diffusion_type=%s blocks=%s", + type(model).__name__, + type(diffusion_model).__name__, + len(blocks), + ) + + block_tokens = _prepare_block_context_tokens(tokens, len(blocks)) + originals = [] + + for block_index, block in enumerate(blocks): + original_forward = block.forward + originals.append((block, original_forward)) + + def _forward_with_context_injection( + self: Any, + x_B_T_H_W_D: Any, + emb_B_T_D: Any, + crossattn_emb: Any, + *args: Any, + _block_index: int = block_index, + _original_forward: Any = original_forward, + **kwargs: Any, + ) -> Any: + injected = block_tokens[_block_index].to( + device=crossattn_emb.device, + dtype=crossattn_emb.dtype, + ) + if scale != 1.0: + injected = injected * scale + if injected.shape[0] == 1 and crossattn_emb.shape[0] != 1: + injected = injected.repeat(crossattn_emb.shape[0], 1, 1) + if injected.shape[0] != crossattn_emb.shape[0]: + raise ValueError( + "context token batch size must be 1 or match cross-attention batch size" + ) + transformer_options = kwargs.get("transformer_options", {}) + cond_or_uncond = transformer_options.get("cond_or_uncond") + if mode == "positive" and cond_or_uncond: + flags = torch.tensor( + cond_or_uncond, + device=crossattn_emb.device, + dtype=injected.dtype, + ) + if crossattn_emb.shape[0] % flags.shape[0] == 0: + flags = flags.repeat_interleave(crossattn_emb.shape[0] // flags.shape[0]) + if flags.shape[0] == crossattn_emb.shape[0]: + positive_mask = (flags == 0).to(injected.dtype).view(-1, 1, 1) + injected = injected * positive_mask + if injected.shape[2] != crossattn_emb.shape[2]: + raise ValueError( + "context token text_dim must match cross-attention context dim " + f"({injected.shape[2]} != {crossattn_emb.shape[2]})" + ) + if debug_state is not None and _block_index == 0: + debug_state["block_calls"] = debug_state.get("block_calls", 0) + 1 + block_calls = debug_state["block_calls"] + if block_calls <= 3: + LOGGER.info( + "Injected Anima context inside block: block=%s call=%s crossattn_before=%s injected=%s injected_norm=%.6f", + _block_index, + block_calls, + tuple(crossattn_emb.shape), + tuple(injected.shape), + float(injected.detach().float().norm().cpu()), + ) + patched_context = torch.cat([injected, crossattn_emb], dim=1) + return _original_forward( + x_B_T_H_W_D, + emb_B_T_D, + patched_context, + *args, + **kwargs, + ) + + block.forward = MethodType(_forward_with_context_injection, block) + + def restore() -> None: + for block, original_forward in originals: + block.forward = original_forward + + return restore + + +def _optimize_reference_context_tokens( + *, + model: Any, + vae: Any, + positive: Any, + reference_images: Any, + initial_tokens: Any, + steps: int, + learning_rate: float, + seed: int, + training_dtype: str, +) -> tuple[Any, list[float]]: + if steps < 1: + raise ValueError("steps must be >= 1") + if learning_rate <= 0: + raise ValueError("learning_rate must be > 0") + + import comfy.model_management + from comfy_extras.nodes_train import TrainGuider + + dtype_map = { + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, + "float16": torch.float16, + "fp16": torch.float16, + "float32": torch.float32, + "fp32": torch.float32, + } + if training_dtype not in dtype_map: + raise ValueError(f"unsupported training_dtype: {training_dtype}") + autocast_dtype = dtype_map[training_dtype] + + encoded = vae.encode(reference_images) + latent_bank = encoded["samples"] if isinstance(encoded, dict) else encoded + if latent_bank.shape[0] < 1: + raise ValueError("encoded reference latents must not be empty") + + losses: list[float] = [] + + with torch.inference_mode(False), torch.enable_grad(): + context_tokens = torch.nn.Parameter(initial_tokens.detach().clone().float()) + optimizer = torch.optim.AdamW([context_tokens], lr=learning_rate) + loss_fn = torch.nn.MSELoss() + + class ContextTokenTrainSampler: + def sample( + self, + model_wrap: Any, + sigmas: Any, + extra_args: dict[str, Any], + callback: Any, + noise: Any, + latent_image: Any = None, + denoise_mask: Any = None, + disable_pbar: bool = False, + ) -> Any: + device = latent_image.device + latent_count = latent_image.shape[0] + + for step in range(steps): + latent_index = step % latent_count + latent_batch = latent_image[latent_index : latent_index + 1] + generator = torch.Generator(device=device) + generator.manual_seed(seed + step * 1009) + batch_noise = torch.randn( + latent_batch.shape, + generator=generator, + device=device, + dtype=latent_batch.dtype, + ) + sigma = model_wrap.inner_model.model_sampling.percent_to_sigma( + torch.rand((1,), device=device).item() + ) + batch_sigmas = torch.tensor([sigma], device=device) + + xt = model_wrap.inner_model.model_sampling.noise_scaling( + batch_sigmas, + batch_noise, + latent_batch, + False, + ) + x0 = model_wrap.inner_model.model_sampling.noise_scaling( + torch.zeros_like(batch_sigmas), + torch.zeros_like(batch_noise), + latent_batch, + False, + ) + + optimizer.zero_grad(set_to_none=True) + with torch.autocast(xt.device.type, dtype=autocast_dtype): + pred = model_wrap( + xt.requires_grad_(True), + batch_sigmas.requires_grad_(True), + **extra_args, + ) + loss = loss_fn(pred.float(), x0.float()) + loss.backward() + optimizer.step() + losses.append(float(loss.detach().cpu())) + if not disable_pbar: + print(f"context step {step + 1}/{steps}: loss={losses[-1]:.6f}") + + return latent_image + + model_patcher = model.clone() + model_patcher.model.requires_grad_(False) + model_patcher.set_model_compute_dtype(autocast_dtype) + token_norm_before = float(context_tokens.detach().float().norm().cpu()) + LOGGER.info( + "Starting Anima context token training: steps=%s lr=%s dtype=%s tokens_shape=%s token_norm_before=%.6f reference_shape=%s grad_enabled=%s inference_mode=%s", + steps, + learning_rate, + training_dtype, + tuple(initial_tokens.shape), + token_norm_before, + tuple(reference_images.shape), + torch.is_grad_enabled(), + torch.is_inference_mode_enabled(), + ) + restore_context_injection = _apply_anima_context_injection( + model_patcher, + context_tokens, + scale=1.0, + mode="positive", + ) + try: + comfy.model_management.load_models_gpu( + [model_patcher], + memory_required=1e20, + force_full_load=True, + ) + guider = TrainGuider(model_patcher, offloading=False) + guider.set_conds(positive) + sigmas = torch.tensor([0]) + dummy_noise = torch.zeros_like(latent_bank) + guider.sample( + dummy_noise, + latent_bank, + ContextTokenTrainSampler(), + sigmas, + disable_pbar=True, + seed=seed, + ) + finally: + restore_context_injection() + + final_loss = losses[-1] if losses else 0.0 + token_norm_after = float(context_tokens.detach().float().norm().cpu()) + LOGGER.info( + "Finished Anima context token training: steps=%s final_loss=%.6f token_norm_before=%.6f token_norm_after=%.6f token_norm_delta=%.6f tokens_device=%s", + steps, + final_loss, + token_norm_before, + token_norm_after, + token_norm_after - token_norm_before, + context_tokens.device, + ) + return context_tokens.detach().cpu(), losses + + +class AnimaFastTrainReferenceContext: + @classmethod + def INPUT_TYPES(cls) -> dict[str, Any]: + return { + "required": { + "model": ("MODEL",), + "clip": ("CLIP",), + "vae": ("VAE",), + "reference_image": ("IMAGE",), + "training_prompt": ( + "STRING", + { + "default": ( + "masterpiece, best quality, score_7, safe, same character, " + "solo, clean white background, clean anime illustration" + ), + "multiline": True, + }, + ), + "seed": ( + "INT", + { + "default": 0, + "min": 0, + "max": 0xffffffffffffffff, + "control_after_generate": "fixed", + }, + ), + "training_steps": ("INT", {"default": 80, "min": 1, "max": 10000}), + "learning_rate": ("FLOAT", {"default": 0.02, "min": 0.000001, "max": 1.0, "step": 0.001}), + "training_image_size": ("INT", {"default": 512, "min": 0, "max": 2048, "step": 64}), + "num_tokens": ("INT", {"default": 16, "min": 1, "max": 256}), + "init_std": ("FLOAT", {"default": 0.02, "min": 0.0, "max": 1.0, "step": 0.001}), + "training_dtype": (["bf16", "fp16", "fp32"], {"default": "bf16"}), + }, + "optional": { + "reference_image_2": ("IMAGE",), + "reference_image_3": ("IMAGE",), + }, + } + + RETURN_TYPES = ("ANIMA_REFERENCE_CONTEXT", "STRING") + RETURN_NAMES = ("context", "info") + FUNCTION = "train" + CATEGORY = "Anima/AnimaFastTrain" + + def train( + self, + model: Any, + clip: Any, + vae: Any, + reference_image: Any, + training_prompt: str, + seed: int, + training_steps: int, + learning_rate: float, + training_image_size: int, + num_tokens: int, + init_std: float, + training_dtype: str, + reference_image_2: Any = None, + reference_image_3: Any = None, + ) -> tuple[AnimaReferenceContext, str]: + LOGGER.info( + "Train node received reference image: shape=%s dtype=%s seed=%s", + tuple(reference_image.shape), + getattr(reference_image, "dtype", None), + seed, + ) + reference_images = _prepare_reference_images( + [reference_image, reference_image_2, reference_image_3], + size=training_image_size, + ) + reference_count = reference_images.shape[0] + LOGGER.info( + "Prepared Anima reference image batch: count=%s shape=%s", + reference_count, + tuple(reference_images.shape), + ) + initial_tokens = _create_reference_context_tokens( + reference_images, + num_blocks=ANIMA_NUM_BLOCKS, + num_tokens=num_tokens, + text_dim=ANIMA_TEXT_DIM, + dtype="float32", + std=init_std, + ) + positive = _encode_prompt(clip, training_prompt) + optimized_tokens, losses = _optimize_reference_context_tokens( + model=model, + vae=vae, + positive=positive, + reference_images=reference_images, + initial_tokens=initial_tokens, + steps=training_steps, + learning_rate=learning_rate, + seed=seed, + training_dtype=training_dtype, + ) + context = AnimaReferenceContext( + tokens=optimized_tokens, + losses=losses, + training_steps=training_steps, + learning_rate=learning_rate, + training_seed=seed, + reference_count=reference_count, + num_blocks=ANIMA_NUM_BLOCKS, + num_tokens=num_tokens, + text_dim=ANIMA_TEXT_DIM, + ) + final_loss = losses[-1] if losses else 0.0 + initial_norm = float(initial_tokens.detach().float().norm().cpu()) + optimized_norm = float(optimized_tokens.detach().float().norm().cpu()) + info = ( + f"Anima context tokens trained in memory: steps={training_steps}, " + f"lr={learning_rate}, references={reference_count}, " + f"tokens={tuple(optimized_tokens.shape)}, " + f"final_loss={final_loss:.6f}, " + f"token_norm_before={initial_norm:.6f}, " + f"token_norm_after={optimized_norm:.6f}, " + f"token_norm_delta={optimized_norm - initial_norm:.6f}" + ) + LOGGER.info(info) + return context, info + + +class AnimaFastTrainPatchModel: + @classmethod + def INPUT_TYPES(cls) -> dict[str, Any]: + return { + "required": { + "model": ("MODEL",), + "context": ("ANIMA_REFERENCE_CONTEXT",), + "reference_context_scale": ("FLOAT", {"default": 0.45, "min": 0.0, "max": 5.0, "step": 0.01}), + "reference_context_mode": (["positive", "all"], {"default": "positive"}), + } + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "patch" + CATEGORY = "Anima/AnimaFastTrain" + + def patch( + self, + model: Any, + context: AnimaReferenceContext, + reference_context_scale: float, + reference_context_mode: str, + ) -> tuple[Any]: + patched_model = model.clone() + previous_wrapper = patched_model.model_options.get("model_function_wrapper") + debug_state = {"block_calls": 0, "wrapper_calls": 0} + LOGGER.info( + "Patch node created AnimaFastTrain model wrapper: scale=%s mode=%s tokens_shape=%s", + reference_context_scale, + reference_context_mode, + tuple(context.tokens.shape), + ) + + def _context_wrapper(apply_model: Any, args: dict[str, Any]) -> Any: + # ComfyUI passes the real BaseModel.apply_model bound method here. + # Patch that bound object so the injection affects the active model + # used by the sampler, even if the ModelPatcher was cloned upstream. + active_model = getattr(apply_model, "__self__", None) or patched_model + context.application_count += 1 + debug_state["wrapper_calls"] = debug_state.get("wrapper_calls", 0) + 1 + if context.application_count <= 3 or context.application_count % 10 == 0: + LOGGER.info( + "Applying AnimaFastTrain context: call=%s active_type=%s scale=%s mode=%s input_shape=%s", + context.application_count, + type(active_model).__name__, + reference_context_scale, + reference_context_mode, + tuple(args["input"].shape), + ) + restore = _apply_anima_context_injection( + active_model, + context.tokens, + scale=reference_context_scale, + mode=reference_context_mode, + debug_state=debug_state, + ) + try: + if previous_wrapper is not None: + return previous_wrapper(apply_model, args) + return apply_model(args["input"], args["timestep"], **args["c"]) + finally: + restore() + + patched_model.set_model_unet_function_wrapper(_context_wrapper) + return (patched_model,) + + +NODE_CLASS_MAPPINGS = { + "AnimaFastTrainReferenceContext": AnimaFastTrainReferenceContext, + "AnimaFastTrainPatchModel": AnimaFastTrainPatchModel, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AnimaFastTrainReferenceContext": "AnimaFastTrain - Train Context Tokens", + "AnimaFastTrainPatchModel": "AnimaFastTrain - Patch Model", +}