Unleash Ace-Step audio (and video?) This ComfyUI sampler is re-engineered for pure noise precision. Say goodbye to extra noise, hello to raw signal. Crafted for dedicated audio/video diffusion, not your grandma's image models. Get ready for some serious waveform power!

PingPongSampler_Custom.py

# ▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄
# ████ PINGPONGSAMPLER V0.8.15 – Optimized for Ace-Step Audio/Video! ████▓▒░
# ░▒▓ Originally cast into the void by: blepping (The OG coder, respect!)
# ▓▒░    Disassembled & warped by: MD (Machine Damage) – Because even perfect
# ▓▒░    code sometimes needs a good, hard whack.
# ▓▒░    Critical fixes & re-re-engineering by: Gemini (Google) based on user
# ▓▒░    feedback. – Turning bug reports into features, like magic... or just
# ▓▒░    really good debugging.
# ▓▒░    Completionist fixups via: devstral / qwen3 (local) – The unsung heroes
# ▓▒░    who tie up loose ends, probably fuelled by questionable energy drinks.
# ▓▒░ License: Apache 2.0 – Because sharing is caring, and we're not greedy
# ▓▒░    like those 3Dfx guys hoarding Voodoo cards.
# ░▒▓ Original Source: https://gist.github.com/blepping/b372ef6c5412080af136aad942d9d76c

# ▓▒░ Description:
#    A sampler node for ComfyUI, **specifically optimized and tested for
#    Ace-Step audio and video diffusion models**. This sampler appreciates
#    noise amplitude like a proper tracker module appreciating a well-timed
#    breakbeat. The core ancestral noise logic keeps it synchronized with
#    the original `blepping` sampler, and we've ejected parameters like
#    `s_noise` and `dynamic_noise`. Why? Because sometimes, you just need
#    the raw, unadulterated signal without extra compression artifacts
#    messing with your vibe. **While it *might* work with image models,
#    results may vary wildly and aren't guaranteed to be pretty. Consider
#    this your audio/video workhorse – it's here to party, not to render
#    your next Instagram selfie.**
#    Warning: May induce flashbacks to 256-byte intros or cause spontaneous
#    desire to optimize everything down to the last byte. Use responsibly,
#    unless you're trying to impress your old demoscene crew.

# ▓▒░ Changes:
# - V0.5 (the wobbly one):
#    * Added noise decay and stepped randomness support. Like a plasma effect,
#      but for your data – mesmerizing, chaotic, and vaguely reminiscent
#      of Windows 95 screensavers.
#    * Removed wave shape selection (sine died for your sins, square rules
#      the world of hardcore digital audio. Deal with it.)

# - V0.8 (consolidated weirdness edition):
#    * External-scheduler-driven noise decay integration. Your sampler, now with
#      a proper tempo – no more off-beat synths or dragging visuals.
#    * Added "step" mode for seed variation on each frame. Like a subtle pixel
#      shimmer, or that one glitch in the matrix you only notice after your
#      third Red Bull.
#    * Optional output clamping (for the safety-conscious. No overbright
#      scanlines here! We respect your eyeballs, mostly.)
#    * Start/end sigma index skipping (controls your entry/exit. Don't waste
#      cycles where they're not needed – efficiency is key, like in your old
#      demos where every byte counted.)
#    * Unified callback and model wrappers (no more dangling hacks or segfaults.
#      Your code should be tighter than your jeans in '98.)
#    * Constant-sampler ("is RF") detection for low-motion cases (for when
#      things get too static and you need to tell the system to chill out.)
#    * Signature cleanup to allow scheduler passthrough support. Plug 'n' play,
#      baby! Like a new ISA card that actually works on the first try.

# - V0.8.1 (Flashback to the 90s! Initial audio experiments. Remember when
#          modems screamed?):
#    * Added comprehensive frequency control parameters (as placeholders
#      initially. We tried to make it sing! Turns out, it just made high-pitched
#      noises, like a forgotten dial-up connection.)
#    * Introduced `audio_type` parameter with optimized settings for: Music,
#      Speech, Voice. (It was a phase. We all have those experimental years.)
#    * Dynamic noise control added for adaptive sampling behavior. (More squelches,
#      less clean. For when you want your output to sound like it's been through
#      a proper bitcrusher.)

# - V0.8.4 (Conceptual Frequency Implementation & UX & Updated Attribution -
#          The Great Purge):
#    * **Farewell, Frequencies!** Removed conceptual frequency masking/boosting.
#      They didn't quite vibe. Sometimes, less is more (or at least, less is less
#      confusing. Our bad, ASCII!).
#    * `Q_width` parameter also packed its bags. (Nobody knew what it did anyway,
#      it was just there, like that weird relative at family gatherings.)
#    * Updated comprehensive tooltips for all *remaining* parameters in the ComfyUI
#      node for better UX. (Because even l33t coders need tooltips. Don't lie.)
#    * **Removed direct attribution to `lks-ai/ComfyUI-StableAudioSampler`
#      as those audio-specific enhancements are no longer utilized.** We're back
#      to raw pixel pushing, baby! The sound was good, but the visuals are where
#      the real magic happens.

# - V0.8.5 (No more audio_type presets - Unplugged Session):
#    * Removed the 'audio_type' preset selector logic. All parameters are now
#      controlled directly by node inputs or explicit YAML settings, without
#      conditional overrides. (No more hand-holding, figure it out yourself,
#      you're an adult now!)

# - V0.8.6 (Critical Fixes & Re-engineering - Initial Attempt):
#    * **Ancestral Step Order Fix:** Corrected logic for `first_ancestral_step`
#      and `last_ancestral_step` to correctly determine active range, even if
#      inputs are inverted. (No more off-by-one errors, that's so last century,
#      like using Netscape Navigator.)
#    * **Noise Scaling Correction:** Eliminated redundant `sigma_next`
#      multiplication for ancestral noise, preventing overblown noise. (We hate
#      clipping! Keep your signal clean, unlike your browser history.)
#    * **Default Noise Sampler Conditioning:** Modified default `noise_sampler`
#      to condition noise by `sigma_next / sigma` for potentially better spectral
#      properties in audio. (This was a fleeting ghost of audio's past, like that
#      one obscure chiptune track you can't find anymore.)
#    * **`effective_noise_strength` Clamping:** Added clamping of
#      `effective_noise_strength` to `1.0` to prevent excessive noise injection.
#      (Keep it within bounds, like a good raster bar, or your internet bill.)
#    * **Clamping Timing:** Moved `enable_clamp_output` to apply only if
#      `sigma_next` is sufficiently small, preventing clamping of intermediate
#      noisy latents. (Don't crush those early vibes! We save the big clamps
#      for the grand finale.)

# - V0.8.7 (Buzzkill Protocol Activated - Final Refinements):
#    * **Ancestral Step Order Fix (Refined):** Ensures `first_ancestral_step`
#      and `last_ancestral_step` are always correctly interpreted as a `min` to
#      `max` range in `__init__`. (Precise pixel placement, every time. No
#      room for fuzziness here, unless it's intentional fuzziness.)
#    * **Default Noise Sampler Conditioning (Refined):** The default noise
#      sampler now scales `randn_like` by `sqrt(sigma_next / sigma)` for better
#      spectral properties, preventing high-frequency buzzing. (Still a little
#      bit of audio nostalgia, but we're mainly visual now! No unwanted modem
#      screeches.)
#    * **Ancestral Update Formula Reworked:** The core ancestral update
#      `x = denoised + sigma_gen * noise_sample * noise_multiplier` is now more
#      robust, ensuring `s_noise` and dynamic noise apply correctly without
#      over-amplifying. `sigma_gen` explicitly handles the noise difference
#      between steps. (Less crackle, more sizzle. Like a good demo party.)
#    * **`is_rf` (CONST Model) Handling:** Adjusted the `is_rf` branch to align
#      with the new ancestral update, ensuring consistent behavior across model
#      types. (No model left behind! We support all our silicon children.)
#    * **Clamping Logic Refined:** `enable_clamp_output` now applies only at the
#      very final stages of denoising, preventing premature clamping of noisy
#      latents. (We save the big clamps for the end credits, when everyone's
#      already impressed.)
#    * **`PingPongSamplerNode.get_sampler` Signature Update:** Corrected the
#      `get_sampler` method to properly accept and pass through `dynamic_noise`,
#      `noise_threshold`, `noise_gain`, `noise_decay_rate`, and `yaml_settings_str`
#      from the ComfyUI node UI. (More parameters, more power! Like a cheat code
#      for reality.)
#    * **YAML Merging Logic Refined:** Improved YAML parsing and merging to
#      correctly prioritize direct node inputs, especially for nested dictionaries
#      like `pingpong_options`. (No more config confusion, YAML is boss now!
#      It's the `.INI` file of the future, without the cryptic numbers.)

# - V0.8.8 (Ace-Step Refit - Unified Noise & Distortion Control):
#    * **Default Noise Sampler De-conditioning:** The `default_noise_sampler`
#      now *always* returns raw, unconditioned `torch.randn_like` noise. All
#      explicit scaling happens in the main sampling loop for better control.
#      (Pure noise, just the way you like it – unfiltered, untamed, like your
#      favorite underground techno.)
#    * **Unified Ancestral Noise Injection:** The `is_rf` (CONST model) branch
#      was removed for ancestral steps. All ancestral noise additions now
#      consistently use the `denoised + sigma_gen * noise_sample * noise_multiplier`
#      formula. This ensured standard ancestral behavior regardless of internal
#      model type detection. (One formula to rule them all! The pixel-perfect
#      algorithm you've always dreamed of.)

# - V0.8.9 (Ace-Step Synchronization - Reverting Noise to Original `blepping` Logic):
#    * **Ancestral Noise Logic Reverted:** The core ancestral noise application
#      logic for `is_rf` and non-`is_rf` models has been reverted to match the
#      original `blepping` implementation. This is based on user feedback that
#      the original code works better for Ace-Step. (If it ain't broke, don't
#      fix it with newfangled ideas. Sometimes, the old ways are the best ways,
#      like using a trackball for gaming.)
#        * `is_rf` (CONST model) branch now uses
#          `self.step_blend_function(denoised, scaled_noise, sigma_next)`.
#        * Non-`is_rf` branch now uses `self.step_blend_function`.
#    * **Internal Blend Modes:** `blend_mode` and `step_blend_mode` are
#      re-introduced as selectable options, defining `torch.lerp`, `a_only`,
#      or `b_only` internally without external dependencies. (Choose your fighter!
#      Like selecting your favorite character in Mortal Kombat.)
#    * **`s_noise` and Dynamic Noise Application:** These are now correctly
#      applied as a multiplier to the raw noise *before* the original `blepping`
#      ancestral formulas are used. (Noise, but with an adjustable volume knob.
#      Because sometimes you want it subtle, sometimes you want it to blow out
#      the speakers.)
#    * **No `pingpong_blend` or `external_sampler`:** These complex features
#      from the original are omitted for simplicity and focus on the core noise
#      issue. (Keep it simple, stupid. K.I.S.S! We're not trying to build a
#      space shuttle here, just awesome visuals.)

# - V0.8.10 (Precision Noise Control - Addressing `s_noise` and Dynamic Noise Sensitivity):
#    * **Refined Noise Multiplier Calculation:** The calculation of `scaled_noise_sample`
#      (which applies `s_noise` and dynamic noise) has been reworked.
#        * `s_noise` now consistently acts as the *base multiplier* for noise amplitude.
#        * Dynamic noise now applies as a *modulating factor* on top of this
#          `s_noise` base, avoiding direct additions or aggressive clamping that
#          caused artifacts. This should allow `s_noise` to work correctly at
#          values other than `1.0`.
#        * This aimed to resolve the "high freq static" from dynamic noise and
#          the "high/mid range static" when `s_noise` is not `1.0`. (This attempt
#          was unsuccessful, leading to V0.8.11. Sometimes, the demo just crashes
#          and you have to reboot.)

# - V0.8.11 (Ace-Step Compatibility - Strict Noise Application):
#    * **Removed `s_noise` and Dynamic Noise:** Based on user feedback that `s_noise`
#      values other than 1.0 and `dynamic_noise` always caused artifacts with Ace-Step,
#      these parameters and their corresponding noise scaling logic have been
#      completely removed for ancestral steps. (No more unwanted pixel snow!
#      We're going for clean, crisp output, not a blurry mess from a low-res GIF.)
#    * **Raw Noise Application:** For ancestral steps, the raw `torch.randn_like`
#      noise is now applied directly without any additional multiplicative scaling.
#      This ensures the sampler provides the exact noise characteristics that
#      Ace-Step expects from the original `blepping` implementation. (Pure,
#      unadulterated noise. Like a fresh .MOD file from the scene, straight
#      to your ears.)

# - V0.8.12 (Fix: Keyword Argument Duplication Error - Stack Overflow Prevention):
#    * **Resolved `TypeError` for `blend_function`:** Corrected how `blend_function`
#      and `step_blend_function` (and other node-specific parameters) are passed
#      from `get_sampler` through `go` to `__init__`. Parameters are now explicitly
#      `pop()`-ed from `kwargs` in `go` before being passed to `__init__`,
#      preventing duplicate keyword arguments when `**kwargs` is unpacked. This
#      resolves the `TypeError: got multiple values for keyword argument
#      'blend_function'`. (No more "multiple definition" errors. We're not coding
#      in Turbo Pascal anymore, where every error was a personal insult.)

# - V0.8.13 (Feature: YAML Overrides Node Inputs - Config Boss Mode):
#    * **Changed Parameter Override Logic:** The `yaml_settings_str` input now
#      takes precedence and *overrides* any conflicting parameters set directly
#      on the node's input fields. If a parameter is defined in both the YAML
#      string and the node input, the YAML value will be used. This provides
#      more flexible configuration management. (YAML is the new `CONFIG.SYS`,
#      controlling everything! Embrace the text-based tyranny.)

# - V0.8.14 (Maintenance: Header Refresh & Humor Injection):
#    * **Updated header attribution and version number.** Because keeping it fresh
#      is key, like a perfectly optimized 64k intro.
#    * **Clarified that audio enhancements from `lks-ai/ComfyUI-StableAudioSampler`
#      are no longer in use,** since `s_noise` and `dynamic_noise` have been
#      removed to prioritize Ace-Step image generation. This sampler is all about
#      the visuals now! (We're focusing on what truly matters: pixels, glorious pixels!)
#    * **Increased demoscene-era humor throughout comments and descriptions.** Get
#      ready for some serious nostalgia and probably a few groans.

# - V0.8.15 (Targeted Optimization - Ace-Step Audio/Video Focus):
#    * **Updated description to specifically highlight optimization for Ace-Step
#      audio and video models.**
#    * **Added a disclaimer that while it *might* work with image models,
#      results are not guaranteed.** (Don't come crying to us if your anime waifu
#      turns into a low-res pixelated mess. You've been warned.)
#    * **Incremented version number.**

# ░▒▓ Suggested usage:
#    **Primary use:** Generating high-quality audio or video with Ace-Step
#    diffusion models. Prepare for digital transcendence – it's like having
#    a supercomputer from the future, but it actually works!
#    **Secondary use (proceed with caution):** Experimenting with visuals if
#    you're feeling adventurous and don't mind unexpected results. Think of
#    it as a digital LSD trip – you might see things.
#    Caution: Excessive use may lead to synesthesia, spontaneous breakdancing,
#    or the urge to re-install Impulse Tracker.
#    If hallucinations persist, consult your nearest demoscene veteran.
#    (They've seen worse, probably while coding in assembly at 3 AM.)
# ▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄

#!/usr/bin/env python3

import torch
from tqdm.auto import trange
from comfy import model_sampling
from comfy.samplers import KSAMPLER
import numpy as np
import yaml # YAML for handling optional string inputs

# --- Internal Blend Modes ---
_INTERNAL_BLEND_MODES = {
    "lerp": torch.lerp,
    "a_only": lambda a, _b, _t: a, # Returns the first input (a)
    "b_only": lambda _a, b, _t: b  # Returns the second input (b)
}

class PingPongSampler:
    """
    A custom sampler implementing a "ping-pong" ancestral noise mixing strategy.
    Steps through provided noise levels (sigmas) and alternates between denoising
    and injecting controlled noise according to scheduler or manual settings.
    This version synchronizes noise application with the original `blepping` sampler.
    """

    CATEGORY = "sampling/custom_sampling" # This is for internal categorization, not node UI
    FUNCTION = "go" # This is the main function that KSAMPLER will call
    RETURN_TYPES = ("SAMPLER",) # This is not directly used for the sampler class itself, but for a node wrapper

    def __init__(
        self,
        model,
        x,
        sigmas,
        extra_args=None,
        callback=None,
        disable=None,
        noise_sampler=None,
        # All parameters are now passed via kwargs, to avoid duplicate arguments.
        **kwargs
    ):
        # Core inputs, directly from KSAMPLER
        self.model_ = model # Underlying diffusion model function
        self.x = x # Initial noisy tensor to sample from
        self.sigmas = sigmas # Array of noise levels to step through
        self.extra_args = extra_args.copy() if extra_args is not None else {} # Additional arguments for model (e.g., conditioning)
        self.callback_ = callback # External callback to report progress (for ComfyUI's progress bar)
        self.disable_pbar = disable # Disable progress bar if truthy (renamed for clarity)

        # Sampling control parameters, extracted from kwargs
        self.start_sigma_index = kwargs.pop("start_sigma_index", 0) # Sigma array index to start sampling from
        self.end_sigma_index = kwargs.pop("end_sigma_index", -1) # Sigma array index to end sampling early (-1 means all steps)
        self.enable_clamp_output = kwargs.pop("enable_clamp_output", False) # Clamp output to [-1,1]

        # Noise injection and random seed controls
        self.step_random_mode = kwargs.pop("step_random_mode", "off") # How to vary RNG seed per step ("off", "block", "reset", "step")
        self.step_size = kwargs.pop("step_size", 5) # Block size for block-based mode or multiplier for reset mode
        self.seed = kwargs.pop("seed", 42) # Base seed for reproducibility

        # Blend functions (already resolved to functions in .go())
        self.blend_function = kwargs.pop("blend_function", torch.lerp)
        self.step_blend_function = kwargs.pop("step_blend_function", torch.lerp)

        # Ancestral (ping-pong) operation boundaries
        pingpong_options = kwargs.pop("pingpong_options", None)
        if pingpong_options is None:
            pingpong_options= {}
        
        # Determine the number of total steps available from the sigmas schedule
        num_steps_available = len(sigmas) - 1 if len(sigmas) > 0 else 0

        # Get raw ancestral step inputs, with sensible defaults if not provided
        raw_first_ancestral_step = pingpong_options.get("first_ancestral_step", 0)
        # Default last_ancestral_step to cover full range if not specified
        raw_last_ancestral_step = pingpong_options.get("last_ancestral_step", num_steps_available - 1)
        
        # Ensure first is always <= last for the ancestral range to be valid
        # And clamp both to be within the actual number of available steps [0, num_steps_available-1]
        self.first_ancestral_step = max(0, min(raw_first_ancestral_step, raw_last_ancestral_step))
        if num_steps_available > 0:
            self.last_ancestral_step = min(num_steps_available - 1, max(raw_first_ancestral_step, raw_last_ancestral_step))
        else:
            self.last_ancestral_step = -1 # No steps available, so no ancestral steps will run


        # Detect if model uses ComfyUI CONST sampling (e.g., for Reflow-like noise injection behavior)
        self.is_rf = False
        current_model_check = model # Use a temporary variable for traversing model wrappers
        try:
            # Common pattern for ComfyUI to wrap the actual model inside `inner_model`
            while hasattr(current_model_check, 'inner_model') and current_model_check.inner_model is not None:
                current_model_check = current_model_check.inner_model
            # Check if the innermost model's sampling type is CONST
            if hasattr(current_model_check, 'model_sampling') and current_model_check.model_sampling is not None:
                self.is_rf = isinstance(current_model_check.model_sampling, model_sampling.CONST)
        except AttributeError:
            # If any attribute is missing along the chain, assume not CONST to be safe
            print("PingPongSampler Warning: Could not definitively determine model_sampling type, assuming not CONST.")
            self.is_rf = False


        # Default noise sampler: NOW RETURNS RAW, UNCONDITIONED UNIT VARIANCE NOISE
        # All scaling will happen explicitly in the __call__ method.
        self.noise_sampler = noise_sampler # Use external noise sampler if provided
        if self.noise_sampler is None:
            def default_noise_sampler(sigma, sigma_next): # Signature kept for compatibility, but sigma/sigma_next are unused here
                return torch.randn_like(x) # Returns raw random noise with the same shape as self.x
            self.noise_sampler = default_noise_sampler

        # Build noise decay array, either from external scheduler or use zeros (no decay)
        scheduler = kwargs.pop("scheduler", None) # Get scheduler from kwargs
        if num_steps_available > 0:
            if scheduler is not None and hasattr(scheduler, 'get_decay'):
                try:
                    # Attempt to get noise decay values from the external scheduler
                    arr = scheduler.get_decay(num_steps_available)
                    decay = np.asarray(arr, dtype=np.float32)
                    assert decay.shape == (num_steps_available,), (
                        f"Expected {num_steps_available} values from scheduler.get_decay, got {decay.shape}"
                    )
                    self.noise_decay = torch.tensor(decay, device=x.device)
                except Exception as e:
                    # Fallback to zeros if scheduler fails (e.g., invalid input or method missing)
                    print(f"PingPongSampler Warning: Could not get decay from scheduler: {e}. Using zeros for noise decay.")
                    self.noise_decay = torch.zeros((num_steps_available,), dtype=torch.float32, device=x.device)
            else:
                # No external scheduler provided or it lacks 'get_decay': no noise decay (all zeros)
                self.noise_decay = torch.zeros((num_steps_available,), dtype=torch.float32, device=x.device)
        else:
            # No steps, so noise_decay should be empty or not accessed
            self.noise_decay = torch.empty((0,), dtype=torch.float32, device=x.device)
        
        # Report any unhandled keyword arguments received during initialization
        if kwargs:
            print(f"PingPongSampler initialized with unused extra_options: {kwargs}")


    def _stepped_seed(self, step: int):
        """
        Determines the RNG seed for the current step based on the selected random mode.
        This provides controlled variations in noise generation across sampling steps.
        """
        if self.step_random_mode == "off":
            return None # No specific seed override, rely on global RNG state if applicable
        
        # Ensure step_size is positive to avoid division by zero
        current_step_size = self.step_size if self.step_size > 0 else 1

        if self.step_random_mode == "block":
            # Seed changes only after a certain number of steps (a "block")
            return self.seed + (step // current_step_size)
        elif self.step_random_mode == "reset":
            # Seed changes by a multiple of step_size, creating large jumps
            return self.seed + (step * current_step_size)
        elif self.step_random_mode == "step":
            # Seed changes incrementally with each step
            return self.seed + step
        else:
            # Unknown mode: fallback to base seed and warn (should not happen with ComfyUI dropdown)
            print(f"PingPongSampler Warning: Unknown step_random_mode '{self.step_random_mode}'. Using base seed.")
            return self.seed

    @classmethod
    def go(cls, # This is the entry point function that ComfyUI's KSAMPLER will call
        model, # The latent tensor representing the current state (noisy or denoised)
        x, # The latent tensor representing the current state (noisy or denoised)
        sigmas, # A tensor of sigma values defining the noise schedule
        extra_args=None, # Additional arguments to pass to the model's forward pass
        callback=None, # A callback function for ComfyUI's progress updates
        disable=None, # Flag to disable the progress bar
        noise_sampler=None, # An optional custom noise sampler function
        # All custom options from PingPongSamplerNode.INPUT_TYPES arrive here as kwargs
        **kwargs
    ):
        """
        Entrypoint for ComfyUI's KSAMPLER to initiate sampling with PingPongSampler.
        This method constructs an instance of PingPongSampler and then executes its main sampling loop.
        """
        # Extract explicit parameters from kwargs that are expected by __init__
        # And also resolve blend functions from strings
        blend_mode = kwargs.pop("blend_mode", "lerp")
        step_blend_mode = kwargs.pop("step_blend_mode", "lerp")
        
        resolved_blend_function = _INTERNAL_BLEND_MODES.get(blend_mode, torch.lerp)
        resolved_step_blend_function = _INTERNAL_BLEND_MODES.get(step_blend_mode, torch.lerp)

        # Other parameters that need to be popped for clarity and explicit passing to __init__
        # These are handled correctly by the final_options merge logic, but need to be popped
        # to prevent them from being passed as redundant kwargs to __init__.
        pingpong_options = kwargs.pop("pingpong_options", None)
        start_sigma_index = kwargs.pop("start_sigma_index", 0)
        end_sigma_index = kwargs.pop("end_sigma_index", -1)
        enable_clamp_output = kwargs.pop("enable_clamp_output", False)
        step_random_mode = kwargs.pop("step_random_mode", "off")
        step_size = kwargs.pop("step_size", 5)
        seed = kwargs.pop("seed", 42)
        scheduler = kwargs.pop("scheduler", None)

        # --- ADD THESE LINES TO POP THE REDUNDANT KEYS ---
        _ = kwargs.pop("verbose", False) # Pop verbose if it exists at top level
        _ = kwargs.pop("first_ancestral_step", -1) # Pop if it exists at top level
        _ = kwargs.pop("last_ancestral_step", -1) # Pop if it exists at top level
        # --- END ADDED LINES ---

        # yaml_settings_str is consumed by get_sampler, but ensure it's popped if it somehow reaches here
        _ = kwargs.pop("yaml_settings_str", "") 

        # Create an instance of the PingPongSampler class.
        # Pass all the required arguments explicitly.
        sampler_instance = cls(
            model=model,
            x=x,
            sigmas=sigmas,
            extra_args=extra_args,
            callback=callback,
            disable=disable,
            noise_sampler=noise_sampler, # This can be None, handled in __init__
            # Pass all the extracted parameters explicitly:
            pingpong_options=pingpong_options,
            start_sigma_index=start_sigma_index,
            end_sigma_index=end_sigma_index,
            enable_clamp_output=enable_clamp_output,
            step_random_mode=step_random_mode,
            step_size=step_size,
            seed=seed,
            scheduler=scheduler, # Pass scheduler here
            blend_function=resolved_blend_function, # Pass resolved functions here
            step_blend_function=resolved_step_blend_function, # Pass resolved functions here
            **kwargs # Pass any remaining kwargs to __init__ (should be empty now for standard operation)
        )
        # Call the __call__ method of the created sampler instance to perform the actual sampling process.
        return sampler_instance()

    def _model_denoise(self, x_tensor, sigma_scalar, **model_call_kwargs):
        """
        Wrapper around the underlying diffusion model's denoising function.
        It expands the scalar sigma to a tensor suitable for batch processing
        and combines any additional arguments before calling the model.
        """
        batch_size = x_tensor.shape[0]
        # Create a sigma tensor of the same batch size as x_tensor
        sigma_tensor = sigma_scalar * x_tensor.new_ones((batch_size,))
        # Merge extra_args from __init__ with any specific to this call
        final_extra_args = {**self.extra_args, **model_call_kwargs}
        # Call the actual diffusion model to predict the denoised latent
        return self.model_(x_tensor, sigma_tensor, **final_extra_args)

    def _do_callback(self, step_idx, current_x, current_sigma, denoised_sample):
        """
        Forwards progress information and intermediate results to ComfyUI's callback system.
        This allows the UI to display progress and intermediate images/latents.
        """
        if self.callback_:
            # Provide dictionary conforming to ComfyUI's expected callback format
            self.callback_({
                "i": step_idx, # Current step index
                "x": current_x, # Current latent state (noisy)
                "sigma": current_sigma, # Current noise level (sigma)
                "sigma_hat": current_sigma, # Denoised sigma (often same as current_sigma for simple samplers)
                "denoised": denoised_sample # The model's prediction of the denoised latent
            })

    def __call__(self): # This is the main sampling loop execution
        """
        Executes the main "ping-pong" sampling loop, iterating through sigma steps
        and applying denoising, ancestral noise injection.
        """
        x_current = self.x.clone() # Create a mutable copy of the initial latent to work on
        num_steps = len(self.sigmas) - 1 # Total number of sampling steps (transitions between sigmas)

        # Handle edge case: no steps to take if only one sigma value or less
        if num_steps <= 0:
            if self.enable_clamp_output:
                x_current = torch.clamp(x_current, -1.0, 1.0) # Apply final clamp if enabled
            return x_current
            
        # Retrieve pre-calculated and clamped ancestral step boundaries from __init__
        astart = self.first_ancestral_step
        aend = self.last_ancestral_step
        
        # Determine the actual end index for the sigma iteration, respecting user input
        # If end_sigma_index is -1, iterate through all steps up to num_steps-1
        actual_iteration_end_idx = self.end_sigma_index if self.end_sigma_index >= 0 else num_steps - 1
        # Ensure the end index does not exceed the actual number of available steps
        actual_iteration_end_idx = min(actual_iteration_end_idx, num_steps - 1)


        # Main sampling loop: iterate through each step (from 0 to num_steps-1)
        for idx in trange(num_steps, disable=self.disable_pbar):
            # Skip steps outside the user-defined start and end sigma indices
            if idx < self.start_sigma_index or idx > actual_iteration_end_idx:
                continue

            # Get the current and next sigma values from the schedule
            sigma_current, sigma_next = self.sigmas[idx], self.sigmas[idx + 1]
            
            # Call the diffusion model to get the denoised prediction for the current step
            denoised_sample = self._model_denoise(x_current, sigma_current)
            
            # Report progress to ComfyUI
            self._do_callback(idx, x_current, sigma_current, denoised_sample)

            # Clamping Logic: Apply clamping only at very low sigma_next (near the end of sampling)
            # This prevents premature clamping of intermediate noisy latents, preserving detail.
            if self.enable_clamp_output and sigma_next < 1e-3: # Threshold 1e-3 can be adjusted
                x_current = torch.clamp(x_current, -1.0, 1.0)
                break # Exit loop after final clamp, as sampling is effectively complete

            # If sigma_next is practically zero, we've reached the final denoised state
            if sigma_next <= 1e-6:
                x_current = denoised_sample # The final result is the denoised prediction
                break # Exit loop, no more noise to add or steps to take


            # Determine whether ancestral noise injection should be used for this specific step
            # `astart` and `aend` are already 0-based and clamped within the valid step range
            use_anc = (astart <= idx <= aend) if astart <= aend else False

            if not use_anc:
                # Non-Ancestral Step (DDIM-like interpolation): Reverting to original's step_blend_function
                # This path does NOT use `noise_sample` or `s_noise` or `dynamic_noise`.
                blend = sigma_next / sigma_current if sigma_current > 0 else 0.0 # Calculate interpolation factor
                x_current = self.step_blend_function(denoised_sample, x_current, blend) # Interpolate using step_blend_function
                continue # Move to the next step

            # --- Ancestral Step Logic: Injecting controlled noise (Reverted to Original `blepping` Behavior) ---
            # Set the RNG seed for reproducibility/variability if a stepped random mode is enabled
            local_seed = self._stepped_seed(idx)
            if local_seed is not None:
                torch.manual_seed(local_seed)

            # Generate new unit variance noise (always raw from default_noise_sampler)
            noise_sample = self.noise_sampler(sigma_current, sigma_next) 

            # IMPORTANT: Based on user feedback, Ace-Step appears to be highly sensitive to the
            # exact magnitude of ancestral noise. 's_noise' and 'dynamic_noise' parameters
            # caused artifacts when applied. Therefore, for ancestral steps, we apply the
            # raw, unscaled noise_sample. The necessary scaling is inherently handled by
            # `sigma_next` in the original blepping formulas or the blend function.
            final_noise_to_add = noise_sample

            # Apply the ancestral step update (ORIGINAL BLEPPING LOGIC)
            if self.is_rf: # If it's a CONST model (like some reflow/consistency models, e.g. from Stable Audio)
                x_current = self.step_blend_function(denoised_sample, final_noise_to_add, sigma_next)
            else: # For other models (e.g., EDM, standard DDPM)
                x_current = denoised_sample + final_noise_to_add * sigma_next

        # Final clamping for the very last step if it wasn't already handled by the `sigma_next < 1e-3` check in the loop.
        if self.enable_clamp_output:
             x_current = torch.clamp(x_current, -1.0, 1.0)

        return x_current


class PingPongSamplerNode:
    """
    ComfyUI node wrapper to register PingPongSampler as a custom sampler.
    This class defines the inputs and outputs that appear in the ComfyUI user interface.
    """
    CATEGORY = "sampling/custom_sampling/samplers" # Standard ComfyUI category for samplers
    RETURN_TYPES = ("SAMPLER",) # This node outputs a SAMPLER object
    FUNCTION = "get_sampler" # The class method that ComfyUI will call to get the sampler object

    @classmethod
    def INPUT_TYPES(cls):
        """
        Defines the input parameters that will be exposed in the ComfyUI node editor.
        These inputs allow users to customize the behavior of the PingPongSampler.
        """
        return {
            "required": {
                # Removed s_noise parameter
                
                # step_random_mode: How the RNG seed varies per sampling step.
                "step_random_mode":    (["off", "block", "reset", "step"], {"default": "block", "tooltip": "Controls how the RNG seed varies per sampling step.\n- 'off': Seed is constant. Predictable, but where's the fun in that?\n- 'block': Seed changes every 'step_size' frames. Like a glitch in the matrix, but intentional.\n- 'reset': Seed is reset based on 'step_size' multiplied by the frame index, offering more varied randomness.\n- 'step': Seed changes incrementally by the frame index at each step, providing subtle variations."}),
                
                # step_size: Interval for 'block' and 'reset' random modes.
                "step_size":           ("INT",    {"default": 4,    "min": 1,        "max": 100, "tooltip": "Used by 'block' and 'reset' step random modes to define the block/reset interval for the seed."}),
                
                # seed: Base random seed.
                "seed":                ("INT",    {"default": 80085,  "min": 0,        "max": 2**32 - 1, "tooltip": "Base random seed. The cosmic initializer. Change it for new universes, keep it for deja vu."}),
                
                # first_ancestral_step: Start index for ancestral noise.
                "first_ancestral_step": ("INT",    {"default": 0, "min": -1, "max": 10000, "tooltip": "The sampler step index (0-based) to begin ancestral\nnoise injection (ping-pong behavior). Use -1 to effectively disable ancestral noise if last_ancestral_step is also -1."}),
                
                # last_ancestral_step: End index for ancestral noise.
                "last_ancestral_step":  ("INT",    {"default": -1, "min": -1, "max": 10000, "tooltip": "The sampler step index (0-based) to end ancestral\nnoise injection (ping-pong behavior). Use -1 to extend ancestral noise to the end of the sampling process."}),
                
                # start_sigma_index: Sigma array index to begin sampling.
                "start_sigma_index":    ("INT",    {"default": 0,    "min": 0,        "max": 10000, "tooltip": "The index in the sigma array (denoising schedule) to begin sampling from. Allows skipping initial high-noise steps, potentially speeding up generation or changing visual character."}),
                
                # end_sigma_index: Sigma array index to end sampling.
                "end_sigma_index":      ("INT",    {"default": -1,   "min": -10000,  "max": 10000, "tooltip": "The index in the sigma array to end sampling at. -1 means sample all steps. To the bitter end, or a graceful exit? You decide."}),
                
                # enable_clamp_output: Clamp final output.
                "enable_clamp_output":  ("BOOLEAN", {"default": False, "tooltip": "If true, clamps the final output latent to the range [-1.0, 1.0]. Useful for preventing extreme values that might lead to artifacts during decoding."}),
                
                # scheduler: External noise decay scheduler.
                "scheduler":           ("SCHEDULER", {"tooltip": "Connect a ComfyUI scheduler node (e.g., KSamplerScheduler) to define the noise decay curve for the sampler. Essential for proper sampling progression. It's the tempo track for your pixels!"}),
                
                # Removed dynamic_noise parameters
                
                # blend_mode: Blend mode for ancestral steps.
                "blend_mode":          (tuple(_INTERNAL_BLEND_MODES.keys()), {"default": "lerp", "tooltip": "Blend mode to use for blending noise in ancestral steps. Defaults to 'lerp' (linear interpolation). Choose your flavor: 'lerp' (smooth blend), 'a_only' (take noise as is), 'b_only' (take other input as is). Fancy, right?"}),
                
                # step_blend_mode: Blend mode for non-ancestral steps.
                "step_blend_mode":     (tuple(_INTERNAL_BLEND_MODES.keys()), {"default": "lerp", "tooltip": "Blend mode to use for non-ancestral steps (regular denoising progression). Changing this from 'lerp' is generally not recommended unless you're feeling particularly chaotic. Like trying to render Doom on a 386SX with 2MB RAM."}),
            },
            "optional": {
                # yaml_settings_str: Optional YAML string for configuration.
                "yaml_settings_str": ("STRING", {"multiline": True, "default": "", "dynamic_prompt": False, "tooltip": "YAML string to configure sampler parameters. Parameters provided directly via the ComfyUI node's inputs will now be **OVERRIDDEN** by any corresponding values set in the YAML string. If the YAML is empty, node inputs are used. YAML is the boss now, respect its authority!"}),
            }
        }

    # This method is called by ComfyUI to get the sampler object
    def get_sampler(
        self,
        step_random_mode: str,
        step_size: int,
        seed: int,
        first_ancestral_step: int,
        last_ancestral_step: int,
        start_sigma_index: int,
        end_sigma_index: int,
        enable_clamp_output: bool,
        scheduler=None, # Scheduler is optional (can be None)
        blend_mode: str = "lerp", # Parameter: blend mode for ancestral steps
        step_blend_mode: str = "lerp", # Parameter: blend mode for non-ancestral steps
        yaml_settings_str: str = "" # Optional YAML string input
    ):
        """
        This method gathers all input parameters from the ComfyUI node,
        merges them with any provided YAML settings (prioritizing YAML),
        and returns a KSAMPLER object configured with the PingPongSampler.
        """
        # Create a dictionary with direct node inputs.
        direct_inputs = {
            "step_random_mode": step_random_mode,
            "step_size": step_size,
            "seed": seed,
            "pingpong_options": { # Nested dictionary for ancestral step control
                "first_ancestral_step": first_ancestral_step,
                "last_ancestral_step": last_ancestral_step,
            },
            "start_sigma_index": start_sigma_index,
            "end_sigma_index": end_sigma_index,
            "enable_clamp_output": enable_clamp_output,
            "scheduler": scheduler,
            "blend_mode": blend_mode, # Pass the string mode here
            "step_blend_mode": step_blend_mode, # Pass the string mode here
            "yaml_settings_str": yaml_settings_str, # Pass this through for `go` to pop it
        }

        # Initialize final_options with direct node inputs first.
        # This acts as the baseline/default if YAML is empty or doesn't specify a key.
        final_options = direct_inputs.copy()

        # Attempt to load YAML data.
        yaml_data = {}
        if yaml_settings_str:
            try:
                temp_yaml_data = yaml.safe_load(yaml_settings_str)
                if isinstance(temp_yaml_data, dict):
                    yaml_data = temp_yaml_data
                    # Convert string representations of booleans from YAML to actual booleans
                    for k, v in yaml_data.items():
                        if isinstance(v, str):
                            if v.lower() == 'true':
                                yaml_data[k] = True
                            elif v.lower() == 'false':
                                yaml_data[k] = False
            except yaml.YAMLError as e:
                print(f"WARNING: PingPongSamplerNode YAML parsing error: {e}. Using direct node inputs only.")
            except Exception as e:
                print(f"WARNING: PingPongSamplerNode unexpected error during YAML parsing: {e}. Using direct node inputs only.")

        # --- MERGE YAML DATA ON TOP OF DIRECT INPUTS ---
        # This is the key change: YAML data now overrides direct inputs.
        for key, value in yaml_data.items():
            if key == "pingpong_options" and isinstance(value, dict) and key in final_options and isinstance(final_options[key], dict):
                # If 'pingpong_options' exists in both YAML and direct inputs and are both dicts,
                # merge their contents (YAML keys now overwrite direct input keys).
                final_options[key].update(value)
            else:
                # For all other parameters, YAML value simply overwrites or adds the value.
                final_options[key] = value

        # The KSAMPLER wrapper takes the sampling function (PingPongSampler.go)
        # and an `extra_options` dictionary, whose contents are passed as **kwargs
        # to the sampling function.
        return (KSAMPLER(PingPongSampler.go, extra_options=final_options),)

# Dictionary mapping class names to identifiers used by ComfyUI
NODE_CLASS_MAPPINGS = {
    "PingPongSampler_Custom": PingPongSamplerNode,
}

# Dictionary mapping class names to display names in the ComfyUI UI
NODE_DISPLAY_NAME_MAPPINGS = {
    "PingPongSampler_Custom": "PingPong Sampler (Custom V0.8.15)", # Updated version number!
}