"""
Pivoine DiffRhythm Node
Custom wrapper for DiffRhythm that disables chunked decoding to prevent
tensor dimension mismatch errors (32 vs 64) in VAE overlap logic.

Author: valknar@pivoine.art
"""

import sys
sys.path.append('/workspace/ComfyUI/custom_nodes/ComfyUI_DiffRhythm')

# Monkey-patch decode_audio from infer_utils to force chunked=False
import infer_utils
_original_decode_audio = infer_utils.decode_audio

def patched_decode_audio(latent, vae_model, chunked=True):
    """Patched version that always uses chunked=False"""
    return _original_decode_audio(latent, vae_model, chunked=False)

# Apply the monkey patch
infer_utils.decode_audio = patched_decode_audio

from DiffRhythmNode import DiffRhythmRun

class PivoineDiffRhythmRun(DiffRhythmRun):
    """
    Pivoine version of DiffRhythmRun with chunked decoding forcibly disabled.

    Changes from original:
    - Monkey-patches the infer() function to always use chunked=False
    - Prevents tensor dimension mismatch in VAE (32 vs 64 channel error)
    - Requires more VRAM (~12-16GB) but works reliably on RTX 4090
    """

    CATEGORY = "🌸Pivoine/Audio"

    @classmethod
    def INPUT_TYPES(cls):
        return super().INPUT_TYPES()

NODE_CLASS_MAPPINGS = {
    "PivoineDiffRhythmRun": PivoineDiffRhythmRun,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PivoineDiffRhythmRun": "Pivoine DiffRhythm Run",
}