# ComfyUI/custom_nodes/ComfyUI_LocalLLMNodes/local_llm_connector.py import os import folder_paths # --- Simple logging utility for this package --- def log(message): print(f"[LocalLLMNodes] {message}") # --- Library for local LLM inference --- # Requires 'transformers' and 'torch'. These are checked in __init__.py. try: from transformers import AutoTokenizer, AutoModelForCausalLM # Optional: For quantization (requires 'bitsandbytes') # from transformers import BitsAndBytesConfig import torch HF_AVAILABLE = True except ImportError: log("Warning: transformers or torch library not found. Local LLM node will not work.") HF_AVAILABLE = False AutoTokenizer = None AutoModelForCausalLM = None torch = None # Define category directly for this package LOCAL_LLM_CATEGORY = "Local LLM Nodes/LLM Connectors" def get_local_llm_model_names(): """Discovers subdirectories within models/LLM that could be local LLM models.""" llm_models_dir = os.path.join(folder_paths.models_dir, "LLM") model_names = [] if os.path.exists(llm_models_dir): try: # List directories inside models/LLM for item in os.listdir(llm_models_dir): item_path = os.path.join(llm_models_dir, item) if os.path.isdir(item_path): # Basic check: assume any directory is a potential model model_names.append(item) except Exception as e: log(f"Error scanning models/LLM directory: {e}") else: log(f"models/LLM directory not found: {llm_models_dir}") if not model_names: model_names = ["No_Local_Models_Found"] # Placeholder if no models found return model_names class SetLocalLLMServiceConnector: """ A node to select and prepare a connection to a local LLM model. Models should be placed in ComfyUI/models/LLM/your_model_name. Requires 'transformers' and 'torch': pip install transformers torch """ @classmethod def INPUT_TYPES(cls): model_names = get_local_llm_model_names() return { "required": { "local_model_name": (model_names, {"default": model_names[0] if model_names else "No_Local_Models_Found"}), }, } # --- CRITICAL: Match the type identifier expected by consuming nodes --- # Based on MieNodes structure (TextTranslator, KontextPromptGenerator), this should be "LLMServiceConnector" # This allows your node to connect to standard MieNodes prompt generators if desired. RETURN_TYPES = ("LLMServiceConnector",) FUNCTION = "get_connector" CATEGORY = LOCAL_LLM_CATEGORY OUTPUT_NODE = False def get_connector(self, local_model_name): """ Returns a connector object that can invoke the selected local model. The actual model loading happens within the connector's invoke method. """ if not HF_AVAILABLE: raise Exception("The 'transformers' and 'torch' libraries are required for the Local LLM node but are not installed.") if local_model_name == "No_Local_Models_Found": raise Exception("No local LLM models found in models/LLM directory. Please place your models there.") model_path = os.path.join(folder_paths.models_dir, "LLM", local_model_name) if not os.path.exists(model_path): raise FileNotFoundError(f"Selected local LLM model path not found: {model_path}") # Create and return the connector instance. Model loading is deferred. connector = LocalLLMServiceConnector(model_path) return (connector,) class LocalLLMServiceConnector: """ Represents the connection to a specific local LLM. Handles loading (on first use) and invocation. Compatible with TextTranslator and KontextPromptGenerator via the 'invoke' method. The type identifier used by nodes expecting this connector is "LLMServiceConnector". """ def __init__(self, model_path): self.model_path = model_path self.model = None self.tokenizer = None self.is_loaded = False def _load_model(self): """Loads the model and tokenizer if not already loaded.""" if self.is_loaded: return # Already loaded try: log(f"Loading local LLM model from: {self.model_path}") # --- Model Loading Configuration --- tokenizer_kwargs = { "trust_remote_code": True # Needed for some non-standard models } # --- Example Quantization Config (Uncomment and use if needed) --- # Requires 'bitsandbytes': pip install bitsandbytes # quantization_config = BitsAndBytesConfig( # load_in_4bit=True, # bnb_4bit_compute_dtype=torch.float16, # bnb_4bit_use_double_quant=True, # bnb_4bit_quant_type="nf4" # ) model_kwargs = { "trust_remote_code": True, # --- Add Quantization Config if using --- # "quantization_config": quantization_config, # <-- Add for quantization # --- Device Placement --- "device_map": "auto", # Automatic device placement (CPU/GPU) - Often crucial with quantization # --- Precision (if not using quantization config) --- # "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32, } # --- Load Tokenizer --- self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, **tokenizer_kwargs) # Handle models without a pad token (common with Llama-based models) if self.tokenizer.pad_token_id is None: self.tokenizer.pad_token_id = self.tokenizer.eos_token_id log(f"Set pad_token_id to eos_token_id ({self.tokenizer.eos_token_id})") # --- Load Model --- # Note: device_map="auto" is often crucial here, especially with quantization self.model = AutoModelForCausalLM.from_pretrained(self.model_path, **model_kwargs) # Manual device placement is usually handled by device_map="auto" # if "device_map" not in model_kwargs and torch.cuda.is_available(): # self.model.to("cuda") # elif "device_map" not in model_kwargs: # self.model.to("cpu") self.is_loaded = True log(f"Local LLM model loaded successfully from: {self.model_path}") except Exception as e: error_msg = f"Failed to load local LLM model from {self.model_path}: {e}" log(error_msg) # Re-raise to stop execution if loading fails raise Exception(error_msg) from e # Chain the exception def invoke(self, messages, **generation_kwargs): """ Generates text using the local LLM based on the messages list. Mimics the API expected by TextTranslator/KontextPromptGenerator. :param messages: List of message dictionaries (like OpenAI format). Example: [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello!"}] :param generation_kwargs: Additional arguments for text generation (e.g., max_new_tokens, temperature, seed). :return: The generated text string. """ try: if not self.is_loaded: self._load_model() # Load model on first invocation if not self.model or not self.tokenizer: raise Exception("Local LLM model or tokenizer failed to load.") # --- Format messages for the local model --- # Try using the tokenizer's chat template if available (more robust) prompt = "" try: # Many models come with a chat template prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) if not isinstance(prompt, str): # Fallback if apply_chat_template behaves unexpectedly raise ValueError("apply_chat_template did not return a string") except Exception as e: # Fallback to simple formatting if chat template fails or isn't available log(f"Falling back to simple prompt formatting: {e}") prompt_parts = [] for message in messages: role = message.get('role', 'user') content = message.get('content', '').strip() # Strip content whitespace if content: # Only add non-empty content lines prompt_parts.append(f"{role.capitalize()}: {content}") prompt = "\n".join(prompt_parts) if prompt: # Add the final prompt indicator only if there's content prompt += "\nAssistant:" # Encourage assistant response else: prompt = "Assistant:" # Fallback if messages were empty if not prompt.strip(): log("Warning: Generated prompt is empty or whitespace.") return "" # Return empty string if prompt is empty # --- Tokenize the prompt --- try: inputs = self.tokenizer(prompt, return_tensors="pt") # Consider moving inputs to model's device if not using device_map="auto" # if hasattr(self.model, 'device') and self.model.device.type != 'meta': # inputs = {k: v.to(self.model.device) for k, v in inputs.items()} except Exception as e: error_msg = f"Error tokenizing prompt: {e}" log(error_msg) raise Exception(error_msg) from e # --- Set default generation parameters --- # These defaults should be reasonable for prompt generation tasks default_kwargs = { "max_new_tokens": 250, # Slightly higher default for complex prompts "temperature": 0.7, # Default creativity "do_sample": True, # Enable sampling for variety "pad_token_id": self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id, # "eos_token_id": self.tokenizer.eos_token_id, # Optional: explicitly set EOS } # Update defaults with any provided kwargs (e.g., from consuming nodes) # Filter out kwargs that might cause issues if not explicitly supported filtered_kwargs = {k: v for k, v in generation_kwargs.items() if k in ['max_new_tokens', 'temperature', 'top_p', 'top_k', 'do_sample', 'num_beams', 'early_stopping', 'pad_token_id', 'eos_token_id', 'seed']} # Handle 'seed' if passed - PyTorch manual seed (affects stochastic operations) if 'seed' in generation_kwargs and generation_kwargs['seed'] is not None: try: torch.manual_seed(generation_kwargs['seed']) except Exception as e: log(f"Warning: Could not set seed: {e}") default_kwargs.update(filtered_kwargs) # --- Generate text --- try: self.model.eval() with torch.no_grad(): outputs = self.model.generate(**inputs, **default_kwargs) except Exception as e: error_msg = f"Error during model.generate: {e}" log(error_msg) raise Exception(error_msg) from e # --- Decode the generated tokens --- try: # Extract only the newly generated part (skip the input prompt tokens) generated_tokens = outputs[:, inputs['input_ids'].shape[-1]:] generated_text = self.tokenizer.decode(generated_tokens[0], skip_special_tokens=True) # Ensure the output is a clean string final_text = generated_text.strip() return final_text # <-- Return ONLY the generated string except Exception as e: error_msg = f"Error decoding generated tokens: {e}" log(error_msg) raise Exception(error_msg) from e except Exception as e: # Catch any error that occurred within the try block and log it error_msg = f"Error in invoke method: {str(e)}" log(error_msg) # Re-raise the exception so the calling node knows it failed raise e # Or raise Exception(error_msg) from e