diff --git a/governance/architecture/soft_reviews.toml b/governance/architecture/soft_reviews.toml index 4e53f17..2f0335d 100644 --- a/governance/architecture/soft_reviews.toml +++ b/governance/architecture/soft_reviews.toml @@ -1,11 +1,10 @@ schema_version = 1 review_by = 2027-03-31 -fingerprint = "sha256:15ead72c26b1b5911aac191d07f36e5d8fd22d583447c2fcb62353be4ad1cd47" +fingerprint = "sha256:a71fe87163eb7585fafbca08b585954be79abd884e3a1df38bc041c2bc934764" cohesive_paths = [ "simple_syrup/masking/prompt_segs_with_sam_service.py", "simple_syrup/nodes/prompt_segs_with_sam.py", - "simple_syrup/nodes_v3/legacy_node_wrappers.py", "simple_syrup/runtime/attention_region_affinity.py", "simple_syrup/runtime/attention_region_capture.py", "simple_syrup/runtime/attention_sampler_lineage.py", diff --git a/simple_syrup/domain/contextual_diffusion.py b/simple_syrup/domain/contextual_diffusion.py index a0dd613..0c3ca2a 100644 --- a/simple_syrup/domain/contextual_diffusion.py +++ b/simple_syrup/domain/contextual_diffusion.py @@ -27,16 +27,37 @@ class ContextualDiffusionControls: global_weight: float global_steps: int global_decay: float + latent_tile_width: int | None = None + latent_tile_height: int | None = None + + @property + def tile_width(self) -> int: + """Use explicit local geometry or the convenience node's context size.""" + return self.latent_tile_width or self.latent_context_size + + @property + def tile_height(self) -> int: + """Keep the global context independent of a rectangular local tile.""" + return self.latent_tile_height or self.latent_context_size def validate(self) -> None: """Reject controls that cannot produce a stable bounded context plan.""" if self.latent_context_size < 16: raise ValueError("latent_context_size must be at least 16 latent pixels.") - if not 0 <= self.latent_context_overlap < self.latent_context_size: + for value in (self.latent_tile_width, self.latent_tile_height): + if value is not None and (type(value) is not int or value < 16): + raise ValueError( + "Local tile dimensions must be at least 16 latent pixels." + ) + if ( + not 0 + <= self.latent_context_overlap + < min(self.tile_width, self.tile_height) + ): raise ValueError( "latent_context_overlap must be non-negative and smaller than " - "latent_context_size." + "both local tile dimensions." ) if self.latent_context_batch_size < 1: raise ValueError("latent_context_batch_size must be at least 1.") @@ -65,6 +86,7 @@ def build_contextual_diffusion_plan( controls: ContextualDiffusionControls, segs: NativeSegs | None, region_masks: torch.Tensor | None = None, + segs_canvas: tuple[int, int] | None = None, ) -> ContextualDiffusionPlan: """Return a global context plus the regular or SEGS-guided context plan.""" @@ -89,27 +111,29 @@ def build_contextual_diffusion_plan( segs=segs, latent_width=latent_width, latent_height=latent_height, - tile_width=controls.latent_context_size, - tile_height=controls.latent_context_size, + tile_width=controls.tile_width, + tile_height=controls.tile_height, overlap=controls.latent_context_overlap, tile_batch_size=controls.latent_context_batch_size, + segs_canvas=segs_canvas, ) elif segs is not None: tile_plan = build_segs_guided_tiled_diffusion_plan( segs=segs, latent_width=latent_width, latent_height=latent_height, - tile_width=controls.latent_context_size, - tile_height=controls.latent_context_size, + tile_width=controls.tile_width, + tile_height=controls.tile_height, overlap=controls.latent_context_overlap, tile_batch_size=controls.latent_context_batch_size, + segs_canvas=segs_canvas, ) else: tile_plan = build_tiled_diffusion_plan( latent_width=latent_width, latent_height=latent_height, - tile_width=controls.latent_context_size, - tile_height=controls.latent_context_size, + tile_width=controls.tile_width, + tile_height=controls.tile_height, overlap=controls.latent_context_overlap, tile_batch_size=controls.latent_context_batch_size, ) diff --git a/simple_syrup/domain/inversion_solver.py b/simple_syrup/domain/inversion_solver.py new file mode 100644 index 0000000..a79d91c --- /dev/null +++ b/simple_syrup/domain/inversion_solver.py @@ -0,0 +1,89 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Integrate source-derived inversion states without ComfyUI dependencies.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import torch + +from .noise_inversion import INVERSION_METHODS, InversionMethod + +InversionVelocity = Callable[[torch.Tensor, torch.Tensor, int], torch.Tensor] +SpatialResize = Callable[[torch.Tensor, int, int], torch.Tensor] + + +@dataclass(slots=True) +class InversionSolverEvidence: + """Count the actual denoiser evaluations performed by an integration stage.""" + + evaluations: int = 0 + + +def integrate_inversion( + source: torch.Tensor, + sigmas: torch.Tensor, + evaluate: InversionVelocity, + *, + method: InversionMethod, + evidence: InversionSolverEvidence | None = None, +) -> torch.Tensor: + """Advance a finite state over strictly increasing positive inversion sigmas.""" + if method not in INVERSION_METHODS: + raise ValueError("Inversion method must be euler or heun.") + if sigmas.ndim != 1 or len(sigmas) < 2 or not bool(torch.isfinite(sigmas).all()): + raise ValueError("A finite one-dimensional inversion schedule is required.") + if not bool(torch.all(sigmas > 0)) or not bool(torch.all(sigmas[1:] > sigmas[:-1])): + raise ValueError("Inversion sigmas must be positive and increasing.") + if not source.is_floating_point() or not bool(torch.isfinite(source).all()): + raise ValueError("Inversion source must contain finite floating-point values.") + record = evidence if evidence is not None else InversionSolverEvidence() + state = source.clone() + + def velocity(x: torch.Tensor, sigma: torch.Tensor, index: int) -> torch.Tensor: + """Count every denoiser evaluation and reject corrupted predictions.""" + record.evaluations += 1 + value = evaluate(x, sigma, index) + if value.shape != x.shape or not bool(torch.isfinite(value).all()): + raise FloatingPointError("Invalid inversion velocity shape or values.") + return value + + for index in range(len(sigmas) - 1): + current, following = sigmas[index], sigmas[index + 1] + delta = following - current + estimate = velocity(state, current, index) + if method == "heun": + corrected = velocity(state + delta * estimate, following, index) + estimate = (estimate + corrected) / 2 + state = state + delta * estimate + if not bool(torch.isfinite(state).all()): + raise FloatingPointError(f"Non-finite inversion state at step {index}.") + return state + + +def lift_inversion_displacement( + full_source: torch.Tensor, + coarse_source: torch.Tensor, + coarse_endpoint: torch.Tensor, + *, + resize: SpatialResize, +) -> torch.Tensor: + """Lift only the inferred change so existing full-size detail survives transfer.""" + if coarse_source.shape != coarse_endpoint.shape: + raise ValueError("Coarse inversion source and endpoint shapes must match.") + if full_source.shape[:-2] != coarse_source.shape[:-2]: + raise ValueError( + "Inversion transfer must preserve batch and channel dimensions." + ) + height, width = full_source.shape[-2:] + lifted = resize(coarse_endpoint - coarse_source, height, width) + if lifted.shape != full_source.shape: + raise ValueError("Inversion displacement resize produced an invalid shape.") + endpoint = full_source + lifted + if not bool(torch.isfinite(endpoint).all()): + raise FloatingPointError("Inversion transfer produced non-finite values.") + return endpoint diff --git a/simple_syrup/domain/noise_inversion.py b/simple_syrup/domain/noise_inversion.py new file mode 100644 index 0000000..bc39f77 --- /dev/null +++ b/simple_syrup/domain/noise_inversion.py @@ -0,0 +1,71 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Define validated source-preserving noise inversion configuration.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Literal + +InversionMethod = Literal["euler", "heun"] +INVERSION_METHODS: tuple[InversionMethod, ...] = ("euler", "heun") + + +@dataclass(frozen=True, slots=True) +class NoiseInversionOptions: + """Configure reduced-resolution inversion and an optional full-size finish. + + The transition is a fraction of the target inversion sigma, not the forward + denoise steps. A full-size inversion uses ``steps`` and needs no transfer. + """ + + method: InversionMethod = "euler" + resolution_scale: float = 0.5 + steps: int = 2 + switch_fraction: float = 0.75 + finishing_steps: int = 1 + finishing_method: InversionMethod = "euler" + + def __post_init__(self) -> None: + """Reject invalid or internally incomplete inversion recipes.""" + if self.method not in INVERSION_METHODS: + raise ValueError("Inversion method must be euler or heun.") + if self.finishing_method not in INVERSION_METHODS: + raise ValueError("Inversion finishing method must be euler or heun.") + if ( + not math.isfinite(self.resolution_scale) + or not 0 < self.resolution_scale <= 1 + ): + raise ValueError("Inversion resolution scale must be in (0, 1].") + if type(self.steps) is not int or not 1 <= self.steps <= 64: + raise ValueError("Inversion steps must be an integer between 1 and 64.") + if type(self.finishing_steps) is not int or not 0 <= self.finishing_steps <= 64: + raise ValueError("Inversion finishing steps must be between 0 and 64.") + if not math.isfinite(self.switch_fraction) or not 0 < self.switch_fraction <= 1: + raise ValueError("Inversion transition must be in (0, 1].") + if self.resolution_scale < 1 and self.finishing_steps > 0: + if self.switch_fraction == 1: + raise ValueError("A full-size finish requires a transition below 100%.") + + @property + def coarse_target_fraction(self) -> float: + """Reach the full target unless an enabled full-size stage follows transfer.""" + return ( + self.switch_fraction + if self.resolution_scale < 1 and self.finishing_steps + else 1.0 + ) + + def coarse_shape(self, height: int, width: int) -> tuple[int, int]: + """Preserve full dimensions or align reduced transformer grids to even sizes.""" + if height < 1 or width < 1: + raise ValueError("Inversion source dimensions must be positive.") + if self.resolution_scale == 1: + return height, width + return ( + max(2, round(height * self.resolution_scale / 2) * 2), + max(2, round(width * self.resolution_scale / 2) * 2), + ) diff --git a/simple_syrup/domain/regional_inversion_geometry.py b/simple_syrup/domain/regional_inversion_geometry.py new file mode 100644 index 0000000..1d1729f --- /dev/null +++ b/simple_syrup/domain/regional_inversion_geometry.py @@ -0,0 +1,57 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Project semantic regional-detailing ownership into each inversion resolution.""" + +from __future__ import annotations + +import math + +import torch.nn.functional as functional + +from .regional_detailing import LatentBox, LatentRegion + + +def project_inversion_regions( + regions: tuple[LatentRegion, ...], + *, + source_width: int, + source_height: int, + target_width: int, + target_height: int, +) -> tuple[LatentRegion, ...]: + """Preserve region identity and conditioning while scaling masks and bounds.""" + if min(source_width, source_height, target_width, target_height) < 1: + raise ValueError("Regional inversion canvases must have positive dimensions.") + if (source_width, source_height) == (target_width, target_height): + return regions + projected: list[LatentRegion] = [] + for region in regions: + box = region.latent_box + if tuple(region.latent_mask.shape) != (source_height, source_width): + raise ValueError("Regional inversion mask must match its canonical canvas.") + if not ( + 0 <= box.x < box.x + box.width <= source_width + and 0 <= box.y < box.y + box.height <= source_height + ): + raise ValueError("Regional inversion bounds must remain inside the canvas.") + left = math.floor(box.x * target_width / source_width) + top = math.floor(box.y * target_height / source_height) + right = math.ceil((box.x + box.width) * target_width / source_width) + bottom = math.ceil((box.y + box.height) * target_height / source_height) + mask = functional.interpolate( + region.latent_mask[None, None].float(), + size=(target_height, target_width), + mode="nearest", + )[0, 0].to(region.latent_mask) + projected.append( + LatentRegion( + region.index, + region.label, + LatentBox(left, top, right - left, bottom - top), + mask, + region.positive, + ) + ) + return tuple(projected) diff --git a/simple_syrup/domain/regional_tiled_diffusion.py b/simple_syrup/domain/regional_tiled_diffusion.py index 6fd2457..a04c155 100644 --- a/simple_syrup/domain/regional_tiled_diffusion.py +++ b/simple_syrup/domain/regional_tiled_diffusion.py @@ -26,6 +26,7 @@ def build_region_constrained_tiled_diffusion_plan( tile_height: int, overlap: int, tile_batch_size: int, + segs_canvas: tuple[int, int] | None = None, ) -> TiledDiffusionPlan: """Build tiles split wherever regional composition or optional SEGS change.""" @@ -37,7 +38,8 @@ def build_region_constrained_tiled_diffusion_plan( ownership_masks = region_ownership if segs is not None: native_segs = coerce_segs(segs) - validate_segs_aspect_ratio(native_segs, latent_height, latent_width) + canvas_height, canvas_width = segs_canvas or (latent_height, latent_width) + validate_segs_aspect_ratio(native_segs, canvas_height, canvas_width) semantic_ownership = segs_ownership_masks( native_segs, latent_height=latent_height, diff --git a/simple_syrup/domain/sampler_options.py b/simple_syrup/domain/sampler_options.py new file mode 100644 index 0000000..ed157f2 --- /dev/null +++ b/simple_syrup/domain/sampler_options.py @@ -0,0 +1,155 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Assemble immutable, order-independent sampler capability configuration.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, replace +from typing import TypeAlias + +from .noise_inversion import NoiseInversionOptions +from .regional_prompting import validate_regional_prompt_weight +from .tiled_diffusion import validate_tiled_diffusion_mode + + +@dataclass(frozen=True, slots=True) +class TilingOptions: + """Own the single local tile layout and blending configuration.""" + + diffusion_mode: str = "multidiffusion" + width: int = 128 + height: int = 128 + overlap: int = 32 + batch_size: int = 4 + differential_diffusion: bool = False + segs: object | None = None + + def __post_init__(self) -> None: + """Reject tile settings that cannot form a bounded prediction plan.""" + validate_tiled_diffusion_mode(self.diffusion_mode) + for name, value in (("width", self.width), ("height", self.height)): + if type(value) is not int or not 16 <= value <= 512: + raise ValueError( + f"Tile {name} must be between 16 and 512 latent pixels." + ) + if type(self.overlap) is not int or not 0 <= self.overlap < min( + self.width, self.height + ): + raise ValueError( + "Tile overlap must be nonnegative and smaller than both dimensions." + ) + if type(self.batch_size) is not int or self.batch_size < 1: + raise ValueError("Tile batch size must be a positive integer.") + if type(self.differential_diffusion) is not bool: + raise TypeError("Differential diffusion must be a boolean.") + + +@dataclass(frozen=True, slots=True) +class ContextualDiffusionOptions: + """Add global scene context to the authoritative local tile predictions. + + With no Tiling options, the context size also supplies the square local + tile size, using 32 latent pixels of overlap and a batch size of four. + """ + + context_size: int = 96 + global_weight: float = 1.0 + global_steps: int = 1 + global_decay: float = 0.5 + + def __post_init__(self) -> None: + """Reject unsupported global-context geometry and schedule values.""" + if type(self.context_size) is not int or not 16 <= self.context_size <= 512: + raise ValueError("Context size must be between 16 and 512 latent pixels.") + if not math.isfinite(self.global_weight) or not 0 <= self.global_weight <= 2: + raise ValueError("Global context weight must be between 0 and 2.") + if type(self.global_steps) is not int or self.global_steps < 0: + raise ValueError("Global context steps must be a nonnegative integer.") + if not math.isfinite(self.global_decay) or not 0 <= self.global_decay <= 1: + raise ValueError("Global context decay must be between 0 and 1.") + + +@dataclass(frozen=True, slots=True) +class AttentionCouplingOptions: + """Carry regional attention inputs without preparing or mutating a model.""" + + region_masks: object + regional_prompt_weight: float = 1.0 + region_mask_feather: int = 0 + + def __post_init__(self) -> None: + """Require a mask source and valid regional attention strengths.""" + if self.region_masks is None: + raise ValueError("Attention Coupling requires region masks.") + validate_regional_prompt_weight(self.regional_prompt_weight) + if type(self.region_mask_feather) is not int or self.region_mask_feather < 0: + raise ValueError("Region mask feather must be a nonnegative integer.") + + +SamplerCapability: TypeAlias = ( + TilingOptions + | ContextualDiffusionOptions + | NoiseInversionOptions + | AttentionCouplingOptions +) + + +@dataclass(frozen=True, slots=True) +class SamplerOptions: + """Own one immutable setting per capability, independent of graph order.""" + + tiling: TilingOptions | None = None + contextual_diffusion: ContextualDiffusionOptions | None = None + noise_inversion: NoiseInversionOptions | None = None + attention_coupling: AttentionCouplingOptions | None = None + + def __post_init__(self) -> None: + """Reject malformed connection payloads at the typed configuration boundary.""" + for name, expected in ( + ("tiling", TilingOptions), + ("contextual_diffusion", ContextualDiffusionOptions), + ("noise_inversion", NoiseInversionOptions), + ("attention_coupling", AttentionCouplingOptions), + ): + value = getattr(self, name) + if value is not None and not isinstance(value, expected): + raise TypeError(f"Sampler option {name} must be {expected.__name__}.") + + def with_capability(self, capability: SamplerCapability) -> SamplerOptions: + """Return a fresh configuration or reject an ambiguous duplicate feature.""" + names: dict[type[object], str] = { + TilingOptions: "tiling", + ContextualDiffusionOptions: "contextual_diffusion", + NoiseInversionOptions: "noise_inversion", + AttentionCouplingOptions: "attention_coupling", + } + name = names.get(type(capability)) + if name is None: + raise TypeError("Unsupported sampler capability configuration.") + if getattr(self, name) is not None: + raise ValueError( + f"Duplicate sampler capability: {name}. Bypass or remove one node." + ) + if isinstance(capability, TilingOptions): + return replace(self, tiling=capability) + if isinstance(capability, ContextualDiffusionOptions): + return replace(self, contextual_diffusion=capability) + if isinstance(capability, NoiseInversionOptions): + return replace(self, noise_inversion=capability) + return replace(self, attention_coupling=capability) + + +def append_sampler_capability( + options: SamplerOptions | None, capability: SamplerCapability +) -> SamplerOptions: + """Start an options chain or append to a validated incoming connection.""" + if options is not None and not isinstance(options, SamplerOptions): + raise TypeError( + "Options input must be a SimpleSyrup sampler options connection." + ) + return (options if options is not None else SamplerOptions()).with_capability( + capability + ) diff --git a/simple_syrup/domain/segs_tiled_diffusion.py b/simple_syrup/domain/segs_tiled_diffusion.py index 601bce4..1e933a3 100644 --- a/simple_syrup/domain/segs_tiled_diffusion.py +++ b/simple_syrup/domain/segs_tiled_diffusion.py @@ -22,16 +22,20 @@ def build_segs_guided_tiled_diffusion_plan( tile_height: int, overlap: int, tile_batch_size: int, + segs_canvas: tuple[int, int] | None = None, ) -> TiledDiffusionPlan: """Build bounded sampling windows whose irregular cores follow supplied SEGS. Every latent pixel receives exactly one ownership core. Each core is sampled through a rectangular window, while its local blend mask retains the irregular boundary and shares a feathered overlap with neighboring cores. + A reduced inversion stage validates proportions against its original canvas + because rounding the reduced dimensions can change their aspect ratio. """ native_segs = coerce_segs(segs) - validate_segs_aspect_ratio(native_segs, latent_height, latent_width) + canvas_height, canvas_width = segs_canvas or (latent_height, latent_width) + validate_segs_aspect_ratio(native_segs, canvas_height, canvas_width) ownership_masks = segs_ownership_masks( native_segs, latent_height=latent_height, diff --git a/simple_syrup/nodes/detail_segs_as_regions.py b/simple_syrup/nodes/detail_segs_as_regions.py index a19a719..6103d3f 100644 --- a/simple_syrup/nodes/detail_segs_as_regions.py +++ b/simple_syrup/nodes/detail_segs_as_regions.py @@ -10,6 +10,7 @@ from typing import Any, ClassVar import torch +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.segs import coerce_segs_group from ..nodes import tooltips from ..nodes.detailer_input_adapters import ( @@ -215,6 +216,7 @@ class DetailSEGSAsRegions: noise_mask_feather: object = 20, tiled_encode: object = False, tiled_decode: object = False, + noise_inversion: NoiseInversionOptions | None = None, ) -> tuple[object]: """Run regional detailing and return the detailed image.""" @@ -237,6 +239,7 @@ class DetailSEGSAsRegions: strict=True, ): result = service.detail( + noise_inversion=noise_inversion, image=single_image, segs=single_segs, model=single_input(model, "model", list_mode, OPERATION), diff --git a/simple_syrup/nodes/detail_segs_by_scale_factor_tiled_diffusion.py b/simple_syrup/nodes/detail_segs_by_scale_factor_tiled_diffusion.py index 12dc109..ccd6648 100644 --- a/simple_syrup/nodes/detail_segs_by_scale_factor_tiled_diffusion.py +++ b/simple_syrup/nodes/detail_segs_by_scale_factor_tiled_diffusion.py @@ -10,6 +10,7 @@ from typing import Any, ClassVar import torch +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.segs import coerce_segs_group from ..domain.tiled_diffusion import TILED_DIFFUSION_MODES from ..nodes import tooltips @@ -263,6 +264,7 @@ class DetailSEGSByScaleFactorTiledDiffusion: latent_tile_height: object = 128, latent_tile_overlap: object = 16, latent_tile_batch_size: object = 4, + noise_inversion: NoiseInversionOptions | None = None, ) -> tuple[object]: """Run tiled diffusion scale-factor detailing and return the image.""" @@ -275,6 +277,7 @@ class DetailSEGSByScaleFactorTiledDiffusion: outputs: list[torch.Tensor] = [] for single_image, single_segs in zip(images, segs_group, strict=True): result = service.detail( + noise_inversion=noise_inversion, image=single_image, segs=single_segs, model=single_input(model, "model", list_mode, OPERATION), diff --git a/simple_syrup/nodes/ksampler_extras.py b/simple_syrup/nodes/ksampler_extras.py index ebc8019..e436722 100644 --- a/simple_syrup/nodes/ksampler_extras.py +++ b/simple_syrup/nodes/ksampler_extras.py @@ -8,6 +8,7 @@ from __future__ import annotations from typing import Any, ClassVar +from ..domain.noise_inversion import NoiseInversionOptions from ..runtime import sampling_samplers, sampling_schedulers from ..services.ksampler_sampling_service import KSamplerSamplingService from . import tooltips @@ -110,6 +111,7 @@ class KSamplerExtras: negative: Any | None = None, latent_image: Latent | None = None, denoise: float = 1.0, + noise_inversion: NoiseInversionOptions | None = None, ) -> tuple[Latent]: """Sample a latent with ComfyUI samplers and extra scheduler sigmas.""" @@ -126,5 +128,6 @@ class KSamplerExtras: negative=negative, latent_image=latent_image, denoise=denoise, + noise_inversion=noise_inversion, ) return (output,) diff --git a/simple_syrup/nodes_v3/__init__.py b/simple_syrup/nodes_v3/__init__.py index 69cc6cc..f9cdffe 100644 --- a/simple_syrup/nodes_v3/__init__.py +++ b/simple_syrup/nodes_v3/__init__.py @@ -14,13 +14,16 @@ def get_nodes() -> list[type[object]]: from .all_prompt_attention_segs import AllPromptAttentionSEGSV3 from .attention_capture_model import AttentionCaptureModelV3 + from .attention_coupling_options import AttentionCouplingOptionsV3 from .attention_masked_conditioning import AttentionMaskedConditioningV3 from .attention_region_mask import AttentionRegionMaskV3 from .batch_region_conditioning import BatchRegionConditioningV3 from .batch_segs import BatchSEGSV3 from .compose_regional_conditioning import ComposeRegionalConditioningV3 from .concept_attention_segs import ConceptAttentionSEGSV3 + from .contextual_diffusion_options import ContextualDiffusionOptionsV3 from .external_llm_prompt import ExternalLLMPromptV3 + from .ksampler import KSamplerV3 from .ksampler_attention_coupling import KSamplerAttentionCouplingV3 from .ksampler_contextual_attention_coupling import ( KSamplerContextualAttentionCouplingV3, @@ -62,6 +65,7 @@ def get_nodes() -> list[type[object]]: from .load_image_list import LoadImageListV3 from .load_mask_batch import LoadMaskBatchV3 from .mask_to_segs import MaskToSEGSV3 + from .noise_inversion_options import NoiseInversionOptionsV3 from .scale_factor import ScaleFactorV3 from .seed_variation import SeedVariationV3 from .simple_load_checkpoint import SimpleLoadCheckpointV3 @@ -71,12 +75,18 @@ def get_nodes() -> list[type[object]]: 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 .vae_decode_options import VAEDecodeOptionsV3 from .vae_encode_options import VAEEncodeOptionsV3 from .wd14_tagger_loader import WD14TaggerLoaderV3 nodes: list[type[object]] = [ AllPromptAttentionSEGSV3, + AttentionCouplingOptionsV3, + ContextualDiffusionOptionsV3, + NoiseInversionOptionsV3, + TilingOptionsV3, + KSamplerV3, AttentionCaptureModelV3, AttentionMaskedConditioningV3, AttentionRegionMaskV3, diff --git a/simple_syrup/nodes_v3/attention_coupling_options.py b/simple_syrup/nodes_v3/attention_coupling_options.py new file mode 100644 index 0000000..7c58903 --- /dev/null +++ b/simple_syrup/nodes_v3/attention_coupling_options.py @@ -0,0 +1,70 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Configure regional attention without applying MODEL patches in the options graph.""" + +from __future__ import annotations + +from typing import Any + +from ..domain.sampler_options import ( + AttentionCouplingOptions, + SamplerOptions, + append_sampler_capability, +) +from .ksampler_schema import attention_coupling_ksampler_inputs +from .sampler_options_schema import ( + COMFY_IO, + OptionsNodeBase, + options_input, + options_output, +) + + +class AttentionCouplingOptionsV3(OptionsNodeBase): + """Pair global-first conditioning on the sampler with ordered region masks.""" + + @classmethod + def define_schema(cls) -> Any: + """Expose regional masks with established strength and feathering controls.""" + controls = attention_coupling_ksampler_inputs(COMFY_IO) + return COMFY_IO.Schema( + node_id="SimpleSyrup.AttentionCouplingOptions", + display_name="Attention Coupling Options", + category="SimpleSyrup/Sampling/Options", + description=( + "Routes global-first sampler conditioning to ordered image " + "regions through regional attention and LoRA hooks." + ), + inputs=[ + options_input(COMFY_IO), + *[ + control + for control in controls + if control.id + in {"region_masks", "regional_prompt_weight", "region_mask_feather"} + ], + ], + outputs=[options_output(COMFY_IO)], + ) + + @classmethod + def execute( + cls, + region_masks: object, + regional_prompt_weight: float = 1.0, + region_mask_feather: int = 0, + options: SamplerOptions | None = None, + ) -> tuple[SamplerOptions]: + """Append regional attention controls while deferring model preparation.""" + return ( + append_sampler_capability( + options, + AttentionCouplingOptions( + region_masks=region_masks, + regional_prompt_weight=regional_prompt_weight, + region_mask_feather=region_mask_feather, + ), + ), + ) diff --git a/simple_syrup/nodes_v3/contextual_diffusion_options.py b/simple_syrup/nodes_v3/contextual_diffusion_options.py new file mode 100644 index 0000000..a1dfa93 --- /dev/null +++ b/simple_syrup/nodes_v3/contextual_diffusion_options.py @@ -0,0 +1,77 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Add whole-scene context without introducing a second local tiling engine.""" + +from __future__ import annotations + +from typing import Any + +from ..domain.sampler_options import ( + ContextualDiffusionOptions, + SamplerOptions, + append_sampler_capability, +) +from .ksampler_schema import contextual_diffusion_inputs +from .sampler_options_schema import ( + COMFY_IO, + OptionsNodeBase, + options_input, + options_output, +) + + +class ContextualDiffusionOptionsV3(OptionsNodeBase): + """Schedule global scene authority over one local tile prediction.""" + + @classmethod + def define_schema(cls) -> Any: + """Expose global controls independently of the optional Tiling contribution.""" + controls = contextual_diffusion_inputs(COMFY_IO) + return COMFY_IO.Schema( + node_id="SimpleSyrup.ContextualDiffusionOptions", + display_name="Contextual Diffusion Options", + category="SimpleSyrup/Sampling/Options", + description=( + "Adds global scene context to local tiles; Tiling options " + "can independently configure local geometry and blending." + ), + inputs=[ + options_input(COMFY_IO), + *[ + control + for control in controls + if control.id + in { + "latent_context_size", + "global_weight", + "global_steps", + "global_decay", + } + ], + ], + outputs=[options_output(COMFY_IO)], + ) + + @classmethod + def execute( + cls, + latent_context_size: int = 96, + global_weight: float = 1.0, + global_steps: int = 1, + global_decay: float = 0.5, + options: SamplerOptions | None = None, + ) -> tuple[SamplerOptions]: + """Append immutable global-context settings in any chain position.""" + return ( + append_sampler_capability( + options, + ContextualDiffusionOptions( + context_size=latent_context_size, + global_weight=global_weight, + global_steps=global_steps, + global_decay=global_decay, + ), + ), + ) diff --git a/simple_syrup/nodes_v3/ksampler.py b/simple_syrup/nodes_v3/ksampler.py new file mode 100644 index 0000000..d7c0ae1 --- /dev/null +++ b/simple_syrup/nodes_v3/ksampler.py @@ -0,0 +1,79 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Expose one native KSampler consuming a composable capability configuration.""" + +from __future__ import annotations + +from typing import Any, ClassVar + +from ..domain.sampler_options import SamplerOptions +from ..nodes import tooltips +from ..services.sampler_options_sampling_service import SamplerOptionsSamplingService +from .ksampler_schema import ksampler_inputs +from .sampler_options_schema import COMFY_IO, OptionsNodeBase, options_input + + +class KSamplerV3(OptionsNodeBase): + """Execute tiling, context, inversion and attention through shared authorities.""" + + service_class: ClassVar[type[SamplerOptionsSamplingService]] = ( + SamplerOptionsSamplingService + ) + + @classmethod + def define_schema(cls) -> Any: + """Declare sampling controls and an optional capabilities connection.""" + return COMFY_IO.Schema( + node_id="SimpleSyrup.KSampler", + display_name="KSampler (SimpleSyrup)", + category="SimpleSyrup/Sampling", + description=( + "Samples latents with connected sampler options for tiling, " + "Contextual Diffusion, noise inversion and Attention Coupling." + ), + inputs=[ + *ksampler_inputs(COMFY_IO, steps_default=20, cfg_default=8.0), + options_input(COMFY_IO), + ], + outputs=[ + COMFY_IO.Latent.Output( + "latent", tooltip=tooltips.DENOISED_LATENT_OUTPUT + ) + ], + ) + + @classmethod + def execute( + cls, + model: Any, + seed: int, + steps: int, + cfg: float, + sampler_name: str, + scheduler: str, + positive: object, + negative: object | None = None, + latent_image: dict[str, Any] | None = None, + denoise: float = 1.0, + options: SamplerOptions | None = None, + ) -> tuple[dict[str, Any]]: + """Delegate sampling without mutating capability configuration.""" + if latent_image is None: + raise TypeError("KSampler requires latent_image.") + return ( + cls.service_class().sample( + 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, + ), + ) diff --git a/simple_syrup/nodes_v3/ksampler_attention_coupling.py b/simple_syrup/nodes_v3/ksampler_attention_coupling.py index 500be78..dcf667e 100644 --- a/simple_syrup/nodes_v3/ksampler_attention_coupling.py +++ b/simple_syrup/nodes_v3/ksampler_attention_coupling.py @@ -17,6 +17,7 @@ from .ksampler_schema import ( ATTENTION_COUPLING_REGIONAL_PROMPT_WEIGHT_DEFAULT, attention_coupling_ksampler_inputs, ) +from .sampler_options_schema import inversion_from_controls, noise_inversion_inputs if TYPE_CHECKING: @@ -68,10 +69,13 @@ class KSamplerAttentionCouplingV3(_ComfyNodeBase): "anima regional prompt", "sdxl regional prompt", ], - inputs=attention_coupling_ksampler_inputs( - _comfy_io, - region_masks_optional=True, - ), + inputs=[ + *attention_coupling_ksampler_inputs( + _comfy_io, + region_masks_optional=True, + ), + *noise_inversion_inputs(_comfy_io, convenience=True), + ], outputs=[ _comfy_io.Latent.Output( "latent", @@ -98,12 +102,28 @@ class KSamplerAttentionCouplingV3(_ComfyNodeBase): ATTENTION_COUPLING_REGIONAL_PROMPT_WEIGHT_DEFAULT ), region_mask_feather: int = 0, + noise_inversion_enabled: bool = False, + inversion_method: str = "euler", + inversion_resolution_scale: float = 0.5, + inversion_steps: int = 2, + inversion_switch_fraction: float = 0.75, + inversion_finishing_steps: int = 1, + inversion_finishing_method: str = "euler", ) -> tuple[dict[str, Any]]: """Delegate ordinary or regional sampling to the routing service.""" if latent_image is None: raise TypeError("KSampler Attention Coupling requires latent_image.") output = cls.sampling_service_class().sample( + noise_inversion=inversion_from_controls( + noise_inversion_enabled=noise_inversion_enabled, + inversion_method=inversion_method, + inversion_resolution_scale=inversion_resolution_scale, + inversion_steps=inversion_steps, + inversion_switch_fraction=inversion_switch_fraction, + inversion_finishing_steps=inversion_finishing_steps, + inversion_finishing_method=inversion_finishing_method, + ), model=model, seed=seed, steps=steps, diff --git a/simple_syrup/nodes_v3/ksampler_contextual_attention_coupling.py b/simple_syrup/nodes_v3/ksampler_contextual_attention_coupling.py index 1cb559d..9228129 100644 --- a/simple_syrup/nodes_v3/ksampler_contextual_attention_coupling.py +++ b/simple_syrup/nodes_v3/ksampler_contextual_attention_coupling.py @@ -17,6 +17,7 @@ from .ksampler_schema import ( attention_coupling_ksampler_inputs, contextual_diffusion_inputs, ) +from .sampler_options_schema import inversion_from_controls, noise_inversion_inputs if TYPE_CHECKING: @@ -75,6 +76,7 @@ class KSamplerContextualAttentionCouplingV3(_ComfyNodeBase): optional=True, tooltip=tooltips.CONTEXTUAL_DIFFUSION_SEGS, ), + *noise_inversion_inputs(_comfy_io, convenience=True), ], outputs=[ _comfy_io.Latent.Output( @@ -112,6 +114,13 @@ class KSamplerContextualAttentionCouplingV3(_ComfyNodeBase): global_steps: int = 1, global_decay: float = 0.5, segs: object | None = None, + noise_inversion_enabled: bool = False, + inversion_method: str = "euler", + inversion_resolution_scale: float = 0.5, + inversion_steps: int = 2, + inversion_switch_fraction: float = 0.75, + inversion_finishing_steps: int = 1, + inversion_finishing_method: str = "euler", ) -> tuple[dict[str, Any], object]: """Delegate the complete request to the combined application service.""" @@ -145,5 +154,14 @@ class KSamplerContextualAttentionCouplingV3(_ComfyNodeBase): global_steps=global_steps, global_decay=global_decay, segs=segs, + noise_inversion=inversion_from_controls( + noise_inversion_enabled=noise_inversion_enabled, + inversion_method=inversion_method, + inversion_resolution_scale=inversion_resolution_scale, + inversion_steps=inversion_steps, + inversion_switch_fraction=inversion_switch_fraction, + inversion_finishing_steps=inversion_finishing_steps, + inversion_finishing_method=inversion_finishing_method, + ), ) return result.latent, result.contexts diff --git a/simple_syrup/nodes_v3/ksampler_contextual_diffusion.py b/simple_syrup/nodes_v3/ksampler_contextual_diffusion.py index 562d227..3e86c98 100644 --- a/simple_syrup/nodes_v3/ksampler_contextual_diffusion.py +++ b/simple_syrup/nodes_v3/ksampler_contextual_diffusion.py @@ -18,6 +18,7 @@ from .ksampler_schema import ( ksampler_inputs, optional_regional_sampling_inputs, ) +from .sampler_options_schema import inversion_from_controls, noise_inversion_inputs if TYPE_CHECKING: @@ -67,6 +68,7 @@ class KSamplerContextualDiffusionV3(_ComfyNodeBase): _comfy_io, segs_tooltip=tooltips.CONTEXTUAL_DIFFUSION_SEGS, ), + *noise_inversion_inputs(_comfy_io, convenience=True), ], outputs=[ _comfy_io.Latent.Output( @@ -104,6 +106,13 @@ class KSamplerContextualDiffusionV3(_ComfyNodeBase): region_masks: object | None = None, regional_prompt_weight: float = 0.5, region_mask_feather: int = 0, + noise_inversion_enabled: bool = False, + inversion_method: str = "euler", + inversion_resolution_scale: float = 0.5, + inversion_steps: int = 2, + inversion_switch_fraction: float = 0.75, + inversion_finishing_steps: int = 1, + inversion_finishing_method: str = "euler", ) -> tuple[dict[str, Any], object]: """Delegate Contextual Diffusion sampling to its application service.""" @@ -131,5 +140,14 @@ class KSamplerContextualDiffusionV3(_ComfyNodeBase): region_masks=region_masks, regional_prompt_weight=regional_prompt_weight, region_mask_feather=region_mask_feather, + noise_inversion=inversion_from_controls( + noise_inversion_enabled=noise_inversion_enabled, + inversion_method=inversion_method, + inversion_resolution_scale=inversion_resolution_scale, + inversion_steps=inversion_steps, + inversion_switch_fraction=inversion_switch_fraction, + inversion_finishing_steps=inversion_finishing_steps, + inversion_finishing_method=inversion_finishing_method, + ), ) return result.latent, result.contexts diff --git a/simple_syrup/nodes_v3/ksampler_prompt_by_region.py b/simple_syrup/nodes_v3/ksampler_prompt_by_region.py index 8971859..53ec77e 100644 --- a/simple_syrup/nodes_v3/ksampler_prompt_by_region.py +++ b/simple_syrup/nodes_v3/ksampler_prompt_by_region.py @@ -13,6 +13,7 @@ from ..nodes import tooltips from ..services.ksampler_sampling_service import KSamplerSamplingService from ..services.regional_conditioning_service import RegionalConditioningService from .ksampler_schema import regional_ksampler_inputs +from .sampler_options_schema import inversion_from_controls, noise_inversion_inputs if TYPE_CHECKING: @@ -51,7 +52,10 @@ class KSamplerPromptByRegionV3(_ComfyNodeBase): "mask-bound regional prompts." ), search_aliases=["ksampler", "regional prompt", "masked prompt"], - inputs=regional_ksampler_inputs(_comfy_io), + inputs=[ + *regional_ksampler_inputs(_comfy_io), + *noise_inversion_inputs(_comfy_io, convenience=True), + ], outputs=[ _comfy_io.Latent.Output( "latent", @@ -76,6 +80,13 @@ class KSamplerPromptByRegionV3(_ComfyNodeBase): region_mask_feather: int = 0, latent_image: dict[str, Any] | None = None, denoise: float = 1.0, + noise_inversion_enabled: bool = False, + inversion_method: str = "euler", + inversion_resolution_scale: float = 0.5, + inversion_steps: int = 2, + inversion_switch_fraction: float = 0.75, + inversion_finishing_steps: int = 1, + inversion_finishing_method: str = "euler", ) -> tuple[dict[str, Any]]: """Assemble regional conditioning and sample the full latent.""" @@ -93,6 +104,15 @@ class KSamplerPromptByRegionV3(_ComfyNodeBase): ) ) output = cls.sampling_service_class().sample( + noise_inversion=inversion_from_controls( + noise_inversion_enabled=noise_inversion_enabled, + inversion_method=inversion_method, + inversion_resolution_scale=inversion_resolution_scale, + inversion_steps=inversion_steps, + inversion_switch_fraction=inversion_switch_fraction, + inversion_finishing_steps=inversion_finishing_steps, + inversion_finishing_method=inversion_finishing_method, + ), model=model, seed=seed, steps=steps, diff --git a/simple_syrup/nodes_v3/ksampler_prompt_by_tiled_region.py b/simple_syrup/nodes_v3/ksampler_prompt_by_tiled_region.py index 80a1ced..54b22fc 100644 --- a/simple_syrup/nodes_v3/ksampler_prompt_by_tiled_region.py +++ b/simple_syrup/nodes_v3/ksampler_prompt_by_tiled_region.py @@ -14,6 +14,7 @@ from ..nodes import tooltips from ..services.regional_conditioning_service import RegionalConditioningService from ..services.tiled_diffusion_sampling_service import TiledDiffusionSamplingService from .ksampler_schema import regional_ksampler_inputs, tiled_diffusion_inputs +from .sampler_options_schema import inversion_from_controls, noise_inversion_inputs if TYPE_CHECKING: @@ -60,6 +61,7 @@ class KSamplerPromptByTiledRegionV3(_ComfyNodeBase): inputs=[ *regional_ksampler_inputs(_comfy_io), *tiled_diffusion_inputs(_comfy_io), + *noise_inversion_inputs(_comfy_io, convenience=True), ], outputs=[ _comfy_io.Latent.Output( @@ -90,6 +92,13 @@ class KSamplerPromptByTiledRegionV3(_ComfyNodeBase): latent_tile_height: int = 128, latent_tile_overlap: int = 16, latent_tile_batch_size: int = 4, + noise_inversion_enabled: bool = False, + inversion_method: str = "euler", + inversion_resolution_scale: float = 0.5, + inversion_steps: int = 2, + inversion_switch_fraction: float = 0.75, + inversion_finishing_steps: int = 1, + inversion_finishing_method: str = "euler", ) -> tuple[dict[str, Any]]: """Assemble regional conditioning and sample overlapping latent tiles.""" @@ -107,6 +116,15 @@ class KSamplerPromptByTiledRegionV3(_ComfyNodeBase): ) ) output = cls.sampling_service_class().sample( + noise_inversion=inversion_from_controls( + noise_inversion_enabled=noise_inversion_enabled, + inversion_method=inversion_method, + inversion_resolution_scale=inversion_resolution_scale, + inversion_steps=inversion_steps, + inversion_switch_fraction=inversion_switch_fraction, + inversion_finishing_steps=inversion_finishing_steps, + inversion_finishing_method=inversion_finishing_method, + ), diffusion_mode=diffusion_mode, model=model, seed=seed, diff --git a/simple_syrup/nodes_v3/ksampler_tiled_attention_coupling.py b/simple_syrup/nodes_v3/ksampler_tiled_attention_coupling.py index c1dcc5e..6be605d 100644 --- a/simple_syrup/nodes_v3/ksampler_tiled_attention_coupling.py +++ b/simple_syrup/nodes_v3/ksampler_tiled_attention_coupling.py @@ -18,6 +18,7 @@ from .ksampler_schema import ( attention_coupling_ksampler_inputs, tiled_diffusion_inputs, ) +from .sampler_options_schema import inversion_from_controls, noise_inversion_inputs if TYPE_CHECKING: @@ -78,6 +79,7 @@ class KSamplerTiledAttentionCouplingV3(_ComfyNodeBase): region_masks_optional=True, ), *tiled_diffusion_inputs(_comfy_io), + *noise_inversion_inputs(_comfy_io, convenience=True), ], outputs=[ _comfy_io.Latent.Output( @@ -105,6 +107,13 @@ class KSamplerTiledAttentionCouplingV3(_ComfyNodeBase): latent_tile_height: int = 128, latent_tile_overlap: int = 16, latent_tile_batch_size: int = 4, + noise_inversion_enabled: bool = False, + inversion_method: str = "euler", + inversion_resolution_scale: float = 0.5, + inversion_steps: int = 2, + inversion_switch_fraction: float = 0.75, + inversion_finishing_steps: int = 1, + inversion_finishing_method: str = "euler", region_masks: object | None = None, regional_prompt_weight: float = ( ATTENTION_COUPLING_REGIONAL_PROMPT_WEIGHT_DEFAULT @@ -116,6 +125,15 @@ class KSamplerTiledAttentionCouplingV3(_ComfyNodeBase): if latent_image is None: raise TypeError("KSampler Tiled Attention Coupling requires latent_image.") output = cls.sampling_service_class().sample( + noise_inversion=inversion_from_controls( + noise_inversion_enabled=noise_inversion_enabled, + inversion_method=inversion_method, + inversion_resolution_scale=inversion_resolution_scale, + inversion_steps=inversion_steps, + inversion_switch_fraction=inversion_switch_fraction, + inversion_finishing_steps=inversion_finishing_steps, + inversion_finishing_method=inversion_finishing_method, + ), diffusion_mode=diffusion_mode, model=model, seed=seed, diff --git a/simple_syrup/nodes_v3/ksampler_tiled_diffusion.py b/simple_syrup/nodes_v3/ksampler_tiled_diffusion.py index f6714cd..fbf546c 100644 --- a/simple_syrup/nodes_v3/ksampler_tiled_diffusion.py +++ b/simple_syrup/nodes_v3/ksampler_tiled_diffusion.py @@ -16,6 +16,7 @@ from .ksampler_schema import ( optional_regional_sampling_inputs, tiled_diffusion_inputs, ) +from .sampler_options_schema import inversion_from_controls, noise_inversion_inputs if TYPE_CHECKING: @@ -65,6 +66,7 @@ class KSamplerTiledDiffusionV3(_ComfyNodeBase): "boundaries while preserving the configured overlap." ), ), + *noise_inversion_inputs(_comfy_io, convenience=True), ], outputs=[ _comfy_io.Latent.Output( @@ -96,6 +98,13 @@ class KSamplerTiledDiffusionV3(_ComfyNodeBase): region_masks: object | None = None, regional_prompt_weight: float = 0.5, region_mask_feather: int = 0, + noise_inversion_enabled: bool = False, + inversion_method: str = "euler", + inversion_resolution_scale: float = 0.5, + inversion_steps: int = 2, + inversion_switch_fraction: float = 0.75, + inversion_finishing_steps: int = 1, + inversion_finishing_method: str = "euler", ) -> tuple[dict[str, Any]]: """Delegate tiled diffusion sampling to its application service.""" @@ -122,5 +131,14 @@ class KSamplerTiledDiffusionV3(_ComfyNodeBase): region_masks=region_masks, regional_prompt_weight=regional_prompt_weight, region_mask_feather=region_mask_feather, + noise_inversion=inversion_from_controls( + noise_inversion_enabled=noise_inversion_enabled, + inversion_method=inversion_method, + inversion_resolution_scale=inversion_resolution_scale, + inversion_steps=inversion_steps, + inversion_switch_fraction=inversion_switch_fraction, + inversion_finishing_steps=inversion_finishing_steps, + inversion_finishing_method=inversion_finishing_method, + ), ) return (output,) diff --git a/simple_syrup/nodes_v3/legacy_inversion_node_adapter.py b/simple_syrup/nodes_v3/legacy_inversion_node_adapter.py new file mode 100644 index 0000000..5f8c608 --- /dev/null +++ b/simple_syrup/nodes_v3/legacy_inversion_node_adapter.py @@ -0,0 +1,86 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Expose shared inversion controls on maintained implementation-backed samplers.""" + +from __future__ import annotations + +from typing import Any + +from ..nodes.detailer_input_adapters import ( + bool_input, + float_input, + int_input, + str_input, +) +from .legacy_node_adapter import LegacyNodeV3Adapter +from .sampler_options_schema import ( + COMFY_IO, + inversion_from_controls, + noise_inversion_inputs, +) + + +class LegacyInversionNodeV3Adapter(LegacyNodeV3Adapter): + """Normalize direct or list-mode widgets into the shared inversion domain value.""" + + @classmethod + def define_schema(cls) -> Any: + """Append optional controls without moving any persisted workflow input.""" + schema = super().define_schema() + schema.inputs.extend(noise_inversion_inputs(COMFY_IO, convenience=True)) + return schema + + @classmethod + def execute(cls, **kwargs: object) -> Any: + """Narrow inversion widgets before delegating normal implementation inputs.""" + values = dict(kwargs) + list_mode = bool(getattr(cls.LEGACY_NODE_CLASS, "INPUT_IS_LIST", False)) + operation = cls.DISPLAY_NAME + inversion = inversion_from_controls( + noise_inversion_enabled=bool_input( + values.pop("noise_inversion_enabled", False), + "noise_inversion_enabled", + list_mode, + operation, + ), + inversion_method=str_input( + values.pop("inversion_method", "euler"), + "inversion_method", + list_mode, + operation, + ), + inversion_resolution_scale=float_input( + values.pop("inversion_resolution_scale", 0.5), + "inversion_resolution_scale", + list_mode, + operation, + ), + inversion_steps=int_input( + values.pop("inversion_steps", 2), + "inversion_steps", + list_mode, + operation, + ), + inversion_switch_fraction=float_input( + values.pop("inversion_switch_fraction", 0.75), + "inversion_switch_fraction", + list_mode, + operation, + ), + inversion_finishing_steps=int_input( + values.pop("inversion_finishing_steps", 1), + "inversion_finishing_steps", + list_mode, + operation, + ), + inversion_finishing_method=str_input( + values.pop("inversion_finishing_method", "euler"), + "inversion_finishing_method", + list_mode, + operation, + ), + ) + values["noise_inversion"] = inversion + return super().execute(**values) diff --git a/simple_syrup/nodes_v3/legacy_node_adapter.py b/simple_syrup/nodes_v3/legacy_node_adapter.py new file mode 100644 index 0000000..42db0fa --- /dev/null +++ b/simple_syrup/nodes_v3/legacy_node_adapter.py @@ -0,0 +1,275 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Translate maintained implementation contracts into native Comfy v3 schemas.""" + +from __future__ import annotations + +from collections.abc import Mapping +from importlib import import_module +from typing import TYPE_CHECKING, Any, ClassVar + +if TYPE_CHECKING: + + class _ComfyNodeBase: + """Type-checking base for Comfy v3 nodes.""" + + hidden: ClassVar[Any] + RETURN_TYPES: ClassVar[list[str]] + RETURN_NAMES: ClassVar[list[str]] + +else: + _ComfyNodeBase = import_module("comfy_api.latest").io.ComfyNode + +_comfy_io: Any = None if TYPE_CHECKING else import_module("comfy_api.latest").io + +_HIDDEN_INPUTS = { + "PROMPT": "prompt", + "DYNPROMPT": "dynprompt", + "EXTRA_PNGINFO": "extra_pnginfo", + "UNIQUE_ID": "unique_id", +} + + +class LegacyNodeV3Adapter(_ComfyNodeBase): + """Build a v3 schema and execution bridge for a legacy implementation class.""" + + LEGACY_NODE_CLASS: ClassVar[type[Any]] + NODE_ID: ClassVar[str] + DISPLAY_NAME: ClassVar[str] + ENABLE_EXPAND: ClassVar[bool] = False + WORKFLOW_INPUT_ORDER: ClassVar[tuple[str, ...] | None] = None + + @classmethod + def define_schema(cls) -> Any: + """Declare a v3 schema from the implementation class contract.""" + + legacy = cls.LEGACY_NODE_CLASS + return _comfy_io.Schema( + node_id=cls.NODE_ID, + display_name=cls.DISPLAY_NAME, + category=str(getattr(legacy, "CATEGORY", "SimpleSyrup")), + description=str(getattr(legacy, "DESCRIPTION", "")), + search_aliases=list(getattr(legacy, "SEARCH_ALIASES", [])), + inputs=_v3_inputs( + legacy.INPUT_TYPES(), + workflow_order=cls.WORKFLOW_INPUT_ORDER, + ), + outputs=_v3_outputs(legacy), + hidden=_v3_hidden_inputs(legacy.INPUT_TYPES()), + is_input_list=bool(getattr(legacy, "INPUT_IS_LIST", False)), + is_output_node=bool(getattr(legacy, "OUTPUT_NODE", False)), + enable_expand=cls.ENABLE_EXPAND, + ) + + @classmethod + def execute(cls, **kwargs: object) -> Any: + """Run the wrapped implementation with v3-provided inputs.""" + + values = dict(kwargs) + for name, hidden_attr in _legacy_hidden_inputs( + cls.LEGACY_NODE_CLASS.INPUT_TYPES() + ).items(): + if name not in values: + values[name] = getattr(cls.hidden, hidden_attr) + + function_name = str(cls.LEGACY_NODE_CLASS.FUNCTION) + implementation = cls.LEGACY_NODE_CLASS() + function = getattr(implementation, function_name) + return function(**values) + + +def _v3_inputs( + input_types: Mapping[str, Mapping[str, object]], + *, + workflow_order: tuple[str, ...] | None = None, +) -> list[Any]: + """Return v3 inputs while preserving any explicit persisted socket order.""" + + declarations: dict[str, tuple[object, bool]] = {} + for section_name, optional in (("required", False), ("optional", True)): + section = input_types.get(section_name, {}) + for name, declaration in section.items(): + if name in declarations: + raise ValueError(f"legacy input {name} is declared more than once.") + declarations[name] = (declaration, optional) + order = tuple(declarations) if workflow_order is None else workflow_order + if len(order) != len(set(order)) or set(order) != set(declarations): + raise ValueError("legacy workflow input order must name every input once.") + return [ + _v3_input(name, declarations[name][0], optional=declarations[name][1]) + for name in order + ] + + +def _v3_input(name: str, declaration: object, *, optional: bool) -> Any: + """Return one v3 input declaration from a legacy field declaration.""" + + if not isinstance(declaration, tuple) or not declaration: + raise TypeError(f"legacy input {name} declaration must be a tuple.") + + io_declaration = declaration[0] + options = _input_options(declaration) + tooltip = _string_option(options, "tooltip") + advanced = _bool_option(options, "advanced") + raw_link = _bool_option(options, "rawLink") or _bool_option(options, "raw_link") + force_input = _bool_option(options, "forceInput") or _bool_option( + options, "force_input" + ) + + if isinstance(io_declaration, (list, tuple)): + return _comfy_io.Combo.Input( + name, + options=list(io_declaration), + optional=optional, + default=options.get("default"), + control_after_generate=options.get("control_after_generate"), + tooltip=tooltip, + raw_link=raw_link, + advanced=advanced, + ) + + if not isinstance(io_declaration, str): + raise TypeError(f"legacy input {name} type must be a string or options list.") + + input_type = io_declaration + input_class = _io_class(input_type) + common_options = { + "optional": optional, + "tooltip": tooltip, + "raw_link": raw_link, + "advanced": advanced, + } + + if input_type == "INT": + return input_class.Input( + name, + default=options.get("default"), + min=options.get("min"), + max=options.get("max"), + step=options.get("step"), + control_after_generate=options.get("control_after_generate"), + **common_options, + ) + if input_type == "FLOAT": + return input_class.Input( + name, + default=options.get("default"), + min=options.get("min"), + max=options.get("max"), + step=options.get("step"), + round=options.get("round"), + **common_options, + ) + if input_type == "STRING": + return input_class.Input( + name, + default=options.get("default"), + multiline=bool(options.get("multiline", False)), + force_input=force_input, + **common_options, + ) + if input_type == "BOOLEAN": + return input_class.Input( + name, + default=options.get("default"), + label_on=options.get("label_on"), + label_off=options.get("label_off"), + **common_options, + ) + + return input_class.Input(name, **common_options) + + +def _v3_outputs(legacy: type[Any]) -> list[Any]: + """Return v3 output declarations from legacy return metadata.""" + + return_types = tuple(getattr(legacy, "RETURN_TYPES", ())) + return_names = getattr(legacy, "RETURN_NAMES", None) + output_tooltips = tuple(getattr(legacy, "OUTPUT_TOOLTIPS", ())) + output_is_list = tuple( + getattr(legacy, "OUTPUT_IS_LIST", (False,) * len(return_types)) + ) + outputs: list[Any] = [] + for index, io_type in enumerate(return_types): + output_name = None + if isinstance(return_names, tuple) and index < len(return_names): + output_name = str(return_names[index]) + tooltip = None + if index < len(output_tooltips): + tooltip = str(output_tooltips[index]) + is_output_list = index < len(output_is_list) and bool(output_is_list[index]) + outputs.append( + _io_class(str(io_type)).Output( + output_name, + tooltip=tooltip, + is_output_list=is_output_list, + ) + ) + return outputs + + +def _v3_hidden_inputs(input_types: Mapping[str, Mapping[str, object]]) -> list[Any]: + """Return v3 hidden declarations requested by legacy hidden inputs.""" + + hidden_values = set(_legacy_hidden_inputs(input_types).values()) + return [getattr(_comfy_io.Hidden, value) for value in sorted(hidden_values)] + + +def _legacy_hidden_inputs( + input_types: Mapping[str, Mapping[str, object]], +) -> dict[str, str]: + """Return legacy hidden input names mapped to v3 hidden holder attributes.""" + + hidden_inputs: dict[str, str] = {} + for name, sentinel in input_types.get("hidden", {}).items(): + if isinstance(sentinel, str) and sentinel in _HIDDEN_INPUTS: + hidden_inputs[name] = _HIDDEN_INPUTS[sentinel] + return hidden_inputs + + +def _io_class(io_type: str) -> Any: + """Return the v3 IO class for a legacy Comfy type string.""" + + known_types = { + "BOOLEAN": _comfy_io.Boolean, + "INT": _comfy_io.Int, + "FLOAT": _comfy_io.Float, + "STRING": _comfy_io.String, + "IMAGE": _comfy_io.Image, + "MASK": _comfy_io.Mask, + "LATENT": _comfy_io.Latent, + "MODEL": _comfy_io.Model, + "CLIP": _comfy_io.Clip, + "VAE": _comfy_io.Vae, + "CONDITIONING": _comfy_io.Conditioning, + "SEGS": _comfy_io.SEGS, + } + return known_types.get(io_type, _comfy_io.Custom(io_type)) + + +def _input_options(declaration: tuple[object, ...]) -> dict[str, object]: + """Return an input options dictionary from a legacy declaration.""" + + if len(declaration) < 2 or not isinstance(declaration[1], dict): + return {} + return dict(declaration[1]) + + +def _string_option(options: Mapping[str, object], name: str) -> str | None: + """Return a string option when present.""" + + value = options.get(name) + if isinstance(value, str): + return value + return None + + +def _bool_option(options: Mapping[str, object], name: str) -> bool | None: + """Return a boolean option when present.""" + + value = options.get(name) + if isinstance(value, bool): + return value + return None diff --git a/simple_syrup/nodes_v3/legacy_node_wrappers.py b/simple_syrup/nodes_v3/legacy_node_wrappers.py index e89a5dc..b9d53b4 100644 --- a/simple_syrup/nodes_v3/legacy_node_wrappers.py +++ b/simple_syrup/nodes_v3/legacy_node_wrappers.py @@ -6,10 +6,6 @@ from __future__ import annotations -from collections.abc import Mapping -from importlib import import_module -from typing import TYPE_CHECKING, Any, ClassVar - from ..nodes.conditioning_batch_pack import ( ConditioningBatchAppend, ConditioningBatchStart, @@ -40,6 +36,8 @@ from ..nodes.segs_from_sam_output import SEGSFromSAMOutput from ..nodes.simple_load_anima import SimpleLoadAnima from ..nodes.simple_preview_segs import SimplePreviewSEGS from ..nodes.vitmatte_model_loader import ViTMatteModelLoader +from .legacy_inversion_node_adapter import LegacyInversionNodeV3Adapter +from .legacy_node_adapter import LegacyNodeV3Adapter from .legacy_workflow_input_order import ( DETAIL_SEGS_AS_REGIONS_INPUT_ORDER, DETAIL_SEGS_BY_SCALE_FACTOR_INPUT_ORDER, @@ -47,75 +45,6 @@ from .legacy_workflow_input_order import ( KSAMPLER_EXTRAS_INPUT_ORDER, ) -if TYPE_CHECKING: - - class _ComfyNodeBase: - """Type-checking base for Comfy v3 nodes.""" - - hidden: ClassVar[Any] - RETURN_TYPES: ClassVar[list[str]] - RETURN_NAMES: ClassVar[list[str]] - -else: - _ComfyNodeBase = import_module("comfy_api.latest").io.ComfyNode - -_comfy_io: Any = None if TYPE_CHECKING else import_module("comfy_api.latest").io - -_HIDDEN_INPUTS = { - "PROMPT": "prompt", - "DYNPROMPT": "dynprompt", - "EXTRA_PNGINFO": "extra_pnginfo", - "UNIQUE_ID": "unique_id", -} - - -class LegacyNodeV3Adapter(_ComfyNodeBase): - """Build a v3 schema and execution bridge for a legacy implementation class.""" - - LEGACY_NODE_CLASS: ClassVar[type[Any]] - NODE_ID: ClassVar[str] - DISPLAY_NAME: ClassVar[str] - ENABLE_EXPAND: ClassVar[bool] = False - WORKFLOW_INPUT_ORDER: ClassVar[tuple[str, ...] | None] = None - - @classmethod - def define_schema(cls) -> Any: - """Declare a v3 schema from the implementation class contract.""" - - legacy = cls.LEGACY_NODE_CLASS - return _comfy_io.Schema( - node_id=cls.NODE_ID, - display_name=cls.DISPLAY_NAME, - category=str(getattr(legacy, "CATEGORY", "SimpleSyrup")), - description=str(getattr(legacy, "DESCRIPTION", "")), - search_aliases=list(getattr(legacy, "SEARCH_ALIASES", [])), - inputs=_v3_inputs( - legacy.INPUT_TYPES(), - workflow_order=cls.WORKFLOW_INPUT_ORDER, - ), - outputs=_v3_outputs(legacy), - hidden=_v3_hidden_inputs(legacy.INPUT_TYPES()), - is_input_list=bool(getattr(legacy, "INPUT_IS_LIST", False)), - is_output_node=bool(getattr(legacy, "OUTPUT_NODE", False)), - enable_expand=cls.ENABLE_EXPAND, - ) - - @classmethod - def execute(cls, **kwargs: object) -> Any: - """Run the wrapped implementation with v3-provided inputs.""" - - values = dict(kwargs) - for name, hidden_attr in _legacy_hidden_inputs( - cls.LEGACY_NODE_CLASS.INPUT_TYPES() - ).items(): - if name not in values: - values[name] = getattr(cls.hidden, hidden_attr) - - function_name = str(cls.LEGACY_NODE_CLASS.FUNCTION) - implementation = cls.LEGACY_NODE_CLASS() - function = getattr(implementation, function_name) - return function(**values) - class ConditioningBatchStartV3(LegacyNodeV3Adapter): """Expose Conditioning Batch Start through Comfy v3 only.""" @@ -149,7 +78,7 @@ class GroundingDINOModelLoaderV3(LegacyNodeV3Adapter): DISPLAY_NAME = "GroundingDINO Model Loader" -class KSamplerExtrasV3(LegacyNodeV3Adapter): +class KSamplerExtrasV3(LegacyInversionNodeV3Adapter): """Expose KSampler Extras through Comfy v3 only.""" LEGACY_NODE_CLASS = KSamplerExtras @@ -224,7 +153,7 @@ class ResizeImageToTargetV3(LegacyNodeV3Adapter): DISPLAY_NAME = "Resize Image to Target" -class DetailSEGSAsRegionsV3(LegacyNodeV3Adapter): +class DetailSEGSAsRegionsV3(LegacyInversionNodeV3Adapter): """Expose Detail SEGS as Regions through Comfy v3 only.""" LEGACY_NODE_CLASS = DetailSEGSAsRegions @@ -242,7 +171,7 @@ class DetailSEGSByScaleFactorV3(LegacyNodeV3Adapter): WORKFLOW_INPUT_ORDER = DETAIL_SEGS_BY_SCALE_FACTOR_INPUT_ORDER -class DetailSEGSByScaleFactorTiledDiffusionV3(LegacyNodeV3Adapter): +class DetailSEGSByScaleFactorTiledDiffusionV3(LegacyInversionNodeV3Adapter): """Expose Detail SEGS by Scale Factor with Tiled Diffusion through Comfy v3.""" LEGACY_NODE_CLASS = DetailSEGSByScaleFactorTiledDiffusion @@ -323,201 +252,6 @@ class ViTMatteModelLoaderV3(LegacyNodeV3Adapter): DISPLAY_NAME = "ViTMatte Model Loader" -def _v3_inputs( - input_types: Mapping[str, Mapping[str, object]], - *, - workflow_order: tuple[str, ...] | None = None, -) -> list[Any]: - """Return v3 inputs while preserving any explicit persisted socket order.""" - - declarations: dict[str, tuple[object, bool]] = {} - for section_name, optional in (("required", False), ("optional", True)): - section = input_types.get(section_name, {}) - for name, declaration in section.items(): - if name in declarations: - raise ValueError(f"legacy input {name} is declared more than once.") - declarations[name] = (declaration, optional) - order = tuple(declarations) if workflow_order is None else workflow_order - if len(order) != len(set(order)) or set(order) != set(declarations): - raise ValueError("legacy workflow input order must name every input once.") - return [ - _v3_input(name, declarations[name][0], optional=declarations[name][1]) - for name in order - ] - - -def _v3_input(name: str, declaration: object, *, optional: bool) -> Any: - """Return one v3 input declaration from a legacy field declaration.""" - - if not isinstance(declaration, tuple) or not declaration: - raise TypeError(f"legacy input {name} declaration must be a tuple.") - - io_declaration = declaration[0] - options = _input_options(declaration) - tooltip = _string_option(options, "tooltip") - advanced = _bool_option(options, "advanced") - raw_link = _bool_option(options, "rawLink") or _bool_option(options, "raw_link") - force_input = _bool_option(options, "forceInput") or _bool_option( - options, "force_input" - ) - - if isinstance(io_declaration, (list, tuple)): - return _comfy_io.Combo.Input( - name, - options=list(io_declaration), - optional=optional, - default=options.get("default"), - control_after_generate=options.get("control_after_generate"), - tooltip=tooltip, - raw_link=raw_link, - advanced=advanced, - ) - - if not isinstance(io_declaration, str): - raise TypeError(f"legacy input {name} type must be a string or options list.") - - input_type = io_declaration - input_class = _io_class(input_type) - common_options = { - "optional": optional, - "tooltip": tooltip, - "raw_link": raw_link, - "advanced": advanced, - } - - if input_type == "INT": - return input_class.Input( - name, - default=options.get("default"), - min=options.get("min"), - max=options.get("max"), - step=options.get("step"), - control_after_generate=options.get("control_after_generate"), - **common_options, - ) - if input_type == "FLOAT": - return input_class.Input( - name, - default=options.get("default"), - min=options.get("min"), - max=options.get("max"), - step=options.get("step"), - round=options.get("round"), - **common_options, - ) - if input_type == "STRING": - return input_class.Input( - name, - default=options.get("default"), - multiline=bool(options.get("multiline", False)), - force_input=force_input, - **common_options, - ) - if input_type == "BOOLEAN": - return input_class.Input( - name, - default=options.get("default"), - label_on=options.get("label_on"), - label_off=options.get("label_off"), - **common_options, - ) - - return input_class.Input(name, **common_options) - - -def _v3_outputs(legacy: type[Any]) -> list[Any]: - """Return v3 output declarations from legacy return metadata.""" - - return_types = tuple(getattr(legacy, "RETURN_TYPES", ())) - return_names = getattr(legacy, "RETURN_NAMES", None) - output_tooltips = tuple(getattr(legacy, "OUTPUT_TOOLTIPS", ())) - output_is_list = tuple( - getattr(legacy, "OUTPUT_IS_LIST", (False,) * len(return_types)) - ) - outputs: list[Any] = [] - for index, io_type in enumerate(return_types): - output_name = None - if isinstance(return_names, tuple) and index < len(return_names): - output_name = str(return_names[index]) - tooltip = None - if index < len(output_tooltips): - tooltip = str(output_tooltips[index]) - is_output_list = index < len(output_is_list) and bool(output_is_list[index]) - outputs.append( - _io_class(str(io_type)).Output( - output_name, - tooltip=tooltip, - is_output_list=is_output_list, - ) - ) - return outputs - - -def _v3_hidden_inputs(input_types: Mapping[str, Mapping[str, object]]) -> list[Any]: - """Return v3 hidden declarations requested by legacy hidden inputs.""" - - hidden_values = set(_legacy_hidden_inputs(input_types).values()) - return [getattr(_comfy_io.Hidden, value) for value in sorted(hidden_values)] - - -def _legacy_hidden_inputs( - input_types: Mapping[str, Mapping[str, object]], -) -> dict[str, str]: - """Return legacy hidden input names mapped to v3 hidden holder attributes.""" - - hidden_inputs: dict[str, str] = {} - for name, sentinel in input_types.get("hidden", {}).items(): - if isinstance(sentinel, str) and sentinel in _HIDDEN_INPUTS: - hidden_inputs[name] = _HIDDEN_INPUTS[sentinel] - return hidden_inputs - - -def _io_class(io_type: str) -> Any: - """Return the v3 IO class for a legacy Comfy type string.""" - - known_types = { - "BOOLEAN": _comfy_io.Boolean, - "INT": _comfy_io.Int, - "FLOAT": _comfy_io.Float, - "STRING": _comfy_io.String, - "IMAGE": _comfy_io.Image, - "MASK": _comfy_io.Mask, - "LATENT": _comfy_io.Latent, - "MODEL": _comfy_io.Model, - "CLIP": _comfy_io.Clip, - "VAE": _comfy_io.Vae, - "CONDITIONING": _comfy_io.Conditioning, - "SEGS": _comfy_io.SEGS, - } - return known_types.get(io_type, _comfy_io.Custom(io_type)) - - -def _input_options(declaration: tuple[object, ...]) -> dict[str, object]: - """Return an input options dictionary from a legacy declaration.""" - - if len(declaration) < 2 or not isinstance(declaration[1], dict): - return {} - return dict(declaration[1]) - - -def _string_option(options: Mapping[str, object], name: str) -> str | None: - """Return a string option when present.""" - - value = options.get(name) - if isinstance(value, str): - return value - return None - - -def _bool_option(options: Mapping[str, object], name: str) -> bool | None: - """Return a boolean option when present.""" - - value = options.get(name) - if isinstance(value, bool): - return value - return None - - __all__ = [ "ConditioningBatchAppendV3", "ConditioningBatchStartV3", diff --git a/simple_syrup/nodes_v3/noise_inversion_options.py b/simple_syrup/nodes_v3/noise_inversion_options.py new file mode 100644 index 0000000..8e03dcf --- /dev/null +++ b/simple_syrup/nodes_v3/noise_inversion_options.py @@ -0,0 +1,62 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Expose inversion resolution and integration controls as sampler options.""" + +from __future__ import annotations + +from typing import Any + +from ..domain.sampler_options import SamplerOptions, append_sampler_capability +from .sampler_options_schema import ( + COMFY_IO, + OptionsNodeBase, + inversion_from_controls, + noise_inversion_inputs, + options_input, + options_output, +) + + +class NoiseInversionOptionsV3(OptionsNodeBase): + """Add source-derived starting noise to an immutable sampler options chain.""" + + @classmethod + def define_schema(cls) -> Any: + """Declare the accepted recipe with independently editable controls.""" + return COMFY_IO.Schema( + node_id="SimpleSyrup.NoiseInversionOptions", + display_name="Noise Inversion Options", + category="SimpleSyrup/Sampling/Options", + description=( + "Derives starting noise from an input image before sampling; " + "control inversion quality and cost independently." + ), + inputs=[options_input(COMFY_IO), *noise_inversion_inputs(COMFY_IO)], + outputs=[options_output(COMFY_IO)], + ) + + @classmethod + def execute( + cls, + inversion_method: str = "euler", + inversion_resolution_scale: float = 0.5, + inversion_steps: int = 2, + inversion_switch_fraction: float = 0.75, + inversion_finishing_steps: int = 1, + inversion_finishing_method: str = "euler", + options: SamplerOptions | None = None, + ) -> tuple[SamplerOptions]: + """Append validated inversion configuration without preparing any model.""" + inversion = inversion_from_controls( + noise_inversion_enabled=True, + inversion_method=inversion_method, + inversion_resolution_scale=inversion_resolution_scale, + inversion_steps=inversion_steps, + inversion_switch_fraction=inversion_switch_fraction, + inversion_finishing_steps=inversion_finishing_steps, + inversion_finishing_method=inversion_finishing_method, + ) + assert inversion is not None + return (append_sampler_capability(options, inversion),) diff --git a/simple_syrup/nodes_v3/sampler_options_schema.py b/simple_syrup/nodes_v3/sampler_options_schema.py new file mode 100644 index 0000000..cc41878 --- /dev/null +++ b/simple_syrup/nodes_v3/sampler_options_schema.py @@ -0,0 +1,163 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Declare bypass-compatible sampler options sockets and inversion controls.""" + +from __future__ import annotations + +from importlib import import_module +from typing import TYPE_CHECKING, Any, ClassVar, cast + +from ..domain.noise_inversion import ( + INVERSION_METHODS, + InversionMethod, + NoiseInversionOptions, +) + +if TYPE_CHECKING: + + class OptionsNodeBase: + """Describe Comfy's host-facing node metadata for strict type checking.""" + + RETURN_TYPES: ClassVar[list[str]] + RETURN_NAMES: ClassVar[list[str]] + +else: + OptionsNodeBase = import_module("comfy_api.latest").io.ComfyNode + +COMFY_IO: Any = None if TYPE_CHECKING else import_module("comfy_api.latest").io +OPTIONS_TYPE = "SIMPLE_SYRUP_SAMPLER_OPTIONS" + + +def options_input(comfy_io: Any) -> Any: + """Allow any capability to start a chain or consume a preceding capability.""" + return comfy_io.Custom(OPTIONS_TYPE).Input( + "options", + optional=True, + tooltip=( + "Optional preceding sampler options; bypass this node " + "to omit its contribution." + ), + ) + + +def options_output(comfy_io: Any) -> Any: + """Match the input type so Comfy can bypass capability nodes natively.""" + return comfy_io.Custom(OPTIONS_TYPE).Output( + "options", + tooltip="Combined sampler options; connect another options node or KSampler.", + ) + + +def noise_inversion_inputs(comfy_io: Any, *, convenience: bool = False) -> list[Any]: + """Use the accepted two coarse plus one finishing-step recipe in both APIs.""" + enabled = ( + [ + comfy_io.Boolean.Input( + "noise_inversion_enabled", + default=False, + optional=True, + tooltip=( + "Derives starting noise from the input image; " + "adds inversion work before denoising." + ), + ) + ] + if convenience + else [] + ) + return [ + *enabled, + comfy_io.Combo.Input( + "inversion_method", + options=list(INVERSION_METHODS), + default="euler", + optional=convenience, + tooltip=( + "Euler uses one model evaluation per inversion step; " + "Heun uses two for greater accuracy." + ), + ), + comfy_io.Float.Input( + "inversion_resolution_scale", + default=0.5, + min=0.01, + max=1.0, + step=0.05, + optional=convenience, + tooltip=( + "Scales inversion width and height; " + "0.5 uses half-sized dimensions for lower cost." + ), + ), + comfy_io.Int.Input( + "inversion_steps", + default=2, + min=1, + max=64, + optional=convenience, + tooltip=( + "Steps at the selected inversion resolution; " + "more steps cost more model evaluations." + ), + ), + comfy_io.Float.Input( + "inversion_switch_fraction", + default=0.75, + min=0.01, + max=1.0, + step=0.05, + optional=convenience, + tooltip=( + "Noise-level fraction reached before the full-resolution finish; " + "0.75 means 75%." + ), + ), + comfy_io.Int.Input( + "inversion_finishing_steps", + default=1, + min=0, + max=64, + optional=convenience, + tooltip=( + "Full-resolution inversion steps after a reduced stage; " + "0 finishes entirely at reduced size." + ), + ), + comfy_io.Combo.Input( + "inversion_finishing_method", + options=list(INVERSION_METHODS), + default="euler", + optional=convenience, + tooltip=( + "Method for full-resolution finishing; " + "Euler costs one evaluation per step, Heun two." + ), + ), + ] + + +def inversion_from_controls( + *, + noise_inversion_enabled: bool, + inversion_method: str = "euler", + inversion_resolution_scale: float = 0.5, + inversion_steps: int = 2, + inversion_switch_fraction: float = 0.75, + inversion_finishing_steps: int = 1, + inversion_finishing_method: str = "euler", +) -> NoiseInversionOptions | None: + """Use domain validation when inversion is selected.""" + if type(noise_inversion_enabled) is not bool: + raise TypeError("noise_inversion_enabled must be a boolean.") + if not noise_inversion_enabled: + return None + return NoiseInversionOptions( + method=cast(InversionMethod, inversion_method), + resolution_scale=inversion_resolution_scale, + steps=inversion_steps, + switch_fraction=inversion_switch_fraction, + finishing_steps=inversion_finishing_steps, + finishing_method=cast(InversionMethod, inversion_finishing_method), + ) diff --git a/simple_syrup/nodes_v3/tiling_options.py b/simple_syrup/nodes_v3/tiling_options.py new file mode 100644 index 0000000..a1247e5 --- /dev/null +++ b/simple_syrup/nodes_v3/tiling_options.py @@ -0,0 +1,88 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Configure the single local tiling authority for a sampler options chain.""" + +from __future__ import annotations + +from typing import Any + +from ..domain.sampler_options import ( + SamplerOptions, + TilingOptions, + append_sampler_capability, +) +from .ksampler_schema import tiled_diffusion_inputs +from .sampler_options_schema import ( + COMFY_IO, + OptionsNodeBase, + options_input, + options_output, +) + + +class TilingOptionsV3(OptionsNodeBase): + """Add bounded local tiles, blend policy and optional semantic guidance.""" + + @classmethod + def define_schema(cls) -> Any: + """Share tiled controls and expose mask-dependent denoising.""" + return COMFY_IO.Schema( + node_id="SimpleSyrup.TilingOptions", + display_name="Tiling Options", + category="SimpleSyrup/Sampling/Options", + description=( + "Samples bounded local tiles; with Contextual Diffusion, " + "configures its sole local tile plan." + ), + inputs=[ + options_input(COMFY_IO), + *tiled_diffusion_inputs(COMFY_IO), + COMFY_IO.Boolean.Input( + "differential_diffusion", + default=False, + tooltip=( + "Uses the noise mask to vary denoising strength spatially; " + "preserves existing model mask behavior." + ), + ), + COMFY_IO.SEGS.Input( + "segs", + optional=True, + tooltip=( + "Optional image regions that guide local tile ownership " + "and boundaries." + ), + ), + ], + outputs=[options_output(COMFY_IO)], + ) + + @classmethod + def execute( + cls, + diffusion_mode: str = "multidiffusion", + latent_tile_width: int = 128, + latent_tile_height: int = 128, + latent_tile_overlap: int = 16, + latent_tile_batch_size: int = 4, + differential_diffusion: bool = False, + segs: object | None = None, + options: SamplerOptions | None = None, + ) -> tuple[SamplerOptions]: + """Append validated tiling without changing the incoming chain.""" + return ( + append_sampler_capability( + options, + TilingOptions( + diffusion_mode=diffusion_mode, + width=latent_tile_width, + height=latent_tile_height, + overlap=latent_tile_overlap, + batch_size=latent_tile_batch_size, + differential_diffusion=differential_diffusion, + segs=segs, + ), + ), + ) diff --git a/simple_syrup/runtime/contextual_diffusion_sampling.py b/simple_syrup/runtime/contextual_diffusion_sampling.py index 88212bb..df4c608 100644 --- a/simple_syrup/runtime/contextual_diffusion_sampling.py +++ b/simple_syrup/runtime/contextual_diffusion_sampling.py @@ -16,16 +16,24 @@ from ..domain.contextual_diffusion import ( ContextualDiffusionControls, ContextualDiffusionPlan, ) +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_features import ( EMPTY_REGIONAL_CAPABILITY_ADMISSION, RegionalCapabilityAdmission, ) +from ..domain.sampler_options import TilingOptions +from ..domain.segs import NativeSegs from ..shared.logging import get_logger from . import sampling_noise, sampling_samplers, sampling_schedulers from .contextual_model_wrapper import ContextualDiffusionModelWrapper +from .differential_diffusion import ( + differential_diffusion_mutation, + has_denoise_mask_function, +) from .guided_sampling import sample_with_optional_negative +from .inversion_model_factory import InversionModelFactory from .model_patcher_mutations import ModelUnetWrapperMutation -from .patcher_lifecycle import PATCHER_LIFECYCLE +from .patcher_lifecycle import PATCHER_LIFECYCLE, ModelMutation from .sampling_model_types import ModelFunctionWrapper from .tiled_sampling_validation import ( Latent, @@ -58,14 +66,18 @@ def sample_contextual_diffusion( capability_admission: RegionalCapabilityAdmission = ( EMPTY_REGIONAL_CAPABILITY_ADMISSION ), + noise_inversion: NoiseInversionOptions | None = None, + inversion_segs: NativeSegs | None = None, + inversion_region_masks: torch.Tensor | None = None, + differential_diffusion: bool = False, ) -> Latent: """Sample one latent through global context and one tiled prediction plan.""" validate_sampling_controls( steps=steps, denoise=denoise, - latent_tile_width=controls.latent_context_size, - latent_tile_height=controls.latent_context_size, + latent_tile_width=controls.tile_width, + latent_tile_height=controls.tile_height, latent_tile_batch_size=controls.latent_context_batch_size, ) controls.validate() @@ -90,8 +102,8 @@ def sample_contextual_diffusion( steps=steps, denoise=denoise, view=sampling_schedulers.SchedulerView( - latent_width=controls.latent_context_size, - latent_height=controls.latent_context_size, + latent_width=controls.tile_width, + latent_height=controls.tile_height, ), ).to(model.load_device) latent_samples = validate_latent_samples(latent_image, sampler_label=SAMPLER_LABEL) @@ -114,8 +126,30 @@ def sample_contextual_diffusion( controls=controls, sigmas=sigmas, diffusion_mode=diffusion_mode, + differential_diffusion=differential_diffusion, ) batch_inds = latent_image.get("batch_index") + inversion_factory = ( + InversionModelFactory( + model=model, + canvas_width=plan.latent_width, + canvas_height=plan.latent_height, + tiling=TilingOptions( + diffusion_mode=diffusion_mode, + width=controls.tile_width, + height=controls.tile_height, + overlap=controls.latent_context_overlap, + batch_size=controls.latent_context_batch_size, + differential_diffusion=differential_diffusion, + ), + context=controls, + forward_sigmas=sigmas, + segs=inversion_segs, + region_masks=inversion_region_masks, + ) + if noise_inversion is not None + else None + ) noise = sampling_noise.prepare_sampling_noise( comfy_sample=comfy_sample, sampler_name=sampler_name, @@ -139,6 +173,8 @@ def sample_contextual_diffusion( callback=callback, disable_pbar=not comfy_utils.PROGRESS_BAR_ENABLED, seed=seed, + noise_inversion=noise_inversion, + inversion_model_factory=inversion_factory, ) LOGGER.info( @@ -178,10 +214,12 @@ def clone_model_with_contextual_diffusion( controls: ContextualDiffusionControls, sigmas: torch.Tensor, diffusion_mode: str, + differential_diffusion: bool = False, + existing_wrapper: ModelFunctionWrapper | None = None, ) -> Any: """Derive a model with one pre-CFG contextual prediction wrapper.""" - old_wrapper = model.model_options.get("model_function_wrapper") + old_wrapper = existing_wrapper or model.model_options.get("model_function_wrapper") if old_wrapper is not None and not callable(old_wrapper): raise ValueError("Existing model_function_wrapper is not callable.") wrapper = ContextualDiffusionModelWrapper( @@ -191,9 +229,13 @@ def clone_model_with_contextual_diffusion( diffusion_mode=diffusion_mode, existing_wrapper=cast(ModelFunctionWrapper | None, old_wrapper), ) + mutations: list[ModelMutation] = [] + if differential_diffusion and not has_denoise_mask_function(model): + mutations.append(differential_diffusion_mutation()) + mutations.append(ModelUnetWrapperMutation(wrapper)) return PATCHER_LIFECYCLE.derive_model( model, - (ModelUnetWrapperMutation(wrapper),), + mutations, operation="SimpleSyrup contextual diffusion", ) diff --git a/simple_syrup/runtime/guided_sampling.py b/simple_syrup/runtime/guided_sampling.py index 49f7fe3..98d6bd3 100644 --- a/simple_syrup/runtime/guided_sampling.py +++ b/simple_syrup/runtime/guided_sampling.py @@ -11,7 +11,9 @@ from typing import Any, cast import torch +from ..domain.noise_inversion import NoiseInversionOptions from ..shared.logging import get_logger +from .noise_inversion import InversionModelFactory, invert_sampling_noise LOGGER = get_logger(__name__) @@ -31,8 +33,25 @@ def sample_with_optional_negative( callback: Any = None, disable_pbar: bool = False, seed: int | None = None, + noise_inversion: NoiseInversionOptions | None = None, + inversion_model_factory: InversionModelFactory | None = None, ) -> torch.Tensor: - """Sample with CFG when negative exists or Comfy's positive-only path otherwise.""" + """Prepare optional source-derived noise and select the actual Comfy guider.""" + + if noise_inversion is not None: + inversion = invert_sampling_noise( + model=model, + latent=latent_image, + forward_sigmas=sigmas, + positive=positive, + negative=negative, + cfg=cfg, + seed=seed, + options=noise_inversion, + model_factory=inversion_model_factory, + noise_mask=noise_mask, + ) + noise = inversion.noise.to(noise) if negative is not None: return cast( diff --git a/simple_syrup/runtime/inversion_model_factory.py b/simple_syrup/runtime/inversion_model_factory.py new file mode 100644 index 0000000..bc7aeab --- /dev/null +++ b/simple_syrup/runtime/inversion_model_factory.py @@ -0,0 +1,164 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Prepare each inversion resolution from the original, spatially unwrapped MODEL.""" + +from __future__ import annotations + +from typing import Any, cast + +import torch +import torch.nn.functional as functional + +from ..domain.contextual_diffusion import ( + ContextualDiffusionControls, + build_contextual_diffusion_plan, +) +from ..domain.regional_tiled_diffusion import ( + build_region_constrained_tiled_diffusion_plan, +) +from ..domain.sampler_options import TilingOptions +from ..domain.segs import NativeSegs +from ..domain.segs_tiled_diffusion import build_segs_guided_tiled_diffusion_plan +from ..domain.tiled_diffusion import TiledDiffusionPlan, build_tiled_diffusion_plan +from .inversion_spatial_context import InversionSpatialContextWrapper +from .model_patcher_mutations import ModelUnetWrapperMutation +from .patcher_lifecycle import PATCHER_LIFECYCLE +from .sampling_model_types import ModelFunctionWrapper + + +class InversionModelFactory: + """Own stage planning while sharing the forward spatial wrapper authorities.""" + + def __init__( + self, + *, + model: Any, + canvas_width: int, + canvas_height: int, + tiling: TilingOptions | None = None, + context: ContextualDiffusionControls | None = None, + forward_sigmas: torch.Tensor | None = None, + segs: NativeSegs | None = None, + region_masks: torch.Tensor | None = None, + ) -> None: + """Retain canonical inputs, not an already wrapped full-resolution model.""" + if context is not None and (tiling is None or forward_sigmas is None): + raise ValueError( + "Contextual inversion requires tile controls and forward sigmas." + ) + self._model = model + self._width = canvas_width + self._height = canvas_height + self._tiling = tiling + self._context = context + self._sigmas = forward_sigmas + self._segs = segs + self._masks = region_masks + + def __call__(self, latent: torch.Tensor) -> Any: + """Replan one resolution with canonical regional-mask coordinates.""" + from .contextual_diffusion_sampling import clone_model_with_contextual_diffusion + from .mixture_of_diffusers_sampling import clone_model_with_mixture_of_diffusers + from .multidiffusion_sampling import clone_model_with_multidiffusion + + width, height = int(latent.shape[-1]), int(latent.shape[-2]) + old_wrapper = self._model.model_options.get("model_function_wrapper") + if old_wrapper is not None and not callable(old_wrapper): + raise TypeError("Existing model_function_wrapper must be callable.") + wrapper = cast(ModelFunctionWrapper | None, old_wrapper) + if wrapper is not None and (width, height) != (self._width, self._height): + wrapper = InversionSpatialContextWrapper( + wrapper, + canvas_width=self._width, + canvas_height=self._height, + stage_width=width, + stage_height=height, + ) + masks = self._stage_masks(height, width) + if self._context is not None: + assert self._tiling is not None and self._sigmas is not None + plan = build_contextual_diffusion_plan( + latent_width=width, + latent_height=height, + controls=self._context, + segs=self._segs, + region_masks=masks, + segs_canvas=(self._height, self._width), + ) + return clone_model_with_contextual_diffusion( + self._model, + plan=plan, + controls=self._context, + sigmas=self._sigmas, + diffusion_mode=self._tiling.diffusion_mode, + differential_diffusion=self._tiling.differential_diffusion, + existing_wrapper=wrapper, + ) + if self._tiling is not None: + tile_plan = self._tile_plan(width, height, masks) + clone = ( + clone_model_with_multidiffusion + if self._tiling.diffusion_mode == "multidiffusion" + else clone_model_with_mixture_of_diffusers + ) + derived, _ = clone( + self._model, + latent_width=width, + latent_height=height, + tile_width=self._tiling.width, + tile_height=self._tiling.height, + overlap=self._tiling.overlap, + tile_batch_size=self._tiling.batch_size, + differential_diffusion=self._tiling.differential_diffusion, + tiled_plan=tile_plan, + existing_wrapper=wrapper, + ) + return derived + if wrapper is old_wrapper: + return self._model + assert wrapper is not None + return PATCHER_LIFECYCLE.derive_model( + self._model, + (ModelUnetWrapperMutation(wrapper),), + operation="SimpleSyrup inversion spatial context", + ) + + def _stage_masks(self, height: int, width: int) -> torch.Tensor | None: + """Resize planning masks once; attention masks keep their canonical bank.""" + if self._masks is None: + return None + if tuple(self._masks.shape[-2:]) == (height, width): + return self._masks + return functional.interpolate( + self._masks.unsqueeze(1).float(), size=(height, width), mode="nearest" + ).squeeze(1) + + def _tile_plan( + self, width: int, height: int, masks: torch.Tensor | None + ) -> TiledDiffusionPlan: + """Share regular, SEGS and regional ownership with forward sampling.""" + assert self._tiling is not None + geometry = { + "latent_width": width, + "latent_height": height, + "tile_width": self._tiling.width, + "tile_height": self._tiling.height, + "overlap": self._tiling.overlap, + "tile_batch_size": self._tiling.batch_size, + } + if masks is not None: + return build_region_constrained_tiled_diffusion_plan( + region_masks=masks, + segs=self._segs, + segs_canvas=(self._height, self._width), + **geometry, + ) + if self._segs is not None: + return build_segs_guided_tiled_diffusion_plan( + segs=self._segs, + segs_canvas=(self._height, self._width), + **geometry, + ) + return build_tiled_diffusion_plan(**geometry) diff --git a/simple_syrup/runtime/inversion_spatial_context.py b/simple_syrup/runtime/inversion_spatial_context.py new file mode 100644 index 0000000..ccf8c2c --- /dev/null +++ b/simple_syrup/runtime/inversion_spatial_context.py @@ -0,0 +1,128 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Project reduced inversion views into the original regional-attention canvas.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from ..domain.spatial_views import SpatialBatchLayout, SpatialView, SpatialViewKind +from .sampling_model_types import ApplyModel, ModelFunctionWrapper +from .spatial_model_arguments import ( + SIMPLE_SYRUP_TRANSFORMER_NAMESPACE, + SPATIAL_BATCH_LAYOUT_KEY, +) + + +def rebase_inversion_layout( + layout: SpatialBatchLayout, *, canvas_width: int, canvas_height: int +) -> SpatialBatchLayout: + """Keep actual model dimensions while mapping source rectangles to full masks.""" + scale_x = canvas_width / layout.canvas_width + scale_y = canvas_height / layout.canvas_height + views: list[SpatialView] = [] + for view in layout.views: + left = round(view.source_x * scale_x) + top = round(view.source_y * scale_y) + right = round(view.source_right * scale_x) + bottom = round(view.source_bottom * scale_y) + views.append( + SpatialView( + kind=( + SpatialViewKind.CONTEXTUAL_GLOBAL + if view.kind is SpatialViewKind.FULL + else view.kind + ), + source_x=left, + source_y=top, + source_width=right - left, + source_height=bottom - top, + model_width=view.model_width, + model_height=view.model_height, + ) + ) + return SpatialBatchLayout( + canvas_width, canvas_height, tuple(views), layout.input_batch_size + ) + + +class InversionSpatialContextWrapper: + """Preserve regional mask coordinates without resizing conditioning twice.""" + + def __init__( + self, + existing_wrapper: ModelFunctionWrapper, + *, + canvas_width: int, + canvas_height: int, + stage_width: int, + stage_height: int, + ) -> None: + """Bind one stage's geometry to the original attention-mask canvas.""" + self._existing_wrapper = existing_wrapper + self._canvas_width = canvas_width + self._canvas_height = canvas_height + self._stage_width = stage_width + self._stage_height = stage_height + + def __call__(self, apply_model: ApplyModel, args: dict[str, Any]) -> torch.Tensor: + """Replace only layout metadata, preserving expanded CFG batch metadata.""" + x = args.get("input") + c = args.get("c", {}) + if not isinstance(x, torch.Tensor) or not isinstance(c, dict): + raise TypeError( + "Inversion spatial context requires tensor input and conditioning." + ) + transformer_options = c.get("transformer_options", {}) + if not isinstance(transformer_options, dict): + raise TypeError("Inversion transformer_options must be a dictionary.") + namespace = transformer_options.get(SIMPLE_SYRUP_TRANSFORMER_NAMESPACE, {}) + if not isinstance(namespace, dict): + raise TypeError( + "Inversion SimpleSyrup transformer namespace must be a dictionary." + ) + layout = namespace.get(SPATIAL_BATCH_LAYOUT_KEY) + if layout is None: + if tuple(x.shape[-2:]) != (self._stage_height, self._stage_width): + raise ValueError( + "Reduced inversion calls require explicit spatial layout." + ) + layout = SpatialBatchLayout( + self._stage_width, + self._stage_height, + ( + SpatialView( + SpatialViewKind.FULL, + 0, + 0, + self._stage_width, + self._stage_height, + self._stage_width, + self._stage_height, + ), + ), + int(x.shape[0]), + ) + if not isinstance(layout, SpatialBatchLayout): + raise TypeError("Inversion spatial layout has an invalid type.") + if (layout.canvas_width, layout.canvas_height) != ( + self._stage_width, + self._stage_height, + ) or layout.expanded_batch_size != int(x.shape[0]): + raise ValueError( + "Inversion layout must describe the current stage and model batch." + ) + rebased = rebase_inversion_layout( + layout, canvas_width=self._canvas_width, canvas_height=self._canvas_height + ) + projected_namespace = {**namespace, SPATIAL_BATCH_LAYOUT_KEY: rebased} + projected_options = { + **transformer_options, + SIMPLE_SYRUP_TRANSFORMER_NAMESPACE: projected_namespace, + } + projected_args = {**args, "c": {**c, "transformer_options": projected_options}} + return self._existing_wrapper(apply_model, projected_args) diff --git a/simple_syrup/runtime/mixture_of_diffusers_sampling.py b/simple_syrup/runtime/mixture_of_diffusers_sampling.py index a79d363..ff3ba3d 100644 --- a/simple_syrup/runtime/mixture_of_diffusers_sampling.py +++ b/simple_syrup/runtime/mixture_of_diffusers_sampling.py @@ -16,10 +16,13 @@ from typing import Any, cast import torch +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_features import ( EMPTY_REGIONAL_CAPABILITY_ADMISSION, RegionalCapabilityAdmission, ) +from ..domain.sampler_options import TilingOptions +from ..domain.segs import NativeSegs from ..domain.tiled_diffusion import ( TiledDiffusionPlan, build_tiled_diffusion_plan, @@ -32,6 +35,7 @@ from .differential_diffusion import ( has_denoise_mask_function, ) from .guided_sampling import sample_with_optional_negative +from .inversion_model_factory import InversionModelFactory from .model_patcher_mutations import ModelUnetWrapperMutation from .patcher_lifecycle import PATCHER_LIFECYCLE, ModelMutation from .sampling_model_types import ( @@ -73,6 +77,9 @@ def sample_mixture_of_diffusers( EMPTY_REGIONAL_CAPABILITY_ADMISSION ), tiled_plan: TiledDiffusionPlan | None = None, + noise_inversion: NoiseInversionOptions | None = None, + inversion_segs: NativeSegs | None = None, + inversion_region_masks: torch.Tensor | None = None, ) -> Latent: """Sample a latent with a cloned model patched for Mixture of Diffusers.""" @@ -159,6 +166,26 @@ def sample_mixture_of_diffusers( callback=callback, disable_pbar=not comfy_utils.PROGRESS_BAR_ENABLED, seed=seed, + noise_inversion=noise_inversion, + inversion_model_factory=( + InversionModelFactory( + model=model, + canvas_width=latent_width, + canvas_height=latent_height, + tiling=TilingOptions( + diffusion_mode="mixture_of_diffusers", + width=latent_tile_width, + height=latent_tile_height, + overlap=latent_tile_overlap, + batch_size=latent_tile_batch_size, + differential_diffusion=differential_diffusion, + ), + segs=inversion_segs, + region_masks=inversion_region_masks, + ) + if noise_inversion is not None + else None + ), ) LOGGER.info( @@ -197,6 +224,7 @@ def clone_model_with_mixture_of_diffusers( tile_batch_size: int, differential_diffusion: bool = False, tiled_plan: TiledDiffusionPlan | None = None, + existing_wrapper: ModelFunctionWrapper | None = None, ) -> tuple[Any, TiledDiffusionPlan]: """Return a derived model patched with a pre-CFG Mixture wrapper.""" @@ -209,7 +237,7 @@ def clone_model_with_mixture_of_diffusers( tile_batch_size=tile_batch_size, ) _validate_supplied_plan(plan, latent_width, latent_height) - old_wrapper = model.model_options.get("model_function_wrapper") + old_wrapper = existing_wrapper or model.model_options.get("model_function_wrapper") if old_wrapper is not None and not callable(old_wrapper): raise ValueError("Existing model_function_wrapper is not callable.") diff --git a/simple_syrup/runtime/multidiffusion_sampling.py b/simple_syrup/runtime/multidiffusion_sampling.py index 6053987..fe52c2e 100644 --- a/simple_syrup/runtime/multidiffusion_sampling.py +++ b/simple_syrup/runtime/multidiffusion_sampling.py @@ -16,10 +16,13 @@ from typing import Any, cast import torch +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_features import ( EMPTY_REGIONAL_CAPABILITY_ADMISSION, RegionalCapabilityAdmission, ) +from ..domain.sampler_options import TilingOptions +from ..domain.segs import NativeSegs from ..domain.tiled_diffusion import ( TiledDiffusionPlan, build_tiled_diffusion_plan, @@ -32,6 +35,7 @@ from .differential_diffusion import ( has_denoise_mask_function, ) from .guided_sampling import sample_with_optional_negative +from .inversion_model_factory import InversionModelFactory from .model_patcher_mutations import ModelUnetWrapperMutation from .patcher_lifecycle import PATCHER_LIFECYCLE, ModelMutation from .sampling_model_types import ( @@ -74,6 +78,9 @@ def sample_multidiffusion( EMPTY_REGIONAL_CAPABILITY_ADMISSION ), tiled_plan: TiledDiffusionPlan | None = None, + noise_inversion: NoiseInversionOptions | None = None, + inversion_segs: NativeSegs | None = None, + inversion_region_masks: torch.Tensor | None = None, ) -> Latent: """Sample a latent with a cloned model patched for MultiDiffusion.""" @@ -161,6 +168,25 @@ def sample_multidiffusion( callback=callback, disable_pbar=not comfy_utils.PROGRESS_BAR_ENABLED, seed=seed, + noise_inversion=noise_inversion, + inversion_model_factory=( + InversionModelFactory( + model=model, + canvas_width=latent_width, + canvas_height=latent_height, + tiling=TilingOptions( + width=latent_tile_width, + height=latent_tile_height, + overlap=latent_tile_overlap, + batch_size=latent_tile_batch_size, + differential_diffusion=differential_diffusion, + ), + segs=inversion_segs, + region_masks=inversion_region_masks, + ) + if noise_inversion is not None + else None + ), ) LOGGER.info( @@ -200,6 +226,7 @@ def clone_model_with_multidiffusion( tile_batch_size: int, differential_diffusion: bool = False, tiled_plan: TiledDiffusionPlan | None = None, + existing_wrapper: ModelFunctionWrapper | None = None, ) -> tuple[Any, TiledDiffusionPlan]: """Return a derived model patched with a pre-CFG MultiDiffusion wrapper.""" @@ -212,7 +239,7 @@ def clone_model_with_multidiffusion( tile_batch_size=tile_batch_size, ) _validate_supplied_plan(plan, latent_width, latent_height) - old_wrapper = model.model_options.get("model_function_wrapper") + old_wrapper = existing_wrapper or model.model_options.get("model_function_wrapper") if old_wrapper is not None and not callable(old_wrapper): raise ValueError("Existing model_function_wrapper is not callable.") diff --git a/simple_syrup/runtime/noise_inversion.py b/simple_syrup/runtime/noise_inversion.py new file mode 100644 index 0000000..f0c22a8 --- /dev/null +++ b/simple_syrup/runtime/noise_inversion.py @@ -0,0 +1,325 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Derive source-dependent sampling noise with measured, uncached inversion stages.""" + +from __future__ import annotations + +import math +import time +from collections.abc import Callable +from dataclasses import dataclass +from importlib import import_module +from typing import Any, TypeAlias + +import torch + +from ..domain.inversion_solver import ( + InversionSolverEvidence, + integrate_inversion, + lift_inversion_displacement, +) +from ..domain.noise_inversion import InversionMethod, NoiseInversionOptions +from ..shared.logging import get_logger +from .spatial_tensor_projection import resize_spatial_tensor +from .tiled_sampling_validation import validate_tensor_shape + +LOGGER = get_logger(__name__) +INVERSION_START_SIGMA = 0.0001 +InversionModelFactory: TypeAlias = Callable[[torch.Tensor], Any] + + +@dataclass(frozen=True, slots=True) +class InversionStageMeasurement: + """Describe actual work and elapsed time for one inversion stage.""" + + name: str + seconds: float + latent_shape: tuple[int, ...] + steps: int + evaluations: int + + +@dataclass(frozen=True, slots=True) +class NoiseInversionResult: + """Return inferred noise with paid inversion cost and reconstruction evidence.""" + + noise: torch.Tensor + seconds: float + stages: tuple[InversionStageMeasurement, ...] + reconstruction_max_error: float + + +def _clock(device: torch.device) -> float: + """Measure completed CUDA work rather than asynchronous kernel submission.""" + if device.type == "cuda": + torch.cuda.synchronize(device) + return time.perf_counter() + + +def _source_in_model_space(model: Any, latent: torch.Tensor) -> torch.Tensor: + """Narrow host latent processing before arithmetic or model execution.""" + source = model.model.process_latent_in(latent) + if not isinstance(source, torch.Tensor) or source.shape != latent.shape: + raise ValueError( + "Model latent processing must preserve the inversion source shape." + ) + if not source.is_floating_point() or not bool(torch.isfinite(source).all()): + raise ValueError( + "Model latent processing must produce finite floating-point values." + ) + return source.detach().float().cpu() + + +def validate_inversion_target(model: Any, sigmas: torch.Tensor) -> float: + """Reject unsupported scaling and singular targets before model execution.""" + sampling_types = import_module("comfy.model_sampling") + sampling = model.get_model_object("model_sampling") + if not isinstance( + sampling, (sampling_types.CONST, sampling_types.EPS) + ) or isinstance( + sampling, (sampling_types.IMG_TO_IMG, sampling_types.IMG_TO_IMG_FLOW) + ): + raise ValueError( + "Noise inversion requires a flow or EPS-compatible image model." + ) + if sigmas.ndim != 1 or len(sigmas) < 2 or not bool(torch.isfinite(sigmas).all()): + raise ValueError("Noise inversion requires a finite forward sampling schedule.") + target = float(sigmas[0]) + if target <= INVERSION_START_SIGMA: + raise ValueError( + "Noise inversion requires a positive partial-denoise start sigma." + ) + if isinstance(sampling, sampling_types.CONST) and target >= 0.9999: + raise ValueError( + "Flow noise inversion requires denoise below the full-noise endpoint." + ) + if isinstance(sampling, sampling_types.EPS): + maximum = float(sampling.sigma_max) + if target > maximum or math.isclose(target, maximum, rel_tol=1e-5): + raise ValueError( + "Noise inversion requires partial denoise, below the model's " + "maximum sigma." + ) + return target + + +def invert_sampling_noise( + *, + model: Any, + latent: torch.Tensor, + forward_sigmas: torch.Tensor, + positive: Any, + negative: Any, + cfg: float, + seed: int | None, + options: NoiseInversionOptions, + model_factory: InversionModelFactory | None = None, + noise_mask: Any = None, +) -> NoiseInversionResult: + """Invert a source latent through Comfy's actual CFG or positive-only guider. + + A spatial sampler supplies a factory that replans each resolution from its + unwrapped prepared model. The final noise recreates the inferred endpoint + under the model's own affine noise scaling; no inference cache is used. + """ + from .guided_sampling import sample_with_optional_negative + + if not isinstance(options, NoiseInversionOptions): + raise TypeError("Noise inversion requires validated NoiseInversionOptions.") + validate_tensor_shape(latent, sampler_label="Noise Inversion") + if not latent.is_floating_point() or not bool(torch.isfinite(latent).all()): + raise ValueError( + "Noise inversion source must contain finite floating-point values." + ) + target = validate_inversion_target(model, forward_sigmas) + coarse_target = target * options.coarse_target_fraction + if coarse_target <= INVERSION_START_SIGMA: + raise ValueError( + "Noise inversion transition must exceed the initial inversion sigma." + ) + device = torch.device(model.load_device) + started = _clock(device) + sampling = model.get_model_object("model_sampling") + source = _source_in_model_space(model, latent) + phases: list[InversionStageMeasurement] = [] + comfy_sample = import_module("comfy.sample") + comfy_samplers = import_module("comfy.samplers") + + def stage( + stage_latent: torch.Tensor, + begin: float, + end: float, + count: int, + initial: torch.Tensor | None, + name: str, + method: InversionMethod, + ) -> torch.Tensor: + """Capture the model-space endpoint before Comfy converts output latents.""" + stage_started = _clock(device) + stage_model = ( + model_factory(stage_latent) if model_factory is not None else model + ) + schedule = torch.linspace(begin, end, count + 1, device=device) + evidence = InversionSolverEvidence() + endpoints: list[torch.Tensor] = [] + + def invert( + model_fn: Any, + state: torch.Tensor, + sigmas: torch.Tensor, + extra_args: dict[str, Any], + callback: Any, + disable: bool, + ) -> torch.Tensor: + """Use Comfy's denoised predictions as the inversion velocity field.""" + if initial is not None: + state = initial.to(state) + + def evaluate( + x: torch.Tensor, sigma: torch.Tensor, index: int + ) -> torch.Tensor: + """Narrow the dynamic Comfy model result before numeric integration.""" + prediction = model_fn( + x, sigma * x.new_ones((x.shape[0],)), **extra_args + ) + if not isinstance(prediction, torch.Tensor): + raise TypeError( + "Noise inversion model must return tensor predictions." + ) + return (x - prediction) / sigma + + endpoint = integrate_inversion( + state, sigmas, evaluate, method=method, evidence=evidence + ) + endpoints.append(endpoint.detach().float().cpu()) + return endpoint + + sample_with_optional_negative( + comfy_sample=comfy_sample, + model=stage_model, + noise=torch.zeros_like(stage_latent), + cfg=cfg, + sampler=comfy_samplers.KSAMPLER(invert), + sigmas=schedule, + positive=positive, + negative=negative, + latent_image=stage_latent, + noise_mask=noise_mask, + seed=seed, + disable_pbar=True, + ) + if len(endpoints) != 1: + raise RuntimeError( + "Noise inversion must produce exactly one endpoint per stage." + ) + phases.append( + InversionStageMeasurement( + name, + _clock(device) - stage_started, + tuple(stage_latent.shape), + count, + evidence.evaluations, + ) + ) + return endpoints[0] + + if options.resolution_scale == 1: + endpoint = stage( + latent, + INVERSION_START_SIGMA, + target, + options.steps, + None, + "full", + options.method, + ) + else: + height, width = options.coarse_shape( + int(latent.shape[-2]), int(latent.shape[-1]) + ) + coarse = resize_spatial_tensor(latent, height=height, width=width, mode="area") + coarse_source = _source_in_model_space(model, coarse) + coarse_endpoint = stage( + coarse, + INVERSION_START_SIGMA, + coarse_target, + options.steps, + None, + "coarse", + options.method, + ) + endpoint = lift_inversion_displacement( + source, + coarse_source, + coarse_endpoint, + resize=lambda x, h, w: resize_spatial_tensor( + x, height=h, width=w, mode="bilinear" + ), + ) + if options.finishing_steps: + endpoint = stage( + latent, + coarse_target, + target, + options.finishing_steps, + endpoint, + "full_finish", + options.finishing_method, + ) + + sigma = torch.tensor(target) + zero = torch.zeros_like(source) + base = sampling.noise_scaling(sigma, zero.clone(), source, max_denoise=False) + amplitude = sampling.noise_scaling( + sigma, torch.ones_like(source), zero, max_denoise=False + ) + if not isinstance(base, torch.Tensor) or not isinstance(amplitude, torch.Tensor): + raise TypeError("Model noise scaling must return tensors.") + if not bool(torch.isfinite(amplitude).all()) or bool(torch.any(amplitude == 0)): + raise ValueError("Model noise scaling is not invertible at the target sigma.") + noise = (endpoint - base) / amplitude + if not bool(torch.isfinite(noise).all()): + raise FloatingPointError("Noise inversion produced non-finite sampling noise.") + reconstructed = sampling.noise_scaling( + sigma, noise.clone(), source, max_denoise=False + ) + if ( + not isinstance(reconstructed, torch.Tensor) + or reconstructed.shape != endpoint.shape + ): + raise ValueError( + "Model noise scaling must preserve the inversion endpoint shape." + ) + if not torch.allclose(reconstructed, endpoint, atol=1e-5, rtol=1e-5): + raise ValueError( + "Model noise scaling cannot reconstruct the inversion endpoint." + ) + error = float((reconstructed - endpoint).abs().max()) + elapsed = _clock(device) - started + LOGGER.info( + "Noise inversion completed in %.3f seconds", + elapsed, + extra={ + "operation": "noise_inversion", + "method": options.method, + "resolution_scale": options.resolution_scale, + "steps": options.steps, + "finishing_steps": options.finishing_steps, + "inversion_seconds": elapsed, + "inversion_stages": [ + { + "name": phase.name, + "seconds": phase.seconds, + "latent_shape": phase.latent_shape, + "steps": phase.steps, + "evaluations": phase.evaluations, + } + for phase in phases + ], + "evaluations": sum(phase.evaluations for phase in phases), + "endpoint_reconstruction_max_error": error, + }, + ) + return NoiseInversionResult(noise, elapsed, tuple(phases), error) diff --git a/simple_syrup/runtime/regional_detail_sampler.py b/simple_syrup/runtime/regional_detail_sampler.py index b8113ac..cf3e9fa 100644 --- a/simple_syrup/runtime/regional_detail_sampler.py +++ b/simple_syrup/runtime/regional_detail_sampler.py @@ -10,6 +10,7 @@ from typing import Any import torch +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_detailing import LatentRegion from . import regional_multidiffusion_sampling from .detail_previews import DetailPreviewContext @@ -51,6 +52,7 @@ class RegionalDetailSampler: global_prompt_weight: float, preview_context: DetailPreviewContext | None = None, differential_diffusion: bool = False, + noise_inversion: NoiseInversionOptions | None = None, ) -> Latent: """Sample one full latent with regional MultiDiffusion.""" @@ -69,4 +71,5 @@ class RegionalDetailSampler: global_prompt_weight=global_prompt_weight, preview_context=preview_context, differential_diffusion=differential_diffusion, + noise_inversion=noise_inversion, ) diff --git a/simple_syrup/runtime/regional_inversion_model_factory.py b/simple_syrup/runtime/regional_inversion_model_factory.py new file mode 100644 index 0000000..fb274d0 --- /dev/null +++ b/simple_syrup/runtime/regional_inversion_model_factory.py @@ -0,0 +1,66 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Rebuild regional prediction ownership from canonical inputs for inversion stages.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from ..domain.regional_detailing import LatentRegion +from ..domain.regional_inversion_geometry import project_inversion_regions +from .inversion_model_factory import InversionModelFactory + + +class RegionalInversionModelFactory: + """Retain one original MODEL and the full-resolution region bank.""" + + def __init__( + self, + *, + model: Any, + canvas_width: int, + canvas_height: int, + regions: tuple[LatentRegion, ...], + global_prompt_weight: float, + differential_diffusion: bool, + ) -> None: + """Keep regional conditioning unchanged across source-sized inversion views.""" + self._base = InversionModelFactory( + model=model, + canvas_width=canvas_width, + canvas_height=canvas_height, + ) + self._width = canvas_width + self._height = canvas_height + self._regions = regions + self._weight = global_prompt_weight + self._differential = differential_diffusion + + def __call__(self, latent: torch.Tensor) -> Any: + """Use the authoritative regional calc-cond-batch wrapper at each stage size.""" + from .regional_multidiffusion_sampling import ( + clone_model_with_regional_multidiffusion, + ) + + height, width = int(latent.shape[-2]), int(latent.shape[-1]) + regions = project_inversion_regions( + self._regions, + source_width=self._width, + source_height=self._height, + target_width=width, + target_height=height, + ) + derived, _ = clone_model_with_regional_multidiffusion( + self._base(latent), + latent_width=width, + latent_height=height, + latent_ndim=latent.ndim, + regions=regions, + global_prompt_weight=self._weight, + differential_diffusion=self._differential, + ) + return derived diff --git a/simple_syrup/runtime/regional_multidiffusion_sampling.py b/simple_syrup/runtime/regional_multidiffusion_sampling.py index 765fa87..3c1111a 100644 --- a/simple_syrup/runtime/regional_multidiffusion_sampling.py +++ b/simple_syrup/runtime/regional_multidiffusion_sampling.py @@ -15,6 +15,7 @@ from importlib import import_module from types import ModuleType from typing import Any, cast +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_detailing import LatentRegion from ..domain.regional_features import EMPTY_REGIONAL_CAPABILITY_ADMISSION from ..shared.logging import get_logger @@ -27,6 +28,7 @@ from .differential_diffusion import ( from .guided_sampling import sample_with_optional_negative from .model_patcher_mutations import ModelCalcCondBatchMutation from .patcher_lifecycle import PATCHER_LIFECYCLE, ModelMutation +from .regional_inversion_model_factory import RegionalInversionModelFactory from .regional_multidiffusion_prediction import ( CalcCondBatchFunction, RegionalMultiDiffusionCalcCondBatch, @@ -72,6 +74,7 @@ def sample_regional_multidiffusion( global_prompt_weight: float, preview_context: DetailPreviewContext | None = None, differential_diffusion: bool = False, + noise_inversion: NoiseInversionOptions | None = None, ) -> Latent: """Sample a latent with regional MultiDiffusion prompt blending.""" @@ -159,6 +162,19 @@ def sample_regional_multidiffusion( callback=callback, disable_pbar=not comfy_utils.PROGRESS_BAR_ENABLED, seed=seed, + noise_inversion=noise_inversion, + inversion_model_factory=( + RegionalInversionModelFactory( + model=model, + canvas_width=latent_width, + canvas_height=latent_height, + regions=regions, + global_prompt_weight=global_prompt_weight, + differential_diffusion=differential_diffusion, + ) + if noise_inversion is not None + else None + ), ) LOGGER.info( diff --git a/simple_syrup/services/attention_coupling_sampling_service.py b/simple_syrup/services/attention_coupling_sampling_service.py index cd73a93..9b6ac6a 100644 --- a/simple_syrup/services/attention_coupling_sampling_service.py +++ b/simple_syrup/services/attention_coupling_sampling_service.py @@ -12,6 +12,7 @@ from ..domain.attention_coupling_request import ( AttentionCouplingRequestMode, classify_attention_coupling_request, ) +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_attention_execution import RegionalAttentionExecutionMode from .attention_coupling_model_preparation_service import ( AttentionCouplingModelPreparationService, @@ -45,6 +46,7 @@ class AttentionCouplingSamplingService: region_mask_feather: int, latent_image: dict[str, Any], denoise: float, + noise_inversion: NoiseInversionOptions | None = None, ) -> dict[str, Any]: """Bypass ordinary requests or prepare one complete regional request.""" @@ -65,6 +67,7 @@ class AttentionCouplingSamplingService: negative=negative, latent_image=latent_image, denoise=denoise, + noise_inversion=noise_inversion, ) prepared = self.model_preparation_service_class().prepare( @@ -88,6 +91,7 @@ class AttentionCouplingSamplingService: negative=prepared.negative, latent_image=latent_image, denoise=denoise, + noise_inversion=noise_inversion, ) diff --git a/simple_syrup/services/contextual_attention_coupling_sampling_service.py b/simple_syrup/services/contextual_attention_coupling_sampling_service.py index 72dd118..d9bbba1 100644 --- a/simple_syrup/services/contextual_attention_coupling_sampling_service.py +++ b/simple_syrup/services/contextual_attention_coupling_sampling_service.py @@ -8,8 +8,10 @@ from __future__ import annotations from typing import Any, ClassVar +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_attention_execution import RegionalAttentionExecutionMode from ..domain.regional_features import RegionalFeature, RegionalFeatureRequest +from ..domain.sampler_options import TilingOptions from .attention_coupling_model_preparation_service import ( AttentionCouplingModelPreparationService, ) @@ -57,6 +59,8 @@ class ContextualAttentionCouplingSamplingService: global_steps: int, global_decay: float, segs: object | None = None, + noise_inversion: NoiseInversionOptions | None = None, + tiling: TilingOptions | None = None, ) -> ContextualDiffusionSamplingResult: """Prepare once and invoke established Contextual local-view execution.""" @@ -94,6 +98,8 @@ class ContextualAttentionCouplingSamplingService: region_mask_feather=0, feature_request=_CONTEXTUAL_ATTENTION_REQUEST, planning_region_masks=prepared.mask_bank.planning_masks, + noise_inversion=noise_inversion, + tiling=tiling, ) diff --git a/simple_syrup/services/contextual_diffusion_sampling_service.py b/simple_syrup/services/contextual_diffusion_sampling_service.py index bac5258..15d88e6 100644 --- a/simple_syrup/services/contextual_diffusion_sampling_service.py +++ b/simple_syrup/services/contextual_diffusion_sampling_service.py @@ -21,6 +21,7 @@ from ..domain.contextual_diffusion import ( ContextualDiffusionControls, build_contextual_diffusion_plan, ) +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_features import ( CONTEXTUAL_DIFFUSION_REGIONAL_SAMPLER_CAPABILITIES, EMPTY_REGIONAL_FEATURE_REQUEST, @@ -28,6 +29,7 @@ from ..domain.regional_features import ( RegionalFeature, RegionalFeatureRequest, ) +from ..domain.sampler_options import TilingOptions from ..domain.segs import NativeSegs, coerce_segs_group from ..domain.tiled_diffusion import validate_tiled_diffusion_mode from ..runtime.contextual_diffusion_sampling import sample_contextual_diffusion @@ -91,10 +93,16 @@ class ContextualDiffusionSamplingService: region_mask_feather: int = 0, feature_request: RegionalFeatureRequest = EMPTY_REGIONAL_FEATURE_REQUEST, planning_region_masks: torch.Tensor | None = None, + noise_inversion: NoiseInversionOptions | None = None, + tiling: TilingOptions | None = None, ) -> ContextualDiffusionSamplingResult: - """Sample a latent with global and bounded detail contexts.""" + """Sample one local tile authority with global context and inversion.""" validate_tiled_diffusion_mode(diffusion_mode) + if tiling is not None: + diffusion_mode = tiling.diffusion_mode + latent_context_overlap = tiling.overlap + latent_context_batch_size = tiling.batch_size controls = ContextualDiffusionControls( latent_context_size=latent_context_size, latent_context_overlap=latent_context_overlap, @@ -102,6 +110,8 @@ class ContextualDiffusionSamplingService: global_weight=global_weight, global_steps=global_steps, global_decay=global_decay, + latent_tile_width=tiling.width if tiling is not None else None, + latent_tile_height=tiling.height if tiling is not None else None, ) controls.validate() regional = self.regional_preparation_service_class().prepare( @@ -172,6 +182,10 @@ class ContextualDiffusionSamplingService: image_width=image_width, region_masks=planning_masks, capability_admission=capability_admission, + noise_inversion=noise_inversion, + differential_diffusion=tiling.differential_diffusion + if tiling is not None + else False, ) outputs: list[torch.Tensor] = [] @@ -207,6 +221,10 @@ class ContextualDiffusionSamplingService: image_width=image_width, region_masks=planning_masks, capability_admission=capability_admission, + noise_inversion=noise_inversion, + differential_diffusion=tiling.differential_diffusion + if tiling is not None + else False, ) samples = item_result.latent.get("samples") if not isinstance(samples, torch.Tensor): @@ -240,8 +258,10 @@ class ContextualDiffusionSamplingService: image_width: int, region_masks: torch.Tensor | None, capability_admission: RegionalCapabilityAdmission, + noise_inversion: NoiseInversionOptions | None, + differential_diffusion: bool, ) -> ContextualDiffusionSamplingResult: - """Build one canvas plan and execute it through the runtime adapter.""" + """Build the forward plan and retain canonical ownership for inversion.""" samples = latent_image.get("samples") if not isinstance(samples, torch.Tensor): @@ -270,6 +290,10 @@ class ContextualDiffusionSamplingService: controls=controls, plan=plan, capability_admission=capability_admission, + noise_inversion=noise_inversion, + inversion_segs=segs, + inversion_region_masks=region_masks, + differential_diffusion=differential_diffusion, ) return ContextualDiffusionSamplingResult( latent=latent, diff --git a/simple_syrup/services/detail_segs_as_regions_service.py b/simple_syrup/services/detail_segs_as_regions_service.py index f178adb..843994f 100644 --- a/simple_syrup/services/detail_segs_as_regions_service.py +++ b/simple_syrup/services/detail_segs_as_regions_service.py @@ -16,6 +16,7 @@ from typing import Any, Protocol import torch from ..domain.conditioning_batch import ConditioningBatch +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_detailing import ( LatentRegion, pair_segments_with_conditioning, @@ -65,6 +66,7 @@ class RegionalDetailSamplingBoundary(Protocol): global_prompt_weight: float, preview_context: DetailPreviewContext | None = None, differential_diffusion: bool = False, + noise_inversion: NoiseInversionOptions | None = None, ) -> Latent: """Sample one full latent with paired regional conditioning.""" @@ -133,6 +135,7 @@ class DetailSEGSAsRegionsService: tiled_encode: bool, tiled_decode: bool, global_prompt_weight: float, + noise_inversion: NoiseInversionOptions | None = None, ) -> DetailSEGSAsRegionsResult: """Run regional MultiDiffusion detailing for provided SEGS.""" @@ -234,6 +237,7 @@ class DetailSEGSAsRegionsService: sampled_region=CropRegion(0, 0, image_width, image_height), ), differential_diffusion=differential_diffusion, + noise_inversion=noise_inversion, ) decoded = self._sampler.decode(vae, sampled, tiled_decode) if decoded.shape[1:3] != image_tensor.shape[1:3]: diff --git a/simple_syrup/services/detail_segs_by_scale_factor_tiled_diffusion_service.py b/simple_syrup/services/detail_segs_by_scale_factor_tiled_diffusion_service.py index 9765d29..e943b61 100644 --- a/simple_syrup/services/detail_segs_by_scale_factor_tiled_diffusion_service.py +++ b/simple_syrup/services/detail_segs_by_scale_factor_tiled_diffusion_service.py @@ -13,6 +13,7 @@ import torch from ..domain.conditioning_batch import select_conditioning from ..domain.detail_geometry import DetailScalePlan, build_detail_scale_plan +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.segs import Segment, coerce_segs from ..domain.segs_mask_ops import ( crop_image, @@ -59,6 +60,7 @@ class TiledDetailSamplingBoundary(Protocol): latent_tile_batch_size: int, preview_context: DetailPreviewContext | None = None, differential_diffusion: bool = False, + noise_inversion: NoiseInversionOptions | None = None, ) -> Latent: """Sample one latent crop with the requested tiled diffusion mode.""" @@ -131,6 +133,7 @@ class DetailSEGSByScaleFactorTiledDiffusionService: latent_tile_height: int, latent_tile_overlap: int, latent_tile_batch_size: int, + noise_inversion: NoiseInversionOptions | None = None, ) -> TiledDetailerResult: """Run crop sampling and composite-back detailing with tiled diffusion.""" @@ -198,6 +201,7 @@ class DetailSEGSByScaleFactorTiledDiffusionService: latent_tile_overlap=latent_tile_overlap, latent_tile_batch_size=latent_tile_batch_size, differential_diffusion=differential_diffusion, + noise_inversion=noise_inversion, ) LOGGER.info( @@ -246,6 +250,7 @@ class DetailSEGSByScaleFactorTiledDiffusionService: latent_tile_overlap: int, latent_tile_batch_size: int, differential_diffusion: bool, + noise_inversion: NoiseInversionOptions | None, ) -> torch.Tensor: """Detail one segment with tiled diffusion and return the updated image.""" @@ -281,6 +286,7 @@ class DetailSEGSByScaleFactorTiledDiffusionService: latent_tile_overlap=latent_tile_overlap, latent_tile_batch_size=latent_tile_batch_size, differential_diffusion=differential_diffusion, + noise_inversion=noise_inversion, preview_context=DetailPreviewContext( image=working_image, work_region=segment.crop_region, diff --git a/simple_syrup/services/ksampler_sampling_service.py b/simple_syrup/services/ksampler_sampling_service.py index df1a176..a8aefa8 100644 --- a/simple_syrup/services/ksampler_sampling_service.py +++ b/simple_syrup/services/ksampler_sampling_service.py @@ -12,9 +12,11 @@ from typing import Any, TypeAlias import torch from ..domain.conditioning_batch import ConditioningBatch, select_conditioning +from ..domain.noise_inversion import NoiseInversionOptions 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 ..shared.logging import get_logger Latent: TypeAlias = dict[str, Any] @@ -37,8 +39,9 @@ class KSamplerSamplingService: negative: Any, latent_image: Latent, denoise: float, + noise_inversion: NoiseInversionOptions | None = None, ) -> Latent: - """Sample a latent with configured SimpleSyrup sampler extensions.""" + """Sample full latents with optional inversion and per-item conditioning.""" sampler = sampling_samplers.resolve_sampler(sampler_name) latent_samples = latent_image["samples"] @@ -93,6 +96,7 @@ class KSamplerSamplingService: callback=callback, disable_pbar=disable_pbar, seed=seed, + noise_inversion=noise_inversion, ) else: samples = sample_with_optional_negative( @@ -109,6 +113,14 @@ class KSamplerSamplingService: callback=callback, disable_pbar=disable_pbar, seed=seed, + noise_inversion=noise_inversion, + inversion_model_factory=InversionModelFactory( + model=model, + canvas_width=int(latent_samples.shape[-1]), + canvas_height=int(latent_samples.shape[-2]), + ) + if noise_inversion is not None + else None, ) if not isinstance(samples, torch.Tensor): raise TypeError("KSampler output samples must be a torch.Tensor.") @@ -147,6 +159,7 @@ class KSamplerSamplingService: callback: Any, disable_pbar: bool, seed: int, + noise_inversion: NoiseInversionOptions | None, ) -> torch.Tensor: """Sample latent items with existing per-item batch selection.""" @@ -171,6 +184,14 @@ class KSamplerSamplingService: callback=callback, disable_pbar=disable_pbar, seed=seed, + noise_inversion=noise_inversion, + inversion_model_factory=InversionModelFactory( + model=model, + canvas_width=int(latent_samples.shape[-1]), + canvas_height=int(latent_samples.shape[-2]), + ) + if noise_inversion is not None + else None, ) ) return torch.cat(sampled, dim=0) diff --git a/simple_syrup/services/regional_tiled_diffusion_sampling_service.py b/simple_syrup/services/regional_tiled_diffusion_sampling_service.py index d5f81fe..63513c2 100644 --- a/simple_syrup/services/regional_tiled_diffusion_sampling_service.py +++ b/simple_syrup/services/regional_tiled_diffusion_sampling_service.py @@ -10,6 +10,7 @@ from typing import Any, TypeAlias import torch +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_features import RegionalCapabilityAdmission from ..domain.regional_tiled_diffusion import ( build_region_constrained_tiled_diffusion_plan, @@ -49,6 +50,7 @@ class RegionalTiledDiffusionSamplingService: preview_context: DetailPreviewContext | None, differential_diffusion: bool, capability_admission: RegionalCapabilityAdmission, + noise_inversion: NoiseInversionOptions | None = None, ) -> Latent: """Sample every latent item with one shared regional composition.""" @@ -106,6 +108,13 @@ class RegionalTiledDiffusionSamplingService: differential_diffusion=differential_diffusion, capability_admission=capability_admission, tiled_plan=plan, + noise_inversion=noise_inversion, + inversion_segs=( + segs_group[0 if len(segs_group) == 1 else index] + if segs_group + else None + ), + inversion_region_masks=region_masks, ) output_samples = output.get("samples") if not isinstance(output_samples, torch.Tensor): diff --git a/simple_syrup/services/sampler_options_sampling_service.py b/simple_syrup/services/sampler_options_sampling_service.py new file mode 100644 index 0000000..a1cbb85 --- /dev/null +++ b/simple_syrup/services/sampler_options_sampling_service.py @@ -0,0 +1,218 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Compile capability configuration into one authoritative sampling execution.""" + +from __future__ import annotations + +from typing import Any, TypedDict + +import torch + +from ..domain.sampler_options import SamplerOptions, TilingOptions +from ..runtime import sampling_schedulers +from ..runtime.noise_inversion import validate_inversion_target +from ..runtime.tiled_sampling_validation import validate_sampling_controls +from .attention_coupling_sampling_service import AttentionCouplingSamplingService +from .contextual_attention_coupling_sampling_service import ( + ContextualAttentionCouplingSamplingService, +) +from .contextual_diffusion_sampling_service import ContextualDiffusionSamplingService +from .ksampler_sampling_service import KSamplerSamplingService +from .tiled_attention_coupling_sampling_service import ( + TiledAttentionCouplingSamplingService, +) +from .tiled_diffusion_sampling_service import TiledDiffusionSamplingService + + +class SamplingArguments(TypedDict): + """Narrow controls while retaining dynamic host MODEL and tensor payloads.""" + + model: Any + seed: int + steps: int + cfg: float + sampler_name: str + scheduler: str + positive: object + negative: object + latent_image: dict[str, Any] + denoise: float + + +class ContextualArguments(TypedDict): + """Describe global context and the sole local tile plan's execution controls.""" + + diffusion_mode: str + latent_context_size: int + latent_context_overlap: int + latent_context_batch_size: int + global_weight: float + global_steps: int + global_decay: float + + +class TiledArguments(TypedDict): + """Describe local geometry and mask-dependent denoising policy.""" + + diffusion_mode: str + latent_tile_width: int + latent_tile_height: int + latent_tile_overlap: int + latent_tile_batch_size: int + differential_diffusion: bool + + +class SamplerOptionsSamplingService: + """Route configuration independently of connection order or node placement.""" + + def sample( + 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, + ) -> dict[str, Any]: + """Apply regional attention once, then one spatial authority and inversion.""" + if options is not None and not isinstance(options, SamplerOptions): + raise TypeError( + "KSampler options must come from SimpleSyrup options nodes." + ) + configured = options if options is not None else SamplerOptions() + context = configured.contextual_diffusion + tiling = configured.tiling + if context is not None and tiling is None: + tiling = TilingOptions( + width=context.context_size, + height=context.context_size, + overlap=min(32, context.context_size - 1), + ) + arguments: SamplingArguments = { + "model": model, + "seed": seed, + "steps": steps, + "cfg": cfg, + "sampler_name": sampler_name, + "scheduler": scheduler, + "positive": positive, + "negative": negative, + "latent_image": latent_image, + "denoise": denoise, + } + self._preflight(arguments, configured, tiling) + attention = configured.attention_coupling + inversion = configured.noise_inversion + if context is not None: + assert tiling is not None + contextual_arguments: ContextualArguments = { + "diffusion_mode": tiling.diffusion_mode, + "latent_context_size": context.context_size, + "latent_context_overlap": tiling.overlap, + "latent_context_batch_size": tiling.batch_size, + "global_weight": context.global_weight, + "global_steps": context.global_steps, + "global_decay": context.global_decay, + } + if attention is not None: + result = ContextualAttentionCouplingSamplingService().sample( + **arguments, + **contextual_arguments, + region_masks=attention.region_masks, + regional_prompt_weight=attention.regional_prompt_weight, + region_mask_feather=attention.region_mask_feather, + segs=tiling.segs, + tiling=tiling, + noise_inversion=inversion, + ) + else: + result = ContextualDiffusionSamplingService().sample( + **arguments, + **contextual_arguments, + segs=tiling.segs, + tiling=tiling, + noise_inversion=inversion, + ) + return result.latent + if tiling is not None: + tiled_arguments: TiledArguments = { + "diffusion_mode": tiling.diffusion_mode, + "latent_tile_width": tiling.width, + "latent_tile_height": tiling.height, + "latent_tile_overlap": tiling.overlap, + "latent_tile_batch_size": tiling.batch_size, + "differential_diffusion": tiling.differential_diffusion, + } + if attention is not None: + return TiledAttentionCouplingSamplingService().sample( + **arguments, + **tiled_arguments, + region_masks=attention.region_masks, + regional_prompt_weight=attention.regional_prompt_weight, + region_mask_feather=attention.region_mask_feather, + segs=tiling.segs, + noise_inversion=inversion, + ) + return TiledDiffusionSamplingService().sample( + **arguments, + **tiled_arguments, + segs=tiling.segs, + noise_inversion=inversion, + ) + if attention is not None: + return AttentionCouplingSamplingService().sample( + **arguments, + region_masks=attention.region_masks, + regional_prompt_weight=attention.regional_prompt_weight, + region_mask_feather=attention.region_mask_feather, + noise_inversion=inversion, + ) + return KSamplerSamplingService().sample(**arguments, noise_inversion=inversion) + + def _preflight( + self, + arguments: SamplingArguments, + options: SamplerOptions, + tiling: TilingOptions | None, + ) -> None: + """Reject unsupported schedules and endpoints before preparing models.""" + samples = arguments["latent_image"].get("samples") + if not isinstance(samples, torch.Tensor): + raise TypeError("KSampler latent samples must be a torch.Tensor.") + validate_sampling_controls( + steps=arguments["steps"], + denoise=arguments["denoise"], + latent_tile_width=tiling.width if tiling is not None else 16, + latent_tile_height=tiling.height if tiling is not None else 16, + latent_tile_batch_size=tiling.batch_size if tiling is not None else 1, + ) + incompatible_unipc = options.contextual_diffusion is not None or ( + tiling is not None and tiling.diffusion_mode == "multidiffusion" + ) + if incompatible_unipc and arguments["sampler_name"] in {"uni_pc", "uni_pc_bh2"}: + raise ValueError( + "Tiling and Contextual Diffusion do not support UniPC samplers." + ) + if options.noise_inversion is not None: + view = ( + sampling_schedulers.SchedulerView(tiling.width, tiling.height) + if tiling is not None + else sampling_schedulers.SchedulerView.from_tensor(samples) + ) + sigmas = sampling_schedulers.calculate_sigmas( + model=arguments["model"], + scheduler_name=arguments["scheduler"], + sampler_name=arguments["sampler_name"], + steps=arguments["steps"], + denoise=arguments["denoise"], + view=view, + ) + validate_inversion_target(arguments["model"], sigmas) diff --git a/simple_syrup/services/segs_guided_tiled_diffusion_sampling_service.py b/simple_syrup/services/segs_guided_tiled_diffusion_sampling_service.py index b55c3ac..bb51dbf 100644 --- a/simple_syrup/services/segs_guided_tiled_diffusion_sampling_service.py +++ b/simple_syrup/services/segs_guided_tiled_diffusion_sampling_service.py @@ -11,6 +11,7 @@ from typing import Any, TypeAlias import torch from ..domain.conditioning_batch import ConditioningBatch, select_conditioning +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_features import RegionalCapabilityAdmission from ..domain.segs import coerce_segs_group from ..domain.segs_tiled_diffusion import build_segs_guided_tiled_diffusion_plan @@ -47,6 +48,7 @@ class SEGSGuidedTiledDiffusionSamplingService: preview_context: DetailPreviewContext | None, differential_diffusion: bool, capability_admission: RegionalCapabilityAdmission, + noise_inversion: NoiseInversionOptions | None = None, ) -> Latent: """Sample every latent batch item using its connected SEGS guide.""" @@ -108,6 +110,8 @@ class SEGSGuidedTiledDiffusionSamplingService: differential_diffusion=differential_diffusion, capability_admission=capability_admission, tiled_plan=plan, + noise_inversion=noise_inversion, + inversion_segs=segs_for_item, ) output_samples = output["samples"] if not isinstance(output_samples, torch.Tensor): diff --git a/simple_syrup/services/tiled_attention_coupling_sampling_service.py b/simple_syrup/services/tiled_attention_coupling_sampling_service.py index ef01e23..44b432a 100644 --- a/simple_syrup/services/tiled_attention_coupling_sampling_service.py +++ b/simple_syrup/services/tiled_attention_coupling_sampling_service.py @@ -12,6 +12,7 @@ from ..domain.attention_coupling_request import ( AttentionCouplingRequestMode, classify_attention_coupling_request, ) +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_attention_execution import RegionalAttentionExecutionMode from ..domain.regional_features import ( EMPTY_REGIONAL_FEATURE_REQUEST, @@ -62,6 +63,8 @@ class TiledAttentionCouplingSamplingService: latent_tile_batch_size: int, preview_context: DetailPreviewContext | None = None, differential_diffusion: bool = False, + segs: object | None = None, + noise_inversion: NoiseInversionOptions | None = None, ) -> dict[str, Any]: """Bypass ordinary requests or prepare one complete regional request.""" @@ -90,6 +93,8 @@ class TiledAttentionCouplingSamplingService: preview_context=preview_context, differential_diffusion=differential_diffusion, feature_request=EMPTY_REGIONAL_FEATURE_REQUEST, + segs=segs, + noise_inversion=noise_inversion, ) prepared = self.model_preparation_service_class().prepare( @@ -121,6 +126,8 @@ class TiledAttentionCouplingSamplingService: preview_context=preview_context, differential_diffusion=differential_diffusion, feature_request=_TILED_ATTENTION_REQUEST, + segs=segs, + noise_inversion=noise_inversion, ) def _sample_tiled( @@ -144,6 +151,8 @@ class TiledAttentionCouplingSamplingService: preview_context: DetailPreviewContext | None, differential_diffusion: bool, feature_request: RegionalFeatureRequest, + segs: object | None, + noise_inversion: NoiseInversionOptions | None, ) -> dict[str, Any]: """Delegate one ordinary or Attention Coupling tiled request.""" @@ -166,10 +175,11 @@ class TiledAttentionCouplingSamplingService: preview_context=preview_context, differential_diffusion=differential_diffusion, feature_request=feature_request, - segs=None, + segs=segs, region_masks=None, regional_prompt_weight=0.5, region_mask_feather=0, + noise_inversion=noise_inversion, ) diff --git a/simple_syrup/services/tiled_detail_sampler.py b/simple_syrup/services/tiled_detail_sampler.py index 19a346f..1011585 100644 --- a/simple_syrup/services/tiled_detail_sampler.py +++ b/simple_syrup/services/tiled_detail_sampler.py @@ -10,6 +10,7 @@ from typing import Any, Protocol import torch +from ..domain.noise_inversion import NoiseInversionOptions from ..runtime.detail_previews import DetailPreviewContext from ..runtime.detail_sampling import DetailSampler, Latent from .tiled_diffusion_sampling_service import TiledDiffusionSamplingService @@ -38,6 +39,7 @@ class TiledDiffusionLatentSamplingBoundary(Protocol): latent_tile_batch_size: int, preview_context: DetailPreviewContext | None = None, differential_diffusion: bool = False, + noise_inversion: NoiseInversionOptions | None = None, ) -> Latent: """Sample a latent using the selected tiled diffusion mode.""" @@ -87,6 +89,7 @@ class TiledDetailSampler: latent_tile_batch_size: int, preview_context: DetailPreviewContext | None = None, differential_diffusion: bool = False, + noise_inversion: NoiseInversionOptions | None = None, ) -> Latent: """Sample one latent crop with the selected tiled diffusion runtime.""" @@ -108,4 +111,5 @@ class TiledDetailSampler: latent_tile_batch_size=latent_tile_batch_size, preview_context=preview_context, differential_diffusion=differential_diffusion, + noise_inversion=noise_inversion, ) diff --git a/simple_syrup/services/tiled_diffusion_conditioning_batch_service.py b/simple_syrup/services/tiled_diffusion_conditioning_batch_service.py index 8253aa9..d4a2087 100644 --- a/simple_syrup/services/tiled_diffusion_conditioning_batch_service.py +++ b/simple_syrup/services/tiled_diffusion_conditioning_batch_service.py @@ -11,6 +11,7 @@ from typing import Any, TypeAlias import torch from ..domain.conditioning_batch import select_conditioning +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_features import RegionalCapabilityAdmission from ..runtime.detail_previews import DetailPreviewContext from .sampling_batch import combine_latent_outputs, single_item_latent @@ -44,6 +45,7 @@ class TiledDiffusionConditioningBatchService: preview_context: DetailPreviewContext | None, differential_diffusion: bool, capability_admission: RegionalCapabilityAdmission, + noise_inversion: NoiseInversionOptions | None = None, ) -> Latent: """Sample each latent item with its selected conditioning values.""" @@ -72,6 +74,7 @@ class TiledDiffusionConditioningBatchService: preview_context=preview_context, differential_diffusion=differential_diffusion, capability_admission=capability_admission, + noise_inversion=noise_inversion, ) output_samples = output.get("samples") if not isinstance(output_samples, torch.Tensor): diff --git a/simple_syrup/services/tiled_diffusion_item_sampling_service.py b/simple_syrup/services/tiled_diffusion_item_sampling_service.py index 4c9cb1c..fa07c5d 100644 --- a/simple_syrup/services/tiled_diffusion_item_sampling_service.py +++ b/simple_syrup/services/tiled_diffusion_item_sampling_service.py @@ -8,7 +8,11 @@ from __future__ import annotations from typing import Any, Protocol, TypeAlias +import torch + +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_features import RegionalCapabilityAdmission +from ..domain.segs import NativeSegs from ..domain.tiled_diffusion import TiledDiffusionPlan from ..runtime import mixture_of_diffusers_sampling, multidiffusion_sampling from ..runtime.detail_previews import DetailPreviewContext @@ -41,6 +45,9 @@ class TiledDiffusionItemSampler(Protocol): differential_diffusion: bool, capability_admission: RegionalCapabilityAdmission, tiled_plan: TiledDiffusionPlan | None = None, + noise_inversion: NoiseInversionOptions | None = None, + inversion_segs: NativeSegs | None = None, + inversion_region_masks: torch.Tensor | None = None, ) -> Latent: """Return one sampled latent item.""" @@ -70,6 +77,9 @@ class TiledDiffusionItemSamplingService: differential_diffusion: bool, capability_admission: RegionalCapabilityAdmission, tiled_plan: TiledDiffusionPlan | None = None, + noise_inversion: NoiseInversionOptions | None = None, + inversion_segs: NativeSegs | None = None, + inversion_region_masks: torch.Tensor | None = None, ) -> Latent: """Invoke exactly one runtime with the unchanged sampling request.""" @@ -97,4 +107,7 @@ class TiledDiffusionItemSamplingService: differential_diffusion=differential_diffusion, capability_admission=capability_admission, tiled_plan=tiled_plan, + noise_inversion=noise_inversion, + inversion_segs=inversion_segs, + inversion_region_masks=inversion_region_masks, ) diff --git a/simple_syrup/services/tiled_diffusion_sampling_service.py b/simple_syrup/services/tiled_diffusion_sampling_service.py index e3ffba0..40fb977 100644 --- a/simple_syrup/services/tiled_diffusion_sampling_service.py +++ b/simple_syrup/services/tiled_diffusion_sampling_service.py @@ -9,6 +9,7 @@ from __future__ import annotations from typing import Any, ClassVar from ..domain.conditioning_batch import ConditioningBatch +from ..domain.noise_inversion import NoiseInversionOptions from ..domain.regional_features import ( EMPTY_REGIONAL_FEATURE_REQUEST, TILED_DIFFUSION_REGIONAL_SAMPLER_CAPABILITIES, @@ -84,6 +85,7 @@ class TiledDiffusionSamplingService: region_masks: object | None = None, regional_prompt_weight: float = 0.5, region_mask_feather: int = 0, + noise_inversion: NoiseInversionOptions | None = None, ) -> Latent: """Sample a latent with the selected tiled diffusion method.""" @@ -129,6 +131,7 @@ class TiledDiffusionSamplingService: preview_context=preview_context, differential_diffusion=differential_diffusion, capability_admission=capability_admission, + noise_inversion=noise_inversion, ) if segs is not None: return self.segs_sampling_service_class().sample( @@ -151,6 +154,7 @@ class TiledDiffusionSamplingService: preview_context=preview_context, differential_diffusion=differential_diffusion, capability_admission=capability_admission, + noise_inversion=noise_inversion, segs=segs, ) if isinstance(positive, ConditioningBatch) or isinstance( @@ -177,6 +181,7 @@ class TiledDiffusionSamplingService: preview_context=preview_context, differential_diffusion=differential_diffusion, capability_admission=capability_admission, + noise_inversion=noise_inversion, ) return self.item_sampling_service_class().sample( diffusion_mode=diffusion_mode, @@ -197,4 +202,5 @@ class TiledDiffusionSamplingService: preview_context=preview_context, differential_diffusion=differential_diffusion, capability_admission=capability_admission, + noise_inversion=noise_inversion, ) diff --git a/tests/crosscutting/contracts/test_registration.py b/tests/crosscutting/contracts/test_registration.py index 91533a0..d909654 100644 --- a/tests/crosscutting/contracts/test_registration.py +++ b/tests/crosscutting/contracts/test_registration.py @@ -18,6 +18,11 @@ from support.repository import REPOSITORY_ROOT BASE_NODE_IDS = [ "SimpleSyrup.AllPromptAttentionSEGS", + "SimpleSyrup.AttentionCouplingOptions", + "SimpleSyrup.ContextualDiffusionOptions", + "SimpleSyrup.NoiseInversionOptions", + "SimpleSyrup.TilingOptions", + "SimpleSyrup.KSampler", "SimpleSyrup.AttentionCaptureModel", "SimpleSyrup.AttentionMaskedConditioning", "SimpleSyrup.AttentionRegionMask", diff --git a/tests/models/patching/test_patcher_lifecycle_policy.py b/tests/models/patching/test_patcher_lifecycle_policy.py index e6c4812..de0a198 100644 --- a/tests/models/patching/test_patcher_lifecycle_policy.py +++ b/tests/models/patching/test_patcher_lifecycle_policy.py @@ -32,6 +32,9 @@ FORBIDDEN_PATCHER_CALLS = frozenset( FORBIDDEN_PATCHER_WRITES = frozenset({"forced_hooks", "use_clip_schedule"}) APPROVED_VALUE_CLONES = Counter( { + ("simple_syrup/domain/inversion_solver.py", "source"): 1, + ("simple_syrup/runtime/noise_inversion.py", "zero"): 1, + ("simple_syrup/runtime/noise_inversion.py", "noise"): 1, ("simple_syrup/domain/semantic_tiled_diffusion.py", "mask"): 2, ("simple_syrup/image/crop_composite.py", "image"): 1, ( diff --git a/tests/node_api/test_legacy_inversion_node_adapter.py b/tests/node_api/test_legacy_inversion_node_adapter.py new file mode 100644 index 0000000..4132b1b --- /dev/null +++ b/tests/node_api/test_legacy_inversion_node_adapter.py @@ -0,0 +1,95 @@ +"""Verify optional inversion widgets on implementation-backed V3 sampler nodes.""" + +from __future__ import annotations + +from typing import Any, ClassVar + +import pytest + +from simple_syrup.domain.noise_inversion import NoiseInversionOptions +from simple_syrup.nodes_v3.legacy_inversion_node_adapter import ( + LegacyInversionNodeV3Adapter, +) +from simple_syrup.nodes_v3.legacy_node_wrappers import ( + DetailSEGSAsRegionsV3, + DetailSEGSByScaleFactorTiledDiffusionV3, + KSamplerExtrasV3, +) + + +class RecordingImplementation: + """Expose a minimal maintained declaration and the execution boundary.""" + + RETURN_TYPES = ("STRING",) + OUTPUT_TOOLTIPS = ("Result.",) + FUNCTION = "run" + CATEGORY = "SimpleSyrup/Test" + DESCRIPTION = "Records sampling inputs." + INPUT_IS_LIST = False + + @classmethod + def INPUT_TYPES(cls) -> dict[str, Any]: + """Keep the persisted input contract independent of inversion widgets.""" + return {"required": {"text": ("STRING", {"default": "", "tooltip": "Text."})}} + + def run( + self, text: object, noise_inversion: NoiseInversionOptions | None = None + ) -> tuple[object, NoiseInversionOptions | None]: + """Return delegated inputs without invoking Comfy neural execution.""" + return text, noise_inversion + + +class RecordingAdapter(LegacyInversionNodeV3Adapter): + """Use the production schema and list-mode normalization implementation.""" + + LEGACY_NODE_CLASS: ClassVar[type[Any]] = RecordingImplementation + NODE_ID = "SimpleSyrup.Recording" + DISPLAY_NAME = "Recording" + + +@pytest.mark.parametrize("list_mode", [False, True]) +def test_inversion_controls_normalize_and_default_off( + monkeypatch: pytest.MonkeyPatch, + list_mode: bool, +) -> None: + """Preserve existing execution while constructing the accepted recipe on request.""" + monkeypatch.setattr(RecordingImplementation, "INPUT_IS_LIST", list_mode) + text = ["prompt"] if list_mode else "prompt" + output, disabled = RecordingAdapter.execute(text=text) + assert output == text and disabled is None + _, enabled = RecordingAdapter.execute( + text=text, + noise_inversion_enabled=[True] if list_mode else True, + ) + assert enabled == NoiseInversionOptions() + + +@pytest.mark.parametrize( + "controls", + [ + {"inversion_steps": 0}, + {"inversion_resolution_scale": 0}, + {"inversion_method": "fireflow"}, + {"inversion_switch_fraction": 1}, + ], +) +def test_invalid_selected_inversion_controls_fail(controls: dict[str, Any]) -> None: + """Apply domain validation instead of handing malformed controls to a sampler.""" + with pytest.raises(ValueError): + RecordingAdapter.execute( + text="prompt", noise_inversion_enabled=True, **controls + ) + + +@pytest.mark.parametrize( + "node", + [KSamplerExtrasV3, DetailSEGSAsRegionsV3, DetailSEGSByScaleFactorTiledDiffusionV3], +) +def test_existing_nodes_append_optional_inversion_without_reordering(node: Any) -> None: + """Preserve every serialized socket position before seven optional new controls.""" + schema = node.define_schema() + order = node.WORKFLOW_INPUT_ORDER + assert [item.id for item in schema.inputs[: len(order)]] == list(order) + new = schema.inputs[len(order) :] + assert len(new) == 7 and all(item.optional and item.tooltip for item in new) + assert new[0].id == "noise_inversion_enabled" and new[0].default is False diff --git a/tests/node_api/test_legacy_node_v3_adapter.py b/tests/node_api/test_legacy_node_v3_adapter.py index 2d141d9..006ba83 100644 --- a/tests/node_api/test_legacy_node_v3_adapter.py +++ b/tests/node_api/test_legacy_node_v3_adapter.py @@ -8,7 +8,7 @@ from __future__ import annotations from typing import Any, ClassVar -from simple_syrup.nodes_v3.legacy_node_wrappers import LegacyNodeV3Adapter +from simple_syrup.nodes_v3.legacy_node_adapter import LegacyNodeV3Adapter class _FakeHidden: diff --git a/tests/regional_generation/attention_coupling/test_attention_coupling_sampling_service.py b/tests/regional_generation/attention_coupling/test_attention_coupling_sampling_service.py index f64df65..2e59685 100644 --- a/tests/regional_generation/attention_coupling/test_attention_coupling_sampling_service.py +++ b/tests/regional_generation/attention_coupling/test_attention_coupling_sampling_service.py @@ -158,6 +158,7 @@ def test_ordinary_request_bypasses_preparation_and_preserves_img2img_inputs( "negative": "negative", "latent_image": latent, "denoise": 0.42, + "noise_inversion": None, } ] diff --git a/tests/regional_generation/regional/test_regional_ksampler_v3_nodes.py b/tests/regional_generation/regional/test_regional_ksampler_v3_nodes.py index 113a0c2..ed345b7 100644 --- a/tests/regional_generation/regional/test_regional_ksampler_v3_nodes.py +++ b/tests/regional_generation/regional/test_regional_ksampler_v3_nodes.py @@ -70,15 +70,23 @@ def test_schemas_expose_exact_names_and_shared_regional_contract() -> None: "region_mask_feather", "latent_image", "denoise", + "noise_inversion_enabled", + "inversion_method", + "inversion_resolution_scale", + "inversion_steps", + "inversion_switch_fraction", + "inversion_finishing_steps", + "inversion_finishing_method", ] - assert tiled_ids[: len(normal_ids)] == normal_ids - assert tiled_ids[len(normal_ids) :] == [ + assert tiled_ids[:13] == normal_ids[:13] + assert tiled_ids[13:18] == [ "diffusion_mode", "latent_tile_width", "latent_tile_height", "latent_tile_overlap", "latent_tile_batch_size", ] + assert tiled_ids[18:] == normal_ids[13:] regional_weight = normal.inputs[9] assert regional_weight.default == 0.5 assert regional_weight.min == 0.0 @@ -95,6 +103,13 @@ def test_schemas_expose_exact_names_and_shared_regional_contract() -> None: "regional_prompt_weight": 0.5, "region_mask_feather": 0, "denoise": 1.0, + "noise_inversion_enabled": False, + "inversion_method": "euler", + "inversion_resolution_scale": 0.5, + "inversion_steps": 2, + "inversion_switch_fraction": 0.75, + "inversion_finishing_steps": 1, + "inversion_finishing_method": "euler", } tiled_defaults = { value.id: value.default for value in tiled.inputs if hasattr(value, "default") diff --git a/tests/regional_generation/spatial/test_contextual_diffusion_domain.py b/tests/regional_generation/spatial/test_contextual_diffusion_domain.py index 56ce896..b6ccd0e 100644 --- a/tests/regional_generation/spatial/test_contextual_diffusion_domain.py +++ b/tests/regional_generation/spatial/test_contextual_diffusion_domain.py @@ -6,6 +6,8 @@ from __future__ import annotations +from dataclasses import replace + import pytest import torch @@ -88,6 +90,27 @@ def test_invalid_overlap_fails_before_planning() -> None: ) +def test_rectangular_tiles_do_not_resize_global_context() -> None: + """Let Tiling configure the sole local plan independently of global context.""" + controls = replace(_controls(), latent_tile_width=48, latent_tile_height=16) + plan = build_contextual_diffusion_plan( + latent_width=96, latent_height=64, controls=controls, segs=None + ) + assert (plan.tile_plan.tile_width, plan.tile_plan.tile_height) == (48, 16) + assert (plan.global_view.model_width, plan.global_view.model_height) == (32, 22) + + +def test_overlap_is_bounded_by_both_explicit_tile_dimensions() -> None: + """Reject an overlap that prevents advancing the shorter local dimension.""" + controls = replace( + _controls(latent_context_overlap=16), + latent_tile_width=48, + latent_tile_height=16, + ) + with pytest.raises(ValueError, match="both local tile dimensions"): + controls.validate() + + def _controls( *, latent_context_overlap: int = 8, diff --git a/tests/sampling/support/__init__.py b/tests/sampling/support/__init__.py new file mode 100644 index 0000000..b912383 --- /dev/null +++ b/tests/sampling/support/__init__.py @@ -0,0 +1 @@ +"""Provide scoped sampling-boundary fixtures without replacing domain behavior.""" diff --git a/tests/sampling/support/inversion_model.py b/tests/sampling/support/inversion_model.py new file mode 100644 index 0000000..274920a --- /dev/null +++ b/tests/sampling/support/inversion_model.py @@ -0,0 +1,36 @@ +"""Represent the external MODEL patcher surface for real spatial wrapper tests.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import torch + +from simple_syrup.runtime.sampling_model_types import ModelFunctionWrapper + + +class SpatialModel: + """Keep independent model options and an explicit lifecycle parent.""" + + def __init__(self, options: dict[str, Any] | None = None) -> None: + """Initialize the minimal host state needed by wrapper derivation.""" + self.model_options = {} if options is None else options + self.load_device = torch.device("cpu") + self.parent: SpatialModel | None = None + + def clone(self) -> SpatialModel: + """Copy host options and retain the direct source identity.""" + derived = SpatialModel(self.model_options.copy()) + derived.parent = self + return derived + + def set_model_unet_function_wrapper(self, wrapper: ModelFunctionWrapper) -> None: + """Implement the public Comfy wrapper installation boundary.""" + self.model_options["model_function_wrapper"] = wrapper + + def set_model_sampler_calc_cond_batch_function( + self, function: Callable[..., object] + ) -> None: + """Implement the external Comfy regional prediction registration boundary.""" + self.model_options["sampler_calc_cond_batch_function"] = function diff --git a/tests/sampling/test_guided_noise_inversion.py b/tests/sampling/test_guided_noise_inversion.py new file mode 100644 index 0000000..9911cf9 --- /dev/null +++ b/tests/sampling/test_guided_noise_inversion.py @@ -0,0 +1,181 @@ +"""Prove inversion composes with both Comfy guider paths exactly once per sample.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch +from comfy import samplers +from comfy.model_sampling import CONST + +from simple_syrup.domain.noise_inversion import NoiseInversionOptions +from simple_syrup.runtime.guided_sampling import sample_with_optional_negative + + +def _model() -> Any: + """Build a dynamic external model with real Comfy flow noise scaling.""" + scaling = CONST() + scaling.sigma_max = 1.0 + return SimpleNamespace( + load_device=torch.device("cpu"), + get_model_object=lambda name: scaling, + model=SimpleNamespace(process_latent_in=lambda x: x), + ) + + +@pytest.fixture +def execution_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict[str, Any]]: + """Replace expensive neural execution, not the shared guider or inversion logic.""" + calls: list[dict[str, Any]] = [] + + def execute( + model: Any, + noise: torch.Tensor, + latent: torch.Tensor, + sampler: Any, + sigmas: torch.Tensor, + **extra: Any, + ) -> torch.Tensor: + """Run custom inversion samplers and record the final supplied forward noise.""" + calls.append({"noise": noise.clone(), "shape": tuple(latent.shape), **extra}) + if not isinstance(sampler, samplers.KSAMPLER): + return noise + state = model.get_model_object("model_sampling").noise_scaling( + sigmas[0], noise, latent, max_denoise=False + ) + + def predict( + x: torch.Tensor, sigma: torch.Tensor, **kwargs: Any + ) -> torch.Tensor: + """Return a known constant velocity through the dynamic model boundary.""" + return x - sigma.reshape((-1,) + (1,) * (x.ndim - 1)) * 0.25 + + return cast( + torch.Tensor, + sampler.sampler_function( + predict, state, sigmas, extra_args={}, callback=None, disable=True + ), + ) + + def sample_custom( + model: Any, + noise: torch.Tensor, + cfg: float, + sampler: Any, + sigmas: torch.Tensor, + positive: Any, + negative: Any, + latent: torch.Tensor, + **extra: Any, + ) -> torch.Tensor: + """Represent the CFG execution boundary with both original branches.""" + return execute( + model, + noise, + latent, + sampler, + sigmas, + positive=positive, + negative=negative, + cfg=cfg, + **extra, + ) + + class PositiveOnlyGuider: + """Represent only the installed guider's external sampling contract.""" + + def __init__(self, model: Any) -> None: + """Retain the requested model for positive-only neural execution.""" + self.model = model + self.conds: dict[str, Any] = {} + + def inner_set_conds(self, conds: dict[str, Any]) -> None: + """Record actual branches selected by the shared production guider.""" + self.conds = conds + + def sample( + self, + noise: torch.Tensor, + latent: torch.Tensor, + sampler: Any, + sigmas: torch.Tensor, + **extra: Any, + ) -> torch.Tensor: + """Run the requested stage with only registered positive conditioning.""" + return execute( + self.model, noise, latent, sampler, sigmas, **self.conds, **extra + ) + + from comfy import sample + + monkeypatch.setattr(sample, "sample_custom", sample_custom) + monkeypatch.setattr(samplers, "CFGGuider", PositiveOnlyGuider) + return calls + + +@pytest.mark.parametrize("negative", [None, []]) +def test_inversion_and_forward_sampling_keep_the_same_guidance_branches( + execution_calls: list[dict[str, Any]], + negative: Any, +) -> None: + """Execute two inversion stages and one forward pass without recursive inversion.""" + from comfy import sample + + positive = [[torch.ones((1, 1, 1)), {}]] + random_noise = torch.full((1, 4, 8, 12), -3.0) + result = sample_with_optional_negative( + comfy_sample=sample, + model=_model(), + noise=random_noise, + cfg=7.5, + sampler=object(), + sigmas=torch.tensor([0.5, 0]), + positive=positive, + negative=negative, + latent_image=torch.zeros_like(random_noise), + seed=42, + noise_inversion=NoiseInversionOptions(), + ) + assert len(execution_calls) == 3 + assert [call["shape"] for call in execution_calls] == [ + (1, 4, 4, 6), + (1, 4, 8, 12), + (1, 4, 8, 12), + ] + for call in execution_calls: + assert call["positive"] is positive + assert call["seed"] == 42 + if negative is None: + assert "negative" not in call and "cfg" not in call + else: + assert call["negative"] is negative and call["cfg"] == 7.5 + assert torch.allclose(result, torch.full_like(result, 0.24995), atol=1e-6) + assert torch.equal(execution_calls[-1]["noise"], result) + assert torch.equal(random_noise, torch.full_like(random_noise, -3)) + + +@pytest.mark.parametrize("negative", [None, []]) +def test_absent_inversion_preserves_supplied_noise_without_extra_execution( + execution_calls: list[dict[str, Any]], + negative: Any, +) -> None: + """Keep existing workflows on their unchanged sampling path.""" + from comfy import sample + + noise = torch.ones((1, 4, 8, 12)) + result = sample_with_optional_negative( + comfy_sample=sample, + model=_model(), + noise=noise, + cfg=7.5, + sampler=object(), + sigmas=torch.tensor([0.5, 0]), + positive=[], + negative=negative, + latent_image=torch.zeros_like(noise), + seed=42, + ) + assert len(execution_calls) == 1 + assert torch.equal(result, noise) diff --git a/tests/sampling/test_inversion_model_factory.py b/tests/sampling/test_inversion_model_factory.py new file mode 100644 index 0000000..6bd136a --- /dev/null +++ b/tests/sampling/test_inversion_model_factory.py @@ -0,0 +1,192 @@ +"""Prove resolution-specific inversion planning and canonical regional masks.""" + +from __future__ import annotations + +from typing import Any + +import pytest +import torch + +from simple_syrup.domain.contextual_diffusion import ContextualDiffusionControls +from simple_syrup.domain.regional_mask_bank import RegionalMaskBank +from simple_syrup.domain.sampler_options import TilingOptions +from simple_syrup.domain.spatial_views import SpatialBatchLayout, SpatialViewKind +from simple_syrup.masking.regional_mask_projection import ( + RegionalMaskForm, + RegionalMaskProjectionMode, + RegionalMaskProjector, +) +from simple_syrup.runtime.contextual_model_wrapper import ( + ContextualDiffusionModelWrapper, +) +from simple_syrup.runtime.inversion_model_factory import InversionModelFactory +from simple_syrup.runtime.inversion_spatial_context import ( + InversionSpatialContextWrapper, +) + +from .support.inversion_model import SpatialModel + + +@pytest.mark.parametrize("mode", ("multidiffusion", "mixture_of_diffusers")) +@pytest.mark.parametrize("contextual", (False, True)) +def test_each_stage_is_planned_from_the_original_model( + mode: str, contextual: bool +) -> None: + """Avoid full-size fallback, nested tiling and mutation of the supplied MODEL.""" + model = SpatialModel() + context = ContextualDiffusionControls(16, 4, 2, 0.5, 1, 0.5) + factory = InversionModelFactory( + model=model, + canvas_width=64, + canvas_height=32, + tiling=TilingOptions( + diffusion_mode=mode, width=16, height=16, overlap=4, batch_size=2 + ), + context=context if contextual else None, + forward_sigmas=torch.tensor([0.5, 0.25, 0]), + ) + model_call_shapes: list[tuple[int, ...]] = [] + + def denoiser(x: torch.Tensor, timestep: torch.Tensor, **c: Any) -> torch.Tensor: + """Provide only the external neural boundary with a deterministic prediction.""" + model_call_shapes.append(tuple(x.shape)) + return x + 1 + + for height, width in ((16, 32), (32, 64)): + stage_model = factory(torch.zeros((1, 4, height, width))) + assert stage_model.parent is model + wrapper = stage_model.model_options["model_function_wrapper"] + if contextual: + assert isinstance(wrapper, ContextualDiffusionModelWrapper) + start = len(model_call_shapes) + x = torch.zeros((1, 4, height, width)) + result = wrapper( + denoiser, {"input": x, "timestep": torch.tensor([0.5]), "c": {}} + ) + assert result.shape == x.shape + assert model_call_shapes[start:] + assert all( + shape[-2] <= 16 and shape[-1] <= 16 for shape in model_call_shapes[start:] + ) + assert model.model_options == {} + + +@pytest.mark.parametrize("contextual", (False, True)) +def test_regional_mask_bank_keeps_original_coordinates_in_coarse_views( + contextual: bool, +) -> None: + """Use real mask projection after coarse local or global context calls.""" + mask = torch.zeros((1, 32, 64)) + mask[:, :, 32:] = 1 + bank = RegionalMaskBank(mask, mask.clone(), 64, 32) + layouts: list[SpatialBatchLayout] = [] + options: list[dict[str, Any]] = [] + + def original(apply_model: Any, args: dict[str, Any]) -> torch.Tensor: + """Exercise the real regional projection authority at the preserved wrapper.""" + transformer = args["c"]["transformer_options"] + layout = transformer["simple_syrup"]["spatial_batch_layout"] + assert isinstance(layout, SpatialBatchLayout) + assert (layout.canvas_width, layout.canvas_height) == (64, 32) + assert layout.expanded_batch_size == int(args["input"].shape[0]) + for index in range(layout.view_count): + projected = RegionalMaskProjector().project_view( + bank=bank, + layout=layout, + view_index=index, + form=RegionalMaskForm.CONDITIONING, + mode=RegionalMaskProjectionMode.NEAREST, + ) + assert tuple(projected.shape[-2:]) == ( + layout.views[index].model_height, + layout.views[index].model_width, + ) + layouts.append(layout) + options.append(transformer) + prediction_input = args["input"] + assert isinstance(prediction_input, torch.Tensor) + return prediction_input + 1 + + model = SpatialModel({"model_function_wrapper": original}) + factory = InversionModelFactory( + model=model, + canvas_width=64, + canvas_height=32, + tiling=TilingOptions(width=16, height=16, overlap=4, batch_size=2), + context=ContextualDiffusionControls(16, 4, 2, 1.0, 1, 0.5) + if contextual + else None, + forward_sigmas=torch.tensor([0.5, 0.25, 0]), + region_masks=mask if contextual else None, + ) + stage = factory(torch.zeros((2, 4, 16, 32))) + wrapper = stage.model_options["model_function_wrapper"] + transformer = {"cond_or_uncond": [0, 1], "uuids": ["positive", "negative"]} + args = { + "input": torch.zeros((2, 4, 16, 32)), + "timestep": torch.tensor([0.5, 0.5]), + "c": {"transformer_options": transformer}, + } + wrapper(lambda *args, **kwargs: None, args) + assert layouts + assert any(layout.views[0].kind is SpatialViewKind.TILE for layout in layouts) + if contextual: + assert any( + layout.views[0].kind is SpatialViewKind.CONTEXTUAL_GLOBAL + for layout in layouts + ) + for layout, transformed in zip(layouts, options, strict=True): + assert len(transformed["cond_or_uncond"]) == 2 * layout.view_count + assert transformer == {"cond_or_uncond": [0, 1], "uuids": ["positive", "negative"]} + assert model.model_options["model_function_wrapper"] is original + + +def test_full_frame_and_single_tile_calls_publish_whole_source_reduced_layout() -> None: + """Preserve regional attention through the one-tile fast path.""" + layouts: list[SpatialBatchLayout] = [] + + def original(apply_model: Any, args: dict[str, Any]) -> torch.Tensor: + """Inspect only the host wrapper boundary's spatial metadata.""" + layout = args["c"]["transformer_options"]["simple_syrup"][ + "spatial_batch_layout" + ] + layouts.append(layout) + prediction_input = args["input"] + assert isinstance(prediction_input, torch.Tensor) + return prediction_input + + model = SpatialModel({"model_function_wrapper": original}) + for tiling in (None, TilingOptions(width=32, height=32, overlap=4)): + factory = InversionModelFactory( + model=model, canvas_width=48, canvas_height=32, tiling=tiling + ) + stage = factory(torch.zeros((1, 4, 16, 24))) + stage.model_options["model_function_wrapper"]( + lambda *args, **kwargs: None, + { + "input": torch.zeros((1, 4, 16, 24)), + "timestep": torch.tensor([0.5]), + "c": {}, + }, + ) + assert len(layouts) == 2 + for layout in layouts: + assert (layout.canvas_width, layout.canvas_height) == (48, 32) + assert layout.views[0].kind is SpatialViewKind.CONTEXTUAL_GLOBAL + assert (layout.views[0].model_width, layout.views[0].model_height) == (24, 16) + + +def test_malformed_reduced_layout_is_not_guessed() -> None: + """Reject calls missing the metadata needed to identify their source crop.""" + wrapper = InversionSpatialContextWrapper( + lambda apply_model, args: args["input"], + canvas_width=64, + canvas_height=32, + stage_width=32, + stage_height=16, + ) + with pytest.raises(ValueError, match="explicit spatial layout"): + wrapper( + lambda *args, **kwargs: torch.empty(0), + {"input": torch.zeros((1, 4, 8, 8)), "c": {}}, + ) diff --git a/tests/sampling/test_inversion_solver.py b/tests/sampling/test_inversion_solver.py new file mode 100644 index 0000000..c9b66bc --- /dev/null +++ b/tests/sampling/test_inversion_solver.py @@ -0,0 +1,130 @@ +"""Prove inversion accuracy, source preservation, and finite-state safety.""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from simple_syrup.domain.inversion_solver import ( + InversionSolverEvidence, + integrate_inversion, + lift_inversion_displacement, +) +from simple_syrup.domain.noise_inversion import InversionMethod +from simple_syrup.runtime.spatial_tensor_projection import resize_spatial_tensor + + +@pytest.mark.parametrize("method", ["euler", "heun"]) +def test_constant_velocity_has_exact_endpoint_and_observed_call_count( + method: InversionMethod, +) -> None: + """Match a closed-form flow path rather than duplicating integration logic.""" + source = torch.zeros((1, 2, 4, 4)) + evidence = InversionSolverEvidence() + calls: list[tuple[float, int]] = [] + + def velocity(x: torch.Tensor, sigma: torch.Tensor, index: int) -> torch.Tensor: + """Record neural-evaluation positions and return a constant flow.""" + calls.append((float(sigma), index)) + return torch.full_like(x, 2) + + endpoint = integrate_inversion( + source, + torch.tensor([0.125, 0.25, 0.5]), + velocity, + method=method, + evidence=evidence, + ) + assert torch.equal(endpoint, torch.full_like(source, 0.75)) + assert torch.count_nonzero(source) == 0 + assert endpoint.data_ptr() != source.data_ptr() + assert evidence.evaluations == len(calls) == (2 if method == "euler" else 4) + + +@pytest.mark.parametrize(("method", "expected"), [("euler", 1.5), ("heun", 1.625)]) +def test_nonconstant_velocity_distinguishes_euler_and_heun( + method: InversionMethod, expected: float +) -> None: + """Prove Heun evaluates its predictor and averages the endpoint velocity.""" + endpoint = integrate_inversion( + torch.ones((1, 1, 2, 2)), + torch.tensor([0.25, 0.75]), + lambda x, sigma, index: x, + method=method, + ) + assert torch.equal(endpoint, torch.full_like(endpoint, expected)) + + +@pytest.mark.parametrize( + "sigmas", + [ + torch.tensor([0.0, 0.5]), + torch.tensor([0.5, 0.5]), + torch.tensor([0.5, 0.25]), + torch.tensor([0.1, float("nan")]), + torch.tensor([0.1]), + torch.ones((2, 2)), + ], +) +def test_invalid_schedule_fails_before_evaluation(sigmas: torch.Tensor) -> None: + """Reject malformed schedules without calling the model.""" + + def forbidden(x: torch.Tensor, sigma: torch.Tensor, index: int) -> torch.Tensor: + """Detect an invalid request reaching the expensive model boundary.""" + raise AssertionError("Invalid schedule reached evaluation.") + + with pytest.raises(ValueError, match="schedule|sigmas"): + integrate_inversion(torch.ones((1, 1, 2, 2)), sigmas, forbidden, method="euler") + + +@pytest.mark.parametrize("wrong_shape", [False, True]) +def test_bad_velocity_is_rejected(wrong_shape: bool) -> None: + """Never continue integrating corrupted predictions.""" + + def velocity(x: torch.Tensor, sigma: torch.Tensor, index: int) -> torch.Tensor: + """Produce one malformed prediction at the runtime boundary.""" + return torch.zeros((1,)) if wrong_shape else torch.full_like(x, float("inf")) + + with pytest.raises(FloatingPointError, match="velocity"): + integrate_inversion( + torch.ones((1, 1, 2, 2)), torch.tensor([0.1, 0.5]), velocity, method="euler" + ) + + +def _resize(tensor: torch.Tensor, height: int, width: int) -> torch.Tensor: + """Use the production spatial projector for displacement transfer.""" + return resize_spatial_tensor(tensor, height=height, width=width, mode="bilinear") + + +def test_displacement_transfer_preserves_full_source_high_frequency_detail() -> None: + """Keep original detail instead of enlarging a blurred low-resolution endpoint.""" + full_source = torch.arange(64, dtype=torch.float32).reshape(1, 1, 8, 8) % 2 + coarse_source = torch.full((1, 1, 4, 4), 0.5) + endpoint = lift_inversion_displacement( + full_source, coarse_source, coarse_source + 2, resize=_resize + ) + assert torch.equal(endpoint, full_source + 2) + + +@pytest.mark.parametrize("singleton_depth", [False, True]) +def test_zero_displacement_is_pixel_exact(singleton_depth: bool) -> None: + """Preserve both supported latent layouts without losing source pixels.""" + shape = (2, 4, 1, 7, 9) if singleton_depth else (2, 4, 7, 9) + source = torch.arange(math.prod(shape), dtype=torch.float32).reshape(shape) + coarse = resize_spatial_tensor(source, height=4, width=4, mode="area") + assert torch.equal( + lift_inversion_displacement(source, coarse, coarse, resize=_resize), source + ) + + +def test_transfer_rejects_mismatched_sources() -> None: + """Reject a transfer that would broadcast across different latent batches.""" + with pytest.raises(ValueError, match="batch"): + lift_inversion_displacement( + torch.zeros((2, 1, 8, 8)), + torch.zeros((1, 1, 4, 4)), + torch.zeros((1, 1, 4, 4)), + resize=_resize, + ) diff --git a/tests/sampling/test_ksampler_attention_coupling_v3_node.py b/tests/sampling/test_ksampler_attention_coupling_v3_node.py index a13ef0a..03c2f1c 100644 --- a/tests/sampling/test_ksampler_attention_coupling_v3_node.py +++ b/tests/sampling/test_ksampler_attention_coupling_v3_node.py @@ -49,6 +49,13 @@ def test_schema_exposes_stable_full_context_contract_and_guidance() -> None: "region_mask_feather", "latent_image", "denoise", + "noise_inversion_enabled", + "inversion_method", + "inversion_resolution_scale", + "inversion_steps", + "inversion_switch_fraction", + "inversion_finishing_steps", + "inversion_finishing_method", ] inputs = {item.id: item for item in schema.inputs} assert inputs["negative"].optional is True @@ -104,6 +111,7 @@ def test_node_delegates_the_complete_request_once() -> None: assert output is _RecordingAttentionService.output assert _RecordingAttentionService.calls == [ { + "noise_inversion": None, "model": "model", "seed": 9, "steps": 30, diff --git a/tests/sampling/test_ksampler_contextual_attention_coupling_v3_node.py b/tests/sampling/test_ksampler_contextual_attention_coupling_v3_node.py index 71b95b9..83f3603 100644 --- a/tests/sampling/test_ksampler_contextual_attention_coupling_v3_node.py +++ b/tests/sampling/test_ksampler_contextual_attention_coupling_v3_node.py @@ -73,6 +73,13 @@ def test_schema_exposes_stable_contextual_attention_coupling_contract() -> None: "global_steps", "global_decay", "segs", + "noise_inversion_enabled", + "inversion_method", + "inversion_resolution_scale", + "inversion_steps", + "inversion_switch_fraction", + "inversion_finishing_steps", + "inversion_finishing_method", ] assert [output.id for output in schema.outputs] == ["latent", "contexts_segs"] assert inputs["segs"].optional @@ -174,6 +181,7 @@ def test_node_delegates_every_contextual_attention_input_once( "global_steps": 4, "global_decay": 0.25, "segs": segs, + "noise_inversion": None, } ] diff --git a/tests/sampling/test_ksampler_contextual_diffusion_node.py b/tests/sampling/test_ksampler_contextual_diffusion_node.py index 240a38d..422df41 100644 --- a/tests/sampling/test_ksampler_contextual_diffusion_node.py +++ b/tests/sampling/test_ksampler_contextual_diffusion_node.py @@ -153,6 +153,7 @@ def test_sample_delegates_every_control_to_service( "region_masks": None, "regional_prompt_weight": 0.5, "region_mask_feather": 0, + "noise_inversion": None, } ] diff --git a/tests/sampling/test_ksampler_tiled_attention_coupling_v3_node.py b/tests/sampling/test_ksampler_tiled_attention_coupling_v3_node.py index f1d8d1c..8385e45 100644 --- a/tests/sampling/test_ksampler_tiled_attention_coupling_v3_node.py +++ b/tests/sampling/test_ksampler_tiled_attention_coupling_v3_node.py @@ -61,6 +61,13 @@ def test_schema_exposes_stable_tiled_attention_coupling_contract() -> None: "latent_tile_height", "latent_tile_overlap", "latent_tile_batch_size", + "noise_inversion_enabled", + "inversion_method", + "inversion_resolution_scale", + "inversion_steps", + "inversion_switch_fraction", + "inversion_finishing_steps", + "inversion_finishing_method", ] assert inputs["region_masks"].optional is True assert inputs["regional_prompt_weight"].default == 1.0 @@ -134,6 +141,7 @@ def test_node_delegates_every_tiled_attention_input_once( assert output is _RecordingTiledAttentionService.output assert _RecordingTiledAttentionService.calls == [ { + "noise_inversion": None, "diffusion_mode": "mixture_of_diffusers", "model": "model", "seed": 17, diff --git a/tests/sampling/test_ksampler_tiled_diffusion_node.py b/tests/sampling/test_ksampler_tiled_diffusion_node.py index 1db601e..6d03df5 100644 --- a/tests/sampling/test_ksampler_tiled_diffusion_node.py +++ b/tests/sampling/test_ksampler_tiled_diffusion_node.py @@ -11,6 +11,7 @@ from typing import Any import pytest import torch +from simple_syrup.domain.noise_inversion import NoiseInversionOptions from simple_syrup.nodes_v3.ksampler_tiled_diffusion import ( KSamplerTiledDiffusionV3, ) @@ -191,6 +192,7 @@ class _FakeTiledDiffusionSamplingService: region_masks: object | None = None, regional_prompt_weight: float = 0.5, region_mask_feather: int = 0, + noise_inversion: NoiseInversionOptions | None = None, ) -> dict[str, Any]: """Record sampling arguments and return a fixed latent.""" @@ -216,6 +218,7 @@ class _FakeTiledDiffusionSamplingService: "region_masks": region_masks, "regional_prompt_weight": regional_prompt_weight, "region_mask_feather": region_mask_feather, + "noise_inversion": noise_inversion, } ) return self.output diff --git a/tests/sampling/test_noise_inversion_options.py b/tests/sampling/test_noise_inversion_options.py new file mode 100644 index 0000000..74002cd --- /dev/null +++ b/tests/sampling/test_noise_inversion_options.py @@ -0,0 +1,90 @@ +"""Safeguard the accepted inversion recipe and independently editable controls.""" + +from __future__ import annotations + +from dataclasses import FrozenInstanceError +from typing import Any + +import pytest + +from simple_syrup.domain.noise_inversion import NoiseInversionOptions + + +def test_defaults_match_accepted_half_euler_recipe() -> None: + """Retain the user-accepted reduced-size path and full-size finishing stage.""" + options = NoiseInversionOptions() + assert options.method == "euler" + assert options.resolution_scale == 0.5 + assert options.steps == 2 + assert options.switch_fraction == 0.75 + assert options.finishing_steps == 1 + assert options.finishing_method == "euler" + assert options.coarse_target_fraction == 0.75 + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"method": "fireflow"}, "method"), + ({"finishing_method": "invalid"}, "finishing method"), + ({"resolution_scale": 0}, "resolution"), + ({"resolution_scale": 1.1}, "resolution"), + ({"resolution_scale": float("nan")}, "resolution"), + ({"steps": 0}, "steps"), + ({"steps": 65}, "steps"), + ({"steps": True}, "steps"), + ({"steps": 1.5}, "steps"), + ({"finishing_steps": -1}, "finishing steps"), + ({"finishing_steps": 65}, "finishing steps"), + ({"switch_fraction": 0}, "transition"), + ({"switch_fraction": float("inf")}, "transition"), + ({"switch_fraction": 1}, "finish"), + ], +) +def test_invalid_configuration_is_rejected( + changes: dict[str, Any], message: str +) -> None: + """Fail before any inversion/model execution for invalid graph controls.""" + with pytest.raises(ValueError, match=message): + NoiseInversionOptions(**changes) + + +def test_no_finish_reaches_target_without_changing_the_saved_transition() -> None: + """Make finishing-step zero a usable coarse-only configuration.""" + options = NoiseInversionOptions(finishing_steps=0) + assert options.coarse_target_fraction == 1 + assert options.switch_fraction == 0.75 + + +def test_full_resolution_preserves_odd_dimensions_and_uses_one_stage() -> None: + """Never round an existing full-size latent onto a different grid.""" + options = NoiseInversionOptions(resolution_scale=1) + assert options.coarse_shape(129, 203) == (129, 203) + assert options.coarse_target_fraction == 1 + + +@pytest.mark.parametrize( + ("height", "width", "scale", "expected"), + [(144, 252, 0.5, (72, 126)), (129, 203, 0.5, (64, 102)), (1, 1, 0.25, (2, 2))], +) +def test_reduced_resolution_matches_experimental_even_grid( + height: int, width: int, scale: float, expected: tuple[int, int] +) -> None: + """Keep reduced transformer-grid rounding consistent with the accepted runs.""" + assert ( + NoiseInversionOptions(resolution_scale=scale).coarse_shape(height, width) + == expected + ) + + +def test_configuration_cannot_be_mutated() -> None: + """Keep branched options chains isolated from downstream edits.""" + field_name = "steps" + with pytest.raises(FrozenInstanceError): + setattr(NoiseInversionOptions(), field_name, 8) + + +def test_invalid_source_shape_is_rejected() -> None: + """Reject impossible source dimensions before spatial projection.""" + with pytest.raises(ValueError, match="positive"): + NoiseInversionOptions().coarse_shape(0, 128) diff --git a/tests/sampling/test_noise_inversion_runtime.py b/tests/sampling/test_noise_inversion_runtime.py new file mode 100644 index 0000000..aad6d9f --- /dev/null +++ b/tests/sampling/test_noise_inversion_runtime.py @@ -0,0 +1,218 @@ +"""Verify real Comfy scaling with isolated external guider/model execution.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch +from comfy.model_sampling import CONST, EPS + +from simple_syrup.domain.noise_inversion import NoiseInversionOptions +from simple_syrup.runtime import guided_sampling, noise_inversion + + +def _flow() -> Any: + """Use the actual Comfy flow scaling object at the dynamic host boundary.""" + sampling = CONST() + sampling.sigma_max = 1.0 + sampling.noise_scale = 1.0 + return sampling + + +def _epsilon() -> Any: + """Use the actual Comfy EPS scaling object at the dynamic host boundary.""" + sampling = EPS() + sampling.sigma_max = 10.0 + return sampling + + +def _model(sampling: Any) -> Any: + """Represent only the dynamic Comfy patcher fields needed by inversion.""" + return SimpleNamespace( + load_device=torch.device("cpu"), + get_model_object=lambda name: sampling, + model=SimpleNamespace(process_latent_in=lambda x: x), + ) + + +@pytest.fixture +def guider_calls(monkeypatch: pytest.MonkeyPatch) -> list[dict[str, Any]]: + """Replace neural execution but retain the actual Comfy noise scaling contract.""" + calls: list[dict[str, Any]] = [] + ticks = iter(range(100)) + monkeypatch.setattr(noise_inversion, "_clock", lambda device: float(next(ticks))) + + def sample_boundary(**kwargs: Any) -> torch.Tensor: + """Drive the registered custom sampler with a closed-form denoiser.""" + calls.append(kwargs) + sigmas = kwargs["sigmas"] + latent = kwargs["latent_image"] + sampling = kwargs["model"].get_model_object("model_sampling") + state = sampling.noise_scaling( + sigmas[0].double(), + kwargs["noise"].double(), + latent.double(), + max_denoise=False, + ) + + def denoiser( + x: torch.Tensor, sigma: torch.Tensor, **extra: Any + ) -> torch.Tensor: + """Expose constant velocity two through Comfy's denoised prediction.""" + return x - 2 * sigma.reshape((-1,) + (1,) * (x.ndim - 1)) + + return cast( + torch.Tensor, + kwargs["sampler"].sampler_function( + denoiser, state, sigmas, extra_args={}, callback=None, disable=True + ), + ) + + monkeypatch.setattr( + guided_sampling, "sample_with_optional_negative", sample_boundary + ) + return calls + + +@pytest.mark.parametrize("negative", [None, [[torch.ones((1, 1, 1)), {}]]]) +@pytest.mark.parametrize("singleton_depth", [False, True]) +def test_default_recipe_executes_coarse_and_full_stages_with_original_conditioning( + guider_calls: list[dict[str, Any]], negative: Any, singleton_depth: bool +) -> None: + """Keep the accepted two-plus-one recipe and positive-only path intact.""" + shape = (1, 4, 1, 8, 12) if singleton_depth else (1, 4, 8, 12) + latent = torch.full(shape, 0.7) + positive = [[torch.ones((1, 1, 1)), {}]] + model = _model(_flow()) + factory_shapes: list[tuple[int, ...]] = [] + + def factory(stage_latent: torch.Tensor) -> Any: + """Record the independently planned stage dimensions.""" + factory_shapes.append(tuple(stage_latent.shape)) + return model + + result = noise_inversion.invert_sampling_noise( + model=model, + latent=latent, + forward_sigmas=torch.tensor([0.5, 0.25, 0]), + positive=positive, + negative=negative, + cfg=3, + seed=123, + options=NoiseInversionOptions(), + model_factory=factory, + ) + assert [phase.name for phase in result.stages] == ["coarse", "full_finish"] + assert [phase.steps for phase in result.stages] == [2, 1] + assert [phase.evaluations for phase in result.stages] == [2, 1] + assert factory_shapes == [shape[:-2] + (4, 6), shape] + assert len(guider_calls) == 2 + for call in guider_calls: + assert call["positive"] is positive + assert call["negative"] is negative + assert call["cfg"] == 3 and call["seed"] == 123 + assert float(guider_calls[0]["sigmas"][-1]) == 0.375 + assert float(guider_calls[1]["sigmas"][0]) == 0.375 + assert float(guider_calls[1]["sigmas"][-1]) == 0.5 + expected_endpoint = latent * (1 - 0.0001) + 2 * (0.5 - 0.0001) + reconstructed = _flow().noise_scaling(torch.tensor(0.5), result.noise, latent) + assert torch.allclose(reconstructed, expected_endpoint, atol=1e-6) + assert result.reconstruction_max_error <= 1e-6 + assert result.seconds == 5 + assert torch.equal(latent, torch.full(shape, 0.7)) + + +@pytest.mark.parametrize("sampling", [_flow(), _epsilon()]) +def test_full_size_euler_and_heun_recreate_endpoint_for_both_noise_scalings( + guider_calls: list[dict[str, Any]], sampling: Any +) -> None: + """Recover actual noise, not a flow-only formula wrongly applied to EPS models.""" + latent = torch.full((1, 2, 7, 9), 0.7) + for method in ("euler", "heun"): + result = noise_inversion.invert_sampling_noise( + model=_model(sampling), + latent=latent, + forward_sigmas=torch.tensor([0.5, 0]), + positive=[], + negative=None, + cfg=1, + seed=4, + options=NoiseInversionOptions(method=method, resolution_scale=1, steps=4), + ) + assert len(result.stages) == 1 + assert result.stages[0].latent_shape == (1, 2, 7, 9) + assert result.stages[0].evaluations == (4 if method == "euler" else 8) + initial = sampling.noise_scaling( + torch.tensor(0.0001), torch.zeros_like(latent), latent + ) + expected = initial + 2 * (0.5 - 0.0001) + assert torch.allclose( + sampling.noise_scaling(torch.tensor(0.5), result.noise.clone(), latent), + expected, + atol=1e-6, + ) + assert len(guider_calls) == 2 + + +def test_coarse_only_inversion_reaches_forward_target( + guider_calls: list[dict[str, Any]], +) -> None: + """Omit a finishing stage without stopping short of the required noise level.""" + result = noise_inversion.invert_sampling_noise( + model=_model(_flow()), + latent=torch.ones((1, 1, 8, 12)), + forward_sigmas=torch.tensor([0.5, 0]), + positive=[], + negative=None, + cfg=1, + seed=1, + options=NoiseInversionOptions(finishing_steps=0), + ) + assert len(result.stages) == len(guider_calls) == 1 + assert float(guider_calls[0]["sigmas"][-1]) == 0.5 + + +@pytest.mark.parametrize( + ("sampling", "target"), + [(_flow(), 1), (_flow(), 0), (_epsilon(), 10), (_epsilon(), 11)], +) +def test_singular_or_full_denoise_is_rejected_before_guider_execution( + guider_calls: list[dict[str, Any]], sampling: Any, target: float +) -> None: + """Do not invent source-preserving inversion at unsupported noise endpoints.""" + with pytest.raises(ValueError, match="denoise"): + noise_inversion.invert_sampling_noise( + model=_model(sampling), + latent=torch.ones((1, 1, 8, 12)), + forward_sigmas=torch.tensor([target, 0]), + positive=[], + negative=None, + cfg=1, + seed=1, + options=NoiseInversionOptions(), + ) + assert guider_calls == [] + + +def test_every_invocation_pays_inversion_cost_without_reusing_prior_noise( + guider_calls: list[dict[str, Any]], +) -> None: + """Keep production work honest even for two identical test invocations.""" + results = [ + noise_inversion.invert_sampling_noise( + model=_model(_flow()), + latent=torch.ones((1, 1, 8, 12)), + forward_sigmas=torch.tensor([0.5, 0]), + positive=[], + negative=None, + cfg=1, + seed=1, + options=NoiseInversionOptions(), + ) + for _ in range(2) + ] + assert len(guider_calls) == 4 + assert results[0].seconds == results[1].seconds == 5 + assert torch.equal(results[0].noise, results[1].noise) diff --git a/tests/sampling/test_regional_inversion_geometry.py b/tests/sampling/test_regional_inversion_geometry.py new file mode 100644 index 0000000..52db263 --- /dev/null +++ b/tests/sampling/test_regional_inversion_geometry.py @@ -0,0 +1,107 @@ +"""Verify regional-detailing masks and bounds at reduced inversion resolutions.""" + +from __future__ import annotations + +from typing import Any + +import pytest +import torch + +from simple_syrup.domain.regional_detailing import LatentBox, LatentRegion +from simple_syrup.domain.regional_inversion_geometry import project_inversion_regions +from simple_syrup.runtime.regional_inversion_model_factory import ( + RegionalInversionModelFactory, +) +from simple_syrup.runtime.regional_multidiffusion_prediction import ( + RegionalMultiDiffusionCalcCondBatch, +) + +from .support.inversion_model import SpatialModel + + +def region() -> LatentRegion: + """Keep an asymmetric region with canonical full-canvas ownership.""" + mask = torch.zeros((9, 13)) + mask[2:8, 3:12] = 1 + return LatentRegion(7, "portrait", LatentBox(3, 2, 9, 6), mask, []) + + +def test_full_resolution_preserves_original_region_identity() -> None: + """Reuse canonical geometry without resampling during full-size finishing.""" + regions = (region(),) + assert ( + project_inversion_regions( + regions, + source_width=13, + source_height=9, + target_width=13, + target_height=9, + ) + is regions + ) + + +def test_coarse_projection_preserves_conditioning_and_coverage() -> None: + """Cover the scaled bounds while retaining the exact regional prompt payload.""" + original = region() + (projected,) = project_inversion_regions( + (original,), + source_width=13, + source_height=9, + target_width=6, + target_height=4, + ) + assert projected.index == original.index and projected.label == original.label + assert projected.positive is original.positive + assert projected.latent_box == LatentBox(1, 0, 5, 4) + assert projected.latent_mask.shape == (4, 6) + assert bool(torch.all((projected.latent_mask == 0) | (projected.latent_mask == 1))) + assert bool(torch.any(projected.latent_mask == 1)) + assert original.latent_mask.shape == (9, 13) + + +@pytest.mark.parametrize("bad_box", [LatentBox(-1, 0, 2, 2), LatentBox(0, 0, 14, 2)]) +def test_invalid_canonical_bounds_fail_closed(bad_box: LatentBox) -> None: + """Do not silently clamp malformed ownership into another region.""" + original = region() + malformed = LatentRegion(7, "portrait", bad_box, original.latent_mask, []) + with pytest.raises(ValueError, match="bounds"): + project_inversion_regions( + (malformed,), + source_width=13, + source_height=9, + target_width=6, + target_height=4, + ) + + +def test_mismatched_mask_canvas_fails_closed() -> None: + """Reject masks from a different coordinate system before stage execution.""" + with pytest.raises(ValueError, match="canonical canvas"): + project_inversion_regions( + (region(),), + source_width=12, + source_height=9, + target_width=6, + target_height=4, + ) + + +def test_factory_rebuilds_each_stage_from_the_unwrapped_model() -> None: + """Install real regional wrappers without nesting a forward regional wrapper.""" + original = SpatialModel() + factory = RegionalInversionModelFactory( + model=original, + canvas_width=13, + canvas_height=9, + regions=(region(),), + global_prompt_weight=0.3, + differential_diffusion=False, + ) + for height, width in ((4, 6), (9, 13)): + derived = factory(torch.zeros((1, 4, height, width))) + assert derived.parent is original + wrapper: Any = derived.model_options["sampler_calc_cond_batch_function"] + assert isinstance(wrapper, RegionalMultiDiffusionCalcCondBatch) + assert wrapper._regions[0].latent_mask.shape == (height, width) + assert original.model_options == {} diff --git a/tests/sampling/test_sampler_options.py b/tests/sampling/test_sampler_options.py new file mode 100644 index 0000000..7c7cea4 --- /dev/null +++ b/tests/sampling/test_sampler_options.py @@ -0,0 +1,104 @@ +"""Prove graph-order independence, branch isolation, and capability conflicts.""" + +from __future__ import annotations + +from itertools import permutations +from typing import Any + +import pytest + +from simple_syrup.domain.noise_inversion import NoiseInversionOptions +from simple_syrup.domain.sampler_options import ( + AttentionCouplingOptions, + ContextualDiffusionOptions, + SamplerCapability, + SamplerOptions, + TilingOptions, + append_sampler_capability, +) + + +def test_all_distinct_capability_orders_produce_identical_configuration() -> None: + """Let users arrange features freely without altering effective settings.""" + capabilities = ( + TilingOptions(), + ContextualDiffusionOptions(), + NoiseInversionOptions(), + AttentionCouplingOptions(region_masks=object()), + ) + expected = SamplerOptions( + tiling=capabilities[0], + contextual_diffusion=capabilities[1], + noise_inversion=capabilities[2], + attention_coupling=capabilities[3], + ) + all_capabilities: tuple[SamplerCapability, ...] = capabilities + for order in permutations(all_capabilities): + options: SamplerOptions | None = None + for capability in order: + options = append_sampler_capability(options, capability) + assert options == expected + + +@pytest.mark.parametrize( + "capability", + [ + TilingOptions(), + ContextualDiffusionOptions(), + NoiseInversionOptions(), + AttentionCouplingOptions(region_masks=object()), + ], +) +def test_every_feature_can_start_a_chain_and_rejects_duplicates( + capability: SamplerCapability, +) -> None: + """Avoid mandatory empty nodes and order-dependent last-wins overrides.""" + options = append_sampler_capability(None, capability) + with pytest.raises(ValueError, match="Duplicate"): + append_sampler_capability(options, capability) + + +def test_branched_options_do_not_mutate_the_shared_upstream_configuration() -> None: + """Allow independent downstream samplers to share a configuration prefix.""" + upstream = SamplerOptions(tiling=TilingOptions()) + contextual = append_sampler_capability(upstream, ContextualDiffusionOptions()) + inverted = append_sampler_capability(upstream, NoiseInversionOptions()) + assert upstream.contextual_diffusion is None + assert upstream.noise_inversion is None + assert contextual.noise_inversion is None + assert inverted.contextual_diffusion is None + assert contextual.tiling is inverted.tiling is upstream.tiling + + +@pytest.mark.parametrize( + ("constructor", "changes"), + [ + (TilingOptions, {"width": 0}), + (TilingOptions, {"overlap": 128}), + (TilingOptions, {"diffusion_mode": "unknown"}), + (TilingOptions, {"batch_size": 0}), + (ContextualDiffusionOptions, {"global_weight": float("nan")}), + (ContextualDiffusionOptions, {"global_decay": 2}), + (ContextualDiffusionOptions, {"global_steps": -1}), + (AttentionCouplingOptions, {"region_masks": None}), + ( + AttentionCouplingOptions, + {"region_masks": object(), "regional_prompt_weight": 1.1}, + ), + ], +) +def test_invalid_feature_settings_fail_at_construction( + constructor: Any, changes: dict[str, Any] +) -> None: + """Validate capability settings before any model mutation or sampling.""" + with pytest.raises(ValueError): + constructor(**changes) + + +def test_malformed_options_connection_is_rejected() -> None: + """Fail closed for foreign objects on the custom connection boundary.""" + malformed: Any = object() + with pytest.raises(TypeError, match="Options input"): + append_sampler_capability(malformed, TilingOptions()) + with pytest.raises(TypeError, match="tiling"): + SamplerOptions(tiling=malformed) diff --git a/tests/sampling/test_sampler_options_nodes.py b/tests/sampling/test_sampler_options_nodes.py new file mode 100644 index 0000000..2236388 --- /dev/null +++ b/tests/sampling/test_sampler_options_nodes.py @@ -0,0 +1,130 @@ +"""Verify native option sockets, defaults, validation and graph composition.""" + +from __future__ import annotations + +from itertools import permutations +from typing import Any + +import pytest +import torch + +from simple_syrup.domain.noise_inversion import NoiseInversionOptions +from simple_syrup.domain.sampler_options import SamplerOptions +from simple_syrup.nodes_v3.attention_coupling_options import AttentionCouplingOptionsV3 +from simple_syrup.nodes_v3.contextual_diffusion_options import ( + ContextualDiffusionOptionsV3, +) +from simple_syrup.nodes_v3.ksampler import KSamplerV3 +from simple_syrup.nodes_v3.ksampler_contextual_diffusion import ( + KSamplerContextualDiffusionV3, +) +from simple_syrup.nodes_v3.noise_inversion_options import NoiseInversionOptionsV3 +from simple_syrup.nodes_v3.sampler_options_schema import OPTIONS_TYPE +from simple_syrup.nodes_v3.tiling_options import TilingOptionsV3 +from simple_syrup.services.sampler_options_sampling_service import SamplingArguments + +OPTION_NODES: tuple[Any, ...] = ( + TilingOptionsV3, + ContextualDiffusionOptionsV3, + NoiseInversionOptionsV3, + AttentionCouplingOptionsV3, +) + + +@pytest.mark.parametrize("node", OPTION_NODES) +def test_each_capability_can_start_a_natively_bypassable_chain(node: Any) -> None: + """Use one matching optional input/output, without an empty or enable node.""" + schema = node.define_schema() + options = next(value for value in schema.inputs if value.id == "options") + assert options.optional and options.io_type == OPTIONS_TYPE + assert len(schema.outputs) == 1 + assert schema.outputs[0].io_type == OPTIONS_TYPE + assert "enabled" not in " ".join(value.id for value in schema.inputs) + assert schema.description + assert all(value.tooltip for value in (*schema.inputs, *schema.outputs)) + + +def test_all_node_orders_preserve_configuration_and_branches() -> None: + """Build all 24 real node chains without touching MODEL preparation.""" + masks = torch.ones((1, 32, 64)) + reference: SamplerOptions | None = None + for order in permutations(OPTION_NODES): + chain: SamplerOptions | None = None + for node in order: + kwargs: dict[str, Any] = {"options": chain} + if node is AttentionCouplingOptionsV3: + kwargs["region_masks"] = masks + previous = chain + (chain,) = node.execute(**kwargs) + assert chain is not previous + assert isinstance(chain, SamplerOptions) + if reference is None: + reference = chain + else: + assert chain == reference + assert reference is not None + assert reference.noise_inversion == NoiseInversionOptions() + + +@pytest.mark.parametrize("node", OPTION_NODES) +def test_duplicate_capability_node_fails_explicitly(node: Any) -> None: + """Reject ambiguous controls rather than choosing a winner by connection order.""" + kwargs = ( + {"region_masks": torch.ones((1, 4, 4))} + if node is AttentionCouplingOptionsV3 + else {} + ) + (first,) = node.execute(**kwargs) + with pytest.raises(ValueError, match="Duplicate sampler capability"): + node.execute(options=first, **kwargs) + + +def test_inversion_node_defaults_match_convenience_controls( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Ensure both entry points construct the same accepted inversion configuration.""" + calls: list[dict[str, Any]] = [] + + class RecordingService: + """Capture the node-to-service boundary without model execution.""" + + def sample(self, **kwargs: Any) -> Any: + """Expose only a stable Comfy-facing result shape.""" + from types import SimpleNamespace + + calls.append(kwargs) + return SimpleNamespace(latent=kwargs["latent_image"], contexts=object()) + + monkeypatch.setattr( + KSamplerContextualDiffusionV3, "service_class", RecordingService + ) + (options,) = NoiseInversionOptionsV3.execute() + standard: SamplingArguments = { + "model": object(), + "seed": 1, + "steps": 4, + "cfg": 1.0, + "sampler_name": "euler", + "scheduler": "simple", + "positive": [], + "negative": None, + "latent_image": {"samples": torch.zeros((1, 4, 32, 64))}, + "denoise": 0.5, + } + KSamplerContextualDiffusionV3.execute(**standard, noise_inversion_enabled=True) + assert ( + calls[-1]["noise_inversion"] + == options.noise_inversion + == NoiseInversionOptions() + ) + KSamplerContextualDiffusionV3.execute(**standard) + assert calls[-1]["noise_inversion"] is None + assert calls[-1]["negative"] is None + + +def test_new_sampler_preserves_optional_negative_and_consumes_typed_options() -> None: + """Keep standard Comfy sampling sockets with one optional capability connection.""" + inputs = {value.id: value for value in KSamplerV3.define_schema().inputs} + assert inputs["negative"].optional + assert inputs["options"].optional and inputs["options"].io_type == OPTIONS_TYPE + assert inputs["steps"].default == 20 and inputs["cfg"].default == 8.0 diff --git a/tests/sampling/test_sampler_options_sampling_service.py b/tests/sampling/test_sampler_options_sampling_service.py new file mode 100644 index 0000000..f05b1be --- /dev/null +++ b/tests/sampling/test_sampler_options_sampling_service.py @@ -0,0 +1,196 @@ +"""Verify capability execution routing, order independence and early rejection.""" + +from __future__ import annotations + +from itertools import permutations, product +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +from comfy.model_sampling import CONST + +from simple_syrup.domain.noise_inversion import NoiseInversionOptions +from simple_syrup.domain.sampler_options import ( + AttentionCouplingOptions, + ContextualDiffusionOptions, + SamplerCapability, + SamplerOptions, + TilingOptions, +) +from simple_syrup.runtime import sampling_schedulers +from simple_syrup.services import sampler_options_sampling_service as routing + + +@pytest.fixture +def routed_calls(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, dict[str, Any]]]: + """Record sampling delegates while exercising the real configuration compiler.""" + calls: list[tuple[str, dict[str, Any]]] = [] + + def boundary(label: str) -> type: + """Represent one expensive application sampling boundary.""" + + class Recorder: + """Keep neural execution outside configuration-routing unit tests.""" + + def sample(self, **kwargs: Any) -> Any: + """Retain complete routed values and the declared result shape.""" + calls.append((label, kwargs)) + result = kwargs["latent_image"] + if label.startswith("context"): + return SimpleNamespace(latent=result, contexts=()) + return result + + return Recorder + + for attribute, label in ( + ("KSamplerSamplingService", "full"), + ("AttentionCouplingSamplingService", "attention"), + ("TiledDiffusionSamplingService", "tiled"), + ("TiledAttentionCouplingSamplingService", "tiled_attention"), + ("ContextualDiffusionSamplingService", "context"), + ("ContextualAttentionCouplingSamplingService", "context_attention"), + ): + monkeypatch.setattr(routing, attribute, boundary(label)) + monkeypatch.setattr( + sampling_schedulers, + "calculate_sigmas", + lambda **kwargs: torch.tensor([0.5, 0.25, 0.0]), + ) + return calls + + +def arguments() -> routing.SamplingArguments: + """Use real flow-target validation with a small dynamic host MODEL boundary.""" + flow = CONST() + flow.noise_scale = 1.0 + return { + "model": SimpleNamespace(get_model_object=lambda name: flow), + "seed": 17, + "steps": 4, + "cfg": 3.0, + "sampler_name": "euler", + "scheduler": "simple", + "positive": [], + "negative": None, + "latent_image": {"samples": torch.zeros((1, 4, 64, 96))}, + "denoise": 0.5, + } + + +def capabilities() -> tuple[SamplerCapability, ...]: + """Keep local geometry independent from the reduced whole-scene context.""" + return ( + TilingOptions(width=64, height=48, overlap=8, batch_size=2), + ContextualDiffusionOptions(context_size=32, global_weight=0.8, global_steps=3), + NoiseInversionOptions(), + AttentionCouplingOptions(torch.ones((1, 64, 96)), 0.7, 2), + ) + + +@pytest.mark.parametrize("enabled", tuple(product((False, True), repeat=4))) +def test_every_capability_combination_uses_one_sampling_authority( + routed_calls: list[tuple[str, dict[str, Any]]], + enabled: tuple[bool, bool, bool, bool], +) -> None: + """Compose all 16 combinations without dropping configuration or negatives.""" + options = SamplerOptions() + for capability, active in zip(capabilities(), enabled, strict=True): + if active: + options = options.with_capability(capability) + values = arguments() + output = routing.SamplerOptionsSamplingService().sample(**values, options=options) + assert output is values["latent_image"] + assert len(routed_calls) == 1 + tile, context, inversion, attention = enabled + expected = ( + ("context_attention" if attention else "context") + if context + else ("tiled_attention" if attention else "tiled") + if tile + else "attention" + if attention + else "full" + ) + label, forwarded = routed_calls[0] + assert label == expected + assert forwarded["negative"] is None + assert forwarded["noise_inversion"] == ( + NoiseInversionOptions() if inversion else None + ) + assert forwarded["model"] is values["model"] + if context: + assert forwarded["latent_context_size"] == 32 + assert forwarded["tiling"].width == (64 if tile else 32) + assert forwarded["tiling"].height == (48 if tile else 32) + elif tile: + assert forwarded["latent_tile_width"] == 64 + assert forwarded["latent_tile_height"] == 48 + if attention: + assert forwarded["regional_prompt_weight"] == 0.7 + assert forwarded["region_mask_feather"] == 2 + + +def test_all_connection_orders_compile_identical_execution( + routed_calls: list[tuple[str, dict[str, Any]]], +) -> None: + """Require all 24 capability orders to reach the same service with exact values.""" + values = arguments() + for order in permutations(capabilities()): + options = SamplerOptions() + for capability in order: + options = options.with_capability(capability) + output = routing.SamplerOptionsSamplingService().sample( + **values, options=options + ) + assert output is values["latent_image"] + assert routed_calls[-1] == routed_calls[0] + assert len(routed_calls) == 24 + + +@pytest.mark.parametrize("options", [object(), {}, "MODEL"]) +def test_foreign_options_fail_before_execution( + routed_calls: list[tuple[str, dict[str, Any]]], + options: Any, +) -> None: + """Reject arbitrary socket payloads rather than interpreting them as controls.""" + with pytest.raises(TypeError, match="options nodes"): + routing.SamplerOptionsSamplingService().sample(**arguments(), options=options) + assert not routed_calls + + +def test_invalid_inversion_target_fails_before_attention_preparation( + routed_calls: list[tuple[str, dict[str, Any]]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Reject a full-noise flow endpoint before any preparation or inference.""" + monkeypatch.setattr( + sampling_schedulers, + "calculate_sigmas", + lambda **kwargs: torch.tensor([1.0, 0.5, 0.0]), + ) + options = SamplerOptions() + for capability in capabilities(): + options = options.with_capability(capability) + with pytest.raises(ValueError, match="full-noise endpoint"): + routing.SamplerOptionsSamplingService().sample(**arguments(), options=options) + assert not routed_calls + + +@pytest.mark.parametrize( + "options", + [ + SamplerOptions(tiling=TilingOptions()), + SamplerOptions(contextual_diffusion=ContextualDiffusionOptions()), + ], +) +def test_unipc_conflicts_fail_before_sampling( + routed_calls: list[tuple[str, dict[str, Any]]], + options: SamplerOptions, +) -> None: + """Do not silently bypass tiling for an incompatible sampling method.""" + values = arguments() + values["sampler_name"] = "uni_pc" + with pytest.raises(ValueError, match="UniPC"): + routing.SamplerOptionsSamplingService().sample(**values, options=options) + assert not routed_calls diff --git a/tests/sampling/test_tiled_diffusion_conditioning_batch_service.py b/tests/sampling/test_tiled_diffusion_conditioning_batch_service.py index 1e7d940..2a8b463 100644 --- a/tests/sampling/test_tiled_diffusion_conditioning_batch_service.py +++ b/tests/sampling/test_tiled_diffusion_conditioning_batch_service.py @@ -86,7 +86,7 @@ def test_batch_service_forwards_sampling_arguments_unchanged() -> None: assert len(calls) == 1 assert { key: value for key, value in calls[0].items() if key != "latent_image" - } == expected_values | {"positive": "positive"} + } == expected_values | {"positive": "positive", "noise_inversion": None} assert torch.equal( calls[0]["latent_image"]["samples"], kwargs["latent_image"]["samples"], diff --git a/tests/sampling/test_tiled_diffusion_item_sampling_service.py b/tests/sampling/test_tiled_diffusion_item_sampling_service.py index 9bf5386..9db1699 100644 --- a/tests/sampling/test_tiled_diffusion_item_sampling_service.py +++ b/tests/sampling/test_tiled_diffusion_item_sampling_service.py @@ -95,6 +95,10 @@ def test_item_service_forwards_every_sampling_argument_unchanged( assert result is output assert calls == { key: value for key, value in kwargs.items() if key != "diffusion_mode" + } | { + "noise_inversion": None, + "inversion_segs": None, + "inversion_region_masks": None, } diff --git a/tests/sampling/test_tiled_diffusion_sampling_routing_service.py b/tests/sampling/test_tiled_diffusion_sampling_routing_service.py index 5c366b8..3d73679 100644 --- a/tests/sampling/test_tiled_diffusion_sampling_routing_service.py +++ b/tests/sampling/test_tiled_diffusion_sampling_routing_service.py @@ -224,6 +224,7 @@ def test_routing_service_sends_ordinary_request_to_item_sampler( } } expected["capability_admission"] = EMPTY_REGIONAL_CAPABILITY_ADMISSION + expected["noise_inversion"] = None assert calls[0] == expected diff --git a/tests/segmentation/segs/test_detail_segs_as_regions_node.py b/tests/segmentation/segs/test_detail_segs_as_regions_node.py index 3a778e0..494453c 100644 --- a/tests/segmentation/segs/test_detail_segs_as_regions_node.py +++ b/tests/segmentation/segs/test_detail_segs_as_regions_node.py @@ -297,6 +297,7 @@ class _FakeRegionalDetailerService: noise_mask_feather: int, tiled_encode: bool, tiled_decode: bool, + noise_inversion: object = None, ) -> DetailSEGSAsRegionsResult: """Return deterministic detailer output and record the call.""" @@ -323,6 +324,7 @@ class _FakeRegionalDetailerService: "noise_mask_feather": noise_mask_feather, "tiled_encode": tiled_encode, "tiled_decode": tiled_decode, + "noise_inversion": noise_inversion, } ) return DetailSEGSAsRegionsResult(image=cast(torch.Tensor, image) + 1.0) diff --git a/tests/segmentation/segs/test_detail_segs_as_regions_service.py b/tests/segmentation/segs/test_detail_segs_as_regions_service.py index b7c8d14..ec31852 100644 --- a/tests/segmentation/segs/test_detail_segs_as_regions_service.py +++ b/tests/segmentation/segs/test_detail_segs_as_regions_service.py @@ -481,6 +481,7 @@ class _FakeRegionalSampler: global_prompt_weight: float, preview_context: DetailPreviewContext | None = None, differential_diffusion: bool = False, + noise_inversion: object = None, ) -> Latent: """Record regional sample options and return the latent unchanged.""" @@ -500,6 +501,7 @@ class _FakeRegionalSampler: "global_prompt_weight": global_prompt_weight, "preview_context": preview_context, "differential_diffusion": differential_diffusion, + "noise_inversion": noise_inversion, } ) return latent_image diff --git a/tests/segmentation/segs/test_detail_segs_by_scale_factor_tiled_diffusion_node.py b/tests/segmentation/segs/test_detail_segs_by_scale_factor_tiled_diffusion_node.py index 0ab2ff1..3c07221 100644 --- a/tests/segmentation/segs/test_detail_segs_by_scale_factor_tiled_diffusion_node.py +++ b/tests/segmentation/segs/test_detail_segs_by_scale_factor_tiled_diffusion_node.py @@ -164,6 +164,7 @@ class _FakeTiledDetailerService: latent_tile_height: int, latent_tile_overlap: int, latent_tile_batch_size: int, + noise_inversion: object = None, ) -> TiledDetailerResult: """Return deterministic output and record tiled detailer inputs.""" @@ -194,6 +195,7 @@ class _FakeTiledDetailerService: "latent_tile_height": latent_tile_height, "latent_tile_overlap": latent_tile_overlap, "latent_tile_batch_size": latent_tile_batch_size, + "noise_inversion": noise_inversion, } ) return TiledDetailerResult(image=cast(torch.Tensor, image) + 1.0) diff --git a/tests/segmentation/segs/test_detail_segs_by_scale_factor_tiled_diffusion_service.py b/tests/segmentation/segs/test_detail_segs_by_scale_factor_tiled_diffusion_service.py index 10aa84e..04b31a5 100644 --- a/tests/segmentation/segs/test_detail_segs_by_scale_factor_tiled_diffusion_service.py +++ b/tests/segmentation/segs/test_detail_segs_by_scale_factor_tiled_diffusion_service.py @@ -287,6 +287,7 @@ class _FakeTiledSampler: latent_tile_batch_size: int, preview_context: DetailPreviewContext | None = None, differential_diffusion: bool = False, + noise_inversion: object = None, ) -> Latent: """Record tiled sample options and return the latent unchanged.""" @@ -309,6 +310,7 @@ class _FakeTiledSampler: "latent_tile_batch_size": latent_tile_batch_size, "preview_context": preview_context, "differential_diffusion": differential_diffusion, + "noise_inversion": noise_inversion, } ) return latent_image diff --git a/tests/segmentation/segs/test_tiled_detail_sampler.py b/tests/segmentation/segs/test_tiled_detail_sampler.py index 672007d..5fb1426 100644 --- a/tests/segmentation/segs/test_tiled_detail_sampler.py +++ b/tests/segmentation/segs/test_tiled_detail_sampler.py @@ -83,6 +83,7 @@ class _FakeTiledSamplingService: latent_tile_batch_size: int, preview_context: DetailPreviewContext | None = None, differential_diffusion: bool = False, + noise_inversion: object = None, ) -> Latent: """Record tiled sampling arguments and return the latent unchanged.""" @@ -105,6 +106,7 @@ class _FakeTiledSamplingService: "latent_tile_batch_size": latent_tile_batch_size, "preview_context": preview_context, "differential_diffusion": differential_diffusion, + "noise_inversion": noise_inversion, } ) return latent_image diff --git a/tests/tooling/test_inversion_experiment_parity.py b/tests/tooling/test_inversion_experiment_parity.py new file mode 100644 index 0000000..ad3a8e0 --- /dev/null +++ b/tests/tooling/test_inversion_experiment_parity.py @@ -0,0 +1,183 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Safeguard the strict live-generation completion gate for inversion promotion.""" + +from __future__ import annotations + +import hashlib +from pathlib import Path + +import pytest +from PIL import Image, PngImagePlugin + +from tools import check_inversion_experiment_parity as parity + + +def _image( + path: Path, + *, + color: tuple[int, int, int, int] = (40, 60, 80, 255), + size: tuple[int, int] = (3, 2), + metadata: str = "", +) -> Path: + """Write a bounded lossless image fixture with optional non-pixel metadata.""" + + path.parent.mkdir(parents=True, exist_ok=True) + png_info = PngImagePlugin.PngInfo() + png_info.add_text("prompt", metadata) + Image.new("RGBA", size, color).save(path, pnginfo=png_info) + return path + + +def _matrix(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path, Path]: + """Provide all six synthetic outputs while fingerprinting their references.""" + + references, candidates = tmp_path / "references", tmp_path / "candidates" + fingerprints: dict[str, str] = {} + for case in ("anima", "krea", "klein"): + reference = _image(references / case / "half2_full1" / "save_final.png") + fingerprints[case] = hashlib.sha256(reference.read_bytes()).hexdigest() + for interface in parity.INTERFACES: + _image(candidates / interface / case / "save_final.png", metadata=interface) + monkeypatch.setattr(parity, "REFERENCE_SHA256", fingerprints) + return references, candidates + + +def test_metadata_differences_do_not_change_decoded_pixel_exactness( + tmp_path: Path, +) -> None: + """Compare generated content, not PNG compression or workflow metadata.""" + + reference = _image(tmp_path / "reference.png", metadata="experimental workflow") + candidate = _image(tmp_path / "candidate.png", metadata="production workflow") + assert reference.read_bytes() != candidate.read_bytes() + observation = parity.compare_images(reference, candidate) + assert observation.exact + assert observation.changed_pixels == 0 + assert observation.maximum_channel_difference == 0 + + +@pytest.mark.parametrize("channel", range(4)) +def test_one_changed_channel_in_one_pixel_fails_without_tolerance( + tmp_path: Path, channel: int +) -> None: + """Reject a single intensity-level change, including alpha-only changes.""" + + reference = _image(tmp_path / "reference.png") + candidate = _image(tmp_path / "candidate.png") + with Image.open(candidate) as source: + changed = source.copy() + original_pixel = changed.getpixel((0, 0)) + assert isinstance(original_pixel, tuple) + pixel = list(original_pixel) + pixel[channel] -= 1 + changed.putpixel((0, 0), tuple(pixel)) + changed.save(candidate) + observation = parity.compare_images(reference, candidate) + assert not observation.exact + assert observation.changed_pixels == 1 + assert observation.maximum_channel_difference == 1 + + +def test_dimension_mismatch_is_not_resized_into_a_match(tmp_path: Path) -> None: + """Reject equal colors at different dimensions rather than resampling them.""" + + observation = parity.compare_images( + _image(tmp_path / "reference.png"), + _image(tmp_path / "candidate.png", size=(6, 4)), + ) + assert not observation.exact + assert observation.changed_pixels is None + + +def test_reference_file_cannot_be_submitted_as_production_output( + tmp_path: Path, +) -> None: + """Require a separate candidate file instead of comparing a reference to itself.""" + + reference = _image(tmp_path / "reference.png") + with pytest.raises(ValueError, match="must not be the experimental reference"): + parity.compare_images(reference, reference) + + +@pytest.mark.parametrize("mode", ("L", "I;16")) +def test_other_bit_depths_are_not_silently_quantized_for_comparison( + tmp_path: Path, mode: str +) -> None: + """Reject formats whose conversion could hide differences in decoded pixels.""" + + reference = _image(tmp_path / "reference.png") + candidate = tmp_path / "candidate.png" + Image.new(mode, (3, 2), 300 if mode == "I;16" else 40).save(candidate) + with pytest.raises(ValueError, match="lossless 8-bit RGB or RGBA PNGs"): + parity.compare_images(reference, candidate) + + +def test_gate_requires_every_case_through_both_interfaces( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Require a six-output matrix rather than one successful showcase generation.""" + + references, candidates = _matrix(tmp_path, monkeypatch) + observations = parity.check_experiment_parity( + experimental_root=references, generated_root=candidates + ) + assert [(case, interface) for case, interface, _ in observations] == [ + (case, interface) + for case in ("anima", "krea", "klein") + for interface in ("convenience", "options_stack") + ] + assert all(observation.exact for _, _, observation in observations) + + +@pytest.mark.parametrize("failure", ("missing", "changed", "reference")) +def test_cli_blocks_completion_on_missing_or_changed_evidence( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, failure: str +) -> None: + """Return failure for incomplete generation, pixel differences or baseline edits.""" + + references, candidates = _matrix(tmp_path, monkeypatch) + candidate = candidates / "options_stack" / "klein" / "save_final.png" + if failure == "missing": + candidate.unlink() + elif failure == "changed": + _image(candidate, color=(41, 60, 80, 255)) + else: + _image( + references / "anima" / "half2_full1" / "save_final.png", metadata="changed" + ) + assert ( + parity.main( + [ + "--experimental-root", + str(references), + "--generated-root", + str(candidates), + ] + ) + == 1 + ) + + +def test_cli_passes_only_the_complete_exact_matrix( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + """Report zero changed pixels for all six required output comparisons.""" + + references, candidates = _matrix(tmp_path, monkeypatch) + assert ( + parity.main( + [ + "--experimental-root", + str(references), + "--generated-root", + str(candidates), + ] + ) + == 0 + ) + output = capsys.readouterr().out + assert output.count("PIXEL-EXACT:") == 6 + assert output.count("changed pixels=0;") == 6 diff --git a/tools/attention_coupling_benchmark/comfy_probe/attention_coupling_phase_node.py b/tools/attention_coupling_benchmark/comfy_probe/attention_coupling_phase_node.py index e262357..80f145e 100644 --- a/tools/attention_coupling_benchmark/comfy_probe/attention_coupling_phase_node.py +++ b/tools/attention_coupling_benchmark/comfy_probe/attention_coupling_phase_node.py @@ -15,6 +15,7 @@ from simple_syrup.nodes_v3.ksampler_attention_coupling import ( from simple_syrup.nodes_v3.ksampler_schema import ( attention_coupling_ksampler_inputs, ) +from simple_syrup.nodes_v3.sampler_options_schema import noise_inversion_inputs from .attention_coupling_phase_profile import ( ProfiledAttentionCouplingSamplingService, @@ -49,10 +50,13 @@ class ProfiledKSamplerAttentionCouplingV3(_ComfyNodeBase): node_id="SimpleSyrupBenchmark.ProfiledKSamplerAttentionCoupling", display_name="Benchmark Profiled KSampler Attention Coupling", category="SimpleSyrup/Benchmark", - inputs=attention_coupling_ksampler_inputs( - _comfy_io, - region_masks_optional=True, - ), + inputs=[ + *attention_coupling_ksampler_inputs( + _comfy_io, + region_masks_optional=True, + ), + *noise_inversion_inputs(_comfy_io, convenience=True), + ], outputs=[_comfy_io.Latent.Output("latent")], is_dev_only=True, ) @@ -73,6 +77,13 @@ class ProfiledKSamplerAttentionCouplingV3(_ComfyNodeBase): region_masks: object | None = None, regional_prompt_weight: float = 1.0, region_mask_feather: int = 0, + noise_inversion_enabled: bool = False, + inversion_method: str = "euler", + inversion_resolution_scale: float = 0.5, + inversion_steps: int = 2, + inversion_switch_fraction: float = 0.75, + inversion_finishing_steps: int = 1, + inversion_finishing_method: str = "euler", ) -> tuple[dict[str, Any]]: """Bridge canonical host batches, then run the inherited exact delegate.""" @@ -94,4 +105,11 @@ class ProfiledKSamplerAttentionCouplingV3(_ComfyNodeBase): region_mask_feather=region_mask_feather, latent_image=latent_image, denoise=denoise, + noise_inversion_enabled=noise_inversion_enabled, + inversion_method=inversion_method, + inversion_resolution_scale=inversion_resolution_scale, + inversion_steps=inversion_steps, + inversion_switch_fraction=inversion_switch_fraction, + inversion_finishing_steps=inversion_finishing_steps, + inversion_finishing_method=inversion_finishing_method, ) diff --git a/tools/attention_coupling_benchmark/comfy_probe/attention_coupling_phase_profile.py b/tools/attention_coupling_benchmark/comfy_probe/attention_coupling_phase_profile.py index 430ed0e..73b33c4 100644 --- a/tools/attention_coupling_benchmark/comfy_probe/attention_coupling_phase_profile.py +++ b/tools/attention_coupling_benchmark/comfy_probe/attention_coupling_phase_profile.py @@ -12,6 +12,7 @@ from simple_syrup.domain.attention_coupling_request import ( AttentionCouplingRequestMode, classify_attention_coupling_request, ) +from simple_syrup.domain.noise_inversion import NoiseInversionOptions from simple_syrup.domain.regional_attention_execution import ( RegionalAttentionExecutionMode, ) @@ -55,6 +56,7 @@ class ProfiledAttentionCouplingSamplingService(AttentionCouplingSamplingService) region_mask_feather: int, latent_image: dict[str, Any], denoise: float, + noise_inversion: NoiseInversionOptions | None = None, ) -> dict[str, Any]: """Time only the production owners selected by request routing.""" @@ -77,6 +79,7 @@ class ProfiledAttentionCouplingSamplingService(AttentionCouplingSamplingService) negative=negative, latent_image=latent_image, denoise=denoise, + noise_inversion=noise_inversion, ) with measure_synchronized_phase("model_preparation_total", device=device): prepared = self.model_preparation_service_class().prepare( @@ -101,4 +104,5 @@ class ProfiledAttentionCouplingSamplingService(AttentionCouplingSamplingService) negative=prepared.negative, latent_image=latent_image, denoise=denoise, + noise_inversion=noise_inversion, ) diff --git a/tools/check_inversion_experiment_parity.py b/tools/check_inversion_experiment_parity.py new file mode 100644 index 0000000..513195e --- /dev/null +++ b/tools/check_inversion_experiment_parity.py @@ -0,0 +1,122 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Require lossless reproduction of accepted inversion images through both APIs.""" + +from __future__ import annotations + +import argparse +import hashlib +import sys +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from PIL import Image + +REFERENCE_SHA256 = { + "anima": "2b9a5932af9209651dc0bb72f317b64598a3cce5b08674eb8f485f98df689608", + "krea": "4f8b256bbb3738004f2e9dc0f677aa2a7660faad1e0a7c89ec68606b1c26ccd0", + "klein": "6485944a7633a2d2f4bf977f5b398f76f08678beea364dd6bc01d84430df7491", +} +INTERFACES = ("convenience", "options_stack") + + +@dataclass(frozen=True, slots=True) +class ImageParity: + """Describe decoded-pixel equality independently of PNG metadata.""" + + reference_size: tuple[int, int] + candidate_size: tuple[int, int] + changed_pixels: int | None + maximum_channel_difference: int | None + + @property + def exact(self) -> bool: + """Accept only equal dimensions with no changed RGBA pixel.""" + + return self.reference_size == self.candidate_size and self.changed_pixels == 0 + + +def compare_images(reference: Path, candidate: Path) -> ImageParity: + """Compare native-resolution decoded RGBA pixels without resizing or tolerance.""" + + if reference.resolve() == candidate.resolve(): + raise ValueError( + "Candidate output must not be the experimental reference file." + ) + with ( + Image.open(reference) as reference_image, + Image.open(candidate) as candidate_image, + ): + for image in (reference_image, candidate_image): + if image.format != "PNG" or image.mode not in {"RGB", "RGBA"}: + raise ValueError( + "Parity images must be lossless 8-bit RGB or RGBA PNGs." + ) + sizes = (reference_image.size, candidate_image.size) + if sizes[0] != sizes[1]: + return ImageParity(*sizes, None, None) + reference_pixels = np.asarray(reference_image.convert("RGBA"), dtype=np.int16) + candidate_pixels = np.asarray(candidate_image.convert("RGBA"), dtype=np.int16) + difference = np.abs(candidate_pixels - reference_pixels) + return ImageParity( + *sizes, + changed_pixels=int(np.count_nonzero(np.any(difference != 0, axis=-1))), + maximum_channel_difference=int(difference.max()), + ) + + +def check_experiment_parity( + *, experimental_root: Path, generated_root: Path +) -> tuple[tuple[str, str, ImageParity], ...]: + """Check all six outputs against fingerprinted, immutable winning references. + + Candidate images must be freshly generated through each production interface + before this checker is run. Missing files, changed references, size differences, + and any changed pixel block completion; visual similarity is insufficient. + """ + + observations: list[tuple[str, str, ImageParity]] = [] + for case, fingerprint in REFERENCE_SHA256.items(): + reference = experimental_root / case / "half2_full1" / "save_final.png" + if hashlib.sha256(reference.read_bytes()).hexdigest() != fingerprint: + raise ValueError(f"Frozen experimental reference changed: {case}.") + for interface in INTERFACES: + candidate = generated_root / interface / case / "save_final.png" + observations.append((case, interface, compare_images(reference, candidate))) + return tuple(observations) + + +def main(argv: Sequence[str] | None = None) -> int: + """Fail the live-generation gate unless every required output is pixel-exact.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--experimental-root", type=Path, required=True) + parser.add_argument("--generated-root", type=Path, required=True) + args = parser.parse_args(argv) + try: + observations = check_experiment_parity( + experimental_root=args.experimental_root, generated_root=args.generated_root + ) + except (OSError, ValueError) as error: + print( + f"FAILED: Inversion experiment parity is incomplete: {error}", + file=sys.stderr, + ) + return 1 + for case, interface, observation in observations: + label = "PIXEL-EXACT" if observation.exact else "FAILED" + print( + f"{label}: {case}/{interface}; " + f"dimensions={observation.candidate_size} vs {observation.reference_size}; " + f"changed pixels={observation.changed_pixels}; " + f"maximum channel difference={observation.maximum_channel_difference}" + ) + return 0 if all(observation.exact for _, _, observation in observations) else 1 + + +if __name__ == "__main__": + raise SystemExit(main())