913 lines
34 KiB
Python
913 lines
34 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 _normalize_reference_mask(reference_mask: Any) -> Any:
|
|
mask = reference_mask.detach().float()
|
|
if len(mask.shape) == 2:
|
|
mask = mask.unsqueeze(0)
|
|
elif len(mask.shape) == 3:
|
|
if mask.shape[-1] in {1, 3, 4} and mask.shape[0] != 1:
|
|
mask = mask[..., 0].unsqueeze(0)
|
|
elif len(mask.shape) == 4:
|
|
if mask.shape[-1] in {1, 3, 4}:
|
|
mask = mask[:1, ..., 0]
|
|
elif mask.shape[1] in {1, 3, 4}:
|
|
mask = mask[:1, 0]
|
|
else:
|
|
raise ValueError("reference mask must be shaped [batch, height, width]")
|
|
else:
|
|
raise ValueError("reference mask must be shaped [height, width] or [batch, height, width]")
|
|
return mask[:1].clamp(0, 1)
|
|
|
|
|
|
def _resize_reference_mask(reference_mask: Any, size: int) -> Any:
|
|
mask = _normalize_reference_mask(reference_mask)
|
|
if size <= 0:
|
|
return mask
|
|
mask = mask.unsqueeze(1)
|
|
mask = torch.nn.functional.interpolate(
|
|
mask,
|
|
size=(size, size),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
return mask.squeeze(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 _prepare_reference_masks(
|
|
reference_masks: list[Any],
|
|
*,
|
|
size: int,
|
|
reference_count: int,
|
|
) -> Any:
|
|
resized_masks = []
|
|
for mask in reference_masks:
|
|
if mask is None:
|
|
continue
|
|
resized_masks.append(_resize_reference_mask(mask, size))
|
|
if not resized_masks:
|
|
raise ValueError("at least one reference mask is required")
|
|
if len(resized_masks) == 1 and reference_count > 1:
|
|
resized_masks = [resized_masks[0].clone() for _ in range(reference_count)]
|
|
if len(resized_masks) != reference_count:
|
|
raise ValueError(
|
|
"reference mask count must be 1 or match reference image count "
|
|
f"({len(resized_masks)} != {reference_count})"
|
|
)
|
|
return torch.cat(resized_masks, dim=0).clamp(0, 1)
|
|
|
|
|
|
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 _apply_rope_split_half_autograd(input_tensor: Any, freqs_cis: Any) -> Any:
|
|
input_dtype = input_tensor.dtype
|
|
freqs_cis = freqs_cis.to(device=input_tensor.device)
|
|
first_half, second_half = input_tensor.to(dtype=freqs_cis.dtype).chunk(2, dim=-1)
|
|
first_out = freqs_cis[..., 0, 0] * first_half + freqs_cis[..., 0, 1] * second_half
|
|
second_out = freqs_cis[..., 1, 0] * first_half + freqs_cis[..., 1, 1] * second_half
|
|
return torch.cat([first_out, second_out], dim=-1).to(dtype=input_dtype)
|
|
|
|
|
|
def _apply_training_autograd_fallbacks() -> Callable[[], None]:
|
|
import comfy.model_management
|
|
|
|
previous_in_training = getattr(comfy.model_management, "in_training", False)
|
|
comfy.model_management.in_training = True
|
|
|
|
original_apply_rope_split_half = None
|
|
ck = None
|
|
try:
|
|
import comfy.quant_ops
|
|
|
|
ck = getattr(comfy.quant_ops, "ck", None)
|
|
original_apply_rope_split_half = getattr(ck, "apply_rope_split_half", None)
|
|
except Exception as exc:
|
|
LOGGER.debug("Could not inspect comfy_kitchen RoPE ops for training fallback: %s", exc)
|
|
|
|
if original_apply_rope_split_half is not None:
|
|
|
|
def _apply_rope_split_half_training(xq: Any, xk: Any, freqs_cis: Any) -> tuple[Any, Any]:
|
|
return (
|
|
_apply_rope_split_half_autograd(xq, freqs_cis),
|
|
_apply_rope_split_half_autograd(xk, freqs_cis),
|
|
)
|
|
|
|
setattr(ck, "apply_rope_split_half", _apply_rope_split_half_training)
|
|
LOGGER.info("Patched comfy_kitchen apply_rope_split_half with autograd fallback for training")
|
|
|
|
def restore() -> None:
|
|
comfy.model_management.in_training = previous_in_training
|
|
if original_apply_rope_split_half is not None:
|
|
setattr(ck, "apply_rope_split_half", original_apply_rope_split_half)
|
|
|
|
return restore
|
|
|
|
|
|
def _prepare_latent_training_mask(reference_masks: Any, latent_bank: Any) -> Any:
|
|
if len(latent_bank.shape) not in {4, 5}:
|
|
raise ValueError("masked context training expects 4D or 5D reference latents")
|
|
if reference_masks.shape[0] == 1 and latent_bank.shape[0] != 1:
|
|
reference_masks = reference_masks.repeat(latent_bank.shape[0], 1, 1)
|
|
if reference_masks.shape[0] != latent_bank.shape[0]:
|
|
raise ValueError(
|
|
"reference mask batch size must be 1 or match reference latent batch size "
|
|
f"({reference_masks.shape[0]} != {latent_bank.shape[0]})"
|
|
)
|
|
mask = reference_masks.to(device=latent_bank.device, dtype=torch.float32).unsqueeze(1)
|
|
mask = torch.nn.functional.interpolate(
|
|
mask,
|
|
size=tuple(latent_bank.shape[-2:]),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
).clamp(0, 1)
|
|
if len(latent_bank.shape) == 5:
|
|
mask = mask.unsqueeze(2)
|
|
if float(mask.detach().sum().cpu()) <= 0:
|
|
raise ValueError("reference mask must contain at least one non-zero pixel")
|
|
return mask
|
|
|
|
|
|
def _masked_mse_loss(pred: Any, target: Any, mask: Any) -> Any:
|
|
diff = (pred.float() - target.float()).pow(2)
|
|
mask = mask.to(device=diff.device, dtype=diff.dtype)
|
|
weighted_diff = diff * mask
|
|
normalizer = mask.expand_as(diff).sum().clamp_min(1.0)
|
|
return weighted_diff.sum() / normalizer
|
|
|
|
|
|
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,
|
|
reference_masks: Any = None,
|
|
) -> 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")
|
|
latent_training_mask = None
|
|
if reference_masks is not None:
|
|
latent_training_mask = _prepare_latent_training_mask(reference_masks, latent_bank)
|
|
|
|
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]
|
|
mask_batch = None
|
|
if latent_training_mask is not None:
|
|
mask_batch = latent_training_mask[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,
|
|
)
|
|
if mask_batch is None:
|
|
loss = loss_fn(pred.float(), x0.float())
|
|
else:
|
|
loss = _masked_mse_loss(pred, x0, mask_batch)
|
|
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 mask_shape=%s mask_mean=%s grad_enabled=%s inference_mode=%s",
|
|
steps,
|
|
learning_rate,
|
|
training_dtype,
|
|
tuple(initial_tokens.shape),
|
|
token_norm_before,
|
|
tuple(reference_images.shape),
|
|
None if latent_training_mask is None else tuple(latent_training_mask.shape),
|
|
None if latent_training_mask is None else f"{float(latent_training_mask.detach().mean().cpu()):.6f}",
|
|
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",
|
|
)
|
|
restore_training_fallbacks = None
|
|
try:
|
|
restore_training_fallbacks = _apply_training_autograd_fallbacks()
|
|
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:
|
|
if restore_training_fallbacks is not None:
|
|
restore_training_fallbacks()
|
|
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 AnimaFastTrainMaskedReferenceContext:
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> dict[str, Any]:
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"clip": ("CLIP",),
|
|
"vae": ("VAE",),
|
|
"reference_image": ("IMAGE",),
|
|
"training_mask": ("MASK",),
|
|
"training_prompt": (
|
|
"STRING",
|
|
{
|
|
"default": (
|
|
"masterpiece, best quality, score_7, safe, same character, "
|
|
"solo, clean white background, clean anime illustration, "
|
|
"multicolored hair, gray eyes"
|
|
),
|
|
"multiline": True,
|
|
},
|
|
),
|
|
"seed": (
|
|
"INT",
|
|
{
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xffffffffffffffff,
|
|
"control_after_generate": "fixed",
|
|
},
|
|
),
|
|
"training_steps": ("INT", {"default": 160, "min": 1, "max": 10000}),
|
|
"learning_rate": ("FLOAT", {"default": 0.010, "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.020, "min": 0.0, "max": 1.0, "step": 0.001}),
|
|
"training_dtype": (["bf16", "fp16", "fp32"], {"default": "bf16"}),
|
|
},
|
|
"optional": {
|
|
"reference_image_2": ("IMAGE",),
|
|
"reference_mask_2": ("MASK",),
|
|
"reference_image_3": ("IMAGE",),
|
|
"reference_mask_3": ("MASK",),
|
|
},
|
|
}
|
|
|
|
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_mask: 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_mask_2: Any = None,
|
|
reference_image_3: Any = None,
|
|
reference_mask_3: Any = None,
|
|
) -> tuple[AnimaReferenceContext, str]:
|
|
LOGGER.info(
|
|
"Masked train node received reference image: shape=%s mask_shape=%s dtype=%s seed=%s",
|
|
tuple(reference_image.shape),
|
|
tuple(training_mask.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]
|
|
reference_masks = _prepare_reference_masks(
|
|
[training_mask, reference_mask_2, reference_mask_3],
|
|
size=training_image_size,
|
|
reference_count=reference_count,
|
|
)
|
|
mask_mean = float(reference_masks.detach().float().mean().cpu())
|
|
LOGGER.info(
|
|
"Prepared masked Anima reference batch: count=%s image_shape=%s mask_shape=%s mask_mean=%.6f",
|
|
reference_count,
|
|
tuple(reference_images.shape),
|
|
tuple(reference_masks.shape),
|
|
mask_mean,
|
|
)
|
|
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,
|
|
reference_masks=reference_masks,
|
|
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 masked context tokens trained in memory: steps={training_steps}, "
|
|
f"lr={learning_rate}, references={reference_count}, "
|
|
f"mask_mean={mask_mean:.6f}, "
|
|
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,
|
|
"AnimaFastTrainMaskedReferenceContext": AnimaFastTrainMaskedReferenceContext,
|
|
"AnimaFastTrainPatchModel": AnimaFastTrainPatchModel,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AnimaFastTrainReferenceContext": "AnimaFastTrain - Train Context Tokens",
|
|
"AnimaFastTrainMaskedReferenceContext": "AnimaFastTrain - MaskTrain Context Tokens",
|
|
"AnimaFastTrainPatchModel": "AnimaFastTrain - Patch Model",
|
|
}
|