From cd22a28d3d7cb468337755b7bab255e294f634f2 Mon Sep 17 00:00:00 2001 From: Dreaming Computers Date: Sun, 1 Feb 2026 12:33:12 +0700 Subject: [PATCH] feat: suno music polling and audio output --- README.md | 4 +- kie_api/audio.py | 37 ++++++++++++ kie_api/suno_music.py | 120 ++++++++++++++++++++++++++++++++++++- nodes.py | 11 ++-- web/docs/KIE_Suno_Music.md | 3 +- 5 files changed, 164 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 1eed283..60cf110 100644 --- a/README.md +++ b/README.md @@ -55,8 +55,8 @@ This node pack currently includes the following nodes: - Supports role selection, media inputs (images/video/audio), optional reasoning output, and Google Search toggle. - **Suno Music (Generate)** - - Creates a music generation task and returns a task ID. - - Uses KIE Suno API `generate` endpoint. + - Generates music and returns AUDIO output. + - Uses KIE Suno API `generate` + `record-info` polling. - **Kling 2.6 Image-to-Video** - Generates video from a single input image. diff --git a/kie_api/audio.py b/kie_api/audio.py index 221cf8c..671a551 100644 --- a/kie_api/audio.py +++ b/kie_api/audio.py @@ -78,3 +78,40 @@ def _waveform_to_wav_bytes(waveform, sample_rate: int) -> bytes: wf.setframerate(sample_rate) wf.writeframes(pcm.tobytes()) return buffer.getvalue() + + +def _audio_bytes_to_comfy_audio(audio_bytes: bytes, filename_hint: str = "audio.mp3"): + """Decode audio bytes into a ComfyUI AUDIO dict.""" + if not isinstance(audio_bytes, (bytes, bytearray)) or not audio_bytes: + raise RuntimeError("audio_bytes must be non-empty bytes.") + + try: + from tempfile import gettempdir + from time import time + from pathlib import Path + tmp_dir = Path(gettempdir()) + suffix = Path(filename_hint).suffix or ".mp3" + tmp_path = tmp_dir / f"kie_audio_{int(time() * 1000)}{suffix}" + tmp_path.write_bytes(audio_bytes) + except Exception as exc: + raise RuntimeError(f"Failed to write temp audio file: {exc}") from exc + + # Try torchaudio first + try: + import torchaudio + waveform, sample_rate = torchaudio.load(str(tmp_path)) + return {"waveform": waveform, "sample_rate": sample_rate, "path": str(tmp_path)} + except Exception: + pass + + # Fallback to soundfile + try: + import soundfile as sf + import torch + data, sample_rate = sf.read(str(tmp_path), always_2d=True) + waveform = torch.from_numpy(data.T).float() + return {"waveform": waveform, "sample_rate": sample_rate, "path": str(tmp_path)} + except Exception as exc: + raise RuntimeError( + "Failed to decode audio. Install torchaudio or soundfile to enable audio output." + ) from exc diff --git a/kie_api/suno_music.py b/kie_api/suno_music.py index 4d2d66f..30ac725 100644 --- a/kie_api/suno_music.py +++ b/kie_api/suno_music.py @@ -1,15 +1,26 @@ """KIE Suno music generation helper (v1 create-only).""" import json +import time from typing import Any from .auth import _load_api_key +from .audio import _audio_bytes_to_comfy_audio from .http import TransientKieError, requests from .log import _log GENERATE_URL = "https://api.kie.ai/api/v1/generate" +RECORD_INFO_URL = "https://api.kie.ai/api/v1/generate/record-info" MODEL_OPTIONS = ["V4", "V4_5", "V4_5PLUS", "V4_5ALL", "V5"] VOCAL_GENDER_OPTIONS = ["m", "f"] +POLLABLE_STATES = {"PENDING", "TEXT_SUCCESS", "FIRST_SUCCESS"} +SUCCESS_STATE = "SUCCESS" +FAIL_STATES = { + "CREATE_TASK_FAILED", + "GENERATE_AUDIO_FAILED", + "CALLBACK_EXCEPTION", + "SENSITIVE_WORD_ERROR", +} def _validate_length(field_name: str, value: str | None, max_len: int) -> None: @@ -27,6 +38,89 @@ def _max_style_len(model: str) -> int: return 200 if model == "V4" else 1000 +def _fetch_music_record(api_key: str, task_id: str) -> dict[str, Any]: + try: + response = requests.get( + RECORD_INFO_URL, + headers={"Authorization": f"Bearer {api_key}"}, + params={"taskId": task_id}, + timeout=30, + ) + except requests.RequestException as exc: + raise RuntimeError(f"Failed to call record-info endpoint: {exc}") from exc + + if response.status_code == 429 or response.status_code >= 500: + raise TransientKieError( + f"record-info returned HTTP {response.status_code}: {response.text}", + status_code=response.status_code, + ) + + try: + payload_json: Any = response.json() + except json.JSONDecodeError as exc: + raise RuntimeError("record-info endpoint did not return valid JSON.") from exc + + if payload_json.get("code") != 200: + message = payload_json.get("msg") or payload_json.get("message") + raise RuntimeError(f"record-info returned error code {payload_json.get('code')}: {message}") + + data = payload_json.get("data") + if data is None: + raise RuntimeError("record-info endpoint returned no data field.") + + return data + + +def _extract_audio_urls(record_data: dict[str, Any]) -> list[str]: + # Expected: record_data["data"] is a list of items with audio_url + items = record_data.get("data") + urls: list[str] = [] + if isinstance(items, list): + for item in items: + if isinstance(item, dict) and item.get("audio_url"): + urls.append(item["audio_url"]) + if not urls and isinstance(record_data, dict) and record_data.get("audio_url"): + urls.append(record_data["audio_url"]) + if not urls: + raise RuntimeError("No audio_url found in record-info response.") + return urls + + +def _poll_music_until_complete( + api_key: str, + task_id: str, + poll_interval_s: float, + timeout_s: int, + log: bool, +) -> dict[str, Any]: + start_time = time.time() + last_state = None + + while True: + elapsed = time.time() - start_time + if elapsed > timeout_s: + raise RuntimeError(f"Task {task_id} timed out after {timeout_s}s.") + + record = _fetch_music_record(api_key, task_id) + state = record.get("status") or record.get("state") or record.get("callbackType") + + if log and state != last_state: + _log(log, f"Suno task {task_id} state: {state}") + last_state = state + + if state == SUCCESS_STATE or state == "complete": + return record + if state in FAIL_STATES or state == "error": + raise RuntimeError(f"Suno task {task_id} failed with state: {state}") + + if state in POLLABLE_STATES or state is None: + time.sleep(poll_interval_s) + continue + + # Unknown state, keep polling conservatively + time.sleep(poll_interval_s) + + def run_suno_generate( *, prompt: str, @@ -42,9 +136,11 @@ def run_suno_generate( weirdness_constraint: float | None = None, audio_weight: float | None = None, persona_id: str | None = None, + poll_interval_s: float = 30.0, + timeout_s: int = 1800, log: bool = True, -) -> tuple[str, str]: - """Create a Suno music generation task. Returns (task_id, raw_json).""" +) -> tuple[dict, str, str]: + """Create a Suno music generation task and return audio output.""" if model not in MODEL_OPTIONS: raise RuntimeError("Invalid model. Use the pinned enum options.") if vocal_gender and vocal_gender not in VOCAL_GENDER_OPTIONS: @@ -128,4 +224,22 @@ def run_suno_generate( if log: _log(log, f"Suno task created: {task_id} (model={model})") - return task_id, json.dumps(payload_json) + record = _poll_music_until_complete( + api_key, + task_id, + poll_interval_s=poll_interval_s, + timeout_s=timeout_s, + log=log, + ) + audio_urls = _extract_audio_urls(record) + audio_url = audio_urls[0] + if log: + _log(log, f"Suno audio URL: {audio_url}") + + try: + audio_bytes = requests.get(audio_url, timeout=180).content + except requests.RequestException as exc: + raise RuntimeError(f"Failed to download audio: {exc}") from exc + + audio_output = _audio_bytes_to_comfy_audio(audio_bytes, "audio.mp3") + return audio_output, task_id, json.dumps(payload_json) diff --git a/nodes.py b/nodes.py index e724be6..2be8fde 100644 --- a/nodes.py +++ b/nodes.py @@ -743,7 +743,8 @@ Inputs: - persona_id: Optional persona (custom mode only) - log: Console logging on/off -Outputs: + Outputs: +- AUDIO: Generated audio - STRING: task_id - STRING: raw_json """ @@ -771,8 +772,8 @@ Outputs: }, } - RETURN_TYPES = ("STRING", "STRING") - RETURN_NAMES = ("task_id", "raw_json") + RETURN_TYPES = ("AUDIO", "STRING", "STRING") + RETURN_NAMES = ("audio", "task_id", "raw_json") FUNCTION = "generate" CATEGORY = "kie/api" @@ -793,7 +794,7 @@ Outputs: persona_id: str = "", log: bool = True, ): - task_id, raw_json = run_suno_generate( + audio_output, task_id, raw_json = run_suno_generate( prompt=prompt, custom_mode=custom_mode, instrumental=instrumental, @@ -809,7 +810,7 @@ Outputs: persona_id=persona_id, log=log, ) - return (task_id, raw_json) + return (audio_output, task_id, raw_json) class KIE_GridSlice: diff --git a/web/docs/KIE_Suno_Music.md b/web/docs/KIE_Suno_Music.md index e6dad01..85d86e4 100644 --- a/web/docs/KIE_Suno_Music.md +++ b/web/docs/KIE_Suno_Music.md @@ -1,6 +1,6 @@ # KIE Suno Music (Generate) -Create a Suno music generation task via KIE API. This v1 node returns a task ID; retrieval of the final audio is handled via callback or a future polling node. +Create a Suno music generation task via KIE API. This node polls until completion and returns an AUDIO output. ## Inputs - **prompt** (STRING, required): Prompt text (lyrics in custom mode when instrumental is false). @@ -21,6 +21,7 @@ Optional: - **log** (BOOLEAN): Enable console logging. ## Outputs +- **audio** (AUDIO): Generated audio. - **task_id** (STRING): Task ID returned by the API. - **raw_json** (STRING): Full API response JSON.