From f9ce68f058d9bfd4b3a4e6defd3dd79bf8e61a7b Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Sat, 7 Feb 2026 18:12:47 -0500 Subject: [PATCH] fix: cfg -> apg, noise --- nodes/acestep/ace_step_guiders.py | 281 ++++++++++++++++++++++++++---- nodes/acestep/ace_step_utils.py | 4 + nodes/acestep/nodes.py | 111 +++++++++--- nodes/acestep/patches.py | 31 +++- 4 files changed, 366 insertions(+), 61 deletions(-) diff --git a/nodes/acestep/ace_step_guiders.py b/nodes/acestep/ace_step_guiders.py index 4055e27..2ed885c 100644 --- a/nodes/acestep/ace_step_guiders.py +++ b/nodes/acestep/ace_step_guiders.py @@ -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) diff --git a/nodes/acestep/ace_step_utils.py b/nodes/acestep/ace_step_utils.py index 6ec2dac..bfbea6b 100644 --- a/nodes/acestep/ace_step_utils.py +++ b/nodes/acestep/ace_step_utils.py @@ -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 diff --git a/nodes/acestep/nodes.py b/nodes/acestep/nodes.py index 044ece7..816e4ea 100644 --- a/nodes/acestep/nodes.py +++ b/nodes/acestep/nodes.py @@ -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, } \ No newline at end of file diff --git a/nodes/acestep/patches.py b/nodes/acestep/patches.py index ffb9714..daa1184 100644 --- a/nodes/acestep/patches.py +++ b/nodes/acestep/patches.py @@ -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