lstein
diff --git a/‎invokeai/app/invocations/anima_denoise.py‎
Lines changed: 41 additions & 68 deletions b/‎invokeai/app/invocations/anima_denoise.py‎
Lines changed: 41 additions & 68 deletions
diff --git a/‎invokeai/backend/anima/scheduler_driver.py‎
Lines changed: 150 additions & 0 deletions b/‎invokeai/backend/anima/scheduler_driver.py‎
Lines changed: 150 additions & 0 deletions
@@ -16,14 +16,12 @@
 - Anima uses 3D latents directly, Z-Image converts 4D -> list of 5D
 """
 
-import inspect
 import math
 from contextlib import ExitStack
 from typing import Callable, Iterator, Optional, Tuple
 
 import torch
 import torchvision.transforms as tv_transforms
-from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from torchvision.transforms.functional import resize as tv_resize
 from tqdm import tqdm
 
@@ -42,7 +40,12 @@
 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.regional_prompting import AnimaRegionalPromptingExtension
-from invokeai.backend.flux.schedulers import ANIMA_SCHEDULER_LABELS, ANIMA_SCHEDULER_MAP, ANIMA_SCHEDULER_NAME_VALUES
+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
@@ -59,8 +62,6 @@
 ANIMA_LATENT_SCALE_FACTOR = 8
 # Anima uses 16 latent channels
 ANIMA_LATENT_CHANNELS = 16
-# Anima uses fixed shift=3.0 for the rectified flow schedule
-ANIMA_SHIFT = 3.0
 # Anima uses raw sigma values as timesteps (no rescaling)
 ANIMA_MULTIPLIER = 1.0
 
@@ -165,7 +166,7 @@ def merge_intermediate_latents_with_init_latents(
     title="Denoise - Anima",
     tags=["image", "anima"],
     category="image",
-    version="1.2.0",
+    version="1.5.0",
     classification=Classification.Prototype,
 )
 class AnimaDenoiseInvocation(BaseInvocation):
@@ -491,22 +492,19 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
 
         step_callback = self._build_step_callback(context)
 
-        # Initialize diffusers scheduler if not using built-in Euler
-        scheduler: SchedulerMixin | None = None
+        # Initialize scheduler driver if not using built-in Euler.
         use_scheduler = self.scheduler != "euler"
-
+        driver: AnimaSchedulerDriver | None = None
         if use_scheduler:
-            scheduler_class = ANIMA_SCHEDULER_MAP[self.scheduler]
-            scheduler = scheduler_class(num_train_timesteps=1000, shift=1.0)
-            is_lcm = self.scheduler == "lcm"
-            set_timesteps_sig = inspect.signature(scheduler.set_timesteps)
-            if not is_lcm and "sigmas" in set_timesteps_sig.parameters:
-                scheduler.set_timesteps(sigmas=sigmas, device=device)
-            else:
-                scheduler.set_timesteps(num_inference_steps=total_steps, device=device)
-            num_scheduler_steps = len(scheduler.timesteps)
-        else:
-            num_scheduler_steps = total_steps
+            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:
             (cached_weights, transformer) = exit_stack.enter_context(transformer_info.model_on_device())
@@ -587,19 +585,12 @@ def _run_transformer(ctx: torch.Tensor, x: torch.Tensor, t: torch.Tensor) -> tor
                     # t5xxl_ids=None skips the LLM Adapter — context is already pre-computed
                 )
 
-            if use_scheduler and scheduler is not None:
-                # Scheduler-based denoising
+            if driver is not None:
                 user_step = 0
                 pbar = tqdm(total=total_steps, desc="Denoising (Anima)")
-                for step_index in range(num_scheduler_steps):
-                    sched_timestep = scheduler.timesteps[step_index]
-                    sigma_curr = sched_timestep.item() / scheduler.config.num_train_timesteps
-
-                    is_heun = hasattr(scheduler, "state_in_first_order")
-                    in_first_order = scheduler.state_in_first_order if is_heun else True
-
+                for it in driver.iterations():
                     timestep = torch.tensor(
-                        [sigma_curr * ANIMA_MULTIPLIER], device=device, dtype=inference_dtype
+                        [it.sigma_curr * ANIMA_MULTIPLIER], device=device, dtype=inference_dtype
                     ).expand(latents.shape[0])
 
                     noise_pred_cond = _run_transformer(pos_context, latents, timestep).float()
@@ -610,48 +601,30 @@ def _run_transformer(ctx: torch.Tensor, x: torch.Tensor, t: torch.Tensor) -> tor
                     else:
                         noise_pred = noise_pred_cond
 
-                    step_output = scheduler.step(model_output=noise_pred, timestep=sched_timestep, sample=latents)
-                    latents = step_output.prev_sample
-
-                    if step_index + 1 < len(scheduler.sigmas):
-                        sigma_prev = scheduler.sigmas[step_index + 1].item()
-                    else:
-                        sigma_prev = 0.0
+                    latents = driver.step(model_output=noise_pred, timestep=it.sched_timestep, sample=latents)
 
-                    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)
+                    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)
 
-                    if is_heun:
-                        if not in_first_order:
-                            user_step += 1
-                            if user_step <= total_steps:
-                                pbar.update(1)
-                                step_callback(
-                                    PipelineIntermediateState(
-                                        step=user_step,
-                                        order=2,
-                                        total_steps=total_steps,
-                                        timestep=int(sigma_curr * 1000),
-                                        latents=latents.squeeze(2),
-                                    )
-                                )
-                    else:
                         user_step += 1
-                        if user_step <= total_steps:
-                            pbar.update(1)
-                            step_callback(
-                                PipelineIntermediateState(
-                                    step=user_step,
-                                    order=1,
-                                    total_steps=total_steps,
-                                    timestep=int(sigma_curr * 1000),
-                                    latents=latents.squeeze(2),
-                                )
+                        pbar.update(1)
+                        step_callback(
+                            PipelineIntermediateState(
+                                step=user_step,
+                                order=it.order,
+                                total_steps=total_steps,
+                                timestep=int(it.sigma_curr * 1000),
+                                latents=latents.squeeze(2),
                             )
+                        )
                 pbar.close()
             else:
                 # Built-in Euler implementation (default for Anima)
 
@@ -0,0 +1,150 @@
+"""Anima scheduler driver.
+
+Encapsulates the per-scheduler API quirks that ``anima_denoise._run_diffusion``
+would otherwise have to know about:
+
+* Schedulers that accept ``set_timesteps(sigmas=...)`` get the pre-shifted
+  Anima schedule passed directly.
+* Schedulers that don't accept ``sigmas=`` use ``set_begin_index()`` over their
+  own internal flow-shifted schedule. For Heun, the doubled-array index
+  translation (logical step ``k`` → doubled index ``2k``) is handled here.
+* SDE-style schedulers receive a seeded ``torch.Generator`` on every step.
+
+The denoise loop iterates :meth:`AnimaSchedulerDriver.iterations` and calls
+:meth:`AnimaSchedulerDriver.step` per iteration; the driver yields the
+``sigma_prev`` and ``completes_user_step`` flags the caller needs for inpaint
+mixing and progress reporting.
+"""
+
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Iterator
+
+import torch
+from diffusers import FlowMatchHeunDiscreteScheduler
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+from invokeai.backend.flux.schedulers import ANIMA_SCHEDULER_MAP
+
+
+@dataclass(frozen=True)
+class AnimaSchedulerIteration:
+    """Per-iteration metadata yielded by :meth:`AnimaSchedulerDriver.iterations`.
+
+    ``sigma_prev`` is the noise level the latents will be at after this iteration's
+    :meth:`AnimaSchedulerDriver.step` call. ``completes_user_step`` is True when
+    this iteration finishes a user-visible step — for Heun, the second-order
+    half of each pair plus the unpaired terminal first-order step; for every
+    other scheduler, always True.
+    """
+
+    sched_timestep: torch.Tensor
+    sigma_curr: float
+    sigma_prev: float
+    completes_user_step: bool
+    order: int
+
+
+class AnimaSchedulerDriver:
+    """Drives a diffusers scheduler over Anima's pre-shifted sigma schedule."""
+
+    def __init__(
+        self,
+        scheduler_name: str,
+        sigmas: list[float],
+        steps: int,
+        denoising_start: float,
+        denoising_end: float,
+        device: torch.device,
+        seed: int,
+    ):
+        scheduler_class, scheduler_kwargs = ANIMA_SCHEDULER_MAP[scheduler_name]
+        self.scheduler: SchedulerMixin = scheduler_class(num_train_timesteps=1000, **scheduler_kwargs)
+        # Heun toggles state_in_first_order during step(); detect by class so we
+        # can read it before set_timesteps has run.
+        self.is_heun: bool = isinstance(self.scheduler, FlowMatchHeunDiscreteScheduler)
+        self._begin_index: int = 0
+        self._step_generator = torch.Generator(device=device).manual_seed(seed)
+
+        is_lcm = scheduler_name == "lcm"
+        accepts_sigmas = "sigmas" in inspect.signature(self.scheduler.set_timesteps).parameters
+        clipped = denoising_start > 0 or denoising_end < 1
+
+        if not is_lcm and accepts_sigmas:
+            self.scheduler.set_timesteps(sigmas=sigmas, device=device)
+            self._num_iterations = len(self.scheduler.timesteps)
+        elif not is_lcm and clipped and hasattr(self.scheduler, "set_begin_index"):
+            k_start = int(denoising_start * steps)
+            k_end = int(denoising_end * steps)
+            self.scheduler.set_timesteps(num_inference_steps=steps, device=device)
+            if self.is_heun:
+                # Heun's timesteps array is 2N-1 entries; logical step k maps to
+                # doubled index 2k. min() clamps denoising_end=1.0 to the
+                # unpaired terminal first-order step.
+                self._begin_index = 2 * k_start
+                self._num_iterations = min(
+                    2 * (k_end - k_start),
+                    len(self.scheduler.timesteps) - self._begin_index,
+                )
+            else:
+                self._begin_index = k_start
+                self._num_iterations = k_end - self._begin_index
+            self.scheduler.set_begin_index(self._begin_index)
+        else:
+            self.scheduler.set_timesteps(num_inference_steps=len(sigmas) - 1, device=device)
+            self._num_iterations = len(self.scheduler.timesteps)
+
+    @property
+    def num_iterations(self) -> int:
+        """Total :meth:`step` calls. For Heun this is roughly 2× the user-visible step count."""
+        return self._num_iterations
+
+    @property
+    def begin_index(self) -> int:
+        return self._begin_index
+
+    def iterations(self) -> Iterator[AnimaSchedulerIteration]:
+        for i in range(self._num_iterations):
+            sched_idx = i + self._begin_index
+            sched_timestep = self.scheduler.timesteps[sched_idx]
+            sigma_curr = sched_timestep.item() / self.scheduler.config.num_train_timesteps
+
+            # Read state_in_first_order before step (Heun toggles it inside step()).
+            in_first_order = self.scheduler.state_in_first_order if self.is_heun else True
+
+            next_idx = sched_idx + 1
+            sigma_prev = self.scheduler.sigmas[next_idx].item() if next_idx < len(self.scheduler.sigmas) else 0.0
+
+            # For Heun, a user step completes on the second-order half of each
+            # pair AND on the unpaired terminal first-order step (sigma_prev==0).
+            is_terminal = sigma_prev == 0.0
+            completes_user_step = (not self.is_heun) or (not in_first_order) or is_terminal
+            order = 2 if self.is_heun else 1
+
+            yield AnimaSchedulerIteration(
+                sched_timestep=sched_timestep,
+                sigma_curr=sigma_curr,
+                sigma_prev=sigma_prev,
+                completes_user_step=completes_user_step,
+                order=order,
+            )
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: torch.Tensor,
+        sample: torch.Tensor,
+    ) -> torch.Tensor:
+        step_output = self.scheduler.step(
+            model_output=model_output,
+            timestep=timestep,
+            sample=sample,
+            generator=self._step_generator,
+        )
+        return step_output.prev_sample
+
+    @property
+    def step_generator(self) -> torch.Generator:
+        return self._step_generator