Artificial Sweetener
027efb8eef
Add provenance recovery, coordinated native sampler-state capture, frequency-completed noise reuse, and optional full-resolution inversion finishing. Support automatic SesquiLSR selection and native bislerp fallback for unknown latents. Verify exported Anima and Krea CD workflows with zero VAE encodes. Finishing outputs reproduce prior images exactly; reuse-only outputs match independent references at the agreed capture coordinates. All 3813 Python tests and 115 frontend tests pass alongside required quality gates.
260 lines
8.6 KiB
Python
260 lines
8.6 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""ComfyUI runtime adapter for contextual diffusion latent sampling."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from importlib import import_module
|
|
from types import ModuleType
|
|
from typing import Any, cast
|
|
|
|
import torch
|
|
|
|
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 .noise_reuse_scope import needs_finishing_model
|
|
from .patcher_lifecycle import PATCHER_LIFECYCLE, ModelMutation
|
|
from .sampling_model_types import ModelFunctionWrapper
|
|
from .tiled_sampling_validation import (
|
|
Latent,
|
|
reject_unsupported_conditioning,
|
|
validate_latent_samples,
|
|
validate_sampling_controls,
|
|
validate_tensor_shape,
|
|
)
|
|
|
|
LOGGER = get_logger(__name__)
|
|
SAMPLER_LABEL = "Contextual Diffusion"
|
|
UNIPC_SAMPLERS = frozenset({"uni_pc", "uni_pc_bh2"})
|
|
|
|
|
|
def sample_contextual_diffusion(
|
|
*,
|
|
model: Any,
|
|
seed: int,
|
|
steps: int,
|
|
cfg: float,
|
|
sampler_name: str,
|
|
scheduler: str,
|
|
positive: Any,
|
|
negative: Any,
|
|
latent_image: Latent,
|
|
denoise: float,
|
|
diffusion_mode: str,
|
|
controls: ContextualDiffusionControls,
|
|
plan: ContextualDiffusionPlan,
|
|
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.tile_width,
|
|
latent_tile_height=controls.tile_height,
|
|
latent_tile_batch_size=controls.latent_context_batch_size,
|
|
)
|
|
controls.validate()
|
|
if sampler_name in UNIPC_SAMPLERS:
|
|
raise ValueError("Contextual Diffusion is not compatible with UniPC samplers.")
|
|
reject_unsupported_conditioning(
|
|
positive,
|
|
sampler_label=SAMPLER_LABEL,
|
|
capability_admission=capability_admission,
|
|
)
|
|
reject_unsupported_conditioning(
|
|
negative,
|
|
sampler_label=SAMPLER_LABEL,
|
|
capability_admission=capability_admission,
|
|
)
|
|
|
|
sampler = sampling_samplers.resolve_sampler(sampler_name)
|
|
sigmas = sampling_schedulers.calculate_sigmas(
|
|
model=model,
|
|
scheduler_name=scheduler,
|
|
sampler_name=sampler_name,
|
|
steps=steps,
|
|
denoise=denoise,
|
|
view=sampling_schedulers.SchedulerView(
|
|
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)
|
|
comfy_sample = _comfy_sample()
|
|
comfy_utils = _comfy_utils()
|
|
latent_samples = comfy_sample.fix_empty_latent_channels(
|
|
model,
|
|
latent_samples,
|
|
latent_image.get("downscale_ratio_spacial", None),
|
|
)
|
|
validate_tensor_shape(latent_samples, sampler_label=SAMPLER_LABEL)
|
|
if latent_samples.shape[-2:] != (plan.latent_height, plan.latent_width):
|
|
raise ValueError(
|
|
"Contextual Diffusion plan dimensions must match the sampled latent shape."
|
|
)
|
|
|
|
sampling_model = clone_model_with_contextual_diffusion(
|
|
model,
|
|
plan=plan,
|
|
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 or needs_finishing_model()
|
|
else None
|
|
)
|
|
noise = sampling_noise.prepare_sampling_noise(
|
|
comfy_sample=comfy_sample,
|
|
sampler_name=sampler_name,
|
|
samples=latent_samples,
|
|
seed=seed,
|
|
batch_indices=batch_inds,
|
|
model=sampling_model,
|
|
)
|
|
callback = _latent_preview().prepare_callback(sampling_model, steps)
|
|
samples = sample_with_optional_negative(
|
|
comfy_sample=comfy_sample,
|
|
model=sampling_model,
|
|
noise=noise,
|
|
cfg=cfg,
|
|
sampler=sampler,
|
|
sigmas=sigmas,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent_image=latent_samples,
|
|
noise_mask=latent_image.get("noise_mask"),
|
|
callback=callback,
|
|
disable_pbar=not comfy_utils.PROGRESS_BAR_ENABLED,
|
|
seed=seed,
|
|
noise_inversion=noise_inversion,
|
|
inversion_model_factory=inversion_factory,
|
|
)
|
|
|
|
LOGGER.info(
|
|
"KSampler Contextual Diffusion pass completed",
|
|
extra={
|
|
"operation": "ksampler_contextual_diffusion",
|
|
"sampler": sampler_name,
|
|
"scheduler": scheduler,
|
|
"steps": steps,
|
|
"denoise": denoise,
|
|
"diffusion_mode": diffusion_mode,
|
|
"latent_width": plan.latent_width,
|
|
"latent_height": plan.latent_height,
|
|
"context_size": controls.latent_context_size,
|
|
"overlap": controls.latent_context_overlap,
|
|
"tile_count": len(plan.tile_plan.tiles),
|
|
"segs_guided": any(
|
|
tile.weight_mask is not None for tile in plan.tile_plan.tiles
|
|
),
|
|
"model_calls_per_prediction": (
|
|
1
|
|
if len(plan.tile_plan.tiles) == 1
|
|
else len(plan.tile_plan.batches) + int(controls.global_weight > 0)
|
|
),
|
|
},
|
|
)
|
|
output = latent_image.copy()
|
|
output.pop("downscale_ratio_spacial", None)
|
|
output["samples"] = samples
|
|
return output
|
|
|
|
|
|
def clone_model_with_contextual_diffusion(
|
|
model: Any,
|
|
*,
|
|
plan: ContextualDiffusionPlan,
|
|
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 = 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(
|
|
plan=plan,
|
|
controls=controls,
|
|
sigmas=sigmas,
|
|
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,
|
|
mutations,
|
|
operation="SimpleSyrup contextual diffusion",
|
|
)
|
|
|
|
|
|
def _comfy_sample() -> ModuleType:
|
|
"""Import ComfyUI sample helpers lazily."""
|
|
|
|
return import_module("comfy.sample")
|
|
|
|
|
|
def _comfy_utils() -> ModuleType:
|
|
"""Import ComfyUI progress state lazily."""
|
|
|
|
return import_module("comfy.utils")
|
|
|
|
|
|
def _latent_preview() -> ModuleType:
|
|
"""Import ComfyUI preview helpers lazily."""
|
|
|
|
return import_module("latent_preview")
|