LLM input

[Reactantvr]

So, I wanted to have my LLM input into vibe voice and be able to use it in ComfyUI. The plan in the end is to add mic input to the LLM, but I haven't gotten around to that yet. The idea is so I can have it become any character then generate the text needed for that persona. I then use a 1-3 minute voice sample and it just becomes whoever that is. Using Kimi K2 locally for this.

The code below adds 3 nodes, one for LLM input and two for saving the conversation and using it in the next loop. It is setup to work with llama.cpp server.

import os
import requests
import json
import folder_paths
import time

class LlamaCPPChat:
    """
    An advanced ComfyUI node for continuous conversation with a local Llama.cpp server API.
    """
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "user_message": ("STRING", {"multiline": True, "default": "Hi"}),
                "api_url": ("STRING", {"default": "http://127.0.0.1:8080/completion"}),
                "prompt_format": ("STRING", {
                    "multiline": True,
                    "default": "\nUSER: {user_message}\nASSISTANT:"
                }),
                "stop_tokens": ("STRING", {"default": "USER:,ASSISTANT:,###", "tooltip": "Comma-separated list of words that stop generation."}),
                "max_tokens": ("INT", { "default": 256, "min": -1, "max": 8192, "step": 64 }),
                "temperature": ("FLOAT", { "default": 0.7, "min": 0.0, "max": 2.0, "step": 0.05 }),
                "repeat_penalty": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 5.0, "step": 0.05}),
                "top_k": ("INT", {"default": 40, "min": 0, "max": 100, "step": 1}),
                "top_p": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.05}),
                "seed": ("INT", { "default": -1, "min": -1, "max": 0xffffffffffffffff }),
            },
            "optional": { "conversation_in": ("STRING", { "forceInput": True }) }
        }

    RETURN_TYPES = ("STRING", "STRING",)
    RETURN_NAMES = ("assistant_response", "new_turn_to_save",)
    FUNCTION = "chat"
    CATEGORY = "LLM"

    def chat(self, user_message, api_url, prompt_format, stop_tokens, max_tokens, temperature, repeat_penalty, top_k, top_p, seed, conversation_in=None):
        history = conversation_in if conversation_in is not None else ""
        
        full_prompt = history + prompt_format.replace("{user_message}", user_message)
        
        stop_list = [token.strip() for token in stop_tokens.split(',') if token.strip()]

        headers = {"Content-Type": "application/json"}
        payload = {
            "prompt": full_prompt, "n_predict": max_tokens, "temperature": temperature,
            "seed": seed, "stop": stop_list, "repeat_penalty": repeat_penalty,
            "top_k": top_k, "top_p": top_p, "stream": False
        }
        
        assistant_response = ""
        try:
            response = requests.post(api_url, headers=headers, data=json.dumps(payload), timeout=20000)
            response.raise_for_status()
            assistant_response = response.json().get("content", "")

            # FIXED: Clean the response to prevent corrupted history.
            # This strips whitespace and removes any stop words the AI might have added at the end.
            assistant_response = assistant_response.strip()
            for token in stop_list:
                if assistant_response.endswith(token):
                    assistant_response = assistant_response[:-len(token)].rstrip()
            
            print(f"[LlamaCPP Chat] Cleaned Assistant Response: {assistant_response}")

        except Exception as e:
            print(f"[LlamaCPP Chat] Error: {e}")

        new_turn_to_save = prompt_format.replace("{user_message}", user_message) + assistant_response
        return (assistant_response, new_turn_to_save,)

# -------------------------------------------------------------------
# NODE: Store Text (Append Mode)
# -------------------------------------------------------------------
class StoreTextFileNode:
    def __init__(self):
        self.output_dir = os.path.join(folder_paths.get_output_directory(), "text_history")
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "text": ("STRING", {"forceInput": True}),
                "filename": ("STRING", {"default": "conversation_history.txt"}),
                "reset_file": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ()
    FUNCTION = "store_text"
    OUTPUT_NODE = True
    CATEGORY = "Text IO"

    def store_text(self, text, filename, reset_file):
        file_path = os.path.join(self.output_dir, filename)
        file_mode = 'w' if reset_file else 'a'
        if reset_file:
            print(f"[Store Text (File)] Resetting and overwriting '{filename}'.")
        else:
            print(f"[Store Text (File)] Appending to '{filename}'.")
        try:
            with open(file_path, file_mode, encoding='utf-8') as f:
                f.write(text)
        except Exception as e:
            print(f"Error saving text file: {e}")
        return {}

# -------------------------------------------------------------------
# NODE: Load Text from a File
# -------------------------------------------------------------------
class LoadTextFileNode:
    def __init__(self):
        self.input_dir = os.path.join(folder_paths.get_output_directory(), "text_history")

    @classmethod
    def INPUT_TYPES(s):
        return { "required": { "filename": ("STRING", {"default": "conversation_history.txt"}) } }
    
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "load_text"
    CATEGORY = "Text IO"

    def IS_CHANGED(self, filename):
        return time.time()

    def load_text(self, filename):
        file_path = os.path.join(self.input_dir, filename)
        text = ""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                text = f.read()
            print(f"[Load Text (File)] Loaded {len(text)} chars from '{filename}'.")
        except FileNotFoundError:
            print(f"[Load Text (File)] File not found, starting fresh.")
            pass
        return (text,)


# -------------------------------------------------------------------
# MAPPINGS
# -------------------------------------------------------------------
NODE_CLASS_MAPPINGS = {
    "LlamaCPP Chat": LlamaCPPChat,
    "StoreTextFileNode": StoreTextFileNode,
    "LoadTextFileNode": LoadTextFileNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LlamaCPP Chat": "💬 Llama.cpp Chat",
    "StoreTextFileNode": "💾 Store Text (File)",
    "LoadTextFileNode": "📄 Load Text (File)",
}