kenyonxu-ComfyUI_KaiUtilities/nodes.py

import json
import os

import folder_paths

# Import parent LoraLoader from ComfyUI built-in nodes
from nodes import LoraLoader

from .cache_manager import cache_manager


class KaiCivitaiLoRALoader(LoraLoader):
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL", {}),
                "clip": ("CLIP", {}),
                "lora_name": (folder_paths.get_filename_list("loras"), {}),
                "strength_model": (
                    "FLOAT",
                    {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01},
                ),
                "strength_clip": (
                    "FLOAT",
                    {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01},
                ),
            }
        }

    RETURN_TYPES = ("MODEL", "CLIP", "STRING", "STRING")
    RETURN_NAMES = ("model", "clip", "trigger_words", "civitai_json")
    FUNCTION = "load_lora_with_info"
    CATEGORY = "KaiUtilities/loaders"

    def load_lora_with_info(self, model, clip, lora_name, strength_model, strength_clip):
        # Load LoRA using parent logic
        model_out, clip_out = self.load_lora(
            model, clip, lora_name, strength_model, strength_clip
        )

        # Get cached info if available (node does NOT auto-fetch; user must click button)
        try:
            lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
        except (FileNotFoundError, KeyError):
            # If path resolution fails, still return LoRA outputs with empty extras
            return (model_out, clip_out, "", "")

        # Use sync cache lookup because ComfyUI node FUNCTIONs run synchronously
        # in worker threads (not in the async event loop). The sync hash function
        # is safe here because worker threads can block on file I/O.
        try:
            cache = cache_manager.get_cache_for_file_sync(lora_path)
        except Exception:
            # Defensive: any unexpected failure in cache lookup should not crash the node
            return (model_out, clip_out, "", "")

        trigger_words = ""
        civitai_json = ""
        if cache:
            words = cache.get("triggerWords")
            if isinstance(words, list):
                trigger_words = ", ".join(str(w) for w in words)
            civitai_json = json.dumps(cache)

        return (model_out, clip_out, trigger_words, civitai_json)


class TextLineCount:
    """Count the number of lines in a multiline text input."""

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "text": ("STRING", {"multiline": True}),
            }
        }

    RETURN_TYPES = ("INT",)
    RETURN_NAMES = ("count",)
    FUNCTION = "count_lines"
    CATEGORY = "KaiUtilities/text"

    def count_lines(self, text):
        if not text:
            return (0,)
        count = text.count("\n") + 1
        return (count,)


class SaveAudioSilent:
    """Save audio without being a terminal output node. For use inside for-loops."""

    def __init__(self):
        self.output_dir = folder_paths.get_output_directory()
        self.type = "output"

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "audio": ("AUDIO",),
                "filename_prefix": ("STRING", {"default": "audio"}),
            },
            "optional": {
                "index": ("INT", {"default": 0, "min": 0, "max": 999999}),
            },
        }

    RETURN_TYPES = ("AUDIO",)
    RETURN_NAMES = ("audio",)
    FUNCTION = "save_audio"
    CATEGORY = "KaiUtilities/audio"

    def save_audio(self, audio, filename_prefix, index=0):
        if audio is None:
            print("[SaveAudioSilent] WARNING: audio input is None, skipping save")
            return (audio,)

        # Extract waveform and sample_rate — handle multiple dict key conventions
        # Use explicit None checks (not "or") because tensors are not bool-safe
        waveform = audio.get("waveform")
        if waveform is None:
            waveform = audio.get("samples")
        if waveform is None:
            waveform = audio.get("data")
        sample_rate = audio.get("sample_rate")
        if sample_rate is None:
            sample_rate = audio.get("sr")
        if sample_rate is None:
            sample_rate = audio.get("sampling_rate")
        if sample_rate is None:
            sample_rate = 44100

        if waveform is None:
            print(f"[SaveAudioSilent] WARNING: no waveform found in audio dict keys={list(audio.keys())}, skipping save")
            return (audio,)

        try:
            import torchaudio

            # Resolve output path with auto-incrementing counter
            full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
                filename_prefix, self.output_dir, 0, 0
            )
            os.makedirs(full_output_folder, exist_ok=True)
            full_path = os.path.join(full_output_folder, f"{filename}_{counter:05d}_.wav")

            # waveform: [batch, channels, samples] -> [channels, samples] for torchaudio
            wav = waveform.cpu()
            if wav.dim() == 3:
                wav = wav.squeeze(0)
            torchaudio.save(full_path, wav, int(sample_rate))

            print(f"[SaveAudioSilent] Saved: {full_path}")
        except Exception as e:
            print(f"[SaveAudioSilent] ERROR saving audio: {e}")
            raise

        return (audio,)