From 951b1318f8958bfc441c0bcbbf2fc0695e1d9278 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Fri, 6 Feb 2026 05:45:16 -0500 Subject: [PATCH] cover fixed --- nodes/acestep/ace_step_guiders.py | 91 ++++++++++++++--- nodes/acestep/nodes.py | 159 ++++++++++++++++++++++++++++-- nodes/acestep/patches.py | 49 ++++++--- 3 files changed, 269 insertions(+), 30 deletions(-) diff --git a/nodes/acestep/ace_step_guiders.py b/nodes/acestep/ace_step_guiders.py index 03a55c4..c8aebff 100644 --- a/nodes/acestep/ace_step_guiders.py +++ b/nodes/acestep/ace_step_guiders.py @@ -228,32 +228,49 @@ class ACEStep15NativeCoverGuider(comfy.samplers.CFGGuider): """ ACE-Step 1.5 cover guider for style transfer/regeneration. - Uses the source audio as context (src_latents) while generating new content - across the entire duration. This is useful for cover songs, style transfer, - or regenerating audio with a new prompt while maintaining the original structure. + Uses semantic tokens (lm_hints) from the source audio as structural guidance + while generating new content with a different style/timbre. Official behavior: - chunk_masks: All 1s (generate everything) - - src_latents: Original audio latents (provides structural guidance) + - is_covers: 1 (use semantic hints instead of raw latents) + - precomputed_lm_hints_25Hz: Semantic representation of source audio + + IMPORTANT: For proper cover behavior, provide semantic_hints from ACEStep15SemanticExtractor. + Without semantic_hints, falls back to using raw VAE latents (less accurate style transfer). """ - def __init__(self, model, positive, negative, cfg, source_latent): + def __init__(self, model, positive, negative, cfg, source_latent, semantic_hints=None): super().__init__(model) self.set_conds(positive, negative) self.set_cfg(cfg) self.source_latent = source_latent + self.semantic_hints = semantic_hints batch_size, channels, total_length = source_latent.shape print(f"[ACE15_NATIVE_COVER] Initializing") print(f"[ACE15_NATIVE_COVER] source_latent.shape: {source_latent.shape}") + print(f"[ACE15_NATIVE_COVER] source_latent stats: mean={source_latent.mean():.4f}, std={source_latent.std():.4f}") + if semantic_hints is not None: + print(f"[ACE15_NATIVE_COVER] semantic_hints.shape: {semantic_hints.shape}") + print(f"[ACE15_NATIVE_COVER] semantic_hints stats: mean={semantic_hints.mean():.4f}, std={semantic_hints.std():.4f}, min={semantic_hints.min():.4f}, max={semantic_hints.max():.4f}") + # Validate semantic hints have reasonable variance + if semantic_hints.std() < 0.01: + print(f"[ACE15_NATIVE_COVER] WARNING: semantic_hints have very low variance ({semantic_hints.std():.6f}) - may be constant/invalid!") + # Validate shape matches source + if semantic_hints.shape != source_latent.shape: + print(f"[ACE15_NATIVE_COVER] WARNING: semantic_hints shape {semantic_hints.shape} != source_latent shape {source_latent.shape}") + print(f"[ACE15_NATIVE_COVER] Using semantic hints for proper cover behavior (is_covers=1)") + else: + print(f"[ACE15_NATIVE_COVER] WARNING: No semantic_hints provided, falling back to VAE latents (is_covers=0)") - # Cover: generate everything but use source audio as context + # Cover: generate everything # chunk_masks: All 1s (generate all frames) self.chunk_masks = torch.ones_like(source_latent) - # src_latents: Original audio (provides structural guidance) + # src_latents: Original audio (used as fallback if is_covers=0) self.src_latents = source_latent.clone() print(f"[ACE15_NATIVE_COVER] chunk_masks.shape: {self.chunk_masks.shape}") @@ -262,12 +279,13 @@ class ACEStep15NativeCoverGuider(comfy.samplers.CFGGuider): self._wrapper_applied = False def _apply_model_wrapper(self): - """Apply model wrapper to inject chunk_masks and src_latents""" + """Apply model wrapper to inject chunk_masks, src_latents, is_covers, and semantic hints""" if self._wrapper_applied: return chunk_masks = self.chunk_masks src_latents = self.src_latents + semantic_hints = self.semantic_hints self.model_patcher.set_model_unet_function_wrapper(None) @@ -286,11 +304,36 @@ class ACEStep15NativeCoverGuider(comfy.samplers.CFGGuider): c["chunk_masks"] = cm c["src_latents"] = sl + # Set is_covers based on whether we have semantic hints + if semantic_hints is not None: + # Proper cover: use semantic hints via is_covers=1 + 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: + # Fallback: use raw VAE latents via is_covers=0 + c["is_covers"] = torch.zeros((input_batch_size,), device=device, dtype=torch.long) + if not hasattr(model_function_wrapper, '_logged'): - print(f"[ACE15_COVER_WRAPPER] Injecting chunk_masks and src_latents") + print(f"[ACE15_COVER_WRAPPER] Injecting cover parameters") print(f"[ACE15_COVER_WRAPPER] input.shape: {args['input'].shape}") + print(f"[ACE15_COVER_WRAPPER] input stats: mean={args['input'].mean():.4f}, std={args['input'].std():.4f}") print(f"[ACE15_COVER_WRAPPER] chunk_masks.shape: {c['chunk_masks'].shape}") + print(f"[ACE15_COVER_WRAPPER] chunk_masks sum: {c['chunk_masks'].sum().item()} (should be all 1s for cover)") print(f"[ACE15_COVER_WRAPPER] src_latents.shape: {c['src_latents'].shape}") + print(f"[ACE15_COVER_WRAPPER] src_latents stats: mean={c['src_latents'].mean():.4f}, std={c['src_latents'].std():.4f}") + print(f"[ACE15_COVER_WRAPPER] is_covers: {c['is_covers']} (1=use semantic hints, 0=use VAE latents)") + if "precomputed_lm_hints_25Hz" in c: + sh = c["precomputed_lm_hints_25Hz"] + print(f"[ACE15_COVER_WRAPPER] precomputed_lm_hints_25Hz.shape: {sh.shape}") + print(f"[ACE15_COVER_WRAPPER] precomputed_lm_hints_25Hz stats: mean={sh.mean():.4f}, std={sh.std():.4f}, min={sh.min():.4f}, max={sh.max():.4f}") + # Check if hints look valid (should have some variance, not constant) + if sh.std() < 0.01: + print(f"[ACE15_COVER_WRAPPER] WARNING: semantic hints have very low variance! May be constant/invalid.") + else: + print(f"[ACE15_COVER_WRAPPER] WARNING: No precomputed_lm_hints_25Hz - will fall back to VAE latents!") model_function_wrapper._logged = True return apply_model(args["input"], args["timestep"], **c) @@ -338,23 +381,31 @@ class ACEStep15NativeExtractGuider(comfy.samplers.CFGGuider): Official behavior: - chunk_masks: All 1s (generate everything) - - src_latents: Original audio latents + - is_covers: 1 (use semantic hints for structure) + - precomputed_lm_hints_25Hz: Semantic representation of source audio - Instruction: "Extract the {TRACK_NAME} track from the audio:" + + IMPORTANT: For proper extract behavior, provide semantic_hints from ACEStep15SemanticExtractor. """ - def __init__(self, model, positive, negative, cfg, source_latent, track_name="vocals"): + def __init__(self, model, positive, negative, cfg, source_latent, track_name="vocals", semantic_hints=None): super().__init__(model) self.set_conds(positive, negative) self.set_cfg(cfg) self.source_latent = source_latent self.track_name = track_name + self.semantic_hints = semantic_hints batch_size, channels, total_length = source_latent.shape print(f"[ACE15_NATIVE_EXTRACT] Initializing") print(f"[ACE15_NATIVE_EXTRACT] source_latent.shape: {source_latent.shape}") print(f"[ACE15_NATIVE_EXTRACT] track_name: {track_name}") + if semantic_hints is not None: + print(f"[ACE15_NATIVE_EXTRACT] semantic_hints.shape: {semantic_hints.shape}") + else: + print(f"[ACE15_NATIVE_EXTRACT] WARNING: No semantic_hints provided, falling back to VAE latents") # Extract: generate everything (the extracted track) self.chunk_masks = torch.ones_like(source_latent) @@ -365,12 +416,13 @@ class ACEStep15NativeExtractGuider(comfy.samplers.CFGGuider): self._wrapper_applied = False def _apply_model_wrapper(self): - """Apply model wrapper to inject chunk_masks and src_latents""" + """Apply model wrapper to inject chunk_masks, src_latents, is_covers, and semantic hints""" if self._wrapper_applied: return chunk_masks = self.chunk_masks src_latents = self.src_latents + semantic_hints = self.semantic_hints self.model_patcher.set_model_unet_function_wrapper(None) @@ -389,8 +441,21 @@ class ACEStep15NativeExtractGuider(comfy.samplers.CFGGuider): c["chunk_masks"] = cm c["src_latents"] = sl + # Set is_covers based on whether we have semantic hints + 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) + if not hasattr(model_function_wrapper, '_logged'): - print(f"[ACE15_EXTRACT_WRAPPER] Injecting chunk_masks and src_latents") + print(f"[ACE15_EXTRACT_WRAPPER] Injecting extract parameters") + print(f"[ACE15_EXTRACT_WRAPPER] is_covers: {c['is_covers']}") + if "precomputed_lm_hints_25Hz" in c: + print(f"[ACE15_EXTRACT_WRAPPER] precomputed_lm_hints_25Hz.shape: {c['precomputed_lm_hints_25Hz'].shape}") model_function_wrapper._logged = True return apply_model(args["input"], args["timestep"], **c) diff --git a/nodes/acestep/nodes.py b/nodes/acestep/nodes.py index e19dccc..9e1117d 100644 --- a/nodes/acestep/nodes.py +++ b/nodes/acestep/nodes.py @@ -122,6 +122,117 @@ class ACEStep15SilenceLatentLoader: return (silence_latent,) +class ACEStep15SemanticExtractor: + """ + Extracts semantic tokens from source audio latents for use with Cover/Extract tasks. + + The ACE-Step 1.5 model uses semantic tokens (lm_hints) as structural guidance + for tasks like Cover and Extract. This node extracts those tokens from VAE latents. + + Flow: VAE latents → tokenizer → indices → quantizer → detokenizer → semantic_hints + + Use the output with ACE-Step 1.5 Cover Guider or Extract Guider for proper + style transfer behavior. + """ + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL",), + "source_latents": ("LATENT",), + } + } + + RETURN_TYPES = ("SEMANTIC_HINTS",) + RETURN_NAMES = ("semantic_hints",) + FUNCTION = "extract" + CATEGORY = "audio/acestep" + + def extract(self, model, source_latents): + 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}") + + print(f"[ACE15_SEMANTIC_EXTRACTOR] Extracting semantic hints") + print(f"[ACE15_SEMANTIC_EXTRACTOR] source_tensor.shape: {source_tensor.shape}") + + # Debug: Print model structure + print(f"[ACE15_SEMANTIC_EXTRACTOR] model type: {type(model)}") + print(f"[ACE15_SEMANTIC_EXTRACTOR] model.__class__.__name__: {model.__class__.__name__}") + if hasattr(model, 'model'): + print(f"[ACE15_SEMANTIC_EXTRACTOR] model.model type: {type(model.model)}") + print(f"[ACE15_SEMANTIC_EXTRACTOR] model.model.__class__.__name__: {model.model.__class__.__name__}") + if hasattr(model.model, 'diffusion_model'): + print(f"[ACE15_SEMANTIC_EXTRACTOR] model.model.diffusion_model type: {type(model.model.diffusion_model)}") + print(f"[ACE15_SEMANTIC_EXTRACTOR] model.model.diffusion_model.__class__.__name__: {model.model.diffusion_model.__class__.__name__}") + + # Get the diffusion model (AceStepConditionGenerationModel) + # It has .tokenizer and .detokenizer attributes + diffusion_model = model.model.diffusion_model + + if not hasattr(diffusion_model, 'tokenizer') or not hasattr(diffusion_model, 'detokenizer'): + raise RuntimeError( + "Model does not have tokenizer/detokenizer. " + "Make sure you're using an ACE-Step 1.5 model." + ) + + # Load model to GPU using ComfyUI's model management + # Pass the ModelPatcher (model), not the raw model (model.model) + print(f"[ACE15_SEMANTIC_EXTRACTOR] Loading model to GPU...") + comfy.model_management.load_model_gpu(model) + + device = comfy.model_management.get_torch_device() + dtype = model.model.get_dtype() + + print(f"[ACE15_SEMANTIC_EXTRACTOR] device: {device}, dtype: {dtype}") + + # Move source tensor to model device/dtype + # Source is in ComfyUI format: [B, D, T] = [B, 64, length] + # Tokenizer expects: [B, T, D] = [B, length, 64] + source_on_device = source_tensor.to(device=device, dtype=dtype) + source_transposed = source_on_device.movedim(-1, -2) # [B, T, D] + + print(f"[ACE15_SEMANTIC_EXTRACTOR] source_transposed.shape: {source_transposed.shape}") + print(f"[ACE15_SEMANTIC_EXTRACTOR] source stats: mean={source_transposed.mean():.4f}, std={source_transposed.std():.4f}, min={source_transposed.min():.4f}, max={source_transposed.max():.4f}") + + # Verify tokenizer/detokenizer weights are on GPU + tokenizer_device = next(diffusion_model.tokenizer.parameters()).device + detokenizer_device = next(diffusion_model.detokenizer.parameters()).device + print(f"[ACE15_SEMANTIC_EXTRACTOR] tokenizer weights on: {tokenizer_device}") + print(f"[ACE15_SEMANTIC_EXTRACTOR] detokenizer weights on: {detokenizer_device}") + + # Step 1: Tokenize - get quantized embeddings at 5Hz + # tokenizer.tokenize() handles the window pooling internally + # Returns: quantized [B, T/5, hidden_size], indices [B, T/5, num_quantizers] + with torch.no_grad(): + quantized, indices = diffusion_model.tokenizer.tokenize(source_transposed) + + print(f"[ACE15_SEMANTIC_EXTRACTOR] quantized.shape: {quantized.shape}") + print(f"[ACE15_SEMANTIC_EXTRACTOR] quantized stats: mean={quantized.mean():.4f}, std={quantized.std():.4f}, min={quantized.min():.4f}, max={quantized.max():.4f}") + print(f"[ACE15_SEMANTIC_EXTRACTOR] indices.shape: {indices.shape}") + + # Step 2: Detokenize - upsample from 5Hz to 25Hz + # FIX: Use quantized directly instead of going through get_output_from_indices + # This avoids potential precision loss from the redundant index lookup + with torch.no_grad(): + lm_hints = diffusion_model.detokenizer(quantized) + + print(f"[ACE15_SEMANTIC_EXTRACTOR] lm_hints.shape: {lm_hints.shape}") + print(f"[ACE15_SEMANTIC_EXTRACTOR] lm_hints stats: mean={lm_hints.mean():.4f}, std={lm_hints.std():.4f}, min={lm_hints.min():.4f}, max={lm_hints.max():.4f}") + + # Transpose back to ComfyUI format: [B, T, D] → [B, D, T] + semantic_hints = lm_hints.movedim(-1, -2) + print(f"[ACE15_SEMANTIC_EXTRACTOR] semantic_hints.shape (output): {semantic_hints.shape}") + + # Move to CPU for storage + semantic_hints = semantic_hints.cpu() + + return (semantic_hints,) + + class ACEStepRepaintGuiderNode: """Node that creates a repaint guider for use with SamplerCustomAdvanced""" @@ -724,8 +835,11 @@ class ACEStep15NativeCoverGuiderNode: """ Creates an ACE-Step 1.5 cover guider for style transfer/regeneration. - Uses the source audio as context while generating new content across - the entire duration. Useful for cover songs or regenerating with a new prompt. + Uses semantic tokens from the source audio as structural guidance while + generating new content with a different style/timbre. + + IMPORTANT: For proper cover behavior, connect semantic_hints from ACEStep15SemanticExtractor. + Without semantic_hints, falls back to using raw VAE latents (less accurate style transfer). """ @classmethod @@ -737,6 +851,9 @@ class ACEStep15NativeCoverGuiderNode: "negative": ("CONDITIONING",), "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}), "source_latents": ("LATENT",), + }, + "optional": { + "semantic_hints": ("SEMANTIC_HINTS",), } } @@ -744,7 +861,7 @@ class ACEStep15NativeCoverGuiderNode: FUNCTION = "get_guider" CATEGORY = "sampling/custom_sampling/guiders" - def get_guider(self, model, positive, negative, cfg, source_latents): + def get_guider(self, model, positive, negative, cfg, source_latents, semantic_hints=None): source_tensor = source_latents["samples"] # Validate v1.5 shape: (batch, 64, length) @@ -753,9 +870,13 @@ class ACEStep15NativeCoverGuiderNode: print(f"[ACE15_COVER_NODE] Creating guider") print(f"[ACE15_COVER_NODE] source_tensor.shape: {source_tensor.shape}") + if semantic_hints is not None: + print(f"[ACE15_COVER_NODE] semantic_hints.shape: {semantic_hints.shape}") + else: + print(f"[ACE15_COVER_NODE] WARNING: No semantic_hints provided - using fallback mode") guider = ACEStep15NativeCoverGuider( - model, positive, negative, cfg, source_tensor + model, positive, negative, cfg, source_tensor, semantic_hints ) return (guider,) @@ -767,6 +888,8 @@ class ACEStep15NativeExtractGuiderNode: Extracts a specific track (vocals, drums, bass, etc.) from the source audio. Use with the ACEStep15TaskTextEncode node with task_type="extract". + + IMPORTANT: For proper extract behavior, connect semantic_hints from ACEStep15SemanticExtractor. """ @classmethod @@ -780,6 +903,9 @@ class ACEStep15NativeExtractGuiderNode: "source_latents": ("LATENT",), "track_name": (["vocals", "drums", "bass", "guitar", "keyboard", "strings", "percussion", "synth", "fx", "brass", "woodwinds", "backing_vocals"],), + }, + "optional": { + "semantic_hints": ("SEMANTIC_HINTS",), } } @@ -787,7 +913,7 @@ class ACEStep15NativeExtractGuiderNode: FUNCTION = "get_guider" CATEGORY = "sampling/custom_sampling/guiders" - def get_guider(self, model, positive, negative, cfg, source_latents, track_name): + def get_guider(self, model, positive, negative, cfg, source_latents, track_name, semantic_hints=None): source_tensor = source_latents["samples"] # Validate v1.5 shape: (batch, 64, length) @@ -797,9 +923,11 @@ class ACEStep15NativeExtractGuiderNode: print(f"[ACE15_EXTRACT_NODE] Creating guider") print(f"[ACE15_EXTRACT_NODE] source_tensor.shape: {source_tensor.shape}") print(f"[ACE15_EXTRACT_NODE] track_name: {track_name}") + if semantic_hints is not None: + print(f"[ACE15_EXTRACT_NODE] semantic_hints.shape: {semantic_hints.shape}") guider = ACEStep15NativeExtractGuider( - model, positive, negative, cfg, source_tensor, track_name + model, positive, negative, cfg, source_tensor, track_name, semantic_hints ) return (guider,) @@ -960,9 +1088,24 @@ class ACEStep15TaskTextEncodeNode: print(f"[ACE15_TEXT_ENCODE] Encoding with task_type={task_type}, track_name={track_name or 'N/A'}") + # Import to get task instruction for verification + from .patches import get_task_instruction, is_patched + instruction = get_task_instruction(task_type, track_name if track_name else None) + print(f"[ACE15_TEXT_ENCODE] patches applied: {is_patched()}") + print(f"[ACE15_TEXT_ENCODE] task instruction: {instruction}") + # Use the patched tokenize_with_weights which accepts task_type # encode_from_tokens_scheduled returns conditioning directly (list format) tokens = clip.tokenize(text, **kwargs) + + # Verify the tokens contain task_type info (patched tokenizer should include this) + if isinstance(tokens, dict): + if "task_type" in tokens: + print(f"[ACE15_TEXT_ENCODE] tokens contain task_type: {tokens.get('task_type')}") + else: + print(f"[ACE15_TEXT_ENCODE] WARNING: tokens dict does not contain task_type - patch may not be applied!") + print(f"[ACE15_TEXT_ENCODE] tokens keys: {list(tokens.keys())}") + conditioning = clip.encode_from_tokens_scheduled(tokens) # For negative conditioning, use ComfyUI's ConditioningZeroOut node @@ -984,6 +1127,8 @@ NODE_CLASS_MAPPINGS = { "ACEStepAudioPostProcessor": ACEStepAudioPostProcessor, # ACE-Step 1.5 loaders "ACEStep15SilenceLatentLoader": ACEStep15SilenceLatentLoader, + # ACE-Step 1.5 semantic extraction + "ACEStep15SemanticExtractor": ACEStep15SemanticExtractor, # ACE-Step 1.5 guiders (model-level mask input) "ACEStep15NativeEditGuider": ACEStep15NativeEditGuiderNode, "ACEStep15NativeCoverGuider": ACEStep15NativeCoverGuiderNode, @@ -1007,6 +1152,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ACEStepAudioPostProcessor": "ACEStep Audio Post Processor", # ACE-Step 1.5 loaders "ACEStep15SilenceLatentLoader": "ACE-Step 1.5 Load Silence Latent", + # ACE-Step 1.5 semantic extraction + "ACEStep15SemanticExtractor": "ACE-Step 1.5 Semantic Extractor", # ACE-Step 1.5 guiders "ACEStep15NativeEditGuider": "ACE-Step 1.5 Edit Guider (Extend/Repaint)", "ACEStep15NativeCoverGuider": "ACE-Step 1.5 Cover Guider", diff --git a/nodes/acestep/patches.py b/nodes/acestep/patches.py index 2bf1fdd..0a68729 100644 --- a/nodes/acestep/patches.py +++ b/nodes/acestep/patches.py @@ -71,29 +71,46 @@ def _create_patched_ace_step15_forward(original_forward): Create a patched forward method for AceStepConditionGenerationModel. The patched version: - 1. Accepts chunk_masks and src_latents as kwargs - 2. Sets is_covers=0 when src_latents is provided (to use provided latents instead of lm_hints) + 1. Accepts chunk_masks, src_latents, is_covers, and precomputed_lm_hints_25Hz as kwargs + 2. Allows guiders to explicitly control is_covers flag for different task types 3. Properly transposes user-provided tensors from (batch, channels, length) to (batch, length, channels) + Task type behavior via is_covers: + - is_covers=0: Use provided src_latents directly (extend/repaint) + - is_covers=1: Use lm_hints from precomputed_lm_hints_25Hz or audio_codes (cover/extract) + - is_covers=None: Default behavior (src_latents = x) + This enables extend, repaint, cover, extract, and lego task types. """ + # One-time logging flag + _forward_logged = [False] + @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, **kwargs): + audio_codes=None, chunk_masks=None, src_latents=None, + is_covers=None, precomputed_lm_hints_25Hz=None, **kwargs): text_attention_mask = None lyric_attention_mask = None refer_audio_order_mask = None attention_mask = None - # chunk_masks and src_latents now come from kwargs for extend/repaint support - # When src_latents is explicitly provided (for extend/repaint), we need is_covers=0 - # so prepare_condition() uses our src_latents instead of replacing with lm_hints - is_covers = None - if src_latents is not None: - # Set is_covers to 0 to indicate we want to use the provided src_latents + # One-time diagnostic log for cover task + if not _forward_logged[0] and precomputed_lm_hints_25Hz is not None: + print(f"[ACE15_PATCHED_FORWARD] Cover mode detected") + print(f"[ACE15_PATCHED_FORWARD] x.shape (before transpose): {x.shape}") + print(f"[ACE15_PATCHED_FORWARD] is_covers: {is_covers}") + print(f"[ACE15_PATCHED_FORWARD] precomputed_lm_hints_25Hz.shape (before transpose): {precomputed_lm_hints_25Hz.shape}") + print(f"[ACE15_PATCHED_FORWARD] precomputed_lm_hints_25Hz stats: mean={precomputed_lm_hints_25Hz.mean():.4f}, std={precomputed_lm_hints_25Hz.std():.4f}") + _forward_logged[0] = True + + # is_covers can now be passed explicitly by guiders: + # - Cover/Extract tasks: is_covers=1 (use lm_hints from precomputed_lm_hints_25Hz) + # - Extend/Repaint tasks: is_covers=0 (use provided src_latents) + # - Default (not provided): fallback to old behavior + if is_covers is None and src_latents is not None: + # Backwards compatibility: default to 0 when src_latents provided but is_covers not specified is_covers = torch.zeros((x.shape[0],), device=x.device, dtype=torch.long) - precomputed_lm_hints_25Hz = None lyric_hidden_states = lyric_embed text_hidden_states = context refer_audio_acoustic_hidden_states_packed = refer_audio.movedim(-1, -2) @@ -103,11 +120,21 @@ def _create_patched_ace_step15_forward(original_forward): if refer_audio_order_mask is None: refer_audio_order_mask = torch.zeros((x.shape[0],), device=x.device, dtype=torch.long) - # Transpose user-provided src_latents/chunk_masks from (batch, channels, length) to (batch, length, channels) + # Transpose user-provided tensors from ComfyUI format (batch, channels, length) + # to model format (batch, length, channels) if src_latents is not None: src_latents = src_latents.movedim(-1, -2) if chunk_masks is not None: chunk_masks = chunk_masks.movedim(-1, -2) + if precomputed_lm_hints_25Hz is not None: + # precomputed_lm_hints_25Hz comes in ComfyUI format [B, D, T], transpose to [B, T, D] + precomputed_lm_hints_25Hz = precomputed_lm_hints_25Hz.movedim(-1, -2) + # Log post-transpose shape (one-time) + if _forward_logged[0] and not hasattr(patched_forward, '_post_transpose_logged'): + print(f"[ACE15_PATCHED_FORWARD] precomputed_lm_hints_25Hz.shape (after transpose): {precomputed_lm_hints_25Hz.shape}") + print(f"[ACE15_PATCHED_FORWARD] x.shape (after transpose): {x.shape}") + print(f"[ACE15_PATCHED_FORWARD] src_latents.shape (after transpose): {src_latents.shape if src_latents is not None else 'None'}") + patched_forward._post_transpose_logged = True if src_latents is None and is_covers is None: src_latents = x