149 lines
4.8 KiB
Python
149 lines
4.8 KiB
Python
"""Anima Edit LoRA ReferenceLatent compatibility node for ComfyUI.
|
|
|
|
ComfyUI's built-in ReferenceLatent workflow sends a list of raw VAE-space
|
|
latents to the model through ``reference_latents``. Anima Edit LoRAs were
|
|
trained by concatenating one or more standardized reference latents on the
|
|
DiT time axis, so this node adds the small model patch that bridges those two
|
|
contracts.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import types
|
|
|
|
import torch
|
|
import comfy.conds
|
|
|
|
|
|
def _as_ref_list(ref_latents):
|
|
if ref_latents is None:
|
|
return []
|
|
if isinstance(ref_latents, (list, tuple, comfy.conds.CONDList)):
|
|
return list(ref_latents)
|
|
return [ref_latents]
|
|
|
|
|
|
def _ensure_time_axis(latent: torch.Tensor) -> torch.Tensor:
|
|
if latent.ndim == 4:
|
|
return latent.unsqueeze(2)
|
|
if latent.ndim == 5:
|
|
return latent
|
|
raise ValueError(
|
|
"AnimaEditLoraReferencePatch expected reference latent shape "
|
|
"[B, C, H, W] or [B, C, T, H, W], got "
|
|
f"{tuple(latent.shape)}"
|
|
)
|
|
|
|
|
|
def _match_batch(latent: torch.Tensor, batch_size: int) -> torch.Tensor:
|
|
ref_batch = latent.shape[0]
|
|
if ref_batch == batch_size:
|
|
return latent
|
|
if ref_batch == 1:
|
|
return latent.expand(batch_size, -1, -1, -1, -1)
|
|
if batch_size % ref_batch == 0:
|
|
return latent.repeat(batch_size // ref_batch, 1, 1, 1, 1)
|
|
raise ValueError(
|
|
"AnimaEditLoraReferencePatch cannot broadcast reference latent batch "
|
|
f"{ref_batch} to sampler batch {batch_size}"
|
|
)
|
|
|
|
|
|
def _append_reference_latents(model_input: torch.Tensor, references) -> torch.Tensor:
|
|
if not references:
|
|
return model_input
|
|
|
|
out = model_input
|
|
for ref in references:
|
|
if not isinstance(ref, torch.Tensor):
|
|
continue
|
|
ref = _ensure_time_axis(ref)
|
|
ref = ref.to(device=out.device, dtype=out.dtype)
|
|
ref = _match_batch(ref, out.shape[0])
|
|
|
|
if ref.shape[1] != out.shape[1] or ref.shape[-2:] != out.shape[-2:]:
|
|
raise ValueError(
|
|
"AnimaEditLoraReferencePatch requires reference latents to "
|
|
"match the target latent channel and spatial size. "
|
|
f"target={tuple(out.shape)}, reference={tuple(ref.shape)}"
|
|
)
|
|
out = torch.cat([out, ref], dim=2)
|
|
return out
|
|
|
|
|
|
def _call_model_apply(model_apply, model_kwargs):
|
|
kwargs = dict(model_kwargs)
|
|
x = kwargs.pop("input")
|
|
timestep = kwargs.pop("timestep")
|
|
cond = kwargs.pop("c", {})
|
|
return model_apply(x, timestep, **cond, **kwargs)
|
|
|
|
|
|
class AnimaEditLora:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"model": ("MODEL",)}}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "patch"
|
|
CATEGORY = "anima/edit/reference"
|
|
DESCRIPTION = (
|
|
"Patch an Anima MODEL so ComfyUI's built-in ReferenceLatent "
|
|
"conditions are consumed as Anima Edit reference frames. Chain after "
|
|
"Anima adapter or LoRA loaders, then feed the patched model to "
|
|
"KSampler."
|
|
)
|
|
|
|
def patch(self, model):
|
|
patched = model.clone()
|
|
|
|
original_extra_conds = patched.model.extra_conds
|
|
|
|
def extra_conds_with_reference_latents(self, **kwargs):
|
|
out = original_extra_conds(**kwargs)
|
|
ref_latents = _as_ref_list(kwargs.get("reference_latents", None))
|
|
if ref_latents:
|
|
out["ref_latents"] = comfy.conds.CONDList(
|
|
[self.process_latent_in(lat) for lat in ref_latents]
|
|
)
|
|
return out
|
|
|
|
patched.add_object_patch(
|
|
"extra_conds",
|
|
types.MethodType(extra_conds_with_reference_latents, patched.model),
|
|
)
|
|
|
|
previous_wrapper = patched.model_options.get("model_function_wrapper", None)
|
|
|
|
def reference_latent_wrapper(model_apply, model_kwargs):
|
|
c_dict = model_kwargs.get("c", {})
|
|
references = _as_ref_list(c_dict.get("ref_latents", None))
|
|
if not references:
|
|
if previous_wrapper is not None:
|
|
return previous_wrapper(model_apply, model_kwargs)
|
|
return _call_model_apply(model_apply, model_kwargs)
|
|
|
|
next_kwargs = dict(model_kwargs)
|
|
original_input = next_kwargs["input"]
|
|
original_t = original_input.shape[2]
|
|
next_kwargs["input"] = _append_reference_latents(original_input, references)
|
|
|
|
if previous_wrapper is not None:
|
|
out = previous_wrapper(model_apply, next_kwargs)
|
|
else:
|
|
out = _call_model_apply(model_apply, next_kwargs)
|
|
|
|
return out[:, :, :original_t, :, :]
|
|
|
|
patched.set_model_unet_function_wrapper(reference_latent_wrapper)
|
|
return (patched,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AnimaEditLora": AnimaEditLora,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AnimaEditLora": "Anima Edit LoRA",
|
|
}
|