fix: cfg -> apg, noise

2026-02-08 10:32:13 -05:00
parent e47d6cc091
commit f9ce68f058
4 changed files with 366 additions and 61 deletions
@@ -1,16 +1,151 @@
 import torch
+import torch.nn.functional as F
 import comfy.samplers
 import comfy.sample
+import comfy.sampler_helpers
+import comfy.model_sampling
+import comfy.model_patcher
 import math
 from .ace_step_utils import ACEStepLatentUtils
 from . import logger

+
+# =============================================================================
+# APG (Adaptive Prompt Guidance) - from ACE-Step reference implementation
+# Reference: https://huggingface.co/ACE-Step/acestep-v15-base/blob/main/apg_guidance.py
+# =============================================================================
+
+class MomentumBuffer:
+    """Momentum buffer for APG guidance smoothing across denoising steps."""
+
+    def __init__(self, momentum: float = -0.75):
+        self.momentum = momentum
+        self.running_average = 0
+
+    def update(self, update_value: torch.Tensor):
+        new_average = self.momentum * self.running_average
+        self.running_average = update_value + new_average
+
+
+def _apg_project(v0: torch.Tensor, v1: torch.Tensor, dims=[-1]):
+    """Project v0 onto v1 and return parallel and orthogonal components."""
+    dtype = v0.dtype
+    device_type = v0.device.type
+    if device_type == "mps":
+        v0, v1 = v0.cpu(), v1.cpu()
+
+    v0, v1 = v0.double(), v1.double()
+    v1 = F.normalize(v1, dim=dims)
+    v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1
+    v0_orthogonal = v0 - v0_parallel
+    return v0_parallel.to(dtype).to(device_type), v0_orthogonal.to(dtype).to(device_type)
+
+
+def _apg_forward(
+    pred_cond: torch.Tensor,
+    pred_uncond: torch.Tensor,
+    guidance_scale: float,
+    momentum_buffer: MomentumBuffer = None,
+    eta: float = 0.0,
+    norm_threshold: float = 2.5,
+    dims=[-1],
+):
+    """
+    Adaptive Prompt Guidance (APG) - replaces standard CFG.
+
+    Instead of standard CFG (uncond + scale * (cond - uncond)), APG:
+    1. Computes guidance direction with momentum smoothing
+    2. Clamps guidance norm to prevent explosion
+    3. Projects guidance perpendicular to conditional prediction
+    4. Applies scaled perpendicular guidance
+
+    This prevents the guided result from drifting off the valid latent manifold,
+    which is the likely cause of silence with standard CFG at high scales.
+    """
+    diff = pred_cond - pred_uncond
+    if momentum_buffer is not None:
+        momentum_buffer.update(diff)
+        diff = momentum_buffer.running_average
+
+    if norm_threshold > 0:
+        ones = torch.ones_like(diff)
+        diff_norm = diff.norm(p=2, dim=dims, keepdim=True)
+        scale_factor = torch.minimum(ones, norm_threshold / diff_norm)
+        diff = diff * scale_factor
+
+    diff_parallel, diff_orthogonal = _apg_project(diff, pred_cond, dims)
+    normalized_update = diff_orthogonal + eta * diff_parallel
+    pred_guided = pred_cond + (guidance_scale - 1) * normalized_update
+    return pred_guided
+
+
+def _create_apg_sampler_cfg_function(momentum_buffer):
+    """
+    Create a sampler_cfg_function that uses APG instead of standard CFG.
+
+    ComfyUI's cfg_function provides args["cond"] = x - denoised = v * sigma (noise pred).
+    The sampler_cfg_function must return a noise prediction (same domain).
+
+    For ComfyUI format [B, C, T], dims=[-1] normalizes along T (time),
+    matching the reference's dims=[1] for [B, T, C] format.
+    """
+    def apg_cfg_function(args):
+        cond = args["cond"]
+        uncond = args["uncond"]
+        guidance_scale = args["cond_scale"]
+
+        guided = _apg_forward(
+            pred_cond=cond,
+            pred_uncond=uncond,
+            guidance_scale=guidance_scale,
+            momentum_buffer=momentum_buffer,
+            eta=0.0,
+            norm_threshold=2.5,
+            dims=[-1],
+        )
+
+        return guided
+
+    return apg_cfg_function
+
+
 # NOTE: gigantic shoutout to the powerful acestep team https://github.com/ace-step/ACE-Step
 # NOTE: And another massive shoutout.... the adapted code here is based on https://github.com/billwuhao/ComfyUI_ACE-Step

 # NOTE: this implementation is experimental and beta with a minimum of testing


+# =============================================================================
+# Reference Audio Timbre Injection Helper
+# =============================================================================
+
+def _inject_reference_audio(c, reference_latent, input_batch_size, device, dtype):
+    """Inject reference audio timbre latent into model conditioning dict.
+
+    Overrides the silence latent that extra_conds places in c["refer_audio"],
+    providing real timbre information to the timbre encoder without triggering
+    the covers-mode logic in extra_conds.
+
+    Truncates or pads the reference to match the existing refer_audio length,
+    and handles CFG batching.
+    """
+    ref = reference_latent.to(device=device, dtype=dtype)
+    # Match length to whatever extra_conds already set (sized to noise length)
+    if "refer_audio" in c:
+        target_len = c["refer_audio"].shape[2]
+    else:
+        target_len = ref.shape[2]
+    ref_len = ref.shape[2]
+    if ref_len > target_len:
+        ref = ref[:, :, :target_len]
+    elif ref_len < target_len:
+        ref = torch.nn.functional.pad(ref, (0, target_len - ref_len))
+    # Handle CFG batching
+    if ref.shape[0] < input_batch_size:
+        ref = ref.repeat(input_batch_size, 1, 1)
+    c["refer_audio"] = ref
+
+
 # =============================================================================
 # ACE-Step 1.5 Native Edit Guider (unified extend + repaint)
 # =============================================================================
@@ -34,12 +169,14 @@ class ACEStep15NativeEditGuider(comfy.samplers.CFGGuider):

    def __init__(self, model, positive, negative, cfg, source_latent,
                 extend_left_seconds=0.0, extend_right_seconds=0.0,
-                 repaint_start_seconds=None, repaint_end_seconds=None):
+                 repaint_start_seconds=None, repaint_end_seconds=None,
+                 reference_latent=None):
        super().__init__(model)
        self.set_conds(positive, negative)
        self.set_cfg(cfg)

        self.source_latent = source_latent
+        self.reference_latent = reference_latent

        # Load silence_latent internally
        from .ace_step_utils import load_silence_latent
@@ -176,6 +313,7 @@ class ACEStep15NativeEditGuider(comfy.samplers.CFGGuider):

        chunk_masks = self.chunk_masks
        src_latents = self.src_latents
+        reference_latent = self.reference_latent

        # Clear any existing wrapper to avoid stale closure values
        self.model_patcher.set_model_unet_function_wrapper(None)
@@ -199,6 +337,10 @@ class ACEStep15NativeEditGuider(comfy.samplers.CFGGuider):
            # This tells the model to use src_latents directly (not semantic hints)
            c["is_covers"] = torch.zeros((input_batch_size,), device=device, dtype=torch.long)

+            # Inject reference audio timbre if provided
+            if reference_latent is not None:
+                _inject_reference_audio(c, reference_latent, input_batch_size, device, dtype)
+
            if not hasattr(model_function_wrapper, '_logged'):
                logger.debug(f"[ACE15_EDIT_WRAPPER] Injecting chunk_masks and src_latents")
                logger.debug(f"[ACE15_EDIT_WRAPPER]   input.shape: {args['input'].shape}")
@@ -206,6 +348,7 @@ class ACEStep15NativeEditGuider(comfy.samplers.CFGGuider):
                logger.debug(f"[ACE15_EDIT_WRAPPER]   src_latents.shape: {c['src_latents'].shape}")
                logger.debug(f"[ACE15_EDIT_WRAPPER]   is_covers: {c['is_covers']} (0=extend/repaint mode)")
                logger.debug(f"[ACE15_EDIT_WRAPPER]   chunk_masks sum: {cm.sum().item():.0f} (generation frames * batch * channels)")
+                logger.debug(f"[ACE15_EDIT_WRAPPER]   reference_audio: {'injected' if reference_latent is not None else 'none (using silence)'}")
                model_function_wrapper._logged = True

            return apply_model(args["input"], args["timestep"], **c)
@@ -216,7 +359,7 @@ class ACEStep15NativeEditGuider(comfy.samplers.CFGGuider):
        logger.debug(f"[ACE15_EDIT] Model wrapper applied")

    def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
-        """Override sample to apply model wrapper and use working latent"""
+        """Override sample to apply model wrapper, APG guidance, and use working latent"""

        logger.debug(f"[ACE15_EDIT] sample() called")
        logger.debug(f"[ACE15_EDIT]   noise.shape: {noise.shape}")
@@ -224,6 +367,14 @@ class ACEStep15NativeEditGuider(comfy.samplers.CFGGuider):

        self._apply_model_wrapper()

+        # Apply APG guidance for base model (cfg > 1)
+        if self.cfg > 1.0:
+            momentum_buffer = MomentumBuffer(momentum=-0.75)
+            apg_cfg_fn = _create_apg_sampler_cfg_function(momentum_buffer)
+            self.model_options = comfy.model_patcher.create_model_options_clone(self.model_options)
+            self.model_options["sampler_cfg_function"] = apg_cfg_fn
+            logger.debug(f"[ACE15_EDIT] APG guidance enabled (cfg={self.cfg})")
+
        # Generate noise for working latent shape
        device = noise.device
        dtype = noise.dtype
@@ -270,12 +421,13 @@ class ACEStep15NativeCoverGuider(comfy.samplers.CFGGuider):
    sources), use ACEStep15SemanticExtractor and ACEStep15SemanticHintsBlend.
    """

-    def __init__(self, model, positive, negative, cfg, source_latent, semantic_hints=None):
+    def __init__(self, model, positive, negative, cfg, source_latent, semantic_hints=None, reference_latent=None):
        super().__init__(model)
        self.set_conds(positive, negative)
        self.set_cfg(cfg)

        self.source_latent = source_latent
+        self.reference_latent = reference_latent

        batch_size, channels, total_length = source_latent.shape

@@ -319,6 +471,7 @@ class ACEStep15NativeCoverGuider(comfy.samplers.CFGGuider):
        chunk_masks = self.chunk_masks
        src_latents = self.src_latents
        semantic_hints = self.semantic_hints
+        reference_latent = self.reference_latent

        self.model_patcher.set_model_unet_function_wrapper(None)

@@ -349,6 +502,10 @@ class ACEStep15NativeCoverGuider(comfy.samplers.CFGGuider):
                # Fallback: use raw VAE latents via is_covers=0
                c["is_covers"] = torch.zeros((input_batch_size,), device=device, dtype=torch.long)

+            # Inject reference audio timbre if provided
+            if reference_latent is not None:
+                _inject_reference_audio(c, reference_latent, input_batch_size, device, dtype)
+
            if not hasattr(model_function_wrapper, '_logged'):
                logger.debug(f"[ACE15_COVER_WRAPPER] Injecting cover parameters")
                logger.debug(f"[ACE15_COVER_WRAPPER]   input.shape: {args['input'].shape}")
@@ -367,6 +524,7 @@ class ACEStep15NativeCoverGuider(comfy.samplers.CFGGuider):
                        logger.debug(f"[ACE15_COVER_WRAPPER]   WARNING: semantic hints have very low variance! May be constant/invalid.")
                else:
                    logger.debug(f"[ACE15_COVER_WRAPPER]   WARNING: No precomputed_lm_hints_25Hz - will fall back to VAE latents!")
+                logger.debug(f"[ACE15_COVER_WRAPPER]   reference_audio: {'injected' if reference_latent is not None else 'none (using silence)'}")
                model_function_wrapper._logged = True

            return apply_model(args["input"], args["timestep"], **c)
@@ -377,7 +535,7 @@ class ACEStep15NativeCoverGuider(comfy.samplers.CFGGuider):
        logger.debug(f"[ACE15_NATIVE_COVER] Model wrapper applied")

    def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
-        """Override sample to apply model wrapper for cover"""
+        """Override sample to apply model wrapper and APG guidance for cover"""

        logger.debug(f"[ACE15_NATIVE_COVER] sample() called")
        logger.debug(f"[ACE15_NATIVE_COVER]   noise.shape: {noise.shape}")
@@ -385,12 +543,16 @@ class ACEStep15NativeCoverGuider(comfy.samplers.CFGGuider):

        self._apply_model_wrapper()

-        device = noise.device
-        dtype = noise.dtype
-        working_noise = torch.randn_like(self.source_latent, device=device, dtype=dtype)
+        # Apply APG guidance for base model (cfg > 1)
+        if self.cfg > 1.0:
+            momentum_buffer = MomentumBuffer(momentum=-0.75)
+            apg_cfg_fn = _create_apg_sampler_cfg_function(momentum_buffer)
+            self.model_options = comfy.model_patcher.create_model_options_clone(self.model_options)
+            self.model_options["sampler_cfg_function"] = apg_cfg_fn
+            logger.debug(f"[ACE15_NATIVE_COVER] APG guidance enabled (cfg={self.cfg})")

        result = super().sample(
-            working_noise,
+            noise,
            self.source_latent,
            sampler,
            sigmas,
@@ -422,13 +584,14 @@ class ACEStep15NativeExtractGuider(comfy.samplers.CFGGuider):
    the source_latent. For advanced workflows, use ACEStep15SemanticExtractor.
    """

-    def __init__(self, model, positive, negative, cfg, source_latent, track_name="vocals", semantic_hints=None):
+    def __init__(self, model, positive, negative, cfg, source_latent, track_name="vocals", semantic_hints=None, reference_latent=None):
        super().__init__(model)
        self.set_conds(positive, negative)
        self.set_cfg(cfg)

        self.source_latent = source_latent
        self.track_name = track_name
+        self.reference_latent = reference_latent

        batch_size, channels, total_length = source_latent.shape

@@ -436,18 +599,9 @@ class ACEStep15NativeExtractGuider(comfy.samplers.CFGGuider):
        logger.debug(f"[ACE15_NATIVE_EXTRACT]   source_latent.shape: {source_latent.shape}")
        logger.debug(f"[ACE15_NATIVE_EXTRACT]   track_name: {track_name}")

-        # Auto-extract semantic hints if not provided
-        if semantic_hints is None:
-            logger.debug(f"[ACE15_NATIVE_EXTRACT]   No semantic_hints provided, auto-extracting...")
-            from .ace_step_utils import extract_semantic_hints
-            semantic_hints = extract_semantic_hints(model, source_latent, verbose=True)
-            logger.debug(f"[ACE15_NATIVE_EXTRACT]   Auto-extracted semantic_hints shape: {semantic_hints.shape}")
-
+        # Store semantic hints if provided (optional, not used for context when is_covers=0)
        self.semantic_hints = semantic_hints

-        # Validate semantic hints
-        logger.debug(f"[ACE15_NATIVE_EXTRACT]   semantic_hints stats: mean={semantic_hints.mean():.4f}, std={semantic_hints.std():.4f}")
-
        # Extract: generate everything (the extracted track)
        self.chunk_masks = torch.ones_like(source_latent)
        self.src_latents = source_latent.clone()
@@ -457,13 +611,22 @@ class ACEStep15NativeExtractGuider(comfy.samplers.CFGGuider):
        self._wrapper_applied = False

    def _apply_model_wrapper(self):
-        """Apply model wrapper to inject chunk_masks, src_latents, is_covers, and semantic hints"""
+        """Apply model wrapper to inject chunk_masks, src_latents, and is_covers for extract.
+
+        Reference behavior for extract (from handler.py):
+        - is_covers = False (tensor zeros) -> prepare_condition keeps src_latents as raw audio
+        - src_latents = VAE-encoded source audio (the mix to extract from)
+        - chunk_masks = all ones (generate everything)
+        - precomputed_lm_hints_25Hz = NOT provided (reference lets prepare_condition
+          tokenize refer_audio internally; hints are computed but unused when is_covers=0)
+        """
        if self._wrapper_applied:
            return

        chunk_masks = self.chunk_masks
        src_latents = self.src_latents
        semantic_hints = self.semantic_hints
+        reference_latent = self.reference_latent

        self.model_patcher.set_model_unet_function_wrapper(None)

@@ -481,22 +644,24 @@ class ACEStep15NativeExtractGuider(comfy.samplers.CFGGuider):

            c["chunk_masks"] = cm
            c["src_latents"] = sl
+            c["is_covers"] = torch.zeros((input_batch_size,), device=device, dtype=torch.long)

-            # Set is_covers based on whether we have semantic hints
+            # Optionally provide precomputed semantic hints to skip internal tokenization
+            # (hints are NOT used for context when is_covers=0, but providing them avoids
+            # a wasteful tokenize/detokenize of the silence latent in prepare_condition)
            if semantic_hints is not None:
                sh = semantic_hints.to(device=device, dtype=dtype)
                if sh.shape[0] < input_batch_size:
                    sh = sh.repeat(input_batch_size, 1, 1)
                c["precomputed_lm_hints_25Hz"] = sh
-                c["is_covers"] = torch.ones((input_batch_size,), device=device, dtype=torch.long)
-            else:
-                c["is_covers"] = torch.zeros((input_batch_size,), device=device, dtype=torch.long)
+
+            # Inject reference audio timbre if provided
+            if reference_latent is not None:
+                _inject_reference_audio(c, reference_latent, input_batch_size, device, dtype)

            if not hasattr(model_function_wrapper, '_logged'):
-                logger.debug(f"[ACE15_EXTRACT_WRAPPER] Injecting extract parameters")
-                logger.debug(f"[ACE15_EXTRACT_WRAPPER]   is_covers: {c['is_covers']}")
-                if "precomputed_lm_hints_25Hz" in c:
-                    logger.debug(f"[ACE15_EXTRACT_WRAPPER]   precomputed_lm_hints_25Hz.shape: {c['precomputed_lm_hints_25Hz'].shape}")
+                logger.debug(f"[ACE15_EXTRACT_WRAPPER] is_covers={c['is_covers']}, "
+                             f"hints={'yes' if 'precomputed_lm_hints_25Hz' in c else 'no'}")
                model_function_wrapper._logged = True

            return apply_model(args["input"], args["timestep"], **c)
@@ -507,18 +672,22 @@ class ACEStep15NativeExtractGuider(comfy.samplers.CFGGuider):
        logger.debug(f"[ACE15_NATIVE_EXTRACT] Model wrapper applied")

    def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
-        """Override sample to apply model wrapper for extract"""
+        """Override sample to apply model wrapper and APG guidance for extract"""

        logger.debug(f"[ACE15_NATIVE_EXTRACT] sample() called")

        self._apply_model_wrapper()

-        device = noise.device
-        dtype = noise.dtype
-        working_noise = torch.randn_like(self.source_latent, device=device, dtype=dtype)
+        # Apply APG guidance for base model (cfg > 1)
+        if self.cfg > 1.0:
+            momentum_buffer = MomentumBuffer(momentum=-0.75)
+            apg_cfg_fn = _create_apg_sampler_cfg_function(momentum_buffer)
+            self.model_options = comfy.model_patcher.create_model_options_clone(self.model_options)
+            self.model_options["sampler_cfg_function"] = apg_cfg_fn
+            logger.debug(f"[ACE15_NATIVE_EXTRACT] APG guidance enabled (cfg={self.cfg})")

        result = super().sample(
-            working_noise,
+            noise,
            self.source_latent,
            sampler,
            sigmas,
@@ -550,13 +719,14 @@ class ACEStep15NativeLegoGuider(comfy.samplers.CFGGuider):
    """

    def __init__(self, model, positive, negative, cfg, source_latent,
-                 track_name="vocals", start_seconds=0.0, end_seconds=None):
+                 track_name="vocals", start_seconds=0.0, end_seconds=None, reference_latent=None):
        super().__init__(model)
        self.set_conds(positive, negative)
        self.set_cfg(cfg)

        self.source_latent = source_latent
        self.track_name = track_name
+        self.reference_latent = reference_latent

        # Load silence_latent internally
        from .ace_step_utils import load_silence_latent
@@ -627,6 +797,7 @@ class ACEStep15NativeLegoGuider(comfy.samplers.CFGGuider):

        chunk_masks = self.chunk_masks
        src_latents = self.src_latents
+        reference_latent = self.reference_latent

        self.model_patcher.set_model_unet_function_wrapper(None)

@@ -644,9 +815,18 @@ class ACEStep15NativeLegoGuider(comfy.samplers.CFGGuider):

            c["chunk_masks"] = cm
            c["src_latents"] = sl
+            # Lego task: is_covers=False (like extract/repaint, not cover)
+            # Using tensor zeros prevents stock prepare_condition's `is_covers is False`
+            # identity check from replacing src_latents with refer_audio.
+            c["is_covers"] = torch.zeros((input_batch_size,), device=device, dtype=torch.long)
+
+            # Inject reference audio timbre if provided
+            if reference_latent is not None:
+                _inject_reference_audio(c, reference_latent, input_batch_size, device, dtype)

            if not hasattr(model_function_wrapper, '_logged'):
-                logger.debug(f"[ACE15_LEGO_WRAPPER] Injecting chunk_masks and src_latents")
+                logger.debug(f"[ACE15_LEGO_WRAPPER] Injecting chunk_masks, src_latents, is_covers=0")
+                logger.debug(f"[ACE15_LEGO_WRAPPER]   reference_audio: {'injected' if reference_latent is not None else 'none (using silence)'}")
                model_function_wrapper._logged = True

            return apply_model(args["input"], args["timestep"], **c)
@@ -657,12 +837,20 @@ class ACEStep15NativeLegoGuider(comfy.samplers.CFGGuider):
        logger.debug(f"[ACE15_NATIVE_LEGO] Model wrapper applied")

    def sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
-        """Override sample to apply model wrapper for lego"""
+        """Override sample to apply model wrapper, APG guidance, and correct sigma schedule for lego"""

        logger.debug(f"[ACE15_NATIVE_LEGO] sample() called")

        self._apply_model_wrapper()

+        # Apply APG guidance for base model (cfg > 1)
+        if self.cfg > 1.0:
+            momentum_buffer = MomentumBuffer(momentum=-0.75)
+            apg_cfg_fn = _create_apg_sampler_cfg_function(momentum_buffer)
+            self.model_options = comfy.model_patcher.create_model_options_clone(self.model_options)
+            self.model_options["sampler_cfg_function"] = apg_cfg_fn
+            logger.debug(f"[ACE15_NATIVE_LEGO] APG guidance enabled (cfg={self.cfg})")
+
        device = noise.device
        dtype = noise.dtype
        working_noise = torch.randn_like(self.source_latent, device=device, dtype=dtype)
@@ -811,6 +999,13 @@ class ACEStepRepaintGuider(comfy.samplers.CFGGuider):
            if callback is not None:
                return callback(step, x0_unused, x, total_steps)
        
+        # Use APG guidance for base model (cfg > 1.0) instead of standard CFG
+        if self.cfg > 1.0:
+            momentum_buffer = MomentumBuffer(momentum=-0.75)
+            apg_cfg_fn = _create_apg_sampler_cfg_function(momentum_buffer)
+            self.model_options = comfy.model_patcher.create_model_options_clone(self.model_options)
+            self.model_options["sampler_cfg_function"] = apg_cfg_fn
+
        # Run sampling with repaint callback
        result = super().sample(noise, latent_image, sampler, sigmas, denoise_mask, repaint_callback, disable_pbar, seed)
        
@@ -1071,6 +1266,13 @@ class ACEStepExtendGuider(comfy.samplers.CFGGuider):
            if callback is not None:
                return callback(step, x0_unused, x, total_steps)
        
+        # Use APG guidance for base model (cfg > 1.0) instead of standard CFG
+        if self.cfg > 1.0:
+            momentum_buffer = MomentumBuffer(momentum=-0.75)
+            apg_cfg_fn = _create_apg_sampler_cfg_function(momentum_buffer)
+            self.model_options = comfy.model_patcher.create_model_options_clone(self.model_options)
+            self.model_options["sampler_cfg_function"] = apg_cfg_fn
+
        # Use extended latent and target_latents as initial noise
        result = super().sample(target_latents, self.extended_latent, sampler, sigmas, denoise_mask, extend_callback, disable_pbar, seed)
        
@@ -1378,6 +1580,13 @@ class ACEStepHybridGuider(comfy.samplers.CFGGuider):
            if callback is not None:
                return callback(step, x0_unused, x, total_steps)
        
+        # Use APG guidance for base model (cfg > 1.0) instead of standard CFG
+        if self.cfg > 1.0:
+            momentum_buffer = MomentumBuffer(momentum=-0.75)
+            apg_cfg_fn = _create_apg_sampler_cfg_function(momentum_buffer)
+            self.model_options = comfy.model_patcher.create_model_options_clone(self.model_options)
+            self.model_options["sampler_cfg_function"] = apg_cfg_fn
+
        # Run sampling with working latent and hybrid callback
        result = super().sample(working_noise, working_latent, sampler, sigmas, denoise_mask, hybrid_callback, disable_pbar, seed)
        
@@ -108,6 +108,8 @@ def extract_semantic_hints(model, source_latent, verbose=True):
    Returns:
        torch.Tensor: Semantic hints in ComfyUI format [B, 64, T]
    """
+    print(f"[DIAG_SEMANTIC_EXTRACT] input latent shape={source_latent.shape}, mean={source_latent.mean():.4f}, std={source_latent.std():.4f}")
+
    prefix = "[SEMANTIC_EXTRACT]"

    # Validate v1.5 shape: (batch, 64, length)
@@ -166,6 +168,8 @@ def extract_semantic_hints(model, source_latent, verbose=True):
    # Transpose back to ComfyUI format: [B, T, D] → [B, D, T]
    semantic_hints = lm_hints.movedim(-1, -2)

+    print(f"[DIAG_SEMANTIC_EXTRACT] output hints shape={semantic_hints.shape}, mean={semantic_hints.mean():.4f}, std={semantic_hints.std():.4f}")
+
    return semantic_hints


@@ -1,5 +1,6 @@
 import torch
 import comfy.model_management
+import comfy.model_sampling
 import comfy.samplers

 # Import patches module to apply ACE-Step 1.5 patches
@@ -10,8 +11,8 @@ from .ace_step_guiders import (
    ACEStepRepaintGuider, ACEStepExtendGuider, ACEStepHybridGuider,
    ACEStep15NativeEditGuider,
    ACEStep15NativeCoverGuider,
-    # TODO: Re-enable when ready
-    # ACEStep15NativeExtractGuider, ACEStep15NativeLegoGuider
+    # ACEStep15NativeExtractGuider,  # TODO: temporarily disabled pending fixes
+    # ACEStep15NativeLegoGuider,     # TODO: temporarily disabled pending fixes
 )
 from .ace_step_utils import ACEStepLatentUtils
 from .audio_mask_nodes import (
@@ -19,6 +20,7 @@ from .audio_mask_nodes import (
    AUDIO_MASK_NODE_CLASS_MAPPINGS, AUDIO_MASK_NODE_DISPLAY_NAME_MAPPINGS
 )

+
 def validate_audio_latent(latents):
    """Validate that latents are audio latents, handling both typed and untyped cases
    Supports both v1.0 (batch, 8, 16, length) and v1.5 (batch, 64, length) shapes
@@ -708,6 +710,7 @@ class ACEStep15NativeEditGuiderNode:
                "extend_right_seconds": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 240.0, "step": 0.1}),
                "repaint_start_seconds": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 1000.0, "step": 0.1}),
                "repaint_end_seconds": ("FLOAT", {"default": -1.0, "min": -1.0, "max": 1000.0, "step": 0.1}),
+                "reference_latent": ("LATENT", {"tooltip": "Optional reference audio latent for timbre conditioning. The model will generate audio with the timbre/instrument character of this reference. Does not affect task behavior (extend/repaint)."}),
            }
        }

@@ -718,7 +721,8 @@ class ACEStep15NativeEditGuiderNode:

    def get_guider(self, model, positive, negative, source_latents, cfg,
                   extend_left_seconds=0.0, extend_right_seconds=0.0,
-                   repaint_start_seconds=-1.0, repaint_end_seconds=-1.0):
+                   repaint_start_seconds=-1.0, repaint_end_seconds=-1.0,
+                   reference_latent=None):

        source_tensor = source_latents["samples"]

@@ -744,6 +748,9 @@ class ACEStep15NativeEditGuiderNode:
        if has_repaint:
            logger.debug(f"[ACE15_EDIT_NODE]   repaint: {repaint_start_seconds}s - {repaint_end_seconds}s")

+        # Extract reference latent tensor if provided
+        ref_tensor = reference_latent["samples"] if reference_latent is not None else None
+
        # Create the unified guider (silence_latent is loaded internally)
        guider = ACEStep15NativeEditGuider(
            model, positive, negative, cfg,
@@ -751,7 +758,8 @@ class ACEStep15NativeEditGuiderNode:
            extend_left_seconds=extend_left_seconds,
            extend_right_seconds=extend_right_seconds,
            repaint_start_seconds=repaint_start,
-            repaint_end_seconds=repaint_end
+            repaint_end_seconds=repaint_end,
+            reference_latent=ref_tensor
        )

        # Return the working latent (may be extended or same as source)
@@ -784,6 +792,7 @@ class ACEStep15NativeCoverGuiderNode:
            },
            "optional": {
                "semantic_hints": ("SEMANTIC_HINTS",),
+                "reference_latent": ("LATENT", {"tooltip": "Optional reference audio latent for timbre conditioning. Decouples timbre from the source audio's semantic content — the cover follows the source's structure but adopts this reference's instrument/voice character."}),
            }
        }

@@ -791,22 +800,27 @@ class ACEStep15NativeCoverGuiderNode:
    FUNCTION = "get_guider"
    CATEGORY = "sampling/custom_sampling/guiders"

-    def get_guider(self, model, positive, negative, cfg, source_latents, semantic_hints=None):
+    def get_guider(self, model, positive, negative, cfg, source_latents, semantic_hints=None, reference_latent=None):
        source_tensor = source_latents["samples"]

        # Validate v1.5 shape: (batch, 64, length)
        if len(source_tensor.shape) != 3 or source_tensor.shape[1] != 64:
            raise ValueError(f"ACE-Step 1.5 requires latent shape (batch, 64, length), got {source_tensor.shape}")

+        # Extract reference latent tensor if provided
+        ref_tensor = reference_latent["samples"] if reference_latent is not None else None
+
        logger.debug(f"[ACE15_COVER_NODE] Creating guider")
        logger.debug(f"[ACE15_COVER_NODE]   source_tensor.shape: {source_tensor.shape}")
        if semantic_hints is not None:
            logger.debug(f"[ACE15_COVER_NODE]   semantic_hints provided externally: {semantic_hints.shape}")
        else:
            logger.debug(f"[ACE15_COVER_NODE]   semantic_hints will be auto-extracted")
+        if ref_tensor is not None:
+            logger.debug(f"[ACE15_COVER_NODE]   reference_latent provided: {ref_tensor.shape}")

        guider = ACEStep15NativeCoverGuider(
-            model, positive, negative, cfg, source_tensor, semantic_hints
+            model, positive, negative, cfg, source_tensor, semantic_hints, reference_latent=ref_tensor
        )

        return (guider,)
@@ -830,13 +844,15 @@ class ACEStep15NativeExtractGuiderNode:
                "model": ("MODEL",),
                "positive": ("CONDITIONING",),
                "negative": ("CONDITIONING",),
-                "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
+                "cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01,
+                         "tooltip": "Classifier-free guidance scale. Base model default is 7.0. Turbo model uses 1.0 (no CFG)."}),
                "source_latents": ("LATENT",),
                "track_name": (["vocals", "drums", "bass", "guitar", "keyboard", "strings",
                               "percussion", "synth", "fx", "brass", "woodwinds", "backing_vocals"],),
            },
            "optional": {
                "semantic_hints": ("SEMANTIC_HINTS",),
+                "reference_latent": ("LATENT", {"tooltip": "Optional reference audio latent for timbre conditioning. Guides what the extracted track should sound like."}),
            }
        }

@@ -844,23 +860,28 @@ class ACEStep15NativeExtractGuiderNode:
    FUNCTION = "get_guider"
    CATEGORY = "sampling/custom_sampling/guiders"

-    def get_guider(self, model, positive, negative, cfg, source_latents, track_name, semantic_hints=None):
+    def get_guider(self, model, positive, negative, cfg, source_latents, track_name, semantic_hints=None, reference_latent=None):
        source_tensor = source_latents["samples"]

        # Validate v1.5 shape: (batch, 64, length)
        if len(source_tensor.shape) != 3 or source_tensor.shape[1] != 64:
            raise ValueError(f"ACE-Step 1.5 requires latent shape (batch, 64, length), got {source_tensor.shape}")

+        # Extract reference latent tensor if provided
+        ref_tensor = reference_latent["samples"] if reference_latent is not None else None
+
        logger.debug(f"[ACE15_EXTRACT_NODE] Creating guider")
        logger.debug(f"[ACE15_EXTRACT_NODE]   source_tensor.shape: {source_tensor.shape}")
        logger.debug(f"[ACE15_EXTRACT_NODE]   track_name: {track_name}")
        if semantic_hints is not None:
            logger.debug(f"[ACE15_EXTRACT_NODE]   semantic_hints provided externally: {semantic_hints.shape}")
        else:
-            logger.debug(f"[ACE15_EXTRACT_NODE]   semantic_hints will be auto-extracted")
+            logger.debug(f"[ACE15_EXTRACT_NODE]   no semantic_hints (not needed for extract)")
+        if ref_tensor is not None:
+            logger.debug(f"[ACE15_EXTRACT_NODE]   reference_latent provided: {ref_tensor.shape}")

        guider = ACEStep15NativeExtractGuider(
-            model, positive, negative, cfg, source_tensor, track_name, semantic_hints
+            model, positive, negative, cfg, source_tensor, track_name, semantic_hints, reference_latent=ref_tensor
        )

        return (guider,)
@@ -885,12 +906,16 @@ class ACEStep15NativeLegoGuiderNode:
                "model": ("MODEL",),
                "positive": ("CONDITIONING",),
                "negative": ("CONDITIONING",),
-                "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
+                "cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01,
+                         "tooltip": "Classifier-free guidance scale. Base model default is 7.0. Turbo model uses 1.0 (no CFG)."}),
                "source_latents": ("LATENT",),
                "track_name": (["vocals", "drums", "bass", "guitar", "keyboard", "strings",
                               "percussion", "synth", "fx", "brass", "woodwinds", "backing_vocals"],),
                "start_seconds": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.1}),
                "end_seconds": ("FLOAT", {"default": 30.0, "min": 0.0, "max": 1000.0, "step": 0.1}),
+            },
+            "optional": {
+                "reference_latent": ("LATENT", {"tooltip": "Optional reference audio latent for timbre conditioning. The generated track will adopt the timbre/instrument character of this reference."}),
            }
        }

@@ -899,7 +924,7 @@ class ACEStep15NativeLegoGuiderNode:
    CATEGORY = "sampling/custom_sampling/guiders"

    def get_guider(self, model, positive, negative, cfg, source_latents,
-                   track_name, start_seconds, end_seconds):
+                   track_name, start_seconds, end_seconds, reference_latent=None):
        if start_seconds >= end_seconds:
            raise ValueError(f"start_seconds ({start_seconds}) must be less than end_seconds ({end_seconds})")

@@ -909,15 +934,20 @@ class ACEStep15NativeLegoGuiderNode:
        if len(source_tensor.shape) != 3 or source_tensor.shape[1] != 64:
            raise ValueError(f"ACE-Step 1.5 requires latent shape (batch, 64, length), got {source_tensor.shape}")

+        # Extract reference latent tensor if provided
+        ref_tensor = reference_latent["samples"] if reference_latent is not None else None
+
        logger.debug(f"[ACE15_LEGO_NODE] Creating guider")
        logger.debug(f"[ACE15_LEGO_NODE]   source_tensor.shape: {source_tensor.shape}")
        logger.debug(f"[ACE15_LEGO_NODE]   track_name: {track_name}")
        logger.debug(f"[ACE15_LEGO_NODE]   region: {start_seconds}s - {end_seconds}s")
+        if ref_tensor is not None:
+            logger.debug(f"[ACE15_LEGO_NODE]   reference_latent provided: {ref_tensor.shape}")

        # Create guider (silence_latent is loaded internally)
        guider = ACEStep15NativeLegoGuider(
            model, positive, negative, cfg, source_tensor,
-            track_name, start_seconds, end_seconds
+            track_name, start_seconds, end_seconds, reference_latent=ref_tensor
        )

        return (guider,)
@@ -973,8 +1003,8 @@ class ACEStep15TaskTextEncodeNode:
            "required": {
                "clip": ("CLIP",),
                "text": ("STRING", {"multiline": True, "dynamicPrompts": True, "default": "A melodic electronic track with soft synths"}),
-                "task_type": (["text2music", "repaint", "cover"],
-                             {"tooltip": "text2music/repaint use LM audio code generation (cfg_scale, temperature, top_p, top_k apply). cover uses precomputed semantic hints from source audio instead."}),  # TODO: Re-enable "extract", "lego" when ready
+                "task_type": (["text2music", "repaint", "cover", "extract", "lego"],
+                             {"tooltip": "text2music/repaint use LM audio code generation (cfg_scale, temperature, top_p, top_k apply). cover/extract/lego use precomputed semantic hints from source audio instead."}),
            },
            "optional": {
                "track_name": (["", "vocals", "drums", "bass", "guitar", "keyboard", "strings",
@@ -1060,6 +1090,43 @@ class ACEStep15TaskTextEncodeNode:
        return (conditioning,)


+class ModelSamplingACEStep:
+    """Override model sampling for ACE-Step 1.5.
+
+    ACE-Step uses ModelSamplingDiscreteFlow with multiplier=1.0.
+    The shift parameter controls the sigma schedule:
+    - shift=3.0 for turbo model (default in supported_models)
+    - shift=1.0 for base model (linear spacing)
+
+    Place this node BEFORE the scheduler so sigmas are computed with the correct shift.
+    """
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "model": ("MODEL",),
+            "shift": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
+        }}
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
+    CATEGORY = "audio/acestep"
+
+    def patch(self, model, shift):
+        m = model.clone()
+
+        sampling_base = comfy.model_sampling.ModelSamplingDiscreteFlow
+        sampling_type = comfy.model_sampling.CONST
+
+        class ModelSamplingAdvanced(sampling_base, sampling_type):
+            pass
+
+        model_sampling = ModelSamplingAdvanced(model.model.model_config)
+        model_sampling.set_parameters(shift=shift, multiplier=1.0)
+        m.add_object_patch("model_sampling", model_sampling)
+        return (m,)
+
+
 # Node class mappings for ComfyUI
 NODE_CLASS_MAPPINGS = {
    # ACE-Step 1.0 guiders (latent-level ODE blending)
@@ -1077,11 +1144,12 @@ NODE_CLASS_MAPPINGS = {
    # ACE-Step 1.5 guiders (model-level mask input)
    "ACEStep15NativeEditGuider": ACEStep15NativeEditGuiderNode,
    "ACEStep15NativeCoverGuider": ACEStep15NativeCoverGuiderNode,
-    # TODO: Re-enable when ready
-    # "ACEStep15NativeExtractGuider": ACEStep15NativeExtractGuiderNode,
-    # "ACEStep15NativeLegoGuider": ACEStep15NativeLegoGuiderNode,
+    # "ACEStep15NativeExtractGuider": ACEStep15NativeExtractGuiderNode,  # TODO: temporarily disabled
+    # "ACEStep15NativeLegoGuider": ACEStep15NativeLegoGuiderNode,      # TODO: temporarily disabled
    # ACE-Step 1.5 text encoder
    "ACEStep15TaskTextEncode": ACEStep15TaskTextEncodeNode,
+    # ACE-Step 1.5 model sampling override
+    "ModelSamplingACEStep": ModelSamplingACEStep,
    **AUDIO_MASK_NODE_CLASS_MAPPINGS,  # Add audio mask nodes
 }

@@ -1102,10 +1170,11 @@ NODE_DISPLAY_NAME_MAPPINGS = {
    # ACE-Step 1.5 guiders
    "ACEStep15NativeEditGuider": "ACE-Step 1.5 Edit Guider (Extend/Repaint)",
    "ACEStep15NativeCoverGuider": "ACE-Step 1.5 Cover Guider",
-    # TODO: Re-enable when ready
-    # "ACEStep15NativeExtractGuider": "ACE-Step 1.5 Extract Guider",
-    # "ACEStep15NativeLegoGuider": "ACE-Step 1.5 Lego Guider",
+    # "ACEStep15NativeExtractGuider": "ACE-Step 1.5 Extract Guider",  # TODO: temporarily disabled
+    # "ACEStep15NativeLegoGuider": "ACE-Step 1.5 Lego Guider",      # TODO: temporarily disabled
    # ACE-Step 1.5 text encoder
    "ACEStep15TaskTextEncode": "ACE-Step 1.5 Task Text Encode",
+    # ACE-Step 1.5 model sampling override
+    "ModelSamplingACEStep": "ACE-Step Model Sampling",
    **AUDIO_MASK_NODE_DISPLAY_NAME_MAPPINGS,
 } 
@@ -102,12 +102,22 @@ def _create_patched_ace_step15_forward(original_forward):
    @functools.wraps(original_forward)
    def patched_forward(self, x, timestep, context, lyric_embed=None, refer_audio=None,
                        audio_codes=None, chunk_masks=None, src_latents=None,
-                        is_covers=None, precomputed_lm_hints_25Hz=None, **kwargs):
+                        is_covers=None, precomputed_lm_hints_25Hz=None,
+                        replace_with_null_embeds=False, **kwargs):
        text_attention_mask = None
        lyric_attention_mask = None
        refer_audio_order_mask = None
        attention_mask = None

+        if not hasattr(patched_forward, '_diag_count'):
+            patched_forward._diag_count = 0
+        patched_forward._diag_count += 1
+        if patched_forward._diag_count <= 2:
+            print(f"[PATCHED_FWD] Call #{patched_forward._diag_count}: "
+                  f"x={x.shape}, t={timestep}, is_covers={is_covers}, "
+                  f"hints={'yes' if precomputed_lm_hints_25Hz is not None else 'no'}, "
+                  f"null_embeds={replace_with_null_embeds}")
+
        # One-time diagnostic log for cover task
        if not _forward_logged[0] and precomputed_lm_hints_25Hz is not None:
            logger.debug(f"[ACE15_PATCHED_FORWARD] Cover mode detected")
@@ -170,6 +180,9 @@ def _create_patched_ace_step15_forward(original_forward):
            if isinstance(is_covers, torch.Tensor) and is_covers.any().item():
                is_covers = True

+        if patched_forward._diag_count <= 2 and precomputed_lm_hints_25Hz is not None:
+            print(f"[PATCHED_FWD]   is_covers_adapted={is_covers}, hints.shape={precomputed_lm_hints_25Hz.shape}")
+
        enc_hidden, enc_mask, context_latents = self.prepare_condition(
            text_hidden_states, text_attention_mask,
            lyric_hidden_states, lyric_attention_mask,
@@ -179,6 +192,13 @@ def _create_patched_ace_step15_forward(original_forward):
            audio_codes=audio_codes
        )

+        # Apply learned null embeddings for uncond CFG pass (matching stock forward behavior)
+        if replace_with_null_embeds:
+            enc_hidden[:] = self.null_condition_emb.to(enc_hidden)
+
+        if patched_forward._diag_count <= 2:
+            print(f"[PATCHED_FWD]   enc_hidden: std={enc_hidden.std():.4f}, ctx_latents: shape={context_latents.shape}")
+
        out = self.decoder(
            hidden_states=x,
            timestep=timestep,
@@ -189,6 +209,9 @@ def _create_patched_ace_step15_forward(original_forward):
            context_latents=context_latents
        )

+        if patched_forward._diag_count <= 2:
+            print(f"[PATCHED_FWD]   decoder out: mean={out.mean():.4f}, std={out.std():.4f}")
+
        return out.movedim(-1, -2)

    return patched_forward
@@ -214,11 +237,11 @@ def _create_patched_tokenizer(original_tokenize_with_weights):

        # Strip our custom kwargs before forwarding to the stock tokenizer
        stock_kwargs = {k: v for k, v in kwargs.items() if k not in ("task_type", "track_name")}
-        print(f"[ACE15_PATCHED_TOKENIZER] stock_kwargs LM params: cfg_scale={stock_kwargs.get('cfg_scale', 'MISSING')}, temperature={stock_kwargs.get('temperature', 'MISSING')}, top_p={stock_kwargs.get('top_p', 'MISSING')}, top_k={stock_kwargs.get('top_k', 'MISSING')}")
+        logger.debug(f"[ACE15_TOKENIZER] task={task_type}, track={track_name}")

        # Call the original tokenizer to get baseline output with all expected keys
        out = original_tokenize_with_weights(self, text, return_word_ids, **stock_kwargs)
-        print(f"[ACE15_PATCHED_TOKENIZER] lm_metadata after stock tokenizer: {out.get('lm_metadata', 'MISSING')}")
+        logger.debug(f"[ACE15_TOKENIZER] generate_audio_codes={out.get('lm_metadata', {}).get('generate_audio_codes', 'N/A')}")

        # Get appropriate instruction for task type
        instruction = get_task_instruction(task_type, track_name)
@@ -249,7 +272,7 @@ def _create_patched_tokenizer(original_tokenize_with_weights):
        out["task_type"] = task_type
        out["track_name"] = track_name

-        print(f"[ACE15_PATCHED_TOKENIZER] FINAL lm_metadata: {out.get('lm_metadata', 'MISSING')}")
+        logger.debug(f"[ACE15_TOKENIZER] final generate_audio_codes={out.get('lm_metadata', {}).get('generate_audio_codes', 'N/A')}")
        return out

    return patched_tokenize_with_weights