feat(latent): add learned upscaling and captured noise reuse

Add provenance recovery, coordinated native sampler-state capture, frequency-completed noise reuse, and optional full-resolution inversion finishing. Support automatic SesquiLSR selection and native bislerp fallback for unknown latents.

Verify exported Anima and Krea CD workflows with zero VAE encodes. Finishing outputs reproduce prior images exactly; reuse-only outputs match independent references at the agreed capture coordinates. All 3813 Python tests and 115 frontend tests pass alongside required quality gates.
This commit is contained in:
Artificial Sweetener
2026-10-04 00:14:54 -04:00
parent 05667f6517
commit 027efb8eef
47 changed files with 3675 additions and 6 deletions
+4
View File
@@ -30,6 +30,9 @@ from .simple_syrup.runtime.attention_region_prompt_handler import ( # noqa: E40
from .simple_syrup.runtime.comfy_safetensors_dtypes import ( # noqa: E402
register_comfy_safetensors_dtypes,
)
from .simple_syrup.runtime.noise_reuse_prompt_handler import ( # noqa: E402
register_noise_reuse_prompt_handler,
)
WEB_DIRECTORY = "./web/dist"
@@ -58,6 +61,7 @@ register_quant_cache_routes()
register_external_llm_routes()
register_mask_batch_preview_routes()
register_attention_region_prompt_handler()
register_noise_reuse_prompt_handler()
__all__ = [
"WEB_DIRECTORY",
+2
View File
@@ -38,6 +38,7 @@ extend-exclude = [
"simple_syrup/third_party/groundingdino_runtime",
"simple_syrup/third_party/res4lyf_runtime",
"simple_syrup/third_party/sam_hq_runtime",
"simple_syrup/third_party/sesqui_lsr",
]
[tool.ruff.lint]
@@ -62,6 +63,7 @@ exclude = [
"simple_syrup/third_party/groundingdino_runtime",
"simple_syrup/third_party/res4lyf_runtime",
"simple_syrup/third_party/sam_hq_runtime",
"simple_syrup/third_party/sesqui_lsr",
]
[[tool.mypy.overrides]]
+55
View File
@@ -0,0 +1,55 @@
"""Define bounded learned scaling stages and verified format recommendations."""
from __future__ import annotations
import math
NATIVE_LATENT_METHODS = ("nearest-exact", "bilinear", "area", "bicubic", "bislerp")
FORMAT_FAMILIES = {
"SDXL": "sdxl",
"Flux": "flux",
"Flux2": "flux2",
"Wan21": "wan",
"Cosmos1CV8x8x8": "wan",
}
def upscale_target(shape: tuple[int, int], factor: float) -> tuple[int, int]:
"""Match Comfy's rounded latent dimensions while rejecting empty outputs."""
if not math.isfinite(factor) or not 0.01 <= factor <= 8:
raise ValueError("Latent scale factor must be between 0.01 and 8.")
target = (round(shape[0] * factor), round(shape[1] * factor))
if min(target) < 1:
raise ValueError("Latent scale factor produces an empty spatial dimension.")
return target
def learned_upscale_stages(
source: tuple[int, int], target: tuple[int, int]
) -> tuple[tuple[int, int], ...]:
"""Reach exact dimensions through stages within the learned model's 1–2x range."""
if min(*source, *target) < 2 or any(
new < old for old, new in zip(source, target, strict=True)
):
raise ValueError(
"Learned latent upscaling requires enlargement and "
"dimensions of at least 2 latent pixels."
)
current = source
result: list[tuple[int, int]] = []
while current != target:
current = (min(current[0] * 2, target[0]), min(current[1] * 2, target[1]))
result.append(current)
return tuple(result)
def recommended_family(latent_format: str) -> str:
"""Select by verified format identity instead of ambiguous channel count."""
family = FORMAT_FAMILIES.get(latent_format)
if family is None:
raise ValueError(
f"Auto latent upscaling has no compatible model for "
f"{latent_format}; select a native method or a compatible "
f"checkpoint."
)
return family
+76
View File
@@ -0,0 +1,76 @@
"""Describe native trajectory captures and independent finishing controls."""
from __future__ import annotations
import math
from dataclasses import dataclass
import torch
from .noise_inversion import InversionMethod
PROVENANCE_KEY = "simple_syrup_sampling_provenance"
@dataclass(frozen=True, slots=True)
class NoiseReuseOptions:
"""Request source noise at the receiving schedule's starting coordinate."""
finishing_steps: int = 0
finishing_method: InversionMethod = "euler"
finishing_start_fraction: float = 0.75
def __post_init__(self) -> None:
"""Reject singular coordinates and unbounded integration requests."""
if (
type(self.finishing_steps) is not int
or not 0 <= self.finishing_steps <= 100
):
raise ValueError("Finishing steps must be between 0 and 100.")
if self.finishing_method not in ("euler", "heun"):
raise ValueError("Finishing method must be Euler or Heun.")
if (
not math.isfinite(self.finishing_start_fraction)
or not 0 < self.finishing_start_fraction <= 1
):
raise ValueError(
"Finishing start fraction must be greater than 0 and at most 1."
)
if self.finishing_steps and self.finishing_start_fraction == 1:
raise ValueError(
"Finishing start fraction must be below 1 when finishing is enabled."
)
def capture_coordinate(self, target: float) -> float:
"""Round the requested coordinate exactly as the native float32 schedule."""
fraction = self.finishing_start_fraction if self.finishing_steps else 1.0
return float(torch.tensor(target * fraction, dtype=torch.float32))
@dataclass(frozen=True, slots=True)
class SamplingProvenance:
"""Carry the source identity and real model-space states beside external latents."""
latent_format: str
clean: torch.Tensor
states: tuple[tuple[float, torch.Tensor], ...] = ()
schedule: torch.Tensor | None = None
residuals: tuple[tuple[float, torch.Tensor], ...] = ()
def state_at(self, coordinate: float) -> torch.Tensor:
"""Require an exact requested schedule coordinate instead of approximating."""
for sigma, state in self.states:
if sigma == coordinate:
return state
raise ValueError(
f"Source sampling did not capture sigma {coordinate:.9g}. "
"Queue the connected workflow again with Noise Reuse Options enabled."
)
def residual_at(self, coordinate: float) -> torch.Tensor:
"""Return noise separated with the originating model's own parameterization."""
self.state_at(coordinate)
for sigma, residual in self.residuals:
if sigma == coordinate:
return residual
raise ValueError("Source capture is missing its native recovered residual.")
+12
View File
@@ -11,6 +11,7 @@ from dataclasses import dataclass, replace
from typing import TypeAlias
from .noise_inversion import NoiseInversionOptions
from .noise_reuse import NoiseReuseOptions
from .regional_prompting import validate_regional_prompt_weight
from .tiled_diffusion import validate_tiled_diffusion_mode
@@ -102,6 +103,7 @@ SamplerCapability: TypeAlias = (
| ContextualDiffusionOptions
| NoiseInversionOptions
| AttentionCouplingOptions
| NoiseReuseOptions
)
@@ -113,6 +115,7 @@ class SamplerOptions:
contextual_diffusion: ContextualDiffusionOptions | None = None
noise_inversion: NoiseInversionOptions | None = None
attention_coupling: AttentionCouplingOptions | None = None
noise_reuse: NoiseReuseOptions | None = None
def __post_init__(self) -> None:
"""Reject malformed connection payloads at the typed configuration boundary."""
@@ -121,10 +124,16 @@ class SamplerOptions:
("contextual_diffusion", ContextualDiffusionOptions),
("noise_inversion", NoiseInversionOptions),
("attention_coupling", AttentionCouplingOptions),
("noise_reuse", NoiseReuseOptions),
):
value = getattr(self, name)
if value is not None and not isinstance(value, expected):
raise TypeError(f"Sampler option {name} must be {expected.__name__}.")
if self.noise_reuse is not None and self.noise_inversion is not None:
raise ValueError(
"Noise Reuse Options cannot be combined with Noise "
"Inversion Options; use the reuse finishing controls."
)
def with_capability(self, capability: SamplerCapability) -> SamplerOptions:
"""Return a fresh configuration or reject an ambiguous duplicate feature."""
@@ -133,6 +142,7 @@ class SamplerOptions:
ContextualDiffusionOptions: "contextual_diffusion",
NoiseInversionOptions: "noise_inversion",
AttentionCouplingOptions: "attention_coupling",
NoiseReuseOptions: "noise_reuse",
}
name = names.get(type(capability))
if name is None:
@@ -147,6 +157,8 @@ class SamplerOptions:
return replace(self, contextual_diffusion=capability)
if isinstance(capability, NoiseInversionOptions):
return replace(self, noise_inversion=capability)
if isinstance(capability, NoiseReuseOptions):
return replace(self, noise_reuse=capability)
return replace(self, attention_coupling=capability)
+61
View File
@@ -0,0 +1,61 @@
"""Lift separated residuals in orthonormal spatial frequency coordinates."""
from __future__ import annotations
from functools import lru_cache
from typing import cast
import torch
@lru_cache(maxsize=16)
def dct_basis(length: int) -> torch.Tensor:
"""Build a CPU DCT-II basis without retaining accelerator allocations."""
if length < 1:
raise ValueError("DCT dimensions must be positive.")
frequency = torch.arange(length, dtype=torch.float64)[:, None]
position = torch.arange(length, dtype=torch.float64)[None, :] + 0.5
matrix = torch.cos(torch.pi * frequency * position / length) * (2 / length) ** 0.5
matrix[0] /= 2**0.5
return cast(torch.Tensor, matrix.float())
def spatial_dct(value: torch.Tensor, *, inverse: bool = False) -> torch.Tensor:
"""Transform spatial axes while preserving batch, channel and frame axes."""
result = value.float()
for axis in (-1, -2):
matrix = dct_basis(result.shape[axis]).to(result.device)
result = result.movedim(axis, -1)
result = result @ (matrix if inverse else matrix.T)
result = result.movedim(-1, axis)
return result
def lift_residual(residual: torch.Tensor, fresh: torch.Tensor) -> torch.Tensor:
"""Retain source coefficients and fill only newly introduced frequency bands."""
if (
residual.ndim not in (4, 5)
or residual.ndim != fresh.ndim
or residual.shape[:-2] != fresh.shape[:-2]
):
raise ValueError(
"Noise reuse requires matching batch, channel and frame dimensions."
)
if any(a > b for a, b in zip(residual.shape[-2:], fresh.shape[-2:], strict=True)):
raise ValueError(
"Noise reuse requires spatial enlargement or unchanged dimensions."
)
if not bool(torch.isfinite(residual).all()) or not bool(
torch.isfinite(fresh).all()
):
raise ValueError(
"Noise reuse requires finite residual and fresh noise tensors."
)
height, width = residual.shape[-2:]
old_spectrum = torch.zeros_like(fresh, dtype=torch.float32)
old_spectrum[..., :height, :width] = spatial_dct(residual)
new_spectrum = spatial_dct(fresh).clone()
new_spectrum[..., :height, :width] = 0
return spatial_dct(old_spectrum, inverse=True) + spatial_dct(
new_spectrum, inverse=True
)
+16
View File
@@ -35,6 +35,7 @@ def get_nodes() -> list[type[object]]:
KSamplerTiledAttentionCouplingV3,
)
from .ksampler_tiled_diffusion import KSamplerTiledDiffusionV3
from .latent_from_image import LatentFromImageV3
from .legacy_node_wrappers import (
ConditioningBatchAppendV3,
ConditioningBatchStartV3,
@@ -66,16 +67,24 @@ def get_nodes() -> list[type[object]]:
from .load_mask_batch import LoadMaskBatchV3
from .mask_to_segs import MaskToSEGSV3
from .noise_inversion_options import NoiseInversionOptionsV3
from .noise_reuse_options import NoiseReuseOptionsV3
from .scale_factor import ScaleFactorV3
from .seed_variation import SeedVariationV3
from .simple_load_checkpoint import SimpleLoadCheckpointV3
from .simple_load_flux import SimpleLoadFluxV3
from .simple_load_flux2 import SimpleLoadFlux2V3
from .simple_load_krea2 import SimpleLoadKrea2V3
from .simple_upscale_latent import SimpleUpscaleLatentV3
from .tag_segs_with_external_llm import TagSEGSWithExternalLLMV3
from .tag_segs_with_wd14 import TagSEGSWithWD14V3
from .tile_and_tag_segs import TileAndTagSEGSV3
from .tiling_options import TilingOptionsV3
from .trajectory_capture import (
AttachSamplingProvenanceV3,
NoiseCaptureRequestV3,
NoiseReuseErrorV3,
TrajectoryCaptureModelV3,
)
from .vae_decode_options import VAEDecodeOptionsV3
from .vae_encode_options import VAEEncodeOptionsV3
from .wd14_tagger_loader import WD14TaggerLoaderV3
@@ -87,6 +96,13 @@ def get_nodes() -> list[type[object]]:
NoiseInversionOptionsV3,
TilingOptionsV3,
KSamplerV3,
LatentFromImageV3,
SimpleUpscaleLatentV3,
NoiseReuseOptionsV3,
NoiseCaptureRequestV3,
TrajectoryCaptureModelV3,
AttachSamplingProvenanceV3,
NoiseReuseErrorV3,
AttentionCaptureModelV3,
AttentionMaskedConditioningV3,
AttentionRegionMaskV3,
@@ -0,0 +1,54 @@
"""Expose original decoded-image latent recovery as a strict provenance bridge."""
from __future__ import annotations
from importlib import import_module
from typing import Any, ClassVar
from ..services.latent_provenance_service import LatentProvenanceService
from .sampler_options_schema import COMFY_IO, OptionsNodeBase
class LatentFromImageV3(OptionsNodeBase):
"""Return the original graph latent while keeping image pixels outside execution."""
hidden: ClassVar[Any]
@classmethod
def define_schema(cls) -> Any:
"""Declare the widget-free bridge with graph metadata supplied by Comfy."""
return COMFY_IO.Schema(
node_id="SimpleSyrup.LatentFromImage",
display_name="Latent From Image",
category="SimpleSyrup/Latent",
enable_expand=True,
description=(
"Recovers the original latent behind an unmodified "
"decoded image using graph provenance."
),
inputs=[
COMFY_IO.Image.Input(
"image",
raw_link=True,
tooltip=(
"Connect an unmodified VAE Decode image; the original "
"latent is recovered through graph provenance."
),
)
],
hidden=[COMFY_IO.Hidden.prompt],
outputs=[
COMFY_IO.Latent.Output(
"latent",
tooltip="Original sampler latent with its sampling provenance.",
)
],
)
@classmethod
def execute(cls, image: object) -> Any:
"""Expand to the original latent link without invoking a VAE encoder."""
link = LatentProvenanceService().recover_link(
cls.hidden.prompt, image, import_module("nodes").NODE_CLASS_MAPPINGS
)
return COMFY_IO.NodeOutput([link[0], link[1]], expand={})
@@ -0,0 +1,99 @@
"""Expose captured noise reuse and optional full-resolution inversion finishing."""
from __future__ import annotations
from typing import Any, cast
from ..domain.noise_inversion import InversionMethod
from ..domain.noise_reuse import NoiseReuseOptions
from ..domain.sampler_options import SamplerOptions, append_sampler_capability
from .sampler_options_schema import (
COMFY_IO,
OptionsNodeBase,
options_input,
options_output,
)
class NoiseReuseOptionsV3(OptionsNodeBase):
"""Contribute one enabled request without changing preceding capabilities."""
@classmethod
def define_schema(cls) -> Any:
"""Place the explicit enable switch above independent finishing controls."""
return COMFY_IO.Schema(
node_id="SimpleSyrup.NoiseReuseOptions",
display_name="Noise Reuse Options",
category="SimpleSyrup/Sampling/Options",
description=(
"Reuses noise from the source sampling run, with optional "
"full-resolution inversion finishing before sampling."
),
inputs=[
options_input(COMFY_IO),
COMFY_IO.Boolean.Input(
"enabled",
default=True,
tooltip=(
"Requests captured source noise when enabled; disabled "
"passes preceding options through unchanged."
),
),
COMFY_IO.Int.Input(
"finishing_steps",
default=0,
min=0,
max=100,
tooltip=(
"Full-resolution inversion steps before sampling; 0 "
"reuses captured noise without finishing."
),
),
COMFY_IO.Combo.Input(
"finishing_method",
options=["euler", "heun"],
default="euler",
tooltip=(
"Euler uses one evaluation per finishing step; Heun uses "
"two for greater accuracy."
),
),
COMFY_IO.Float.Input(
"finishing_start_fraction",
default=0.75,
min=0.01,
max=0.99,
step=0.01,
tooltip=(
"With finishing enabled, captures at this fraction of the "
"receiving sampler's starting noise coordinate."
),
),
],
outputs=[options_output(COMFY_IO)],
)
@classmethod
def execute(
cls,
enabled: bool = True,
finishing_steps: int = 0,
finishing_method: str = "euler",
finishing_start_fraction: float = 0.75,
options: SamplerOptions | None = None,
) -> tuple[SamplerOptions]:
"""Validate enabled controls and preserve the exact disabled input object."""
if type(enabled) is not bool:
raise TypeError("Noise reuse enabled must be a boolean.")
if finishing_method not in ("euler", "heun"):
raise ValueError("Finishing method must be Euler or Heun.")
reuse = (
NoiseReuseOptions(
finishing_steps,
cast(InversionMethod, finishing_method),
finishing_start_fraction,
)
if enabled
else None
)
return (append_sampler_capability(options, reuse),)
@@ -0,0 +1,73 @@
"""Expose automatic learned and native latent scaling through one method control."""
from __future__ import annotations
from typing import Any
from ..runtime.latent_upscale_catalog import latent_model_choices
from ..services.latent_upscale_service import LatentUpscaleService
from .sampler_options_schema import COMFY_IO, OptionsNodeBase
class SimpleUpscaleLatentV3(OptionsNodeBase):
"""Resize the clean latent while retaining its source noise provenance."""
@classmethod
def define_schema(cls) -> Any:
"""Declare one method dropdown and a scale multiplier without loading
weights.
"""
return COMFY_IO.Schema(
node_id="SimpleSyrup.SimpleUpscaleLatent",
display_name="Simple Upscale Latent",
category="SimpleSyrup/Latent",
description=(
"Upscales clean latents with a compatible learned model "
"or native interpolation, preserving sampling provenance."
),
inputs=[
COMFY_IO.Latent.Input(
"latent",
tooltip=(
"Clean latent to enlarge, carrying source sampling "
"provenance when noise reuse is needed."
),
),
COMFY_IO.Combo.Input(
"method",
options=latent_model_choices(),
default="Auto",
tooltip=(
"Auto selects a compatible learned model from sampling "
"provenance, or Comfy's bislerp for unknown latent formats."
),
),
COMFY_IO.Float.Input(
"scale_factor",
default=2.0,
min=0.01,
max=8.0,
step=0.01,
tooltip=(
"Width and height multiplier; learned models enlarge "
"through internal stages of at most 2x."
),
),
],
outputs=[
COMFY_IO.Latent.Output(
"latent",
tooltip=(
"Resized clean latent retaining source sampling "
"provenance for downstream noise reuse."
),
)
],
)
@classmethod
def execute(
cls, latent: dict[str, Any], method: str = "Auto", scale_factor: float = 2.0
) -> tuple[dict[str, Any]]:
"""Delegate method selection and inference to the latent upscale service."""
return (LatentUpscaleService().upscale(latent, method, scale_factor),)
+213
View File
@@ -0,0 +1,213 @@
"""Expose internal cache-keyed capture planning and handoff nodes."""
from __future__ import annotations
from typing import Any
from ..domain.sampler_options import SamplerOptions
from ..runtime.trajectory_capture import TrajectoryCapture
from ..services.noise_capture_request_service import NoiseCaptureRequestService
from .sampler_options_schema import COMFY_IO, OptionsNodeBase, options_input
REQUEST_TYPE = "SIMPLE_SYRUP_CAPTURE_REQUEST"
CAPTURE_TYPE = "SIMPLE_SYRUP_TRAJECTORY_CAPTURE"
class NoiseCaptureRequestV3(OptionsNodeBase):
"""Resolve the receiver's live schedule before source generation."""
@classmethod
def define_schema(cls) -> Any:
"""Include every schedule-affecting dependency in capture caching."""
return COMFY_IO.Schema(
node_id="SimpleSyrup.NoiseCaptureRequest",
display_name="Noise Capture Request (Internal)",
category="SimpleSyrup/Internal",
is_dev_only=True,
description=(
"Coordinates upstream state capture using receiving sampler settings."
),
inputs=[
COMFY_IO.Model.Input(
"model",
tooltip="Receiving sampler model and schedule parameterization.",
),
options_input(COMFY_IO),
COMFY_IO.Latent.Input(
"source_latent",
tooltip="Original source input used only for schedule geometry.",
),
COMFY_IO.Int.Input("steps", tooltip="Receiving sampling step count."),
COMFY_IO.String.Input(
"sampler_name", tooltip="Receiving sampler identifier."
),
COMFY_IO.String.Input(
"scheduler", tooltip="Receiving scheduler identifier."
),
COMFY_IO.Float.Input(
"denoise", tooltip="Receiving sampler denoise strength."
),
COMFY_IO.String.Input(
"geometry_json",
tooltip="Validated resizing operations for schedule planning.",
),
COMFY_IO.Custom(REQUEST_TYPE).Input(
"previous",
optional=True,
tooltip="Other receiving requests sharing this source sampler.",
),
],
outputs=[
COMFY_IO.Custom(REQUEST_TYPE).Output(
"requests", tooltip="Exact source capture coordinates."
)
],
)
@classmethod
def execute(
cls,
model: Any,
source_latent: dict[str, Any],
steps: int,
sampler_name: str,
scheduler: str,
denoise: float,
geometry_json: str,
options: SamplerOptions | None = None,
previous: tuple[float, ...] = (),
) -> tuple[tuple[float, ...]]:
"""Delegate to the shared receiver schedule authority."""
return (
NoiseCaptureRequestService().coordinates(
model=model,
options=options or SamplerOptions(),
steps=steps,
sampler_name=sampler_name,
scheduler=scheduler,
denoise=denoise,
source_latent=source_latent,
geometry_json=geometry_json,
previous=previous,
),
)
class TrajectoryCaptureModelV3(OptionsNodeBase):
"""Install capture on a source MODEL without modifying shared model instances."""
@classmethod
def define_schema(cls) -> Any:
"""Expose only internal derivation dependencies."""
return COMFY_IO.Schema(
node_id="SimpleSyrup.TrajectoryCaptureModel",
display_name="Trajectory Capture Model (Internal)",
category="SimpleSyrup/Internal",
is_dev_only=True,
description=(
"Captures real sampler states for downstream latent noise reuse."
),
inputs=[
COMFY_IO.Model.Input(
"model", tooltip="Source MODEL cloned for capture."
),
COMFY_IO.String.Input(
"sampler_name", tooltip="Source solver identifier."
),
COMFY_IO.Custom(REQUEST_TYPE).Input(
"requests",
optional=True,
tooltip=(
"Requested real-state coordinates; omitted records format "
"provenance only."
),
),
],
outputs=[
COMFY_IO.Model.Output(
"model", tooltip="Source MODEL carrying its capture observer."
),
COMFY_IO.Custom(CAPTURE_TYPE).Output(
"capture", tooltip="Source-local capture lifecycle."
),
],
)
@classmethod
def execute(
cls, model: Any, sampler_name: str, requests: tuple[float, ...] = ()
) -> tuple[Any, TrajectoryCapture]:
"""Keep observer state scoped to this source and cache identity."""
capture = TrajectoryCapture(model, sampler_name, requests)
return capture.prepare_model(), capture
class AttachSamplingProvenanceV3(OptionsNodeBase):
"""Attach a verified source capture before latents enter resizing or decode."""
@classmethod
def define_schema(cls) -> Any:
"""Require both completed source output and its precise capture owner."""
return COMFY_IO.Schema(
node_id="SimpleSyrup.AttachSamplingProvenance",
display_name="Attach Sampling Provenance (Internal)",
category="SimpleSyrup/Internal",
is_dev_only=True,
description=(
"Associates original sampler output with matching captured states."
),
inputs=[
COMFY_IO.Latent.Input(
"latent", tooltip="Completed source sampler output."
),
COMFY_IO.Custom(CAPTURE_TYPE).Input(
"capture",
tooltip="Capture belonging to this exact source sampler execution.",
),
],
outputs=[
COMFY_IO.Latent.Output(
"latent",
tooltip="Original latent with verified sampling provenance.",
)
],
)
@classmethod
def execute(
cls, latent: dict[str, Any], capture: TrajectoryCapture
) -> tuple[dict[str, Any]]:
"""Validate source identity before publishing capture metadata."""
return (capture.attach(latent),)
class NoiseReuseErrorV3(OptionsNodeBase):
"""Surface graph planning failures through normal Comfy execution errors."""
@classmethod
def define_schema(cls) -> Any:
"""Avoid prompt-handler exception swallowing by making failure explicit."""
return COMFY_IO.Schema(
node_id="SimpleSyrup.NoiseReuseError",
display_name="Noise Reuse Validation (Internal)",
category="SimpleSyrup/Internal",
is_dev_only=True,
description=(
"Reports invalid sampling provenance before downstream sampling."
),
inputs=[
COMFY_IO.String.Input(
"message", tooltip="Actionable source graph validation error."
)
],
outputs=[
COMFY_IO.Latent.Output(
"latent", tooltip="Validation must succeed before sampling."
)
],
)
@classmethod
def execute(cls, message: str) -> tuple[dict[str, Any]]:
"""Fail closed instead of substituting fresh noise or encoded pixels."""
raise ValueError(message)
@@ -33,6 +33,7 @@ from .differential_diffusion import (
from .guided_sampling import sample_with_optional_negative
from .inversion_model_factory import InversionModelFactory
from .model_patcher_mutations import ModelUnetWrapperMutation
from .noise_reuse_scope import needs_finishing_model
from .patcher_lifecycle import PATCHER_LIFECYCLE, ModelMutation
from .sampling_model_types import ModelFunctionWrapper
from .tiled_sampling_validation import (
@@ -147,7 +148,7 @@ def sample_contextual_diffusion(
segs=inversion_segs,
region_masks=inversion_region_masks,
)
if noise_inversion is not None
if noise_inversion is not None or needs_finishing_model()
else None
)
noise = sampling_noise.prepare_sampling_noise(
+46
View File
@@ -14,6 +14,8 @@ import torch
from ..domain.noise_inversion import NoiseInversionOptions
from ..shared.logging import get_logger
from .noise_inversion import InversionModelFactory, invert_sampling_noise
from .noise_reuse_sampling import prepare_reused_noise
from .noise_reuse_scope import current_noise_reuse
LOGGER = get_logger(__name__)
@@ -38,6 +40,50 @@ def sample_with_optional_negative(
) -> torch.Tensor:
"""Prepare optional source-derived noise and select the actual Comfy guider."""
reuse = current_noise_reuse()
if reuse is not None:
if noise_inversion is not None:
raise ValueError(
"Noise reuse and full noise inversion cannot own the same "
"sampling pass."
)
noise, expected = prepare_reused_noise(
execution=reuse,
model=model,
fresh=noise,
latent=latent_image,
sigmas=sigmas,
positive=positive,
negative=negative,
cfg=cfg,
seed=seed,
noise_mask=noise_mask,
model_factory=inversion_model_factory,
)
original_callback = callback
def observe_start(
step: int, denoised: torch.Tensor, state: torch.Tensor, total: int
) -> None:
"""Check the actual first state to catch double-noising or schedule
mismatch.
"""
if step == 0 and noise_mask is None:
error = float((state.detach().float().cpu() - expected).abs().max())
if error > 2e-5:
raise ValueError(
f"Noise reuse starting state differs from native "
f"reconstruction: {error:.6g}."
)
LOGGER.info(
"Noise reuse starting state verified",
extra={"operation": "noise_reuse_start", "max_error": error},
)
if original_callback is not None:
original_callback(step, denoised, state, total)
callback = observe_start
if noise_inversion is not None:
inversion = invert_sampling_noise(
model=model,
@@ -0,0 +1,72 @@
"""List checksum-pinned SesquiLSR artifacts in Comfy's standard latent model folder."""
from __future__ import annotations
from importlib import import_module
from .auto_model_artifact import AutoModelArtifact
SESQUI_REVISION = "befae004248c403f38b76b9f65fd43b901ea3eaa"
SESQUI_REPOSITORY = "https://github.com/LoganBooker/SesquiLSR"
_CHECKPOINTS = (
(
"wan",
"Wan21",
12373048,
"d1b4c2c1d80b2af81e6bbfd38501ffb5adffb9f7c0ed2599555b837d0d2d015a",
),
(
"sdxl",
"SDXL",
12293696,
"392c0e36fc0389a170cf6c1682556972d91d750c31aa392cdad1b13cad4e18e4",
),
(
"flux",
"Flux",
12374768,
"dbc6362c910a179a22eca26f656fdacde17eded876ee8f05b8f5b73630e4ddbe",
),
(
"flux2",
"Flux2",
12525400,
"f385ff2250b3f2c72ad422768558414d617b9a006cc0b1ce37a296cc99d17177",
),
)
LATENT_UPSCALE_ARTIFACTS = {
family: AutoModelArtifact(
cache_id=f"sesqui_lsr_{family}_{SESQUI_REVISION}",
filename=f"upscaler_{name}.safetensors",
folder_name="latent_upscale_models",
canonical_subfolder=family,
source_url=(
f"https://raw.githubusercontent.com/LoganBooker/SesquiLSR/"
f"{SESQUI_REVISION}/models/upscaler_{name}.safetensors"
),
source_repo=SESQUI_REPOSITORY,
description=f"SesquiLSR {name} latent upscaler",
sha256=checksum,
file_size_bytes=size,
)
for family, name, size, checksum in _CHECKPOINTS
}
def latent_model_choices() -> list[str]:
"""Advertise pinned downloads and installed checkpoints without hashing or
IO loading.
"""
from ..domain.latent_upscale import NATIVE_LATENT_METHODS
known = [
f"{artifact.canonical_subfolder}/{artifact.filename}"
for artifact in LATENT_UPSCALE_ARTIFACTS.values()
]
installed = [
name.replace("\\", "/")
for name in import_module("folder_paths").get_filename_list(
"latent_upscale_models"
)
]
return list(dict.fromkeys(["Auto", *known, *installed, *NATIVE_LATENT_METHODS]))
@@ -0,0 +1,93 @@
"""Resolve trusted learned checkpoints through registered Comfy model paths."""
from __future__ import annotations
from importlib import import_module
from pathlib import Path
from typing import Any
import torch
from ..domain.latent_upscale import recommended_family
from ..domain.noise_reuse import PROVENANCE_KEY, SamplingProvenance
from .auto_model_resolver import AutoModelResolver
from .latent_upscale_catalog import LATENT_UPSCALE_ARTIFACTS
from .model_downloads import ComfyProgressReporter
def resolve_latent_upscale_model(
latent: dict[str, Any], method: str
) -> tuple[str, Path]:
"""Resolve verified Auto identity or a checkpoint with a recognized architecture."""
provenance = latent.get(PROVENANCE_KEY)
if method == "Auto":
if not isinstance(provenance, SamplingProvenance):
raise ValueError(
"Auto latent upscaling requires sampling provenance; "
"connect a source sampler or select a compatible "
"checkpoint explicitly."
)
family = recommended_family(provenance.latent_format)
return family, AutoModelResolver().resolve(
LATENT_UPSCALE_ARTIFACTS[family], ComfyProgressReporter()
).path
for family, artifact in LATENT_UPSCALE_ARTIFACTS.items():
if method.replace("\\", "/") == f"{family}/{artifact.filename}":
if (
isinstance(provenance, SamplingProvenance)
and recommended_family(provenance.latent_format) != family
):
raise ValueError(
"Selected latent upscale checkpoint does not match the "
"source latent format."
)
return family, AutoModelResolver().resolve(
artifact, ComfyProgressReporter()
).path
folders = import_module("folder_paths")
path = Path(folders.get_full_path_or_raise("latent_upscale_models", method))
state = import_module("comfy.utils").load_torch_file(str(path), safe_load=True)
head = state.get("head_conv.0.weight")
if (
not isinstance(head, torch.Tensor)
or head.ndim != 4
or "upsample.base_taps" not in state
):
raise ValueError(
"Selected checkpoint is not a supported SesquiLSR latent upscaler."
)
named_family = next(
(
family
for family, artifact in LATENT_UPSCALE_ARTIFACTS.items()
if path.name.casefold() == artifact.filename.casefold()
),
None,
)
if named_family is not None:
expected_channels = {"sdxl": 4, "flux": 16, "wan": 16, "flux2": 32}[
named_family
]
if head.shape[1] != expected_channels:
raise ValueError(
"Selected checkpoint filename and latent channels disagree."
)
detected_family = named_family or {4: "sdxl", 32: "flux2"}.get(int(head.shape[1]))
if detected_family is None and head.shape[1] == 16:
if not isinstance(provenance, SamplingProvenance):
raise ValueError(
"A 16-channel checkpoint requires verified Flux or Wan "
"sampling provenance."
)
detected_family = recommended_family(provenance.latent_format)
if detected_family not in LATENT_UPSCALE_ARTIFACTS:
raise ValueError("Selected checkpoint has unsupported latent channels.")
if (
isinstance(provenance, SamplingProvenance)
and recommended_family(provenance.latent_format) != detected_family
):
raise ValueError(
"Selected latent upscale checkpoint does not match "
"the source latent format."
)
return detected_family, path
@@ -37,6 +37,7 @@ from .differential_diffusion import (
from .guided_sampling import sample_with_optional_negative
from .inversion_model_factory import InversionModelFactory
from .model_patcher_mutations import ModelUnetWrapperMutation
from .noise_reuse_scope import needs_finishing_model
from .patcher_lifecycle import PATCHER_LIFECYCLE, ModelMutation
from .sampling_model_types import (
ApplyModel,
@@ -183,7 +184,7 @@ def sample_mixture_of_diffusers(
segs=inversion_segs,
region_masks=inversion_region_masks,
)
if noise_inversion is not None
if noise_inversion is not None or needs_finishing_model()
else None
),
)
@@ -37,6 +37,7 @@ from .differential_diffusion import (
from .guided_sampling import sample_with_optional_negative
from .inversion_model_factory import InversionModelFactory
from .model_patcher_mutations import ModelUnetWrapperMutation
from .noise_reuse_scope import needs_finishing_model
from .patcher_lifecycle import PATCHER_LIFECYCLE, ModelMutation
from .sampling_model_types import (
ApplyModel,
@@ -184,7 +185,7 @@ def sample_multidiffusion(
segs=inversion_segs,
region_masks=inversion_region_masks,
)
if noise_inversion is not None
if noise_inversion is not None or needs_finishing_model()
else None
),
)
@@ -0,0 +1,47 @@
"""Separate clean references and residuals with the model's actual parameterization."""
from __future__ import annotations
from typing import Any
import torch
def noise_coefficients(
scaling: Any, clean: torch.Tensor, coordinate: float
) -> tuple[torch.Tensor, torch.Tensor]:
"""Validate affine noise mixing instead of assuming DDPM alpha/sigma units."""
sigma = clean.new_tensor(coordinate)
base = scaling.noise_scaling(sigma, torch.zeros_like(clean), clean, False)
amplitude = (
scaling.noise_scaling(sigma, torch.ones_like(clean), clean, False) - base
)
if not isinstance(base, torch.Tensor) or not isinstance(amplitude, torch.Tensor):
raise TypeError("Native noise scaling must produce tensor coefficients.")
if base.shape != clean.shape or amplitude.shape != clean.shape:
raise ValueError("Native noise scaling must preserve the clean latent shape.")
if not bool(torch.isfinite(amplitude).all()) or bool((amplitude == 0).any()):
raise ValueError(
"Capture coordinate must have finite invertible noise amplitude."
)
probe = scaling.noise_scaling(sigma, torch.full_like(clean, 0.37), clean, False)
if not isinstance(probe, torch.Tensor) or not torch.allclose(
probe, base + amplitude * 0.37, atol=2e-5, rtol=2e-5
):
raise ValueError("Native noise scaling must be affine in its noise argument.")
return base, amplitude
def recover_residual(
scaling: Any, clean: torch.Tensor, state: torch.Tensor, coordinate: float
) -> torch.Tensor:
"""Require the recovered residual to reconstruct its actual captured state."""
if clean.shape != state.shape:
raise ValueError("Captured state and original clean reference shapes differ.")
base, amplitude = noise_coefficients(scaling, clean, coordinate)
residual = (state - base) / amplitude
if not bool(torch.isfinite(residual).all()) or not torch.allclose(
base + amplitude * residual, state, atol=2e-5, rtol=2e-5
):
raise ValueError("Recovered residual cannot reconstruct the captured state.")
return residual
@@ -0,0 +1,106 @@
"""Advance an already reconstructed state through full-resolution inversion
finishing.
"""
from __future__ import annotations
from importlib import import_module
from typing import Any
import torch
from ..domain.inversion_solver import InversionSolverEvidence, integrate_inversion
from ..domain.noise_reuse import NoiseReuseOptions
from ..shared.logging import get_logger
from .noise_inversion import InversionModelFactory
from .noise_reuse_scope import noise_reuse_scope
LOGGER = get_logger(__name__)
def finish_reused_state(
*,
model: Any,
latent: torch.Tensor,
initial: torch.Tensor,
begin: float,
target: float,
options: NoiseReuseOptions,
positive: Any,
negative: Any,
cfg: float,
seed: int | None,
noise_mask: Any,
model_factory: InversionModelFactory | None,
) -> torch.Tensor:
"""Run only the requested finishing evaluations, using native CFG and
spatial policy.
"""
from .guided_sampling import sample_with_optional_negative
if options.finishing_steps < 1 or begin >= target:
raise ValueError(
"Finishing requires positive steps and a start below the "
"sampler coordinate."
)
device = torch.device(model.load_device)
schedule = torch.linspace(begin, target, options.finishing_steps + 1, device=device)
evidence = InversionSolverEvidence()
endpoints: list[torch.Tensor] = []
def integrate(
model_fn: Any,
state: torch.Tensor,
sigmas: torch.Tensor,
extra_args: dict[str, Any],
callback: Any,
disable: bool,
) -> torch.Tensor:
"""Consume the captured initial state without adding a second layer of noise."""
state = initial.to(state)
def evaluate(x: torch.Tensor, sigma: torch.Tensor, index: int) -> torch.Tensor:
"""Use the production inversion velocity from native denoised
predictions.
"""
prediction = model_fn(x, sigma * x.new_ones((x.shape[0],)), **extra_args)
if not isinstance(prediction, torch.Tensor):
raise TypeError("Finishing model must return tensor predictions.")
return (x - prediction) / sigma
endpoint = integrate_inversion(
state, sigmas, evaluate, method=options.finishing_method, evidence=evidence
)
endpoints.append(endpoint.detach().float().cpu())
return endpoint
stage_model = model_factory(latent) if model_factory is not None else model
with noise_reuse_scope(None):
sample_with_optional_negative(
comfy_sample=import_module("comfy.sample"),
model=stage_model,
noise=torch.zeros_like(latent),
cfg=cfg,
sampler=import_module("comfy.samplers").KSAMPLER(integrate),
sigmas=schedule,
positive=positive,
negative=negative,
latent_image=latent,
noise_mask=noise_mask,
seed=seed,
disable_pbar=True,
)
if len(endpoints) != 1:
raise RuntimeError("Finishing must produce exactly one native endpoint.")
LOGGER.info(
"Captured noise finishing completed",
extra={
"operation": "noise_reuse_finish",
"steps": options.finishing_steps,
"evaluations": evidence.evaluations,
"begin_sigma": begin,
"target_sigma": target,
"coarse_inversion_evaluations": 0,
},
)
return endpoints[0]
+238
View File
@@ -0,0 +1,238 @@
"""Resolve strict latent lineage and inject cache-keyed upstream capture
dependencies.
"""
from __future__ import annotations
import json
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any
from ..domain.graph_provenance import BrokenProvenance, GraphLink
from .comfy_graph_provenance import (
NodeRegistry,
parse_graph_link,
trace_vae_decode_provenance,
)
SOURCE_SAMPLERS = frozenset(
{
"KSampler",
"KSamplerAdvanced",
"SimpleSyrup.KSampler",
"SimpleSyrup.KSamplerExtras",
"SimpleSyrup.KSamplerContextualDiffusion",
"SimpleSyrup.KSamplerTiledDiffusion",
}
)
@dataclass(frozen=True, slots=True)
class CaptureConsumer:
"""Describe source capture requirements without storing runtime model objects."""
source: str
receiver: str | None
geometry_json: str
def source_lineage(
prompt: Mapping[str, Any], start: object, registry: NodeRegistry
) -> tuple[str, str]:
"""Trace only latent-preserving resizing and unmodified decoded-image bridges."""
current = parse_graph_link(start)
visited: set[GraphLink] = set()
geometry: list[dict[str, object]] = []
while current is not None and len(visited) < 128:
if current in visited:
raise ValueError("Noise reuse latent provenance contains a cycle.")
visited.add(current)
node = prompt.get(current[0])
if not isinstance(node, Mapping) or current[1] != 0:
raise ValueError(
"Noise reuse source must be a supported sampler's latent output."
)
kind, inputs = node.get("class_type"), node.get("inputs")
if not isinstance(inputs, Mapping):
raise ValueError("Noise reuse source node has missing inputs.")
if kind in SOURCE_SAMPLERS:
return current[0], json.dumps(list(reversed(geometry)))
if kind == "SimpleSyrup.LatentFromImage":
decoded = trace_vae_decode_provenance(prompt, inputs.get("image"), registry)
if isinstance(decoded, BrokenProvenance):
raise ValueError(
f"Noise reuse image provenance is broken: {decoded.reason}."
)
current = decoded.samples_link
continue
if kind in {"SimpleSyrup.SimpleUpscaleLatent", "LatentUpscaleBy"}:
key = (
"scale_factor"
if kind == "SimpleSyrup.SimpleUpscaleLatent"
else "scale_by"
)
scale = inputs.get(key, 2.0 if key == "scale_factor" else 1.5)
geometry.append({"scale": scale})
current = parse_graph_link(
inputs.get("latent" if key == "scale_factor" else "samples")
)
continue
if kind == "LatentUpscale":
if inputs.get("crop", "disabled") != "disabled":
raise ValueError(
"Noise reuse requires latent upscaling with crop disabled."
)
geometry.append({"pixels": [inputs.get("height"), inputs.get("width")]})
current = parse_graph_link(inputs.get("samples"))
continue
class_def = registry.get(str(kind))
rules = getattr(class_def, "GRAPH_PASSTHROUGH_OUTPUTS", None)
if isinstance(rules, Mapping) and isinstance(rules.get(0), str):
current = parse_graph_link(inputs.get(rules[0]))
continue
raise ValueError(
f"Noise reuse cannot resolve latent provenance through "
f"{kind}. Use a direct sampler latent, Simple Upscale "
f"Latent, or Latent From Image."
)
raise ValueError("Noise reuse requires a graph-connected source sampler latent.")
def has_noise_reuse(prompt: Mapping[str, Any], start: object) -> bool:
"""Detect potentially enabled reuse without evaluating options or loading models."""
current = parse_graph_link(start)
visited: set[GraphLink] = set()
while current is not None and current not in visited and len(visited) < 128:
visited.add(current)
node = prompt.get(current[0])
if not isinstance(node, Mapping):
return False
inputs = node.get("inputs", {})
if not isinstance(inputs, Mapping):
return False
if (
node.get("class_type") == "SimpleSyrup.NoiseReuseOptions"
and inputs.get("enabled", True) is not False
):
return True
current = parse_graph_link(inputs.get("options"))
return False
def rewrite_noise_reuse_graph(prompt: dict[str, Any], registry: NodeRegistry) -> None:
"""Group receiving requests and redirect every source latent consumer once."""
if any(
isinstance(node, Mapping)
and node.get("class_type")
in {"SimpleSyrup.TrajectoryCaptureModel", "SimpleSyrup.NoiseReuseError"}
for node in prompt.values()
):
return
consumers: list[CaptureConsumer] = []
failures: dict[str, str] = {}
for node_id, node in tuple(prompt.items()):
if not isinstance(node, Mapping):
continue
kind, inputs = node.get("class_type"), node.get("inputs", {})
if not isinstance(inputs, Mapping):
continue
receiver = kind == "SimpleSyrup.KSampler" and has_noise_reuse(
prompt, inputs.get("options")
)
auto = (
kind == "SimpleSyrup.SimpleUpscaleLatent"
and inputs.get("method", "Auto") == "Auto"
)
if not receiver and not auto:
continue
try:
source, geometry = source_lineage(
prompt, inputs.get("latent_image" if receiver else "latent"), registry
)
consumers.append(
CaptureConsumer(source, node_id if receiver else None, geometry)
)
except ValueError as error:
if receiver:
failures[node_id] = str(error)
for source in sorted({consumer.source for consumer in consumers}):
_inject_source(
prompt,
source,
[consumer for consumer in consumers if consumer.source == source],
)
for node_id, message in failures.items():
error_id = _unique_id(prompt, f"{node_id}.noise_reuse_error")
prompt[error_id] = {
"class_type": "SimpleSyrup.NoiseReuseError",
"inputs": {"message": message},
}
inputs = prompt[node_id]["inputs"]
inputs[
"latent_image"
if prompt[node_id]["class_type"] == "SimpleSyrup.KSampler"
else "latent"
] = [error_id, 0]
def _inject_source(
prompt: dict[str, Any], source: str, consumers: list[CaptureConsumer]
) -> None:
"""Put receiver settings in the source MODEL dependency chain, then attach
its result.
"""
source_inputs = prompt[source]["inputs"]
request: list[object] | None = None
for consumer in consumers:
if consumer.receiver is None:
continue
receiver = prompt[consumer.receiver]["inputs"]
request_id = _unique_id(prompt, f"{consumer.receiver}.noise_capture_request")
values = {
key: receiver[key]
for key in ("model", "steps", "sampler_name", "scheduler")
}
values.update(
options=receiver["options"],
denoise=receiver.get("denoise", 1.0),
source_latent=source_inputs["latent_image"],
geometry_json=consumer.geometry_json,
)
if request is not None:
values["previous"] = request
prompt[request_id] = {
"class_type": "SimpleSyrup.NoiseCaptureRequest",
"inputs": values,
}
request = [request_id, 0]
capture_id = _unique_id(prompt, f"{source}.trajectory_capture")
values = {
"model": source_inputs["model"],
"sampler_name": source_inputs["sampler_name"],
}
if request is not None:
values["requests"] = request
prompt[capture_id] = {
"class_type": "SimpleSyrup.TrajectoryCaptureModel",
"inputs": values,
}
source_inputs["model"] = [capture_id, 0]
attach_id = _unique_id(prompt, f"{source}.sampling_provenance")
for node in prompt.values():
inputs = node.get("inputs", {})
for key, value in tuple(inputs.items()):
if parse_graph_link(value) == (source, 0):
inputs[key] = [attach_id, 0]
prompt[attach_id] = {
"class_type": "SimpleSyrup.AttachSamplingProvenance",
"inputs": {"latent": [source, 0], "capture": [capture_id, 1]},
}
def _unique_id(prompt: Mapping[str, Any], base: str) -> str:
"""Avoid collisions with user-authored or earlier injected graph identifiers."""
result = base
while result in prompt:
result += "_"
return result
@@ -0,0 +1,32 @@
"""Register capture coordination before Comfy validates or executes source sampling."""
from __future__ import annotations
from importlib import import_module
from typing import Any
from ..shared.logging import get_logger
from .noise_reuse_graph import rewrite_noise_reuse_graph
LOGGER = get_logger(__name__)
_MARKER = "_simple_syrup_noise_reuse_prompt_handler"
def register_noise_reuse_prompt_handler() -> None:
"""Register once when the native PromptServer becomes available."""
try:
instance = getattr(import_module("server").PromptServer, "instance", None)
except (ImportError, AttributeError):
return
if instance is not None and not getattr(instance, _MARKER, False):
instance.add_on_prompt_handler(_rewrite_prompt)
setattr(instance, _MARKER, True)
def _rewrite_prompt(data: Any) -> Any:
"""Use normal error-producing nodes for invalid source provenance."""
if isinstance(data, dict) and isinstance(data.get("prompt"), dict):
rewrite_noise_reuse_graph(
data["prompt"], import_module("nodes").NODE_CLASS_MAPPINGS
)
return data
@@ -0,0 +1,84 @@
"""Prepare captured residuals for the shared native sampler starting-noise input."""
from __future__ import annotations
from typing import Any
import torch
from ..domain.spectral_noise import lift_residual
from ..shared.logging import get_logger
from .native_noise_scaling import noise_coefficients, recover_residual
from .noise_inversion import InversionModelFactory, validate_inversion_target
from .noise_reuse_finishing import finish_reused_state
from .noise_reuse_scope import NoiseReuseExecution
LOGGER = get_logger(__name__)
def prepare_reused_noise(
*,
execution: NoiseReuseExecution,
model: Any,
fresh: torch.Tensor,
latent: torch.Tensor,
sigmas: torch.Tensor,
positive: Any,
negative: Any,
cfg: float,
seed: int | None,
noise_mask: Any,
model_factory: InversionModelFactory | None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Lift source coefficients and return the expected actual model-space
starting state.
"""
if execution.consumed:
raise ValueError(
"Noise reuse supports one native sampling pass per "
"receiver; use separate samplers for conditioning "
"batches."
)
execution.consumed = True
target = validate_inversion_target(model, sigmas)
coordinate = execution.options.capture_coordinate(target)
provenance = execution.provenance
if type(model.model.latent_format).__name__ != provenance.latent_format:
raise ValueError(
"Noise reuse source and receiving model use different latent formats."
)
scaling = model.get_model_object("model_sampling")
residual = provenance.residual_at(coordinate)
lifted = lift_residual(residual, fresh.detach().float().cpu())
clean = model.model.process_latent_in(latent).detach().float().cpu()
if execution.options.finishing_steps:
base, amplitude = noise_coefficients(scaling, clean, coordinate)
endpoint = finish_reused_state(
model=model,
latent=latent,
initial=base + amplitude * lifted,
begin=coordinate,
target=target,
options=execution.options,
positive=positive,
negative=negative,
cfg=cfg,
seed=seed,
noise_mask=noise_mask,
model_factory=model_factory,
)
lifted = recover_residual(scaling, clean, endpoint, target)
base, amplitude = noise_coefficients(scaling, clean, target)
expected = base + amplitude * lifted
LOGGER.info(
"Captured sampling noise prepared",
extra={
"operation": "noise_reuse",
"capture_sigma": coordinate,
"target_sigma": target,
"source_shape": tuple(residual.shape),
"target_shape": tuple(lifted.shape),
"finishing_steps": execution.options.finishing_steps,
},
)
return lifted.to(fresh), expected
+45
View File
@@ -0,0 +1,45 @@
"""Own one request-scoped noise handoff across existing sampler routing."""
from __future__ import annotations
from collections.abc import Iterator
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass
from ..domain.noise_reuse import NoiseReuseOptions, SamplingProvenance
@dataclass(slots=True)
class NoiseReuseExecution:
"""Keep a single receiver's immutable capture and consumption lifecycle."""
options: NoiseReuseOptions
provenance: SamplingProvenance
consumed: bool = False
_CURRENT: ContextVar[NoiseReuseExecution | None] = ContextVar(
"simple_syrup_noise_reuse", default=None
)
def current_noise_reuse() -> NoiseReuseExecution | None:
"""Resolve only the receiving service's current execution, never cached globals."""
return _CURRENT.get()
def needs_finishing_model() -> bool:
"""Request a spatial inversion factory only for enabled finishing work."""
current = _CURRENT.get()
return current is not None and current.options.finishing_steps > 0
@contextmanager
def noise_reuse_scope(execution: NoiseReuseExecution | None) -> Iterator[None]:
"""Restore the caller's context even when validation or native inference fails."""
token = _CURRENT.set(execution)
try:
yield
finally:
_CURRENT.reset(token)
@@ -27,6 +27,7 @@ from .differential_diffusion import (
)
from .guided_sampling import sample_with_optional_negative
from .model_patcher_mutations import ModelCalcCondBatchMutation
from .noise_reuse_scope import needs_finishing_model
from .patcher_lifecycle import PATCHER_LIFECYCLE, ModelMutation
from .regional_inversion_model_factory import RegionalInversionModelFactory
from .regional_multidiffusion_prediction import (
@@ -172,7 +173,7 @@ def sample_regional_multidiffusion(
global_prompt_weight=global_prompt_weight,
differential_diffusion=differential_diffusion,
)
if noise_inversion is not None
if noise_inversion is not None or needs_finishing_model()
else None
),
)
@@ -0,0 +1,93 @@
"""Adapt SesquiLSR inference to Comfy external latent layouts and device management."""
from __future__ import annotations
from importlib import import_module
from pathlib import Path
from typing import Any
import torch
from ..domain.latent_upscale import learned_upscale_stages
from ..shared.logging import get_logger
LOGGER = get_logger(__name__)
_CHANNELS = {"sdxl": 4, "flux": 16, "wan": 16, "flux2": 32}
def upscale_with_sesqui(
samples: torch.Tensor, *, path: Path, family: str, target: tuple[int, int]
) -> torch.Tensor:
"""Run float32 learned stages in raw VAE latent units, without image decoding."""
module: Any = import_module("simple_syrup.third_party.sesqui_lsr")
management = import_module("comfy.model_management")
utils = import_module("comfy.utils")
channels = _CHANNELS[family]
external_channels = (
128 if family == "flux2" and samples.shape[1] == 128 else channels
)
if samples.ndim not in (4, 5) or samples.shape[1] != external_channels:
raise ValueError(
f"SesquiLSR {family} requires BCHW or BCTHW latents with "
f"{external_channels} channels."
)
if not samples.is_floating_point() or not bool(torch.isfinite(samples).all()):
raise ValueError("Latent upscaling requires finite floating-point samples.")
network = module.LatentUpscaler(in_channels=channels)
network.load_state_dict(
utils.load_torch_file(str(path), safe_load=True), strict=True
)
device = management.get_torch_device()
management.free_memory(
management.module_size(network) + samples.numel() * 64, device
)
network.eval().requires_grad_(False).to(device=device, dtype=torch.float32)
batch, frames = (
int(samples.shape[0]),
int(samples.shape[2]) if samples.ndim == 5 else 1,
)
flat = (
samples.permute(0, 2, 1, 3, 4).reshape(
batch * frames, external_channels, *samples.shape[-2:]
)
if samples.ndim == 5
else samples
)
adaptor = (
module.make_flux2()
if external_channels == 128
else module.make_identity(channels)
)
current: torch.Tensor = adaptor.to_vae_latent(
flat.to(device=device, dtype=torch.float32)
)
vae_target = adaptor.vae_target_size(target)
stages = learned_upscale_stages(
(int(current.shape[-2]), int(current.shape[-1])), vae_target
)
try:
with torch.inference_mode():
for size in stages:
current = network(current, size)
output: torch.Tensor = adaptor.from_vae_latent(current.float()).to(samples)
finally:
network.to(management.unet_offload_device())
if output.shape[-2:] != target:
raise RuntimeError(
"Learned latent upscaler did not produce the requested dimensions."
)
if samples.ndim == 5:
output = output.reshape(batch, frames, external_channels, *target).permute(
0, 2, 1, 3, 4
)
LOGGER.info(
"Learned latent upscale completed",
extra={
"operation": "latent_upscale",
"family": family,
"model": path.name,
"stages": stages,
"output_shape": tuple(output.shape),
},
)
return output
+192
View File
@@ -0,0 +1,192 @@
"""Observe one upstream sampler through Comfy's per-model sampler wrapper."""
from __future__ import annotations
from copy import copy
from dataclasses import dataclass, field
from importlib import import_module
from typing import Any
import torch
from ..domain.noise_reuse import PROVENANCE_KEY, SamplingProvenance
from ..shared.logging import get_logger
from .model_patcher_mutations import ModelKeyedWrapperMutation
from .native_noise_scaling import recover_residual
from .patcher_lifecycle import PATCHER_LIFECYCLE
LOGGER = get_logger(__name__)
SUPPORTED_CAPTURE_SAMPLERS = frozenset({"euler", "euler_ancestral"})
def insert_capture_coordinates(
sigmas: torch.Tensor, coordinates: tuple[float, ...]
) -> torch.Tensor:
"""Insert exact float32 boundaries without replacing native scheduler points."""
if sigmas.ndim != 1 or len(sigmas) < 2 or not bool(torch.isfinite(sigmas).all()):
raise ValueError(
"Source capture requires a finite descending sampling schedule."
)
if not bool((sigmas[:-1] > sigmas[1:]).all()):
raise ValueError(
"Source capture requires strictly descending sigma boundaries."
)
if float(sigmas[-1]) != 0:
raise ValueError(
"Source capture requires a schedule ending at zero noise; "
"disable leftover noise on the originating sampler."
)
result = sigmas
for value in sorted(set(coordinates), reverse=True):
coordinate = sigmas.new_tensor(value)
if not 0 < float(coordinate) < float(sigmas[0]):
raise ValueError(
f"Requested source sigma {value:.9g} lies outside the "
f"source sampling schedule."
)
if not bool((result == coordinate).any()):
index = int((result > coordinate).sum())
result = torch.cat((result[:index], coordinate.reshape(1), result[index:]))
return result
@dataclass(slots=True)
class TrajectoryCapture:
"""Keep capture storage local to an injected, cache-keyed MODEL derivation."""
model: Any
sampler_name: str
coordinates: tuple[float, ...]
states: dict[float, torch.Tensor] = field(default_factory=dict)
clean: torch.Tensor | None = None
schedule: torch.Tensor | None = None
def __post_init__(self) -> None:
"""Store the exact float32 coordinates used by native sampler schedules."""
self.coordinates = tuple(
sorted(
{float(torch.tensor(value)) for value in self.coordinates}, reverse=True
)
)
def prepare_model(self) -> Any:
"""Clone only this source MODEL and leave unrelated samplers untouched."""
if self.coordinates and self.sampler_name not in SUPPORTED_CAPTURE_SAMPLERS:
raise ValueError(
"Noise reuse source capture supports Euler and Euler "
"ancestral; select a supported source sampler."
)
if not self.coordinates:
return self.model
wrappers = import_module("comfy.patcher_extension").WrappersMP
return PATCHER_LIFECYCLE.derive_model(
self.model,
(
ModelKeyedWrapperMutation(
wrappers.SAMPLER_SAMPLE,
"simple_syrup.trajectory_capture",
self.observe,
),
),
operation="SimpleSyrup trajectory capture",
)
def observe(
self, executor: Any, *arguments: Any, **keyword_arguments: Any
) -> torch.Tensor:
"""Wrap the individual native solver instance and record its pre-update x."""
local = copy(executor.class_obj)
solver = local.sampler_function
self.states.clear()
self.clean = None
def recorded(
model: Any, state: torch.Tensor, sigmas: torch.Tensor, **kwargs: Any
) -> torch.Tensor:
"""Use native solver randomness and preserve its original callback."""
schedule = insert_capture_coordinates(sigmas, self.coordinates)
self.schedule = schedule.detach().float().cpu().clone()
original_callback = kwargs.get("callback")
def observe_state(record: dict[str, Any]) -> None:
"""Capture actual noisy state, rejecting churned or ambiguous
coordinates.
"""
coordinate = float(schedule[int(record["i"])])
if coordinate in self.coordinates:
if float(record.get("sigma_hat", record["sigma"])) != coordinate:
raise ValueError(
"Noise reuse cannot capture churned source states."
)
self.states[coordinate] = record["x"].detach().float().cpu().clone()
if original_callback is not None:
original_callback(record)
kwargs["callback"] = observe_state
result: torch.Tensor = solver(model, state, schedule, **kwargs)
self.clean = result.detach().float().cpu().clone()
return result
local.sampler_function = recorded
isolated = copy(executor)
isolated.original = local.sample
result: torch.Tensor = isolated(*arguments, **keyword_arguments)
return result
def attach(self, latent: dict[str, Any]) -> dict[str, Any]:
"""Associate captured states with their matching external clean sampler
output.
"""
samples = latent.get("samples")
if not isinstance(samples, torch.Tensor):
raise TypeError(
"Source sampler output must contain a latent samples tensor."
)
actual = self.model.model.process_latent_in(samples).detach().float().cpu()
if self.coordinates:
if (
self.clean is None
or self.schedule is None
or set(self.states) != set(self.coordinates)
):
raise ValueError(
"Source trajectory capture is incomplete; run the full "
"connected workflow again."
)
if not torch.allclose(actual, self.clean, atol=1e-5, rtol=1e-5):
raise ValueError(
"Captured trajectory endpoint does not match the source "
"sampler output."
)
actual = self.clean
result = latent.copy()
result[PROVENANCE_KEY] = SamplingProvenance(
type(self.model.model.latent_format).__name__,
actual,
tuple(
(coordinate, self.states[coordinate]) for coordinate in self.coordinates
),
self.schedule,
tuple(
(
coordinate,
recover_residual(
self.model.get_model_object("model_sampling"),
actual,
self.states[coordinate],
coordinate,
),
)
for coordinate in self.coordinates
),
)
LOGGER.info(
"Sampling provenance attached",
extra={
"operation": "trajectory_capture",
"sampler": self.sampler_name,
"coordinates": self.coordinates,
"source_shape": tuple(actual.shape),
},
)
return result
@@ -17,6 +17,7 @@ from ..runtime import sampling_noise, sampling_samplers, sampling_schedulers
from ..runtime.comfy_latent_normalization import COMFY_LATENT_NORMALIZER
from ..runtime.guided_sampling import sample_with_optional_negative
from ..runtime.inversion_model_factory import InversionModelFactory
from ..runtime.noise_reuse_scope import needs_finishing_model
from ..shared.logging import get_logger
Latent: TypeAlias = dict[str, Any]
@@ -119,7 +120,7 @@ class KSamplerSamplingService:
canvas_width=int(latent_samples.shape[-1]),
canvas_height=int(latent_samples.shape[-2]),
)
if noise_inversion is not None
if noise_inversion is not None or needs_finishing_model()
else None,
)
if not isinstance(samples, torch.Tensor):
@@ -190,7 +191,7 @@ class KSamplerSamplingService:
canvas_width=int(latent_samples.shape[-1]),
canvas_height=int(latent_samples.shape[-2]),
)
if noise_inversion is not None
if noise_inversion is not None or needs_finishing_model()
else None,
)
)
@@ -0,0 +1,26 @@
"""Recover original latents from image graph provenance without pixel encoding."""
from __future__ import annotations
from collections.abc import Mapping
from typing import Any
from ..domain.graph_provenance import BrokenProvenance, GraphLink
from ..runtime.comfy_graph_provenance import NodeRegistry, trace_vae_decode_provenance
class LatentProvenanceService:
"""Admit only unmodified image paths backed by a VAE decode source."""
def recover_link(
self, prompt: Mapping[str, Any], image: object, registry: NodeRegistry
) -> GraphLink:
"""Return the original samples link or a specific actionable graph error."""
provenance = trace_vae_decode_provenance(prompt, image, registry)
if isinstance(provenance, BrokenProvenance):
raise ValueError(
f"Latent From Image requires an unmodified VAE Decode "
f"source: {provenance.reason}. Connect the original "
f"decoded image or a direct latent connection."
)
return provenance.samples_link
@@ -0,0 +1,78 @@
"""Own latent upscale method selection, provenance preservation and mask geometry."""
from __future__ import annotations
from importlib import import_module
from typing import Any
import torch
from ..domain.latent_upscale import (
FORMAT_FAMILIES,
NATIVE_LATENT_METHODS,
upscale_target,
)
from ..domain.noise_reuse import PROVENANCE_KEY, SamplingProvenance
from ..runtime.latent_upscale_models import resolve_latent_upscale_model
from ..runtime.sesqui_latent_upscale import upscale_with_sesqui
from ..shared.logging import get_logger
LOGGER = get_logger(__name__)
class LatentUpscaleService:
"""Resolve learned checkpoints or delegate native scaling to Comfy's node."""
def upscale(
self, latent: dict[str, Any], method: str, scale_factor: float
) -> dict[str, Any]:
"""Use native bislerp for unknown Auto formats and preserve provenance."""
samples = latent.get("samples")
if not isinstance(samples, torch.Tensor) or samples.ndim not in (4, 5):
raise ValueError("Simple Upscale Latent requires BCHW or BCTHW samples.")
if not samples.is_floating_point() or not bool(torch.isfinite(samples).all()):
raise ValueError(
"Simple Upscale Latent requires finite floating-point samples."
)
mask = latent.get("noise_mask")
if mask is not None and (not isinstance(mask, torch.Tensor) or mask.ndim < 2):
raise TypeError("Latent noise mask must contain spatial tensor dimensions.")
target = upscale_target(
(int(samples.shape[-2]), int(samples.shape[-1])), scale_factor
)
if method == "Auto":
provenance = latent.get(PROVENANCE_KEY)
latent_format = (
provenance.latent_format
if isinstance(provenance, SamplingProvenance)
else None
)
if latent_format not in FORMAT_FAMILIES:
method = "bislerp"
LOGGER.info(
"Simple Upscale Latent selected native bislerp: "
"latent_format=%s scale_factor=%s",
latent_format or "unknown",
scale_factor,
)
if method in NATIVE_LATENT_METHODS:
result: dict[str, Any] = (
import_module("nodes")
.LatentUpscaleBy()
.upscale(latent, method, scale_factor)[0]
)
else:
family, path = resolve_latent_upscale_model(latent, method)
result = latent.copy()
result["samples"] = upscale_with_sesqui(
samples, path=path, family=family, target=target
)
if mask is not None:
assert isinstance(mask, torch.Tensor)
resized = torch.nn.functional.interpolate(
mask.reshape(-1, 1, *mask.shape[-2:]).float(),
size=target,
mode="nearest",
)
result["noise_mask"] = resized.reshape(*mask.shape[:-2], *target).to(mask)
return result
@@ -0,0 +1,111 @@
"""Resolve upstream capture requests from the receiving sampler's live settings."""
from __future__ import annotations
import json
from typing import Any
import torch
from ..domain.sampler_options import SamplerOptions
from ..runtime.comfy_latent_normalization import COMFY_LATENT_NORMALIZER
from ..runtime.noise_inversion import validate_inversion_target
from ..runtime.sampling_schedulers import SchedulerView, calculate_sigmas
class NoiseCaptureRequestService:
"""Evaluate schedule planning before the originating sampler can execute."""
def coordinates(
self,
*,
model: Any,
options: SamplerOptions,
steps: int,
sampler_name: str,
scheduler: str,
denoise: float,
source_latent: dict[str, Any],
geometry_json: str,
previous: tuple[float, ...] = (),
) -> tuple[float, ...]:
"""Calculate the same model-view schedule as the receiving sampling route."""
if not isinstance(options, SamplerOptions):
raise TypeError("Capture planning requires SimpleSyrup sampler options.")
if options.noise_reuse is None:
return previous
samples = source_latent.get("samples")
if not isinstance(samples, torch.Tensor):
raise TypeError("Capture planning requires source latent geometry.")
context, tiling = options.contextual_diffusion, options.tiling
if context is not None:
view = SchedulerView(context.context_size, context.context_size)
elif tiling is not None:
view = SchedulerView(tiling.width, tiling.height)
elif scheduler == "Flux2":
samples = COMFY_LATENT_NORMALIZER.normalize(
model=model,
samples=samples,
spatial_downscale_ratio=source_latent.get("downscale_ratio_spacial"),
)
view = self._view(samples, geometry_json)
else:
view = SchedulerView.from_tensor(samples)
sigmas = calculate_sigmas(
model, scheduler, sampler_name, steps, denoise, view=view
)
target = validate_inversion_target(model, sigmas)
coordinate = options.noise_reuse.capture_coordinate(target)
return tuple(sorted({*previous, coordinate}, reverse=True))
def _view(self, samples: torch.Tensor, geometry_json: str) -> SchedulerView:
"""Resolve only validated native sizing operations for resolution-
dependent schedules.
"""
operations = json.loads(geometry_json)
if not isinstance(operations, list):
raise ValueError("Noise capture geometry must be an operation list.")
height, width = int(samples.shape[-2]), int(samples.shape[-1])
for operation in operations:
if isinstance(operation, dict) and set(operation) == {"pixels"}:
pixels = operation["pixels"]
if (
not isinstance(pixels, list)
or len(pixels) != 2
or not all(type(value) is int and value >= 0 for value in pixels)
):
raise ValueError(
"Resolution-dependent noise capture requires literal "
"non-negative pixel dimensions."
)
pixel_height, pixel_width = pixels
if pixel_height == 0 and pixel_width == 0:
continue
if pixel_width == 0:
pixel_height = max(64, pixel_height)
pixel_width = max(64, round(width * pixel_height / height))
elif pixel_height == 0:
pixel_width = max(64, pixel_width)
pixel_height = max(64, round(height * pixel_width / width))
else:
pixel_width, pixel_height = (
max(64, pixel_width),
max(64, pixel_height),
)
height, width = pixel_height // 8, pixel_width // 8
continue
if not isinstance(operation, dict) or set(operation) != {"scale"}:
raise ValueError(
"Noise capture geometry contains an unsupported operation."
)
scale = operation["scale"]
if (
not isinstance(scale, (int, float))
or isinstance(scale, bool)
or not 0.01 <= scale <= 8
):
raise ValueError(
"Noise capture geometry scale must be between 0.01 and 8."
)
height, width = round(height * scale), round(width * scale)
return SchedulerView(width, height)
@@ -14,9 +14,11 @@ from ..domain.attention_coupling_request import (
AttentionCouplingRequestMode,
classify_attention_coupling_request,
)
from ..domain.noise_reuse import PROVENANCE_KEY, SamplingProvenance
from ..domain.sampler_options import SamplerOptions, TilingOptions
from ..runtime import sampling_schedulers
from ..runtime.noise_inversion import validate_inversion_target
from ..runtime.noise_reuse_scope import NoiseReuseExecution, noise_reuse_scope
from ..runtime.tiled_sampling_validation import validate_sampling_controls
from ..shared.logging import get_logger
from .attention_coupling_sampling_service import AttentionCouplingSamplingService
@@ -90,6 +92,57 @@ class SamplerOptionsSamplingService:
options: SamplerOptions | None = None,
segs: object | None = None,
region_masks: object | None = None,
) -> dict[str, Any]:
"""Own the receiving request lifecycle across every capability route."""
reuse = options.noise_reuse if isinstance(options, SamplerOptions) else None
execution = None
if reuse is not None:
provenance = latent_image.get(PROVENANCE_KEY)
if not isinstance(provenance, SamplingProvenance):
raise ValueError(
"Noise Reuse Options requires captured sampling provenance. "
"Connect a supported sampler's latent directly or through "
"Latent From Image."
)
execution = NoiseReuseExecution(reuse, provenance)
with noise_reuse_scope(execution):
result = self._sample_configured(
model=model,
seed=seed,
steps=steps,
cfg=cfg,
sampler_name=sampler_name,
scheduler=scheduler,
positive=positive,
negative=negative,
latent_image=latent_image,
denoise=denoise,
options=options,
segs=segs,
region_masks=region_masks,
)
if execution is not None and not execution.consumed:
raise RuntimeError(
"Receiving sampler did not consume the captured noise request."
)
return result
def _sample_configured(
self,
*,
model: Any,
seed: int,
steps: int,
cfg: float,
sampler_name: str,
scheduler: str,
positive: object,
negative: object,
latent_image: dict[str, Any],
denoise: float,
options: SamplerOptions | None = None,
segs: object | None = None,
region_masks: object | None = None,
) -> dict[str, Any]:
"""Admit connected region data only through its enabled sampling capability."""
if options is not None and not isinstance(options, SamplerOptions):
+27
View File
@@ -0,0 +1,27 @@
"""
SesquiLSR - Latent Super-Resolution Upscaler.
Operates on VAE latents at arbitrary scales in [1x, 2x].
"""
from .inference_adaptors import (
LatentFormatAdaptor,
make_flux,
make_flux2,
make_ideogram4,
make_identity,
make_sdxl,
make_wan21,
)
from .upscaler import LatentUpscaler
__all__ = [
"LatentUpscaler",
"LatentFormatAdaptor",
"make_identity",
"make_sdxl",
"make_flux",
"make_flux2",
"make_ideogram4",
"make_wan21",
]
@@ -0,0 +1,355 @@
"""Latent-format adaptors for Sesqui inference.
Each adaptor converts between a pipeline's external latent space (what the
diffusion model passes around) and the canonical VAE latent space that
Sesqui models are trained on (as defined by the VAE adaptor).
Adaptors are only needed when the pipeline modifies the VAE output.
"""
from __future__ import annotations
import torch
import torch.nn.functional as F
from torch import Tensor
# ── Patchify helpers (shared by Flux2-family formats) ─────────────────────
def _patchify(z: Tensor) -> Tensor:
"""32ch -> 128ch via 2x2 spatial packing."""
B, C, H, W = z.shape
return (
z.float()
.reshape(B, C, H // 2, 2, W // 2, 2)
.permute(0, 1, 3, 5, 2, 4)
.reshape(B, C * 4, H // 2, W // 2)
)
def _unpatchify(z: Tensor) -> Tensor:
"""128ch -> 32ch via 2x2 spatial unpacking."""
B, C, H, W = z.shape
return (
z.float()
.reshape(B, C // 4, 2, 2, H, W)
.permute(0, 1, 4, 2, 5, 3)
.reshape(B, C // 4, H * 2, W * 2)
)
# ── Base class ────────────────────────────────────────────────────────────
class LatentFormatAdaptor:
"""Convert between an external pipeline's latent and VAE latent.
Subclass or instantiate via the ``make_*`` constructors below.
Parameters
----------
external_channels : int
Channel count of the external/pipeline latent.
spatial_scale : int, default 1
Spatial multiplier between external and VAE-latent coords.
Flux2-family formats use ``2`` (VAE latent is 2x spatially larger
because the pipeline packs 2x2 into channels). All others use ``1``.
"""
def __init__(self, external_channels: int, spatial_scale: int = 1):
self.external_channels = external_channels
self.spatial_scale = spatial_scale
def to_vae_latent(self, z: Tensor) -> Tensor:
"""External/pipeline latent -> VAE latent."""
return z.float()
def from_vae_latent(self, z: Tensor) -> Tensor:
"""VAE latent -> external/pipeline latent."""
return z
def vae_target_size(self, target_size: tuple[int, int]) -> tuple[int, int]:
"""Target size in external coords -> VAE-latent coords."""
h, w = target_size
s = self.spatial_scale
return (h * s, w * s) if s != 1 else target_size
# ── Constructors ──────────────────────────────────────────────────────────
def make_identity(external_channels: int) -> LatentFormatAdaptor:
"""Adaptor for formats where the external latent IS the model latent.
Wan 2.1, Anima, Qwen Image, Krea 2, and any other model where ComfyUI's
latent_format does not apply a scaling transform (process_in/out are
identity or cancel out before reaching custom LATENT nodes).
"""
return LatentFormatAdaptor(external_channels=external_channels)
def make_sdxl() -> LatentFormatAdaptor:
"""SDXL - 4ch VAE latent with 0.13025 scaling.
ComfyUI stores raw VAE latents between nodes; KSampler applies
process_in (x0.13025) / process_out (/0.13025) around the model.
Sesqui was trained on the scaled latents.
"""
return _AffineAdaptor(external_channels=4, scale=0.13025, shift=0.0)
def make_flux() -> LatentFormatAdaptor:
"""Flux / Z-Image Turbo / Lumina - 16ch VAE latent with shift/scale.
ComfyUI stores raw VAE latents between nodes; KSampler applies
process_in ((z - 0.1159) x 0.3611) / process_out (z / 0.3611 + 0.1159).
Sesqui was trained on the scaled latents.
"""
return _AffineAdaptor(external_channels=16, scale=0.3611, shift=0.1159)
def make_flux2(
running_mean: Tensor | None = None,
running_var: Tensor | None = None,
eps: float = 1e-4,
) -> LatentFormatAdaptor:
"""Flux2 - 32ch raw VAE ↔ 128ch patched + BN pipeline latent.
Uses BN statistics from the Flux2 VAE's ``bn`` module.
Defaults are extracted from the official Flux2 VAE checkpoint.
"""
if running_mean is None:
running_mean = torch.tensor([
-0.06762, -0.07152, -0.07534, -0.07449, 0.02228, 0.01800, 0.01420, 0.01836,
-0.00006, -0.00625, -0.00021, -0.00314, -0.02720, -0.02811, -0.02765, -0.02903,
-0.07689, -0.06717, -0.09019, -0.08921, 0.01684, 0.01521, 0.00790, 0.00858,
0.00835, 0.00154, 0.00026, -0.00428, -0.04388, -0.04190, -0.04378, -0.04315,
-0.01025, -0.01319, -0.00662, -0.00477, -0.03106, -0.03055, -0.02790, -0.01795,
0.00302, 0.00150, 0.01259, 0.01447, 0.03472, 0.03377, 0.03366, 0.02830,
0.00198, 0.00473, 0.00465, 0.00496, 0.01227, 0.00810, 0.00806, 0.01458,
0.06811, 0.06790, 0.07665, 0.07319, -0.04621, -0.04739, -0.03919, -0.05109,
-0.05278, -0.04774, -0.04700, -0.05172, -0.03171, -0.03163, -0.03447, -0.02826,
0.05097, 0.04450, 0.05781, 0.04580, -0.04116, -0.04583, -0.04874, -0.04674,
-0.00884, -0.01063, -0.00881, -0.00461, -0.03758, -0.04322, -0.04357, -0.04989,
0.01185, 0.01664, 0.02028, 0.02790, 0.01127, 0.01290, 0.00156, 0.00716,
-0.01180, -0.00184, -0.01414, -0.00537, -0.00910, -0.01380, -0.01447, -0.01870,
0.03225, 0.03050, 0.02587, 0.02996, 0.05400, 0.06144, 0.04954, 0.05899,
-0.05108, -0.06033, -0.04778, -0.05240, -0.02268, -0.02742, -0.01537, -0.02546,
-0.05721, -0.05648, -0.05176, -0.04956, 0.01159, 0.00542, 0.01630, 0.01038,
])
if running_var is None:
running_var = torch.tensor([
3.25021, 3.16341, 3.19243, 3.18137, 3.13891, 3.09414, 3.10118, 3.05509,
3.00518, 3.01795, 3.00676, 3.00764, 3.46902, 3.43252, 3.47023, 3.45539,
3.09498, 3.07138, 3.08192, 3.09134, 3.01471, 3.02746, 3.01199, 3.02529,
3.00746, 2.97413, 3.02488, 2.99405, 3.08042, 3.06691, 3.08315, 3.05815,
3.40362, 3.40553, 3.44087, 3.43550, 3.32671, 3.17300, 3.18745, 3.22017,
3.25698, 3.19532, 3.13096, 3.12421, 3.16203, 3.12096, 3.21296, 3.18538,
3.09027, 3.03003, 3.05658, 3.01625, 3.22585, 3.23913, 3.21108, 3.21309,
3.16103, 3.14950, 3.14238, 3.15017, 3.07164, 3.04400, 3.11775, 3.06079,
3.15937, 3.13995, 3.17299, 3.17302, 3.29846, 3.24451, 3.24831, 3.25173,
3.07203, 3.00360, 3.08447, 3.05619, 3.10095, 3.06496, 3.12614, 3.10201,
3.12051, 3.07826, 3.17810, 3.14189, 3.20242, 3.23967, 3.19097, 3.15400,
3.10219, 3.10638, 3.08341, 3.08930, 3.16211, 3.12266, 3.17199, 3.16812,
2.95874, 2.91292, 2.98084, 2.92094, 3.16569, 3.08972, 3.06321, 3.04655,
3.09284, 3.06227, 3.07098, 3.01419, 3.10315, 3.08778, 3.04287, 3.03801,
3.06550, 3.10085, 3.10954, 3.10174, 2.97687, 2.93585, 2.99999, 2.96735,
3.12007, 3.10587, 3.13934, 3.12008, 3.04748, 3.04194, 3.08653, 3.07292,
])
return _Flux2BNAdaptor(running_mean=running_mean, running_var=running_var, eps=eps)
def make_ideogram4(
shift: Tensor | None = None,
scale: Tensor | None = None,
) -> LatentFormatAdaptor:
"""Ideogram 4 - 32ch raw VAE ↔ 128ch patched + shift/scale pipeline latent.
Uses the same Flux2 VAE with 2x2 patchify, but applies per-channel
shift+scale normalization instead of BatchNorm.
"""
if shift is None:
shift = torch.tensor([
0.01984364, 0.10149707, 0.29689495, 0.27188619, -0.21445648, -0.15979549,
0.05021099, -0.15083604, -0.15360136, -0.20131799, 0.01922352, 0.06226260,
0.10140969, -0.06739428, 0.37582610, -0.23371200, 0.35164491, -0.02590912,
-0.02719350, -0.10833897, -0.14768480, -0.01130957, -0.22983720, 0.23526423,
-0.10893522, 0.11957631, 0.04047799, 0.31345890, -0.17225064, -0.18646109,
-0.34691978, -0.03571246, 0.02583857, 0.10190072, 0.28402294, 0.26952152,
-0.21634675, -0.17938656, 0.04358909, -0.15007621, -0.15485020, -0.18971131,
0.02710861, 0.05609494, 0.10697846, -0.06854968, 0.38167698, -0.24269937,
0.35705471, -0.03063305, -0.02946109, -0.11244286, -0.14336038, -0.01362137,
-0.21863696, 0.23228983, -0.11739769, 0.11693044, 0.02563311, 0.31356594,
-0.17420591, -0.19006285, -0.34905377, -0.04025005, 0.01924137, 0.07652984,
0.29956080, 0.26280570, -0.22011674, -0.12715361, 0.04879879, -0.14075719,
-0.15935895, -0.21235840, 0.01974813, 0.05523547, 0.10011992, -0.06428964,
0.37781868, -0.21491644, 0.34254215, -0.03153528, -0.03100820, -0.10761415,
-0.14730405, -0.02475182, -0.22855880, 0.25150810, -0.10445128, 0.12446000,
0.07062869, 0.30880162, -0.18016875, -0.18869164, -0.34533499, -0.01291770,
0.02578168, 0.07993659, 0.28642181, 0.26038408, -0.22459419, -0.14820155,
0.04059549, -0.14043529, -0.16111187, -0.20203050, 0.02602069, 0.04852717,
0.10432153, -0.06309942, 0.38402443, -0.22397003, 0.34814481, -0.03774432,
-0.03381438, -0.11245691, -0.14128767, -0.02853208, -0.21752016, 0.24872463,
-0.11399775, 0.12226870, 0.05620835, 0.30917800, -0.18065738, -0.19401479,
-0.34495114, -0.01760592,
])
if scale is None:
scale = torch.tensor([
1.63933691, 1.70204478, 1.73642566, 1.90004803, 1.66753160, 1.69059584,
1.56853198, 1.62314944, 1.89106626, 1.58086668, 1.60822129, 1.60962993,
1.63322129, 1.56074359, 1.73419528, 1.79192650, 1.64040632, 1.66802808,
1.60390303, 1.75480492, 1.63187587, 1.64334594, 1.61722884, 1.60146046,
1.63459219, 1.55291476, 1.68771497, 1.68415657, 1.78966054, 1.66631641,
1.65626686, 1.65976433, 1.63487607, 1.69513249, 1.72933756, 1.91310663,
1.67035057, 1.72286863, 1.56719251, 1.61934825, 1.88628859, 1.56911539,
1.59455129, 1.60829869, 1.62470611, 1.56052853, 1.73677003, 1.77563606,
1.63732541, 1.66370527, 1.59508952, 1.75153949, 1.63029275, 1.64517667,
1.61659342, 1.59722044, 1.64103121, 1.54085310, 1.68610394, 1.67772755,
1.78998563, 1.66621713, 1.65458955, 1.66041308, 1.64710857, 1.68163503,
1.74000294, 1.92784786, 1.67411194, 1.67395548, 1.57406532, 1.62199356,
1.87618195, 1.55843750, 1.57438785, 1.61711053, 1.63094305, 1.55644029,
1.73124302, 1.80666627, 1.64636210, 1.65932006, 1.60816188, 1.75682671,
1.64695873, 1.63121722, 1.61380832, 1.60478651, 1.63396035, 1.53505068,
1.65534289, 1.67132281, 1.80317197, 1.67673140, 1.65700938, 1.68426259,
1.65339716, 1.67540638, 1.73298504, 1.94067348, 1.67893609, 1.70635117,
1.57309060, 1.61928553, 1.87148809, 1.56244866, 1.56697152, 1.61584394,
1.62759496, 1.55480378, 1.73484107, 1.79055143, 1.64688773, 1.66121492,
1.60135887, 1.75254572, 1.64798332, 1.62989921, 1.61381592, 1.60792883,
1.63939668, 1.53075757, 1.65371318, 1.66801185, 1.80029087, 1.67591476,
1.65655173, 1.68533454,
])
return _ShiftScalePatchAdaptor(shift=shift, scale=scale)
def make_wan21(
latents_mean: Tensor | None = None,
latents_std: Tensor | None = None,
) -> LatentFormatAdaptor:
"""Wan21 / Anima / QwenImage - per-channel z-score.
Pipeline latent = ``(VAE_raw - mean) / std``.
Default stats from the Wan21 latent format.
"""
if latents_mean is None:
latents_mean = torch.tensor([
-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921,
])
if latents_std is None:
latents_std = torch.tensor([
2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160,
])
return _ZScoreAdaptor(
external_channels=16,
mean=latents_mean,
std=latents_std,
)
# ── Internal implementations ──────────────────────────────────────────────
class _AffineAdaptor(LatentFormatAdaptor):
"""Affine transform: (z - shift) * scale and inverse.
Used by SDXL and Flux where ComfyUI stores raw VAE latents between nodes
but the diffusion model (and Sesqui) operates on scaled latents.
"""
def __init__(self, external_channels: int, scale: float, shift: float):
super().__init__(external_channels=external_channels)
self.scale = scale
self.shift = shift
def to_vae_latent(self, z: Tensor) -> Tensor:
return (z.float() - self.shift) * self.scale
def from_vae_latent(self, z: Tensor) -> Tensor:
return z.float() / self.scale + self.shift
class _ZScoreAdaptor(LatentFormatAdaptor):
"""Per-channel z-score normalization (Wan21 / Anima / QwenImage)."""
def __init__(self, external_channels: int, mean: Tensor, std: Tensor):
super().__init__(external_channels=external_channels)
self.mean = mean.clone().view(1, -1, 1, 1)
self.std = std.clone().view(1, -1, 1, 1)
def _cast(self, z: Tensor) -> tuple[Tensor, Tensor]:
d = z.device
return (
self.mean.to(dtype=z.dtype, device=d),
self.std.to(dtype=z.dtype, device=d),
)
def to_vae_latent(self, z: Tensor) -> Tensor:
# Pipeline latent = (VAE_raw - mean) / std
# -> VAE_raw = z * std + mean
m, s = self._cast(z)
return z.float() * s + m
def from_vae_latent(self, z: Tensor) -> Tensor:
m, s = self._cast(z)
return (z.float() - m) / s
class _Flux2BNAdaptor(LatentFormatAdaptor):
"""Flux2: 32ch raw VAE ↔ 128ch patched + BN-normalized pipeline."""
def __init__(self, running_mean: Tensor, running_var: Tensor, eps: float = 1e-4):
super().__init__(external_channels=128, spatial_scale=2)
self.bn_mean = running_mean.clone().float().view(1, -1, 1, 1)
self.bn_var = running_var.clone().float().view(1, -1, 1, 1)
self.bn_eps = eps
def to_vae_latent(self, z: Tensor) -> Tensor:
# Pipeline = BN(patchify(VAE_raw))
# Undo BN: x * sqrt(var) + mean
var = self.bn_var.to(dtype=z.dtype, device=z.device)
mean = self.bn_mean.to(dtype=z.dtype, device=z.device)
z = z.float() * torch.sqrt(var + self.bn_eps) + mean
return _unpatchify(z)
def from_vae_latent(self, z: Tensor) -> Tensor:
z = _patchify(z)
var = self.bn_var.to(dtype=z.dtype, device=z.device)
mean = self.bn_mean.to(dtype=z.dtype, device=z.device)
return F.batch_norm(
z,
running_mean=mean.squeeze(),
running_var=var.squeeze(),
momentum=0.0,
eps=self.bn_eps,
)
class _ShiftScalePatchAdaptor(LatentFormatAdaptor):
"""Ideogram 4: 32ch raw VAE <-> 128ch patched + per-channel shift/scale.
Pipeline latent = ``patchify(VAE_raw) * scale + shift``.
"""
def __init__(self, shift: Tensor, scale: Tensor):
super().__init__(external_channels=128, spatial_scale=2)
self.shift = shift.clone().float().view(1, -1, 1, 1)
self.scale = scale.clone().float().view(1, -1, 1, 1)
def _cast(self, z: Tensor) -> tuple[Tensor, Tensor]:
d = z.device
return (
self.shift.to(dtype=z.dtype, device=d),
self.scale.to(dtype=z.dtype, device=d),
)
def to_vae_latent(self, z: Tensor) -> Tensor:
# Pipeline = patchify(VAE_raw) * scale + shift
# -> VAE_raw = unpatchify((z - shift) / scale)
s, sc = self._cast(z)
return _unpatchify((z.float() - s) / sc)
def from_vae_latent(self, z: Tensor) -> Tensor:
s, sc = self._cast(z)
return _patchify(z) * sc + s
+290
View File
@@ -0,0 +1,290 @@
"""
SesquiLSR — Multi-scale latent upscaler.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
# I like Mish, but any decent activation (like SiLU) works fine.
Activation = nn.Mish
# IR LayerNorm from https://arxiv.org/abs/2504.06629
# Removed bias to reduce chances of mean drift, also it didn't really make
# much of a difference with it in.
class IRLayerNorm(nn.Module):
def __init__(self, channels: int):
super().__init__()
self.weight = nn.Parameter(torch.ones(channels))
def forward(self, x: Tensor) -> tuple[Tensor, Tensor]:
var, mean = torch.var_mean(x, dim=(1, 2, 3), unbiased=False, keepdim=True)
std = torch.sqrt(var + 1e-6)
y = (x - mean) / std * self.weight.view(1, -1, 1, 1)
return y, std
# Inspired by WDSR / wide activation: https://arxiv.org/abs/1808.08718
# Uses a lopsided gate design to maximise values, 8:1 ratio works well
# in practise.
class GatedWDSR(nn.Module):
def __init__(
self,
channels: int,
expand_ratio: float = 2.0,
shrink_ratio: float = 0.8,
values_per_gate: int = 8,
):
super().__init__()
if values_per_gate < 1:
raise ValueError(f"values_per_gate must be >= 1, got {values_per_gate}")
gate_ch = math.ceil(expand_ratio * channels / (values_per_gate + 1))
gate_ch = math.ceil(gate_ch / values_per_gate) * values_per_gate
value_ch = gate_ch * values_per_gate
mid_ch = math.ceil(channels * shrink_ratio)
hidden = value_ch + gate_ch
self.value_ch = value_ch
self.gate_ch = gate_ch
self.values_per_gate = values_per_gate
self.norm = IRLayerNorm(channels)
self.up = nn.Conv2d(channels, hidden, 1, bias=False)
self.down = nn.Conv2d(value_ch, mid_ch, 1, bias=False)
self.mix = nn.Conv2d(
mid_ch, channels, kernel_size=3, padding=1,
padding_mode="reflect", bias=False,
)
self.act = Activation()
def forward(self, x: Tensor) -> Tensor:
y, s = self.norm(x)
v, g = torch.split(self.up(y), [self.value_ch, self.gate_ch], dim=1)
v = self.act(v)
g = torch.tanh(0.5 * g) # Bound residual and slow down saturation.
B, _, H, W = v.shape
v = v.view(B, self.gate_ch, self.values_per_gate, H, W)
y = (v * g[:, :, None, :, :]).reshape(B, self.value_ch, H, W)
y = self.down(y)
y = self.mix(y)
return x + y * s
# Originally the model uses AdaLN in every block for scale conditioning, but this
# ended up exploding. Turns out we can just do conditioning in the head, and the
# model works fine.
class PeriodicEncoder1D(nn.Module):
def __init__(self, n_freqs: int = 4, max_freq: float = 4.0):
super().__init__()
self.out_dim = 2 * n_freqs
self.register_buffer("freqs", 2.0 ** torch.linspace(0, math.log2(max_freq), n_freqs))
def forward(self, x: Tensor) -> Tensor:
angles = x.unsqueeze(-1) * self.freqs * math.pi
return torch.cat([angles.sin(), angles.cos()], dim=-1)
# The "sesqui" part. Inspired by rational/fractional resampling, we use PS to get to 2x,
# then use a learned downsample to reach our target. The original model was a fixed 1.5x;
# sesqui means "one and a half". While the fixed upscale is gone, the name remains.
class LowRankReassemblyHead(nn.Module):
def __init__(self, channels: int, refine_expand: float = 1.5, rank: int = 2):
super().__init__()
self.channels = channels
self.rank = rank
self.K = 4
self.register_buffer(
"tap_offsets", torch.tensor([-1, 0, 1, 2], dtype=torch.long),
persistent=False,
)
self.ps_expand = nn.Conv2d(channels, channels * 4, 1, bias=True)
self._init_pixel_shuffle()
lift_ch = min(64, channels // 2)
self.lift_expand = nn.Conv2d(channels, lift_ch * 4, 1, bias=False)
self.lift_mix = nn.Conv2d(lift_ch, channels, 3, padding=1, padding_mode="reflect", bias=False)
nn.init.trunc_normal_(self.lift_expand.weight, std=0.02)
nn.init.trunc_normal_(self.lift_mix.weight, std=0.02)
# You could probably use a simple conv stack here, but another resblock works well so why not.
self.refine = GatedWDSR(channels, expand_ratio=refine_expand, shrink_ratio=1.0, values_per_gate=8)
taps = torch.zeros(rank, 2, channels, self.K)
taps[0, :, :, :] = torch.tensor([-0.0625, 0.5625, 0.5625, -0.0625], dtype=taps.dtype) # ~Bicubic start
self.base_taps = nn.Parameter(taps)
self.pe = PeriodicEncoder1D(n_freqs=4)
self.axis_emb = nn.Embedding(2, 8)
self.rank_emb = nn.Embedding(rank, 8)
cond_dim = self.pe.out_dim + 1 + 8 + 8
self.filter_net = nn.Sequential(
nn.Linear(cond_dim, 64, bias=True),
Activation(),
nn.Linear(64, channels * self.K, bias=True),
)
nn.init.trunc_normal_(self.filter_net[-1].weight, std=0.02)
nn.init.zeros_(self.filter_net[-1].bias)
self.rank_gain = nn.Parameter(torch.tensor([2.0] + [0.0] * (rank - 1), dtype=torch.float32))
def _init_pixel_shuffle(self) -> None:
with torch.no_grad():
self.ps_expand.weight.zero_()
self.ps_expand.bias.zero_()
w = self.ps_expand.weight.view(self.channels, 4, self.channels)
idx = torch.arange(self.channels, device=w.device)
w[idx, :, idx] = 1.0
@staticmethod
def _axis_coords(n_src: int, n_tgt: int, device: torch.device) -> tuple[Tensor, Tensor]:
pos = (torch.arange(n_tgt, device=device, dtype=torch.float32) + 0.5)
pos = pos * (n_src / n_tgt) - 0.5
base = torch.floor(pos).long().clamp_(0, n_src - 1)
frac = pos - base.float()
return frac, base
def _predict_weights(self, fy: Tensor, fx: Tensor,
log_foot_y: float, log_foot_x: float,
batch: int) -> tuple[Tensor, Tensor]:
ny, nx = fy.numel(), fx.numel()
R = self.rank
n_axis = ny + nx
frac_axis = torch.cat([fy, fx])
log_foot = fy.new_full((n_axis, 1), log_foot_x)
log_foot[:ny] = log_foot_y
axis_id = torch.empty(n_axis, dtype=torch.long, device=fy.device)
axis_id[:ny] = 0
axis_id[ny:] = 1
cond_axis = torch.cat([
self.pe(frac_axis), log_foot, self.axis_emb(axis_id),
], dim=-1)
rank_id = torch.arange(R, device=fy.device).repeat_interleave(n_axis)
cond = torch.cat([cond_axis.repeat(R, 1), self.rank_emb(rank_id)], dim=-1)
cond = cond.to(dtype=self.base_taps.dtype)
residuals = self.filter_net(cond).view(R, n_axis, self.channels, self.K)
residuals = 0.2 * torch.tanh(residuals) # Again, bound the residual otherwise kaboom...
residuals = residuals - residuals.mean(dim=-1, keepdim=True)
wy = residuals[:, :ny] + self.base_taps[:, 0].unsqueeze(1)
wx = residuals[:, ny:] + self.base_taps[:, 1].unsqueeze(1)
ky = wy.permute(0, 2, 1, 3).reshape(R, self.channels * ny, self.K).unsqueeze(1).expand(-1, batch, -1, -1)
kx = wx.permute(0, 2, 1, 3).reshape(R, self.channels * nx, self.K).unsqueeze(1).expand(-1, batch, -1, -1)
return ky, kx
def _filter_axis(self, x: Tensor, centers: Tensor, weights: Tensor, dim: int) -> Tensor:
B, C = x.shape[0], x.shape[1]
n_out = centers.numel()
limit = x.shape[dim] - 1
idx = (centers.unsqueeze(0) + self.tap_offsets.unsqueeze(1)).clamp(0, limit)
w = weights.view(B, C, n_out, self.K)
if dim == 2:
idx_exp = idx.view(1, 1, self.K * n_out, 1).expand(B, C, -1, x.shape[3])
sample = torch.gather(x, 2, idx_exp).view(B, C, self.K, n_out, x.shape[3])
return (sample * w.permute(0, 1, 3, 2).unsqueeze(-1)).sum(dim=2)
idx_exp = idx.view(1, 1, 1, self.K * n_out).expand(B, C, x.shape[2], -1)
sample = torch.gather(x, 3, idx_exp).view(B, C, x.shape[2], self.K, n_out)
return (sample * w.permute(0, 1, 3, 2).unsqueeze(2)).sum(dim=3)
def forward(self, x: Tensor, target_size: tuple[int, int]) -> Tensor:
h_tgt, w_tgt = target_size
batch = x.shape[0]
x_base = F.pixel_shuffle(self.ps_expand(x), 2)
x_lift = self.lift_mix(F.pixel_shuffle(self.lift_expand(x), 2))
x = x_base + x_lift
h_src, w_src = x.shape[-2:]
x = self.refine(x)
fy, iy = self._axis_coords(h_src, h_tgt, x.device)
fx, ix = self._axis_coords(w_src, w_tgt, x.device)
ky, kx = self._predict_weights(
fy, fx,
math.log2(max(h_src / h_tgt, 0.01)),
math.log2(max(w_src / w_tgt, 0.01)),
batch,
)
gains = self.rank_gain.sigmoid()
out = torch.zeros(batch, self.channels, h_tgt, w_tgt, device=x.device, dtype=x.dtype)
for r in range(self.rank):
x_r = self._filter_axis(x, ix, kx[r], dim=3)
x_r = self._filter_axis(x_r, iy, ky[r], dim=2)
out = out + gains[r] * x_r
return out
# Some models/VAEs can get away with very little in terms of trunks. For example, the Qwen/Wan21 VAE
# was completely happy with 2/4 instead of 8/6. However, to keep the arch consistent, we use the same
# structure for all upscalers.
class LatentUpscaler(nn.Module):
def __init__(
self,
in_depth: int = 8,
out_depth: int = 6,
in_channels: int = 4,
width: int = 128,
in_expand: float = 1.5,
out_expand: float = 2.0,
):
super().__init__()
self.head_conv = nn.Sequential(
nn.Conv2d(in_channels, width // 2, 3, padding=1, padding_mode="reflect", bias=False),
Activation(),
nn.Conv2d(width // 2, width, 3, padding=1, padding_mode="reflect", bias=True),
)
# No difference with a slight bottleneck here and smaller expand, saves a few parameters.
self.in_blocks = nn.Sequential(*[
GatedWDSR(width, expand_ratio=in_expand, shrink_ratio=0.75, values_per_gate=8)
for _ in range(in_depth)
])
self.upsample = LowRankReassemblyHead(width, refine_expand=2.0)
# Traditional SR says do most of the work at LR, then upscale and output. The problem is
# latents are delicate structures where no amount of front-loaded work can make up for the
# damage introduced by upscaling. Post-correction is a must, and all variants tested without
# post-correction were much worse.
self.out_blocks = nn.Sequential(*[
GatedWDSR(width, expand_ratio=out_expand, shrink_ratio=1.0, values_per_gate=8)
for _ in range(out_depth)
])
self.tail_conv = nn.Sequential(
nn.Conv2d(width, width // 2, 1, bias=False),
Activation(),
nn.Conv2d(width // 2, in_channels, 3, padding=1, padding_mode="reflect", bias=True),
)
# Give the model a way to do light correction on the skip so the trunk/head can focus on
# detail/damage control.
self.skip_conv = nn.Conv2d(in_channels, in_channels, 5, padding=2, padding_mode="reflect", bias=False)
nn.init.dirac_(self.skip_conv.weight)
def forward(self, x: Tensor, target_size: tuple[int, int]) -> Tensor:
h = self.head_conv(x)
h = self.in_blocks(h)
h = self.upsample(h, target_size)
h = self.out_blocks(h)
h = self.tail_conv(h)
# Magic sauce. Learn correction on crude upsample.
x_up = self.skip_conv(F.interpolate(x, size=target_size, mode="bicubic", align_corners=False))
return h + x_up
@@ -0,0 +1,165 @@
"""Prove receiver planning, strict image provenance and source cache dependencies."""
from copy import deepcopy
from typing import Any
import pytest
from simple_syrup.runtime.noise_reuse_graph import (
rewrite_noise_reuse_graph,
source_lineage,
)
def graph() -> dict[str, Any]:
"""Build shared source branches without invoking model IO."""
return {
"source": {
"class_type": "KSampler",
"inputs": {
"model": ["loader", 0],
"latent_image": ["empty", 0],
"sampler_name": "euler",
},
},
"decode": {
"class_type": "VAEDecode",
"inputs": {"samples": ["source", 0], "vae": ["loader", 2]},
},
"recover": {
"class_type": "SimpleSyrup.LatentFromImage",
"inputs": {"image": ["decode", 0]},
},
"upscale": {
"class_type": "SimpleSyrup.SimpleUpscaleLatent",
"inputs": {"latent": ["recover", 0], "method": "Auto", "scale_factor": 3.0},
},
"reuse": {
"class_type": "SimpleSyrup.NoiseReuseOptions",
"inputs": {"enabled": True, "finishing_steps": 0},
},
"receive": {
"class_type": "SimpleSyrup.KSampler",
"inputs": {
"model": ["loader", 0],
"latent_image": ["upscale", 0],
"steps": 16,
"sampler_name": "euler",
"scheduler": "normal",
"denoise": 0.2,
"options": ["reuse", 0],
},
},
}
def test_image_bridge_resolves_original_latent_and_both_consumers_share_capture() -> (
None
):
"""Coordinate before source execution and preserve direct latent consumers."""
prompt = graph()
prompt["finish"] = {
"class_type": "SimpleSyrup.NoiseReuseOptions",
"inputs": {"finishing_steps": 1},
}
prompt["receive_finish"] = deepcopy(prompt["receive"])
prompt["receive_finish"]["inputs"]["options"] = ["finish", 0]
assert source_lineage(prompt, ["upscale", 0], {}) == ("source", '[{"scale": 3.0}]')
rewrite_noise_reuse_graph(prompt, {})
captures = [
node
for node in prompt.values()
if node["class_type"] == "SimpleSyrup.TrajectoryCaptureModel"
]
requests = [
node
for node in prompt.values()
if node["class_type"] == "SimpleSyrup.NoiseCaptureRequest"
]
assert len(captures) == 1 and len(requests) == 2
assert requests[1]["inputs"]["previous"] == ["receive.noise_capture_request", 0]
assert prompt["source"]["inputs"]["model"] == ["source.trajectory_capture", 0]
assert prompt["decode"]["inputs"]["samples"] == ["source.sampling_provenance", 0]
assert prompt["source.sampling_provenance"]["inputs"]["latent"] == ["source", 0]
assert requests[0]["inputs"]["source_latent"] == ["empty", 0]
assert not any(node["class_type"] == "VAEEncode" for node in prompt.values())
def test_missing_or_edited_provenance_creates_actionable_error() -> None:
"""Reject pixel edits before any receiving sampler can substitute a fallback."""
prompt = graph()
prompt["decode"]["class_type"] = "ImageScale"
with pytest.raises(ValueError, match="image provenance"):
source_lineage(prompt, ["upscale", 0], {})
rewrite_noise_reuse_graph(prompt, {})
assert prompt["receive"]["inputs"]["latent_image"] == [
"receive.noise_reuse_error",
0,
]
assert (
prompt["receive.noise_reuse_error"]["class_type"]
== "SimpleSyrup.NoiseReuseError"
)
def test_disabled_reuse_leaves_native_graph_unchanged() -> None:
"""Avoid upstream capture or new GPU work when the option is switched off."""
prompt = graph()
prompt["upscale"]["inputs"]["method"] = "bilinear"
prompt["reuse"]["inputs"]["enabled"] = False
before = deepcopy(prompt)
rewrite_noise_reuse_graph(prompt, {})
assert prompt == before
def test_auto_only_requests_format_provenance_without_state_capture() -> None:
"""Identify compatible models without perturbing the original source schedule."""
prompt = graph()
prompt["reuse"]["inputs"]["enabled"] = False
rewrite_noise_reuse_graph(prompt, {})
assert "requests" not in prompt["source.trajectory_capture"]["inputs"]
assert not any(
node["class_type"] == "SimpleSyrup.NoiseCaptureRequest"
for node in prompt.values()
)
def test_auto_without_sampler_lineage_keeps_native_fallback_input() -> None:
"""Allow unknown encoded latents while keeping noise reuse source checks strict."""
prompt = graph()
prompt["reuse"]["inputs"]["enabled"] = False
prompt["upscale"]["inputs"]["latent"] = ["encode", 0]
prompt["encode"] = {"class_type": "VAEEncode", "inputs": {}}
before = deepcopy(prompt)
rewrite_noise_reuse_graph(prompt, {})
assert prompt == before
prompt["reuse"]["inputs"]["enabled"] = True
rewrite_noise_reuse_graph(prompt, {})
assert prompt["upscale"]["inputs"]["latent"] == ["encode", 0]
assert prompt["receive"]["inputs"]["latent_image"] == [
"receive.noise_reuse_error",
0,
]
def test_requeued_coordinated_prompt_is_idempotent() -> None:
"""Keep history requeues from installing duplicate keyed source wrappers."""
prompt = graph()
rewrite_noise_reuse_graph(prompt, {})
before = deepcopy(prompt)
rewrite_noise_reuse_graph(prompt, {})
assert prompt == before
def test_receiving_settings_are_live_source_cache_dependencies() -> None:
"""Invalidate source capture when the receiving schedule or options change."""
prompt = graph()
prompt["receive"]["inputs"]["steps"] = ["live_steps", 0]
prompt["receive"]["inputs"]["denoise"] = ["live_denoise", 0]
rewrite_noise_reuse_graph(prompt, {})
request = prompt["receive.noise_capture_request"]["inputs"]
assert request["steps"] == ["live_steps", 0]
assert request["denoise"] == ["live_denoise", 0]
assert request["options"] == ["reuse", 0]
capture = prompt["source.trajectory_capture"]["inputs"]
assert capture["requests"] == ["receive.noise_capture_request", 0]
@@ -23,6 +23,13 @@ BASE_NODE_IDS = [
"SimpleSyrup.NoiseInversionOptions",
"SimpleSyrup.TilingOptions",
"SimpleSyrup.KSampler",
"SimpleSyrup.LatentFromImage",
"SimpleSyrup.SimpleUpscaleLatent",
"SimpleSyrup.NoiseReuseOptions",
"SimpleSyrup.NoiseCaptureRequest",
"SimpleSyrup.TrajectoryCaptureModel",
"SimpleSyrup.AttachSamplingProvenance",
"SimpleSyrup.NoiseReuseError",
"SimpleSyrup.AttentionCaptureModel",
"SimpleSyrup.AttentionMaskedConditioning",
"SimpleSyrup.AttentionRegionMask",
+217
View File
@@ -0,0 +1,217 @@
"""Verify native geometry, capture preservation and learned format selection."""
from importlib import import_module
from pathlib import Path
from typing import Any
import pytest
import torch
from simple_syrup.domain.latent_upscale import (
learned_upscale_stages,
recommended_family,
)
from simple_syrup.domain.noise_reuse import PROVENANCE_KEY, SamplingProvenance
from simple_syrup.runtime import latent_upscale_models
from simple_syrup.runtime.latent_upscale_catalog import latent_model_choices
from simple_syrup.runtime.latent_upscale_models import resolve_latent_upscale_model
from simple_syrup.services.latent_upscale_service import LatentUpscaleService
def test_stages_reach_three_times_without_exceeding_trained_range() -> None:
"""Use 2x then 1.5x and preserve exact rounded endpoint dimensions."""
assert learned_upscale_stages((96, 168), (288, 504)) == ((192, 336), (288, 504))
assert learned_upscale_stages((192, 256), (288, 384)) == ((288, 384),)
assert learned_upscale_stages((4, 4), (4, 4)) == ()
with pytest.raises(ValueError, match="enlargement"):
learned_upscale_stages((8, 8), (4, 4))
def test_native_comfy_resize_preserves_provenance_and_resizes_mask() -> None:
"""Exercise the real Comfy native node and preserve capture ownership."""
samples = torch.arange(96).float().reshape(1, 4, 4, 6)
provenance = SamplingProvenance("SDXL", samples)
latent: dict[str, Any] = {
"samples": samples,
PROVENANCE_KEY: provenance,
"noise_mask": torch.ones(1, 4, 6),
}
output = LatentUpscaleService().upscale(latent, "nearest-exact", 1.5)
assert output["samples"].shape == (1, 4, 6, 9)
assert output["noise_mask"].shape == (1, 6, 9)
assert output[PROVENANCE_KEY] is provenance
assert latent["samples"] is samples and latent["noise_mask"].shape == (1, 4, 6)
@pytest.mark.parametrize(
"latent_format,family",
(
("SDXL", "sdxl"),
("Flux", "flux"),
("Flux2", "flux2"),
("Wan21", "wan"),
("Cosmos1CV8x8x8", "wan"),
),
)
def test_auto_uses_verified_format_instead_of_channel_guess(
latent_format: str, family: str
) -> None:
"""Distinguish Flux and Wan even though both can have 16 channels."""
assert recommended_family(latent_format) == family
def test_auto_and_wrong_family_fail_before_download() -> None:
"""Reject ambiguous source identity and incompatible explicit checkpoints."""
with pytest.raises(ValueError, match="provenance"):
resolve_latent_upscale_model({"samples": torch.zeros(1, 16, 4, 4)}, "Auto")
with pytest.raises(ValueError, match="does not match"):
resolve_latent_upscale_model(
{PROVENANCE_KEY: SamplingProvenance("Flux", torch.zeros(1, 16, 4, 4))},
"wan/upscaler_Wan21.safetensors",
)
with pytest.raises(ValueError, match="compatible"):
recommended_family("SD15")
def test_method_choices_deduplicate_registered_windows_paths(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Keep one choice per installed checkpoint while retaining native algorithms."""
folder_paths = import_module("folder_paths")
monkeypatch.setattr(
folder_paths,
"get_filename_list",
lambda category: [
r"wan\upscaler_Wan21.safetensors",
r"custom\example.safetensors",
],
)
choices = latent_model_choices()
assert choices[0] == "Auto"
assert choices.count("wan/upscaler_Wan21.safetensors") == 1
assert "custom/example.safetensors" in choices and "nearest-exact" in choices
def test_invalid_samples_and_masks_fail_before_model_resolution(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Avoid downloads or inference when tensor inputs are invalid."""
from simple_syrup.services import latent_upscale_service
def forbidden(*args: Any) -> Any:
"""Detect side effects before input validation completes."""
raise AssertionError("Model resolution must not occur.")
monkeypatch.setattr(
latent_upscale_service, "resolve_latent_upscale_model", forbidden
)
service = LatentUpscaleService()
with pytest.raises(ValueError, match="finite"):
service.upscale({"samples": torch.full((1, 4, 8, 12), float("nan"))}, "Auto", 2)
with pytest.raises(TypeError, match="mask"):
service.upscale(
{"samples": torch.zeros(1, 4, 8, 12), "noise_mask": "invalid"}, "Auto", 2
)
def test_registered_original_filename_retains_family_identity(
monkeypatch: pytest.MonkeyPatch,
tmp_path: Path,
) -> None:
"""Distinguish same-channel Flux and Wan weights outside the default folders."""
folders = import_module("folder_paths")
utils = import_module("comfy.utils")
monkeypatch.setattr(
folders,
"get_full_path_or_raise",
lambda category, name: str(
tmp_path / "custom_models" / "upscaler_Flux.safetensors"
),
)
monkeypatch.setattr(
utils,
"load_torch_file",
lambda path, **kwargs: {
"head_conv.0.weight": torch.zeros(8, 16, 3, 3),
"upsample.base_taps": torch.ones(1),
},
)
latent = {PROVENANCE_KEY: SamplingProvenance("Wan21", torch.zeros(1, 16, 4, 4))}
with pytest.raises(ValueError, match="does not match"):
latent_upscale_models.resolve_latent_upscale_model(
latent, "custom/upscaler_Flux.safetensors"
)
@pytest.mark.parametrize("latent_format", (None, "SD15", "UnknownLatentFormat"))
@pytest.mark.parametrize("shape", ((2, 16, 4, 6), (1, 16, 2, 4, 6)))
def test_auto_unknown_latents_use_native_bislerp_without_model_io(
monkeypatch: pytest.MonkeyPatch,
latent_format: str | None,
shape: tuple[int, ...],
) -> None:
"""Match native bislerp exactly for unrecognized image and video latents."""
from simple_syrup.services import latent_upscale_service
def forbidden(*args: Any) -> Any:
"""Reject checkpoint lookup or downloads on the native fallback path."""
raise AssertionError("Unknown latents must use native bislerp.")
monkeypatch.setattr(
latent_upscale_service, "resolve_latent_upscale_model", forbidden
)
samples = torch.randn(shape, generator=torch.Generator().manual_seed(53))
latent: dict[str, Any] = {
"samples": samples,
"noise_mask": torch.ones(1, 4, 6),
"batch_index": [0],
}
if latent_format is not None:
latent[PROVENANCE_KEY] = SamplingProvenance(latent_format, samples)
expected = (
import_module("nodes").LatentUpscaleBy().upscale(latent, "bislerp", 1.5)[0]
)
output = LatentUpscaleService().upscale(latent, "Auto", 1.5)
torch.testing.assert_close(output["samples"], expected["samples"], rtol=0, atol=0)
assert output["noise_mask"].shape == (1, 6, 9)
assert output["batch_index"] is latent["batch_index"]
assert output.get(PROVENANCE_KEY) is latent.get(PROVENANCE_KEY)
assert latent["samples"] is samples
@pytest.mark.parametrize(
"latent_format", ("SDXL", "Flux", "Flux2", "Wan21", "Cosmos1CV8x8x8")
)
def test_auto_recognized_latents_keep_learned_model_selection(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path, latent_format: str
) -> None:
"""Keep existing learned Auto selection for every supported format."""
from simple_syrup.services import latent_upscale_service
samples = torch.ones(1, 4, 4, 6)
latent = {
"samples": samples,
PROVENANCE_KEY: SamplingProvenance(latent_format, samples),
}
family = recommended_family(latent_format)
checkpoint = tmp_path / "upscaler.safetensors"
def resolve(value: dict[str, Any], method: str) -> tuple[str, Path]:
"""Verify known formats reach the authoritative learned resolver."""
assert value is latent and method == "Auto"
return family, checkpoint
def infer(
value: torch.Tensor, *, path: Path, family: str, target: tuple[int, int]
) -> torch.Tensor:
"""Verify the learned inference boundary receives the chosen checkpoint."""
assert value is samples and path == checkpoint
assert family == recommended_family(latent_format) and target == (6, 9)
return torch.ones(1, 4, *target)
monkeypatch.setattr(latent_upscale_service, "resolve_latent_upscale_model", resolve)
monkeypatch.setattr(latent_upscale_service, "upscale_with_sesqui", infer)
output = LatentUpscaleService().upscale(latent, "Auto", 1.5)
assert output["samples"].shape == (1, 4, 6, 9)
assert output[PROVENANCE_KEY] is latent[PROVENANCE_KEY]
+55
View File
@@ -0,0 +1,55 @@
"""Exercise actual learned inference in Comfy's external latent units and layouts."""
from importlib import import_module
from pathlib import Path
from typing import Any
import pytest
import torch
from safetensors.torch import save_file
from simple_syrup.runtime.sesqui_latent_upscale import upscale_with_sesqui
@pytest.mark.parametrize(
"family,channels,video",
(("sdxl", 4, False), ("flux", 16, False), ("wan", 16, True), ("flux2", 128, False)),
)
def test_learned_adapter_matches_direct_raw_vae_inference(
tmp_path: Path, family: str, channels: int, video: bool
) -> None:
"""Detect extra latent normalization, packed Flux2 mistakes and frame transposes."""
vendor: Any = import_module("simple_syrup.third_party.sesqui_lsr")
raw_channels = 32 if channels == 128 else channels
with torch.random.fork_rng():
torch.manual_seed(42)
network = vendor.LatentUpscaler(in_channels=raw_channels).eval()
samples = torch.randn(
(2, channels, 3, 8, 12) if video else (2, channels, 8, 12)
)
path = tmp_path / "learned.safetensors"
save_file(network.state_dict(), str(path))
flat = (
samples.permute(0, 2, 1, 3, 4).reshape(6, channels, 8, 12) if video else samples
)
adaptor = vendor.make_flux2() if channels == 128 else vendor.make_identity(channels)
with torch.inference_mode():
expected = adaptor.from_vae_latent(
network(adaptor.to_vae_latent(flat), adaptor.vae_target_size((12, 18)))
)
if video:
expected = expected.reshape(2, 3, channels, 12, 18).permute(0, 2, 1, 3, 4)
output = upscale_with_sesqui(samples, path=path, family=family, target=(12, 18))
assert output.shape == expected.shape and output.dtype == samples.dtype
assert torch.equal(output, expected)
def test_invalid_channels_fail_before_checkpoint_loading(tmp_path: Path) -> None:
"""Reject incompatible learned format layouts before model IO occurs."""
with pytest.raises(ValueError, match="16 channels"):
upscale_with_sesqui(
torch.zeros(1, 4, 8, 12),
path=tmp_path / "missing.safetensors",
family="wan",
target=(12, 18),
)
@@ -32,6 +32,19 @@ FORBIDDEN_PATCHER_CALLS = frozenset(
FORBIDDEN_PATCHER_WRITES = frozenset({"forced_hooks", "use_clip_schedule"})
APPROVED_VALUE_CLONES = Counter(
{
("simple_syrup/domain/spectral_noise.py", "spatial_dct(fresh)"): 1,
(
"simple_syrup/runtime/trajectory_capture.py",
"record['x'].detach().float().cpu()",
): 1,
(
"simple_syrup/runtime/trajectory_capture.py",
"result.detach().float().cpu()",
): 1,
(
"simple_syrup/runtime/trajectory_capture.py",
"schedule.detach().float().cpu()",
): 1,
("simple_syrup/domain/inversion_solver.py", "source"): 1,
("simple_syrup/runtime/noise_inversion.py", "zero"): 1,
("simple_syrup/runtime/noise_inversion.py", "noise"): 1,
+74
View File
@@ -0,0 +1,74 @@
"""Verify the public latent reuse node contracts and v3 registration."""
from typing import Any
import pytest
from simple_syrup.domain.sampler_options import (
ContextualDiffusionOptions,
SamplerOptions,
)
from simple_syrup.nodes_v3 import get_nodes
from simple_syrup.nodes_v3.latent_from_image import LatentFromImageV3
from simple_syrup.nodes_v3.noise_reuse_options import NoiseReuseOptionsV3
from simple_syrup.nodes_v3.simple_upscale_latent import SimpleUpscaleLatentV3
from simple_syrup.services.latent_provenance_service import LatentProvenanceService
@pytest.mark.parametrize(
"node", (LatentFromImageV3, NoiseReuseOptionsV3, SimpleUpscaleLatentV3)
)
def test_v3_registration_and_tooltips_cover_public_contract(node: Any) -> None:
"""Export all three nodes with descriptions and actionable socket tooltips."""
assert node in get_nodes()
schema = node.define_schema()
assert schema.node_id.startswith("SimpleSyrup.")
assert schema.description
assert all(value.tooltip for value in (*schema.inputs, *schema.outputs))
def test_upscale_exposes_one_dropdown_and_scale_default() -> None:
"""Keep format and model loading choices behind the agreed method control."""
schema = SimpleUpscaleLatentV3.define_schema()
inputs = {value.id: value for value in schema.inputs}
assert list(inputs) == ["latent", "method", "scale_factor"]
assert inputs["method"].default == "Auto"
assert "bislerp" in inputs["method"].tooltip.lower()
assert inputs["scale_factor"].default == 2.0
assert schema.display_name == "Simple Upscale Latent"
def test_reuse_enabled_is_first_widget_and_disabled_passes_same_prefix() -> None:
"""Prove zero finishing is reuse-only and switching off preserves prior options."""
schema = NoiseReuseOptionsV3.define_schema()
inputs = {value.id: value for value in schema.inputs}
assert list(inputs) == [
"options",
"enabled",
"finishing_steps",
"finishing_method",
"finishing_start_fraction",
]
assert inputs["enabled"].default is True
assert inputs["finishing_steps"].default == 0
prefix = SamplerOptions(contextual_diffusion=ContextualDiffusionOptions())
assert NoiseReuseOptionsV3.execute(enabled=False, options=prefix)[0] is prefix
enabled = NoiseReuseOptionsV3.execute(options=prefix)[0]
assert enabled.contextual_diffusion is prefix.contextual_diffusion
assert enabled.noise_reuse is not None and enabled.noise_reuse.finishing_steps == 0
def test_provenance_recovery_has_no_vae_or_widgets_and_never_encodes() -> None:
"""Return graph lineage directly instead of producing a re-encoded approximation."""
schema = LatentFromImageV3.define_schema()
assert [value.id for value in schema.inputs] == ["image"]
service = LatentProvenanceService()
prompt = {
"decode": {
"class_type": "VAEDecode",
"inputs": {"samples": ["sample", 0], "vae": ["loader", 2]},
}
}
assert service.recover_link(prompt, ["decode", 0], {}) == ("sample", 0)
with pytest.raises(ValueError, match="unmodified"):
service.recover_link({}, ["loaded_image", 0], {})
@@ -0,0 +1,123 @@
"""Verify receiving prediction views and model-aware schedule geometry."""
from typing import Any, TypedDict
import pytest
import torch
from simple_syrup.domain.noise_reuse import NoiseReuseOptions
from simple_syrup.domain.sampler_options import (
ContextualDiffusionOptions,
SamplerOptions,
)
from simple_syrup.runtime.comfy_latent_normalization import COMFY_LATENT_NORMALIZER
from simple_syrup.runtime.sampling_schedulers import SchedulerView
from simple_syrup.services import noise_capture_request_service as requests
class CaptureArguments(TypedDict):
"""Describe graph controls at the external MODEL and LATENT boundary."""
model: object
steps: int
sampler_name: str
scheduler: str
denoise: float
source_latent: dict[str, Any]
geometry_json: str
def test_cd_coordinate_ignores_unresolved_scale_and_coalesces_finish_requests(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Use the CD context schedule to select source capture coordinates."""
views: list[SchedulerView] = []
def schedule(*args: Any, view: SchedulerView) -> torch.Tensor:
"""Represent the external native schedule boundary with a known coordinate."""
views.append(view)
return torch.tensor([0.4, 0.2, 0])
monkeypatch.setattr(requests, "calculate_sigmas", schedule)
monkeypatch.setattr(
requests, "validate_inversion_target", lambda model, sigmas: float(sigmas[0])
)
arguments: CaptureArguments = {
"model": object(),
"steps": 16,
"sampler_name": "euler",
"scheduler": "normal",
"denoise": 0.2,
"source_latent": {"samples": torch.zeros(1, 16, 1, 96, 168)},
"geometry_json": '[{"scale": ["dynamic_factor", 0]}]',
}
prefix = ContextualDiffusionOptions(context_size=128)
service = requests.NoiseCaptureRequestService()
original = service.coordinates(
**arguments,
options=SamplerOptions(
contextual_diffusion=prefix, noise_reuse=NoiseReuseOptions()
),
)
finished = service.coordinates(
**arguments,
previous=original,
options=SamplerOptions(
contextual_diffusion=prefix,
noise_reuse=NoiseReuseOptions(finishing_steps=1),
),
)
assert finished == (
float(torch.tensor(0.4)),
float(torch.tensor(float(torch.tensor(0.4)) * 0.75)),
)
assert views == [SchedulerView(128, 128)] * 2
def test_flux2_planning_normalizes_empty_layout_before_scale(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Match native packed Flux2 dimensions before its resolution-dependent schedule."""
views: list[SchedulerView] = []
normalized: list[object] = []
def normalize(**kwargs: Any) -> torch.Tensor:
"""Stand in for Comfy's empty latent normalization boundary."""
normalized.append(kwargs["spatial_downscale_ratio"])
return torch.zeros(1, 128, 8, 12)
def schedule(*args: Any, view: SchedulerView) -> torch.Tensor:
"""Record the prediction canvas passed to the native scheduler."""
views.append(view)
return torch.tensor([0.5, 0])
monkeypatch.setattr(COMFY_LATENT_NORMALIZER, "normalize", normalize)
monkeypatch.setattr(requests, "calculate_sigmas", schedule)
monkeypatch.setattr(
requests, "validate_inversion_target", lambda model, sigmas: float(sigmas[0])
)
result = requests.NoiseCaptureRequestService().coordinates(
model=object(),
steps=4,
sampler_name="euler",
scheduler="Flux2",
denoise=0.25,
options=SamplerOptions(noise_reuse=NoiseReuseOptions()),
source_latent={
"samples": torch.zeros(1, 4, 16, 24),
"downscale_ratio_spacial": 8,
},
geometry_json='[{"scale": 1.5}]',
)
assert result == (0.5,) and normalized == [8]
assert views == [SchedulerView(18, 12)]
def test_stock_fixed_size_matches_comfy_pixel_and_aspect_rules() -> None:
"""Match stock resizing's eight-pixel convention and aspect handling."""
samples = torch.zeros(1, 128, 8, 12)
service = requests.NoiseCaptureRequestService()
assert service._view(samples, '[{"pixels": [0, 288]}]') == SchedulerView(36, 24)
assert service._view(samples, '[{"pixels": [0, 0]}]') == SchedulerView(12, 8)
with pytest.raises(ValueError, match="scale"):
service._view(samples, '[{"scale": ["dynamic", 0]}]')
+72
View File
@@ -0,0 +1,72 @@
"""Verify captured coefficient identity, coordinate selection and option isolation."""
import pytest
import torch
from simple_syrup.domain.noise_inversion import NoiseInversionOptions
from simple_syrup.domain.noise_reuse import NoiseReuseOptions, SamplingProvenance
from simple_syrup.domain.sampler_options import (
SamplerOptions,
append_sampler_capability,
)
from simple_syrup.domain.spectral_noise import lift_residual, spatial_dct
def test_lifting_keeps_source_coefficients_and_completes_only_new_bands() -> None:
"""Detect interpolation, whitening and replacement of captured coefficients."""
generator = torch.Generator().manual_seed(17)
source = torch.randn((1, 16, 1, 5, 7), generator=generator) * 0.7 + 0.2
fresh = torch.randn((1, 16, 1, 13, 18), generator=generator)
lifted = spatial_dct(lift_residual(source, fresh))
assert torch.allclose(lifted[..., :5, :7], spatial_dct(source), atol=2e-6)
difference = lifted - spatial_dct(fresh)
difference[..., :5, :7] = 0
assert float(difference.abs().max()) < 2e-6
def test_same_size_reuses_residual_and_wrong_layout_fails() -> None:
"""Avoid extra noise at unchanged size and reject silently broadcast batches."""
source = torch.arange(24).float().reshape(1, 1, 4, 6)
assert torch.allclose(
lift_residual(source, torch.zeros_like(source)), source, atol=3e-6
)
with pytest.raises(ValueError, match="batch"):
lift_residual(source, torch.zeros(2, 1, 8, 12))
with pytest.raises(ValueError, match="enlargement"):
lift_residual(source, torch.zeros(1, 1, 2, 3))
def test_finishing_selects_its_actual_start_and_zero_selects_sampler_start() -> None:
"""Keep sampler denoise separate from inversion finishing work."""
target = float(torch.tensor(0.41665297746658325))
assert NoiseReuseOptions().capture_coordinate(target) == target
assert NoiseReuseOptions(finishing_steps=1).capture_coordinate(target) == float(
torch.tensor(target * 0.75)
)
with pytest.raises(ValueError, match="fraction"):
NoiseReuseOptions(finishing_start_fraction=0)
with pytest.raises(ValueError, match="steps"):
NoiseReuseOptions(finishing_steps=-1)
def test_options_branch_without_mutating_prefix_and_reject_full_inversion() -> None:
"""Reject two competing starting-noise owners in either connection order."""
prefix = SamplerOptions()
reuse = append_sampler_capability(prefix, NoiseReuseOptions())
assert prefix.noise_reuse is None
assert reuse.noise_reuse == NoiseReuseOptions()
with pytest.raises(ValueError, match="combined"):
append_sampler_capability(reuse, NoiseInversionOptions())
with pytest.raises(ValueError, match="combined"):
append_sampler_capability(
SamplerOptions(noise_inversion=NoiseInversionOptions()), NoiseReuseOptions()
)
def test_missing_capture_does_not_fall_back_to_fresh_noise() -> None:
"""Require the requested actual coordinate instead of the closest checkpoint."""
state = torch.zeros(1, 4, 2, 2)
provenance = SamplingProvenance("SDXL", state, ((0.5, state),))
assert provenance.state_at(0.5) is state
with pytest.raises(ValueError, match="did not capture"):
provenance.state_at(0.4)
@@ -0,0 +1,126 @@
"""Verify native residual parameterization and independent full-resolution finishing."""
from types import SimpleNamespace
from typing import Any
import pytest
import torch
from comfy import sample as comfy_sample
from simple_syrup.domain.noise_reuse import NoiseReuseOptions, SamplingProvenance
from simple_syrup.runtime.native_noise_scaling import (
noise_coefficients,
recover_residual,
)
from simple_syrup.runtime.noise_reuse_finishing import finish_reused_state
from simple_syrup.runtime.noise_reuse_scope import (
NoiseReuseExecution,
current_noise_reuse,
noise_reuse_scope,
)
@pytest.mark.parametrize("flow", (False, True))
def test_native_coefficients_reconstruct_eps_and_flow_states(flow: bool) -> None:
"""Detect treating flow noise as an additive epsilon residual."""
clean = torch.arange(16).reshape(1, 4, 2, 2).float() / 4
noise = torch.linspace(-2, 2, 16).reshape_as(clean)
scaling = SimpleNamespace(
noise_scaling=lambda sigma, residual, reference, maximum: (
(1 - sigma) * reference + sigma * residual
if flow
else reference + sigma * residual
)
)
state = scaling.noise_scaling(clean.new_tensor(0.4), noise, clean, False)
recovered = recover_residual(scaling, clean, state, 0.4)
base, amplitude = noise_coefficients(scaling, clean, 0.4)
assert torch.allclose(recovered, noise, atol=2e-6)
assert torch.allclose(base + amplitude * recovered, state, atol=2e-6)
with pytest.raises(ValueError, match="invertible"):
recover_residual(scaling, clean, state, 0)
nonlinear = SimpleNamespace(
noise_scaling=lambda sigma, residual, reference, maximum: (
reference + sigma * residual.square()
)
)
with pytest.raises(ValueError, match="affine"):
noise_coefficients(nonlinear, clean, 0.4)
@pytest.mark.parametrize("method,evaluations", (("euler", 2), ("heun", 4)))
def test_finisher_consumes_supplied_state_and_native_cfg_once(
monkeypatch: pytest.MonkeyPatch, method: Any, evaluations: int
) -> None:
"""Exercise the real integration solver without coarse inversion or double reuse."""
initial = torch.ones(1, 4, 4, 6)
latent = torch.full_like(initial, 7)
model = SimpleNamespace(load_device="cpu")
stage_model = object()
positive, negative = object(), object()
states: list[torch.Tensor] = []
calls: list[dict[str, Any]] = []
factory_shapes: list[tuple[int, ...]] = []
def factory(value: torch.Tensor) -> Any:
"""Return one full-resolution model without resizing the clean reference."""
factory_shapes.append(tuple(value.shape))
return stage_model
def sample(*args: Any, **kwargs: Any) -> torch.Tensor:
"""Stand in only for the external Comfy guider, retaining real solver math."""
assert current_noise_reuse() is None
assert args[0] is stage_model and args[2] == 6.0
assert args[5] is positive and args[6] is negative and args[7] is latent
assert torch.count_nonzero(args[1]) == 0
calls.append(kwargs)
def predict(x: torch.Tensor, sigma: torch.Tensor, **extra: Any) -> torch.Tensor:
"""Expose recognizable native denoised predictions for integration."""
states.append(x.clone())
return torch.zeros_like(x)
result: torch.Tensor = args[3].sampler_function(
predict, torch.full_like(initial, 99), args[4], {}, None, True
)
return result
monkeypatch.setattr(comfy_sample, "sample_custom", sample)
execution = NoiseReuseExecution(
NoiseReuseOptions(), SamplingProvenance("SDXL", latent)
)
with noise_reuse_scope(execution):
output = finish_reused_state(
model=model,
latent=latent,
initial=initial,
begin=0.3,
target=0.4,
options=NoiseReuseOptions(finishing_steps=2, finishing_method=method),
positive=positive,
negative=negative,
cfg=6.0,
seed=17,
noise_mask=None,
model_factory=factory,
)
assert current_noise_reuse() is execution
assert current_noise_reuse() is None
assert len(states) == evaluations and len(calls) == 1
assert torch.equal(states[0], initial)
assert torch.allclose(output, initial * (0.4 / 0.3), atol=2e-6)
assert factory_shapes == [tuple(latent.shape)]
assert calls[0]["seed"] == 17
def test_scope_restores_receiver_after_nested_failure() -> None:
"""Prevent failed finishing from leaving another workflow's reuse context active."""
execution = NoiseReuseExecution(
NoiseReuseOptions(), SamplingProvenance("SDXL", torch.zeros(1, 4, 2, 2))
)
with noise_reuse_scope(execution):
with pytest.raises(RuntimeError, match="boundary"):
with noise_reuse_scope(None):
raise RuntimeError("boundary")
assert current_noise_reuse() is execution
assert current_noise_reuse() is None
+121
View File
@@ -0,0 +1,121 @@
"""Exercise capture through native Euler solver states and verified endpoint
identity.
"""
from types import SimpleNamespace
from typing import Any
import pytest
import torch
from comfy.k_diffusion.sampling import sample_euler
from simple_syrup.domain.noise_reuse import PROVENANCE_KEY, SamplingProvenance
from simple_syrup.runtime.trajectory_capture import (
TrajectoryCapture,
insert_capture_coordinates,
)
class NativeSampler:
"""Exercise the same sampler-instance method binding as native Comfy."""
def __init__(self) -> None:
"""Keep native solver ownership local to the sampler instance."""
self.sampler_function: Any = sample_euler
def sample(self, *args: Any) -> torch.Tensor:
"""Invoke the currently owned solver using a deterministic denoised
prediction.
"""
def model(
state: torch.Tensor, sigma: torch.Tensor, **kwargs: Any
) -> torch.Tensor:
"""Return a clean zero reference through the native model boundary."""
return torch.zeros_like(state)
result: torch.Tensor = self.sampler_function(model, *args, disable=True)
return result
class Executor:
"""Represent Comfy's external sampler-wrapper invocation contract."""
def __init__(self) -> None:
"""Retain one externally owned native sampler."""
self.class_obj = NativeSampler()
self.original: Any = self.class_obj.sample
def __call__(self, *args: Any, **kwargs: Any) -> torch.Tensor:
"""Bind replacement sampler methods like the native WrapperExecutor does."""
result: torch.Tensor = self.original(*args, **kwargs)
return result
def test_native_euler_captures_actual_state_at_inserted_coordinate() -> None:
"""Detect accidental capture of the denoised prediction instead of noisy x."""
format_type = type("Wan21", (), {})
model = SimpleNamespace(
model=SimpleNamespace(
latent_format=format_type(), process_latent_in=lambda x: x
)
)
model.get_model_object = lambda name: SimpleNamespace(
noise_scaling=lambda sigma, noise, clean, maximum: clean + sigma * noise
)
capture = TrajectoryCapture(model, "euler", (0.6, 0.3))
initial = torch.ones(1, 16, 1, 3, 4)
result = capture.observe(Executor(), initial, torch.tensor([1.0, 0.5, 0.0]))
assert torch.allclose(capture.states[float(torch.tensor(0.6))], initial * 0.6)
assert torch.allclose(capture.states[float(torch.tensor(0.3))], initial * 0.3)
attached = capture.attach({"samples": result})
provenance = attached[PROVENANCE_KEY]
assert isinstance(provenance, SamplingProvenance)
assert provenance.latent_format == "Wan21"
assert torch.equal(provenance.clean, result)
assert torch.allclose(provenance.residual_at(float(torch.tensor(0.6))), initial)
assert initial.sum() > 0 and result.sum() == 0
with pytest.raises(ValueError, match="endpoint"):
capture.attach({"samples": result + 1})
def test_capture_uses_native_clean_endpoint_before_external_format_roundoff() -> None:
"""Preserve the original trajectory's reference across lossy format conversions."""
initial = torch.ones(1, 16, 1, 3, 4)
model = SimpleNamespace(
model=SimpleNamespace(
latent_format=type("Wan21", (), {})(),
process_latent_in=lambda value: value + 1e-7,
),
get_model_object=lambda name: SimpleNamespace(
noise_scaling=lambda sigma, noise, clean, maximum: clean + sigma * noise
),
)
capture = TrajectoryCapture(model, "euler", (0.5,))
result = capture.observe(Executor(), initial, torch.tensor([1.0, 0.5, 0.0]))
provenance = capture.attach({"samples": result})[PROVENANCE_KEY]
assert isinstance(provenance, SamplingProvenance)
assert torch.equal(provenance.clean, result)
assert torch.equal(provenance.residual_at(0.5), initial)
def test_coordinate_insertion_keeps_boundaries_and_rejects_invalid_ranges() -> None:
"""Avoid changing source schedule points or silently accepting unavailable
noise levels.
"""
original = torch.tensor([1.0, 0.7, 0.4, 0.0])
schedule = insert_capture_coordinates(original, (0.6, 0.2))
assert torch.equal(schedule, torch.tensor([1.0, 0.7, 0.6, 0.4, 0.2, 0.0]))
assert torch.equal(original, torch.tensor([1.0, 0.7, 0.4, 0.0]))
assert torch.equal(insert_capture_coordinates(schedule, (0.6,)), schedule)
with pytest.raises(ValueError, match="outside"):
insert_capture_coordinates(original, (1.5,))
with pytest.raises(ValueError, match="leftover noise"):
insert_capture_coordinates(torch.tensor([1.0, 0.7, 0.4]), (0.6,))
def test_incomplete_capture_and_unsupported_solver_fail_closed() -> None:
"""Keep unavailable data from masquerading as usable source noise."""
capture = TrajectoryCapture(object(), "dpmpp_2m", (0.5,))
with pytest.raises(ValueError, match="Euler"):
capture.prepare_model()
+4
View File
@@ -1,5 +1,9 @@
# Third-Party Notices
## SesquiLSR latent upscale runtime
The latent upscaler and inference adaptors come from [LoganBooker/SesquiLSR](https://github.com/LoganBooker/SesquiLSR) at revision `befae004248c403f38b76b9f65fd43b901ea3eaa`, under the MIT license. The retained license is `third_party/licenses/sesqui_lsr.LICENSE.txt`.
This repository vendors selected third-party behavior for runtime use. Each vendored component is recorded in `third_party/manifest.toml`, and the corresponding license text is stored in `third_party/licenses/`.
## SAM-HQ and MobileSAM runtime
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2026 SesquiLSR Contributors
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
+13
View File
@@ -192,3 +192,16 @@ vendored_files = [
"simple_syrup/runtime/negpip/krea2.py",
"simple_syrup/services/negpip_model_service.py",
]
[[component]]
name = "SesquiLSR latent upscale runtime"
license = "MIT"
license_file = "third_party/licenses/sesqui_lsr.LICENSE.txt"
source = "https://github.com/LoganBooker/SesquiLSR"
revision = "befae004248c403f38b76b9f65fd43b901ea3eaa"
source_paths = ["sesqui_lsr/__init__.py", "sesqui_lsr/upscaler.py", "sesqui_lsr/inference_adaptors.py"]
vendored_files = [
"simple_syrup/third_party/sesqui_lsr/__init__.py",
"simple_syrup/third_party/sesqui_lsr/upscaler.py",
"simple_syrup/third_party/sesqui_lsr/inference_adaptors.py",
]