From 5096e3ffb5e6e365880f7db722d519f6bff93933 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20Kr=C3=BCger?= Date: Mon, 24 Nov 2025 16:28:54 +0100 Subject: [PATCH] feat: add Pivoine custom ComfyUI nodes for DiffRhythm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add custom node wrapper PivoineDiffRhythmRun that fixes tensor dimension mismatch error by disabling chunked VAE decoding. The original DiffRhythm node's overlap=32 parameter conflicts with the VAE's 64-channel architecture. Changes: - Add comfyui/nodes/pivoine_diffrhythm.py: Custom node wrapper - Add comfyui/nodes/__init__.py: Package initialization - Add arty.yml setup/pivoine-nodes: Deployment script for symlink - Update all 4 DiffRhythm workflows to use PivoineDiffRhythmRun Technical details: - Inherits from DiffRhythmRun to avoid upstream patching - Forces chunked=False in diffrhythmgen() override - Requires more VRAM (~12-16GB) but RTX 4090 has 24GB - Category: 🌸Pivoine/Audio for easy identification 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- arty.yml | 28 ++++++++ comfyui/nodes/__init__.py | 20 ++++++ comfyui/nodes/pivoine_diffrhythm.py | 69 +++++++++++++++++++ .../diffrhythm-full-length-t2m-v1.json | 2 +- .../diffrhythm-random-generation-v1.json | 2 +- .../diffrhythm-reference-based-v1.json | 2 +- .../diffrhythm-simple-t2m-v1.json | 2 +- 7 files changed, 121 insertions(+), 4 deletions(-) create mode 100644 comfyui/nodes/__init__.py create mode 100644 comfyui/nodes/pivoine_diffrhythm.py diff --git a/arty.yml b/arty.yml index f1670a6..f5bf2e6 100644 --- a/arty.yml +++ b/arty.yml @@ -368,6 +368,34 @@ scripts: exit 1 fi + setup/pivoine-nodes: | + echo "=========================================" + echo " Linking Pivoine Custom Nodes" + echo "=========================================" + echo "" + + NODES_SRC="/workspace/ai/comfyui/nodes" + NODES_DEST="/workspace/ComfyUI/custom_nodes/ComfyUI_Pivoine" + + # Remove existing symlink if present + if [ -L "$NODES_DEST" ] || [ -d "$NODES_DEST" ]; then + echo "Removing existing: $NODES_DEST" + rm -rf "$NODES_DEST" + fi + + # Create symlink + ln -s "$NODES_SRC" "$NODES_DEST" + + echo "" + echo "✓ Pivoine custom nodes linked" + echo " Source: $NODES_SRC" + echo " Linked: $NODES_DEST" + echo "" + echo "Available Pivoine nodes:" + echo " 🌸 PivoineDiffRhythmRun - DiffRhythm with chunked disabled" + echo "" + echo "Category: 🌸Pivoine/Audio" + setup/comfyui-extensions-deps: | echo "=========================================" echo " Installing ComfyUI Extensions Dependencies" diff --git a/comfyui/nodes/__init__.py b/comfyui/nodes/__init__.py new file mode 100644 index 0000000..f28d2d5 --- /dev/null +++ b/comfyui/nodes/__init__.py @@ -0,0 +1,20 @@ +""" +Pivoine Custom ComfyUI Nodes +Custom node extensions and wrappers for RunPod deployment + +Author: valknar@pivoine.art +""" + +from .pivoine_diffrhythm import NODE_CLASS_MAPPINGS as DIFFRHYTHM_MAPPINGS +from .pivoine_diffrhythm import NODE_DISPLAY_NAME_MAPPINGS as DIFFRHYTHM_DISPLAY + +# Combine all node mappings +NODE_CLASS_MAPPINGS = { + **DIFFRHYTHM_MAPPINGS, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + **DIFFRHYTHM_DISPLAY, +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/comfyui/nodes/pivoine_diffrhythm.py b/comfyui/nodes/pivoine_diffrhythm.py new file mode 100644 index 0000000..008909f --- /dev/null +++ b/comfyui/nodes/pivoine_diffrhythm.py @@ -0,0 +1,69 @@ +""" +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", +} diff --git a/comfyui/workflows/text-to-music/diffrhythm-full-length-t2m-v1.json b/comfyui/workflows/text-to-music/diffrhythm-full-length-t2m-v1.json index 85db7d4..0f47f2b 100644 --- a/comfyui/workflows/text-to-music/diffrhythm-full-length-t2m-v1.json +++ b/comfyui/workflows/text-to-music/diffrhythm-full-length-t2m-v1.json @@ -4,7 +4,7 @@ "nodes": [ { "id": 1, - "type": "DiffRhythmRun", + "type": "PivoineDiffRhythmRun", "pos": [100, 100], "size": [400, 400], "flags": {}, diff --git a/comfyui/workflows/text-to-music/diffrhythm-random-generation-v1.json b/comfyui/workflows/text-to-music/diffrhythm-random-generation-v1.json index 55ba62f..5ebaa9c 100644 --- a/comfyui/workflows/text-to-music/diffrhythm-random-generation-v1.json +++ b/comfyui/workflows/text-to-music/diffrhythm-random-generation-v1.json @@ -4,7 +4,7 @@ "nodes": [ { "id": 1, - "type": "DiffRhythmRun", + "type": "PivoineDiffRhythmRun", "pos": [100, 100], "size": [400, 400], "flags": {}, diff --git a/comfyui/workflows/text-to-music/diffrhythm-reference-based-v1.json b/comfyui/workflows/text-to-music/diffrhythm-reference-based-v1.json index efe8b93..b4f924a 100644 --- a/comfyui/workflows/text-to-music/diffrhythm-reference-based-v1.json +++ b/comfyui/workflows/text-to-music/diffrhythm-reference-based-v1.json @@ -27,7 +27,7 @@ }, { "id": 2, - "type": "DiffRhythmRun", + "type": "PivoineDiffRhythmRun", "pos": [500, 100], "size": [400, 450], "flags": {}, diff --git a/comfyui/workflows/text-to-music/diffrhythm-simple-t2m-v1.json b/comfyui/workflows/text-to-music/diffrhythm-simple-t2m-v1.json index 3f422cb..a1b069d 100644 --- a/comfyui/workflows/text-to-music/diffrhythm-simple-t2m-v1.json +++ b/comfyui/workflows/text-to-music/diffrhythm-simple-t2m-v1.json @@ -4,7 +4,7 @@ "nodes": [ { "id": 1, - "type": "DiffRhythmRun", + "type": "PivoineDiffRhythmRun", "pos": [100, 100], "size": [400, 400], "flags": {},