Add files via upload

Added SUNO Music generator and thinking model improvememts

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)",
 } 

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):

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)",
+} 

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]:
         """

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]]:
         """

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 <think>...</think> blocks from text, including the tags themselves."""
+    import re
+    # Pattern to match <think>...</think> blocks (including nested content and newlines)
+    pattern = r'<think>.*?</think>'
+    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 <think> 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.")

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]:

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)",
+} 

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 <think>...</think> blocks from text, including the tags themselves."""
+    import re
+    # Pattern to match <think>...</think> blocks (including nested content and newlines)
+    pattern = r'<think>.*?</think>'
+    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 <think> 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 <think> 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] == '<think>':
+                                        inside_thinking = True
+                                        thinking_buffer = '<think>'
+                                        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('</think>'):
+                                        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
-}
+}

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):

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

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]]:
         """

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)"} 

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]:
         """

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):

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 []

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,
+    )