wehub-resource-sync
cddb07a176
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
922 lines
42 KiB
Python
922 lines
42 KiB
Python
"""Anima denoising invocation.
|
|
|
|
Implements the rectified flow denoising loop for Anima models:
|
|
- Direct prediction: denoised = input - output * sigma
|
|
- Fixed shift=3.0 via loglinear_timestep_shift (Flux paper by Black Forest Labs)
|
|
- Timestep convention: timestep = sigma * 1.0 (raw sigma, NOT 1-sigma like Z-Image)
|
|
- NO v-prediction negation (unlike Z-Image)
|
|
- 3D latent space: [B, C, T, H, W] with T=1 for images
|
|
- 16 latent channels, 8x spatial compression
|
|
|
|
Key differences from Z-Image denoise:
|
|
- Anima uses fixed shift=3.0, Z-Image uses dynamic shift based on resolution
|
|
- Anima: timestep = sigma (raw), Z-Image: model_t = 1.0 - sigma
|
|
- Anima: noise_pred = model_output (direct), Z-Image: noise_pred = -model_output (v-pred)
|
|
- Anima transformer takes (x, timesteps, context, t5xxl_ids, t5xxl_weights)
|
|
- Anima uses 3D latents directly, Z-Image converts 4D -> list of 5D
|
|
"""
|
|
|
|
import math
|
|
import sys
|
|
from contextlib import ExitStack
|
|
from typing import Callable, Iterator, Optional, Tuple
|
|
|
|
import torch
|
|
import torchvision.transforms as tv_transforms
|
|
from torch.nn.attention import SDPBackend, sdpa_kernel
|
|
from torchvision.transforms.functional import resize as tv_resize
|
|
from torchvision.transforms.functional import to_tensor
|
|
from tqdm import tqdm
|
|
|
|
from invokeai.app.invocations.anima_lllite import AnimaLLLiteField
|
|
from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
|
|
from invokeai.app.invocations.fields import (
|
|
AnimaConditioningField,
|
|
DenoiseMaskField,
|
|
FieldDescriptions,
|
|
Input,
|
|
InputField,
|
|
LatentsField,
|
|
)
|
|
from invokeai.app.invocations.latent_noise import validate_noise_tensor_shape
|
|
from invokeai.app.invocations.model import TransformerField
|
|
from invokeai.app.invocations.primitives import LatentsOutput
|
|
from invokeai.app.services.shared.invocation_context import InvocationContext
|
|
from invokeai.backend.anima.anima_transformer_patch import patch_anima_for_regional_prompting
|
|
from invokeai.backend.anima.conditioning_data import AnimaRegionalTextConditioning, AnimaTextConditioning
|
|
from invokeai.backend.anima.control_net_lllite import (
|
|
AnimaControlNetLLLite,
|
|
build_inpaint_cond_image,
|
|
prepare_cond_image,
|
|
prepare_mask,
|
|
)
|
|
from invokeai.backend.anima.regional_prompting import AnimaRegionalPromptingExtension
|
|
from invokeai.backend.anima.scheduler_driver import AnimaSchedulerDriver
|
|
from invokeai.backend.flux.schedulers import (
|
|
ANIMA_SCHEDULER_LABELS,
|
|
ANIMA_SCHEDULER_NAME_VALUES,
|
|
ANIMA_SHIFT,
|
|
)
|
|
from invokeai.backend.model_manager.taxonomy import BaseModelType
|
|
from invokeai.backend.patches.layer_patcher import LayerPatcher
|
|
from invokeai.backend.patches.lora_conversions.anima_lora_constants import ANIMA_LORA_TRANSFORMER_PREFIX
|
|
from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
|
|
from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import (
|
|
RectifiedFlowInpaintExtension,
|
|
assert_broadcastable,
|
|
)
|
|
from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
|
|
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import AnimaConditioningInfo, Range
|
|
from invokeai.backend.util.devices import TorchDevice
|
|
|
|
# Anima uses 8x spatial compression (VAE downsamples by 2^3)
|
|
ANIMA_LATENT_SCALE_FACTOR = 8
|
|
# Anima uses 16 latent channels
|
|
ANIMA_LATENT_CHANNELS = 16
|
|
# Anima uses raw sigma values as timesteps (no rescaling)
|
|
ANIMA_MULTIPLIER = 1.0
|
|
|
|
|
|
def loglinear_timestep_shift(alpha: float, t: float) -> float:
|
|
"""Apply log-linear timestep shift to a noise schedule value.
|
|
|
|
This shift biases the noise schedule toward higher noise levels, as described
|
|
in the Flux model (Black Forest Labs, 2024). With alpha > 1, the model spends
|
|
proportionally more denoising steps at higher noise levels.
|
|
|
|
Formula: sigma = alpha * t / (1 + (alpha - 1) * t)
|
|
|
|
Args:
|
|
alpha: Shift factor (3.0 for Anima, resolution-dependent for Flux).
|
|
t: Timestep value in [0, 1].
|
|
|
|
Returns:
|
|
Shifted timestep value.
|
|
"""
|
|
if alpha == 1.0:
|
|
return t
|
|
return alpha * t / (1 + (alpha - 1) * t)
|
|
|
|
|
|
def inverse_loglinear_timestep_shift(alpha: float, sigma: float) -> float:
|
|
"""Recover linear t from a shifted sigma value.
|
|
|
|
Inverse of loglinear_timestep_shift: given sigma = alpha * t / (1 + (alpha-1) * t),
|
|
solve for t = sigma / (alpha - (alpha-1) * sigma).
|
|
|
|
This is needed for the inpainting extension, which expects linear t values
|
|
for gradient mask thresholding. With Anima's shift=3.0, the difference
|
|
between shifted sigma and linear t is large (e.g. at t=0.5, sigma=0.75),
|
|
causing overly aggressive mask thresholding if sigma is used directly.
|
|
|
|
Args:
|
|
alpha: Shift factor (3.0 for Anima).
|
|
sigma: Shifted sigma value in [0, 1].
|
|
|
|
Returns:
|
|
Linear t value in [0, 1].
|
|
"""
|
|
if alpha == 1.0:
|
|
return sigma
|
|
denominator = alpha - (alpha - 1) * sigma
|
|
if abs(denominator) < 1e-8:
|
|
return 1.0
|
|
return sigma / denominator
|
|
|
|
|
|
class AnimaInpaintExtension(RectifiedFlowInpaintExtension):
|
|
"""Inpaint extension for Anima that accounts for the time-SNR shift.
|
|
|
|
Anima uses a fixed shift=3.0 which makes sigma values significantly larger
|
|
than the corresponding linear t values. The base RectifiedFlowInpaintExtension
|
|
uses t_prev for both gradient mask thresholding and noise mixing, which assumes
|
|
linear t values.
|
|
|
|
This subclass:
|
|
- Uses the LINEAR t for gradient mask thresholding (correct progressive reveal)
|
|
- Uses the SHIFTED sigma for noise mixing (matches the denoiser's noise level)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
init_latents: torch.Tensor,
|
|
inpaint_mask: torch.Tensor,
|
|
noise: torch.Tensor,
|
|
shift: float = ANIMA_SHIFT,
|
|
):
|
|
assert_broadcastable(init_latents.shape, inpaint_mask.shape, noise.shape)
|
|
self._init_latents = init_latents
|
|
self._inpaint_mask = inpaint_mask
|
|
self._noise = noise
|
|
self._shift = shift
|
|
|
|
def merge_intermediate_latents_with_init_latents(
|
|
self, intermediate_latents: torch.Tensor, sigma_prev: float
|
|
) -> torch.Tensor:
|
|
"""Merge intermediate latents with init latents, correcting for Anima's shift.
|
|
|
|
Args:
|
|
intermediate_latents: The denoised latents at the current step.
|
|
sigma_prev: The SHIFTED sigma value for the next step.
|
|
"""
|
|
# Recover linear t from shifted sigma for gradient mask thresholding.
|
|
# This ensures the gradient mask is revealed at the correct pace.
|
|
t_prev = inverse_loglinear_timestep_shift(self._shift, sigma_prev)
|
|
mask = self._apply_mask_gradient_adjustment(t_prev)
|
|
|
|
# Use shifted sigma for noise mixing to match the denoiser's noise level.
|
|
# The Euler step produces latents at noise level sigma_prev, so the
|
|
# preserved regions must also be at sigma_prev noise level.
|
|
noised_init_latents = self._noise * sigma_prev + (1.0 - sigma_prev) * self._init_latents
|
|
|
|
return intermediate_latents * mask + noised_init_latents * (1.0 - mask)
|
|
|
|
|
|
@invocation(
|
|
"anima_denoise",
|
|
title="Denoise - Anima",
|
|
tags=["image", "anima"],
|
|
category="image",
|
|
version="1.8.0",
|
|
classification=Classification.Prototype,
|
|
)
|
|
class AnimaDenoiseInvocation(BaseInvocation):
|
|
"""Run the denoising process with an Anima model.
|
|
|
|
Uses rectified flow sampling with shift=3.0 and the Cosmos Predict2 DiT
|
|
backbone with integrated LLM Adapter for text conditioning.
|
|
|
|
Supports txt2img, img2img (via latents input), and inpainting (via denoise_mask).
|
|
"""
|
|
|
|
# If latents is provided, this means we are doing image-to-image.
|
|
latents: Optional[LatentsField] = InputField(
|
|
default=None, description=FieldDescriptions.latents, input=Input.Connection
|
|
)
|
|
noise: Optional[LatentsField] = InputField(
|
|
default=None, description=FieldDescriptions.noise, input=Input.Connection
|
|
)
|
|
# denoise_mask is used for inpainting. Only the masked region is modified.
|
|
denoise_mask: Optional[DenoiseMaskField] = InputField(
|
|
default=None, description=FieldDescriptions.denoise_mask, input=Input.Connection
|
|
)
|
|
denoising_start: float = InputField(default=0.0, ge=0, le=1, description=FieldDescriptions.denoising_start)
|
|
denoising_end: float = InputField(default=1.0, ge=0, le=1, description=FieldDescriptions.denoising_end)
|
|
add_noise: bool = InputField(default=True, description="Add noise based on denoising start.")
|
|
transformer: TransformerField = InputField(
|
|
description="Anima transformer model.", input=Input.Connection, title="Transformer"
|
|
)
|
|
positive_conditioning: AnimaConditioningField | list[AnimaConditioningField] = InputField(
|
|
description=FieldDescriptions.positive_cond, input=Input.Connection
|
|
)
|
|
negative_conditioning: AnimaConditioningField | list[AnimaConditioningField] | None = InputField(
|
|
default=None, description=FieldDescriptions.negative_cond, input=Input.Connection
|
|
)
|
|
guidance_scale: float = InputField(
|
|
default=4.5,
|
|
ge=1.0,
|
|
description="Guidance scale for classifier-free guidance. Recommended: 4.0-5.0 for Anima.",
|
|
title="Guidance Scale",
|
|
)
|
|
width: int = InputField(default=1024, multiple_of=8, description="Width of the generated image.")
|
|
height: int = InputField(default=1024, multiple_of=8, description="Height of the generated image.")
|
|
steps: int = InputField(default=30, gt=0, description="Number of denoising steps. 30 recommended for Anima.")
|
|
seed: int = InputField(default=0, description="Randomness seed for reproducibility.")
|
|
# ControlNet-LLLite support (e.g. model-level inpaint conditioning, control layers)
|
|
control_lllite: AnimaLLLiteField | list[AnimaLLLiteField] | None = InputField(
|
|
default=None,
|
|
description="Anima ControlNet-LLLite conditioning (e.g. inpaint adapter, control layers). Adapters are "
|
|
"applied in a deterministic order (sorted by model key); each model may be used at most once.",
|
|
input=Input.Connection,
|
|
)
|
|
scheduler: ANIMA_SCHEDULER_NAME_VALUES = InputField(
|
|
default="euler",
|
|
description="Scheduler (sampler) for the denoising process.",
|
|
ui_choice_labels=ANIMA_SCHEDULER_LABELS,
|
|
)
|
|
|
|
@torch.no_grad()
|
|
def invoke(self, context: InvocationContext) -> LatentsOutput:
|
|
latents = self._run_diffusion(context)
|
|
latents = latents.detach().to("cpu")
|
|
name = context.tensors.save(tensor=latents)
|
|
return LatentsOutput.build(latents_name=name, latents=latents, seed=None)
|
|
|
|
def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor) -> torch.Tensor | None:
|
|
"""Prepare the inpaint mask for Anima.
|
|
|
|
Anima uses 3D latents [B, C, T, H, W] internally but the mask operates
|
|
on the spatial dimensions [B, C, H, W] which match the squeezed output.
|
|
"""
|
|
if self.denoise_mask is None:
|
|
return None
|
|
mask = context.tensors.load(self.denoise_mask.mask_name)
|
|
|
|
# Invert mask: 0.0 = regions to denoise, 1.0 = regions to preserve
|
|
mask = 1.0 - mask
|
|
|
|
_, _, latent_height, latent_width = latents.shape
|
|
mask = tv_resize(
|
|
img=mask,
|
|
size=[latent_height, latent_width],
|
|
interpolation=tv_transforms.InterpolationMode.BILINEAR,
|
|
antialias=False,
|
|
)
|
|
|
|
mask = mask.to(device=latents.device, dtype=latents.dtype)
|
|
return mask
|
|
|
|
@staticmethod
|
|
def _normalize_control_lllite(
|
|
control_lllite: AnimaLLLiteField | list[AnimaLLLiteField] | None,
|
|
) -> list[AnimaLLLiteField]:
|
|
"""Normalize the control_lllite input to a sorted list and reject duplicate models.
|
|
|
|
The model cache returns ONE shared AnimaControlNetLLLite instance per
|
|
model key, so two adapters using the same model in one run would share
|
|
cond/multiplier state and clobber each other's bindings.
|
|
|
|
The list is sorted by model key: the frontend fans adapters into a
|
|
`collect` node whose output order follows graph node ids (random
|
|
UUIDs), not user intent, and composition is weakly order-sensitive
|
|
(each adapter's delta sees the perturbations of adapters applied after
|
|
it). Sorting makes the cascade deterministic and reproducible.
|
|
"""
|
|
if control_lllite is None:
|
|
lllite_fields: list[AnimaLLLiteField] = []
|
|
elif isinstance(control_lllite, AnimaLLLiteField):
|
|
lllite_fields = [control_lllite]
|
|
elif isinstance(control_lllite, list):
|
|
lllite_fields = control_lllite
|
|
else:
|
|
raise ValueError(f"Unsupported control_lllite type: {type(control_lllite)}")
|
|
|
|
seen_keys: set[str] = set()
|
|
for lllite_field in lllite_fields:
|
|
key = lllite_field.control_model.key
|
|
if key in seen_keys:
|
|
raise ValueError(
|
|
f"The Anima ControlNet-LLLite model '{lllite_field.control_model.name}' is used by more than "
|
|
"one control input. Each LLLite model can only be applied once per generation — remove the "
|
|
"duplicate, or select a different model for it."
|
|
)
|
|
seen_keys.add(key)
|
|
return sorted(lllite_fields, key=lambda f: f.control_model.key)
|
|
|
|
def _build_lllite_cond_image(
|
|
self,
|
|
context: InvocationContext,
|
|
lllite_field: AnimaLLLiteField,
|
|
lllite_model: AnimaControlNetLLLite,
|
|
latents: torch.Tensor,
|
|
patch_spatial: int = 2,
|
|
) -> torch.Tensor:
|
|
"""Build one adapter's LLLite conditioning image tensor (once per generation).
|
|
|
|
The cond image is sized from the ACTUAL latent H/W (mirroring the DiT's
|
|
patch padding) — see target_cond_hw in the backend module.
|
|
"""
|
|
latent_h, latent_w = latents.shape[-2], latents.shape[-1]
|
|
|
|
image_pil = context.images.get_pil(lllite_field.image_name, "RGB")
|
|
rgb_01 = to_tensor(image_pil).unsqueeze(0) # (1, 3, H, W) in [0, 1]
|
|
rgb_pm1 = prepare_cond_image(rgb_01, latent_h, latent_w, patch_spatial)
|
|
|
|
if lllite_model.cond_in_channels == 4:
|
|
if lllite_field.mask_name is None:
|
|
raise ValueError(
|
|
"This Anima ControlNet-LLLite adapter is an inpainting adapter (4-channel conditioning) and "
|
|
"requires a mask. Connect a mask (white = inpaint area) to the Anima ControlNet-LLLite node."
|
|
)
|
|
mask_pil = context.images.get_pil(lllite_field.mask_name, "L")
|
|
mask_01 = to_tensor(mask_pil).unsqueeze(0) # (1, 1, H, W) in [0, 1]
|
|
mask_01 = prepare_mask(mask_01, latent_h, latent_w, patch_spatial)
|
|
return build_inpaint_cond_image(rgb_pm1, mask_01, lllite_model.inpaint_masked_input)
|
|
|
|
if lllite_model.cond_in_channels != 3:
|
|
raise ValueError(
|
|
f"Unsupported Anima ControlNet-LLLite adapter: expected 3 or 4 conditioning channels, got "
|
|
f"{lllite_model.cond_in_channels}."
|
|
)
|
|
if lllite_field.mask_name is not None:
|
|
context.logger.warning(
|
|
"The selected Anima ControlNet-LLLite adapter does not use a mask (3-channel conditioning); the "
|
|
"connected mask will be ignored."
|
|
)
|
|
return rgb_pm1
|
|
|
|
@staticmethod
|
|
def _get_lllite_multiplier(lllite_field: AnimaLLLiteField, step_index: int, total_steps: int) -> float:
|
|
"""Step-range gate for one LLLite adapter's multiplier.
|
|
|
|
Uses the same user-facing step-index/percent convention as
|
|
BaseControlNetExtension._get_weight.
|
|
"""
|
|
first_step = math.floor(lllite_field.begin_step_percent * total_steps)
|
|
last_step = math.ceil(lllite_field.end_step_percent * total_steps)
|
|
|
|
if step_index < first_step or step_index > last_step:
|
|
return 0.0
|
|
|
|
return lllite_field.weight
|
|
|
|
def _get_noise(
|
|
self,
|
|
height: int,
|
|
width: int,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
seed: int,
|
|
) -> torch.Tensor:
|
|
"""Generate initial noise tensor in 3D latent space [B, C, T, H, W]."""
|
|
rand_device = "cpu"
|
|
return torch.randn(
|
|
1,
|
|
ANIMA_LATENT_CHANNELS,
|
|
1, # T=1 for single image
|
|
height // ANIMA_LATENT_SCALE_FACTOR,
|
|
width // ANIMA_LATENT_SCALE_FACTOR,
|
|
device=rand_device,
|
|
dtype=torch.float32,
|
|
generator=torch.Generator(device=rand_device).manual_seed(seed),
|
|
).to(device=device, dtype=dtype)
|
|
|
|
def _get_sigmas(self, num_steps: int) -> list[float]:
|
|
"""Generate sigma schedule with fixed shift=3.0.
|
|
|
|
Uses the log-linear timestep shift from the Flux model (Black Forest Labs)
|
|
with a fixed shift factor of 3.0 (no dynamic resolution-based shift).
|
|
|
|
Returns:
|
|
List of num_steps + 1 sigma values from ~1.0 (noise) to 0.0 (clean).
|
|
"""
|
|
sigmas = []
|
|
for i in range(num_steps + 1):
|
|
t = 1.0 - i / num_steps
|
|
sigma = loglinear_timestep_shift(ANIMA_SHIFT, t)
|
|
sigmas.append(sigma)
|
|
return sigmas
|
|
|
|
def _load_conditioning(
|
|
self,
|
|
context: InvocationContext,
|
|
cond_field: AnimaConditioningField,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
) -> AnimaConditioningInfo:
|
|
"""Load Anima conditioning data from storage."""
|
|
cond_data = context.conditioning.load(cond_field.conditioning_name)
|
|
assert len(cond_data.conditionings) == 1
|
|
cond_info = cond_data.conditionings[0]
|
|
assert isinstance(cond_info, AnimaConditioningInfo)
|
|
return cond_info.to(dtype=dtype, device=device)
|
|
|
|
def _load_text_conditionings(
|
|
self,
|
|
context: InvocationContext,
|
|
cond_field: AnimaConditioningField | list[AnimaConditioningField],
|
|
img_token_height: int,
|
|
img_token_width: int,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
) -> list[AnimaTextConditioning]:
|
|
"""Load Anima text conditioning with optional regional masks.
|
|
|
|
Args:
|
|
context: The invocation context.
|
|
cond_field: Single conditioning field or list of fields.
|
|
img_token_height: Height of the image token grid (H // patch_size).
|
|
img_token_width: Width of the image token grid (W // patch_size).
|
|
dtype: Target dtype.
|
|
device: Target device.
|
|
|
|
Returns:
|
|
List of AnimaTextConditioning objects with optional masks.
|
|
"""
|
|
cond_list = cond_field if isinstance(cond_field, list) else [cond_field]
|
|
|
|
text_conditionings: list[AnimaTextConditioning] = []
|
|
for cond in cond_list:
|
|
cond_info = self._load_conditioning(context, cond, dtype, device)
|
|
|
|
# Load the mask, if provided
|
|
mask: torch.Tensor | None = None
|
|
if cond.mask is not None:
|
|
mask = context.tensors.load(cond.mask.tensor_name)
|
|
mask = mask.to(device=device)
|
|
mask = AnimaRegionalPromptingExtension.preprocess_regional_prompt_mask(
|
|
mask, img_token_height, img_token_width, dtype, device
|
|
)
|
|
|
|
text_conditionings.append(
|
|
AnimaTextConditioning(
|
|
qwen3_embeds=cond_info.qwen3_embeds,
|
|
t5xxl_ids=cond_info.t5xxl_ids,
|
|
t5xxl_weights=cond_info.t5xxl_weights,
|
|
mask=mask,
|
|
)
|
|
)
|
|
|
|
return text_conditionings
|
|
|
|
def _run_llm_adapter_for_regions(
|
|
self,
|
|
transformer,
|
|
text_conditionings: list[AnimaTextConditioning],
|
|
dtype: torch.dtype,
|
|
) -> AnimaRegionalTextConditioning:
|
|
"""Run the LLM Adapter separately for each regional conditioning and concatenate.
|
|
|
|
Args:
|
|
transformer: The AnimaTransformer instance (must be on device).
|
|
text_conditionings: List of per-region conditioning data.
|
|
dtype: Inference dtype.
|
|
|
|
Returns:
|
|
AnimaRegionalTextConditioning with concatenated context and masks.
|
|
"""
|
|
context_embeds_list: list[torch.Tensor] = []
|
|
context_ranges: list[Range] = []
|
|
image_masks: list[torch.Tensor | None] = []
|
|
cur_len = 0
|
|
|
|
for tc in text_conditionings:
|
|
qwen3_embeds = tc.qwen3_embeds.unsqueeze(0) # (1, seq_len, 1024)
|
|
t5xxl_ids = tc.t5xxl_ids.unsqueeze(0) # (1, seq_len)
|
|
t5xxl_weights = None
|
|
if tc.t5xxl_weights is not None:
|
|
t5xxl_weights = tc.t5xxl_weights.unsqueeze(0).unsqueeze(-1) # (1, seq_len, 1)
|
|
|
|
# Run the LLM Adapter to produce context for this region
|
|
context = transformer.preprocess_text_embeds(
|
|
qwen3_embeds.to(dtype=dtype),
|
|
t5xxl_ids,
|
|
t5xxl_weights=t5xxl_weights.to(dtype=dtype) if t5xxl_weights is not None else None,
|
|
)
|
|
# context shape: (1, 512, 1024) — squeeze batch dim
|
|
context_2d = context.squeeze(0) # (512, 1024)
|
|
|
|
context_embeds_list.append(context_2d)
|
|
context_ranges.append(Range(start=cur_len, end=cur_len + context_2d.shape[0]))
|
|
image_masks.append(tc.mask)
|
|
cur_len += context_2d.shape[0]
|
|
|
|
concatenated_context = torch.cat(context_embeds_list, dim=0)
|
|
|
|
return AnimaRegionalTextConditioning(
|
|
context_embeds=concatenated_context,
|
|
image_masks=image_masks,
|
|
context_ranges=context_ranges,
|
|
)
|
|
|
|
def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
|
|
device = TorchDevice.choose_torch_device()
|
|
inference_dtype = TorchDevice.choose_anima_inference_dtype(device)
|
|
|
|
if self.denoising_start >= self.denoising_end:
|
|
raise ValueError(
|
|
f"denoising_start ({self.denoising_start}) must be less than denoising_end ({self.denoising_end})."
|
|
)
|
|
|
|
lllite_fields = self._normalize_control_lllite(self.control_lllite)
|
|
|
|
transformer_info = context.models.load(self.transformer.transformer)
|
|
|
|
# Compute image token grid dimensions for regional prompting
|
|
# Anima: 8x VAE compression, 2x patch size → 16x total
|
|
patch_size = 2
|
|
latent_height = self.height // ANIMA_LATENT_SCALE_FACTOR
|
|
latent_width = self.width // ANIMA_LATENT_SCALE_FACTOR
|
|
img_token_height = latent_height // patch_size
|
|
img_token_width = latent_width // patch_size
|
|
img_seq_len = img_token_height * img_token_width
|
|
|
|
# Load positive conditioning with optional regional masks
|
|
pos_text_conditionings = self._load_text_conditionings(
|
|
context=context,
|
|
cond_field=self.positive_conditioning,
|
|
img_token_height=img_token_height,
|
|
img_token_width=img_token_width,
|
|
dtype=inference_dtype,
|
|
device=device,
|
|
)
|
|
has_regional = len(pos_text_conditionings) > 1 or any(tc.mask is not None for tc in pos_text_conditionings)
|
|
|
|
# Load negative conditioning if CFG is enabled
|
|
do_cfg = not math.isclose(self.guidance_scale, 1.0) and self.negative_conditioning is not None
|
|
neg_text_conditionings: list[AnimaTextConditioning] | None = None
|
|
if do_cfg:
|
|
assert self.negative_conditioning is not None
|
|
neg_text_conditionings = self._load_text_conditionings(
|
|
context=context,
|
|
cond_field=self.negative_conditioning,
|
|
img_token_height=img_token_height,
|
|
img_token_width=img_token_width,
|
|
dtype=inference_dtype,
|
|
device=device,
|
|
)
|
|
|
|
# Generate sigma schedule
|
|
sigmas = self._get_sigmas(self.steps)
|
|
|
|
# Apply denoising_start and denoising_end clipping (for img2img/inpaint)
|
|
if self.denoising_start > 0 or self.denoising_end < 1:
|
|
total_sigmas = len(sigmas)
|
|
start_idx = int(self.denoising_start * (total_sigmas - 1))
|
|
end_idx = int(self.denoising_end * (total_sigmas - 1)) + 1
|
|
sigmas = sigmas[start_idx:end_idx]
|
|
|
|
total_steps = len(sigmas) - 1
|
|
|
|
# Load input latents if provided (image-to-image)
|
|
init_latents = context.tensors.load(self.latents.latents_name) if self.latents else None
|
|
if init_latents is not None:
|
|
init_latents = init_latents.to(device=device, dtype=inference_dtype)
|
|
# Anima denoiser works in 3D: add temporal dim if needed
|
|
if init_latents.ndim == 4:
|
|
init_latents = init_latents.unsqueeze(2) # [B, C, H, W] -> [B, C, 1, H, W]
|
|
|
|
# Generate initial noise (3D latent: [B, C, T, H, W]).
|
|
# If noise will never be consumed, avoid validating/loading it.
|
|
should_ignore_noise = init_latents is not None and not self.add_noise and self.denoise_mask is None
|
|
noise: torch.Tensor | None
|
|
if should_ignore_noise:
|
|
noise = None
|
|
else:
|
|
noise = self._prepare_noise_tensor(context, inference_dtype, device)
|
|
|
|
# Prepare input latents
|
|
if init_latents is not None:
|
|
if self.add_noise:
|
|
assert noise is not None
|
|
# Noise the init latents using the first sigma from the clipped
|
|
# InvokeAI schedule.
|
|
#
|
|
# Known limitation: if the selected scheduler later starts from a
|
|
# different first effective sigma/timestep than sigmas[0], the
|
|
# img2img preblend below may not match that scheduler exactly.
|
|
# This is an existing pipeline limitation and affects both
|
|
# internally generated noise and externally supplied noise.
|
|
s_0 = sigmas[0]
|
|
latents = s_0 * noise + (1.0 - s_0) * init_latents
|
|
else:
|
|
latents = init_latents
|
|
else:
|
|
if self.denoising_start > 1e-5:
|
|
raise ValueError("denoising_start should be 0 when initial latents are not provided.")
|
|
assert noise is not None
|
|
latents = noise
|
|
|
|
if total_steps <= 0:
|
|
return latents.squeeze(2)
|
|
|
|
# Prepare inpaint extension
|
|
inpaint_mask = self._prep_inpaint_mask(context, latents.squeeze(2))
|
|
inpaint_extension: AnimaInpaintExtension | None = None
|
|
if inpaint_mask is not None:
|
|
if init_latents is None:
|
|
raise ValueError("Initial latents are required when using an inpaint mask (image-to-image inpainting)")
|
|
assert noise is not None
|
|
inpaint_extension = AnimaInpaintExtension(
|
|
init_latents=init_latents.squeeze(2),
|
|
inpaint_mask=inpaint_mask,
|
|
noise=noise.squeeze(2),
|
|
shift=ANIMA_SHIFT,
|
|
)
|
|
|
|
step_callback = self._build_step_callback(context)
|
|
|
|
# Initialize scheduler driver if not using built-in Euler.
|
|
use_scheduler = self.scheduler != "euler"
|
|
driver: AnimaSchedulerDriver | None = None
|
|
if use_scheduler:
|
|
driver = AnimaSchedulerDriver(
|
|
scheduler_name=self.scheduler,
|
|
sigmas=sigmas,
|
|
steps=self.steps,
|
|
denoising_start=self.denoising_start,
|
|
denoising_end=self.denoising_end,
|
|
device=device,
|
|
seed=self.seed,
|
|
)
|
|
|
|
with ExitStack() as exit_stack:
|
|
# On Windows + CUDA, opt into cuDNN's flash-attention kernel. torch on
|
|
# Windows does not ship the native flash-attention backend, so SDPA falls
|
|
# back to the slower memory-efficient kernel; cuDNN's kernel is ~1.7x faster
|
|
# for Anima's attention shapes (~6% per denoising step at 1024x1024). The
|
|
# cuDNN SDP backend sits low in torch's default priority and is never
|
|
# selected, so we opt in explicitly.
|
|
#
|
|
# This is a priority *list*, not a forced kernel: torch picks the first
|
|
# eligible backend per call, and MATH is always eligible, so it degrades
|
|
# safely on any GPU/dtype where cuDNN is unavailable. Scoped to Windows +
|
|
# CUDA — elsewhere torch's default backend is already fast.
|
|
if device.type == "cuda" and sys.platform == "win32":
|
|
exit_stack.enter_context(
|
|
sdpa_kernel(
|
|
[
|
|
SDPBackend.CUDNN_ATTENTION,
|
|
SDPBackend.FLASH_ATTENTION,
|
|
SDPBackend.EFFICIENT_ATTENTION,
|
|
SDPBackend.MATH,
|
|
],
|
|
set_priority=True,
|
|
)
|
|
)
|
|
(cached_weights, transformer) = exit_stack.enter_context(transformer_info.model_on_device())
|
|
|
|
# Prepare the ControlNet-LLLite adapters if provided. Each adapter's
|
|
# conditioning image is built ONCE per generation (not per step).
|
|
lllite_adapters: list[tuple[AnimaLLLiteField, AnimaControlNetLLLite, torch.Tensor]] = []
|
|
for lllite_field in lllite_fields:
|
|
lllite_info = context.models.load(lllite_field.control_model)
|
|
(_, lllite_model) = exit_stack.enter_context(lllite_info.model_on_device())
|
|
assert isinstance(lllite_model, AnimaControlNetLLLite)
|
|
lllite_cond = self._build_lllite_cond_image(
|
|
context,
|
|
lllite_field,
|
|
lllite_model,
|
|
latents,
|
|
patch_spatial=int(getattr(transformer, "patch_spatial", 2)),
|
|
)
|
|
lllite_adapters.append((lllite_field, lllite_model, lllite_cond))
|
|
|
|
# Apply LoRA models to the transformer.
|
|
# Note: We apply the LoRA after the transformer has been moved to its target device for faster patching.
|
|
exit_stack.enter_context(
|
|
LayerPatcher.apply_smart_model_patches(
|
|
model=transformer,
|
|
patches=self._lora_iterator(context),
|
|
prefix=ANIMA_LORA_TRANSFORMER_PREFIX,
|
|
dtype=inference_dtype,
|
|
cached_weights=cached_weights,
|
|
)
|
|
)
|
|
|
|
# Run LLM Adapter for each regional conditioning to produce context vectors.
|
|
# This must happen with the transformer on device since it uses the adapter weights.
|
|
if has_regional:
|
|
pos_regional = self._run_llm_adapter_for_regions(transformer, pos_text_conditionings, inference_dtype)
|
|
pos_context = pos_regional.context_embeds.unsqueeze(0) # (1, total_ctx_len, 1024)
|
|
|
|
# Build regional prompting extension with cross-attention mask
|
|
regional_extension = AnimaRegionalPromptingExtension.from_regional_conditioning(
|
|
pos_regional, img_seq_len
|
|
)
|
|
|
|
# For negative, concatenate all regions without masking (matches Z-Image behavior)
|
|
neg_context = None
|
|
if do_cfg and neg_text_conditionings is not None:
|
|
neg_regional = self._run_llm_adapter_for_regions(
|
|
transformer, neg_text_conditionings, inference_dtype
|
|
)
|
|
neg_context = neg_regional.context_embeds.unsqueeze(0)
|
|
else:
|
|
# Single conditioning — run LLM Adapter via normal forward path
|
|
tc = pos_text_conditionings[0]
|
|
pos_qwen3_embeds = tc.qwen3_embeds.unsqueeze(0)
|
|
pos_t5xxl_ids = tc.t5xxl_ids.unsqueeze(0)
|
|
pos_t5xxl_weights = None
|
|
if tc.t5xxl_weights is not None:
|
|
pos_t5xxl_weights = tc.t5xxl_weights.unsqueeze(0).unsqueeze(-1)
|
|
|
|
# Pre-compute context via LLM Adapter
|
|
pos_context = transformer.preprocess_text_embeds(
|
|
pos_qwen3_embeds.to(dtype=inference_dtype),
|
|
pos_t5xxl_ids,
|
|
t5xxl_weights=pos_t5xxl_weights.to(dtype=inference_dtype)
|
|
if pos_t5xxl_weights is not None
|
|
else None,
|
|
)
|
|
|
|
neg_context = None
|
|
if do_cfg and neg_text_conditionings is not None:
|
|
ntc = neg_text_conditionings[0]
|
|
neg_qwen3 = ntc.qwen3_embeds.unsqueeze(0)
|
|
neg_ids = ntc.t5xxl_ids.unsqueeze(0)
|
|
neg_weights = None
|
|
if ntc.t5xxl_weights is not None:
|
|
neg_weights = ntc.t5xxl_weights.unsqueeze(0).unsqueeze(-1)
|
|
neg_context = transformer.preprocess_text_embeds(
|
|
neg_qwen3.to(dtype=inference_dtype),
|
|
neg_ids,
|
|
t5xxl_weights=neg_weights.to(dtype=inference_dtype) if neg_weights is not None else None,
|
|
)
|
|
|
|
regional_extension = None
|
|
|
|
# Apply regional prompting patch if we have regional masks
|
|
exit_stack.enter_context(patch_anima_for_regional_prompting(transformer, regional_extension))
|
|
|
|
# Helper to run transformer with pre-computed context (bypasses LLM Adapter)
|
|
def _run_transformer(ctx: torch.Tensor, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
|
|
return transformer(
|
|
x=x.to(transformer.dtype if hasattr(transformer, "dtype") else inference_dtype),
|
|
timesteps=t,
|
|
context=ctx,
|
|
# t5xxl_ids=None skips the LLM Adapter — context is already pre-computed
|
|
)
|
|
|
|
try:
|
|
# Bind AFTER LoRA patching so the LLLite modules wrap the patched
|
|
# forwards. List order = apply order; restore must be the reverse.
|
|
for _, lllite_model, lllite_cond in lllite_adapters:
|
|
lllite_model.apply_to(transformer)
|
|
lllite_model.set_cond_image(lllite_cond)
|
|
|
|
if driver is not None:
|
|
user_step = 0
|
|
pbar = tqdm(total=total_steps, desc="Denoising (Anima)")
|
|
for it in driver.iterations():
|
|
# Gate on the user-facing step index so both halves of a
|
|
# multi-pass step (e.g. Heun pairs) share one gate value.
|
|
for lllite_field, lllite_model, _ in lllite_adapters:
|
|
lllite_model.set_multiplier(
|
|
self._get_lllite_multiplier(lllite_field, user_step, total_steps)
|
|
)
|
|
|
|
timestep = torch.tensor(
|
|
[it.sigma_curr * ANIMA_MULTIPLIER], device=device, dtype=inference_dtype
|
|
).expand(latents.shape[0])
|
|
|
|
noise_pred_cond = _run_transformer(pos_context, latents, timestep).float()
|
|
|
|
if do_cfg and neg_context is not None:
|
|
noise_pred_uncond = _run_transformer(neg_context, latents, timestep).float()
|
|
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond)
|
|
else:
|
|
noise_pred = noise_pred_cond
|
|
|
|
latents_preview = self._estimate_preview_latents(
|
|
latents=latents,
|
|
sigma=it.sigma_curr,
|
|
noise_pred=noise_pred,
|
|
)
|
|
|
|
latents = driver.step(model_output=noise_pred, timestep=it.sched_timestep, sample=latents)
|
|
|
|
if it.completes_user_step:
|
|
# RectifiedFlowInpaintExtension expects this once per user step (its
|
|
# docstring), so for Heun we skip the FO half of each pair to avoid
|
|
# corrupting the second-order corrector's input.
|
|
if inpaint_extension is not None:
|
|
latents_4d = latents.squeeze(2)
|
|
latents_4d = inpaint_extension.merge_intermediate_latents_with_init_latents(
|
|
latents_4d, it.sigma_prev
|
|
)
|
|
latents = latents_4d.unsqueeze(2)
|
|
|
|
user_step += 1
|
|
pbar.update(1)
|
|
step_callback(
|
|
PipelineIntermediateState(
|
|
step=user_step,
|
|
order=it.order,
|
|
total_steps=total_steps,
|
|
timestep=int(it.sigma_curr * 1000),
|
|
latents=latents_preview.squeeze(2),
|
|
)
|
|
)
|
|
pbar.close()
|
|
else:
|
|
# Built-in Euler implementation (default for Anima)
|
|
for step_idx in tqdm(range(total_steps), desc="Denoising (Anima)"):
|
|
for lllite_field, lllite_model, _ in lllite_adapters:
|
|
lllite_model.set_multiplier(
|
|
self._get_lllite_multiplier(lllite_field, step_idx, total_steps)
|
|
)
|
|
|
|
sigma_curr = sigmas[step_idx]
|
|
sigma_prev = sigmas[step_idx + 1]
|
|
|
|
timestep = torch.tensor(
|
|
[sigma_curr * ANIMA_MULTIPLIER], device=device, dtype=inference_dtype
|
|
).expand(latents.shape[0])
|
|
|
|
noise_pred_cond = _run_transformer(pos_context, latents, timestep).float()
|
|
|
|
if do_cfg and neg_context is not None:
|
|
noise_pred_uncond = _run_transformer(neg_context, latents, timestep).float()
|
|
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond)
|
|
else:
|
|
noise_pred = noise_pred_cond
|
|
|
|
latents_dtype = latents.dtype
|
|
latents = latents.to(dtype=torch.float32)
|
|
latents = latents + (sigma_prev - sigma_curr) * noise_pred
|
|
latents = latents.to(dtype=latents_dtype)
|
|
latents_preview = self._estimate_preview_latents(
|
|
latents=latents, sigma=sigma_prev, noise_pred=noise_pred
|
|
)
|
|
|
|
if inpaint_extension is not None:
|
|
latents_4d = latents.squeeze(2)
|
|
latents_4d = inpaint_extension.merge_intermediate_latents_with_init_latents(
|
|
latents_4d, sigma_prev
|
|
)
|
|
latents = latents_4d.unsqueeze(2)
|
|
|
|
step_callback(
|
|
PipelineIntermediateState(
|
|
step=step_idx + 1,
|
|
order=1,
|
|
total_steps=total_steps,
|
|
timestep=int(sigma_curr * 1000),
|
|
latents=latents_preview.squeeze(2),
|
|
),
|
|
)
|
|
finally:
|
|
# The adapter models are shared via the model cache — always undo
|
|
# the forward swaps and drop the per-run cond state. unbind() is
|
|
# only correct LIFO, so restore in REVERSE apply order (see
|
|
# AnimaControlNetLLLite.restore). Each restore is isolated so one
|
|
# failure cannot leave the remaining adapters bound to the
|
|
# cache-shared transformer.
|
|
for lllite_field, lllite_model, _ in reversed(lllite_adapters):
|
|
try:
|
|
lllite_model.restore()
|
|
lllite_model.clear_cond_image()
|
|
except Exception as e:
|
|
context.logger.error(
|
|
f"Failed to restore Anima ControlNet-LLLite adapter "
|
|
f"'{lllite_field.control_model.name}': {e}"
|
|
)
|
|
|
|
# Remove temporal dimension for output: [B, C, 1, H, W] -> [B, C, H, W]
|
|
return latents.squeeze(2)
|
|
|
|
def _prepare_noise_tensor(
|
|
self, context: InvocationContext, inference_dtype: torch.dtype, device: torch.device
|
|
) -> torch.Tensor:
|
|
if self.noise is not None:
|
|
noise = context.tensors.load(self.noise.latents_name).to(device=device, dtype=inference_dtype)
|
|
validate_noise_tensor_shape(noise, "Anima", self.width, self.height)
|
|
return noise
|
|
|
|
return self._get_noise(self.height, self.width, inference_dtype, device, self.seed)
|
|
|
|
def _estimate_preview_latents(self, latents: torch.Tensor, sigma: float, noise_pred: torch.Tensor) -> torch.Tensor:
|
|
latents_dtype = latents.dtype
|
|
latents_fp32 = latents.to(dtype=torch.float32)
|
|
preview = latents_fp32 - sigma * noise_pred.to(dtype=torch.float32)
|
|
return preview.to(dtype=latents_dtype)
|
|
|
|
def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]:
|
|
def step_callback(state: PipelineIntermediateState) -> None:
|
|
context.util.sd_step_callback(state, BaseModelType.Anima)
|
|
|
|
return step_callback
|
|
|
|
def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[ModelPatchRaw, float]]:
|
|
"""Iterate over LoRA models to apply to the transformer."""
|
|
for lora in self.transformer.loras:
|
|
lora_info = context.models.load(lora.lora)
|
|
if not isinstance(lora_info.model, ModelPatchRaw):
|
|
raise TypeError(
|
|
f"Expected ModelPatchRaw for LoRA '{lora.lora.key}', got {type(lora_info.model).__name__}. "
|
|
"The LoRA model may be corrupted or incompatible."
|
|
)
|
|
yield (lora_info.model, lora.weight)
|
|
del lora_info
|