"""
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')

from DiffRhythmNode import DiffRhythmRun

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

    Changes from original:
    - chunked parameter defaults to False (was True)
    - Prevents tensor dimension mismatch in VAE
    - Requires more VRAM (~12-16GB) but works reliably on RTX 4090
    """

    CATEGORY = "🌸Pivoine/Audio"

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

    def diffrhythmgen(
        self,
        edit,
        model: str,
        style_prompt: str = None,
        lyrics_or_edit_lyrics: str = "",
        style_audio_or_edit_song = None,
        edit_segments: str = "",
        chunked: bool = False,  # Changed from True to False
        odeint_method: str = "euler",
        steps: int = 30,
        cfg: int = 4,
        quality_or_speed: str = "speed",
        unload_model: bool = False,
        seed: int = 0
    ):
        # Force chunked=False to avoid dimension mismatch
        return super().diffrhythmgen(
            edit=edit,
            model=model,
            style_prompt=style_prompt,
            lyrics_or_edit_lyrics=lyrics_or_edit_lyrics,
            style_audio_or_edit_song=style_audio_or_edit_song,
            edit_segments=edit_segments,
            chunked=False,
            odeint_method=odeint_method,
            steps=steps,
            cfg=cfg,
            quality_or_speed=quality_or_speed,
            unload_model=unload_model,
            seed=seed
        )

NODE_CLASS_MAPPINGS = {
    "PivoineDiffRhythmRun": PivoineDiffRhythmRun,
}

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