"""
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 the infer function to force chunked=False
import DiffRhythmNode
_original_infer = DiffRhythmNode.infer

def patched_infer(*args, **kwargs):
    # Force chunked to False if present
    if 'chunked' in kwargs:
        kwargs['chunked'] = False
    return _original_infer(*args, chunked=False, **kwargs)

# Apply the monkey patch
DiffRhythmNode.infer = patched_infer

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