diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..62fd899 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,31 @@ +name: Tests + +on: + push: + branches: [main, master] + pull_request: + workflow_dispatch: + +jobs: + pytest: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + # Floor matches requires-python; the newest release guards against + # deprecations landing in a future ComfyUI runtime. + python-version: ["3.10", "3.13"] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + - name: Install test dependencies + # requirements-dev.txt pulls in requirements.txt (requests, Pillow, + # numpy, lmstudio). The suite stubs lmstudio and comfy when they are + # missing, so it also has to pass without a running LM Studio - the + # startup model fetch simply fails to connect and is handled. + run: pip install -r requirements-dev.txt + - name: Run tests + run: pytest -q diff --git a/.gitignore b/.gitignore index cf6d177..01683c7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,9 @@ # User configuration (survives updates) lms_config/user_config.json +# Local development worklogs / scratch notes (never published) +worklogs/ + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..4ad1d29 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,121 @@ +# EA_LMStudio — agent notes + +A single ComfyUI custom node (`EA_LMStudio`, displayed as **EA LM Studio**) that +runs text and vision generation against a local [LM Studio](https://lmstudio.ai/) +server through the official `lmstudio` Python SDK. It returns three STRING +outputs — `response`, `reasoning`, `troubleshooting` — and renders the response +inside the node. + +## Current state + +_Last verified: 2026-08-09_ + +- **Status:** v2.0.0, published to the Comfy Registry from `main` on every + `pyproject.toml` version change. +- **Works:** model discovery + refresh (startup fetch and a live API route), + text generation, multi-image VLM input, reasoning extraction (LM Studio's own + split, with tag-regex fallback), structured JSON output, stop strings, + context-overflow policy, speculative decoding with acceptance stats, + cancellable streaming with a queue progress bar, VRAM unload of both the LLM + and ComfyUI's own models, and migration of workflows saved by 1.x. +- **In progress:** nothing outstanding. +- **Known gaps:** LM Studio's `ttl`, tool/function calling (`.act()`), GBNF + grammars, and load-time `seed` (real determinism, but it only applies when the + node actually loads the model) are all supported by the SDK and not exposed. + There is no automated frontend test — `web/ea_lmstudio.js` is verified by hand + in a browser. +- **Deep docs:** user-facing behaviour lives in `README.md`; nothing is + duplicated here. + +## Build / test / run + +```bash +pip install -r requirements-dev.txt # pulls requirements.txt too +pytest -q # 100+ tests, no LM Studio needed +``` + +The suite is dependency-light by design. `tests/conftest.py` stubs `lmstudio` +and `comfy.model_management` **only when they are genuinely missing**, so a real +install is never shadowed, and registers the repo root as a synthetic package so +`LMStudio.py` (which uses relative imports) is importable without the checkout +having to be named `EA_LMStudio`. + +CI runs the same command on Python 3.10 and 3.13 (`.github/workflows/test.yml`). + +To exercise the node for real, point a scratch script at a running LM Studio and +stub `comfy.model_management` / `comfy.utils` — the node only needs +`unload_all_models`, `soft_empty_cache`, `processing_interrupted`, +`throw_exception_if_processing_interrupted`, `InterruptProcessingException` and +`ProgressBar`. + +## Layout + +| File | Responsibility | +|------|----------------| +| `LMStudio.py` | The node: INPUT_TYPES, streaming, diagnostics, reasoning split | +| `lms_params.py` | Pure widget-string → SDK-value helpers (stop strings, schema, fences) | +| `lms_reasoning.py` | Tag/harmony reasoning regexes (fallback path) | +| `lms_image.py` | ComfyUI IMAGE tensor → JPEG-safe PIL | +| `model_fetcher.py` | `/v1/models` discovery, validation, cache | +| `lms_config/` | `default_config.json` + gitignored `user_config.json` | +| `web/ea_lmstudio.js` | Refresh toggle, in-node preview, 1.x workflow migration | +| `example_workflows/` | Shipped examples, loadable by drag-and-drop | + +Everything except `LMStudio.py` is deliberately free of the `lmstudio` SDK and +`comfy` imports so it stays unit-testable. + +## Things that will bite you + +**The SDK silently discards unknown prediction-config keys.** `LlmPredictionConfig` +is built with msgspec and drops anything not in `LlmPredictionConfigDict` rather +than raising, so a wrong or wished-for key becomes a no-op that looks like it +worked. v1.x shipped `presencePenalty` and `enableThinking` this way for +releases. Before adding a parameter, check it exists in `LlmPredictionConfigDict` +in the installed SDK, then confirm it round-trips in +`PredictionResult.prediction_config`. `generate()` performs that diff on every +run and warns — do not remove it. + +**LM Studio has no presence penalty, no frequency penalty, no inference-time +seed, and no thinking on/off flag.** `seed` exists only in the *load* config. +The `seed` widget is a ComfyUI cache-buster and nothing more. + +**Removing or reordering a widget corrupts saved workflows.** ComfyUI serialises +`widgets_values` positionally, so a removal or a regroup shifts every later +value. The migration in `web/ea_lmstudio.js` keys the old array by the v1.5.x +widget-name order (two variants, with and without the `control_after_generate` +widget ComfyUI inserts after an INT named `seed`) and writes values onto current +widgets **by name**, which is why v2.0.0 could both drop two widgets and regroup +the rest. Any future removal or reorder needs the same treatment, and the legacy +order table must be kept. + +**A stored node size is restored verbatim and is not re-checked against the +widgets.** Adding a widget therefore leaves every previously saved workflow too +short, and the overflow draws outside the node frame. `growToFitWidgets` in the +frontend extension grows (never shrinks) the node on configure and after +execution. Note that a stock ComfyUI `Note` node's textarea overhangs its own +frame by ~13 units at any size — that is upstream behaviour, not a symptom of +this, so don't chase it. + +**A streamed prediction must be drained, not broken out of.** Breaking the `for` +loop closes the generator and `stream.result()` then raises `GeneratorExit`. +Call `stream.cancel()` and keep iterating; it ends promptly with +`stop_reason == "userStopped"`. + +**Every `LlmPredictionStats` field except `stop_reason` is Optional.** Formatting +one with `:.2f` without a None check raised `TypeError` *after* a successful +generation, which the outer handler then reported as a failure — throwing away +text the model had already produced. + +**`{"type": "json"}` without a schema does not constrain decoding.** Models +routinely answer with a ```` ```json ```` fence. Only `jsonSchema` constrains the +sampler. + +## Conventions + +- Registry publishing is driven by the `version` in `pyproject.toml`; the + workflow compares it against `HEAD^` and skips when unchanged. Bump it in the + same commit as any change worth shipping. +- `lms_config/user_config.json` and `worklogs/` are gitignored. Never commit + either, and never put a server address or token in a tracked file. +- Keep the `CUSTOM_MODEL_OPTION` literal in `web/ea_lmstudio.js` in sync with + `model_fetcher.py`, which is the source of truth. diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..43c994c --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/LMStudio.py b/LMStudio.py index 4d3e22b..ab30dbf 100644 --- a/LMStudio.py +++ b/LMStudio.py @@ -3,7 +3,7 @@ Provides text generation using local LLM/VLM models via LM Studio server. """ import logging -from typing import Optional, Tuple, List +from typing import Any, Dict, List, Optional, Tuple import os import time from tempfile import NamedTemporaryFile @@ -15,16 +15,23 @@ # ComfyUI imports import comfy.model_management as model_management +try: # Optional: only used to drive the queue progress bar. + from comfy.utils import ProgressBar +except Exception: # pragma: no cover - ComfyUI always provides this at runtime + ProgressBar = None + # Local imports from .lms_config.config_manager import ConfigManager from .model_fetcher import ( get_model_choices, + get_default_model_choice, refresh_model_cache, initialize_model_cache, validate_model_identifier, get_last_fetch_error, get_last_fetch_success, get_cached_model_count, + get_last_rejected_models, CUSTOM_MODEL_OPTION, ) from .lms_reasoning import ( @@ -32,6 +39,15 @@ extract_reasoning_custom, looks_like_leaked_thinking, ) +from .lms_params import ( + CONTEXT_OVERFLOW_OPTIONS, + CONTEXT_OVERFLOW_POLICIES, + OUTPUT_FORMAT_OPTIONS, + build_structured_setting, + missing_config_keys, + parse_stop_strings, + strip_json_code_fence, +) from .lms_image import convert_image_to_pil # Setup logging @@ -98,6 +114,15 @@ def _quiet_dependency_logs() -> None: excluded_patterns=_config_manager.get_excluded_patterns(_startup_config), ) +# ComfyUI's cancel button sets an interrupt flag. We poll it while streaming so a +# runaway generation can actually be stopped, then re-raise so the queue reports +# a cancellation rather than a node error. The class is looked up defensively so +# the module still imports under a stubbed ``comfy`` (tests, registry scanners). +_INTERRUPT_EXCEPTION = getattr(model_management, "InterruptProcessingException", None) +_INTERRUPT_EXCEPTIONS: Tuple[type, ...] = ( + (_INTERRUPT_EXCEPTION,) if isinstance(_INTERRUPT_EXCEPTION, type) else () +) + # Image resize options IMAGE_RESIZE_OPTIONS = [ "No Resize", @@ -123,19 +148,30 @@ def _quiet_dependency_logs() -> None: "Custom tags", ] -# enable_thinking toggle options. "Model default" leaves the model's own -# behavior untouched (nothing is sent); the other two force thinking on/off -# for hybrid reasoning models (e.g. Qwen3) via the SDK's enableThinking flag. -ENABLE_THINKING_OPTIONS = [ - "Model default", - "Enabled", - "Disabled", -] - # Reasoning/thinking extraction (regexes + helpers) lives in lms_reasoning.py # so the pure text-processing logic stays importable and unit-testable without # pulling in the lmstudio SDK, ComfyUI, or the startup network fetch. +# ``draftModel`` is excluded from the applied-config check: LM Studio only echoes +# speculative-decoding settings back when the draft model was actually accepted, +# and we have no way to distinguish "not echoed" from "not applied" here. A false +# "this did nothing" warning would be worse than no warning. +_UNVERIFIABLE_CONFIG_KEYS = frozenset({"draftModel"}) + +# Keep long values (stop strings, JSON schemas) from flooding the config summary. +_CONFIG_SUMMARY_VALUE_LIMIT = 120 + + +def _summarize_config(gen_config: Dict[str, Any]) -> str: + """Render the generation config as one readable line.""" + parts = [] + for key, value in gen_config.items(): + text = repr(value) if isinstance(value, (str, list, dict)) else str(value) + if len(text) > _CONFIG_SUMMARY_VALUE_LIMIT: + text = text[:_CONFIG_SUMMARY_VALUE_LIMIT] + "...(truncated)" + parts.append(f"{key}={text}") + return ", ".join(parts) + class EALMStudio: """ @@ -148,13 +184,25 @@ class EALMStudio: CATEGORY = "EA/LMStudio" RETURN_TYPES = ("STRING", "STRING", "STRING") RETURN_NAMES = ("response", "reasoning", "troubleshooting") + # OUTPUT_NODE lets the node run as a graph terminal (so it can be queued + # without wiring its outputs anywhere) AND carries the ``ui`` payload that + # renders the response inside the node. Both halves matter: without the + # payload the flag would only cost a wasted inference on disconnected nodes. OUTPUT_NODE = True FUNCTION = "generate" + DESCRIPTION = ( + "Generate text with a local LM Studio model. Supports vision models, " + "reasoning extraction, structured JSON output and VRAM management." + ) @classmethod def INPUT_TYPES(cls): model_choices = get_model_choices() - default_model = model_choices[0] if model_choices else CUSTOM_MODEL_OPTION + # Default to a real model when discovery worked, so a freshly dropped + # node is runnable instead of landing on the "Custom" sentinel with an + # empty identifier. The draft dropdown keeps the sentinel: speculative + # decoding must stay opt-in. + default_model = get_default_model_choice() return { "required": { @@ -197,36 +245,15 @@ def INPUT_TYPES(cls): "default": 0, "min": 0, "max": 0xffffffffffffffff, - "tooltip": "Seed for ComfyUI workflow reproducibility. Note: LM Studio SDK does not support inference-time seeding." + "tooltip": "Re-roll control only. LM Studio has no inference-time seed, so this does NOT make output reproducible - changing it simply tells ComfyUI the node is dirty so it generates again instead of reusing the cached response. Set control_after_generate to 'randomize' for a fresh answer every queue, or 'fixed' to keep the cached one." }), }, + # Widget order below is the on-node layout. Grouped most-used first: + # sampling -> output shaping -> reasoning -> vision -> speculative + # decoding -> management. Each group's dependent fields follow the + # control that switches them on, and every dependent field's tooltip + # names that control, so the layout reads top-to-bottom. "optional": { - # --- Image inputs (for VLMs) --- - "image_resize": (IMAGE_RESIZE_OPTIONS, { - "default": "Medium (768px)", - "tooltip": "Resize images before processing. Smaller = faster inference. 'No Resize' keeps original size. Only applies when images are connected." - }), - "image1": ("IMAGE", { - "tooltip": "First image input for vision models (VLMs). Leave unconnected for text-only inference." - }), - "image2": ("IMAGE", { - "tooltip": "Second image input for multi-image VLMs. Not all VLMs support multiple images." - }), - "image3": ("IMAGE", { - "tooltip": "Third image input for multi-image VLMs. Not all VLMs support multiple images." - }), - "image4": ("IMAGE", { - "tooltip": "Fourth image input for multi-image VLMs. Not all VLMs support multiple images." - }), - # --- Advanced model options --- - "draft_model_selection": (model_choices, { - "default": default_model, - "tooltip": "Optional draft model for speculative decoding (faster inference). Select 'Custom' and leave empty to disable." - }), - "custom_draft_model": ("STRING", { - "default": "", - "tooltip": "Manual draft model identifier. Only used when draft 'Custom' is selected. Leave empty to disable." - }), # --- Sampling parameters --- "top_p": ("FLOAT", { "default": 1.0, @@ -247,7 +274,7 @@ def INPUT_TYPES(cls): "min": 0.0, "max": 2.0, "step": 0.05, - "tooltip": "Penalizes tokens that already appeared, scaled by how often (default 1.0 = disabled). Raising (1.1-1.3) reduces repetition/loops; too high can hurt coherence. Below 1.0 encourages repetition." + "tooltip": "Penalizes tokens that already appeared, scaled by how often (default 1.0 = disabled). Raising (1.1-1.3) reduces repetition/loops; too high can hurt coherence. Below 1.0 encourages repetition. LM Studio has no presence or frequency penalty - this and min_p are the repetition controls it offers." }), "min_p": ("FLOAT", { "default": 0.0, @@ -256,21 +283,29 @@ def INPUT_TYPES(cls): "step": 0.01, "tooltip": "Min-P sampling: drop tokens below this fraction of the top token's probability (default 0.0 = disabled). Raising (e.g. 0.05-0.1) = more focused/coherent; lowering toward 0 = more diverse. A modern alternative to top_p." }), - "presence_penalty": ("FLOAT", { - "default": 0.0, - "min": -2.0, - "max": 2.0, - "step": 0.05, - "tooltip": "Flat penalty on any token already used, encouraging new topics (default 0.0 = disabled). Raising (e.g. 0.3-0.8) reduces repetition / broadens topics; negative values encourage reuse. Distinct from repeat_penalty. Note: LM Studio has no frequency_penalty." + # --- Output shaping --- + "stop_strings": ("STRING", { + "multiline": True, + "default": "", + "tooltip": "Stop generation when any of these strings appears. One per line; blank lines ignored. Leading/trailing spaces are kept, and \\n \\r \\t \\\\ are expanded - so a line of '\\nUser:' stops at a newline followed by 'User:'. Empty = no stop strings. Useful to stop a chatty model running on past the answer." + }), + "context_overflow": (CONTEXT_OVERFLOW_OPTIONS, { + "default": "Truncate middle", + "tooltip": "What LM Studio does when prompt + response exceed the model's context window. 'Truncate middle' (default) silently drops the middle of the conversation. 'Rolling window' drops from the start. 'Stop at limit (error)' fails loudly instead - pick it if a silently shortened prompt would be worse than no answer." + }), + "output_format": (OUTPUT_FORMAT_OPTIONS, { + "default": "Text", + "tooltip": "'Text' = normal prose. 'JSON (schema below)' constrains decoding to the schema in json_schema and is the reliable choice when a downstream node must parse the response. 'JSON (no schema)' only asks for JSON - it does NOT constrain decoding, and many models answer with a ```json fenced block (which is unwrapped automatically when the contents are valid JSON). Structured output and thinking models mix poorly." }), - "enable_thinking": (ENABLE_THINKING_OPTIONS, { - "default": "Model default", - "tooltip": "Force thinking/reasoning on hybrid models like Qwen3 without the '/think' prompt hack (default 'Model default' = leave the model's own behavior untouched). 'Enabled' turns thinking on; 'Disabled' turns it off. Pairs with reasoning_mode. Ignored by models/backends that don't support it." + "json_schema": ("STRING", { + "multiline": True, + "default": "", + "tooltip": "JSON Schema object, used only when output_format is 'JSON (schema below)'. Example: {\"type\": \"object\", \"properties\": {\"caption\": {\"type\": \"string\"}}, \"required\": [\"caption\"]}" }), # --- Reasoning extraction --- "reasoning_mode": (REASONING_MODE_OPTIONS, { "default": "Auto-detect (recommended)", - "tooltip": "How to extract reasoning/thinking from model output. Auto-detect works with DeepSeek, Qwen, QwQ, GLM, GPT-OSS and similar models. Note: Models don't always produce thinking output for simple queries. For Qwen3, add '/think' to your prompt to force thinking mode." + "tooltip": "How to split thinking from the final answer. When LM Studio's own Reasoning Parsing is configured for the model, its tagging is used directly and this setting is not needed. Otherwise Auto-detect handles DeepSeek, Qwen, QwQ, GLM, GPT-OSS and similar tag formats. Models don't always think for simple queries." }), "custom_open_tag": ("STRING", { "default": "", @@ -280,10 +315,36 @@ def INPUT_TYPES(cls): "default": "", "tooltip": "Custom closing tag for reasoning extraction. Only used when reasoning_mode is 'Custom tags'." }), + # --- Vision (only relevant when an image input is connected) --- + "image_resize": (IMAGE_RESIZE_OPTIONS, { + "default": "Medium (768px)", + "tooltip": "Resize images before processing. Smaller = faster inference. 'No Resize' keeps original size. Only applies when images are connected." + }), + "image1": ("IMAGE", { + "tooltip": "First image input for vision models (VLMs). Leave unconnected for text-only inference." + }), + "image2": ("IMAGE", { + "tooltip": "Second image input for multi-image VLMs. Not all VLMs support multiple images." + }), + "image3": ("IMAGE", { + "tooltip": "Third image input for multi-image VLMs. Not all VLMs support multiple images." + }), + "image4": ("IMAGE", { + "tooltip": "Fourth image input for multi-image VLMs. Not all VLMs support multiple images." + }), + # --- Speculative decoding --- + "draft_model_selection": (model_choices, { + "default": CUSTOM_MODEL_OPTION, + "tooltip": "Optional draft model for speculative decoding (faster inference). Must share a tokenizer with the main model. Leave on 'Custom' with an empty box to disable. Acceptance stats are reported in troubleshooting." + }), + "custom_draft_model": ("STRING", { + "default": "", + "tooltip": "Manual draft model identifier. Only used when draft 'Custom' is selected. Leave empty to disable." + }), # --- Management --- "unload_llm": ("BOOLEAN", { "default": True, - "tooltip": "Unload the LLM from LM Studio after generation. Recommended to free VRAM for image generation." + "tooltip": "Unload the LLM from LM Studio after generation. Recommended to free VRAM for image generation. Turn off to keep the model warm across runs (this also unloads a model you loaded by hand in LM Studio)." }), "unload_comfy_models": ("BOOLEAN", { "default": False, @@ -308,6 +369,30 @@ def IS_CHANGED(cls, **kwargs): return float("nan") # Always different return "" + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + @staticmethod + def _output( + response: str, reasoning: str, troubleshooting_lines: List[str] + ) -> Dict[str, Any]: + """Build the node return value. + + Because this is an OUTPUT_NODE, the ``ui`` half is what makes the flag + worth having: the response is rendered inside the node itself, so the + common "just enhance my prompt" workflow needs no extra preview node. + """ + troubleshooting = "\n".join(troubleshooting_lines) + return { + "ui": { + "text": [response], + "reasoning": [reasoning], + "troubleshooting": [troubleshooting], + }, + "result": (response, reasoning, troubleshooting), + } + def _resolve_model_identifier( self, selection: str, @@ -336,6 +421,94 @@ def _resolve_model_identifier( return model_id, None + @staticmethod + def _prepare_images(client, pil_images: List[Image.Image]) -> List[Any]: + """Upload PIL images to LM Studio and return its file handles. + + Each image is written to a temp file, closed, uploaded, then deleted. + The file is closed before uploading because on Windows a second reader + of a still-open handle is not guaranteed, and deletion of an open file + fails outright. + """ + handles = [] + for pil_img in pil_images: + temp_path = None + try: + with NamedTemporaryFile(suffix=".jpg", delete=False) as temp: + temp_path = temp.name + pil_img.save(temp, format="JPEG", quality=95) + handles.append(client.files.prepare_image(temp_path)) + finally: + if temp_path: + try: + os.unlink(temp_path) + except OSError: + pass + return handles + + @staticmethod + def _stats_lines(stats, elapsed: float) -> List[str]: + """Format inference statistics, tolerating fields the backend omits. + + Every count on LlmPredictionStats except stop_reason is Optional, so a + bare f-string format spec (e.g. ``:.2f`` on None) raises TypeError. That + used to surface as "Generation failed" *after* a successful generation, + discarding the text the model had already produced. + """ + lines = [] + + def value(name, default=None): + return getattr(stats, name, default) + + tokens_per_sec = value("tokens_per_second") + if tokens_per_sec is not None: + lines.append(f"[INFO] Tokens per second: {tokens_per_sec:.2f}") + + for label, attr in ( + ("Input tokens", "prompt_tokens_count"), + ("Output tokens", "predicted_tokens_count"), + ("Total tokens", "total_tokens_count"), + ): + count = value(attr) + if count is not None: + lines.append(f"[INFO] {label}: {count}") + + ttft = value("time_to_first_token_sec") + if ttft is not None: + lines.append(f"[INFO] Time to first token: {ttft:.3f}s") + + gpu_layers = value("num_gpu_layers") + if gpu_layers is not None: + lines.append(f"[INFO] GPU layers: {gpu_layers:g}") + + # Speculative decoding: without acceptance numbers there is no way to + # tell whether a draft model is helping or just burning time. + drafted = value("total_draft_tokens_count") + if drafted: + accepted = value("accepted_draft_tokens_count") or 0 + rejected = value("rejected_draft_tokens_count") or 0 + rate = (accepted / drafted * 100.0) if drafted else 0.0 + draft_key = value("used_draft_model_key") + if draft_key: + lines.append(f"[INFO] Draft model used: {draft_key}") + lines.append( + f"[INFO] Speculative decoding: {accepted}/{drafted} draft tokens accepted " + f"({rate:.0f}%), {rejected} rejected" + ) + if rate < 30.0: + lines.append( + "[HINT] Low draft acceptance - this draft model may be slowing " + "generation down. Try a smaller/closer-matched draft model or disable it." + ) + + lines.append(f"[INFO] Stop reason: {value('stop_reason', 'unknown')}") + lines.append(f"[INFO] Total time: {elapsed:.2f}s") + return lines + + # ------------------------------------------------------------------ + # Execution + # ------------------------------------------------------------------ + def generate( self, system_message: str, @@ -356,22 +529,25 @@ def generate( top_k: int = 0, repeat_penalty: float = 1.0, min_p: float = 0.0, - presence_penalty: float = 0.0, - enable_thinking: str = "Model default", + stop_strings: str = "", + context_overflow: str = "Truncate middle", + output_format: str = "Text", + json_schema: str = "", reasoning_mode: str = "Auto-detect (recommended)", custom_open_tag: str = "", custom_close_tag: str = "", unload_llm: bool = True, unload_comfy_models: bool = False, refresh_models: bool = False - ) -> Tuple[str, str, str]: + ) -> Dict[str, Any]: """ Generate text using LM Studio. Returns: - Tuple of (response_text, reasoning_text, troubleshooting_info) + A ComfyUI node result dict carrying both the (response, reasoning, + troubleshooting) tuple and the ``ui`` payload rendered in the node. """ - troubleshooting_lines = [] + troubleshooting_lines: List[str] = [] # Get current config (read the file once and derive everything from it) config = _config_manager.get_config() @@ -395,18 +571,31 @@ def generate( if last_error and not get_last_fetch_success(): troubleshooting_lines.append(f"[WARNING] Startup model fetch: {last_error}") + # Models LM Studio offered but we could not accept: without this the + # model just silently isn't in the dropdown and nobody knows why. + # Informational, not a warning: it is a standing condition the user + # usually cannot act on, and repeating it as a WARNING every single run + # would just train people to ignore the warning prefix. + rejected = get_last_rejected_models() + if rejected: + troubleshooting_lines.append( + f"[INFO] {len(rejected)} model(s) hidden from the dropdown - " + "LM Studio reported an identifier with unsupported characters: " + + ", ".join(repr(m) for m in rejected) + ) + # Resolve main model model_identifier, error = self._resolve_model_identifier( model_selection, custom_model_name, "model" ) if error: troubleshooting_lines.append(f"[ERROR] {error}") - return "", "", "\n".join(troubleshooting_lines) + return self._output("", "", troubleshooting_lines) if not model_identifier: error_msg = "No model selected. Choose a model from dropdown or enter a custom model name." troubleshooting_lines.append(f"[ERROR] {error_msg}") - return "", "", "\n".join(troubleshooting_lines) + return self._output("", "", troubleshooting_lines) troubleshooting_lines.append(f"[INFO] Model: {model_identifier}") @@ -417,9 +606,24 @@ def generate( if error: troubleshooting_lines.append(f"[WARNING] Draft model error: {error}") draft_model = None + elif draft_model and draft_model == model_identifier: + troubleshooting_lines.append( + "[WARNING] Draft model is the same as the main model - speculative " + "decoding disabled (it would only load the same weights twice)" + ) + draft_model = None elif draft_model: troubleshooting_lines.append(f"[INFO] Draft model: {draft_model}") + # Structured output must be resolved before any model work: a broken + # schema should fail instantly, not after loading a model. + structured, structured_error = build_structured_setting(output_format, json_schema) + if structured_error: + troubleshooting_lines.append(f"[ERROR] {structured_error}") + return self._output("", "", troubleshooting_lines) + + stops = parse_stop_strings(stop_strings) + # Unload ComfyUI models if requested if unload_comfy_models: troubleshooting_lines.append("[INFO] Unloading ComfyUI models...") @@ -469,41 +673,26 @@ def generate( # Add user message (with optional images) if pil_images: - # Prepare all images for the SDK - image_handles = [] - 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) + chat.add_user_message(prompt, images=self._prepare_images(client, pil_images)) else: chat.add_user_message(prompt) - # Build generation config - gen_config = { + # Build generation config. + # Keys are validated against LlmPredictionConfigDict in the + # lmstudio SDK. Anything not in that type is DISCARDED SILENTLY + # by the SDK before the request leaves the machine, so a wrong + # key looks like it worked - the applied-config check below is + # what makes that visible. + gen_config: Dict[str, Any] = { "temperature": temperature, "maxTokens": max_tokens, - # Handle context overflow by truncating middle of conversation - "contextOverflowPolicy": "truncateMiddle", + "contextOverflowPolicy": CONTEXT_OVERFLOW_POLICIES.get( + context_overflow, "truncateMiddle" + ), } # Add optional parameters only when set away from their disabled - # default, so default workflows send nothing extra (and stay safe - # on older LM Studio backends, which silently ignore params they - # don't support rather than erroring). - # Key names validated against LM Studio 0.4.17 / LLMPredictionConfigInput. + # default, so default workflows send nothing extra. if top_p < 1.0: gen_config["topPSampling"] = top_p if top_k > 0: @@ -512,69 +701,93 @@ def generate( gen_config["repeatPenalty"] = repeat_penalty if min_p > 0.0: gen_config["minPSampling"] = min_p - if presence_penalty != 0.0: - gen_config["presencePenalty"] = presence_penalty - if enable_thinking != "Model default": - gen_config["enableThinking"] = (enable_thinking == "Enabled") - # Note: seed is not a valid inference-time parameter in LM Studio SDK + if stops: + gen_config["stopStrings"] = stops + if structured: + gen_config["structured"] = structured if draft_model: gen_config["draftModel"] = draft_model - # Show every parameter actually being sent, so the user can confirm - # exactly what is applied. Params left at their disabled default are - # omitted above and therefore won't appear here. - config_summary = ", ".join(f"{k}={v}" for k, v in gen_config.items()) - troubleshooting_lines.append(f"[INFO] Config: {config_summary}") + troubleshooting_lines.append(f"[INFO] Config: {_summarize_config(gen_config)}") troubleshooting_lines.append("[INFO] Generating...") start_time = time.time() - # Generate response - response = model.respond(chat, config=gen_config) - response_text = str(response) + response, native_reasoning, plain_content, interrupted = self._stream( + model, chat, gen_config, max_tokens + ) + elapsed = time.time() - start_time - troubleshooting_lines.append("[INFO] Generation complete") - troubleshooting_lines.append(f"[INFO] Raw response length: {len(response_text)} chars") + if interrupted: + troubleshooting_lines.append( + "[WARNING] Cancelled from ComfyUI - partial output returned" + ) + else: + troubleshooting_lines.append("[INFO] Generation complete") - # Extract inference statistics - tokens_per_sec = getattr(response.stats, 'tokens_per_second', 0.0) - input_tokens = getattr(response.stats, 'prompt_tokens_count', 0) - output_tokens = getattr(response.stats, 'predicted_tokens_count', 0) - time_to_first_token = getattr(response.stats, 'time_to_first_token_sec', None) - stop_reason = getattr(response.stats, 'stop_reason', 'unknown') - elapsed = time.time() - start_time + response_text = response.content + troubleshooting_lines.append(f"[INFO] Raw response length: {len(response_text)} chars") - troubleshooting_lines.append(f"[INFO] Tokens per second: {tokens_per_sec:.2f}") - troubleshooting_lines.append(f"[INFO] Input tokens: {input_tokens}") - troubleshooting_lines.append(f"[INFO] Output tokens: {output_tokens}") - if time_to_first_token is not None: - troubleshooting_lines.append(f"[INFO] Time to first token: {time_to_first_token:.3f}s") - troubleshooting_lines.append(f"[INFO] Stop reason: {stop_reason}") - troubleshooting_lines.append(f"[INFO] Total time: {elapsed:.2f}s") - - # Extract reasoning based on mode - final_response = response_text - reasoning = "" - - if reasoning_mode == "Auto-detect (recommended)": - final_response, reasoning, detected_pattern = extract_reasoning_auto(response_text) - if detected_pattern: - troubleshooting_lines.append(f"[INFO] Auto-detected reasoning format: {detected_pattern}") - elif looks_like_leaked_thinking(response_text): - # The model thought, but in a tagless plain-text format that - # neither LM Studio's parser nor our tag-based extractor caught, - # so the reasoning leaked into the response output. - troubleshooting_lines.append("[WARNING] Output looks like tagless thinking that leaked into the response (no -style tags found)") - if enable_thinking == "Disabled": - troubleshooting_lines.append("[HINT] This model kept thinking despite enable_thinking=Disabled - its chat template likely ignores the enableThinking flag (common for community merges/finetunes)") - troubleshooting_lines.append("[HINT] To fix in LM Studio: set this model's Reasoning Parsing delimiters, or edit its Jinja template to hard-disable thinking ({%- set enable_thinking = false %})") - troubleshooting_lines.append("[HINT] Or, if the model uses a consistent marker, switch reasoning_mode to 'Custom tags' and set the open/close tags") + # Verify the server actually applied what we asked for. The SDK + # drops unknown keys without a word, so this is the only way a + # parameter that quietly does nothing becomes visible. + try: + applied = response.prediction_config.to_dict() + except Exception: # pragma: no cover - defensive + applied = {} + if applied: + ignored = [ + key for key in missing_config_keys(gen_config, applied) + if key not in _UNVERIFIABLE_CONFIG_KEYS + ] + if ignored: + troubleshooting_lines.append( + f"[WARNING] LM Studio did not apply: {', '.join(ignored)} - " + "your installed lmstudio package or LM Studio build may be too old for these" + ) + + troubleshooting_lines.extend(self._stats_lines(response.stats, elapsed)) + + if structured: + if response.structured: + troubleshooting_lines.append("[INFO] Structured output: valid JSON") else: - troubleshooting_lines.append("[INFO] No reasoning tags detected (model may not have used thinking for this query)") - elif reasoning_mode == "Custom tags": - final_response, reasoning = extract_reasoning_custom( - response_text, custom_open_tag, custom_close_tag - ) - # else: "Disabled" - no extraction + # "JSON (no schema)" does not constrain decoding, so models + # commonly answer with a ```json fenced block. Unwrapping it + # (only when the contents really parse) is the difference + # between a usable output and one no downstream node can read. + unfenced, stripped = strip_json_code_fence(response_text) + if stripped: + response_text = unfenced + troubleshooting_lines.append( + "[INFO] Structured output: removed a ```json code fence - " + "the response is valid JSON underneath" + ) + else: + troubleshooting_lines.append( + "[WARNING] Structured output requested but the response did not parse as JSON" + ) + if getattr(response.stats, "stop_reason", None) == "maxPredictedTokensReached": + troubleshooting_lines.append( + "[HINT] The response hit max_tokens mid-object, so the JSON is " + "truncated. Raise max_tokens." + ) + else: + troubleshooting_lines.append( + "[HINT] 'JSON (no schema)' asks for JSON but does not constrain " + "decoding. Use 'JSON (schema below)' with an explicit schema when " + "a downstream node has to parse the response." + ) + + # Split reasoning from the answer. + final_response, reasoning = self._split_reasoning( + response_text, + native_reasoning, + plain_content, + reasoning_mode, + custom_open_tag, + custom_close_tag, + troubleshooting_lines, + ) if reasoning: troubleshooting_lines.append(f"[INFO] Extracted reasoning: {len(reasoning)} chars") @@ -588,8 +801,15 @@ def generate( except Exception as e: troubleshooting_lines.append(f"[WARNING] Failed to unload LLM: {e}") - return final_response, reasoning, "\n".join(troubleshooting_lines) + if interrupted: + # Report the cancellation to ComfyUI *after* unloading, so a + # cancelled run still frees VRAM. + model_management.throw_exception_if_processing_interrupted() + return self._output(final_response, reasoning, troubleshooting_lines) + + except _INTERRUPT_EXCEPTIONS: + raise # user cancellation is not a node failure except Exception as e: error_msg = f"Generation failed: {type(e).__name__}: {e}" troubleshooting_lines.append(f"[ERROR] {error_msg}") @@ -601,13 +821,103 @@ def generate( elif "context" in error_str or "length" in error_str or "2048" in error_str: troubleshooting_lines.append("[HINT] Context length exceeded. In LM Studio, increase the model's context length setting") troubleshooting_lines.append("[HINT] Note: maxTokens limits OUTPUT tokens; contextLength limits TOTAL tokens (input + output)") + elif "schema" in error_str or "json" in error_str: + troubleshooting_lines.append("[HINT] Check json_schema is a valid JSON Schema object, or set output_format back to 'Text'") elif "not found" in error_str or "model" in error_str: troubleshooting_lines.append("[HINT] Check model identifier matches LM Studio exactly") elif "image" in error_str or "vision" in error_str or "multi" in error_str: troubleshooting_lines.append("[HINT] This model may not support images or multiple image inputs. Try with a single image or text-only.") logger.exception("EA_LMStudio generation error") - return "", "", "\n".join(troubleshooting_lines) + return self._output("", "", troubleshooting_lines) + + def _stream(self, model, chat, gen_config: Dict[str, Any], max_tokens: int): + """Run the prediction as a stream. + + Streaming (rather than a blocking ``respond``) buys three things: + ComfyUI's cancel button can actually stop a runaway generation, the + queue progress bar moves instead of the node looking hung, and LM + Studio's own per-fragment ``reasoning_type`` tagging becomes available - + which is more reliable than any tag regex when the model has Reasoning + Parsing configured in LM Studio. + + Returns ``(result, native_reasoning, plain_content, interrupted)``. + """ + pbar = ProgressBar(max_tokens) if ProgressBar is not None else None + interrupted = False + reasoning_chunks: List[str] = [] + content_chunks: List[str] = [] + tokens_seen = 0 + + stream = model.respond_stream(chat, config=gen_config) + # Note: the stream must be drained rather than broken out of. Breaking + # closes the underlying generator and .result() then raises GeneratorExit; + # cancel() ends it promptly with stop_reason "userStopped" instead. + for fragment in stream: + reasoning_type = getattr(fragment, "reasoning_type", "none") + if reasoning_type == "reasoning": + reasoning_chunks.append(fragment.content) + elif reasoning_type == "none": + content_chunks.append(fragment.content) + # reasoningStartTag / reasoningEndTag fragments are the delimiters + # themselves and belong in neither output. + + tokens_seen += getattr(fragment, "tokens_count", 0) or 0 + if pbar is not None: + pbar.update_absolute(min(tokens_seen, max_tokens), max_tokens) + + if not interrupted and model_management.processing_interrupted(): + stream.cancel() + interrupted = True + + return stream.result(), "".join(reasoning_chunks), "".join(content_chunks), interrupted + + @staticmethod + def _split_reasoning( + response_text: str, + native_reasoning: str, + plain_content: str, + reasoning_mode: str, + custom_open_tag: str, + custom_close_tag: str, + troubleshooting_lines: List[str], + ) -> Tuple[str, str]: + """Separate thinking from the final answer. + + LM Studio's own reasoning parser wins when it fired, because it works + off the model's configured delimiters rather than a guess. The tag + regexes remain the fallback for the (very common) case of a model whose + thinking LM Studio was never told how to parse. + """ + if native_reasoning.strip(): + troubleshooting_lines.append( + "[INFO] Reasoning separated by LM Studio's own parser (reasoning_type fragments)" + ) + # plain_content is the same text minus the reasoning; prefer it, but + # fall back to the full content if the backend sent no plain fragments. + answer = plain_content.strip() or response_text.strip() + return answer, native_reasoning.strip() + + if reasoning_mode == "Auto-detect (recommended)": + final_response, reasoning, detected_pattern = extract_reasoning_auto(response_text) + if detected_pattern: + troubleshooting_lines.append(f"[INFO] Auto-detected reasoning format: {detected_pattern}") + elif looks_like_leaked_thinking(response_text): + # The model thought, but in a tagless plain-text format that + # neither LM Studio's parser nor our tag-based extractor caught, + # so the reasoning leaked into the response output. + troubleshooting_lines.append("[WARNING] Output looks like tagless thinking that leaked into the response (no -style tags found)") + troubleshooting_lines.append("[HINT] To fix in LM Studio: set this model's Reasoning Parsing delimiters, or edit its Jinja template to hard-disable thinking ({%- set enable_thinking = false %})") + troubleshooting_lines.append("[HINT] Or, if the model uses a consistent marker, switch reasoning_mode to 'Custom tags' and set the open/close tags") + else: + troubleshooting_lines.append("[INFO] No reasoning tags detected (model may not have used thinking for this query)") + return final_response, reasoning + + if reasoning_mode == "Custom tags": + return extract_reasoning_custom(response_text, custom_open_tag, custom_close_tag) + + # "Disabled" - no extraction + return response_text, "" # Node registration diff --git a/README.md b/README.md index 9f69775..19b43c8 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,12 @@ # EA_LMStudio - The most fully-featured [LM Studio](https://lmstudio.ai/) integration for ComfyUI — auto model discovery, multi-image vision, reasoning extraction, full sampler control, and built-in VRAM management. + The most fully-featured [LM Studio](https://lmstudio.ai/) integration for ComfyUI — auto model discovery, multi-image vision, reasoning extraction, structured JSON output, full sampler control, cancellable streaming, and built-in VRAM management. - image - _Example of using as a Text-only LLM, with automated reasoning extraction_ +![EA LM Studio turning a short idea into a full text-to-image prompt, with the response shown inside the node and inference stats in the troubleshooting output](docs/images/node-text-prompt-enhancer.png) +_Text-only: a prompt enhancer. The response renders inside the node, and `troubleshooting` reports exactly which parameters were applied plus full inference stats._ -image - _Example of using as a VLM, with multiple image inputs_ +![EA LM Studio captioning an image with a vision model, constrained to a JSON schema](docs/images/node-vision-structured-json.png) +_Vision: captioning an image with a VLM into schema-constrained JSON that downstream nodes can parse directly._ This node is ready to be integrated into your workflows in dozens of ways! @@ -14,10 +14,13 @@ - **Auto Model Discovery** - Models populate automatically from LM Studio at startup - **Vision Support** - Up to 4 image inputs with smart auto-resize to prevent OOM -- **Reasoning Extraction** - Separates thinking from final response (DeepSeek R1, Qwen3, QwQ, GLM, etc.), including GPT-OSS "harmony" channel output — leaked `<|channel|>` / `<|message|>` / `<|end|>` / `<|return|>` markers are extracted and cleaned automatically, even when only partial markers survive -- **Advanced Controls** - Temperature, top-k/p, min-p, repetition & presence penalties, thinking toggle, speculative decoding -- **Smart Troubleshooting** - Helpful error messages with specific hints -- **Detailed Stats** - Tokens/sec, time to first token, stop reason, and token counts in troubleshooting output +- **Reasoning Extraction** - Uses LM Studio's own reasoning split when the model has Reasoning Parsing configured, and falls back to tag detection (DeepSeek R1, Qwen3, QwQ, GLM, GPT-OSS "harmony" channels, and similar) when it does not +- **Structured JSON Output** - Constrain decoding to a JSON Schema so downstream nodes can parse the response +- **Cancellable & Streamed** - ComfyUI's cancel button really stops a runaway generation, and the queue progress bar tracks tokens instead of the node looking frozen +- **In-Node Response Preview** - The generated text appears inside the node; no extra preview node required +- **Advanced Controls** - Temperature, top-k/p, min-p, repetition penalty, stop strings, context-overflow policy, speculative decoding +- **Honest Diagnostics** - Every parameter is checked against the config LM Studio reports back, so a setting that silently does nothing is reported instead of hidden +- **Detailed Stats** - Tokens/sec, time to first token, stop reason, token counts, and speculative-decoding acceptance rates - **VRAM Management** - Auto-unload after generation (enabled by default) ## Installation @@ -33,10 +36,10 @@ pip install -r EA_LMStudio/requirements.txt ## Requirements -- **LM Studio** with the local server enabled. Validated against **LM Studio 0.4.17**. -- The **`lmstudio`** Python package (installed via `requirements.txt`). +- **LM Studio** with the local server enabled. Validated against **LM Studio 0.4.19**. +- The **`lmstudio`** Python package (installed via `requirements.txt`). Requires Python 3.10+. -> **Keep LM Studio and the `lmstudio` package updated.** Newer sampling controls (e.g. min-p, presence penalty, thinking toggle) are applied by the backend only when it supports them — older backends silently ignore parameters they don't recognize rather than erroring, so updating ensures every setting actually takes effect. +> **Keep LM Studio and the `lmstudio` package updated.** The SDK discards prediction parameters it does not recognise instead of raising, so an out-of-date package can turn a setting into a silent no-op. The node now compares what it sent against the config LM Studio echoes back and reports any parameter that was dropped, in the `troubleshooting` output. ## Quick Start @@ -44,14 +47,23 @@ pip install -r EA_LMStudio/requirements.txt 2. Start ComfyUI 3. Find the node: **EA -> LMStudio** +Ready-made workflows live in [`example_workflows/`](example_workflows) — drag one onto the ComfyUI canvas: + +| Workflow | What it shows | +|----------|---------------| +| `01-prompt-enhancer.json` | Turning a short idea into a rich t2i prompt, with reasoning and stats previews | +| `02-vision-caption-json.json` | Captioning an image with a VLM into schema-constrained JSON | + ## Tips - **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:** Set the `enable_thinking` toggle to `Enabled` for hybrid models like Qwen3 (no prompt hacks needed), or add `/think` to prompts / "Think step by step" for others. -- **`enable_thinking` not working?** The toggle sets the `enableThinking` flag, which only takes effect when the model's chat template honors the `enable_thinking` Jinja variable. Many community finetunes/merges (and models that emit a tagless `Thinking Process:` preamble instead of `` tags) ignore it and keep thinking regardless. To actually disable or cleanly separate thinking for those, fix it in LM Studio: either edit the model's Jinja template (`{%- set enable_thinking = false %}`) or configure its **Reasoning Parsing** delimiters so LM Studio knows how to split reasoning from the answer. -- **Reducing repetition:** Raise `repeat_penalty` (1.1-1.3) and/or `presence_penalty` (0.3-0.8). These are distinct controls; LM Studio has no separate frequency penalty. `min_p` (0.05-0.1) is a modern alternative to `top_p` for keeping output coherent. +- **Context errors?** Increase context length in LM Studio settings (not max_tokens). Set `context_overflow` to `Stop at limit (error)` if a silently shortened prompt would be worse than no answer. +- **VLM issues?** Try a smaller image resize option or a single image if multi-image fails. +- **Force thinking mode:** Add `/think` to the prompt (Qwen3 family) or "Think step by step" for others. LM Studio's API has no thinking on/off flag — see below. +- **Thinking leaking into the response?** LM Studio splits reasoning from the answer only when the model's **Reasoning Parsing** delimiters are configured. Set them in LM Studio (per model) and the node uses that split directly. Otherwise it falls back to tag detection, and reports it when a tagless `Thinking Process:` preamble appears to have leaked through. +- **Reducing repetition:** Raise `repeat_penalty` (1.1-1.3) and/or use `min_p` (0.05-0.1) as a modern alternative to `top_p`. LM Studio has no presence or frequency penalty. +- **Guaranteed parseable output:** Set `output_format` to `JSON (schema below)` and supply a schema. `JSON (no schema)` only *asks* for JSON — it does not constrain decoding, and many models answer with a ```` ```json ```` fenced block (the node unwraps that automatically when the contents are valid JSON). +- **Speculative decoding not helping?** The troubleshooting output reports the draft-token acceptance rate. Below ~30% the draft model is usually costing more than it saves. ## Custom Server @@ -79,13 +91,25 @@ Exclude models from the dropdown by adding patterns to `lms_config/user_config.j - The `user_config.json` file is gitignored so your settings survive updates. - Restart ComfyUI or toggle the **refresh_models** checkbox after changing patterns. +A model whose identifier contains characters the node will not accept (LM Studio occasionally serves ids like `some-model@?`) is hidden from the dropdown and named in the `troubleshooting` output, so it is never a silent disappearance. + ## Outputs | Output | Description | |--------|-------------| | response | Generated text (reasoning removed if extracted) | | reasoning | Extracted thinking content | -| troubleshooting | Status messages, debug hints, and inference stats (tokens/sec, input/output tokens, time to first token, stop reason, total time) | +| troubleshooting | Status messages, debug hints, and inference stats (tokens/sec, input/output/total tokens, time to first token, stop reason, speculative-decoding acceptance, total time) | + +The response also renders inside the node itself, so a preview node is optional. + +## Upgrading from 1.x + +**Version 2.0.0 removed the `presence_penalty` and `enable_thinking` widgets.** Neither ever did anything: the `lmstudio` SDK silently discards prediction-config keys it does not recognise, and LM Studio has no presence penalty and no thinking on/off flag. They were dropped before the request ever left your machine. + +You do not need to rebuild anything. When a workflow saved by 1.x is loaded, the node detects the old layout and realigns every remaining setting to the right widget, showing a toast when it does. Without that, removing two mid-list widgets would have shifted every later value by one or two slots. + +If you relied on `enable_thinking`, the equivalents that genuinely work are configuring **Reasoning Parsing** for the model in LM Studio, editing the model's Jinja template (`{%- set enable_thinking = false %}`), or a `/think` / `/no_think` marker in the prompt for Qwen3-family models. ## License diff --git a/docs/images/node-text-prompt-enhancer.png b/docs/images/node-text-prompt-enhancer.png new file mode 100644 index 0000000..f8bb25b Binary files /dev/null and b/docs/images/node-text-prompt-enhancer.png differ diff --git a/docs/images/node-vision-structured-json.png b/docs/images/node-vision-structured-json.png new file mode 100644 index 0000000..cf5d61b Binary files /dev/null and b/docs/images/node-vision-structured-json.png differ diff --git a/example_workflows/01-prompt-enhancer.json b/example_workflows/01-prompt-enhancer.json new file mode 100644 index 0000000..0755515 --- /dev/null +++ b/example_workflows/01-prompt-enhancer.json @@ -0,0 +1,122 @@ +{ + "last_node_id": 5, + "last_link_id": 3, + "nodes": [ + { + "id": 1, + "type": "Note", + "pos": [40, 40], + "size": [420, 390], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "EA LM Studio - prompt enhancer\n\nSetup:\n1. Start LM Studio and enable its local server (default 127.0.0.1:1234).\n2. Start ComfyUI AFTER LM Studio so the model list is populated, or flip\n the refresh_models toggle on the node to fetch it now.\n3. Pick your model in model_selection.\n\nThe response is shown inside the node itself and is also wired to the\nPreview as Text nodes so you can see the thinking and the diagnostics.\n\nTo feed an image generator, wire the 'response' output into a\nCLIP Text Encode node's text input instead of the preview.\n\nunload_llm is ON by default so the LLM frees its VRAM before you\nrender - turn it off to keep the model warm between runs." + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 2, + "type": "EA_LMStudio", + "pos": [500, 40], + "size": [480, 900], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { "name": "image1", "type": "IMAGE", "link": null, "shape": 7 }, + { "name": "image2", "type": "IMAGE", "link": null, "shape": 7 }, + { "name": "image3", "type": "IMAGE", "link": null, "shape": 7 }, + { "name": "image4", "type": "IMAGE", "link": null, "shape": 7 } + ], + "outputs": [ + { "name": "response", "type": "STRING", "links": [1], "slot_index": 0 }, + { "name": "reasoning", "type": "STRING", "links": [2], "slot_index": 1 }, + { "name": "troubleshooting", "type": "STRING", "links": [3], "slot_index": 2 } + ], + "properties": { "Node name for S&R": "EA_LMStudio" }, + "widgets_values": [ + "You are a prompt engineer for text-to-image models. Rewrite the user's idea as a single vivid prompt of at most 60 words. Describe subject, setting, lighting, composition and lens. Output the prompt only - no preamble, no quotes, no bullet points.", + "a lighthouse in a storm", + "-- Custom (enter below) --", + "", + 300, + 0.8, + 0, + "randomize", + 0.95, + 0, + 1.1, + 0.05, + "", + "Truncate middle", + "Text", + "", + "Auto-detect (recommended)", + "", + "", + "Medium (768px)", + "-- Custom (enter below) --", + "", + true, + false, + false + ] + }, + { + "id": 3, + "type": "PreviewAny", + "pos": [1020, 40], + "size": [400, 220], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [{ "name": "source", "type": "*", "link": 1 }], + "outputs": [], + "properties": { "Node name for S&R": "PreviewAny" }, + "widgets_values": [], + "title": "Enhanced prompt" + }, + { + "id": 4, + "type": "PreviewAny", + "pos": [1020, 300], + "size": [400, 220], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [{ "name": "source", "type": "*", "link": 2 }], + "outputs": [], + "properties": { "Node name for S&R": "PreviewAny" }, + "widgets_values": [], + "title": "Reasoning (empty unless the model thinks)" + }, + { + "id": 5, + "type": "PreviewAny", + "pos": [1020, 560], + "size": [400, 300], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [{ "name": "source", "type": "*", "link": 3 }], + "outputs": [], + "properties": { "Node name for S&R": "PreviewAny" }, + "widgets_values": [], + "title": "Troubleshooting + inference stats" + } + ], + "links": [ + [1, 2, 0, 3, 0, "STRING"], + [2, 2, 1, 4, 0, "STRING"], + [3, 2, 2, 5, 0, "STRING"] + ], + "groups": [], + "config": {}, + "extra": {}, + "version": 0.4 +} diff --git a/example_workflows/02-vision-caption-json.json b/example_workflows/02-vision-caption-json.json new file mode 100644 index 0000000..8bb7320 --- /dev/null +++ b/example_workflows/02-vision-caption-json.json @@ -0,0 +1,124 @@ +{ + "last_node_id": 5, + "last_link_id": 3, + "nodes": [ + { + "id": 1, + "type": "Note", + "pos": [40, 40], + "size": [460, 460], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "EA LM Studio - vision captioning with structured JSON output\n\nSetup:\n1. Load a VISION model (VLM) in LM Studio - e.g. a Qwen-VL or\n LLaVA-family model. A text-only model will fail on the image input.\n2. Pick it in model_selection (flip refresh_models if the list is stale).\n3. Choose an image in Load Image.\n\nThis example uses output_format = 'JSON (schema below)', so decoding is\nconstrained to the schema and the response is always parseable JSON.\nThat is the reliable mode - 'JSON (no schema)' only asks for JSON and\nmodels often answer with a ```json fenced block instead.\n\nimage_resize caps the longest edge before upload. Smaller is faster and\navoids out-of-memory on large VLMs; use 'No Resize' only when fine detail\nmatters.\n\nUp to four images can be connected, but not every VLM accepts more than one." + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 2, + "type": "LoadImage", + "pos": [40, 560], + "size": [460, 460], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { "name": "IMAGE", "type": "IMAGE", "links": [1], "slot_index": 0 }, + { "name": "MASK", "type": "MASK", "links": null, "slot_index": 1 } + ], + "properties": { "Node name for S&R": "LoadImage" }, + "widgets_values": ["example.png", "image"] + }, + { + "id": 3, + "type": "EA_LMStudio", + "pos": [540, 40], + "size": [480, 940], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { "name": "image1", "type": "IMAGE", "link": 1, "shape": 7 }, + { "name": "image2", "type": "IMAGE", "link": null, "shape": 7 }, + { "name": "image3", "type": "IMAGE", "link": null, "shape": 7 }, + { "name": "image4", "type": "IMAGE", "link": null, "shape": 7 } + ], + "outputs": [ + { "name": "response", "type": "STRING", "links": [2], "slot_index": 0 }, + { "name": "reasoning", "type": "STRING", "links": null, "slot_index": 1 }, + { "name": "troubleshooting", "type": "STRING", "links": [3], "slot_index": 2 } + ], + "properties": { "Node name for S&R": "EA_LMStudio" }, + "widgets_values": [ + "You are an image captioning assistant. Describe only what is visible. Do not speculate.", + "Describe this image.", + "-- Custom (enter below) --", + "", + 400, + 0.3, + 0, + "randomize", + 1, + 0, + 1, + 0, + "", + "Truncate middle", + "JSON (schema below)", + "{\n \"type\": \"object\",\n \"properties\": {\n \"caption\": { \"type\": \"string\" },\n \"subjects\": { \"type\": \"array\", \"items\": { \"type\": \"string\" } },\n \"setting\": { \"type\": \"string\" },\n \"mood\": { \"type\": \"string\" }\n },\n \"required\": [\"caption\", \"subjects\", \"setting\", \"mood\"]\n}", + "Auto-detect (recommended)", + "", + "", + "Medium (768px)", + "-- Custom (enter below) --", + "", + true, + false, + false + ] + }, + { + "id": 4, + "type": "PreviewAny", + "pos": [1060,40], + "size": [400, 320], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [{ "name": "source", "type": "*", "link": 2 }], + "outputs": [], + "properties": { "Node name for S&R": "PreviewAny" }, + "widgets_values": [], + "title": "Caption (JSON)" + }, + { + "id": 5, + "type": "PreviewAny", + "pos": [1060,400], + "size": [400, 420], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [{ "name": "source", "type": "*", "link": 3 }], + "outputs": [], + "properties": { "Node name for S&R": "PreviewAny" }, + "widgets_values": [], + "title": "Troubleshooting + inference stats" + } + ], + "links": [ + [1, 2, 0, 3, 0, "IMAGE"], + [2, 3, 0, 4, 0, "STRING"], + [3, 3, 2, 5, 0, "STRING"] + ], + "groups": [], + "config": {}, + "extra": {}, + "version": 0.4 +} diff --git a/lms_params.py b/lms_params.py new file mode 100644 index 0000000..99d0bac --- /dev/null +++ b/lms_params.py @@ -0,0 +1,167 @@ +""" +Prediction-parameter helpers for EA_LMStudio. + +Pure functions that turn widget strings into the shapes the LM Studio SDK +expects. Kept out of ``LMStudio.py`` (which imports the lmstudio SDK and ComfyUI +and performs a network fetch at import time) so they stay unit-testable on their +own. + +Everything here is validated against ``LlmPredictionConfigDict`` from the +``lmstudio`` Python SDK. That type is the authoritative list of accepted keys — +and the SDK *silently discards* keys it does not recognise rather than raising, +so a typo or a wished-for parameter becomes a no-op that looks like it worked. +``LMStudio.py`` guards against that by diffing the requested config against the +config the server echoes back on every run. +""" +import json +from typing import Any, Dict, List, Optional, Tuple + +# Widget label -> LM Studio contextOverflowPolicy value. +CONTEXT_OVERFLOW_OPTIONS = [ + "Truncate middle", + "Rolling window", + "Stop at limit (error)", +] + +CONTEXT_OVERFLOW_POLICIES = { + "Truncate middle": "truncateMiddle", + "Rolling window": "rollingWindow", + "Stop at limit (error)": "stopAtLimit", +} + +# Widget label -> structured output mode. +# +# "JSON (no schema)" maps to LM Studio's ``{"type": "json"}``, which *asks* for +# JSON but does not constrain decoding: models frequently answer with a +# ```json fenced block. Only the schema form actually constrains the sampler, +# which is why the labels say so rather than promising valid JSON either way. +OUTPUT_FORMAT_OPTIONS = [ + "Text", + "JSON (no schema)", + "JSON (schema below)", +] + +# Escape sequences honoured inside a stop string, so a user can stop on a +# newline (which they cannot type into a single line of the widget). +_STOP_STRING_ESCAPES = ( + ("\\\\", "\x00"), # protect literal backslashes first + ("\\n", "\n"), + ("\\r", "\r"), + ("\\t", "\t"), +) + + +def unescape_stop_string(raw: str) -> str: + """Expand ``\\n``/``\\r``/``\\t``/``\\\\`` in a single stop string.""" + out = raw + for token, replacement in _STOP_STRING_ESCAPES: + out = out.replace(token, replacement) + return out.replace("\x00", "\\") + + +def parse_stop_strings(raw: str) -> List[str]: + """Parse the ``stop_strings`` widget into a list for the SDK. + + One stop string per line. Blank lines are ignored so a trailing newline in + the textarea does not become an empty stop string (which would stop the + prediction immediately). Leading/trailing spaces are significant and kept — + stopping on ``"\\nUser:"`` or ``" ###"`` is a real use case — so only the + line terminator is stripped. + """ + if not raw: + return [] + stops = [] + for line in raw.replace("\r\n", "\n").split("\n"): + if not line.strip(): + continue + stops.append(unescape_stop_string(line)) + return stops + + +def parse_json_schema(raw: str) -> Tuple[Optional[Any], Optional[str]]: + """Parse the ``json_schema`` widget. + + Returns: + ``(schema, error)`` — exactly one is non-None. A schema must be a JSON + object; LM Studio rejects bare arrays/scalars, so we catch that here + where we can explain it, rather than letting the server fail mid-run. + """ + text = (raw or "").strip() + if not text: + return None, "Output format is 'JSON (schema below)' but json_schema is empty" + + try: + schema = json.loads(text) + except json.JSONDecodeError as e: + return None, f"json_schema is not valid JSON: {e}" + + if not isinstance(schema, dict): + return None, ( + f"json_schema must be a JSON object (got {type(schema).__name__}). " + 'Example: {"type": "object", "properties": {"caption": {"type": "string"}}}' + ) + + return schema, None + + +def build_structured_setting( + output_format: str, json_schema_text: str +) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: + """Build the ``structured`` prediction-config value for the chosen format. + + Returns ``(setting, error)``. ``(None, None)`` means plain text — send + nothing, which is what a default workflow does. + """ + if output_format == "JSON (no schema)": + return {"type": "json"}, None + + if output_format == "JSON (schema below)": + schema, error = parse_json_schema(json_schema_text) + if error: + return None, error + return {"type": "json", "jsonSchema": schema}, None + + return None, None + + +def strip_json_code_fence(text: str) -> Tuple[str, bool]: + """Unwrap a ```json ... ``` fence around an otherwise-valid JSON document. + + Returns ``(text, stripped)``. Only used when the caller asked for JSON + output and the raw response did not parse: models routinely answer a JSON + request with a fenced markdown block, which breaks every downstream parser. + The unwrapped text is returned only if it actually parses as JSON, so this + can never turn a prose answer into something that looks structured. + """ + candidate = (text or "").strip() + if not candidate.startswith("```"): + return text, False + + # Drop the opening fence line (```json / ```JSON / bare ```) and the closer. + newline = candidate.find("\n") + if newline == -1: + return text, False + inner = candidate[newline + 1:] + if inner.rstrip().endswith("```"): + inner = inner.rstrip()[: -3] + + inner = inner.strip() + try: + json.loads(inner) + except (json.JSONDecodeError, ValueError): + return text, False + + return inner, True + + +def missing_config_keys( + requested: Dict[str, Any], applied: Dict[str, Any] +) -> List[str]: + """Keys we asked for that LM Studio did not echo back as applied. + + The SDK drops unrecognised prediction-config keys without complaining, so a + parameter can silently do nothing. Comparing what we sent against the + ``prediction_config`` the server returns turns that class of bug into a + visible warning instead of a mystery. + """ + return [key for key in requested if key not in applied] diff --git a/model_fetcher.py b/model_fetcher.py index e72c0e3..c87c26e 100644 --- a/model_fetcher.py +++ b/model_fetcher.py @@ -14,6 +14,10 @@ _cached_models: List[str] = [] _last_fetch_error: Optional[str] = None _last_fetch_success: bool = False +# Model IDs the server offered but validate_model_identifier refused. Kept so the +# node can say *why* a model is missing from the dropdown instead of it just not +# being there (LM Studio does hand out ids like "some-model@?" in practice). +_last_rejected_models: List[str] = [] # Constants CUSTOM_MODEL_OPTION = "-- Custom (enter below) --" @@ -69,7 +73,7 @@ def fetch_models_from_server( server_url: str, timeout: float = 5.0, excluded_patterns: Optional[List[str]] = None, -) -> tuple[List[str], Optional[str]]: +) -> tuple[List[str], Optional[str], List[str]]: """ Fetch available models from LM Studio server. @@ -80,11 +84,13 @@ def fetch_models_from_server( If None, uses default ["embedding"]. Pass an empty list to include all models. Returns: - Tuple of (model_list, error_message) + Tuple of (model_list, error_message, rejected_models) - model_list: Filtered list of model IDs, empty on failure - error_message: None on success, descriptive error on failure + - rejected_models: IDs the server offered that failed validation """ models: List[str] = [] + rejected: List[str] = [] error: Optional[str] = None # Default to ["embedding"] if no patterns specified @@ -106,7 +112,7 @@ def fetch_models_from_server( if "data" not in data: error = "Unexpected response format from LM Studio (missing 'data' field)" logger.warning(f"EA_LMStudio: {error}") - return models, error + return models, error, rejected for model in data["data"]: model_id = model.get("id", "") @@ -120,9 +126,14 @@ def fetch_models_from_server( continue # Validate the model ID before adding - is_valid, _ = validate_model_identifier(model_id) + is_valid, reason = validate_model_identifier(model_id) if is_valid: models.append(model_id) + else: + rejected.append(model_id) + logger.warning( + f"EA_LMStudio: hiding model {model_id!r} from the dropdown - {reason}" + ) # Sort alphabetically for easier navigation models.sort(key=str.lower) @@ -148,7 +159,7 @@ def fetch_models_from_server( error = f"Unexpected error fetching models: {type(e).__name__}: {str(e)}" logger.error(f"EA_LMStudio: {error}") - return models, error + return models, error, rejected def get_model_choices() -> List[str]: @@ -168,6 +179,17 @@ def get_model_choices() -> List[str]: return choices +def get_default_model_choice() -> str: + """Default selection for the main model dropdown. + + The first *real* model when discovery worked, so a freshly added node is + runnable straight away. Previously this resolved to the "Custom" sentinel in + every case (it took choices[0], which is always the sentinel), so a new node + always failed its first run with "No model selected". + """ + return _cached_models[0] if _cached_models else CUSTOM_MODEL_OPTION + + def refresh_model_cache( server_url: str, timeout: float = 5.0, @@ -185,9 +207,9 @@ def refresh_model_cache( Returns: Tuple of (success, message) """ - global _cached_models, _last_fetch_error, _last_fetch_success + global _cached_models, _last_fetch_error, _last_fetch_success, _last_rejected_models - models, error = fetch_models_from_server(server_url, timeout, excluded_patterns) + models, error, rejected = fetch_models_from_server(server_url, timeout, excluded_patterns) if error: _last_fetch_error = error @@ -195,13 +217,19 @@ def refresh_model_cache( return False, error _cached_models = models + _last_rejected_models = rejected _last_fetch_error = None _last_fetch_success = True + suffix = f" ({len(rejected)} hidden - unsafe identifier)" if rejected else "" + if models: - return True, f"Successfully loaded {len(models)} models from LM Studio" + return True, f"Successfully loaded {len(models)} models from LM Studio{suffix}" else: - return True, "Connected to LM Studio but no models found (embedding models are excluded)" + return True, ( + "Connected to LM Studio but no models found " + f"(embedding models are excluded){suffix}" + ) def initialize_model_cache(server_url: str, timeout: float = 5.0, excluded_patterns=None) -> None: @@ -231,3 +259,8 @@ def get_last_fetch_success() -> bool: def get_cached_model_count() -> int: """Get the number of currently cached models.""" return len(_cached_models) + + +def get_last_rejected_models() -> List[str]: + """Model IDs the last successful fetch refused as unsafe identifiers.""" + return list(_last_rejected_models) diff --git a/pyproject.toml b/pyproject.toml index a75cb3e..b69267b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,10 +1,16 @@ [project] name = "EA_LMStudio" description = "Fully-featured LM Studio integration for ComfyUI: model discovery, vision, reasoning extraction, and full sampler control." -version = "1.5.1" +version = "2.0.0" license = {file = "LICENSE"} -requires-python = ">=3.9" +# The lmstudio SDK itself declares Requires-Python >=3.10, so 3.9 was never +# actually installable regardless of what this said. +requires-python = ">=3.10" dependencies = [ + # LlmPredictionFragment.reasoning_type (LM Studio's own reasoning split) and + # PredictionResult.prediction_config (the applied-config check) are read + # defensively, so an older SDK degrades to tag-regex extraction rather than + # failing. Verified present from 1.0.1 onwards. "lmstudio>=1.0.0", "requests>=2.27.0", "Pillow>=9.0.0", diff --git a/tests/conftest.py b/tests/conftest.py index a0af481..17d9743 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -40,3 +40,17 @@ def _ensure_stub(name: str) -> None: # time, so bare module objects are enough to let the package import succeed. _ensure_stub("lmstudio") _ensure_stub("comfy.model_management") + + +# ``LMStudio.py`` is a package module (``from .lms_config import ...``), so it +# cannot be imported as a top-level module the way the dependency-light modules +# are. Register the repo root as a synthetic package instead of relying on the +# checkout directory being named "EA_LMStudio" - a clone into any other folder +# name would otherwise break the suite. ``__init__.py`` is deliberately not +# executed: it registers ComfyUI server routes we have no use for here. +NODE_PACKAGE = "ea_lmstudio_under_test" + +if NODE_PACKAGE not in sys.modules: + _package = types.ModuleType(NODE_PACKAGE) + _package.__path__ = [_REPO_ROOT] + sys.modules[NODE_PACKAGE] = _package diff --git a/tests/test_model_fetcher.py b/tests/test_model_fetcher.py index 18a8799..7cb42ec 100644 --- a/tests/test_model_fetcher.py +++ b/tests/test_model_fetcher.py @@ -1,5 +1,12 @@ -"""Tests for model_fetcher: identifier validation and exclusion matching.""" -from model_fetcher import validate_model_identifier, _is_excluded +"""Tests for model_fetcher: identifier validation, exclusion and discovery.""" +import model_fetcher +from model_fetcher import ( + CUSTOM_MODEL_OPTION, + _is_excluded, + fetch_models_from_server, + get_default_model_choice, + validate_model_identifier, +) # --- validate_model_identifier ------------------------------------------ @@ -63,3 +70,67 @@ def test_not_excluded(): def test_empty_patterns_excludes_nothing(): assert _is_excluded("anything", []) is False + + +# --- fetch_models_from_server ------------------------------------------- + +class _FakeResponse: + def __init__(self, payload): + self._payload = payload + + def raise_for_status(self): + pass + + def json(self): + return self._payload + + +def _fake_server(monkeypatch, ids): + payload = {"data": [{"id": model_id} for model_id in ids]} + monkeypatch.setattr( + model_fetcher.requests, "get", lambda *a, **kw: _FakeResponse(payload) + ) + + +def test_fetch_sorts_and_excludes(monkeypatch): + _fake_server(monkeypatch, ["zeta-7b", "text-embedding-small", "Alpha-3b"]) + models, error, rejected = fetch_models_from_server("http://x", 1.0, ["embedding"]) + assert error is None + assert models == ["Alpha-3b", "zeta-7b"] + assert rejected == [] + + +def test_unsafe_identifier_is_reported_not_just_dropped(monkeypatch): + """LM Studio really does serve ids like "model@?". + + Those fail validation, and before v2.0.0 they vanished from the dropdown + with no explanation anywhere. + """ + _fake_server(monkeypatch, ["good-model", "gemma4-31b-balanced-mtp@?"]) + models, error, rejected = fetch_models_from_server("http://x", 1.0, []) + assert error is None + assert models == ["good-model"] + assert rejected == ["gemma4-31b-balanced-mtp@?"] + + +def test_missing_data_field_is_an_error(monkeypatch): + monkeypatch.setattr( + model_fetcher.requests, "get", lambda *a, **kw: _FakeResponse({"oops": []}) + ) + models, error, rejected = fetch_models_from_server("http://x", 1.0, []) + assert models == [] + assert rejected == [] + assert "data" in error + + +# --- get_default_model_choice ------------------------------------------- + +def test_default_is_a_real_model_when_available(monkeypatch): + """A fresh node must land on a usable model, not the Custom sentinel.""" + monkeypatch.setattr(model_fetcher, "_cached_models", ["alpha-3b", "zeta-7b"]) + assert get_default_model_choice() == "alpha-3b" + + +def test_default_falls_back_to_custom_when_discovery_failed(monkeypatch): + monkeypatch.setattr(model_fetcher, "_cached_models", []) + assert get_default_model_choice() == CUSTOM_MODEL_OPTION diff --git a/tests/test_node_helpers.py b/tests/test_node_helpers.py new file mode 100644 index 0000000..cda0108 --- /dev/null +++ b/tests/test_node_helpers.py @@ -0,0 +1,194 @@ +"""Tests for the node's pure helpers. + +Importing ``LMStudio`` pulls in the ``lmstudio``/``comfy`` stubs registered by +conftest and performs the startup model fetch, which fails harmlessly when no LM +Studio server is listening. The helpers exercised here never touch either. +""" +import importlib + +# Synthetic package registered by conftest.py (pytest's importlib mode means +# conftest itself is not importable by name, so the name is repeated here). +NODE_PACKAGE = "ea_lmstudio_under_test" + +_node = importlib.import_module(f"{NODE_PACKAGE}.LMStudio") +EALMStudio = _node.EALMStudio +_summarize_config = _node._summarize_config + + +class _Stats: + """Stand-in for LlmPredictionStats. + + Every field except ``stop_reason`` is Optional on the real type, which is + exactly what these tests pin down. + """ + + def __init__(self, **fields): + self.stop_reason = "eosFound" + self.tokens_per_second = None + self.prompt_tokens_count = None + self.predicted_tokens_count = None + self.total_tokens_count = None + self.time_to_first_token_sec = None + self.num_gpu_layers = None + self.total_draft_tokens_count = None + self.accepted_draft_tokens_count = None + self.rejected_draft_tokens_count = None + self.used_draft_model_key = None + for key, value in fields.items(): + setattr(self, key, value) + + +# --- _stats_lines --------------------------------------------------------- + +def test_stats_survive_a_backend_that_reports_nothing(): + """The regression that discarded successful generations. + + tokens_per_second is Optional; formatting None with ":.2f" raised TypeError, + which the outer handler reported as "Generation failed" *after* the model had + already produced the text. + """ + lines = EALMStudio._stats_lines(_Stats(), 1.5) + assert any("Stop reason" in line for line in lines) + assert any("Total time" in line for line in lines) + assert not any("Tokens per second" in line for line in lines) + + +def test_stats_reported_when_present(): + lines = EALMStudio._stats_lines( + _Stats( + tokens_per_second=42.123, + prompt_tokens_count=30, + predicted_tokens_count=7, + total_tokens_count=37, + time_to_first_token_sec=0.1122, + ), + 2.0, + ) + joined = "\n".join(lines) + assert "Tokens per second: 42.12" in joined + assert "Input tokens: 30" in joined + assert "Output tokens: 7" in joined + assert "Total tokens: 37" in joined + assert "Time to first token: 0.112s" in joined + + +def test_zero_output_tokens_still_reported(): + """0 is a meaningful count and must not be swallowed as falsy.""" + lines = EALMStudio._stats_lines(_Stats(predicted_tokens_count=0), 1.0) + assert any("Output tokens: 0" in line for line in lines) + + +def test_speculative_decoding_acceptance_reported(): + lines = EALMStudio._stats_lines( + _Stats( + total_draft_tokens_count=6, + accepted_draft_tokens_count=5, + rejected_draft_tokens_count=1, + used_draft_model_key="tiny-draft", + ), + 1.0, + ) + joined = "\n".join(lines) + assert "tiny-draft" in joined + assert "5/6 draft tokens accepted (83%)" in joined + assert "HINT" not in joined # good acceptance rate, no nagging + + +def test_poor_draft_acceptance_gets_a_hint(): + lines = EALMStudio._stats_lines( + _Stats( + total_draft_tokens_count=100, + accepted_draft_tokens_count=5, + rejected_draft_tokens_count=95, + ), + 1.0, + ) + assert any("HINT" in line for line in lines) + + +# --- _split_reasoning ----------------------------------------------------- + +def test_native_reasoning_wins_over_regex(): + """LM Studio's own parser is authoritative when it fired.""" + log = [] + answer, reasoning = EALMStudio._split_reasoning( + "The answer is 42.", + "let me work this out", + "The answer is 42.", + "Auto-detect (recommended)", + "", + "", + log, + ) + assert answer == "The answer is 42." + assert reasoning == "let me work this out" + assert any("LM Studio's own parser" in line for line in log) + + +def test_native_reasoning_falls_back_when_no_plain_fragments(): + answer, reasoning = EALMStudio._split_reasoning( + "only content here", "thinking", "", "Auto-detect (recommended)", + "", "", [], + ) + assert answer == "only content here" + assert reasoning == "thinking" + + +def test_regex_used_when_lmstudio_did_not_tag(): + answer, reasoning = EALMStudio._split_reasoning( + "hmmFinal answer.", "", "", "Auto-detect (recommended)", + "", "", [], + ) + assert answer == "Final answer." + assert reasoning == "hmm" + + +def test_disabled_mode_passes_text_through(): + answer, reasoning = EALMStudio._split_reasoning( + "hmmFinal answer.", "", "", "Disabled", + "", "", [], + ) + assert answer == "hmmFinal answer." + assert reasoning == "" + + +def test_custom_tags_mode(): + answer, reasoning = EALMStudio._split_reasoning( + "[R]secret[/R]Answer.", "", "", "Custom tags", "[R]", "[/R]", [], + ) + assert answer == "Answer." + assert reasoning == "secret" + + +def test_whitespace_only_native_reasoning_does_not_hijack(): + """A stream of blank reasoning fragments must not suppress the regex path.""" + answer, reasoning = EALMStudio._split_reasoning( + "hmmFinal answer.", " \n ", "", "Auto-detect (recommended)", + "", "", [], + ) + assert answer == "Final answer." + assert reasoning == "hmm" + + +# --- ui payload ----------------------------------------------------------- + +def test_output_carries_both_ui_and_result(): + """OUTPUT_NODE is only worth having because of the ui half.""" + out = EALMStudio._output("hello", "thinking", ["[INFO] a", "[INFO] b"]) + assert out["result"] == ("hello", "thinking", "[INFO] a\n[INFO] b") + assert out["ui"]["text"] == ["hello"] + assert out["ui"]["reasoning"] == ["thinking"] + + +# --- _summarize_config ---------------------------------------------------- + +def test_config_summary_truncates_long_values(): + summary = _summarize_config({"structured": {"jsonSchema": {"x": "y" * 500}}}) + assert "truncated" in summary + assert len(summary) < 300 + + +def test_config_summary_lists_every_key(): + summary = _summarize_config({"temperature": 0.7, "maxTokens": 10}) + assert "temperature=0.7" in summary + assert "maxTokens=10" in summary diff --git a/tests/test_params.py b/tests/test_params.py new file mode 100644 index 0000000..166772c --- /dev/null +++ b/tests/test_params.py @@ -0,0 +1,220 @@ +"""Tests for lms_params: widget strings -> LM Studio prediction config values.""" +import pytest + +from lms_params import ( + CONTEXT_OVERFLOW_OPTIONS, + CONTEXT_OVERFLOW_POLICIES, + OUTPUT_FORMAT_OPTIONS, + build_structured_setting, + missing_config_keys, + parse_json_schema, + parse_stop_strings, + strip_json_code_fence, + unescape_stop_string, +) + + +# --- parse_stop_strings --------------------------------------------------- + +def test_empty_gives_no_stop_strings(): + assert parse_stop_strings("") == [] + assert parse_stop_strings("\n\n \n") == [] + + +def test_one_per_line(): + assert parse_stop_strings("END\nUser:") == ["END", "User:"] + + +def test_blank_lines_dropped(): + """A trailing newline must not become an empty stop string. + + An empty stop string matches immediately, so it would truncate every + response to nothing. + """ + assert parse_stop_strings("END\n\n") == ["END"] + assert "" not in parse_stop_strings("A\n\n\nB\n") + + +def test_crlf_handled(): + assert parse_stop_strings("A\r\nB") == ["A", "B"] + + +def test_significant_whitespace_preserved(): + assert parse_stop_strings(" ###") == [" ###"] + assert parse_stop_strings("END ") == ["END "] + + +def test_escape_sequences_expanded(): + assert parse_stop_strings("\\nUser:") == ["\nUser:"] + assert parse_stop_strings("a\\tb") == ["a\tb"] + + +def test_literal_backslash_not_eaten_by_escape(): + # "\\n" (an escaped backslash followed by n) must stay a backslash + n, + # not become a newline. + assert unescape_stop_string("\\\\n") == "\\n" + + +# --- parse_json_schema ---------------------------------------------------- + +def test_empty_schema_is_an_error(): + schema, error = parse_json_schema(" ") + assert schema is None + assert "empty" in error.lower() + + +def test_invalid_json_reports_error(): + schema, error = parse_json_schema("{not json}") + assert schema is None + assert "not valid json" in error.lower() + + +def test_non_object_schema_rejected(): + schema, error = parse_json_schema('["a", "b"]') + assert schema is None + assert "object" in error.lower() + + +def test_valid_schema_parsed(): + schema, error = parse_json_schema('{"type": "object", "properties": {}}') + assert error is None + assert schema["type"] == "object" + + +# --- build_structured_setting -------------------------------------------- + +def test_text_sends_nothing(): + setting, error = build_structured_setting("Text", "") + assert setting is None + assert error is None + + +def test_free_form_json(): + setting, error = build_structured_setting("JSON (no schema)", "") + assert error is None + assert setting == {"type": "json"} + + +def test_schema_mode_includes_schema(): + setting, error = build_structured_setting( + "JSON (schema below)", '{"type": "object"}' + ) + assert error is None + assert setting["type"] == "json" + assert setting["jsonSchema"] == {"type": "object"} + + +def test_schema_mode_without_schema_errors(): + setting, error = build_structured_setting("JSON (schema below)", "") + assert setting is None + assert error + + +def test_free_form_json_ignores_schema_box(): + """Leftover text in json_schema must not break the 'any shape' mode.""" + setting, error = build_structured_setting("JSON (no schema)", "{not json}") + assert error is None + assert setting == {"type": "json"} + + +# --- strip_json_code_fence ------------------------------------------------ + +def test_fenced_json_is_unwrapped(): + """Verified live: 'JSON (no schema)' really does come back fenced.""" + text, stripped = strip_json_code_fence('```json\n{"sky": "blue"}\n```') + assert stripped is True + assert text == '{"sky": "blue"}' + + +def test_bare_fence_without_language_tag(): + text, stripped = strip_json_code_fence('```\n{"a": 1}\n```') + assert stripped is True + assert text == '{"a": 1}' + + +def test_unterminated_fence_still_unwrapped_if_valid(): + # Truncated responses often lose the closing fence. + text, stripped = strip_json_code_fence('```json\n{"a": 1}') + assert stripped is True + assert text == '{"a": 1}' + + +def test_plain_json_left_alone(): + text, stripped = strip_json_code_fence('{"a": 1}') + assert stripped is False + assert text == '{"a": 1}' + + +def test_prose_never_becomes_structured(): + """Unwrapping must never dress up a non-JSON answer as JSON.""" + original = "```json\nI could not answer that.\n```" + text, stripped = strip_json_code_fence(original) + assert stripped is False + assert text == original + + +def test_fenced_non_json_code_left_alone(): + original = "```python\nprint(1)\n```" + text, stripped = strip_json_code_fence(original) + assert stripped is False + assert text == original + + +def test_empty_input_is_safe(): + assert strip_json_code_fence("") == ("", False) + + +# --- missing_config_keys -------------------------------------------------- + +def test_nothing_missing(): + assert missing_config_keys({"temperature": 0.7}, {"temperature": 0.7}) == [] + + +def test_dropped_key_detected(): + """The check that would have caught presencePenalty/enableThinking. + + The lmstudio SDK discards config keys it does not know instead of raising, + so the only evidence is the config the server echoes back. + """ + requested = {"temperature": 0.7, "presencePenalty": 0.5} + applied = {"temperature": 0.7, "cpuThreads": 8} + assert missing_config_keys(requested, applied) == ["presencePenalty"] + + +def test_extra_applied_keys_are_not_reported(): + # LM Studio echoes back defaults we never asked for; those are not a problem. + assert missing_config_keys({"temperature": 1.0}, {"temperature": 1.0, "rawTools": {}}) == [] + + +# --- option tables -------------------------------------------------------- + +@pytest.mark.parametrize("label", CONTEXT_OVERFLOW_OPTIONS) +def test_every_context_option_maps_to_a_policy(label): + """A widget label with no mapping would silently fall back to the default.""" + assert label in CONTEXT_OVERFLOW_POLICIES + + +def test_context_policies_are_valid_sdk_values(): + # LlmPredictionConfigDict["contextOverflowPolicy"] literal values. + assert set(CONTEXT_OVERFLOW_POLICIES.values()) <= { + "stopAtLimit", + "truncateMiddle", + "rollingWindow", + } + + +def test_default_context_option_preserves_v1_behaviour(): + assert CONTEXT_OVERFLOW_POLICIES[CONTEXT_OVERFLOW_OPTIONS[0]] == "truncateMiddle" + + +def test_default_output_format_is_text(): + assert OUTPUT_FORMAT_OPTIONS[0] == "Text" + + +@pytest.mark.parametrize("label", OUTPUT_FORMAT_OPTIONS) +def test_every_output_format_is_handled(label): + setting, error = build_structured_setting(label, '{"type": "object"}') + # Either it produces a setting or it is the plain-text mode; never an error + # for a well-formed schema. + assert error is None + assert setting is None or setting["type"] == "json" diff --git a/web/ea_lmstudio.js b/web/ea_lmstudio.js index caf3a3f..5964e03 100644 --- a/web/ea_lmstudio.js +++ b/web/ea_lmstudio.js @@ -1,26 +1,249 @@ import { app } from "../../scripts/app.js"; import { api } from "../../scripts/api.js"; +import { ComfyWidgets } from "../../scripts/widgets.js"; /* - * EA_LMStudio Model Refresh Extension + * EA_LMStudio frontend extension * - * Intercepts the refresh_models toggle to fetch an updated model list from - * the LM Studio server and update the dropdown widgets in-place. + * Three jobs: + * 1. Model refresh - the refresh_models toggle re-fetches the model list and + * updates both dropdowns in place, with a visible toast. + * 2. Response preview - renders the generated text inside the node (the node + * is an OUTPUT_NODE and sends a ui payload for this). + * 3. Legacy migration - v1.x saved presence_penalty and enable_thinking into + * widgets_values. Both were removed in v2.0.0 because the + * LM Studio SDK silently discarded them, so every widget + * after them would otherwise load one or two slots out of + * alignment. See migrateLegacyWidgetValues below. + */ + +const NODE_CLASS = "EA_LMStudio"; +const CUSTOM_MODEL_OPTION = "-- Custom (enter below) --"; // keep in sync with model_fetcher.py +const PREVIEW_WIDGET_NAME = "response_preview"; + +/* + * Widget order as serialised by v1.5.x, in LiteGraph order (required first, + * then optional; IMAGE inputs are links, not widgets, so they never appear). + * ComfyUI inserts a control_after_generate widget straight after an INT named + * "seed" on some frontend versions, hence the two variants. + */ +const LEGACY_ORDER = [ + "system_message", + "prompt", + "model_selection", + "custom_model_name", + "max_tokens", + "temperature", + "seed", + "image_resize", + "draft_model_selection", + "custom_draft_model", + "top_p", + "top_k", + "repeat_penalty", + "min_p", + "presence_penalty", + "enable_thinking", + "reasoning_mode", + "custom_open_tag", + "custom_close_tag", + "unload_llm", + "unload_comfy_models", + "refresh_models", +]; +// v2.0.0 also regrouped the widgets, so the migration below cannot assume the +// old and new orders line up - it maps every legacy value onto the current +// widget of the same name, which is order-independent by construction. + +const LEGACY_ORDER_WITH_SEED_CONTROL = [ + ...LEGACY_ORDER.slice(0, 7), + "control_after_generate", + ...LEGACY_ORDER.slice(7), +]; + +// v1.x values of the removed enable_thinking widget. Used to confirm an array of +// the right length really is a legacy layout before rewriting anything. +const LEGACY_ENABLE_THINKING_VALUES = ["Model default", "Enabled", "Disabled"]; + +function toast(severity, summary, detail) { + const manager = app.extensionManager?.toast; + if (manager?.add) { + manager.add({ severity, summary, detail, life: 5000 }); + } else if (severity === "error") { + console.error(`[EA_LMStudio] ${summary}: ${detail}`); + } else { + console.log(`[EA_LMStudio] ${summary}: ${detail}`); + } +} + +/** + * Collect each input's declared default so widgets that did not exist in v1.x + * can be reset after a legacy remap (LiteGraph will already have filled them + * with whatever landed at their index in the old array). + */ +function collectDefaults(nodeData) { + const defaults = {}; + for (const group of ["required", "optional"]) { + const inputs = nodeData?.input?.[group] ?? {}; + for (const [name, spec] of Object.entries(inputs)) { + const [type, options] = Array.isArray(spec) ? spec : [spec, undefined]; + if (options && Object.prototype.hasOwnProperty.call(options, "default")) { + defaults[name] = options.default; + } else if (Array.isArray(type)) { + defaults[name] = type[0]; // combo: first entry + } + } + } + return defaults; +} + +/** + * Repair widget values loaded from a workflow saved by v1.x. + * + * Returns true when a remap happened. A v2 workflow has at least 24 widget + * values, so the 22/23 length test cannot collide with a current save; the + * enable_thinking value check guards against a coincidence anyway. + */ +function migrateLegacyWidgetValues(node, widgetValues, defaults) { + if (!Array.isArray(widgetValues)) return false; + + let order = null; + if (widgetValues.length === LEGACY_ORDER.length) { + order = LEGACY_ORDER; + } else if (widgetValues.length === LEGACY_ORDER_WITH_SEED_CONTROL.length) { + order = LEGACY_ORDER_WITH_SEED_CONTROL; + } + if (!order) return false; + + const enableThinkingValue = widgetValues[order.indexOf("enable_thinking")]; + if (!LEGACY_ENABLE_THINKING_VALUES.includes(enableThinkingValue)) return false; + + const byName = {}; + order.forEach((name, index) => { + byName[name] = widgetValues[index]; + }); + + for (const widget of node.widgets ?? []) { + if (widget.name === PREVIEW_WIDGET_NAME) continue; + if (Object.prototype.hasOwnProperty.call(byName, widget.name)) { + widget.value = byName[widget.name]; + } else if (Object.prototype.hasOwnProperty.call(defaults, widget.name)) { + // New in v2.0.0 - it holds a shifted v1 value right now. + widget.value = defaults[widget.name]; + } + } + + toast( + "info", + "EA LM Studio workflow updated", + "Loaded a v1 workflow: presence_penalty and enable_thinking were removed " + + "(LM Studio never applied them) and the remaining settings were realigned." + ); + return true; +} + +/** + * Grow a node that is stored smaller than its widgets need. * - * 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. + * A workflow stores the node's size, and LiteGraph restores it verbatim - it + * does not re-check that the widgets still fit. Any release that adds a widget + * therefore leaves every previously saved workflow a row or two too short, and + * the overflowing widgets draw outside the node's frame. v2.0.0 has a net two + * more widgets than v1.5.x, so this affects every upgraded workflow, not just + * an unlucky few. Only ever grows - a deliberately widened node is preserved. */ +function growToFitWidgets(node) { + try { + const [minWidth, minHeight] = node.computeSize(); + if (node.size[0] < minWidth || node.size[1] < minHeight) { + node.setSize([ + Math.max(node.size[0], minWidth), + Math.max(node.size[1], minHeight), + ]); + } + } catch (err) { + console.error("[EA_LMStudio] Could not resize node to fit widgets:", err); + } +} + +function getPreviewWidget(node) { + let widget = node.widgets?.find((w) => w.name === PREVIEW_WIDGET_NAME); + if (widget) return widget; + + widget = ComfyWidgets["STRING"]( + node, + PREVIEW_WIDGET_NAME, + ["STRING", { multiline: true }], + app + ).widget; + + if (widget.inputEl) { + widget.inputEl.readOnly = true; + widget.inputEl.style.opacity = "0.8"; + widget.inputEl.placeholder = "Response appears here after the node runs"; + } + // Never let the preview reach the prompt or the saved widget values - it is + // display only, and a stray extra value is exactly the kind of drift this + // extension has to migrate away from. + widget.options = { ...(widget.options ?? {}), serialize: false }; + widget.serializeValue = () => undefined; + return widget; +} app.registerExtension({ - name: "EA_LMStudio.ModelRefresh", + name: "EA_LMStudio.NodeExtras", + + async beforeRegisterNodeDef(nodeType, nodeData) { + if (nodeData.name !== NODE_CLASS) return; + + const defaults = collectDefaults(nodeData); + + const onConfigure = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function (info) { + onConfigure?.apply(this, arguments); + try { + migrateLegacyWidgetValues(this, info?.widgets_values, defaults); + } catch (err) { + console.error("[EA_LMStudio] Legacy workflow migration failed:", err); + } + growToFitWidgets(this); + }; + + /* + * Keep the preview out of the saved workflow. + * + * The preview is display-only and is excluded from the API prompt, but + * LiteGraph's serialize() still writes its text into widgets_values - + * which would bake the last generated response into every workflow file + * the user shares. It is always the last widget, so dropping its entry + * cannot disturb the index of any real widget. + */ + const onSerialize = nodeType.prototype.onSerialize; + nodeType.prototype.onSerialize = function (info) { + onSerialize?.apply(this, arguments); + const index = this.widgets?.findIndex((w) => w.name === PREVIEW_WIDGET_NAME); + if (index >= 0 && Array.isArray(info?.widgets_values) && info.widgets_values.length > index) { + info.widgets_values.splice(index, 1); + } + }; + + const onExecuted = nodeType.prototype.onExecuted; + nodeType.prototype.onExecuted = function (message) { + onExecuted?.apply(this, arguments); + + const text = Array.isArray(message?.text) ? message.text.join("") : ""; + const widget = getPreviewWidget(this); + widget.value = text; + + growToFitWidgets(this); + app.graph.setDirtyCanvas(true, false); + }; + }, async nodeCreated(node) { - if (node.comfyClass !== "EA_LMStudio") return; + if (node.comfyClass !== NODE_CLASS) return; - const refreshWidget = node.widgets?.find(w => w.name === "refresh_models"); + const refreshWidget = node.widgets?.find((w) => w.name === "refresh_models"); if (!refreshWidget) return; const originalCallback = refreshWidget.callback; @@ -38,12 +261,10 @@ app.registerExtension({ const data = await resp.json(); if (data.success && data.models) { - // Keep this literal in sync with CUSTOM_MODEL_OPTION in - // model_fetcher.py (the Python side is the source of truth). - const choices = ["-- Custom (enter below) --", ...data.models]; + const choices = [CUSTOM_MODEL_OPTION, ...data.models]; for (const widgetName of ["model_selection", "draft_model_selection"]) { - const w = node.widgets?.find(ww => ww.name === widgetName); + 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) @@ -54,16 +275,16 @@ app.registerExtension({ } } - console.log( - `[EA_LMStudio] Refreshed models: ${data.models.length} found` - ); + toast("success", "LM Studio models refreshed", data.message); } else { - console.warn( - `[EA_LMStudio] Model refresh failed: ${data.message || "unknown error"}` + toast( + "warn", + "LM Studio model refresh failed", + data.message || "Unknown error - is LM Studio running with the server enabled?" ); } } catch (err) { - console.error("[EA_LMStudio] Failed to refresh models:", err); + toast("error", "LM Studio model refresh failed", String(err)); } // Toggle back off so it acts like a one-shot button