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:
@@ -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",
|
||||
|
||||
@@ -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]]
|
||||
|
||||
@@ -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
|
||||
@@ -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.")
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
)
|
||||
@@ -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),)
|
||||
@@ -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(
|
||||
|
||||
@@ -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]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
@@ -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
@@ -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",
|
||||
|
||||
@@ -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]
|
||||
@@ -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,
|
||||
|
||||
@@ -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]}]')
|
||||
@@ -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
|
||||
@@ -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()
|
||||
Vendored
+4
@@ -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
@@ -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.
|
||||
Vendored
+13
@@ -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",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user