AEmotionStudio-ComfyUI-ShaderNoiseKSampler/core/schedule.py

"""
Sampling schedule for the standard pipeline.

One sigma schedule is built for the whole run and split into segments at the
points where new shader noise enters. Legacy instead ran every stage as its own
complete sampling pass: each stage restarted from maximum noise, so on flow
models (noise_scaling is `sigma * noise + (1 - sigma) * latent`, and sigma is
1.0 at the start) the previous stage's result was multiplied by zero. Two
sequential stages were therefore a half-length generation, not two halves of
one.

The boundaries here are indices into that single schedule, so stages become
segments of one trajectory.
"""
import math
from typing import List, Optional, Sequence, Tuple

import torch

import comfy.samplers

# A segment shorter than this cannot denoise anything meaningful. Legacy's
# injection points produced ranges like (19, 20) for 20 steps with 3 stages,
# making the final image a single step from full noise.
MIN_SEGMENT_STEPS = 2

SUPPORTED_DISTRIBUTIONS = (
    "uniform", "linear_decrease", "linear_increase", "gaussian",
    "first_stronger", "last_stronger",
)


def stage_strengths(base_strength: float, num_stages: int, distribution: str) -> List[float]:
    """Shader strength per stage, following the chosen distribution curve."""
    if num_stages <= 0:
        return []
    if num_stages == 1:
        return [base_strength]

    span = num_stages - 1
    min_factor = 0.25  # keep every stage audible rather than fading one to nothing

    if distribution == "linear_decrease":
        factors = [max(1.0 - i / span, min_factor) for i in range(num_stages)]
    elif distribution == "linear_increase":
        factors = [max(i / span, min_factor) for i in range(num_stages)]
    elif distribution == "gaussian":
        mid, std_dev = span / 2.0, num_stages / 3.0
        factors = [math.exp(-((i - mid) ** 2) / (2 * std_dev ** 2)) for i in range(num_stages)]
    elif distribution == "first_stronger":
        factors = [math.exp(-i) for i in range(num_stages)]
    elif distribution == "last_stronger":
        factors = [math.exp(i - span) for i in range(num_stages)]
    else:  # "uniform" and anything unrecognised
        factors = [1.0] * num_stages

    return [base_strength * f for f in factors]


def sequential_starts(total_steps: int, num_stages: int) -> List[int]:
    """First step of each sequential stage, spread evenly across the schedule."""
    if num_stages <= 1:
        return [0]
    return [int(i * total_steps / num_stages) for i in range(num_stages)]


def injection_points(total_steps: int, num_stages: int) -> List[int]:
    """
    Steps at which injection stages add shader noise.

    These are interior points: injecting at step 0 is what a sequential stage
    already does, and injecting at the last step leaves nothing to sample.
    """
    if num_stages <= 0:
        return []
    return [int(round((i + 1) * total_steps / (num_stages + 1))) for i in range(num_stages)]


def merge_boundaries(
    total_steps: int,
    starts: Sequence[int],
    points: Sequence[int],
    min_segment: int = MIN_SEGMENT_STEPS,
) -> List[int]:
    """
    Combine stage starts and injection points into ascending segment boundaries.

    Always begins at 0, drops duplicates, and discards any boundary that would
    leave a segment shorter than `min_segment` steps.
    """
    boundaries: List[int] = []
    for boundary in sorted({0, *starts, *points}):
        if boundary < 0 or boundary >= total_steps:
            continue
        if boundaries and boundary - boundaries[-1] < min_segment:
            continue
        boundaries.append(boundary)

    if not boundaries:
        return [0]
    # The tail is a segment too: drop trailing boundaries that would truncate it.
    while len(boundaries) > 1 and total_steps - boundaries[-1] < min_segment:
        boundaries.pop()
    return boundaries


def segments(boundaries: Sequence[int], total_steps: int) -> List[Tuple[int, int]]:
    """Turn boundaries into (start, end) step ranges covering the whole schedule."""
    if not boundaries:
        return [(0, total_steps)]
    ends = list(boundaries[1:]) + [total_steps]
    return [(start, end) for start, end in zip(boundaries, ends) if end > start]


def build_sigmas(
    model,
    steps: int,
    sampler_name: str,
    scheduler: str,
    denoise: float = 1.0,
    custom_sigmas: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """
    Build the sigma schedule for the whole run.

    Custom sigmas are used as given (they define their own trajectory). Legacy
    wrapped the model to override `model_sampling`, but ComfyUI reads the
    schedule through `model.get_model_object("model_sampling")`, which resolves
    past such a wrapper, so the supplied values never reached the sampler -- only
    their count did.

    Otherwise KSampler builds the schedule, which applies `denoise` (legacy
    hard-coded 1.0 per stage) and the discard-penultimate-sigma samplers.
    """
    if custom_sigmas is not None and torch.is_tensor(custom_sigmas) and custom_sigmas.numel() > 1:
        sigmas = custom_sigmas.detach().clone().float()
        if sigmas[0] < sigmas[-1]:  # accept ascending input, sample descending
            sigmas = sigmas.flip(0)
        return sigmas

    sampler = comfy.samplers.KSampler(
        model,
        steps=steps,
        device=model.load_device,
        sampler=sampler_name,
        scheduler=scheduler,
        denoise=denoise,
        model_options=getattr(model, "model_options", {}),
    )
    return sampler.sigmas