diff --git a/__init__.py b/__init__.py index 5870185..c749b76 100644 --- a/__init__.py +++ b/__init__.py @@ -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", diff --git a/pyproject.toml b/pyproject.toml index 7d09591..d6eb87f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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]] diff --git a/simple_syrup/domain/latent_upscale.py b/simple_syrup/domain/latent_upscale.py new file mode 100644 index 0000000..55272fd --- /dev/null +++ b/simple_syrup/domain/latent_upscale.py @@ -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 diff --git a/simple_syrup/domain/noise_reuse.py b/simple_syrup/domain/noise_reuse.py new file mode 100644 index 0000000..252922d --- /dev/null +++ b/simple_syrup/domain/noise_reuse.py @@ -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.") diff --git a/simple_syrup/domain/sampler_options.py b/simple_syrup/domain/sampler_options.py index 18629a8..cf45851 100644 --- a/simple_syrup/domain/sampler_options.py +++ b/simple_syrup/domain/sampler_options.py @@ -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) diff --git a/simple_syrup/domain/spectral_noise.py b/simple_syrup/domain/spectral_noise.py new file mode 100644 index 0000000..8e58d1f --- /dev/null +++ b/simple_syrup/domain/spectral_noise.py @@ -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 + ) diff --git a/simple_syrup/nodes_v3/__init__.py b/simple_syrup/nodes_v3/__init__.py index f9cdffe..378be2a 100644 --- a/simple_syrup/nodes_v3/__init__.py +++ b/simple_syrup/nodes_v3/__init__.py @@ -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, diff --git a/simple_syrup/nodes_v3/latent_from_image.py b/simple_syrup/nodes_v3/latent_from_image.py new file mode 100644 index 0000000..1a6b9ba --- /dev/null +++ b/simple_syrup/nodes_v3/latent_from_image.py @@ -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={}) diff --git a/simple_syrup/nodes_v3/noise_reuse_options.py b/simple_syrup/nodes_v3/noise_reuse_options.py new file mode 100644 index 0000000..dc088d1 --- /dev/null +++ b/simple_syrup/nodes_v3/noise_reuse_options.py @@ -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),) diff --git a/simple_syrup/nodes_v3/simple_upscale_latent.py b/simple_syrup/nodes_v3/simple_upscale_latent.py new file mode 100644 index 0000000..41c43ec --- /dev/null +++ b/simple_syrup/nodes_v3/simple_upscale_latent.py @@ -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),) diff --git a/simple_syrup/nodes_v3/trajectory_capture.py b/simple_syrup/nodes_v3/trajectory_capture.py new file mode 100644 index 0000000..632c651 --- /dev/null +++ b/simple_syrup/nodes_v3/trajectory_capture.py @@ -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) diff --git a/simple_syrup/runtime/contextual_diffusion_sampling.py b/simple_syrup/runtime/contextual_diffusion_sampling.py index df4c608..c1fb700 100644 --- a/simple_syrup/runtime/contextual_diffusion_sampling.py +++ b/simple_syrup/runtime/contextual_diffusion_sampling.py @@ -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( diff --git a/simple_syrup/runtime/guided_sampling.py b/simple_syrup/runtime/guided_sampling.py index 98d6bd3..b128009 100644 --- a/simple_syrup/runtime/guided_sampling.py +++ b/simple_syrup/runtime/guided_sampling.py @@ -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, diff --git a/simple_syrup/runtime/latent_upscale_catalog.py b/simple_syrup/runtime/latent_upscale_catalog.py new file mode 100644 index 0000000..a1fa829 --- /dev/null +++ b/simple_syrup/runtime/latent_upscale_catalog.py @@ -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])) diff --git a/simple_syrup/runtime/latent_upscale_models.py b/simple_syrup/runtime/latent_upscale_models.py new file mode 100644 index 0000000..87f8cba --- /dev/null +++ b/simple_syrup/runtime/latent_upscale_models.py @@ -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 diff --git a/simple_syrup/runtime/mixture_of_diffusers_sampling.py b/simple_syrup/runtime/mixture_of_diffusers_sampling.py index ff3ba3d..bb344e0 100644 --- a/simple_syrup/runtime/mixture_of_diffusers_sampling.py +++ b/simple_syrup/runtime/mixture_of_diffusers_sampling.py @@ -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 ), ) diff --git a/simple_syrup/runtime/multidiffusion_sampling.py b/simple_syrup/runtime/multidiffusion_sampling.py index fe52c2e..17662e1 100644 --- a/simple_syrup/runtime/multidiffusion_sampling.py +++ b/simple_syrup/runtime/multidiffusion_sampling.py @@ -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 ), ) diff --git a/simple_syrup/runtime/native_noise_scaling.py b/simple_syrup/runtime/native_noise_scaling.py new file mode 100644 index 0000000..4b10125 --- /dev/null +++ b/simple_syrup/runtime/native_noise_scaling.py @@ -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 diff --git a/simple_syrup/runtime/noise_reuse_finishing.py b/simple_syrup/runtime/noise_reuse_finishing.py new file mode 100644 index 0000000..285ada7 --- /dev/null +++ b/simple_syrup/runtime/noise_reuse_finishing.py @@ -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] diff --git a/simple_syrup/runtime/noise_reuse_graph.py b/simple_syrup/runtime/noise_reuse_graph.py new file mode 100644 index 0000000..d18e271 --- /dev/null +++ b/simple_syrup/runtime/noise_reuse_graph.py @@ -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 diff --git a/simple_syrup/runtime/noise_reuse_prompt_handler.py b/simple_syrup/runtime/noise_reuse_prompt_handler.py new file mode 100644 index 0000000..85a324d --- /dev/null +++ b/simple_syrup/runtime/noise_reuse_prompt_handler.py @@ -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 diff --git a/simple_syrup/runtime/noise_reuse_sampling.py b/simple_syrup/runtime/noise_reuse_sampling.py new file mode 100644 index 0000000..492db68 --- /dev/null +++ b/simple_syrup/runtime/noise_reuse_sampling.py @@ -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 diff --git a/simple_syrup/runtime/noise_reuse_scope.py b/simple_syrup/runtime/noise_reuse_scope.py new file mode 100644 index 0000000..cc6fdbc --- /dev/null +++ b/simple_syrup/runtime/noise_reuse_scope.py @@ -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) diff --git a/simple_syrup/runtime/regional_multidiffusion_sampling.py b/simple_syrup/runtime/regional_multidiffusion_sampling.py index 3c1111a..f816409 100644 --- a/simple_syrup/runtime/regional_multidiffusion_sampling.py +++ b/simple_syrup/runtime/regional_multidiffusion_sampling.py @@ -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 ), ) diff --git a/simple_syrup/runtime/sesqui_latent_upscale.py b/simple_syrup/runtime/sesqui_latent_upscale.py new file mode 100644 index 0000000..9b9813f --- /dev/null +++ b/simple_syrup/runtime/sesqui_latent_upscale.py @@ -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 diff --git a/simple_syrup/runtime/trajectory_capture.py b/simple_syrup/runtime/trajectory_capture.py new file mode 100644 index 0000000..49b2e13 --- /dev/null +++ b/simple_syrup/runtime/trajectory_capture.py @@ -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 diff --git a/simple_syrup/services/ksampler_sampling_service.py b/simple_syrup/services/ksampler_sampling_service.py index a8aefa8..a2ed623 100644 --- a/simple_syrup/services/ksampler_sampling_service.py +++ b/simple_syrup/services/ksampler_sampling_service.py @@ -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, ) ) diff --git a/simple_syrup/services/latent_provenance_service.py b/simple_syrup/services/latent_provenance_service.py new file mode 100644 index 0000000..5c50afa --- /dev/null +++ b/simple_syrup/services/latent_provenance_service.py @@ -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 diff --git a/simple_syrup/services/latent_upscale_service.py b/simple_syrup/services/latent_upscale_service.py new file mode 100644 index 0000000..e60aaff --- /dev/null +++ b/simple_syrup/services/latent_upscale_service.py @@ -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 diff --git a/simple_syrup/services/noise_capture_request_service.py b/simple_syrup/services/noise_capture_request_service.py new file mode 100644 index 0000000..fc0a205 --- /dev/null +++ b/simple_syrup/services/noise_capture_request_service.py @@ -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) diff --git a/simple_syrup/services/sampler_options_sampling_service.py b/simple_syrup/services/sampler_options_sampling_service.py index 552c38a..14472c4 100644 --- a/simple_syrup/services/sampler_options_sampling_service.py +++ b/simple_syrup/services/sampler_options_sampling_service.py @@ -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): diff --git a/simple_syrup/third_party/sesqui_lsr/__init__.py b/simple_syrup/third_party/sesqui_lsr/__init__.py new file mode 100644 index 0000000..576f7c3 --- /dev/null +++ b/simple_syrup/third_party/sesqui_lsr/__init__.py @@ -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", +] diff --git a/simple_syrup/third_party/sesqui_lsr/inference_adaptors.py b/simple_syrup/third_party/sesqui_lsr/inference_adaptors.py new file mode 100644 index 0000000..e935dd4 --- /dev/null +++ b/simple_syrup/third_party/sesqui_lsr/inference_adaptors.py @@ -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 diff --git a/simple_syrup/third_party/sesqui_lsr/upscaler.py b/simple_syrup/third_party/sesqui_lsr/upscaler.py new file mode 100644 index 0000000..23a693a --- /dev/null +++ b/simple_syrup/third_party/sesqui_lsr/upscaler.py @@ -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 diff --git a/tests/comfy_integration/test_noise_reuse_graph.py b/tests/comfy_integration/test_noise_reuse_graph.py new file mode 100644 index 0000000..3242b20 --- /dev/null +++ b/tests/comfy_integration/test_noise_reuse_graph.py @@ -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] diff --git a/tests/crosscutting/contracts/test_registration.py b/tests/crosscutting/contracts/test_registration.py index d909654..36bbd0c 100644 --- a/tests/crosscutting/contracts/test_registration.py +++ b/tests/crosscutting/contracts/test_registration.py @@ -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", diff --git a/tests/media/test_latent_upscale.py b/tests/media/test_latent_upscale.py new file mode 100644 index 0000000..8a74eaa --- /dev/null +++ b/tests/media/test_latent_upscale.py @@ -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] diff --git a/tests/media/test_sesqui_latent_adapter.py b/tests/media/test_sesqui_latent_adapter.py new file mode 100644 index 0000000..e6791de --- /dev/null +++ b/tests/media/test_sesqui_latent_adapter.py @@ -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), + ) diff --git a/tests/models/patching/test_patcher_lifecycle_policy.py b/tests/models/patching/test_patcher_lifecycle_policy.py index de0a198..1fc7b87 100644 --- a/tests/models/patching/test_patcher_lifecycle_policy.py +++ b/tests/models/patching/test_patcher_lifecycle_policy.py @@ -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, diff --git a/tests/node_api/test_latent_reuse_nodes.py b/tests/node_api/test_latent_reuse_nodes.py new file mode 100644 index 0000000..2a5d9a8 --- /dev/null +++ b/tests/node_api/test_latent_reuse_nodes.py @@ -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], {}) diff --git a/tests/sampling/test_noise_capture_request.py b/tests/sampling/test_noise_capture_request.py new file mode 100644 index 0000000..2dab1ec --- /dev/null +++ b/tests/sampling/test_noise_capture_request.py @@ -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]}]') diff --git a/tests/sampling/test_noise_reuse.py b/tests/sampling/test_noise_reuse.py new file mode 100644 index 0000000..d294493 --- /dev/null +++ b/tests/sampling/test_noise_reuse.py @@ -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) diff --git a/tests/sampling/test_noise_reuse_execution.py b/tests/sampling/test_noise_reuse_execution.py new file mode 100644 index 0000000..ed3f44a --- /dev/null +++ b/tests/sampling/test_noise_reuse_execution.py @@ -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 diff --git a/tests/sampling/test_trajectory_capture.py b/tests/sampling/test_trajectory_capture.py new file mode 100644 index 0000000..13f9f10 --- /dev/null +++ b/tests/sampling/test_trajectory_capture.py @@ -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() diff --git a/third_party/NOTICE.md b/third_party/NOTICE.md index 651ab29..421f616 100644 --- a/third_party/NOTICE.md +++ b/third_party/NOTICE.md @@ -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 diff --git a/third_party/licenses/sesqui_lsr.LICENSE.txt b/third_party/licenses/sesqui_lsr.LICENSE.txt new file mode 100644 index 0000000..cb38dc5 --- /dev/null +++ b/third_party/licenses/sesqui_lsr.LICENSE.txt @@ -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. diff --git a/third_party/manifest.toml b/third_party/manifest.toml index ec0c9ae..1b22a41 100644 --- a/third_party/manifest.toml +++ b/third_party/manifest.toml @@ -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", +]