cover fixed

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)

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",

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