commit 679126592d37f80922c1b6ecb840d16e81986e16 Author: Sayo Date: Sat Feb 14 08:29:39 2026 +0800 feat: init diff --git a/.github/workflows/publish_action.yml b/.github/workflows/publish_action.yml new file mode 100644 index 0000000..a9b2587 --- /dev/null +++ b/.github/workflows/publish_action.yml @@ -0,0 +1,19 @@ +name: Publish to Comfy Registry + +on: + push: + branches: + - main + - master + paths: + - "pyproject.toml" + workflow_dispatch: + +jobs: + publish: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: Comfy-Org/publish-node-action@main + with: + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..03da3a1 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +.config/ +__pycache__/ +*.pyc diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..88f6757 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 ACE-Step + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..1889fc0 --- /dev/null +++ b/__init__.py @@ -0,0 +1,13 @@ +from .nodes import AceStepMusicGen + +NODE_CLASS_MAPPINGS = { + "AceStepMusicGen": AceStepMusicGen, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AceStepMusicGen": "ACE-Step Music Generate", +} + +WEB_DIRECTORY = "./js" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/js/ace_button.js b/js/ace_button.js new file mode 100644 index 0000000..9606f9b --- /dev/null +++ b/js/ace_button.js @@ -0,0 +1,86 @@ +import { app } from "../../scripts/app.js"; + +function hideWidget(node, widget) { + if (widget._hidden) return; + widget._hidden = true; + widget._origType = widget.type; + widget._origComputeSize = widget.computeSize; + widget.type = "hidden"; + widget.computeSize = () => [0, -4]; +} + +function showWidget(node, widget) { + if (!widget._hidden) return; + widget._hidden = false; + widget.type = widget._origType; + widget.computeSize = widget._origComputeSize; + delete widget._origType; + delete widget._origComputeSize; +} + +app.registerExtension({ + name: "AceStep.NodeExtensions", + + nodeCreated(node) { + if (node.comfyClass !== "AceStepMusicGen") return; + + const REPAINT_WIDGETS = ["repainting_start", "repainting_end"]; + const COVER_WIDGETS = ["audio_cover_strength"]; + + function updateVisibility() { + const taskWidget = node.widgets.find(w => w.name === "task_type"); + if (!taskWidget) return; + + const task = taskWidget.value; + + for (const w of node.widgets) { + if (REPAINT_WIDGETS.includes(w.name)) { + task === "repaint" ? showWidget(node, w) : hideWidget(node, w); + } else if (COVER_WIDGETS.includes(w.name)) { + task === "cover" ? showWidget(node, w) : hideWidget(node, w); + } + } + + const sz = node.computeSize(); + node.setSize([Math.max(sz[0], node.size[0]), sz[1]]); + } + + requestAnimationFrame(() => { + const taskWidget = node.widgets.find(w => w.name === "task_type"); + if (taskWidget) { + const origCallback = taskWidget.callback; + taskWidget.callback = function (...args) { + if (origCallback) origCallback.apply(this, args); + updateVisibility(); + }; + } + updateVisibility(); + }); + + // --- API Key masking --- + const apiKeyWidget = node.widgets.find(w => w.name === "api_key"); + if (apiKeyWidget) { + let realKey = ""; + + // Return real key for execution, not the masked display value + apiKeyWidget.serializeValue = async () => realKey; + + apiKeyWidget.callback = function (value) { + if (!value) { + realKey = ""; + return; + } + // Ignore if user submitted the mask unchanged + if (/^●+$/.test(value)) return; + realKey = value; + apiKeyWidget.value = "●".repeat(Math.min(value.length, 16)); + }; + } + + // --- Get API Key button --- + const btn = node.addWidget("button", "🔑 Get API Key", null, () => { + window.open("https://acemusic.ai/api-key", "_blank"); + }); + btn.serialize = false; + }, +}); diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..82c982a --- /dev/null +++ b/nodes.py @@ -0,0 +1,387 @@ +import json +import base64 +import io +import struct +import wave +import os + +import numpy as np +import torch +import requests + + +_CONFIG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".config") +_API_KEY_FILE = os.path.join(_CONFIG_DIR, "api_key") + + +def _save_api_key(key): + os.makedirs(_CONFIG_DIR, exist_ok=True) + with open(_API_KEY_FILE, "w", encoding="utf-8") as f: + f.write(key) + + +def _load_api_key(): + try: + with open(_API_KEY_FILE, "r", encoding="utf-8") as f: + return f.read().strip() + except FileNotFoundError: + return "" + + +class AceStepMusicGen: + """Generate music via ACE-Step OpenRouter-compatible API server.""" + + CATEGORY = "api node/audio/AceStep" + FUNCTION = "generate" + RETURN_TYPES = ("AUDIO", "STRING") + RETURN_NAMES = ("audio", "info") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "server_url": ("STRING", { + "default": "https://api.acemusic.ai", + "tooltip": "ACE-Step API server base URL", + }), + "prompt": ("STRING", { + "default": "", + "multiline": True, + "tooltip": "Music description / caption", + }), + }, + "optional": { + "lyrics": ("STRING", { + "default": "", + "multiline": True, + "tooltip": "Song lyrics (leave empty for instrumental)", + }), + "task_type": (["text2music", "cover", "repaint"], { + "default": "text2music", + }), + "instrumental": ("BOOLEAN", { + "default": False, + "tooltip": "Generate instrumental only (no vocals)", + }), + "duration": ("FLOAT", { + "default": 30.0, + "min": 1.0, + "max": 300.0, + "step": 1.0, + "tooltip": "Audio duration in seconds", + }), + "bpm": ("INT", { + "default": 0, + "min": 0, + "max": 300, + "step": 1, + "tooltip": "Beats per minute (0 = auto)", + }), + "key_scale": ("STRING", { + "default": "", + "tooltip": "Musical key and scale (e.g. 'C major', 'A minor')", + }), + "time_signature": ("STRING", { + "default": "", + "tooltip": "Time signature (e.g. '4/4', '3/4', '6/8')", + }), + "vocal_language": (["en", "zh", "ja", "ko", "es", "fr", "de", "unknown"], { + "default": "en", + }), + "sample_mode": ("BOOLEAN", { + "default": False, + "tooltip": "Use LLM to generate prompt/lyrics from query", + }), + "use_format": ("BOOLEAN", { + "default": False, + "tooltip": "Use LLM to enhance caption and lyrics", + }), + "batch_size": ("INT", { + "default": 1, + "min": 1, + "max": 2, + "step": 1, + }), + "seed": ("STRING", { + "default": "-1", + "tooltip": "Random seed (-1 = random). Comma-separated for batch (e.g. '42,123')", + }), + "temperature": ("FLOAT", { + "default": 0.85, + "min": 0.0, + "max": 2.0, + "step": 0.05, + }), + "guidance_scale": ("FLOAT", { + "default": 7.0, + "min": 0.0, + "max": 20.0, + "step": 0.5, + }), + "src_audio": ("AUDIO", { + "tooltip": "Source audio for cover/repaint tasks", + }), + "ref_audio": ("AUDIO", { + "tooltip": "Reference audio for style/timbre guidance", + }), + "repainting_start": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 300.0, + "step": 0.1, + "tooltip": "Repaint region start time in seconds", + }), + "repainting_end": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 300.0, + "step": 0.1, + "tooltip": "Repaint region end time in seconds (0 = full length)", + }), + "audio_cover_strength": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 1.0, + "step": 0.05, + "tooltip": "Cover strength (1.0 = full cover)", + }), + "api_key": ("STRING", { + "default": "", + "tooltip": "API key (auto-saved locally after first input, also reads ACESTEP_API_KEY env var)", + }), + }, + } + + @classmethod + def IS_CHANGED(cls, **kwargs): + return float("NaN") + + def generate( + self, + server_url, + prompt, + lyrics="", + task_type="text2music", + instrumental=False, + duration=30.0, + bpm=0, + key_scale="", + time_signature="", + vocal_language="en", + sample_mode=False, + use_format=False, + batch_size=1, + seed="-1", + temperature=0.85, + guidance_scale=7.0, + src_audio=None, + ref_audio=None, + repainting_start=0.0, + repainting_end=0.0, + audio_cover_strength=1.0, + api_key="", + ): + url = server_url.rstrip("/") + "/v1/chat/completions" + + # Build message content with audio inputs + # API routing: cover/repaint → audio[0]=src, audio[1]=ref; text2music → audio[0]=ref + audio_list = [] + if task_type in ("cover", "repaint"): + if src_audio is not None: + audio_list.append(src_audio) + if ref_audio is not None: + audio_list.append(ref_audio) + else: + if ref_audio is not None: + audio_list.append(ref_audio) + + content = self._build_content(prompt, audio_list) + + # Build request body + body = { + "model": "acemusic/acestep-v15-turbo", + "messages": [{"role": "user", "content": content}], + "modalities": ["audio"], + "stream": False, + "temperature": temperature, + "guidance_scale": guidance_scale, + "batch_size": batch_size, + "task_type": task_type, + "sample_mode": sample_mode, + "use_format": use_format, + "audio_config": { + "duration": duration, + "format": "wav", + "vocal_language": vocal_language, + "instrumental": instrumental, + }, + } + + # Seed: "-1" means random, otherwise pass as-is (supports comma-separated) + seed_str = seed.strip() + if seed_str and seed_str != "-1": + body["seed"] = seed_str + if bpm > 0: + body["audio_config"]["bpm"] = bpm + body["audio_config"]["key_scale"] = key_scale.strip() or None + body["audio_config"]["time_signature"] = time_signature.strip() or None + if lyrics: + body["lyrics"] = lyrics + + # Repaint/cover specific parameters (only when relevant) + if task_type == "repaint": + body["repainting_start"] = repainting_start + if repainting_end > 0.0: + body["repainting_end"] = repainting_end + if task_type == "cover": + body["audio_cover_strength"] = audio_cover_strength + + # HTTP request – resolve API key: widget input > env var > saved config + resolved_key = api_key.strip() if api_key else "" + if resolved_key: + _save_api_key(resolved_key) + os.environ["ACESTEP_API_KEY"] = resolved_key + else: + resolved_key = os.environ.get("ACESTEP_API_KEY", "") or _load_api_key() + headers = { + "Content-Type": "application/json", + "User-Agent": "AceMusic-ComfyUI/1.0", + } + if resolved_key: + headers["Authorization"] = f"Bearer {resolved_key}" + + try: + resp = requests.post( + url, + data=json.dumps(body, ensure_ascii=False).encode("utf-8"), + headers=headers, + timeout=600, + ) + resp.raise_for_status() + result = resp.json() + except requests.exceptions.HTTPError as e: + raise RuntimeError(f"API error {resp.status_code}: {resp.text}") + except requests.exceptions.ConnectionError as e: + raise RuntimeError(f"Cannot connect to {server_url}: {e}") + except requests.exceptions.Timeout: + raise RuntimeError(f"Request to {server_url} timed out") + + # Parse response + choices = result.get("choices", []) + if not choices: + raise RuntimeError("API returned no choices") + + message = choices[0].get("message", {}) + text_content = message.get("content", "") + audio_items = message.get("audio", []) + + if not audio_items: + raise RuntimeError(f"API returned no audio. Response: {text_content}") + + # Decode all audio items and batch them + waveforms = [] + sample_rate = None + for item in audio_items: + audio_url = item.get("audio_url", {}).get("url", "") + if not audio_url: + continue + wf, sr = self._decode_audio_data_url(audio_url) + # wf shape: (1, channels, samples) + waveforms.append(wf) + sample_rate = sr + + if not waveforms: + raise RuntimeError("API returned no valid audio data") + + # Pad to same length and concat along batch dim + max_len = max(w.shape[-1] for w in waveforms) + padded = [] + for wf in waveforms: + if wf.shape[-1] < max_len: + pad_size = max_len - wf.shape[-1] + wf = torch.nn.functional.pad(wf, (0, pad_size)) + padded.append(wf) + + # (N, channels, samples) + audio_tensor = torch.cat(padded, dim=0) + + return ( + {"waveform": audio_tensor, "sample_rate": sample_rate}, + text_content, + ) + + def _build_content(self, prompt, audio_list): + """Build message content, optionally with multimodal audio inputs.""" + if not audio_list: + return prompt + + # Multimodal: text + input_audio(s) + parts = [] + + if prompt: + parts.append({"type": "text", "text": prompt}) + + for audio in audio_list: + waveform = audio["waveform"].squeeze(0) # (channels, samples) + sr = audio["sample_rate"] + b64 = self._tensor_to_wav_b64(waveform, sr) + parts.append({ + "type": "input_audio", + "input_audio": {"data": b64, "format": "wav"}, + }) + + return parts + + def _decode_audio_data_url(self, data_url): + """Decode base64 data URL to torch audio tensor using stdlib wave.""" + if "," in data_url: + b64_data = data_url.split(",", 1)[1] + else: + b64_data = data_url + + audio_bytes = base64.b64decode(b64_data) + buf = io.BytesIO(audio_bytes) + + with wave.open(buf, "rb") as wf: + n_channels = wf.getnchannels() + sampwidth = wf.getsampwidth() + sample_rate = wf.getframerate() + n_frames = wf.getnframes() + raw = wf.readframes(n_frames) + + # Use numpy for fast PCM decoding + if sampwidth == 2: + samples = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 + elif sampwidth == 4: + samples = np.frombuffer(raw, dtype=np.int32).astype(np.float32) / 2147483648.0 + else: + raise RuntimeError(f"Unsupported WAV sample width: {sampwidth}") + + # Interleaved -> (n_frames, n_channels) -> (n_channels, n_frames) + samples = samples.reshape(n_frames, n_channels).T + # Add batch dim: (1, n_channels, n_frames) + tensor = torch.from_numpy(samples).unsqueeze(0) + return tensor, sample_rate + + def _tensor_to_wav_b64(self, waveform, sample_rate): + """Encode torch tensor to WAV base64 string using stdlib wave.""" + # waveform shape: (channels, samples) + n_channels = waveform.shape[0] + n_frames = waveform.shape[1] + + # Clamp and convert to 16-bit PCM + clamped = waveform.clamp(-1.0, 1.0) + pcm = (clamped * 32767).to(torch.int16) + + # Interleave channels: (channels, samples) -> (samples * channels,) + interleaved = pcm.T.contiguous().reshape(-1) + raw = struct.pack(f"<{interleaved.numel()}h", *interleaved.tolist()) + + buf = io.BytesIO() + with wave.open(buf, "wb") as wf: + wf.setnchannels(n_channels) + wf.setsampwidth(2) + wf.setframerate(sample_rate) + wf.writeframes(raw) + + return base64.b64encode(buf.getvalue()).decode("utf-8") diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..ba08d1f --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,28 @@ +[project] +name = "ACE-Step-ComfyUI" +description = "Official ComfyUI node for ACE-Step AI music generation via ACE-Step API" +version = "0.0.1" +license = {file = "LICENSE"} +dependencies = [ + "torch", + "numpy", + "requests", +] +classifiers = [ + "Operating System :: OS Independent", +] + + +[project.urls] +Repository = "https://github.com/ace-step/ACE-Step-ComfyUI" +# Used by Comfy Registry https://registry.comfy.org +Documentation = "https://github.com/ace-step/ACE-Step-ComfyUI/wiki" +"Bug Tracker" = "https://github.com/ace-step/ACE-Step-ComfyUI/issues" + +[tool.comfy] +PublisherId = "ace-step" +DisplayName = "ACE-Step-ComfyUI" +Icon = "" +includes = ["__init__.py", "nodes.py", "requirements.txt", "js/"] +# "requires-comfyui" = ">=1.0.0" # ComfyUI version compatibility + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..8ad30cc --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +torch +numpy +requests