feat: Add audio_cover_strength input to ACEStep 1.5 text encoder node#12510
Open
Subarasheese wants to merge 6 commits intoComfy-Org:masterfrom
Open
feat: Add audio_cover_strength input to ACEStep 1.5 text encoder node#12510Subarasheese wants to merge 6 commits intoComfy-Org:masterfrom
Subarasheese wants to merge 6 commits intoComfy-Org:masterfrom
Conversation
Subarasheese
requested review from
Kosinkadink,
comfyanonymous and
guill
as code owners
Subarasheese
changed the title
ChuxiJ
approved these changes
Feb 21, 2026
nick-morhun
reviewed
Feb 23, 2026
|
|
||
| audio_cover_strength = kwargs.get('audio_cover_strength', 1.0) | ||
| is_cover_mode = out.get('is_covers', comfy.conds.CONDConstant(None)).cond != False | ||
| if audio_cover_strength < 1.0 and is_cover_mode and self.current_patcher is not None: |
There was a problem hiding this comment.
Can audio_cover_strength be less than 0?
nick-morhun
reviewed
Feb 23, 2026
| if audio_cover_strength < 1.0 and is_cover_mode and self.current_patcher is not None: | ||
| if not self.current_patcher.get_wrappers(comfy.patcher_extension.WrappersMP.PREDICT_NOISE, 'ace_step_cover_strength'): | ||
| _strength = audio_cover_strength | ||
| def audio_cover_strength_wrapper(executor, x, timestep, model_options={}, seed=None): |
There was a problem hiding this comment.
Is it possible to make this an independent method?
Member
Test Evidence CheckIf this PR changes user-facing behavior, visual proof (screen recording or screenshot) is required. PRs without applicable visual documentation may not be reviewed until provided. You can add it by:
|
📝 WalkthroughWalkthroughAdded an 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@comfy/model_base.py`:
- Around line 1668-1677: The wrapper currently computes progress from sigma
magnitude; change it to use denoise step counts instead: read total denoise
steps (infer_steps) and the current denoise index (current_step) from
model_options (or compute current_step from timestep if model passes step
index), compute cover_steps = int(infer_steps * strength) and set the cutoff as
current_step >= cover_steps; update ace_audio_cover_strength_wrapper to use
infer_steps and current_step (rather than
sample_sigmas/current_sigma/max_sigma/min_sigma) when deciding whether progress
>= strength so behavior matches the step-based reference.
- Around line 1391-1394: The code calls a non-existent API
comfy.context_windows.slice_cond(...) — replace those calls with the actual
slicing function exported by comfy.context_windows (the function that accepts
parameters (cond_value, window, x_in, device, temporal_dim, retain_index_list));
update the call inside resize_cond_for_context_window (involving cond_key
"vace_context") and the similar sites handling "audio_embed",
"face_pixel_values", and "pose_latents" to call that real function (or import it
at top and invoke it directly) and pass temporal_dim=3 (or the appropriate
temporal dim) and retain_index_list as before so the context-window resize uses
the correct comfy.context_windows API.
- Around line 1731-1739: The code currently adds a wrapper bound to
audio_cover_strength but never removes or replaces it if audio_cover_strength
changes; update the logic in the block guarding on audio_cover_strength and
is_cover_mode so that before adding the new bound wrapper
(ace_audio_cover_strength_wrapper) you call
self.current_patcher.remove_wrappers_with_key(comfy.patcher_extension.WrappersMP.PREDICT_NOISE,
'ace_step_cover_strength') to clear any previous wrapper, or alternatively track
the previous strength and only rebind when it differs; target the
current_patcher.get_wrappers / add_wrapper_with_key usage to ensure the wrapper
is removed/replaced when audio_cover_strength changes.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 104c1eaa-a60d-4a49-aae9-c541a748f867
📒 Files selected for processing (2)
comfy/model_base.pycomfy_extras/nodes_ace.py
Comment on lines
+1668
to
+1677
| def ace_audio_cover_strength_wrapper(strength, executor, x, timestep, model_options={}, seed=None): | ||
| sample_sigmas = model_options.get('transformer_options', {}).get('sample_sigmas', None) | ||
| if sample_sigmas is not None: | ||
| current_sigma = float(timestep.max()) | ||
| max_sigma = float(sample_sigmas[0]) | ||
| min_sigma = float(sample_sigmas[-1]) | ||
| sigma_range = max_sigma - min_sigma | ||
| if sigma_range > 0: | ||
| progress = 1.0 - (current_sigma - min_sigma) / sigma_range | ||
| if progress >= strength: |
There was a problem hiding this comment.
Use denoise step count for the cutoff, not normalized sigma.
This does not match the reference behavior described in the PR. cover_steps = int(infer_steps * strength) is step-based, but Lines 1671-1677 derive progress from sigma magnitude, so 0.5 will flip at different points on non-linear schedules.
💡 Suggested direction
- current_sigma = float(timestep.max())
- max_sigma = float(sample_sigmas[0])
- min_sigma = float(sample_sigmas[-1])
- sigma_range = max_sigma - min_sigma
- if sigma_range > 0:
- progress = 1.0 - (current_sigma - min_sigma) / sigma_range
- if progress >= strength:
+ infer_steps = max(len(sample_sigmas) - 1, 1)
+ cover_steps = int(infer_steps * strength)
+ current_step = min(
+ range(infer_steps),
+ key=lambda i: abs(float(sample_sigmas[i]) - float(timestep.max()))
+ )
+ if current_step >= cover_steps:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@comfy/model_base.py` around lines 1668 - 1677, The wrapper currently computes
progress from sigma magnitude; change it to use denoise step counts instead:
read total denoise steps (infer_steps) and the current denoise index
(current_step) from model_options (or compute current_step from timestep if
model passes step index), compute cover_steps = int(infer_steps * strength) and
set the cutoff as current_step >= cover_steps; update
ace_audio_cover_strength_wrapper to use infer_steps and current_step (rather
than sample_sigmas/current_sigma/max_sigma/min_sigma) when deciding whether
progress >= strength so behavior matches the step-based reference.
Comment on lines
+1731
to
+1739
| if audio_cover_strength < 1.0 and is_cover_mode and self.current_patcher is not None: | ||
| if not self.current_patcher.get_wrappers(comfy.patcher_extension.WrappersMP.PREDICT_NOISE, 'ace_step_cover_strength'): | ||
| import functools | ||
| bound_wrapper = functools.partial(ace_audio_cover_strength_wrapper, audio_cover_strength) | ||
| self.current_patcher.add_wrapper_with_key( | ||
| comfy.patcher_extension.WrappersMP.PREDICT_NOISE, | ||
| 'ace_step_cover_strength', | ||
| bound_wrapper | ||
| ) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== current_patcher lifecycle =="
rg -n -C2 '\bcurrent_patcher\b' comfy
echo
echo "== wrapper storage / replacement APIs =="
fd -a 'model_patcher\.py|patcher_extension\.py' comfy | xargs rg -n -C3 'def (add_wrapper_with_key|get_wrappers|remove_wrapper_with_key|clone)\b|wrappers'
echo
echo "== ace_step_cover_strength registrations / cleanup =="
rg -n -C2 'ace_step_cover_strength|add_wrapper_with_key|remove_wrapper_with_key' comfyRepository: Comfy-Org/ComfyUI
Length of output: 18448
🏁 Script executed:
# Check if model_base.py lines 1731-1739 were recently modified
cd /home/jailuser/git/comfy
git log --oneline -n 20 -- comfy/model_base.py | head -20
# Check the context around line 1731 to see if it's new or existing
git show HEAD:comfy/model_base.py | sed -n '1720,1750p'
# Look for where prepare_sampling is called to understand patcher lifecycle
rg -n 'prepare_sampling' comfy --type py -B2 -A2 | head -80Repository: Comfy-Org/ComfyUI
Length of output: 1735
🏁 Script executed:
# Check git history for this specific code block
git log --all -p -- comfy/model_base.py | grep -A 20 -B 5 'ace_step_cover_strength' | head -50
# Check if this is a new file or recent addition
git log --oneline -n 30 -- comfy/model_base.py
# Search for where wrappers are cleared or reset
rg -n 'remove_wrappers_with_key|clean_wrappers|reset' comfy --type pyRepository: Comfy-Org/ComfyUI
Length of output: 6735
🏁 Script executed:
# See what changed in that commit
git show d1aeda9 --stat
# Check if there's any cleanup of wrappers when patcher is reused
rg -n 'cleanup\|pre_run' comfy/model_patcher.py -B3 -A10
# Check how prepare_sampling is called and whether new patcher is created each time
rg -n 'prepare_sampling\|current_patcher = ' comfy --type py -B3 -A3 | head -100Repository: Comfy-Org/ComfyUI
Length of output: 41390
🏁 Script executed:
# Check cleanup and pre_run flow
git show d1aeda9:comfy/model_patcher.py | sed -n '1036,1042p; 1164,1170p'
# Check if clone() is called before reuse
rg -n 'prepare_sampling\|\.clone()' comfy --type py -B2 -A2 | head -100
# Look for usage pattern - is the patcher cleaned/cloned between calls?
rg -n 'prepare_sampling\|pre_run\|cleanup' comfy/samplers.py -B1 -A1Repository: Comfy-Org/ComfyUI
Length of output: 566
🏁 Script executed:
# Search for where prepare_sampling is called
rg -n 'prepare_sampling' comfy --type py -B3 -A3
# Check how prepare_state is called and if new patcher is created
rg -n 'prepare_state\|\.prepare_sampling' comfy/samplers.py -B2 -A2
# Check if patcher is cloned or recreated
rg -n 'ModelPatcher\(' comfy/samplers.py
# Search for where current_patcher is actually used in extra_conds
rg -n 'extra_conds\|prepare_sampling' comfy/samplers.py -B5 -A5Repository: Comfy-Org/ComfyUI
Length of output: 4653
🏁 Script executed:
# Check how extra_conds is called and if audio_cover_strength can change
rg -n 'extra_conds' comfy --type py -B3 -A3 | head -100
# Look at the full _prepare_sampling function to see when extra_conds is called
rg -n '_prepare_sampling' comfy/sampler_helpers.py -A 50 | head -80
# Check if multiple calls to extra_conds happen in same sampling
rg -n 'model.extra_conds\|\.extra_conds' comfy --type py -B2 -A2Repository: Comfy-Org/ComfyUI
Length of output: 10316
🏁 Script executed:
# Check if extra_conds is called multiple times during a single sampling run
rg -n 'encode_model_conds' comfy/samplers.py -B5 -A10
# Check the flow - is extra_conds called once per cond or multiple times
rg -n 'for k in conds' comfy/samplers.py -B3 -A5
# Look for any mechanism that might update wrappers between calls
rg -n 'remove_wrappers_with_key\|wrappers.pop' comfy --type pyRepository: Comfy-Org/ComfyUI
Length of output: 4946
🏁 Script executed:
# Check if wrappers are cleared in cleanup or if they persist
git show d1aeda9:comfy/model_patcher.py | sed -n '1030,1045p'
# Check the clone() method to see if wrappers are copied
git show d1aeda9:comfy/model_patcher.py | sed -n '310,350p'
# Check if there's any test or documentation about patcher reuse
rg -n 'current_patcher\|patcher.*reuse' comfy --type py | grep -i test
# Final check: see what happens when audio_cover_strength differs
rg -n 'audio_cover_strength' comfy --type py -B5 -A5Repository: Comfy-Org/ComfyUI
Length of output: 4901
Remove or update the wrapper when audio_cover_strength changes.
The guard check prevents duplicate registration but also means a wrapper with a stale audio_cover_strength value persists if the patcher is reused across sampling runs. When audio_cover_strength differs on a subsequent call, the wrapper bound to the previous value is never replaced. Call self.current_patcher.remove_wrappers_with_key(comfy.patcher_extension.WrappersMP.PREDICT_NOISE, 'ace_step_cover_strength') before registering, or store the previous strength and only add/replace when it changes.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@comfy/model_base.py` around lines 1731 - 1739, The code currently adds a
wrapper bound to audio_cover_strength but never removes or replaces it if
audio_cover_strength changes; update the logic in the block guarding on
audio_cover_strength and is_cover_mode so that before adding the new bound
wrapper (ace_audio_cover_strength_wrapper) you call
self.current_patcher.remove_wrappers_with_key(comfy.patcher_extension.WrappersMP.PREDICT_NOISE,
'ace_step_cover_strength') to clear any previous wrapper, or alternatively track
the previous strength and only rebind when it differs; target the
current_patcher.get_wrappers / add_wrapper_with_key usage to ensure the wrapper
is removed/replaced when audio_cover_strength changes.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
comfy/model_base.py (1)
1592-1594: Consider clampingaudio_cover_strengthdefensively.While the node input in
nodes_ace.pyenforces the 0.0–1.0 range, the model code doesn't validate. A negative value would causeprogress >= _strengthto always be true (since progress is 0.0–1.0), immediately disabling cover mode. Adding a simple clamp would make this more robust.🛡️ Defensive clamp
out['refer_audio'] = comfy.conds.CONDRegular(refer_audio) - audio_cover_strength = kwargs.get('audio_cover_strength', 1.0) + audio_cover_strength = max(0.0, min(1.0, kwargs.get('audio_cover_strength', 1.0))) is_cover_mode = out.get('is_covers', comfy.conds.CONDConstant(None)).cond != False🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@comfy/model_base.py` around lines 1592 - 1594, Clamp the audio_cover_strength value defensively before it's used: replace the direct read of kwargs.get('audio_cover_strength', 1.0) with a bounded value (e.g. audio_cover_strength = max(0.0, min(1.0, kwargs.get('audio_cover_strength', 1.0)))). This ensures the subsequent is_cover_mode and self.current_patcher logic (the check using audio_cover_strength in the progress comparison) cannot be bypassed by negative or >1 values; keep the rest of the checks (is_cover_mode, self.current_patcher) unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@comfy/model_base.py`:
- Around line 1592-1594: Clamp the audio_cover_strength value defensively before
it's used: replace the direct read of kwargs.get('audio_cover_strength', 1.0)
with a bounded value (e.g. audio_cover_strength = max(0.0, min(1.0,
kwargs.get('audio_cover_strength', 1.0)))). This ensures the subsequent
is_cover_mode and self.current_patcher logic (the check using
audio_cover_strength in the progress comparison) cannot be bypassed by negative
or >1 values; keep the rest of the checks (is_cover_mode, self.current_patcher)
unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 75fb5c34-1927-4610-a626-281a9eac441a
📒 Files selected for processing (2)
comfy/model_base.pycomfy_extras/nodes_ace.py
🚧 Files skipped from review as they are similar to previous changes (1)
- comfy_extras/nodes_ace.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What issue does this address?
The
TextEncodeAceStepAudio1.5node had no way to control LM code conditioning strength during the denoising process. The ACE-Step 1.5 reference implementation exposes this asaudio_cover_strength, which controls how many denoising steps use LM code/cover conditioning before switching to unconditioned generation. This parameter was missing from the ComfyUI node.How does this PR address it?
Files changed:
comfy_extras/nodes_ace.pyaudio_cover_strengthfloat input (range 0.0–1.0, default 1.0, advanced) toTextEncodeAceStepAudio1.5conditioning_set_valuescomfy/model_base.pyACEStep15.extra_conds(), readsaudio_cover_strengthfrom conditioning kwargsPREDICT_NOISEwrapper on the model patcher that monitors sampling progress viasample_sigmasaudio_cover_strength, flipsis_coverstoFalsefor remaining steps, matching the reference implementation behavior (cover_steps = int(infer_steps * audio_cover_strength))Testing
Tested with
audio_cover_strength = 0.5on ACE-Step 1.5 with LM code generation enabled. The wrapper correctly toggles conditioning at the midpoint of the denoising schedule. Whenaudio_cover_strength = 1.0(default), no wrapper is registered and behavior is identical to before.Potential concerns
condsinmodel_optionsin-place to flipis_covers. This is the same pattern used by other conditioning overrides in ComfyUI, but worth reviewing.get_wrapperscheck) and scoped to the model patcher, so it should not leak across runs.