diff --git a/comfy-nodes/audio_duration_frames.py b/comfy-nodes/audio_duration_frames.py
index a7799a5..6357ea5 100644
--- a/comfy-nodes/audio_duration_frames.py
+++ b/comfy-nodes/audio_duration_frames.py
@@ -42,7 +42,7 @@ class AudioDurationFrames:
"end_time",
)
FUNCTION = "calculate"
- CATEGORY = "llm-toolkit"
+ CATEGORY = "llm_toolkit/utils"
# ---------------------------------------------------------------------
# Core logic
@@ -117,5 +117,5 @@ NODE_CLASS_MAPPINGS = {
}
NODE_DISPLAY_NAME_MAPPINGS = {
- "AudioDurationFrames": "Audio Duration → Frames",
+ "AudioDurationFrames": "Audio Duration → Frames (LLMToolkit)",
}
\ No newline at end of file
diff --git a/comfy-nodes/bfl_provider.py b/comfy-nodes/bfl_provider.py
index ef973ec..ef70887 100644
--- a/comfy-nodes/bfl_provider.py
+++ b/comfy-nodes/bfl_provider.py
@@ -58,7 +58,7 @@ class BFLProviderNode:
RETURN_TYPES = ("*",)
RETURN_NAMES = ("context",)
FUNCTION = "configure_bfl"
- CATEGORY = "llm_toolkit"
+ CATEGORY = "llm_toolkit/providers"
@classmethod
def IS_CHANGED(cls, llm_model: str, context: Any = None):
diff --git a/comfy-nodes/blank_image.py b/comfy-nodes/blank_image.py
new file mode 100644
index 0000000..915078f
--- /dev/null
+++ b/comfy-nodes/blank_image.py
@@ -0,0 +1,66 @@
+import logging
+from typing import Tuple
+import torch
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+class BlankImage:
+ """ComfyUI node that creates a solid-color image tensor of the desired size.
+
+ The color is chosen via a color-picker UI (hex string) and size via width/height ints.
+ Output format: torch.Tensor shaped (1, 3, H, W) with values in [0,1].
+ """
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "width": ("INT", {"default": 512, "min": 1, "max": 4096, "step": 1, "tooltip": "Image width in pixels"}),
+ "height": ("INT", {"default": 512, "min": 1, "max": 4096, "step": 1, "tooltip": "Image height in pixels"}),
+ "color": ("COLOR", {"default": "#000000", "tooltip": "Solid fill color"}),
+ },
+ }
+
+ RETURN_TYPES = ("IMAGE",)
+ RETURN_NAMES = ("image",)
+ FUNCTION = "create"
+ CATEGORY = "llm_toolkit/utils"
+
+ # ------------------------------------------------------------------
+ # Core logic
+ # ------------------------------------------------------------------
+ def _hex_to_rgb(self, hex_color: str) -> Tuple[float, float, float]:
+ """Convert hex string (6- or 8-digit) to 0-1 float RGB tuple."""
+ hex_color = hex_color.lstrip("#")
+ if len(hex_color) not in {6, 8}:
+ raise ValueError("Invalid hex color string")
+ r = int(hex_color[0:2], 16)
+ g = int(hex_color[2:4], 16)
+ b = int(hex_color[4:6], 16)
+ return r / 255.0, g / 255.0, b / 255.0
+
+ def create(self, width: int, height: int, color: str):
+ try:
+ r, g, b = self._hex_to_rgb(color)
+ arr = np.stack([
+ np.full((height, width), r, dtype=np.float32),
+ np.full((height, width), g, dtype=np.float32),
+ np.full((height, width), b, dtype=np.float32),
+ ], axis=0) # shape (3, H, W)
+ img_tensor = torch.from_numpy(arr).unsqueeze(0) # (1,3,H,W)
+ return (img_tensor,)
+ except Exception as e:
+ logger.error("BlankImage: Failed to create image – %s", e, exc_info=True)
+ return (torch.zeros((1, 3, height, width), dtype=torch.float32),)
+
+# --------------------------------------------------------------------
+# Mappings for ComfyUI
+# --------------------------------------------------------------------
+NODE_CLASS_MAPPINGS = {
+ "BlankImage": BlankImage,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+ "BlankImage": "Blank Image (llm_toolkit)",
+}
\ No newline at end of file
diff --git a/comfy-nodes/check_image_empty.py b/comfy-nodes/check_image_empty.py
index 96a7842..561e9c6 100644
--- a/comfy-nodes/check_image_empty.py
+++ b/comfy-nodes/check_image_empty.py
@@ -28,7 +28,7 @@ class CheckImageEmpty:
RETURN_NAMES = ("is_empty", "status_message")
FUNCTION = "check_image_empty"
OUTPUT_NODE = False # This is a utility node, not an output node
- CATEGORY = "llm_toolkit"
+ CATEGORY = "llm_toolkit/utils"
def check_image_empty(self, image: Optional[torch.Tensor] = None) -> Tuple[bool, str]:
"""
diff --git a/comfy-nodes/config_generate_image.py b/comfy-nodes/config_generate_image.py
index 12a2658..2c6abbe 100644
--- a/comfy-nodes/config_generate_image.py
+++ b/comfy-nodes/config_generate_image.py
@@ -68,7 +68,7 @@ class ConfigGenerateImage:
RETURN_TYPES = ("*",)
RETURN_NAMES = ("context",)
FUNCTION = "configure"
- CATEGORY = "llm_toolkit"
+ CATEGORY = "llm_toolkit/utils"
def configure(self, context: Optional[Dict[str, Any]] = None, **kwargs) -> Tuple[Dict[str, Any]]:
"""
diff --git a/comfy-nodes/display_text.py b/comfy-nodes/display_text.py
index f076fdf..c965e1d 100644
--- a/comfy-nodes/display_text.py
+++ b/comfy-nodes/display_text.py
@@ -16,6 +16,16 @@ logger = logging.getLogger(__name__)
# Helper to extract context from payload objects
from context_payload import extract_context
+def _remove_thinking_tags(text: str) -> str:
+ """Remove ... blocks from text, including the tags themselves."""
+ import re
+ # Pattern to match ... blocks (including nested content and newlines)
+ pattern = r'.*?'
+ cleaned = re.sub(pattern, '', text, flags=re.DOTALL)
+ # Clean up any extra whitespace/newlines left behind
+ cleaned = re.sub(r'\n\s*\n\s*\n', '\n\n', cleaned) # Replace multiple newlines with double
+ return cleaned.strip()
+
class Display_Text:
"""
Displays text extracted from a wildcard input type, typically containing LLM responses.
@@ -33,6 +43,7 @@ class Display_Text:
"default": "0",
"tooltip": "Select which line to output (cycles through available lines)"
}),
+ "hide_thinking": ("BOOLEAN", {"default": True, "tooltip": "Hide model thinking process (content between tags)"})
},
"optional": {},
"hidden": {},
@@ -44,9 +55,9 @@ class Display_Text:
OUTPUT_IS_LIST = (False, True, False, False, False) # text_list is the only list output
FUNCTION = "display_llm_text"
OUTPUT_NODE = True
- CATEGORY = "llm_toolkit" # Changed category to llm_toolkit
+ CATEGORY = "llm_toolkit/utils" # Changed category to llm_toolkit
- def display_llm_text(self, context: Any, select: str):
+ def display_llm_text(self, context: Any, select: str, hide_thinking: bool):
# --- Safe conversion for 'select' input string ---
select_int = 0 # Default value
try:
@@ -116,6 +127,10 @@ class Display_Text:
text_to_display = ""
# --- End Text Extraction ---
+ # Apply thinking tag removal if requested
+ if hide_thinking and text_to_display:
+ text_to_display = _remove_thinking_tags(text_to_display)
+
# Use the extracted text for display logic
if text_to_display is None: # Should not happen with default ""
logger.error("text_to_display is None unexpectedly.")
diff --git a/comfy-nodes/generate_image.py b/comfy-nodes/generate_image.py
index a5e2afc..ce989cf 100644
--- a/comfy-nodes/generate_image.py
+++ b/comfy-nodes/generate_image.py
@@ -72,7 +72,7 @@ listen to the heartbeat of a baby otter.
RETURN_TYPES = ("*", "IMAGE",) # Keep IMAGE for direct use/preview
RETURN_NAMES = ("context", "image",)
FUNCTION = "generate"
- CATEGORY = "llm_toolkit"
+ CATEGORY = "llm_toolkit/generators"
OUTPUT_NODE = True
def generate(self, prompt: str, mode: str = "generate", context: Optional[Dict[str, Any]] = None) -> Tuple[Dict[str, Any], Any]:
diff --git a/comfy-nodes/generate_music.py b/comfy-nodes/generate_music.py
new file mode 100644
index 0000000..f69e71a
--- /dev/null
+++ b/comfy-nodes/generate_music.py
@@ -0,0 +1,253 @@
+import os
+import sys
+import logging
+import tempfile
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+import requests
+import torchaudio
+
+# Ensure parent directory (project root) is on sys.path for imports
+current_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(current_dir)
+if parent_dir not in sys.path:
+ sys.path.insert(0, parent_dir)
+
+from suno_api import (
+ send_suno_music_generation_request,
+ send_suno_lyrics_generation_request,
+)
+from llmtoolkit_utils import get_api_key
+from context_payload import extract_context
+from send_request import run_async
+
+logger = logging.getLogger(__name__)
+
+
+class GenerateMusic: # noqa: N801 – follow existing naming pattern
+ """Generate music with Suno API and return an AUDIO tensor compatible with nodes_audio."""
+
+ DEFAULT_PROVIDER = "suno"
+ DEFAULT_MODEL = "V3_5"
+ DEFAULT_PROMPT = "A calm and relaxing piano track with soft melodies"
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "prompt": (
+ "STRING",
+ {"multiline": True, "default": cls.DEFAULT_PROMPT},
+ ),
+ },
+ "optional": {
+ "context": ("*", {}),
+ },
+ }
+
+ RETURN_TYPES = ("*", "AUDIO")
+ RETURN_NAMES = ("context", "audio")
+ FUNCTION = "generate"
+ CATEGORY = "llm_toolkit/generators"
+ OUTPUT_NODE = True
+
+ # ---------------------------------------------------------------------
+ # Helper utilities
+ # ---------------------------------------------------------------------
+ def _download_audio_to_tensor(self, audio_url: str) -> Dict[str, Any]:
+ """Download a remote audio file (MP3/OGG/WAV) and convert to ComfyUI audio tensor."""
+ try:
+ with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
+ logger.info("Downloading generated audio from %s", audio_url)
+ with requests.get(audio_url, stream=True, timeout=60) as resp:
+ resp.raise_for_status()
+ for chunk in resp.iter_content(chunk_size=8192):
+ if chunk:
+ tmp.write(chunk)
+ tmp.flush()
+ temp_path = tmp.name
+
+ waveform, sample_rate = torchaudio.load(temp_path)
+ os.unlink(temp_path)
+
+ # Ensure batch dimension [B, C, T]
+ if waveform.dim() == 2:
+ waveform = waveform.unsqueeze(0)
+
+ return {"waveform": waveform, "sample_rate": sample_rate}
+ except Exception as e:
+ logger.error("Failed to download or decode audio: %s", e, exc_info=True)
+ # Return 1-second silent placeholder tensor on failure
+ return {
+ "waveform": torch.zeros((1, 1, 44100)),
+ "sample_rate": 44100,
+ }
+
+ # ---------------------------------------------------------------------
+ # Main generate() implementation
+ # ---------------------------------------------------------------------
+ def generate(self, prompt: str, context: Optional[Dict[str, Any]] = None) -> Tuple[Dict[str, Any], Any]:
+ if context is None:
+ context = {}
+ elif not isinstance(context, dict):
+ context = extract_context(context) or {}
+
+ provider_cfg = context.get("provider_config", {})
+ llm_provider = provider_cfg.get("provider_name", self.DEFAULT_PROVIDER)
+
+ if llm_provider != "suno":
+ err = f"GenerateMusic node supports only provider 'suno'. Received '{llm_provider}'."
+ logger.error(err)
+ context["error"] = err
+ placeholder = {"waveform": torch.zeros((1, 1, 44100)), "sample_rate": 44100}
+ return (context, placeholder)
+
+ # ------------------------------------------------------------------
+ # API Key resolution
+ # ------------------------------------------------------------------
+ api_key = provider_cfg.get("api_key", "")
+ if not api_key or api_key in {"", "1234"}:
+ try:
+ api_key = get_api_key("SUNO_API_KEY", "suno")
+ except ValueError:
+ api_key = ""
+
+ if not api_key:
+ err = "Missing Suno API key. Please set SUNO_API_KEY environment variable or pass via ProviderSelector."
+ logger.error(err)
+ context["error"] = err
+ placeholder = {"waveform": torch.zeros((1, 1, 44100)), "sample_rate": 44100}
+ return (context, placeholder)
+
+ # ------------------------------------------------------------------
+ # Read generation_config overrides from context (optional)
+ # ------------------------------------------------------------------
+ gen_cfg = context.get("generation_config", {})
+ style = gen_cfg.get("style")
+ title = gen_cfg.get("title")
+ custom_mode = gen_cfg.get("custom_mode", False)
+ instrumental = gen_cfg.get("instrumental", False)
+ model_name = provider_cfg.get("llm_model") or gen_cfg.get("model") or self.DEFAULT_MODEL
+ negative_tags = gen_cfg.get("negative_tags")
+ callback_url = gen_cfg.get("callback_url", "https://example.com/callback")
+
+ # ------------------------------------------------------------------
+ # Call Suno async helper via run_async
+ # ------------------------------------------------------------------
+ api_response = run_async(
+ send_suno_music_generation_request(
+ api_key=api_key,
+ prompt=prompt,
+ style=style,
+ title=title,
+ custom_mode=custom_mode,
+ instrumental=instrumental,
+ model=model_name,
+ negative_tags=negative_tags,
+ callback_url=callback_url,
+ poll=True,
+ )
+ )
+
+ audio_tensor_dict: Optional[Dict[str, Any]] = None
+ if isinstance(api_response, dict):
+ data_block = api_response.get("data", {})
+ tracks = data_block.get("data", []) if isinstance(data_block, dict) else []
+ if tracks:
+ first_track = tracks[0]
+ audio_url = first_track.get("audio_url") or first_track.get("source_audio_url")
+ if audio_url:
+ audio_tensor_dict = self._download_audio_to_tensor(audio_url)
+ context["suno_audio_url"] = audio_url
+
+ if audio_tensor_dict is None:
+ logger.warning("Audio not ready or failed to download. Returning silent placeholder.")
+ audio_tensor_dict = {"waveform": torch.zeros((1, 1, 44100)), "sample_rate": 44100}
+ context["error"] = "Audio generation failed or timed out."
+
+ context["suno_raw_response"] = api_response
+ return (context, audio_tensor_dict)
+
+
+class GenerateLyrics: # noqa: N801
+ """Generate song lyrics with Suno API."""
+
+ DEFAULT_PROMPT = "A song about peaceful night in the city"
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "prompt": (
+ "STRING",
+ {"multiline": True, "default": cls.DEFAULT_PROMPT},
+ ),
+ },
+ "optional": {
+ "context": ("*", {}),
+ },
+ }
+
+ RETURN_TYPES = ("*", "STRING")
+ RETURN_NAMES = ("context", "lyrics")
+ FUNCTION = "generate_lyrics"
+ CATEGORY = "llm_toolkit/generators"
+ OUTPUT_NODE = True
+
+ def generate_lyrics(self, prompt: str, context: Optional[Dict[str, Any]] = None):
+ if context is None:
+ context = {}
+ elif not isinstance(context, dict):
+ context = extract_context(context) or {}
+
+ provider_cfg = context.get("provider_config", {})
+ llm_provider = provider_cfg.get("provider_name", "suno")
+ if llm_provider != "suno":
+ err = f"GenerateLyrics node supports only provider 'suno'. Got '{llm_provider}'."
+ context["error"] = err
+ return (context, err)
+
+ api_key = provider_cfg.get("api_key", "")
+ if not api_key or api_key in {"", "1234"}:
+ try:
+ api_key = get_api_key("SUNO_API_KEY", "suno")
+ except ValueError:
+ api_key = ""
+ if not api_key:
+ err = "Missing Suno API key."
+ context["error"] = err
+ return (context, err)
+
+ api_response = run_async(
+ send_suno_lyrics_generation_request(
+ api_key=api_key,
+ prompt=prompt,
+ callback_url="https://example.com/callback",
+ poll=True,
+ )
+ )
+
+ lyrics_text = ""
+ if isinstance(api_response, dict):
+ data_block = api_response.get("data", {})
+ lyrics_list = data_block.get("lyricsData", [])
+ if lyrics_list:
+ lyrics_text = lyrics_list[0].get("text", "")
+
+ context["suno_lyrics_response"] = api_response
+ return (context, lyrics_text)
+
+
+# -------------------------------------------------------------------------
+# Node registration for ComfyUI
+# -------------------------------------------------------------------------
+NODE_CLASS_MAPPINGS = {
+ "GenerateMusic": GenerateMusic,
+ "GenerateLyrics": GenerateLyrics,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+ "GenerateMusic": "Generate Music (LLMToolkit)",
+ "GenerateLyrics": "Generate Lyrics (LLMToolkit)",
+}
\ No newline at end of file
diff --git a/comfy-nodes/generate_text.py b/comfy-nodes/generate_text.py
index 165a849..216b92f 100644
--- a/comfy-nodes/generate_text.py
+++ b/comfy-nodes/generate_text.py
@@ -315,6 +315,16 @@ async def send_request_stream(
yield f"[Error: {e}]"
return # Stop generation after yielding the fallback response
+def _remove_thinking_tags(text: str) -> str:
+ """Remove ... blocks from text, including the tags themselves."""
+ import re
+ # Pattern to match ... blocks (including nested content and newlines)
+ pattern = r'.*?'
+ cleaned = re.sub(pattern, '', text, flags=re.DOTALL)
+ # Clean up any extra whitespace/newlines left behind
+ cleaned = re.sub(r'\n\s*\n\s*\n', '\n\n', cleaned) # Replace multiple newlines with double
+ return cleaned.strip()
+
# --- Original Node (for reference or non-streaming use) ---
class LLMToolkitTextGenerator:
DEFAULT_PROVIDER = "openai"
@@ -328,7 +338,8 @@ class LLMToolkitTextGenerator:
return {
"required": {
"llm_model": (cls.MODEL_LIST, {"default": cls.DEFAULT_MODEL}),
- "prompt": ("STRING", {"multiline": False, "default": "Write a short story about a robot learning to paint."})
+ "prompt": ("STRING", {"multiline": False, "default": "Write a short story about a robot learning to paint."}),
+ "hide_thinking": ("BOOLEAN", {"default": True, "tooltip": "Hide model thinking process (content between tags)"})
},
"optional": {
"context": ("*", {})
@@ -338,10 +349,10 @@ class LLMToolkitTextGenerator:
RETURN_TYPES = ("*", "STRING")
RETURN_NAMES = ("context", "text")
FUNCTION = "generate"
- CATEGORY = "llm_toolkit"
+ CATEGORY = "llm_toolkit/generators"
OUTPUT_NODE = True # Keeps the text widget for non-streaming version
- def generate(self, llm_model, prompt, context=None):
+ def generate(self, llm_model, prompt, hide_thinking, context=None):
# ... (original generate logic using run_async(send_request(...))) ...
# This function remains mostly the same as the user provided,
# calling the original non-streaming send_request.
@@ -480,16 +491,26 @@ class LLMToolkitTextGenerator:
logger.error(f"Error in non-streaming send_request call: {e}", exc_info=True)
response_data = {"choices": [{"message": {"content": f"Error calling send_request: {str(e)}"}}]}
- if response_data is None: content = "Error: Received None response"
+ if response_data is None:
+ content = "Error: Received None response"
elif isinstance(response_data, dict):
if "choices" in response_data and response_data["choices"]:
message = response_data["choices"][0].get("message", {})
content = message.get("content", "")
- if content is None: content = "Error: Null content in response"
- elif "response" in response_data: content = response_data["response"]
- else: content = f"Error: Unexpected format: {str(response_data)}"
- elif isinstance(response_data, str): content = response_data
- else: content = f"Error: Unexpected type: {type(response_data)}"
+ if content is None:
+ content = "Error: Null content in response"
+ elif "response" in response_data:
+ content = response_data["response"]
+ else:
+ content = f"Error: Unexpected format: {str(response_data)}"
+ elif isinstance(response_data, str):
+ content = response_data
+ else:
+ content = f"Error: Unexpected type: {type(response_data)}"
+
+ # Apply thinking tag removal if requested
+ if hide_thinking and content:
+ content = _remove_thinking_tags(content)
if context is not None and isinstance(context, dict):
context_out = context.copy()
@@ -529,7 +550,7 @@ def _sanitize_params_for_log(d: dict) -> dict:
class LLMToolkitTextGeneratorStream:
DEFAULT_PROVIDER = "openai"
# For default OpenAI we will use GPT-4o mini (or 4o-mini) – alias may differ
- DEFAULT_MODEL: str = "gpt-4.1-mini"
+ DEFAULT_MODEL: str = "gpt-4o-mini"
MODEL_LIST: List[str] = [DEFAULT_MODEL]
@@ -538,7 +559,8 @@ class LLMToolkitTextGeneratorStream:
return {
"required": {
"llm_model": (cls.MODEL_LIST, {"default": cls.DEFAULT_MODEL}),
- "prompt": ("STRING", {"multiline": False, "default": "Write a detailed description of a futuristic city."})
+ "prompt": ("STRING", {"multiline": False, "default": "Write a detailed description of a futuristic city."}),
+ "hide_thinking": ("BOOLEAN", {"default": True, "tooltip": "Hide model thinking process (content between tags)"})
},
"optional": {
"context": ("*", {})
@@ -551,10 +573,10 @@ class LLMToolkitTextGeneratorStream:
RETURN_TYPES = ("*", "STRING")
RETURN_NAMES = ("context", "text")
FUNCTION = "generate_stream" # <-- Use new function name
- CATEGORY = "llm_toolkit"
+ CATEGORY = "llm_toolkit/generators"
OUTPUT_NODE = True # Keep the JS widget logic
- def generate_stream(self, llm_model, prompt, unique_id, context=None, **kwargs):
+ def generate_stream(self, llm_model, prompt, hide_thinking, unique_id, context=None, **kwargs):
"""
Generates text using the specified provider and streams the response back
to the UI via websocket messages. (synchronous wrapper)
@@ -571,6 +593,8 @@ class LLMToolkitTextGeneratorStream:
server = PromptServer.instance
full_response_text = ""
+ thinking_buffer = ""
+ inside_thinking = False
try:
# --- Parameter processing logic (same as non-streaming) ---
@@ -728,43 +752,84 @@ class LLMToolkitTextGeneratorStream:
async for chunk in stream_generator:
if chunk:
full_response_text += chunk
- # Send chunk to frontend via websocket
- server.send_sync(
- "llmtoolkit.stream.chunk",
- {"node": unique_id, "text": chunk},
- sid=server.client_id,
- )
+
+ # Handle thinking tag filtering for streaming
+ if hide_thinking:
+ # Track thinking state and buffer content
+ chunk_to_send = ""
+ i = 0
+ while i < len(chunk):
+ char = chunk[i]
+
+ if not inside_thinking:
+ # Check for start of thinking tag
+ if char == '<' and chunk[i:i+7] == '':
+ inside_thinking = True
+ thinking_buffer = ''
+ i += 7
+ continue
+ else:
+ chunk_to_send += char
+ else:
+ # Inside thinking block, buffer until we find closing tag
+ thinking_buffer += char
+ if thinking_buffer.endswith(''):
+ inside_thinking = False
+ thinking_buffer = ""
+ i += 1
+ continue
+ i += 1
+
+ # Only send non-thinking content
+ if chunk_to_send:
+ server.send_sync(
+ "llmtoolkit.stream.chunk",
+ {"node": unique_id, "text": chunk_to_send},
+ sid=server.client_id,
+ )
+ else:
+ # Send all content if not hiding thinking
+ server.send_sync(
+ "llmtoolkit.stream.chunk",
+ {"node": unique_id, "text": chunk},
+ sid=server.client_id,
+ )
await asyncio.sleep(0.001)
+ # Apply thinking tag removal to final text if requested
+ final_text = full_response_text
+ if hide_thinking:
+ final_text = _remove_thinking_tags(full_response_text)
+
logger.info(
- f"[Streaming] Finished for node {unique_id}. Total length: {len(full_response_text)}"
+ f"[Streaming] Finished for node {unique_id}. Total length: {len(final_text)}"
)
server.send_sync(
"llmtoolkit.stream.end",
- {"node": unique_id, "final_text": full_response_text},
+ {"node": unique_id, "final_text": final_text},
sid=server.client_id,
)
# --- Prepare final context output ---
if context is not None and isinstance(context, dict):
context_out = context.copy()
- context_out["llm_response"] = full_response_text
+ context_out["llm_response"] = final_text
context_out["llm_raw_response"] = {
"status": "Streamed successfully",
- "final_length": len(full_response_text),
+ "final_length": len(final_text),
}
else:
context_out = {
- "llm_response": full_response_text,
+ "llm_response": final_text,
"llm_raw_response": {
"status": "Streamed successfully",
- "final_length": len(full_response_text),
+ "final_length": len(final_text),
},
"passthrough_data": context,
}
- payload = ContextPayload(full_response_text, context_out)
- return {"ui": {"string": [full_response_text]}, "result": (payload, str(payload))}
+ payload = ContextPayload(final_text, context_out)
+ return {"ui": {"string": [final_text]}, "result": (payload, str(payload))}
except Exception as e:
error_message = f"Error during streaming generation: {str(e)}"
@@ -800,4 +865,4 @@ NODE_CLASS_MAPPINGS = {
NODE_DISPLAY_NAME_MAPPINGS = {
"LLMToolkitTextGenerator": "Generate Text (LLMToolkit)",
"LLMToolkitTextGeneratorStream": "Generate Text Stream (LLMToolkit)" # New display name
-}
+}
\ No newline at end of file
diff --git a/comfy-nodes/llmtoolkit_providers.py b/comfy-nodes/llmtoolkit_providers.py
index 9d3fa62..9665f2b 100644
--- a/comfy-nodes/llmtoolkit_providers.py
+++ b/comfy-nodes/llmtoolkit_providers.py
@@ -177,7 +177,7 @@ class LLMToolkitProviderSelector:
RETURN_NAMES = ("context",)
FUNCTION = "select_provider"
- CATEGORY = "llm_toolkit"
+ CATEGORY = "llm_toolkit/providers"
@classmethod
def IS_CHANGED(cls, llm_provider, llm_model, base_ip="localhost", port="11434", external_api_key="", context=None):
diff --git a/comfy-nodes/openai_provider.py b/comfy-nodes/openai_provider.py
index 38b6d61..6bee256 100644
--- a/comfy-nodes/openai_provider.py
+++ b/comfy-nodes/openai_provider.py
@@ -99,7 +99,7 @@ class OpenAIProviderNode:
RETURN_TYPES = ("*",)
RETURN_NAMES = ("context",)
FUNCTION = "configure_openai"
- CATEGORY = "llm_toolkit"
+ CATEGORY = "llm_toolkit/providers"
# This allows ComfyUI to invalidate the node if inputs that matter change.
@classmethod
diff --git a/comfy-nodes/prompt_manager.py b/comfy-nodes/prompt_manager.py
index 3e1744a..fc0d15a 100644
--- a/comfy-nodes/prompt_manager.py
+++ b/comfy-nodes/prompt_manager.py
@@ -70,7 +70,7 @@ class PromptManager:
RETURN_TYPES = ("*",)
RETURN_NAMES = ("context",)
FUNCTION = "manage_prompt"
- CATEGORY = "llm_toolkit"
+ CATEGORY = "llm_toolkit/utils"
def manage_prompt(self, context: Optional[Dict[str, Any]] = None, **kwargs) -> Tuple[Dict[str, Any]]:
"""
diff --git a/comfy-nodes/suno_provider.py b/comfy-nodes/suno_provider.py
new file mode 100644
index 0000000..2f08076
--- /dev/null
+++ b/comfy-nodes/suno_provider.py
@@ -0,0 +1,71 @@
+import logging
+from typing import Any, Dict, Tuple
+
+from llmtoolkit_utils import get_models, get_api_key
+
+logger = logging.getLogger(__name__)
+
+
+class SunoProviderSelector:
+ """Dedicated provider selector for Suno music API."""
+
+ SUNO_MODELS = ["V3_5", "V4", "V4_5"]
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "llm_model": (cls.SUNO_MODELS, {"default": "V3_5"}),
+ },
+ "optional": {
+ "external_api_key": (
+ "STRING",
+ {
+ "default": "",
+ "multiline": False,
+ "tooltip": "Optionally supply SUNO API key directly. Otherwise SUNO_API_KEY env var will be used.",
+ },
+ ),
+ "context": ("*", {}),
+ },
+ }
+
+ RETURN_TYPES = ("*",)
+ RETURN_NAMES = ("context",)
+ FUNCTION = "select_provider"
+ CATEGORY = "llm_toolkit/providers"
+
+ def select_provider(
+ self,
+ llm_model: str,
+ external_api_key: str = "",
+ context: Any = None,
+ ) -> Tuple[Any]:
+ api_key = external_api_key.strip() or ""
+ if not api_key:
+ try:
+ api_key = get_api_key("SUNO_API_KEY", "suno")
+ except ValueError:
+ api_key = ""
+
+ if not api_key:
+ logger.warning("SunoProviderSelector: API key not provided or found. Calls will fail until set.")
+
+ provider_config: Dict[str, Any] = {
+ "provider_name": "suno",
+ "llm_model": llm_model,
+ "api_key": api_key or "",
+ }
+
+ if context is not None:
+ if isinstance(context, dict):
+ context["provider_config"] = provider_config
+ return (context,)
+ else:
+ wrapped = {"provider_config": provider_config, "passthrough_data": context}
+ return (wrapped,)
+ return (provider_config,)
+
+
+NODE_CLASS_MAPPINGS = {"SunoProviderSelector": SunoProviderSelector}
+NODE_DISPLAY_NAME_MAPPINGS = {"SunoProviderSelector": "Suno Provider Selector (LLMToolkit)"}
\ No newline at end of file
diff --git a/comfy-nodes/switch_any.py b/comfy-nodes/switch_any.py
index 152862f..9c38427 100644
--- a/comfy-nodes/switch_any.py
+++ b/comfy-nodes/switch_any.py
@@ -18,7 +18,7 @@ class SwitchAny:
Useful for conditional logic in ComfyUI workflows.
"""
def __init__(self):
- self.type = "utility"
+ self.type = "llm_toolkit/utils"
@classmethod
def INPUT_TYPES(cls):
@@ -37,7 +37,7 @@ class SwitchAny:
RETURN_NAMES = ("output",)
FUNCTION = "switch"
OUTPUT_NODE = False # This is a utility node, not an output node
- CATEGORY = "llm_toolkit"
+ CATEGORY = "llm_toolkit/utils"
def switch(self, selector: bool, input_a: Any = None, input_b: Any = None) -> Tuple[Any]:
"""
diff --git a/comfy-nodes/transformers_provider.py b/comfy-nodes/transformers_provider.py
index f97d36f..d4d00a5 100644
--- a/comfy-nodes/transformers_provider.py
+++ b/comfy-nodes/transformers_provider.py
@@ -223,7 +223,7 @@ class LocalTransformersProviderNode:
RETURN_TYPES = ("*",)
RETURN_NAMES = ("context",)
FUNCTION = "configure_transformers"
- CATEGORY = "llm_toolkit"
+ CATEGORY = "llm_toolkit/providers"
@classmethod
def IS_CHANGED(cls, llm_model: str, context=None):
diff --git a/llmtoolkit_utils.py b/llmtoolkit_utils.py
index 763bdbd..9331b2d 100644
--- a/llmtoolkit_utils.py
+++ b/llmtoolkit_utils.py
@@ -1665,6 +1665,10 @@ def get_models(engine, base_ip, port, api_key):
# BFL currently exposes a single public model via the Kontekst endpoint.
return ["flux-kontext-max"]
+ elif engine == "suno":
+ # Suno offers three primary model versions
+ return ["V3_5", "V4", "V4_5"]
+
else:
print(f"Unsupported engine - {engine}")
return []
diff --git a/suno_api.py b/suno_api.py
new file mode 100644
index 0000000..b6358e2
--- /dev/null
+++ b/suno_api.py
@@ -0,0 +1,230 @@
+import asyncio
+import time
+import logging
+from typing import Dict, Any, Optional
+
+import aiohttp
+
+logger = logging.getLogger(__name__)
+
+BASE_URL = "https://api.sunoapi.org"
+
+
+async def _request(method: str, url: str, api_key: str, *, json_payload: Optional[Dict[str, Any]] = None,
+ params: Optional[Dict[str, str]] = None, timeout: int = 60) -> Dict[str, Any]:
+ """Internal helper for performing authenticated HTTP requests."""
+ headers = {
+ "Authorization": f"Bearer {api_key}",
+ "Accept": "application/json",
+ "Content-Type": "application/json",
+ }
+ async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=timeout)) as session:
+ async with session.request(method, url, headers=headers, json=json_payload, params=params) as response:
+ response.raise_for_status()
+ return await response.json()
+
+
+async def send_suno_music_generation_request(
+ *,
+ api_key: str,
+ prompt: str,
+ style: Optional[str] = None,
+ title: Optional[str] = None,
+ custom_mode: bool = False,
+ instrumental: bool = False,
+ model: str = "V3_5",
+ negative_tags: Optional[str] = None,
+ callback_url: str = "https://example.com/callback",
+ poll: bool = True,
+ poll_interval: float = 5.0,
+ max_poll_seconds: int = 300,
+) -> Dict[str, Any]:
+ """Create music with Suno API and optionally poll until complete.
+
+ Returns the final record-info payload (or the initial response if polling is disabled).
+ """
+ payload = {
+ "prompt": prompt,
+ "style": style,
+ "title": title,
+ "customMode": custom_mode,
+ "instrumental": instrumental,
+ "model": model,
+ "negativeTags": negative_tags,
+ "callBackUrl": callback_url,
+ }
+ # Remove None values that the API may reject
+ payload = {k: v for k, v in payload.items() if v is not None}
+
+ logger.info(f"Suno music request payload: {payload}")
+ initial_url = f"{BASE_URL}/api/v1/generate"
+ init_resp = await _request("POST", initial_url, api_key, json_payload=payload)
+ logger.info(f"Suno initial response: {init_resp}")
+
+ if not poll:
+ return init_resp
+
+ task_id = init_resp.get("data", {}).get("task_id") or init_resp.get("data", {}).get("taskId")
+ if not task_id:
+ logger.warning("No task_id returned by Suno, cannot poll.")
+ return init_resp
+
+ info_url = f"{BASE_URL}/api/v1/generate/record-info"
+ start_time = time.time()
+
+ while True:
+ await asyncio.sleep(poll_interval)
+ elapsed = time.time() - start_time
+ if elapsed > max_poll_seconds:
+ logger.warning("Polling for Suno music generation timed out.")
+ break
+ try:
+ info_resp = await _request("GET", info_url, api_key, params={"taskId": task_id})
+ except Exception as e:
+ logger.error(f"Error polling Suno record-info: {e}")
+ continue
+
+ status_code = info_resp.get("code", 0)
+ if status_code == 200:
+ status_data = info_resp.get("data", {})
+ callback_type = status_data.get("callbackType") or status_data.get("status")
+ if callback_type in {"SUCCESS", "complete", "FIRST_SUCCESS", "FIRST", "COMPLETE"}:
+ logger.info("Suno music generation completed.")
+ return info_resp
+ # Continue polling otherwise
+
+ return {"error": "Timeout waiting for Suno music generation", "initial_response": init_resp}
+
+
+async def send_suno_lyrics_generation_request(
+ *,
+ api_key: str,
+ prompt: str,
+ callback_url: str = "https://example.com/callback",
+ poll: bool = True,
+ poll_interval: float = 5.0,
+ max_poll_seconds: int = 180,
+) -> Dict[str, Any]:
+ """Generate lyrics with Suno API and optionally poll until complete."""
+ payload = {
+ "prompt": prompt,
+ "callBackUrl": callback_url,
+ }
+ lyrics_url = f"{BASE_URL}/api/v1/lyrics"
+ init_resp = await _request("POST", lyrics_url, api_key, json_payload=payload)
+
+ if not poll:
+ return init_resp
+
+ task_id = init_resp.get("data", {}).get("task_id") or init_resp.get("data", {}).get("taskId")
+ if not task_id:
+ logger.warning("No task_id returned by Suno lyrics endpoint, cannot poll.")
+ return init_resp
+
+ info_url = f"{BASE_URL}/api/v1/lyrics/record-info"
+ start_time = time.time()
+
+ while True:
+ await asyncio.sleep(poll_interval)
+ if time.time() - start_time > max_poll_seconds:
+ logger.warning("Polling for Suno lyrics generation timed out.")
+ break
+ try:
+ info_resp = await _request("GET", info_url, api_key, params={"taskId": task_id})
+ except Exception as e:
+ logger.error(f"Error polling Suno lyrics record-info: {e}")
+ continue
+
+ if info_resp.get("code") == 200 and info_resp.get("data", {}).get("status") == "SUCCESS":
+ return info_resp
+
+ return {"error": "Timeout waiting for Suno lyrics generation", "initial_response": init_resp}
+
+
+async def send_suno_upload_cover_request(
+ *,
+ api_key: str,
+ upload_url: str,
+ prompt: str,
+ style: Optional[str],
+ title: Optional[str],
+ custom_mode: bool,
+ instrumental: bool,
+ model: str,
+ negative_tags: Optional[str],
+ callback_url: str,
+ poll: bool = True,
+) -> Dict[str, Any]:
+ """Call Suno upload-cover endpoint."""
+ payload = {
+ "uploadUrl": upload_url,
+ "prompt": prompt,
+ "style": style,
+ "title": title,
+ "customMode": custom_mode,
+ "instrumental": instrumental,
+ "model": model,
+ "negativeTags": negative_tags,
+ "callBackUrl": callback_url,
+ }
+ payload = {k: v for k, v in payload.items() if v is not None}
+ resp = await _request("POST", f"{BASE_URL}/api/v1/generate/upload-cover", api_key, json_payload=payload)
+ if not poll:
+ return resp
+ # Re-use record-info polling
+ task_id = resp.get("data", {}).get("task_id")
+ if not task_id:
+ return resp
+ return await send_suno_music_generation_request(
+ api_key=api_key,
+ prompt=prompt,
+ custom_mode=custom_mode,
+ instrumental=instrumental,
+ model=model,
+ poll=True,
+ callback_url=callback_url,
+ )
+
+
+async def send_suno_upload_extend_request(
+ *,
+ api_key: str,
+ upload_url: str,
+ default_param_flag: bool,
+ prompt: Optional[str],
+ style: Optional[str],
+ title: Optional[str],
+ continue_at: Optional[int],
+ instrumental: bool,
+ model: str,
+ negative_tags: Optional[str],
+ callback_url: str,
+ poll: bool = True,
+) -> Dict[str, Any]:
+ payload = {
+ "uploadUrl": upload_url,
+ "defaultParamFlag": default_param_flag,
+ "instrumental": instrumental,
+ "prompt": prompt,
+ "style": style,
+ "title": title,
+ "continueAt": continue_at,
+ "model": model,
+ "negativeTags": negative_tags,
+ "callBackUrl": callback_url,
+ }
+ payload = {k: v for k, v in payload.items() if v is not None}
+ resp = await _request("POST", f"{BASE_URL}/api/v1/generate/upload-extend", api_key, json_payload=payload)
+ if not poll:
+ return resp
+ task_id = resp.get("data", {}).get("task_id")
+ if not task_id:
+ return resp
+ # Poll via record-info
+ return await send_suno_music_generation_request(
+ api_key=api_key,
+ prompt=prompt or "",
+ model=model,
+ poll=True,
+ callback_url=callback_url,
+ )
\ No newline at end of file