nodes/quinteroac-ComfyUI-AnimaFastTrain
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add masked context token training

Browse Source
This commit is contained in:
quinteroac
2026-06-02 12:54:42 -06:00
parent 3fac898aa6
commit 722a49178c
3 changed files with 308 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+11 -1
View File
@@ -52,6 +52,15 @@ Restart ComfyUI after updating.
- `learning_rate=0.02`
- internal Anima shape: `num_blocks=28`, `text_dim=1024`
- **AnimaFastTrain - MaskTrain Context Tokens**
- Inputs: `MODEL`, `CLIP`, `VAE`, a reference `IMAGE`, and a `MASK`.
- Trains the same in-memory per-block context tokens as the regular train
node, but weights the reconstruction loss by the mask in latent space.
- White/1.0 mask areas contribute to training; black/0.0 areas are ignored.
- Optional second and third reference images can use matching optional masks.
If only one mask is connected, it is reused for every reference image.
- This is mask-weighted training, not a runtime cross-attention mask.
- **AnimaFastTrain - Patch Model**
- Inputs: `MODEL` and the in-memory context object.
- Outputs a patched `MODEL`.
@@ -76,7 +85,8 @@ and drop this experimental wrapper.
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**.
3. Run **Train Context Tokens**, or **MaskTrain Context Tokens** if you want the
training loss to focus on a masked region.
4. Run **Patch Model** with the returned context.
5. Use the patched model with normal conditioning, latent, sampler, and VAE
decode nodes.
nodes.py
+296 -2
View File
@@ -69,6 +69,39 @@ def _resize_reference_image(reference_image: Any, size: int) -> Any:
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],
*,
@@ -84,6 +117,29 @@ def _prepare_reference_images(
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,
*,
@@ -253,6 +309,82 @@ def _apply_anima_context_injection(
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,
@@ -264,6 +396,7 @@ def _optimize_reference_context_tokens(
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")
@@ -289,6 +422,9 @@ def _optimize_reference_context_tokens(
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] = []
@@ -315,6 +451,9 @@ def _optimize_reference_context_tokens(
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(
@@ -348,7 +487,10 @@ def _optimize_reference_context_tokens(
batch_sigmas.requires_grad_(True),
**extra_args,
)
loss = loss_fn(pred.float(), x0.float())
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()))
@@ -362,13 +504,15 @@ def _optimize_reference_context_tokens(
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",
"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(),
)
@@ -378,7 +522,9 @@ def _optimize_reference_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,
@@ -397,6 +543,8 @@ def _optimize_reference_context_tokens(
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
@@ -539,6 +687,150 @@ class AnimaFastTrainReferenceContext:
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]:
@@ -609,10 +901,12 @@ class AnimaFastTrainPatchModel:
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",
}
pyproject.toml
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-AnimaFastTrain"
description = "Experimental ComfyUI nodes for training in-memory Anima reference context tokens."
version = "0.1.0"
version = "0.1.1"
license = {text = "MIT License"}
classifiers = [
"Operating System :: OS Independent",