diff --git a/README.md b/README.md index 9b63972..1f6cdcb 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/nodes.py b/nodes.py index 7652b87..633fd5d 100644 --- a/nodes.py +++ b/nodes.py @@ -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", } diff --git a/pyproject.toml b/pyproject.toml index 011b6ea..67eb805 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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",