316 lines
9.6 KiB
Python
316 lines
9.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
|
|
#
|
|
# Portions of this file incorporate behavior derived from
|
|
# multidiffusion-upscaler-for-automatic1111. See third_party/manifest.toml and
|
|
# third_party/NOTICE.md.
|
|
|
|
"""Shared runtime helpers for tiled latent denoising samplers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Sequence
|
|
from typing import Any, TypeAlias
|
|
|
|
import torch
|
|
|
|
from ..domain.tiled_diffusion import LatentTile, TiledDiffusionPlan
|
|
|
|
Latent: TypeAlias = dict[str, Any]
|
|
ApplyModel: TypeAlias = Callable[..., torch.Tensor]
|
|
ModelFunctionWrapper: TypeAlias = Callable[[ApplyModel, dict[str, Any]], torch.Tensor]
|
|
|
|
UNSUPPORTED_CONDITIONING_KEYS = frozenset({"area", "mask", "control", "gligen"})
|
|
|
|
|
|
def validate_sampling_controls(
|
|
*,
|
|
steps: int,
|
|
denoise: float,
|
|
latent_tile_width: int,
|
|
latent_tile_height: int,
|
|
latent_tile_batch_size: int,
|
|
) -> None:
|
|
"""Reject invalid KSampler and tile controls before runtime side effects."""
|
|
|
|
if steps < 1:
|
|
raise ValueError("steps must be at least 1.")
|
|
if not 0.0 <= denoise <= 1.0:
|
|
raise ValueError("denoise must be between 0 and 1.")
|
|
if latent_tile_width < 4:
|
|
raise ValueError("latent_tile_width must be at least 4.")
|
|
if latent_tile_height < 4:
|
|
raise ValueError("latent_tile_height must be at least 4.")
|
|
if latent_tile_batch_size < 1:
|
|
raise ValueError("latent_tile_batch_size must be at least 1.")
|
|
|
|
|
|
def validate_latent_samples(
|
|
latent_image: Latent,
|
|
*,
|
|
sampler_label: str,
|
|
) -> torch.Tensor:
|
|
"""Return validated samples from a ComfyUI latent dictionary."""
|
|
|
|
samples = latent_image.get("samples")
|
|
if not isinstance(samples, torch.Tensor):
|
|
raise ValueError("latent samples must be a torch tensor.")
|
|
validate_tensor_shape(samples, sampler_label=sampler_label)
|
|
return samples
|
|
|
|
|
|
def validate_tensor_shape(samples: torch.Tensor, *, sampler_label: str) -> None:
|
|
"""Reject unsupported latent tensor shapes before spatial tiling."""
|
|
|
|
if getattr(samples, "is_nested", False):
|
|
raise ValueError(
|
|
f"{sampler_label} requires non-nested latent samples shaped "
|
|
"[batch, channels, height, width] or "
|
|
"[batch, channels, 1, height, width]."
|
|
)
|
|
if samples.ndim == 4:
|
|
return
|
|
if samples.ndim == 5 and int(samples.shape[2]) == 1:
|
|
return
|
|
if samples.ndim == 5:
|
|
raise ValueError(
|
|
f"{sampler_label} 5D latent support requires a singleton third "
|
|
"axis shaped [batch, channels, 1, height, width]."
|
|
)
|
|
raise ValueError(
|
|
f"{sampler_label} requires latent samples shaped "
|
|
"[batch, channels, height, width] or "
|
|
"[batch, channels, 1, height, width]."
|
|
)
|
|
|
|
|
|
def reject_unsupported_conditioning(
|
|
conditioning: object,
|
|
*,
|
|
sampler_label: str,
|
|
) -> None:
|
|
"""Reject regional and external-control conditioning for basic tiled samplers."""
|
|
|
|
if contains_unsupported_conditioning_key(conditioning):
|
|
raise ValueError(
|
|
f"{sampler_label} does not support regional conditioning or "
|
|
"ControlNet in the first implementation."
|
|
)
|
|
|
|
|
|
def contains_unsupported_conditioning_key(value: object) -> bool:
|
|
"""Return whether a nested conditioning object contains unsupported keys."""
|
|
|
|
if isinstance(value, dict):
|
|
if any(key in UNSUPPORTED_CONDITIONING_KEYS for key in value):
|
|
return True
|
|
return any(
|
|
contains_unsupported_conditioning_key(item) for item in value.values()
|
|
)
|
|
if isinstance(value, list | tuple):
|
|
return any(contains_unsupported_conditioning_key(item) for item in value)
|
|
return False
|
|
|
|
|
|
def spatial_tile_slicer(tile: LatentTile, tensor_ndim: int) -> tuple[slice, ...]:
|
|
"""Return a slicer that crops only a tensor's final height and width axes."""
|
|
|
|
return (
|
|
(slice(None),) * (tensor_ndim - 2)
|
|
+ (slice(tile.y, tile.y + tile.height),)
|
|
+ (slice(tile.x, tile.x + tile.width),)
|
|
)
|
|
|
|
|
|
def new_spatial_weight_buffer(
|
|
x: torch.Tensor,
|
|
plan: TiledDiffusionPlan,
|
|
) -> torch.Tensor:
|
|
"""Create a float32 spatial weight buffer broadcastable over non-spatial axes."""
|
|
|
|
return torch.zeros(
|
|
(1,) * (x.ndim - 2) + (plan.latent_height, plan.latent_width),
|
|
device=x.device,
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
|
|
def make_tiled_model_args(
|
|
*,
|
|
args: dict[str, Any],
|
|
tiles: Sequence[LatentTile],
|
|
input_batch_size: int,
|
|
latent_height: int,
|
|
latent_width: int,
|
|
) -> dict[str, Any]:
|
|
"""Create apply-model args for one spatial tile batch."""
|
|
|
|
x = args["input"]
|
|
timestep = args["timestep"]
|
|
conditioning = args.get("c", {})
|
|
if not isinstance(x, torch.Tensor):
|
|
raise ValueError("tiled sampler model input must be a tensor.")
|
|
if not isinstance(timestep, torch.Tensor):
|
|
raise ValueError("tiled sampler timestep must be a tensor.")
|
|
if not isinstance(conditioning, dict):
|
|
raise ValueError("tiled sampler conditioning must be a dict.")
|
|
|
|
tiled_x = torch.cat(
|
|
[x[spatial_tile_slicer(tile, x.ndim)] for tile in tiles],
|
|
dim=0,
|
|
)
|
|
tiled_timestep = torch.cat([timestep] * len(tiles), dim=0)
|
|
tiled_conditioning = tile_conditioning(
|
|
conditioning=conditioning,
|
|
tiles=tiles,
|
|
input_batch_size=input_batch_size,
|
|
latent_height=latent_height,
|
|
latent_width=latent_width,
|
|
tiled_timestep=tiled_timestep,
|
|
)
|
|
tiled_args = args.copy()
|
|
tiled_args["input"] = tiled_x
|
|
tiled_args["timestep"] = tiled_timestep
|
|
tiled_args["c"] = tiled_conditioning
|
|
if "cond_or_uncond" in args:
|
|
tiled_args["cond_or_uncond"] = repeat_sequence(
|
|
args["cond_or_uncond"],
|
|
len(tiles),
|
|
)
|
|
return tiled_args
|
|
|
|
|
|
def tile_conditioning(
|
|
*,
|
|
conditioning: dict[str, Any],
|
|
tiles: Sequence[LatentTile],
|
|
input_batch_size: int,
|
|
latent_height: int,
|
|
latent_width: int,
|
|
tiled_timestep: torch.Tensor,
|
|
) -> dict[str, Any]:
|
|
"""Tile ComfyUI apply-model conditioning alongside latent tiles."""
|
|
|
|
tiled: dict[str, Any] = {}
|
|
for key, value in conditioning.items():
|
|
if key == "transformer_options" and isinstance(value, dict):
|
|
tiled[key] = tile_transformer_options(
|
|
value,
|
|
tile_count=len(tiles),
|
|
tiled_timestep=tiled_timestep,
|
|
)
|
|
continue
|
|
tiled[key] = tile_value(
|
|
value,
|
|
tiles=tiles,
|
|
input_batch_size=input_batch_size,
|
|
latent_height=latent_height,
|
|
latent_width=latent_width,
|
|
)
|
|
return tiled
|
|
|
|
|
|
def tile_transformer_options(
|
|
options: dict[str, Any],
|
|
*,
|
|
tile_count: int,
|
|
tiled_timestep: torch.Tensor,
|
|
) -> dict[str, Any]:
|
|
"""Repeat transformer metadata that is aligned to the model batch."""
|
|
|
|
tiled_options = options.copy()
|
|
if "cond_or_uncond" in options:
|
|
tiled_options["cond_or_uncond"] = repeat_sequence(
|
|
options["cond_or_uncond"],
|
|
tile_count,
|
|
)
|
|
if "uuids" in options:
|
|
tiled_options["uuids"] = repeat_sequence(options["uuids"], tile_count)
|
|
if "sigmas" in options:
|
|
tiled_options["sigmas"] = tiled_timestep
|
|
return tiled_options
|
|
|
|
|
|
def tile_value(
|
|
value: Any,
|
|
*,
|
|
tiles: Sequence[LatentTile],
|
|
input_batch_size: int,
|
|
latent_height: int,
|
|
latent_width: int,
|
|
) -> Any:
|
|
"""Tile tensors inside a conditioning value while preserving containers."""
|
|
|
|
if isinstance(value, torch.Tensor):
|
|
return tile_tensor(
|
|
value,
|
|
tiles=tiles,
|
|
input_batch_size=input_batch_size,
|
|
latent_height=latent_height,
|
|
latent_width=latent_width,
|
|
)
|
|
if isinstance(value, list):
|
|
return [
|
|
tile_value(
|
|
item,
|
|
tiles=tiles,
|
|
input_batch_size=input_batch_size,
|
|
latent_height=latent_height,
|
|
latent_width=latent_width,
|
|
)
|
|
for item in value
|
|
]
|
|
if isinstance(value, tuple):
|
|
return tuple(
|
|
tile_value(
|
|
item,
|
|
tiles=tiles,
|
|
input_batch_size=input_batch_size,
|
|
latent_height=latent_height,
|
|
latent_width=latent_width,
|
|
)
|
|
for item in value
|
|
)
|
|
return value
|
|
|
|
|
|
def tile_tensor(
|
|
tensor: torch.Tensor,
|
|
*,
|
|
tiles: Sequence[LatentTile],
|
|
input_batch_size: int,
|
|
latent_height: int,
|
|
latent_width: int,
|
|
) -> torch.Tensor:
|
|
"""Repeat or crop a conditioning tensor for a tile batch."""
|
|
|
|
if tensor.ndim >= 4 and tensor.shape[-2:] == (latent_height, latent_width):
|
|
return torch.cat(
|
|
[
|
|
tensor[
|
|
...,
|
|
tile.y : tile.y + tile.height,
|
|
tile.x : tile.x + tile.width,
|
|
]
|
|
for tile in tiles
|
|
],
|
|
dim=0,
|
|
)
|
|
if tensor.ndim >= 1 and tensor.shape[0] == input_batch_size:
|
|
return torch.cat([tensor] * len(tiles), dim=0)
|
|
if tensor.ndim >= 1 and tensor.shape[0] == 1:
|
|
repeat_shape = [input_batch_size * len(tiles)] + [1] * (tensor.ndim - 1)
|
|
return tensor.repeat(repeat_shape)
|
|
return tensor
|
|
|
|
|
|
def repeat_sequence(value: object, times: int) -> object:
|
|
"""Repeat a metadata sequence once per tile."""
|
|
|
|
if isinstance(value, list):
|
|
return value * times
|
|
if isinstance(value, tuple):
|
|
return value * times
|
|
return value
|