From 764cb4d70cd388011ff46006c687fd7ab7fb85fe Mon Sep 17 00:00:00 2001 From: ImpactFrames Date: Wed, 9 Jul 2025 04:34:51 +0100 Subject: [PATCH] Add files via upload --- comfy-nodes/audio_duration_frames.py | 2 +- comfy-nodes/blank_image.py | 2 +- comfy-nodes/check_image_empty.py | 2 +- comfy-nodes/config_generate_image.py | 2 +- comfy-nodes/config_generate_music.py | 2 +- comfy-nodes/config_generate_speech.py | 93 +++++++++++++ comfy-nodes/config_generate_video.py | 83 +++++++++++ comfy-nodes/display_text.py | 2 +- comfy-nodes/gemini_provider.py | 166 ++++++++++++++++++++++ comfy-nodes/generate_image.py | 42 +++++- comfy-nodes/generate_speech.py | 188 +++++++++++++++++++++++++ comfy-nodes/generate_video.py | 189 ++++++++++++++++++++++++++ comfy-nodes/llmtoolkit_providers.py | 2 +- comfy-nodes/load_audio_from_path.py | 83 +++++++++++ comfy-nodes/load_video_from_path.py | 127 +++++++++++++++++ comfy-nodes/logic_preview_image.py | 132 ++++++++++++++++++ comfy-nodes/preview_outputs.py | 2 +- comfy-nodes/switch_any.py | 140 ++++++++++++++++--- 18 files changed, 1229 insertions(+), 30 deletions(-) create mode 100644 comfy-nodes/config_generate_speech.py create mode 100644 comfy-nodes/config_generate_video.py create mode 100644 comfy-nodes/gemini_provider.py create mode 100644 comfy-nodes/generate_speech.py create mode 100644 comfy-nodes/generate_video.py create mode 100644 comfy-nodes/load_audio_from_path.py create mode 100644 comfy-nodes/load_video_from_path.py create mode 100644 comfy-nodes/logic_preview_image.py diff --git a/comfy-nodes/audio_duration_frames.py b/comfy-nodes/audio_duration_frames.py index 6357ea5..4474602 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/utils" + CATEGORY = "llm_toolkit/utils/audio" # --------------------------------------------------------------------- # Core logic diff --git a/comfy-nodes/blank_image.py b/comfy-nodes/blank_image.py index 822008a..20720a9 100644 --- a/comfy-nodes/blank_image.py +++ b/comfy-nodes/blank_image.py @@ -25,7 +25,7 @@ class BlankImage: RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image",) FUNCTION = "create" - CATEGORY = "llm_toolkit/utils" + CATEGORY = "llm_toolkit/utils/image" # ------------------------------------------------------------------ # Core logic diff --git a/comfy-nodes/check_image_empty.py b/comfy-nodes/check_image_empty.py index 561e9c6..aca5f79 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/utils" + CATEGORY = "llm_toolkit/utils/logic" 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 2c6abbe..39ecab2 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/utils" + CATEGORY = "llm_toolkit/config" def configure(self, context: Optional[Dict[str, Any]] = None, **kwargs) -> Tuple[Dict[str, Any]]: """ diff --git a/comfy-nodes/config_generate_music.py b/comfy-nodes/config_generate_music.py index 3e56d54..7e0e526 100644 --- a/comfy-nodes/config_generate_music.py +++ b/comfy-nodes/config_generate_music.py @@ -177,7 +177,7 @@ class ConfigGenerateMusic: RETURN_TYPES = ("*",) RETURN_NAMES = ("context",) FUNCTION = "configure" - CATEGORY = "llm_toolkit/utils" + CATEGORY = "llm_toolkit/config" # ------------------------------------------------------------------ # Execution diff --git a/comfy-nodes/config_generate_speech.py b/comfy-nodes/config_generate_speech.py new file mode 100644 index 0000000..bbd3520 --- /dev/null +++ b/comfy-nodes/config_generate_speech.py @@ -0,0 +1,93 @@ +# comfy-nodes/config_generate_speech.py +"""ComfyUI helper node – prepares generation_config for Gemini TTS. + +The node mirrors *config_generate_image.py* but specialises on the +parameters needed by the Gemini 2.5 TTS models. +""" + +import os +import sys +import logging +from typing import Any, Dict, Optional, Tuple + +from context_payload import extract_context + +# Ensure parent dir in path +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) + +logger = logging.getLogger(__name__) + +# Pre-built voice names published by Google (see docs) +VOICE_OPTIONS = [ + "Zephyr", "Puck", "Charon", "Kore", "Fenrir", "Leda", "Orus", "Aoede", + "Callirrhoe", "Autonoe", "Enceladus", "Iapetus", "Umbriel", "Algieba", + "Despina", "Erinome", "Algenib", "Rasalgethi", "Laomedeia", "Achernar", + "Alnilam", "Schedar", "Gacrux", "Pulcherrima", "Achird", "Zubenelgenubi", + "Vindemiatrix", "Sadachbia", "Sadaltager", "Sulafat", +] + + +class ConfigGenerateSpeech: + """Adds text-to-speech parameters to the context for downstream nodes.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": {}, + "optional": { + "context": ("*", {}), + "voice_name": (VOICE_OPTIONS, {"default": "Kore"}), + # Sample rate & channels kept here for future flexibility + "sample_rate": ("INT", {"default": 24000, "min": 8000, "max": 48000}), + "channels": ("INT", {"default": 1, "min": 1, "max": 2}), + }, + } + + RETURN_TYPES = ("*",) + RETURN_NAMES = ("context",) + FUNCTION = "configure" + CATEGORY = "llm_toolkit/config" + + def configure( + self, + context: Optional[Dict[str, Any]] = None, + voice_name: str = "Kore", + sample_rate: int = 24000, + channels: int = 1, + ) -> Tuple[Dict[str, Any]]: + logger.info("ConfigGenerateSpeech executing…") + + # Unwrap / copy context + if context is None: + output_context: Dict[str, Any] = {} + elif isinstance(context, dict): + output_context = context.copy() + else: + output_context = extract_context(context) + if not isinstance(output_context, dict): + output_context = {"passthrough_data": context} + + gen_cfg = output_context.get("generation_config", {}) + if not isinstance(gen_cfg, dict): + gen_cfg = {} + + gen_cfg.update( + { + "voice_name": voice_name, + "sample_rate": sample_rate, + "channels": channels, + } + ) + + output_context["generation_config"] = gen_cfg + logger.info("ConfigGenerateSpeech: Updated generation_config with voice %s", voice_name) + return (output_context,) + + +NODE_CLASS_MAPPINGS = {"ConfigGenerateSpeech": ConfigGenerateSpeech} +NODE_DISPLAY_NAME_MAPPINGS = { + "ConfigGenerateSpeech": "Configure Speech Generation (LLMToolkit)", +} \ No newline at end of file diff --git a/comfy-nodes/config_generate_video.py b/comfy-nodes/config_generate_video.py new file mode 100644 index 0000000..66bdb77 --- /dev/null +++ b/comfy-nodes/config_generate_video.py @@ -0,0 +1,83 @@ +# comfy-nodes/config_generate_video.py +"""Configure Video Generation (Veo) – adds GenerateVideosConfig values to context.""" + +import os +import sys +import logging +from typing import Any, Dict, Optional, Tuple + +from context_payload import extract_context + +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) + +logger = logging.getLogger(__name__) + +ASPECT_OPTIONS = ["16:9", "9:16"] +PERSON_OPTIONS = ["dont_allow", "allow_adult"] # Removed "allow_all" - not supported by Veo API + +class ConfigGenerateVideo: + @classmethod + def INPUT_TYPES(cls): + return { + "required": {}, + "optional": { + "context": ("*", {}), + "aspect_ratio": (ASPECT_OPTIONS, {"default": "16:9"}), + "person_generation": (PERSON_OPTIONS, {"default": "dont_allow"}), + "number_of_videos": ("INT", {"default": 1, "min": 1, "max": 2}), + "duration_seconds": ("INT", {"default": 6, "min": 5, "max": 8}), + "negative_prompt": ("STRING", {"default": "", "multiline": True}), + "enhance_prompt": ("BOOLEAN", {"default": True}), + }, + } + + RETURN_TYPES = ("*",) + RETURN_NAMES = ("context",) + FUNCTION = "configure" + CATEGORY = "llm_toolkit/config" + + def configure( + self, + context: Optional[Dict[str, Any]] = None, + aspect_ratio: str = "16:9", + person_generation: str = "dont_allow", + number_of_videos: int = 1, + duration_seconds: int = 6, + negative_prompt: str = "", + enhance_prompt: bool = True, + ) -> Tuple[Dict[str, Any]]: + logger.info("ConfigGenerateVideo executing…") + + if context is None: + out_ctx: Dict[str, Any] = {} + elif isinstance(context, dict): + out_ctx = context.copy() + else: + out_ctx = extract_context(context) + if not isinstance(out_ctx, dict): + out_ctx = {"passthrough_data": context} + + gen_cfg = out_ctx.get("generation_config", {}) + if not isinstance(gen_cfg, dict): + gen_cfg = {} + + gen_cfg.update( + { + "aspect_ratio": aspect_ratio, + "person_generation": person_generation, + "number_of_videos": number_of_videos, + "duration_seconds": duration_seconds, + "negative_prompt": negative_prompt.strip(), + "enhance_prompt": enhance_prompt, + } + ) + + out_ctx["generation_config"] = gen_cfg + logger.info("ConfigGenerateVideo: saved config %s", gen_cfg) + return (out_ctx,) + +NODE_CLASS_MAPPINGS = {"ConfigGenerateVideo": ConfigGenerateVideo} +NODE_DISPLAY_NAME_MAPPINGS = {"ConfigGenerateVideo": "Configure Video Generation (LLMToolkit)"} \ No newline at end of file diff --git a/comfy-nodes/display_text.py b/comfy-nodes/display_text.py index c965e1d..f8fc353 100644 --- a/comfy-nodes/display_text.py +++ b/comfy-nodes/display_text.py @@ -55,7 +55,7 @@ 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/utils" # Changed category to llm_toolkit + CATEGORY = "llm_toolkit/utils/text" # Changed category to llm_toolkit def display_llm_text(self, context: Any, select: str, hide_thinking: bool): # --- Safe conversion for 'select' input string --- diff --git a/comfy-nodes/gemini_provider.py b/comfy-nodes/gemini_provider.py new file mode 100644 index 0000000..f0e496a --- /dev/null +++ b/comfy-nodes/gemini_provider.py @@ -0,0 +1,166 @@ +# comfy-nodes/gemini_provider.py +import os +import sys +import logging +from typing import Any, Dict, List, Optional, Tuple + +# Ensure project root utilities are importable +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) + +try: + from llmtoolkit_utils import ( + get_api_key, + get_models, + validate_gemini_key, + ) +except ImportError as e: + raise ImportError( + "GeminiProviderNode: failed to import helpers from llmtoolkit_utils – " + + str(e) + ) + +logger = logging.getLogger(__name__) + + +class GeminiProviderNode: + """Dedicated ComfyUI node that prepares *provider_config* for Gemini. + + The node fetches available Gemini models (if an API key is available), + exposes them as a dropdown, and returns a context dictionary containing a + *provider_config* block so that downstream generator nodes can send the + request without having to know any Gemini specifics. + """ + + # Fallback list used when live fetching fails or no API key is present. + _DEFAULT_MODELS: List[str] = [ + "gemini-2.5-pro", + "gemini-2.5-flash", + "gemini-2.0-pro", + "gemini-2.0-flash", + "gemini-pro", + "gemini-pro-vision", + "imagen-3-light-alpha", + ] + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + @classmethod + def _fetch_models(cls, api_key: Optional[str]) -> List[str]: + """Return list of model names, falling back to defaults on error.""" + try: + models = get_models("gemini", base_ip="localhost", port="0", api_key=api_key) + if models and isinstance(models, list): + return models + except Exception as e: + logger.warning(f"GeminiProviderNode: model fetch failed – {e}") + return cls._DEFAULT_MODELS + + # ------------------------------------------------------------------ + # ComfyUI node interface + # ------------------------------------------------------------------ + @classmethod + def INPUT_TYPES(cls): + # Attempt to fetch models using env key at import time so the dropdown + # has a useful list even before the user enters a key in the UI. + api_key_env = os.getenv("GEMINI_API_KEY", "").strip() + model_list = cls._fetch_models(api_key_env) or cls._DEFAULT_MODELS + default_model = model_list[0] if model_list else "gemini-pro" + return { + "required": { + "llm_model": (model_list, {"default": default_model}), + }, + "optional": { + "external_api_key": ( + "STRING", + { + "default": "", + "multiline": False, + "tooltip": "Gemini API key (optional – overrides env/.env)", + }, + ), + "context": ("*", {}), + }, + } + + RETURN_TYPES = ("*",) + RETURN_NAMES = ("context",) + FUNCTION = "select_gemini" + CATEGORY = "llm_toolkit/providers" + + # ------------------------------------------------------------------ + # IS_CHANGED → forces UI to refresh dropdown when API key changes + # ------------------------------------------------------------------ + @classmethod + def IS_CHANGED(cls, llm_model, external_api_key="", context=None): + import hashlib + + key_hash = hashlib.md5(external_api_key.encode()).hexdigest() if external_api_key else "nokey" + state = f"{llm_model}-{key_hash}" + return hashlib.md5(state.encode()).hexdigest() + + # ------------------------------------------------------------------ + # Main execution + # ------------------------------------------------------------------ + def select_gemini( + self, + llm_model: str, + external_api_key: str = "", + context: Optional[Any] = None, + ) -> Tuple[Any]: + """Validate key, build *provider_config*, and merge it into context.""" + + external_api_key = external_api_key.strip() + final_api_key = "" + + # 1) Validate external key if supplied + if external_api_key: + if validate_gemini_key(external_api_key): + final_api_key = external_api_key + logger.info("GeminiProviderNode: using valid external API key") + else: + logger.warning( + "GeminiProviderNode: provided external API key failed validation" + ) + + # 2) Fallback to env or .env file + if not final_api_key: + try: + final_api_key = get_api_key("GEMINI_API_KEY", "gemini") + except ValueError: + final_api_key = "" + + # 3) Allow dummy key for local/offline testing so other providers stay unaffected + if not final_api_key: + final_api_key = "1234" + + provider_config: Dict[str, Any] = { + "provider_name": "gemini", + "llm_model": llm_model, + "api_key": final_api_key, + "base_ip": None, # Cloud provider – no IP/port + "port": None, + } + + # Merge with incoming context if present + if context is not None: + if isinstance(context, dict): + merged = context.copy() + merged["provider_config"] = provider_config + result: Any = merged + else: + result = {"provider_config": provider_config, "passthrough_data": context} + else: + result = provider_config + + return (result,) + + +# ------------------------------------------------------------------ +# Node registration – discovered automatically by ComfyUI +# ------------------------------------------------------------------ +NODE_CLASS_MAPPINGS = {"GeminiProviderNode": GeminiProviderNode} +NODE_DISPLAY_NAME_MAPPINGS = {"GeminiProviderNode": "Gemini Provider Selector (LLMToolkit)"} \ No newline at end of file diff --git a/comfy-nodes/generate_image.py b/comfy-nodes/generate_image.py index ce989cf..21cd117 100644 --- a/comfy-nodes/generate_image.py +++ b/comfy-nodes/generate_image.py @@ -13,7 +13,8 @@ if parent_dir not in sys.path: # Relative imports try: from llmtoolkit_utils import tensor_to_base64, process_images_for_comfy, TENSOR_SUPPORT, get_api_key - from send_request import run_async # Assuming send_request.py exists for run_async + # Import helpers from send_request + from send_request import run_async, send_request # Added send_request for Gemini image generation # Import the specific API call function from root directory from openai_api import send_openai_image_generation_request except ImportError: @@ -294,6 +295,45 @@ listen to the heartbeat of a baby otter. error_message = "API response did not contain expected image data." logger.error(f"{error_message} Response: {raw_api_response}") + elif llm_provider in {"gemini", "google"}: + # ------------------------------------------------------------------ + # Gemini / Imagen – image generation via OpenAI-compat endpoint + # ------------------------------------------------------------------ + + if edit_mode or variation_mode: + error_message = ( + "Gemini image editing / variation not yet supported by GenerateImage node. " + "Please use 'generate' mode." + ) + logger.error(error_message) + else: + # Batch generation handled via batch_count (n) + try: + raw_api_response = run_async( + send_request( + llm_provider=llm_provider, + base_ip="localhost", + port="0", + images=None, # Gemini generation does not require prior image + llm_model=llm_model, + system_message="You are a helpful assistant.", + user_message=prompt_text, + messages=[], + llm_api_key=api_key, + batch_count=n, + strategy="create", + ) + ) + + if raw_api_response and raw_api_response.get("data"): + output_image_tensor, _ = process_images_for_comfy(raw_api_response) + else: + error_message = "Gemini API response did not contain expected image data." + logger.error("%s Response: %s", error_message, raw_api_response) + except Exception as g_exc: + error_message = f"Gemini image generation error: {g_exc}" + logger.error(error_message, exc_info=True) + elif llm_provider == "bfl": # ------------------------------------------------------------------ # BFL (Flux Kontext MAX) provider diff --git a/comfy-nodes/generate_speech.py b/comfy-nodes/generate_speech.py new file mode 100644 index 0000000..ebd255a --- /dev/null +++ b/comfy-nodes/generate_speech.py @@ -0,0 +1,188 @@ +# comfy-nodes/generate_speech.py +"""Generate speech (text-to-speech) using Gemini 2.5 TTS models. + +The node consumes *provider_config* and *generation_config* from the +context (populated e.g. via GeminiProviderNode & ConfigGenerateSpeech) and +saves a WAV file with the audio. +""" + +import os +import sys +import logging +import wave +from pathlib import Path +from typing import Any, Dict, Optional, Tuple + +from context_payload import extract_context + +# Add parent dir to path for utils +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) + +try: + from llmtoolkit_utils import get_api_key +except ImportError: + get_api_key = None # type: ignore + +try: + import folder_paths +except ImportError: + folder_paths = None # type: ignore + +logger = logging.getLogger(__name__) + + +class GenerateSpeech: + """Node that converts text to speech with Gemini TTS models.""" + + DEFAULT_MODEL = "gemini-2.5-flash-preview-tts" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"multiline": True, "default": "Hello world!"}), + }, + "optional": { + "context": ("*", {}), + }, + } + + RETURN_TYPES = ("*", "STRING") + RETURN_NAMES = ("context", "audio_path") + FUNCTION = "generate" + CATEGORY = "llm_toolkit/generators" + + def _save_wav(self, pcm_bytes: bytes, file_path: Path, channels: int = 1, rate: int = 24000, sample_width: int = 2): + file_path.parent.mkdir(parents=True, exist_ok=True) + with wave.open(str(file_path), "wb") as wf: + wf.setnchannels(channels) + wf.setsampwidth(sample_width) + wf.setframerate(rate) + wf.writeframes(pcm_bytes) + + # ------------------------------------------------------------------ + # Main node function + # ------------------------------------------------------------------ + def generate(self, prompt: str, context: Optional[Any] = None): + logger.info("GenerateSpeech node executing…") + + # -------------------- Unwrap context --------------------------- + if context is None: + context_dict: Dict[str, Any] = {} + elif isinstance(context, dict): + context_dict = context.copy() + else: + context_dict = extract_context(context) + if not isinstance(context_dict, dict): + context_dict = {"passthrough_data": context} + + provider_cfg = context_dict.get("provider_config", {}) + gen_cfg = context_dict.get("generation_config", {}) + + llm_provider = provider_cfg.get("provider_name", "gemini").lower() + if llm_provider not in {"gemini", "google"}: + err = f"GenerateSpeech currently supports Gemini only, received provider '{llm_provider}'." + logger.error(err) + context_dict["error"] = err + return (context_dict, "") + + llm_model = provider_cfg.get("llm_model", self.DEFAULT_MODEL) + if not llm_model: + llm_model = self.DEFAULT_MODEL + + # -------------------- API key resolution ---------------------- + api_key = provider_cfg.get("api_key", "").strip() + if (not api_key or api_key == "1234") and get_api_key is not None: + try: + api_key = get_api_key("GEMINI_API_KEY", "gemini") + logger.info("GenerateSpeech: fetched API key from environment/.env") + except ValueError: + pass + if not api_key: + err = "GenerateSpeech: missing Gemini API key." + logger.error(err) + context_dict["error"] = err + return (context_dict, "") + + # -------------------- Import Google GenAI --------------------- + try: + from google import genai + from google.genai import types + except ImportError: + err = "google-generativeai package not installed. Run 'pip install google-generativeai'." + logger.error(err) + context_dict["error"] = err + return (context_dict, "") + + # -------------------- Build configs --------------------------- + voice_name = gen_cfg.get("voice_name", "Kore") + sample_rate = int(gen_cfg.get("sample_rate", 24000)) + channels = int(gen_cfg.get("channels", 1)) + + speech_config = types.SpeechConfig( + voice_config=types.VoiceConfig( + prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name=voice_name) + ) + ) + + generation_config = types.GenerateContentConfig( + response_modalities=["AUDIO"], + speech_config=speech_config, + ) + + # -------------------- Call API ------------------------------- + try: + client = genai.Client(api_key=api_key) + response = client.models.generate_content( + model=llm_model, + contents=prompt, + config=generation_config, + ) + except Exception as api_exc: + err = f"Gemini TTS API error: {api_exc}" + logger.error(err, exc_info=True) + context_dict["error"] = err + return (context_dict, "") + + # Extract PCM bytes + try: + audio_bytes: bytes = response.candidates[0].content.parts[0].inline_data.data # type: ignore + except Exception as e: + err = f"Unexpected TTS response format: {e}" + logger.error(err, exc_info=True) + context_dict["error"] = err + return (context_dict, "") + + # Save WAV + # Use ComfyUI's standard output directory + if folder_paths and hasattr(folder_paths, 'get_output_directory'): + base_output_dir = folder_paths.get_output_directory() + out_dir = Path(base_output_dir) / "tts_audio" + else: + # Fallback to current directory if folder_paths not available + out_dir = Path(os.getcwd()) / "output" / "tts_audio" + + out_dir.mkdir(parents=True, exist_ok=True) + logger.info(f"Saving audio to: {out_dir}") + + filename = f"gemini_tts_{voice_name.lower()}_{os.getpid()}.wav" + file_path = out_dir / filename + try: + self._save_wav(audio_bytes, file_path, channels=channels, rate=sample_rate) + except Exception as save_exc: + err = f"Failed saving WAV: {save_exc}" + logger.error(err, exc_info=True) + context_dict["error"] = err + return (context_dict, "") + + logger.info("GenerateSpeech: saved audio to %s", file_path) + context_dict["generated_audio_path"] = str(file_path) + + return (context_dict, str(file_path)) + + +NODE_CLASS_MAPPINGS = {"GenerateSpeech": GenerateSpeech} +NODE_DISPLAY_NAME_MAPPINGS = {"GenerateSpeech": "Generate Speech (LLMToolkit)"} \ No newline at end of file diff --git a/comfy-nodes/generate_video.py b/comfy-nodes/generate_video.py new file mode 100644 index 0000000..9ea3d92 --- /dev/null +++ b/comfy-nodes/generate_video.py @@ -0,0 +1,189 @@ +# comfy-nodes/generate_video.py +"""Generate Video using Veo 2 via Gemini API. + +The node supports text-to-video for now. It polls the long-running +operation until completion and downloads the resulting MP4(s). +""" + +import os +import sys +import time +import logging +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +from context_payload import extract_context + +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) + +try: + from llmtoolkit_utils import get_api_key +except ImportError: + get_api_key = None # type: ignore + +try: + import folder_paths +except ImportError: + folder_paths = None # type: ignore + +logger = logging.getLogger(__name__) + + +class GenerateVideo: + DEFAULT_MODEL = "veo-2.0-generate-001" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"multiline": True, "default": "Cinematic shot of sunrise over mountains"}), + }, + "optional": { + "context": ("*", {}), + }, + } + + RETURN_TYPES = ("*", "STRING") + RETURN_NAMES = ("context", "video_paths") + FUNCTION = "generate" + CATEGORY = "llm_toolkit/generators" + + # ------------------------------------------------------------------ + def _save_videos(self, client, videos: List[Any], out_dir: Path, base_name: str) -> List[str]: + paths: List[str] = [] + out_dir.mkdir(parents=True, exist_ok=True) + for idx, vid in enumerate(videos): + fname = f"{base_name}_{idx}.mp4" + tgt = out_dir / fname + try: + client.files.download(file=vid.video) # ensures local cache + vid.video.save(str(tgt)) + paths.append(str(tgt)) + except Exception as dl_exc: + logger.error("Failed saving video %s – %s", fname, dl_exc, exc_info=True) + return paths + + # ------------------------------------------------------------------ + def generate(self, prompt: str, context: Optional[Any] = None): + logger.info("GenerateVideo node executing…") + + # ------------ Context handling ------------- + if context is None: + ctx: Dict[str, Any] = {} + elif isinstance(context, dict): + ctx = context.copy() + else: + ctx = extract_context(context) + if not isinstance(ctx, dict): + ctx = {"passthrough_data": context} + + provider_cfg = ctx.get("provider_config", {}) + gen_cfg = ctx.get("generation_config", {}) + + # Default to gemini provider if none + llm_provider = provider_cfg.get("provider_name", "gemini").lower() + if llm_provider not in {"gemini", "google"}: + err = f"GenerateVideo supports Gemini/Veo only, provider '{llm_provider}' unsupported." + logger.error(err) + ctx["error"] = err + return (ctx, "") + + llm_model = provider_cfg.get("llm_model", self.DEFAULT_MODEL) or self.DEFAULT_MODEL + + # API key + api_key = provider_cfg.get("api_key", "").strip() + if (not api_key or api_key == "1234") and get_api_key is not None: + try: + api_key = get_api_key("GEMINI_API_KEY", "gemini") + except ValueError: + pass + if not api_key: + err = "GenerateVideo: missing Gemini API key" + logger.error(err) + ctx["error"] = err + return (ctx, "") + + # Import SDK + try: + from google import genai + from google.genai import types + except ImportError: + err = "google-generativeai not installed – run 'pip install google-generativeai'" + logger.error(err) + ctx["error"] = err + return (ctx, "") + + aspect_ratio = gen_cfg.get("aspect_ratio", "16:9") + person_generation = gen_cfg.get("person_generation", "dont_allow") + + # Map unsupported "allow_all" to "allow_adult" for backward compatibility + if person_generation == "allow_all": + logger.warning("person_generation 'allow_all' is not supported by Veo API, using 'allow_adult' instead") + person_generation = "allow_adult" + + number_of_videos = int(gen_cfg.get("number_of_videos", 1)) + duration_seconds = int(gen_cfg.get("duration_seconds", 6)) + negative_prompt = gen_cfg.get("negative_prompt", "") + enhance_prompt = bool(gen_cfg.get("enhance_prompt", True)) + + video_cfg = types.GenerateVideosConfig( + aspect_ratio=aspect_ratio, + person_generation=person_generation, + number_of_videos=number_of_videos, + duration_seconds=duration_seconds, + negative_prompt=negative_prompt if negative_prompt else None, + enhance_prompt=enhance_prompt, + ) + + try: + client = genai.Client(api_key=api_key) + operation = client.models.generate_videos( + model=llm_model, + prompt=prompt, + config=video_cfg, + ) + + # Poll + poll_sec = 20 + while not operation.done: + logger.info("GenerateVideo: operation not done yet – sleeping %s sec", poll_sec) + time.sleep(poll_sec) + operation = client.operations.get(operation) + + videos = operation.response.generated_videos # type: ignore + if not videos: + raise RuntimeError("Veo generation returned no videos") + + # Use ComfyUI's standard output directory + if folder_paths and hasattr(folder_paths, 'get_output_directory'): + base_output_dir = folder_paths.get_output_directory() + out_dir = Path(base_output_dir) / "veo_videos" + else: + # Fallback to current directory if folder_paths not available + out_dir = Path(os.getcwd()) / "output" / "veo_videos" + + out_dir.mkdir(parents=True, exist_ok=True) + logger.info(f"Saving videos to: {out_dir}") + + base_name = f"veo_{os.getpid()}" + paths = self._save_videos(client, videos, out_dir, base_name) + if not paths: + raise RuntimeError("Failed saving any generated videos") + + ctx["generated_video_paths"] = paths + logger.info("GenerateVideo: saved %d video(s)", len(paths)) + # Return first path for convenience + return (ctx, paths[0] if len(paths) == 1 else "|".join(paths)) + + except Exception as exc: + err = f"Veo API error: {exc}" + logger.error(err, exc_info=True) + ctx["error"] = err + return (ctx, "") + + +NODE_CLASS_MAPPINGS = {"GenerateVideo": GenerateVideo} +NODE_DISPLAY_NAME_MAPPINGS = {"GenerateVideo": "Generate Video (LLMToolkit)"} \ No newline at end of file diff --git a/comfy-nodes/llmtoolkit_providers.py b/comfy-nodes/llmtoolkit_providers.py index 9665f2b..a9ba6a8 100644 --- a/comfy-nodes/llmtoolkit_providers.py +++ b/comfy-nodes/llmtoolkit_providers.py @@ -152,7 +152,7 @@ class LLMToolkitProviderSelector: SUPPORTED_PROVIDERS = [ "transformers", "llamacpp", "ollama", "kobold", "lmstudio", "textgen", "groq", "gemini", "openai", "bfl", "anthropic", "mistral", "deepseek", "xai", - "huggingface", "vllm" + "huggingface", "vllm", "google" ] REQUIRES_IP_PORT = ["ollama", "llamacpp", "kobold", "lmstudio", "textgen", "vllm"] diff --git a/comfy-nodes/load_audio_from_path.py b/comfy-nodes/load_audio_from_path.py new file mode 100644 index 0000000..f2c3c63 --- /dev/null +++ b/comfy-nodes/load_audio_from_path.py @@ -0,0 +1,83 @@ +# comfy-nodes/load_audio_from_path.py +"""Load Audio From Path - Bridge node to load audio files into ComfyUI format.""" + +import os +import torch +import logging +from typing import Tuple, Dict, Any + +# Try to import torchaudio for audio loading +try: + import torchaudio +except ImportError: + print("Warning: torchaudio is not installed. The LoadAudioFromPath node will not work.") + print("Please run: pip install torchaudio") + torchaudio = None + +logger = logging.getLogger(__name__) + + +class LoadAudioFromPath: + """Load an audio file from a file path string and convert to ComfyUI audio format.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + # This input will connect to the 'audio_path' output of GenerateSpeech + "audio_path": ("STRING", {"multiline": False, "default": ""}), + } + } + + # We output the special AUDIO type that other audio nodes understand + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + FUNCTION = "load_audio" + CATEGORY = "llm_toolkit/utils/audio" # Consistent with LoadVideoFromPath + + def load_audio(self, audio_path: str) -> Tuple[Dict[str, Any]]: + """Load audio file and return as ComfyUI audio format.""" + + if not torchaudio: + raise ImportError("torchaudio is not installed. Cannot run LoadAudioFromPath.") + + # Validate the input path + if not audio_path or not os.path.exists(audio_path): + logger.error(f"LoadAudioFromPath: Audio path is empty or file does not exist: {audio_path}") + # Return a silent, empty audio object to prevent crashing the workflow + return ({"waveform": torch.zeros((1, 1, 1)), "sample_rate": 44100},) + + try: + # Load the audio file using torchaudio + # This is the most reliable way to get the waveform and sample rate + waveform, sample_rate = torchaudio.load(audio_path) + + # Ensure the waveform tensor has a batch dimension for ComfyUI compatibility + # Standard audio nodes expect [batch_size, num_channels, num_samples] + if waveform.dim() == 2: + waveform = waveform.unsqueeze(0) # Adds the batch dimension + + # Create the AUDIO dictionary object + audio_output = { + "waveform": waveform, + "sample_rate": sample_rate + } + + logger.info(f"Loaded audio from {audio_path} with shape {waveform.shape} and sample rate {sample_rate}") + + return (audio_output,) + + except Exception as e: + logger.error(f"Error loading audio from path '{audio_path}': {e}", exc_info=True) + # Return a silent, empty audio object on failure + return ({"waveform": torch.zeros((1, 1, 1)), "sample_rate": 44100},) + + +# Add the new node to ComfyUI's list of recognized nodes +NODE_CLASS_MAPPINGS = { + "LoadAudioFromPath": LoadAudioFromPath +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LoadAudioFromPath": "Load Audio From Path (LLMToolkit)" +} \ No newline at end of file diff --git a/comfy-nodes/load_video_from_path.py b/comfy-nodes/load_video_from_path.py new file mode 100644 index 0000000..73dc7a0 --- /dev/null +++ b/comfy-nodes/load_video_from_path.py @@ -0,0 +1,127 @@ +# comfy-nodes/load_video_from_path.py +"""Load Video From Path - Bridge node to load video files into ComfyUI format.""" + +import os +import torch +import numpy as np +import logging +from typing import Tuple + +# A popular library for video processing +try: + import cv2 +except ImportError: + print("Warning: opencv-python is not installed. The LoadVideoFromPath node will not work.") + print("Please run: pip install opencv-python") + cv2 = None + +logger = logging.getLogger(__name__) + + +class LoadVideoFromPath: + """Load a video file from a file path string and convert to ComfyUI image format.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + # The video_paths string from your GenerateVideo node + "video_path": ("STRING", {"multiline": False, "default": ""}), + }, + "optional": { + "max_frames": ("INT", {"default": 0, "min": 0, "max": 10000, + "tooltip": "Maximum frames to load (0 = all frames)"}), + "skip_frames": ("INT", {"default": 0, "min": 0, "max": 100, + "tooltip": "Load every Nth frame (1 = all, 2 = every other, etc.)"}), + } + } + + RETURN_TYPES = ("IMAGE", "INT", "FLOAT", "INT", "INT") + RETURN_NAMES = ("images", "frame_count", "fps", "width", "height") + FUNCTION = "load_video" + CATEGORY = "llm_toolkit/utils/video" + + def load_video(self, video_path: str, max_frames: int = 0, skip_frames: int = 0) -> Tuple: + """Load video and return frames as ComfyUI image tensor.""" + + if not cv2: + raise ImportError("opencv-python is not installed. Cannot run LoadVideoFromPath.") + + # Handle multiple paths separated by | + if "|" in video_path: + video_path = video_path.split("|")[0].strip() + logger.info(f"Multiple paths provided, using first: {video_path}") + + if not video_path or not os.path.exists(video_path): + logger.error(f"LoadVideoFromPath: Video path is empty or file does not exist: {video_path}") + # Return empty tensors to prevent crashing the workflow + return (torch.zeros((1, 64, 64, 3)), 0, 0.0, 64, 64) + + try: + # Open the video file + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise IOError(f"Cannot open video file: {video_path}") + + # Get video properties + fps = cap.get(cv2.CAP_PROP_FPS) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + logger.info(f"Video info: {width}x{height} @ {fps}fps, {total_frames} total frames") + + frames = [] + frame_idx = 0 + frames_loaded = 0 + skip_interval = max(1, skip_frames) + + while True: + # Read one frame + ret, frame = cap.read() + if not ret: + break # End of video + + # Apply frame skipping + if frame_idx % skip_interval == 0: + # OpenCV loads images in BGR format, so we convert to RGB + frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + + # Convert the numpy array frame to a torch tensor + # and normalize from [0, 255] to [0.0, 1.0] + frame_tensor = torch.from_numpy(frame_rgb).float() / 255.0 + frames.append(frame_tensor) + frames_loaded += 1 + + # Check if we've loaded enough frames + if max_frames > 0 and frames_loaded >= max_frames: + break + + frame_idx += 1 + + cap.release() + + if not frames: + raise ValueError("No frames could be read from the video.") + + # Stack all frame tensors into a single tensor (batch of images) + video_tensor = torch.stack(frames) + frame_count = len(frames) + + logger.info(f"Loaded {frame_count} frames from {video_path}") + + return (video_tensor, frame_count, fps, width, height) + + except Exception as e: + logger.error(f"Error loading video from path '{video_path}': {e}", exc_info=True) + return (torch.zeros((1, 64, 64, 3)), 0, 0.0, 64, 64) + + +# Add the new node to ComfyUI's list of recognized nodes +NODE_CLASS_MAPPINGS = { + "LoadVideoFromPath": LoadVideoFromPath +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LoadVideoFromPath": "Load Video From Path (LLMToolkit)" +} \ No newline at end of file diff --git a/comfy-nodes/logic_preview_image.py b/comfy-nodes/logic_preview_image.py new file mode 100644 index 0000000..c710598 --- /dev/null +++ b/comfy-nodes/logic_preview_image.py @@ -0,0 +1,132 @@ +import os +import logging +from typing import Tuple, Optional, List + +import torch +import numpy as np + +# ComfyUI utility for temp/output directories (if available) +try: + import folder_paths # type: ignore +except ImportError: + folder_paths = None # Fallback: will use working dir + +logger = logging.getLogger(__name__) + + +def _hex_to_rgb_floats(hex_color: str) -> Tuple[float, float, float]: + """Convert #RRGGBB hex to normalized floats (0-1).""" + hex_color = hex_color.lstrip("#") + if len(hex_color) != 6: + raise ValueError("Color must be #RRGGBB") + r = int(hex_color[0:2], 16) / 255.0 + g = int(hex_color[2:4], 16) / 255.0 + b = int(hex_color[4:6], 16) / 255.0 + return r, g, b + + +class PreviewImageLogic: + """Preview/save an image but fall back to a generated blank if no image provided. + + Useful in logic branches where an IMAGE might be None (e.g., from a switch). + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "save_image": ("BOOLEAN", {"default": False, "tooltip": "Save the resulting image to disk"}), + "width": ("INT", {"default": 1024, "min": 1, "max": 4096, "step": 1}), + "height": ("INT", {"default": 1024, "min": 1, "max": 4096, "step": 1}), + "color": ("COLOR", {"default": "#FFFFFF", "tooltip": "Fill color when generating blank image"}), + }, + "optional": { + "file_name": ("STRING", {"default": "logic_preview", "multiline": False, "tooltip": "Base filename when saving (no extension)"}), + "image": ("IMAGE", {"tooltip": "Optional image tensor; if None blank is generated"}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "preview" + OUTPUT_NODE = False + CATEGORY = "llm_toolkit/utils/logic" + + # ------------------------------------------------------------------ + def _make_blank(self, width: int, height: int, color: str) -> torch.Tensor: + r, g, b = _hex_to_rgb_floats(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=-1) # (H,W,3) + arr = arr[np.newaxis, ...] # (1,H,W,3) + return torch.from_numpy(arr) + + def _save(self, img_tensor: torch.Tensor, name_prefix: str) -> str: + # Lazy import PIL only when saving + from PIL import Image # type: ignore + + # Accept (B,H,W,3) or (B,3,H,W). Convert to (H,W,3) for PIL. + if img_tensor.ndim != 4: + raise ValueError("Expected image tensor (B,3,H,W) or (B,H,W,3)") + t = img_tensor[0].clamp(0, 1).cpu() + if t.shape[0] == 3: # (3,H,W) + img_np = (t.numpy() * 255).astype(np.uint8).transpose(1, 2, 0) + else: # (H,W,3) + img_np = (t.numpy() * 255).astype(np.uint8) + h, w = img_np.shape[:2] + img = Image.fromarray(img_np) + + output_dir = folder_paths.get_output_directory() if folder_paths else os.getcwd() + os.makedirs(output_dir, exist_ok=True) + filename = f"{name_prefix}_{os.getpid()}_{np.random.randint(0,1e6):06d}.png" + path = os.path.join(output_dir, filename) + img.save(path) + return path + + # ------------------------------------------------------------------ + def preview( + self, + save_image: bool, + width: int, + height: int, + color: str, + file_name: str = "logic_preview", + image: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor]: + try: + # Handle ComfyUI IMAGE formats: could be torch.Tensor, list[torch.Tensor], or None + if image is None: + out_img = None + elif isinstance(image, list): + out_img = image[0] if image else None + else: + out_img = image # assume torch.Tensor + + if out_img is None: + logger.debug("PreviewImageLogic: Creating blank image %dx%d color %s", width, height, color) + out_img = self._make_blank(width, height, color) + + if save_image: + try: + path = self._save(out_img, file_name or "logic_preview") + logger.info("PreviewImageLogic: Saved image to %s", path) + except Exception as save_err: + logger.error("PreviewImageLogic: Failed to save image – %s", save_err, exc_info=True) + return (out_img,) + except Exception as e: + logger.error("PreviewImageLogic: %s", e, exc_info=True) + # On catastrophic error, return blank 1x1 black to keep graph alive + fallback = torch.zeros((1, 3, 1, 1), dtype=torch.float32) + return (fallback,) + + +# --------------------------------------------------------------------------- +NODE_CLASS_MAPPINGS = { + "PreviewImageLogic": PreviewImageLogic, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PreviewImageLogic": "Preview Image Logic (LLMToolkit)", +} \ No newline at end of file diff --git a/comfy-nodes/preview_outputs.py b/comfy-nodes/preview_outputs.py index 4fdd968..f66b7eb 100644 --- a/comfy-nodes/preview_outputs.py +++ b/comfy-nodes/preview_outputs.py @@ -40,7 +40,7 @@ class PreviewOutputs: # noqa: N801 – stay consistent with other nodes RETURN_TYPES = ("*", "AUDIO", "IMAGE", "STRING") RETURN_NAMES = ("context", "audio", "image", "lyrics") FUNCTION = "preview" - CATEGORY = "llm_toolkit/utils" + CATEGORY = "llm_toolkit/utils/audio" OUTPUT_NODE = True # ------------------------------------------------------------------ diff --git a/comfy-nodes/switch_any.py b/comfy-nodes/switch_any.py index 2963c6c..2672c8d 100644 --- a/comfy-nodes/switch_any.py +++ b/comfy-nodes/switch_any.py @@ -7,6 +7,8 @@ class AnyType(str): return False WILDCARD = AnyType("*") +# Alias used for readability +any = WILDCARD # Initialize logger logger = logging.getLogger(__name__) @@ -18,7 +20,7 @@ class SwitchAny: Useful for conditional logic in ComfyUI workflows. """ def __init__(self): - self.type = "llm_toolkit/utils" + self.type = "llm_toolkit/utils/logic" @classmethod def INPUT_TYPES(cls): @@ -37,7 +39,7 @@ class SwitchAny: RETURN_NAMES = ("output",) FUNCTION = "switch" OUTPUT_NODE = False # This is a utility node, not an output node - CATEGORY = "llm_toolkit/utils" + CATEGORY = "llm_toolkit/utils/logic" def switch(self, selector: bool, input_a: Any = None, input_b: Any = None) -> Tuple[Any]: """ @@ -72,23 +74,33 @@ class SwitchAny: # In case of error, return None return (None,) + +# --------------------------------------------------------------------------- +# Original SwitchAnyRoute (two outputs, maintains exact type) +# --------------------------------------------------------------------------- + + class SwitchAnyRoute: """ - A reverse switch node that takes a single input of any type and a boolean selector. - It diverts (routes) the input to one of two outputs depending on the selector. - - If selector is True -> the value is emitted on output_true and output_false is None. - If selector is False -> the value is emitted on output_false and output_true is None. + Reverse switch node that takes a single input (any type) and a boolean selector. + Routes the input to either `output_true` or `output_false` so that downstream + nodes can branch based on the boolean flag. """ def __init__(self): - self.type = "llm_toolkit/utils" + self.type = "llm_toolkit/utils/logic" @classmethod def INPUT_TYPES(cls): return { "required": { - "selector": ("BOOLEAN", {"default": True, "tooltip": "True = route to output_true, False = route to output_false"}), + "selector": ( + "BOOLEAN", + { + "default": True, + "tooltip": "True = route to output_true, False = output_false", + }, + ), "input": ("*", {"tooltip": "Input to be routed"}), }, "hidden": {}, @@ -97,30 +109,116 @@ class SwitchAnyRoute: RETURN_TYPES = (WILDCARD, WILDCARD) RETURN_NAMES = ("output_true", "output_false") FUNCTION = "route" - OUTPUT_NODE = False # Utility node - CATEGORY = "llm_toolkit/utils" + OUTPUT_NODE = False + CATEGORY = "llm_toolkit/utils/logic" - def route(self, selector: bool, input: Any) -> Tuple[Any, Any]: - """Route the input to one of the outputs based on selector.""" - logger.info(f"SwitchAnyRoute executing with selector={selector}") + def route(self, selector: bool, input: Any): # noqa: D401 + logger.info("SwitchAnyRoute executing with selector=%s", selector) try: if selector: - logger.info("SwitchAnyRoute: Sending input to output_true (selector=True)") return (input, None) - else: - logger.info("SwitchAnyRoute: Sending input to output_false (selector=False)") - return (None, input) + return (None, input) except Exception as e: - logger.error(f"Error in SwitchAnyRoute: {e}", exc_info=True) + logger.error("Error in SwitchAnyRoute: %s", e, exc_info=True) + return (None, None) + + +# =============================================================== +# _wANY variants – preserve input/output dynamic typing +# =============================================================== + + +class SwitchAny_wANY: + """ + Identical to SwitchAny but declares inputs/outputs using the `any` wildcard + instance so ComfyUI treats the data type dynamically (avoids validation + errors when routing IMAGE/LATENT/etc.). + """ + + def __init__(self): + self.type = "llm_toolkit/utils/logic" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "selector": ("BOOLEAN", {"default": True, "tooltip": "True = output input_a, False = output input_b"}), + }, + "optional": { + "input_a": (any, {"tooltip": "First input to output when selector is True"}), + "input_b": (any, {"tooltip": "Second input to output when selector is False"}), + }, + "hidden": {}, + } + + RETURN_TYPES = (WILDCARD,) + RETURN_NAMES = ("output",) + FUNCTION = "switch" + OUTPUT_NODE = False + CATEGORY = "llm_toolkit/utils/logic" + + def switch(self, selector: bool, input_a: Any = None, input_b: Any = None): # type: ignore[override] + logger.info(f"SwitchAny_wANY executing with selector={selector}") + try: + selected_output = input_a if selector else input_b + logger.info("SwitchAny_wANY selected %s", "input_a" if selector else "input_b") + return (selected_output,) + except Exception as e: + logger.error("Error in SwitchAny_wANY: %s", e, exc_info=True) + return (None,) + + +class SwitchAnyRoute_wANY: + """ + Reverse switch that routes a single input to two outputs using wildcard type. + """ + + def __init__(self): + self.type = "llm_toolkit/utils/logic" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "selector": ( + "BOOLEAN", + { + "default": True, + "tooltip": "True = route to output_true, False = output_false", + }, + ), + "input": (any, {"tooltip": "Input to be routed"}), + }, + "hidden": {}, + } + + RETURN_TYPES = (WILDCARD, WILDCARD) + RETURN_NAMES = ("output_true", "output_false") + FUNCTION = "route" + OUTPUT_NODE = False + CATEGORY = "llm_toolkit/utils/logic" + + def route(self, selector: bool, input: Any): # type: ignore[override] + logger.info("SwitchAnyRoute_wANY executing with selector=%s", selector) + try: + if selector: + return (input, None) + return (None, input) + except Exception as e: + logger.error("Error in SwitchAnyRoute_wANY: %s", e, exc_info=True) return (None, None) # --- Node Mappings for ComfyUI --- NODE_CLASS_MAPPINGS = { "SwitchAny": SwitchAny, - "SwitchAnyRoute": SwitchAnyRoute + "SwitchAnyRoute": SwitchAnyRoute, + "SwitchAny_wANY": SwitchAny_wANY, + "SwitchAnyRoute_wANY": SwitchAnyRoute_wANY, } NODE_DISPLAY_NAME_MAPPINGS = { "SwitchAny": "Switch Any (LLMToolkit)", - "SwitchAnyRoute": "Switch Any Route (LLMToolkit)" + "SwitchAnyRoute": "Switch Any Route (LLMToolkit)", + "SwitchAny_wANY": "Switch Any _wANY (LLMToolkit)", + "SwitchAnyRoute_wANY": "Switch Any Route _wANY (LLMToolkit)", } \ No newline at end of file