Artificial Sweetener
f89b436a09
Capture the nearest existing source model evaluation without inserting sampling steps, retain its actual noise coordinate, and recover residuals with stable native scaling. Observe source and receiver states at ComfyUI's shared boundaries, discover compatible named/custom source nodes, and report graph node identities and selected samplers in failures. Admit the complete sampler catalog through CD and tiled routes with regression coverage for future additions. Verify the isolated CD-only tree: 4322 Python tests passed, 48 skipped; 115 frontend tests passed; format, lint, strict typing, architecture, and test governance passed. Contextual Attention remains outside this commit.
258 lines
8.4 KiB
Python
258 lines
8.4 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"
|
|
|
|
|
|
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()
|
|
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,
|
|
sampler_name=sampler_name,
|
|
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")
|