From 4403b41f323122f37e93bf6eb7be5bc03000ec7c Mon Sep 17 00:00:00 2001
From: Joel Trauger <joel.trauger@gmail.com>
Date: Sun, 16 Feb 2025 19:06:30 -0600
Subject: [PATCH] initial upload

---
 .github/workflows/publish_action.yml | 20 +++++++++++++
 README.md                            | 14 +++++++++-
 __init__.py                          | 39 ++++++++++++++++++++++++++
 nodes.py                             | 42 ++++++++++++++++++++++++++++
 pyproject.toml                       | 15 ++++++++++
 requirements.txt                     |  6 ++++
 6 files changed, 135 insertions(+), 1 deletion(-)
 create mode 100644 .github/workflows/publish_action.yml
 create mode 100644 __init__.py
 create mode 100644 nodes.py
 create mode 100644 pyproject.toml
 create mode 100644 requirements.txt

diff --git a/.github/workflows/publish_action.yml b/.github/workflows/publish_action.yml
new file mode 100644
index 0000000..122a1ae
--- /dev/null
+++ b/.github/workflows/publish_action.yml
@@ -0,0 +1,20 @@
+name: Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+    paths:
+      - "pyproject.toml"
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Publish Custom Node
+        uses: Comfy-Org/publish-node-action@main
+        with:
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} ## Add your own personal access token to your Github Repository secrets and reference it here.
\ No newline at end of file
diff --git a/README.md b/README.md
index 420cadf..6b78aa5 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,14 @@
 # ComfyUI-Riffusion
-Riffusion extension for ComfyUI to generate audio from text
+
+A ComfyUI extension for Riffusion audio generation.
+
+## Installation
+
+1. Clone this repository into your `ComfyUI/custom_nodes` directory:
+```git clone https://github.com/your-username/ComfyUI-Riffusion```
+
+2. Restart ComfyUI
+
+## Usage
+
+The extension adds a new node "Riffusion" under the audio category.
\ No newline at end of file
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..21b57ff
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,39 @@
+import sys
+import subprocess
+from pathlib import Path
+
+### GET PACKAGES ###
+def packages(versions=False):
+    try:
+        result = subprocess.check_output([sys.executable, '-m', 'pip', 'freeze'], stderr=subprocess.STDOUT)
+        lines = result.decode().splitlines()
+        return [line if versions else line.split('==')[0] for line in lines]
+    except subprocess.CalledProcessError as e:
+        print("An error occurred while fetching packages:", e.output.decode())
+        return []
+
+### INSTALL PACKAGE ###
+def install_package(package):
+    print(f"Installing package {package}...")
+    subprocess.check_call([sys.executable, '-s', '-m', 'pip', '-q', 'install', package])
+
+### READ REQUIREMENTS ###
+def read_requirements():
+    req_file = Path(__file__).parent / 'requirements.txt'
+    if not req_file.exists():
+        return []
+    with open(req_file, 'r') as f:
+        return [line.strip().split('>=')[0] for line in f if line.strip() and not line.startswith('#')]
+
+### Ensure all required packages are installed ###
+installed_packages = packages()
+required_packages = read_requirements()
+
+for package in required_packages:
+    if package.lower() not in [pkg.lower() for pkg in installed_packages]:
+        install_package(package)
+
+### Don't import nodes until all packages are installed ###
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
\ No newline at end of file
diff --git a/nodes.py b/nodes.py
new file mode 100644
index 0000000..691d21f
--- /dev/null
+++ b/nodes.py
@@ -0,0 +1,42 @@
+import io
+import typing as T
+import os
+import numpy as np
+from PIL import Image, ImageDraw
+from scipy.io import wavfile
+import torch
+import torchaudio
+from pedalboard.io import AudioFile
+import glob
+import wave
+
+MAX_BATCH_SIZE = 8
+
+class RiffusionNode:
+    def __init__(self):
+        pass
+        
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                'spectrogram': ("IMAGE",),
+                'output_type': ("STRING", {"default": "audio/wav"}),
+            },
+        }
+    
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "process"
+    CATEGORY = "audio"
+
+    def process(self, spectrogram, output_type):
+        # Your processing logic here
+        pass
+
+NODE_CLASS_MAPPINGS = {
+    "RiffusionNode": RiffusionNode
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "RiffusionNode": "Riffusion"
+}
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..5d5e232
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,15 @@
+[project]
+name = "riffusion"
+description = "A ComfyUI extension for Riffusion audio generation."
+version = "1.0.0"
+license = {file = "LICENSE"}
+dependencies = ["numpy>=1.21.0", "Pillow>=9.0.0", "scipy>=1.7.0", "torch>=1.9.0", "torchaudio>=0.9.0", "pedalboard>=0.7.0"]
+
+[project.urls]
+Repository = "https://github.com/traugdor/ComfyUI-Riffusion"
+#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.comfy]
+PublisherId = "traugdor"
+DisplayName = "ComfyUI-Riffusion"
+Icon = ""
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..1be5b19
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,6 @@
+numpy>=1.21.0
+Pillow>=9.0.0
+scipy>=1.7.0
+torch>=1.9.0
+torchaudio>=0.9.0
+pedalboard>=0.7.0