From c01d0c3bbd768c50c580e1b99bda24cd2c4e3fa2 Mon Sep 17 00:00:00 2001 From: "Eric W. Burns" Date: Mon, 27 Jan 2025 00:23:23 -0500 Subject: [PATCH] initial commit initial commit --- __init__.py | 3 + flux_example_ebu_lmstudio.json | 1119 ++++++++++++++++++++++++++++++++ nodes.py | 309 +++++++++ pyproject.toml | 16 + readme.md | 104 +++ requirements.txt | 1 + 6 files changed, 1552 insertions(+) create mode 100644 __init__.py create mode 100644 flux_example_ebu_lmstudio.json create mode 100644 nodes.py create mode 100644 pyproject.toml create mode 100644 readme.md create mode 100644 requirements.txt diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..d721463 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/flux_example_ebu_lmstudio.json b/flux_example_ebu_lmstudio.json new file mode 100644 index 0000000..2f34475 --- /dev/null +++ b/flux_example_ebu_lmstudio.json @@ -0,0 +1,1119 @@ +{ + "last_node_id": 110, + "last_link_id": 245, + "nodes": [ + { + "id": 13, + "type": "SamplerCustomAdvanced", + "pos": [ + 842, + 215 + ], + "size": [ + 355.20001220703125, + 326 + ], + "flags": {}, + "order": 20, + "mode": 0, + "inputs": [ + { + "name": "noise", + "type": "NOISE", + "link": 37, + "slot_index": 0 + }, + { + "name": "guider", + "type": "GUIDER", + "link": 30, + "slot_index": 1 + }, + { + "name": "sampler", + "type": "SAMPLER", + "link": 19, + "slot_index": 2 + }, + { + "name": "sigmas", + "type": "SIGMAS", + "link": 20, + "slot_index": 3 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 23, + "slot_index": 4 + } + ], + "outputs": [ + { + "name": "output", + "type": "LATENT", + "links": [ + 24 + ], + "slot_index": 0, + "shape": 3 + }, + { + "name": "denoised_output", + "type": "LATENT", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "SamplerCustomAdvanced" + }, + "widgets_values": [], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 25, + "type": "RandomNoise", + "pos": [ + 470, + 611 + ], + "size": [ + 315, + 82 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "NOISE", + "type": "NOISE", + "links": [ + 37 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "RandomNoise" + }, + "widgets_values": [ + 627278456893124, + "randomize" + ] + }, + { + "id": 10, + "type": "VAELoader", + "pos": [ + 866, + 649 + ], + "size": [ + 315, + 58 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 12 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "VAELoader" + }, + "widgets_values": [ + "ae.sft" + ] + }, + { + "id": 5, + "type": "EmptyLatentImage", + "pos": [ + 473, + 450 + ], + "size": [ + 315, + 106 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 23 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "EmptyLatentImage" + }, + "widgets_values": [ + 1280, + 768, + 1 + ] + }, + { + "id": 17, + "type": "BasicScheduler", + "pos": [ + 468, + 867 + ], + "size": [ + 315, + 106 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 38, + "slot_index": 0 + } + ], + "outputs": [ + { + "name": "SIGMAS", + "type": "SIGMAS", + "links": [ + 20 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "BasicScheduler" + }, + "widgets_values": [ + "simple", + 8, + 1 + ] + }, + { + "id": 8, + "type": "VAEDecode", + "pos": [ + 1248, + 192 + ], + "size": [ + 210, + 46 + ], + "flags": {}, + "order": 21, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 24 + }, + { + "name": "vae", + "type": "VAE", + "link": 12 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 41 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + }, + "widgets_values": [] + }, + { + "id": 26, + "type": "PreviewImage", + "pos": [ + 1499.792236328125, + 187.1358642578125 + ], + "size": [ + 528.840576171875, + 480.5522155761719 + ], + "flags": {}, + "order": 22, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 41 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage" + }, + "widgets_values": [], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 6, + "type": "CLIPTextEncode", + "pos": [ + 361.58343505859375, + 219.5135040283203 + ], + "size": [ + 422.84503173828125, + 164.31304931640625 + ], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 10 + }, + { + "name": "text", + "type": "STRING", + "link": 242, + "widget": { + "name": "text" + } + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 60 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "an oil painting of a duck" + ], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 43, + "type": "FluxGuidance", + "pos": [ + 419.49072265625, + -168.48008728027344 + ], + "size": [ + 320, + 58 + ], + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 60 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 61 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "FluxGuidance" + }, + "widgets_values": [ + 3 + ] + }, + { + "id": 22, + "type": "BasicGuider", + "pos": [ + 449.3519287109375, + 1.1387512683868408 + ], + "size": [ + 241.79998779296875, + 46 + ], + "flags": {}, + "order": 19, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 39, + "slot_index": 0 + }, + { + "name": "conditioning", + "type": "CONDITIONING", + "link": 61, + "slot_index": 1 + } + ], + "outputs": [ + { + "name": "GUIDER", + "type": "GUIDER", + "links": [ + 30 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "BasicGuider" + }, + "widgets_values": [] + }, + { + "id": 16, + "type": "KSamplerSelect", + "pos": [ + 470, + 749 + ], + "size": [ + 315, + 58 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "SAMPLER", + "type": "SAMPLER", + "links": [ + 19 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "KSamplerSelect" + }, + "widgets_values": [ + "dpmpp_2m" + ] + }, + { + "id": 12, + "type": "UNETLoader", + "pos": [ + -62.69803237915039, + 114.85723876953125 + ], + "size": [ + 315, + 82 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 38, + 39 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "UNETLoader" + }, + "widgets_values": [ + "flux1-schnell-fp8.safetensors", + "default" + ], + "color": "#323", + "bgcolor": "#535" + }, + { + "id": 11, + "type": "DualCLIPLoader", + "pos": [ + 101.36359405517578, + 490.7603759765625 + ], + "size": [ + 315, + 106 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 10 + ], + "slot_index": 0, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "DualCLIPLoader" + }, + "widgets_values": [ + "clip_l.safetensors", + "t5xxl_fp16.safetensors", + "flux", + "default" + ] + }, + { + "id": 31, + "type": "String-🔬", + "pos": [ + -1577.001220703125, + 421.1225891113281 + ], + "size": [ + 463.93511962890625, + 183.8080596923828 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": [ + 236 + ], + "slot_index": 0 + } + ], + "title": "LLM Prompt", + "properties": { + "Node name for S&R": "String-🔬" + }, + "widgets_values": [ + "Give me a detailed description of a high definition professional photographic image of a duck. The duck is a real duck, not a cartoon. It is on the table \ninside a sunny cafe next to a steaming cup of coffee. \n\nDescribe everything in detail in a 200 or so word prompt.\n\n" + ], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 38, + "type": "ShowText|pysssss", + "pos": [ + -602.6541748046875, + -407.1946105957031 + ], + "size": [ + 498.1600341796875, + 461.8752136230469 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "text", + "type": "STRING", + "link": 234, + "widget": { + "name": "text" + } + } + ], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": [ + 241 + ], + "shape": 6, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ShowText|pysssss" + }, + "widgets_values": [ + "", + "Create a vivid photographic image description of a high-resolution shot for a café interior. The focus is on a realistic, life-size duck sitting confidently on a sturdy ceramic plate placed upon a round, well-lit wooden table. This table sits near a large, steaming cappuccino or espresso coffee cup, its rich brown color contrasting elegantly with the white ceramic mug and the rustic wood surface. Sunlight streams through a nearby window, casting warm golden rays onto the scene. The duck's feathers are meticulously rendered in detail, showcasing a variety of shades from iridescent greens to subtle browns, reflecting the ambient light as if it were a miniature watercolor painting. The café itself is depicted with attention to architectural details such as gleaming brass accents on the bar, polished marble countertops, and a well-stocked display case filled with pastries and bread baskets. Soft shadows from the window frame accentuate the scene, creating depth and dimensionality. The coffee cup and the duck are the focal points of this composition, with the duck's head slightly tilted towards the viewer, as if it were observing the scene with an air of curiosity or mild amusement. The overall atmosphere is cozy yet sophisticated, capturing the essence of a relaxed European-style café." + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 108, + "type": "EbuLMStudioUnload", + "pos": [ + -57.53535842895508, + -196.06173706054688 + ], + "size": [ + 315, + 106 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "input_string", + "type": "STRING", + "link": 241, + "widget": { + "name": "input_string" + } + } + ], + "outputs": [ + { + "name": "export_string", + "type": "STRING", + "links": [ + 242 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "EbuLMStudioUnload" + }, + "widgets_values": [ + "", + 962001017751484, + "randomize" + ], + "color": "#222", + "bgcolor": "#000" + }, + { + "id": 105, + "type": "String-🔬", + "pos": [ + -48.6381721496582, + -455.9993896484375 + ], + "size": [ + 400, + 200 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": null + } + ], + "properties": { + "Node name for S&R": "String-🔬" + }, + "widgets_values": [ + "STEP THREE\n\nWe unload the LLM to free up memory and pass the generated prompt to Flux." + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 104, + "type": "String-🔬", + "pos": [ + -1128.3271484375, + -423.2717590332031 + ], + "size": [ + 400, + 200 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": null + } + ], + "properties": { + "Node name for S&R": "String-🔬" + }, + "widgets_values": [ + "STEP TWO\n\nWe query LMStudio with our request." + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 103, + "type": "String-🔬", + "pos": [ + -909.1731567382812, + 448.65771484375 + ], + "size": [ + 400, + 200 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": null + } + ], + "properties": { + "Node name for S&R": "String-🔬" + }, + "widgets_values": [ + "STEP ONE\n\nWe pass our prompt through the Loader to trigger the LLM load. It frees up VRAM by unloading both LLMs and image models." + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 106, + "type": "EbuLMStudioMakeRequest", + "pos": [ + -1097.0440673828125, + -152.1427459716797 + ], + "size": [ + 400, + 318 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "prompt", + "type": "STRING", + "link": 237, + "widget": { + "name": "prompt" + } + } + ], + "outputs": [ + { + "name": "generated_text", + "type": "STRING", + "links": [ + 234 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "EbuLMStudioMakeRequest" + }, + "widgets_values": [ + "", + "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.", + "http://127.0.0.1:1234/v1/chat/completions", + 4096, + 324960009893632, + "randomize", + 300, + 0.7000000000000001, + 0.9500000000000001, + false + ], + "color": "#222", + "bgcolor": "#000" + }, + { + "id": 109, + "type": "ShowText|pysssss", + "pos": [ + -359.1662292480469, + 731.47705078125 + ], + "size": [ + 472.24285888671875, + 366.3433837890625 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "text", + "type": "STRING", + "link": 243, + "widget": { + "name": "text" + } + } + ], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": null, + "shape": 6 + } + ], + "properties": { + "Node name for S&R": "ShowText|pysssss" + }, + "widgets_values": [ + "" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 110, + "type": "ShowText|pysssss", + "pos": [ + -354.2527160644531, + 456.302001953125 + ], + "size": [ + 403.4491271972656, + 191.902099609375 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "text", + "type": "STRING", + "link": 245, + "widget": { + "name": "text" + } + } + ], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": null, + "shape": 6 + } + ], + "properties": { + "Node name for S&R": "ShowText|pysssss" + }, + "widgets_values": [ + "" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 107, + "type": "EbuLMStudioLoadModel", + "pos": [ + -874.427001953125, + 701.475830078125 + ], + "size": [ + 315, + 218 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "input_string", + "type": "STRING", + "link": 236, + "widget": { + "name": "input_string" + } + } + ], + "outputs": [ + { + "name": "export_string", + "type": "STRING", + "links": [ + 237 + ], + "slot_index": 0 + }, + { + "name": "model_selected", + "type": "STRING", + "links": [ + 245 + ], + "slot_index": 1 + }, + { + "name": "models_found", + "type": "STRING", + "links": [ + 243 + ], + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "EbuLMStudioLoadModel" + }, + "widgets_values": [ + "", + "Llama-3.2", + 4096, + 1077434342249168, + "randomize", + false + ], + "color": "#222", + "bgcolor": "#000" + } + ], + "links": [ + [ + 10, + 11, + 0, + 6, + 0, + "CLIP" + ], + [ + 12, + 10, + 0, + 8, + 1, + "VAE" + ], + [ + 19, + 16, + 0, + 13, + 2, + "SAMPLER" + ], + [ + 20, + 17, + 0, + 13, + 3, + "SIGMAS" + ], + [ + 23, + 5, + 0, + 13, + 4, + "LATENT" + ], + [ + 24, + 13, + 0, + 8, + 0, + "LATENT" + ], + [ + 30, + 22, + 0, + 13, + 1, + "GUIDER" + ], + [ + 37, + 25, + 0, + 13, + 0, + "NOISE" + ], + [ + 38, + 12, + 0, + 17, + 0, + "MODEL" + ], + [ + 39, + 12, + 0, + 22, + 0, + "MODEL" + ], + [ + 41, + 8, + 0, + 26, + 0, + "IMAGE" + ], + [ + 60, + 6, + 0, + 43, + 0, + "CONDITIONING" + ], + [ + 61, + 43, + 0, + 22, + 1, + "CONDITIONING" + ], + [ + 234, + 106, + 0, + 38, + 0, + "STRING" + ], + [ + 236, + 31, + 0, + 107, + 0, + "STRING" + ], + [ + 237, + 107, + 0, + 106, + 0, + "STRING" + ], + [ + 241, + 38, + 0, + 108, + 0, + "STRING" + ], + [ + 242, + 108, + 0, + 6, + 1, + "STRING" + ], + [ + 243, + 107, + 2, + 109, + 0, + "STRING" + ], + [ + 245, + 107, + 1, + 110, + 0, + "STRING" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.41772481694161123, + "offset": [ + 1145.2442489625016, + 417.2069241225013 + ] + }, + "node_versions": { + "comfy-core": "0.3.12", + "ComfyUI-Logic": "42d4f3df45fb7f0dd6e2201a14c07d4dd09f235d", + "ComfyUI-Custom-Scripts": "bc8922deff73f59311c05cef27b9d4caaf43e87b" + }, + "ue_links": [] + }, + "version": 0.4 +} \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..03eaa3b --- /dev/null +++ b/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", +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..994ebb6 --- /dev/null +++ b/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" \ No newline at end of file diff --git a/readme.md b/readme.md new file mode 100644 index 0000000..573c588 --- /dev/null +++ b/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. diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..b8b0b57 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +# No additional requirements beyond Python standard library \ No newline at end of file