initial commit

initial commit

__init__.py

@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

nodes.py

@@ -0,0 +1,309 @@
+import subprocess
+import sys
+import re
+import comfy.model_management as model_management
+import gc
+import torch
+import requests
+import time
+
+class EbuLMStudioLoadModel:
+    ERROR_NO_MODEL_FOUND = "no model by that name found"
+    ERROR_CANT_CONNECT = "can't connect to lmstudio"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "input_string": ("STRING",),
+                "model_search_string": ("STRING", {"default": "llama"}),
+                "context_length": ("INT", {"default": 4096, "min": 1000, "max": 200000}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "unload_image_models_first": ("BOOLEAN", {"default": False}),
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "STRING", "STRING",)
+    RETURN_NAMES = ("export_string", "model_selected", "models_found",)
+    FUNCTION = "load_model"
+    CATEGORY = "LMStudio"
+
+    def check_models_loaded(self):
+        """Check if any models are currently loaded in LMStudio"""
+        try:
+            # Run the 'lms status' command
+            command = ['lms', 'status']
+            status_process = self.run_command(command)
+
+            if not status_process:
+                return True  # Assume models are loaded if we can't check
+
+            # Look for the "No Models Loaded" indicator in the output
+            return "No Models Loaded" not in status_process.stdout
+
+        except Exception as e:
+            print(f"Error checking model status: {str(e)}", file=sys.stderr)
+            return True  # Assume models are loaded if check fails
+
+    def run_command(self, command):
+        try:
+            # Use utf-8 encoding with errors='replace' to handle problematic characters
+            process = subprocess.Popen(
+                command,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                universal_newlines=False  # Get bytes instead of string
+            )
+
+            stdout_bytes, stderr_bytes = process.communicate()
+
+            # Decode with 'utf-8' and replace invalid characters
+            stdout = stdout_bytes.decode('utf-8', errors='replace')
+            stderr = stderr_bytes.decode('utf-8', errors='replace')
+
+            if process.returncode != 0:
+                raise subprocess.CalledProcessError(process.returncode, command, stdout, stderr)
+
+            result = subprocess.CompletedProcess(
+                args=command,
+                returncode=process.returncode,
+                stdout=stdout,
+                stderr=stderr
+            )
+            return result
+
+        except subprocess.CalledProcessError as e:
+            print(f"Error executing command {command}: {e}", file=sys.stderr)
+            return None
+
+    def unload_image_models(self):
+        print("Unload Image Models:")
+        print(" - Unloading all image models...")
+        model_management.unload_all_models()
+        model_management.soft_empty_cache(True)
+        try:
+            print(" - Clearing Cache...")
+            gc.collect()
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+        except:
+            print("   - Unable to clear cache")
+
+    def find_matching_model(self, search_term):
+        try:
+            print("Fetching model list...")
+
+            # Run the 'lms ls --detailed' command to get available models
+            command = ['lms', 'ls', '--detailed']
+            list_process = self.run_command(command)
+
+            if not list_process:
+                return self.ERROR_CANT_CONNECT, ""
+
+            result = list_process.stdout.replace("\r\n", "\n").replace("\r", "\n")
+
+            # Build a case-insensitive search pattern from the search term
+            search_pattern = re.escape(search_term.lower())
+            matched_models = []
+
+            # Iterate through the output and find matching models
+            for line in result.split("\n"):
+                line = line.strip()
+
+                # Only match lines that start with '/', these are the llm full names
+                if line.startswith('/'):
+                    model_path = line.split()[0].lstrip("/")
+                    if re.search(search_pattern, model_path.lower()):
+                        matched_models.append(model_path)
+
+            # Return the first matched model or None if no match
+            return matched_models[0] if matched_models else self.ERROR_NO_MODEL_FOUND, result
+
+        except subprocess.CalledProcessError:
+            print("Error: Unable to fetch the model list. Is LMStudio running?", file=sys.stderr)
+            return self.ERROR_CANT_CONNECT, ""
+        except Exception as e:
+            print(f"Unexpected error: {e}", file=sys.stderr)
+            return self.ERROR_CANT_CONNECT, ""
+
+    def load_model(self, input_string, model_search_string, context_length, seed, unload_image_models_first):
+        try:
+            # Only unload image models if specified
+            if unload_image_models_first:
+                print("Unloading ComfyUI image models to save memory (as requested)")
+                self.unload_image_models()
+
+            # Always unload all LLM models first if there are any in memory
+            if self.check_models_loaded():
+	            print("Unloading LLMs in memory...")
+	            subprocess.run(['lms', 'unload', '--all'],
+	                           stdout=subprocess.DEVNULL,
+	                           stderr=subprocess.DEVNULL,
+	                           check=True)
+
+            print(f"Loading model: {model_search_string}")
+            model_path, models_found = self.find_matching_model(model_search_string)
+            if model_path in [self.ERROR_NO_MODEL_FOUND, self.ERROR_CANT_CONNECT]:
+                return (input_string, model_path, models_found)
+
+            # Run the 'lms load' command to load the model
+            command = ['lms', 'load', model_path, '-y', f'--context-length={context_length}', f'--gpu=1']
+            load_process = self.run_command(command)
+
+            if load_process and load_process.returncode == 0:  # Check return code for success
+                print(f"Model loaded: {model_path}")
+                return (input_string, model_path, models_found)
+            else:
+                print(f"Warning: Issue loading model: {model_path}")
+                return (input_string, self.ERROR_NO_MODEL_FOUND, models_found)
+
+        except Exception as e:
+            print(f"Error: {str(e)}", file=sys.stderr)
+
+class EbuLMStudioUnload:
+    ERROR_CANT_CONNECT = "can't connect to lmstudio"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "input_string": ("STRING", {"default": ""}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+             },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("export_string",)
+    FUNCTION = "unload_all"
+    CATEGORY = "LMStudio"
+
+    def unload_all(self, input_string, seed):
+        try:
+            print("Calling lms unload --all")
+
+            # Run the command 'lms unload --all' to unload all llms
+            result = subprocess.run(['lms', 'unload', '--all'],
+                                    capture_output=True,
+                                    text=True,
+                                    check=True)
+
+            # Print the output of the command
+            print("Command output:", result.stdout)
+
+            # If there's any error output, print it as well
+            if result.stderr:
+                print("", result.stderr, file=sys.stderr)
+
+            # Return the input string unchanged
+            return (input_string,)
+
+        except subprocess.CalledProcessError as e:
+            print(f"An error occurred while running the command: {e}", file=sys.stderr)
+            return (input_string,)
+        except Exception as e:
+            print(f"An unexpected error occurred: {e}", file=sys.stderr)
+            return (input_string,)
+
+
+# initially inspired by the excellent LLM node from CrasH Utils Custom Nodes
+class EbuLMStudioMakeRequest:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": ("STRING", { "multiline": True, "default": "" }),
+                "system_message": ("STRING", { "multiline": True, "default": "You are an assistant designed to craft AI image prompts for an AI image generator that uses natural language prompts. Follow the instructions you are given, or use the guidelines, to create a detailed prompt that includes creativity and amazing visual details for an unforgetable image. Respond with just your new prompt." }),
+                "url": ("STRING", { "multiline": False, "default": "http://127.0.0.1:1234/v1/chat/completions" }),
+                "context_length": ("INT", { "default": 4096, "min": 512, "max": 8192, "display": "slider" }),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "max_tokens": ("INT", {"default": 300, "min": 10, "max": 100000}),
+                "temp": ("FLOAT", {
+                    "default": 0.7,
+                    "min": 0,
+                    "max": 1,
+                    "step": 0.05,
+                    "display": "number"
+                }),
+                "top_p": ("FLOAT", {
+                    "default": 0.95,
+                    "min": 0,
+                    "max": 1,
+                    "step": 0.01,
+                    "display": "number"
+                }),
+                "utf8_safe_replace": ("BOOLEAN", {"default": False}),
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("generated_text",)
+    FUNCTION = "generateText"
+    OUTPUT_NODE = False
+    CATEGORY = "LMStudio"
+
+    def sanitize_utf8(self, text):
+        # Sanitize text to ensure it's valid UTF-8 using unicode normalization
+        print("Sanitizing the response into UTF-8 safe characters as requested...")
+        try:
+            from unicodedata import normalize
+            # Normalize to decomposed form and convert to ASCII, dropping diacritics
+            return normalize('NFKD', text).encode('ascii', 'ignore').decode('ascii')
+        except Exception as e:
+            print(f"Warning: UTF-8 sanitization failed: {str(e)}")
+            # Fallback to simple replacement if something goes wrong
+            return text.encode('ascii', 'replace').decode('ascii')
+
+    def generateText(self, prompt, system_message, url, context_length, seed, max_tokens, temp, top_p, utf8_safe_replace):
+        description = self.call_api(prompt, system_message, url, context_length, seed, max_tokens, temp, top_p, utf8_safe_replace)
+        return (description,)
+
+    def call_api(self, prompt_text, system_message, url, context_length, seed, max_tokens, temp, top_p, utf8_safe_replace):
+        payload = {
+            "messages": [
+            { "role": "system", "content": f"{system_message}\n" },
+            { "role": "user", "content": f"{prompt_text}\n" }],
+            "max_tokens": max_tokens,
+            "temperature": temp,
+            "top_p": top_p,
+            "seed": seed
+        }
+        start_time = time.time()
+        try:
+            response = requests.post(url, json=payload)
+            if response.status_code == 200:
+                result_json = response.json()
+                generatedText = result_json["choices"][0]["message"]["content"]
+
+                if utf8_safe_replace:
+                    generatedText = self.sanitize_utf8(generatedText)
+                    print("Response (UTF-8 sanitized): " + generatedText)
+                else:
+                    print("Response: " + generatedText)
+
+                end_time = time.time()
+                elapsed_time = end_time - start_time
+                print(f"\n[Response generation time: {elapsed_time:.2f} seconds]\n")
+                return generatedText
+            else:
+                error_message = f"Error {response.status_code}: {response.text}"
+                print(error_message)
+                return error_message
+        except requests.exceptions.RequestException as e:
+            error_message = f"Request failed: {str(e)}"
+            print(error_message)
+            return error_message
+
+NODE_CLASS_MAPPINGS = {
+    "EbuLMStudioLoadModel": EbuLMStudioLoadModel,
+    "EbuLMStudioUnload": EbuLMStudioUnload,
+    "EbuLMStudioMakeRequest": EbuLMStudioMakeRequest,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "EbuLMStudioLoadModel": "EBU LMStudio Load",
+    "EbuLMStudioUnload": "EBU LMStudio Unload All",
+    "EbuLMStudioMakeRequest": "EBU LMStudio Make Request",
+}

pyproject.toml

@@ -0,0 +1,16 @@
+[project]
+name = "comfyui-ebu-lmstudio"
+description = "This ComfyUI extension provides custom nodes for integrating with LMStudio, allowing for loading, managing, and making requests of LLM models through the LMStudio local server and command-line interface."
+version = "1.0.0"
+authors = [
+    {name = "burnsbert"}
+]
+keywords = ["comfyui", "lmstudio", "llm", "stable-diffusion", "flux"]
+license = {file = "LICENSE"}
+
+[project.urls]
+Repository = "https://github.com/burnsbert/ComfyUI-EBU-LMStudio"
+
+[tool.comfy]
+PublisherId = "burnsbert"
+DisplayName = "ComfyUI-EBU-LMStudio"

readme.md

@@ -0,0 +1,104 @@
+# ComfyUI-EBU-LMStudio
+
+A collection of custom nodes for ComfyUI that provide integration with LMStudio. These nodes allow you to load, unload, and interact with LLMs through the LMStudio command-line interface. This can be used to facilitate the automatic creation of descriptive prompts using a LLM, which is helpful for modern image models like Flux that respond well to detailed prompts in natural language.
+
+## Features
+
+- Load and manage LLM models directly from ComfyUI
+- Fuzzy model name searching for easy model selection
+- Configurable context length and parameters
+- Optional image model unloading to free up VRAM
+- Simple prompt submission interface
+- Option for UTF-8 safe LLM text responses
+- Detailed error reporting and status messages
+
+## Installation
+
+1. Copy this directory to your `ComfyUI/custom_nodes/` directory
+2. Install LMStudio and make sure 'lms' command line tools are working and that 'Developer Mode' is enabled to run LMStudio as a local server.
+3. Make sure to have at least one LLM installed inside of LMStudio.
+4. Test this by running `lms ls --detailed` from the command line. You should see a list of all installed LLMs. If this doesn't work you may need to modify your PATH.
+5. Restart ComfyUI
+
+## Nodes
+
+### EBU LMStudio Load Model
+
+Unloads any loaded LLMs and then loads an LLM model into LMStudio with the following parameters:
+
+- `input_string`: Input text (passed through unchanged), serves to trigger the action
+- `model_search_string`: Part of the model name to search for (default: "llama"). Spaces are treated as wild cards. First matching LLM is loaded. So "nemo q5" will match "Mistral-Nemo-Instruct-2407-abliterated.Q5_K_M.gguf" for example.
+- `context_length`: Model context length (default: 4096)
+- `seed`: Random seed to ensure that this runs on each request
+- `unload_image_models_first`: Whether to unload ComfyUI image models before loading LLM
+
+Returns:
+- `export_string`: The input string (unchanged)
+- `model_selected`: Full path of the loaded model or error message
+- `models_found`: Complete list of available models
+
+### EBU LMStudio Unload Models
+
+Unloads all currently loaded LLM models from LMStudio.
+
+Parameters:
+- `input_string`: Input text (passed through unchanged), serves to trigger the action
+- `seed`: Random seed to ensure that this runs on each request
+
+Returns:
+- `export_string`: The input string (unchanged)
+
+### EBU LMStudio Submit Prompt
+
+Submits a prompt to the currently loaded LLM and returns the response.
+
+Parameters:
+- `prompt`: The text prompt to send to the LLM
+- `system_message`: System context message (default provides image prompt generation context)
+- `url`: LMStudio API endpoint (default: http://127.0.0.1:1234/v1/chat/completions)
+- `context_length`: Context window size (default: 4096)
+- `seed`: Random seed for reproducibility
+- `max_tokens`: Maximum response length (default: 300)
+- `temp`: Temperature parameter (default: 0.5)
+- `top_p`: Top-p sampling parameter (default: 0.9)
+- `utf8_safe_replace`: Convert response to safe ASCII characters (default: false)
+
+Returns:
+- `generated_text`: The LLM's response
+
+## Requirements
+
+- Python 3.x
+- ComfyUI
+- LMStudio with CLI access and local server configured
+
+## Common Issues and Solutions
+
+### "can't connect to lmstudio"
+- Ensure LMStudio is running
+- Check that LMS CLI Access is healthy
+- Verify the port number matches your LMStudio settings
+- Confirm LMStudio is in your system PATH
+
+### "no model by that name found"
+- Check that the model name search string matches part of your model's name
+- Use `lms ls --detailed` in terminal to see available models
+- Models must be added to LMStudio before they can be loaded
+
+### Model Loading Issues
+- Ensure you have enough VRAM available
+- Try unloading image models first
+- Check LMStudio logs for detailed error messages
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details
+
+## Acknowledgments
+
+- Initially inspired by the LLM node from CrasH Utils Custom Nodes
+- Thanks to the ComfyUI and LMStudio teams
+
+## For more help
+
+See the included example workflow `flux_example_ebu_lmstudio.json` to see how these three nodes can work together for VRAM efficient prompt generation and execution.

requirements.txt

@@ -0,0 +1 @@
+# No additional requirements beyond Python standard library