Merge pull request #12 from EnragedAntelope/claude/fix-lmstudio-model-refresh-VIoV5

Add instant model refresh with frontend extension

Author: EnragedAntelope

LMStudio.py

@@ -5,6 +5,7 @@
 import logging
 import re
 from typing import Optional, Tuple, List
+import os
 from tempfile import NamedTemporaryFile
 import numpy as np
 from PIL import Image
@@ -220,7 +221,7 @@ def INPUT_TYPES(cls):
                 }),
                 "refresh_models": ("BOOLEAN", {
                     "default": False,
-                    "tooltip": "Re-fetch model list from LM Studio server. Enable, queue once, then disable. Updates dropdown on next node load."
+                    "tooltip": "Toggle ON to re-fetch the model list from LM Studio and update the dropdowns instantly. Automatically toggles back off."
                 }),
             }
         }
@@ -468,7 +469,6 @@ def generate(
             success, message = refresh_model_cache(server_url, timeout)
             if success:
                 troubleshooting_lines.append(f"[INFO] Model refresh: {message}")
-                troubleshooting_lines.append("[INFO] Dropdown will update on next node load")
             else:
                 troubleshooting_lines.append(f"[WARNING] Model refresh failed: {message}")
 
@@ -549,12 +549,21 @@ def generate(
                 if pil_images:
                     # Prepare all images for the SDK
                     image_handles = []
-                    for pil_img in pil_images:
-                        with NamedTemporaryFile(suffix=".jpg", delete=False) as temp:
-                            pil_img.save(temp, format="JPEG", quality=95)
-                            temp.flush()
-                            image_handle = client.files.prepare_image(temp.name)
-                            image_handles.append(image_handle)
+                    temp_paths = []
+                    try:
+                        for pil_img in pil_images:
+                            with NamedTemporaryFile(suffix=".jpg", delete=False) as temp:
+                                pil_img.save(temp, format="JPEG", quality=95)
+                                temp.flush()
+                                temp_paths.append(temp.name)
+                                image_handle = client.files.prepare_image(temp.name)
+                                image_handles.append(image_handle)
+                    finally:
+                        for path in temp_paths:
+                            try:
+                                os.unlink(path)
+                            except OSError:
+                                pass
 
                     chat.add_user_message(prompt, images=image_handles)
                 else:

README.md

@@ -38,7 +38,7 @@ pip install -r EA_LMStudio/requirements.txt
 
 ## Tips
 
-- **Models not showing?** LM Studio must be running before ComfyUI starts. Use `refresh_models` checkbox to re-fetch.
+- **Models not showing?** LM Studio must be running before ComfyUI starts. Toggle the `refresh_models` checkbox to instantly re-fetch and update the dropdowns.
 - **Context errors?** Increase context length in LM Studio settings (not max_tokens).
 - **VLM issues?** Try a smaller image resize option or single image if multi-image fails.
 - **Force thinking mode:** Add `/think` to prompts for Qwen3, or "Think step by step" for others.

__init__.py

@@ -5,4 +5,31 @@
 """
 from .LMStudio import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
 
-__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+# Serve web extension for frontend model refresh
+WEB_DIRECTORY = "./web"
+
+# Register server API route for model refresh
+from aiohttp import web
+from server import PromptServer
+
+from .model_fetcher import refresh_model_cache, get_model_choices, CUSTOM_MODEL_OPTION
+from .lms_config.config_manager import ConfigManager
+
+@PromptServer.instance.routes.post("/ea_lmstudio/refresh_models")
+async def _refresh_models_route(request):
+    """API endpoint to refresh model list from LM Studio server."""
+    config_manager = ConfigManager()
+    server_url = config_manager.get_server_url()
+    timeout = config_manager.get_timeout()
+
+    success, message = refresh_model_cache(server_url, timeout)
+    choices = get_model_choices()
+    models = [m for m in choices if m != CUSTOM_MODEL_OPTION]
+
+    return web.json_response({
+        "success": success,
+        "message": message,
+        "models": models,
+    })
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "EA_LMStudio"
 description = "A custom node for LM Studio integration into ComfyUI."
-version = "1.2.2"
+version = "1.2.3"
 license = {file = "LICENSE"}
 
 [project.urls]

web/ea_lmstudio.js

@@ -0,0 +1,73 @@
+import { app } from "../../scripts/app.js";
+import { api } from "../../scripts/api.js";
+
+/*
+ * EA_LMStudio Model Refresh Extension
+ *
+ * Intercepts the refresh_models toggle to fetch an updated model list from
+ * the LM Studio server and update the dropdown widgets in-place.
+ *
+ * Compatibility:
+ *   - Legacy LiteGraph frontend: full support via widget.options.values
+ *   - Nodes 2.0 / Vue frontend: the server-side cache is always updated,
+ *     so a browser refresh (F5) will pick up new models even if the
+ *     in-place widget update doesn't propagate in a future Vue renderer.
+ */
+
+app.registerExtension({
+    name: "EA_LMStudio.ModelRefresh",
+
+    async nodeCreated(node) {
+        if (node.comfyClass !== "EA_LMStudio") return;
+
+        const refreshWidget = node.widgets?.find(w => w.name === "refresh_models");
+        if (!refreshWidget) return;
+
+        const originalCallback = refreshWidget.callback;
+
+        refreshWidget.callback = async function (value) {
+            if (!value) {
+                if (originalCallback) originalCallback.call(this, value);
+                return;
+            }
+
+            try {
+                const resp = await api.fetchApi("/ea_lmstudio/refresh_models", {
+                    method: "POST",
+                });
+                const data = await resp.json();
+
+                if (data.success && data.models) {
+                    const choices = ["-- Custom (enter below) --", ...data.models];
+
+                    for (const widgetName of ["model_selection", "draft_model_selection"]) {
+                        const w = node.widgets?.find(ww => ww.name === widgetName);
+                        if (w && w.options) {
+                            // Replace values array (breaks shared reference intentionally
+                            // so other node instances also get the update)
+                            w.options.values = choices;
+                            if (!choices.includes(w.value)) {
+                                w.value = choices[0];
+                            }
+                        }
+                    }
+
+                    console.log(
+                        `[EA_LMStudio] Refreshed models: ${data.models.length} found`
+                    );
+                } else {
+                    console.warn(
+                        `[EA_LMStudio] Model refresh failed: ${data.message || "unknown error"}`
+                    );
+                }
+            } catch (err) {
+                console.error("[EA_LMStudio] Failed to refresh models:", err);
+            }
+
+            // Toggle back off so it acts like a one-shot button
+            refreshWidget.value = false;
+
+            if (originalCallback) originalCallback.call(this, false);
+        };
+    },
+});