Files
quinteroac-ComfyUI-AnimaFas…/nodes.py
T

619 lines
22 KiB
Python

"""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",
}