SolitaryThinker
4a33540c08
The torch adapter boundary marshalled the latent host<->device on EVERY denoise step: _t uploaded the latent (and re-uploaded the text embeds) and _n downloaded the velocity with a forced CUDA sync — 2*N PCIe copies + N syncs per generation, buying nothing, since the latent could stay resident on the GPU the whole loop. Root cause was a numpy loop surface. But the loop MATH is already array-agnostic (CFG combine + flow-match Euler are pure arithmetic; the solver kernel already passes torch through). So introduce a per-platform array namespace (v2/platform/ array_ns): numpy on CPU (torch-free — the parity mini is unchanged), torch-on- device on cuda. The latent is seeded with numpy and uploaded ONCE; it then stays resident through forward -> CFG combine -> solver -> next step; a single host marshal happens at the request/output boundary (engine._to_artifact). Opt-in per recipe via ModelCard.device_io (set on the Wan cards). When set on a GPU box, build_component flips the components' TorchComponent.device_io so _out keeps tensors on-device (in fp32, matching the old _to_numpy cast so the combine dtype is unchanged). Un-migrated families and the CPU toy keep numpy in/out. Also: PrecisionPolicy.cast is array-preserving; _t accepts resident tensors. Verified BIT-IDENTICAL on real Wan2.1-1.3B / H100: the on-device latent equals the pre-change numpy-path latent exactly (max_abs_diff 0.0, np.array_equal True). CPU mini holds 237 passed / 5 pre-existing; pre-commit clean. Other WanDenoiseLoop families can flip device_io next (per-family GPU re-verify); non-Wan loops migrate to the xp namespace later.
114 lines
4.6 KiB
Python
114 lines
4.6 KiB
Python
"""Policy base classes — the *default* step decomposition.
|
|
|
|
Policies (CFG, expert routing, precision, flow-shift, conditioning) delete duplication for
|
|
the families that fit; they are never required. A family whose math is braided ships a custom
|
|
``next``/``advance`` and uses these as a library (LTX-2's multi-pass guidance does exactly
|
|
that). Policy *bindings* resolve at build; policy *state* is per-request in ``LoopState``
|
|
(the adaptive-gate cached delta).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
from v2.loop.contracts import StepContext
|
|
from v2.loop.sampler import build_flow_sigmas
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# FlowShiftPolicy — resolution-bucket shift + sigma schedule #
|
|
# --------------------------------------------------------------------------- #
|
|
class FlowShiftPolicy:
|
|
"""config-driven flow-shift lookup (e.g. Wan 480p shift=3.0, 720p=5.0)."""
|
|
|
|
def __init__(self, shift: float = 3.0, bucket_lookup: dict[int, float] | None = None):
|
|
self.shift = shift
|
|
self.bucket_lookup = bucket_lookup or {}
|
|
|
|
def shift_for(self, height: int = 0, width: int = 0) -> float:
|
|
return self.bucket_lookup.get(height * width, self.shift)
|
|
|
|
def build_schedule(self,
|
|
num_steps: int,
|
|
height: int = 0,
|
|
width: int = 0,
|
|
sigmas: list[float] | None = None) -> np.ndarray:
|
|
if sigmas is not None: # explicit distilled schedule (LTX-2)
|
|
return np.asarray(sigmas, dtype=np.float64)
|
|
return build_flow_sigmas(num_steps, shift=self.shift_for(height, width))
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# PrecisionPolicy — autocast / scheduler-step dtype control #
|
|
# --------------------------------------------------------------------------- #
|
|
class PrecisionPolicy:
|
|
"""Replaces ``prefix=='Flux'`` autocast hacks + ``scheduler_step_in_fp32``."""
|
|
|
|
def __init__(self, compute_dtype: str = "float32", scheduler_step_in_fp32: bool = True):
|
|
self.compute_dtype = compute_dtype
|
|
self.scheduler_step_in_fp32 = scheduler_step_in_fp32
|
|
|
|
def cast(self, arr: Any) -> Any:
|
|
# Array-preserving: a device (torch) tensor is cast in place on its device — never pulled to
|
|
# host. numpy on CPU is unchanged. (torch is imported lazily, only on the tensor path.)
|
|
if isinstance(arr, np.ndarray):
|
|
return np.asarray(arr, dtype=np.dtype(self.compute_dtype))
|
|
import torch
|
|
return arr.to(getattr(torch, self.compute_dtype, torch.float32))
|
|
|
|
@property
|
|
def scheduler_dtype(self):
|
|
return np.float32 if self.scheduler_step_in_fp32 else np.dtype(self.compute_dtype)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# ExpertRouting — Wan2.2 boundary-timestep transformer switch #
|
|
# --------------------------------------------------------------------------- #
|
|
class ExpertRouting(ABC):
|
|
|
|
@abstractmethod
|
|
def expert_for(self, ctx: StepContext) -> str:
|
|
...
|
|
|
|
|
|
class NoRouting(ExpertRouting):
|
|
"""Single-expert models (Wan2.1 1.3B): always the same component."""
|
|
|
|
def __init__(self, component_id: str = "transformer"):
|
|
self.component_id = component_id
|
|
|
|
def expert_for(self, ctx: StepContext) -> str:
|
|
return self.component_id
|
|
|
|
|
|
class BoundaryTimestepRouting(ExpertRouting):
|
|
"""Wan2.2 ``boundary_ratio`` switch between two transformers."""
|
|
|
|
def __init__(self, high_noise: str, low_noise: str, boundary: float = 0.5):
|
|
self.high_noise = high_noise
|
|
self.low_noise = low_noise
|
|
self.boundary = boundary
|
|
|
|
def expert_for(self, ctx: StepContext) -> str:
|
|
return self.high_noise if ctx.sigma >= self.boundary else self.low_noise
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# ConditioningInjector — writes RequestState.cond; the loop stays agnostic #
|
|
# --------------------------------------------------------------------------- #
|
|
class ConditioningInjector(ABC):
|
|
|
|
@abstractmethod
|
|
def inject(self, state: Any, request: Any) -> None:
|
|
...
|
|
|
|
|
|
class PassthroughConditioning(ConditioningInjector):
|
|
"""Copies precomputed encoder outputs (text/image embeds) into ``state.cond``."""
|
|
|
|
def inject(self, state: Any, request: Any) -> None:
|
|
# encoder ComponentNodes write into state.scratch["cond"] upstream of the loop
|
|
state.cond.update(state.scratch.get("cond", {}))
|