diff --git a/lmstudio_config.json b/lmstudio_config.json
index a2dd991..ae2b52d 100644
--- a/lmstudio_config.json
+++ b/lmstudio_config.json
@@ -1,10 +1,10 @@
{
"temperature": 0.15,
- "max_tokens": 768,
+ "max_tokens": -1,
"seed": 0,
"unload_after_use": true,
"default_model": "qwen/qwen2.5-vl-3b",
"image_max_sizes": [256, 512, 1024, 2048],
"default_image_max_size": 1024,
- "allowed_root_directories": ["/mnt/k", "/mnt/m"]
+ "allowed_root_directories": ["/mnt/k", "/mnt/m", "/mnt/e", "/home/was/projects"]
}
diff --git a/nodes.py b/nodes.py
index e281f7c..32d658d 100644
--- a/nodes.py
+++ b/nodes.py
@@ -3,6 +3,8 @@ import json
import io
import uuid
import shutil
+import base64
+import requests as _requests
import numpy as np
@@ -85,6 +87,17 @@ def image_tensor_to_temp_png_path(image_tensor, max_edge: int = 1024) -> str:
return p
+def resize_disk_image_to_temp(path: str, max_edge: int = 1024) -> str:
+ """Resize an on-disk image to max_edge (keeping aspect ratio) and return a temp PNG path."""
+ img = Image.open(path)
+ if max(img.size) > max_edge:
+ img.thumbnail((max_edge, max_edge), Image.Resampling.LANCZOS)
+ d = folder_paths.get_temp_directory()
+ out = os.path.join(d, f"{uuid.uuid4().hex}.png")
+ img.save(out, format="PNG")
+ return out
+
+
def listify(x):
"""Convert input to list if it's not already."""
if x is None:
@@ -105,6 +118,11 @@ def strip_thinking_tags(text: str) -> str:
# Remove thinking tags and everything between them (including newlines)
result = re.sub(r'.*?', '', text, flags=re.DOTALL)
result = re.sub(r'.*?', '', result, flags=re.DOTALL)
+ # Handle missing opening tag: strip everything before a lone or
+ result = re.sub(r'^.*?', '', result, flags=re.DOTALL)
+ result = re.sub(r'^.*?', '', result, flags=re.DOTALL)
+ # Handle SDK "Thinking Process:" prefix when context was exhausted before
+ result = re.sub(r'^Thinking Process:.*', '', result, flags=re.DOTALL)
# Clean up any extra whitespace left behind
result = re.sub(r'\n\s*\n\s*\n', '\n\n', result) # Replace 3+ newlines with 2
return result.strip()
@@ -115,8 +133,13 @@ def map_lms_params(params: Dict[str, Any]) -> Dict[str, Any]:
for k, v in params.items():
if v is None:
continue
+ if k == "strip_thinking_tags":
+ continue
if k == "max_tokens":
- out["maxTokens"] = v
+ if v != -1:
+ out["maxTokens"] = v
+ else:
+ print("[LMStudio] max_tokens=-1 → unlimited generation (maxTokens omitted from SDK config)")
elif k == "top_p":
out["topP"] = v
elif k == "top_k":
@@ -132,6 +155,157 @@ def map_lms_params(params: Dict[str, Any]) -> Dict[str, Any]:
return out
+def encode_image_base64(image_path: str) -> str:
+ with open(image_path, "rb") as f:
+ return base64.b64encode(f.read()).decode("utf-8")
+
+
+def query_lmstudio_rest(
+ model_id: str,
+ messages: List[Dict[str, Any]],
+ params: Optional[Dict[str, Any]] = None,
+ base_url: str = "http://127.0.0.1:1234",
+) -> Tuple[str, Optional[str], str]:
+ payload: Dict[str, Any] = {"model": model_id, "messages": messages}
+ if params:
+ for k, v in params.items():
+ if k in ("strip_thinking_tags",):
+ continue
+ if v is None:
+ continue
+ if k == "max_tokens":
+ if v != -1:
+ payload["max_tokens"] = v
+ elif k == "seed":
+ if v != 0:
+ payload["seed"] = v
+ elif k == "stop":
+ payload["stop"] = v
+ elif k == "repeat_penalty":
+ payload["repeat_penalty"] = v
+ else:
+ payload[k] = v
+ resp = _requests.post(f"{base_url}/v1/chat/completions", json=payload, timeout=600)
+ resp.raise_for_status()
+ data = resp.json()
+ choice = data["choices"][0]
+ msg = choice["message"]
+ content = msg.get("content", "") or ""
+ reasoning = msg.get("reasoning_content", None)
+ finish_reason = choice.get("finish_reason", "unknown")
+ return (content, reasoning, finish_reason)
+
+
+def build_openai_messages(
+ system_prompt: str,
+ user_prompt: str,
+ image_paths: Optional[List[str]] = None,
+ history: Optional[List[Dict[str, str]]] = None,
+) -> List[Dict[str, Any]]:
+ messages: List[Dict[str, Any]] = []
+ if history:
+ messages.extend(history)
+ elif system_prompt:
+ messages.append({"role": "system", "content": system_prompt})
+ if image_paths:
+ content_parts: List[Dict[str, Any]] = []
+ if user_prompt:
+ content_parts.append({"type": "text", "text": user_prompt})
+ for path in image_paths:
+ b64 = encode_image_base64(path)
+ content_parts.append({
+ "type": "image_url",
+ "image_url": {"url": f"data:image/png;base64,{b64}"},
+ })
+ messages.append({"role": "user", "content": content_parts})
+ else:
+ messages.append({"role": "user", "content": user_prompt or ""})
+ return messages
+
+
+def predict_text(
+ model_dict: Dict[str, Any],
+ system_prompt: str,
+ user_prompt: str,
+ image_paths: Optional[List[str]] = None,
+ params: Optional[Dict[str, Any]] = None,
+ model_handle: Any = None,
+ history: Optional[List[Dict[str, str]]] = None,
+) -> str:
+ model_id = model_dict.get("model_id", "")
+ if model_dict.get("use_rest_api", False):
+ cfg = read_config()
+ base_url = cfg.get("base_url", "http://127.0.0.1:1234")
+ messages = build_openai_messages(system_prompt, user_prompt, image_paths, history)
+ content, reasoning, finish_reason = query_lmstudio_rest(model_id, messages, params, base_url)
+ print(f"[LMStudio REST] finish={finish_reason} content_len={len(content)}, reasoning_len={len(reasoning) if reasoning else 0}")
+ result = content.strip()
+ if not result:
+ # Model produced reasoning but no content — retry with /no_think to
+ # force a direct answer without the thinking phase.
+ nothink_prompt = ((user_prompt or "") + " /no_think").strip()
+ retry_messages = build_openai_messages(system_prompt, nothink_prompt, image_paths, history)
+ content2, _, finish2 = query_lmstudio_rest(model_id, retry_messages, params, base_url)
+ print(f"[LMStudio REST] retry /no_think: finish={finish2} content_len={len(content2)}")
+ result = content2.strip()
+ print(f"[LMStudio REST] result_len={len(result)}")
+ return result
+ else:
+ import lmstudio as lms
+ import tempfile
+ handle = model_handle or (lms.llm(model_id) if model_id else lms.llm())
+ strip = (params or {}).get("strip_thinking_tags", False)
+ prompt = user_prompt or ""
+ sdk_config = map_lms_params(params or {})
+
+ # When dealing with a thinking model (strip_thinking_tags=True) and
+ # max_tokens is unlimited (-1), cap generation to prevent the model
+ # from spending the entire context window on thinking.
+ if strip and "maxTokens" not in sdk_config:
+ sdk_config["maxTokens"] = 4096
+ print("[LMStudio SDK] Thinking model with unlimited max_tokens — capping generation at 4096 tokens")
+
+ # Retry loop: on context overflow, downscale images and retry
+ current_paths = list(image_paths) if image_paths else []
+ max_retries = 3
+ for attempt in range(max_retries + 1):
+ if history:
+ chat = lms.Chat.from_history({"messages": history})
+ else:
+ chat = lms.Chat(system_prompt) if system_prompt else lms.Chat()
+ if current_paths:
+ img_handles = [lms.prepare_image(p) for p in current_paths]
+ chat.add_user_message(prompt, images=img_handles)
+ else:
+ chat.add_user_message(prompt)
+ try:
+ pred = handle.respond(chat, config=sdk_config)
+ raw = str(pred)
+ has_think = "" in raw or "" in raw
+ print(f"[LMStudio SDK] raw_len={len(raw)} has_think_tags={has_think} strip_requested={strip}")
+ return raw
+ except Exception as e:
+ err_msg = str(e)
+ if "Context size" in err_msg and current_paths and attempt < max_retries:
+ # Downscale images by 75% and retry
+ new_paths = []
+ for p in current_paths:
+ img = Image.open(p)
+ w, h = img.size
+ nw, nh = int(w * 0.75), int(h * 0.75)
+ if nw < 64 or nh < 64:
+ print(f"[LMStudio SDK] Image too small to downscale further ({w}x{h})")
+ raise
+ img = img.resize((nw, nh), Image.LANCZOS)
+ tmp = os.path.join(tempfile.gettempdir(), f"lms_ctx_{attempt}_{os.path.basename(p)}")
+ img.save(tmp, "PNG")
+ new_paths.append(tmp)
+ current_paths = new_paths
+ print(f"[LMStudio SDK] Context overflow — downscaled images to ~{int(100 * 0.75**(attempt+1))}% and retrying (attempt {attempt+2}/{max_retries+1})")
+ else:
+ raise
+
+
def check_chat_fits_in_context(model_handle: Any, chat: Any) -> Tuple[bool, int, int]:
"""
Check if a chat conversation fits within the model's context window.
@@ -386,9 +560,9 @@ class WASLMStudioModel:
"INT",
{
"default": max_tokens_default,
- "min": 1,
+ "min": -1,
"max": 32768,
- "tooltip": "Maximum new tokens to generate for the assistant reply.",
+ "tooltip": "Maximum new tokens to generate for the assistant reply. Use -1 for unlimited (recommended for thinking models).",
},
),
"seed": (
@@ -407,6 +581,13 @@ class WASLMStudioModel:
"tooltip": "Maximum edge size for input images (keeps aspect ratio). Images larger than this are downscaled before encoding and sending to LM Studio.",
},
),
+ "use_rest_api": (
+ "BOOLEAN",
+ {
+ "default": False,
+ "tooltip": "Use the OpenAI-compatible REST API instead of the LM Studio SDK for predictions. Required for proper thinking model support (separates reasoning from content).",
+ },
+ ),
},
"optional": {
},
@@ -426,6 +607,7 @@ class WASLMStudioModel:
max_tokens: int,
seed: int,
image_max_size: str,
+ use_rest_api: bool = False,
):
import lmstudio as lms
@@ -450,6 +632,7 @@ class WASLMStudioModel:
"max_tokens": int(max_tokens),
"seed": int(seed) if seed != 0 else None,
"image_max_size": int(selected_max),
+ "use_rest_api": bool(use_rest_api),
}
return (result,)
@@ -543,9 +726,6 @@ class WASLMStudioQuery:
for idx, img in enumerate(imgs):
path = image_tensor_to_temp_png_path(img, max_edge=max_edge)
tmp_img_paths.append(path)
- img_handle = lms.prepare_image(path)
- chat = lms.Chat(system_prompt) if system_prompt else lms.Chat()
- chat.add_user_message(user_prompt or "", images=[img_handle])
params: Dict[str, Any] = {
"temperature": temperature,
"max_tokens": max_tokens,
@@ -558,18 +738,16 @@ class WASLMStudioQuery:
params[k] = v
except Exception:
pass
- pred = model_handle.respond(chat, config=map_lms_params(params))
- out.append(str(pred))
+ text = predict_text(model, system_prompt, user_prompt or "", image_paths=[path], params=params, model_handle=model_handle)
+ out.append(text)
pbar.update_absolute(idx + 1)
responses_out = out
else:
- img_handles = []
+ all_paths = []
for i in imgs:
path = image_tensor_to_temp_png_path(i, max_edge=max_edge)
tmp_img_paths.append(path)
- img_handles.append(lms.prepare_image(path))
- chat = lms.Chat(system_prompt) if system_prompt else lms.Chat()
- chat.add_user_message(user_prompt or "", images=img_handles)
+ all_paths.append(path)
params: Dict[str, Any] = {
"temperature": temperature,
"max_tokens": max_tokens,
@@ -582,11 +760,9 @@ class WASLMStudioQuery:
params[k] = v
except Exception:
pass
- pred = model_handle.respond(chat, config=map_lms_params(params))
- responses_out = [str(pred)]
+ text = predict_text(model, system_prompt, user_prompt or "", image_paths=all_paths, params=params, model_handle=model_handle)
+ responses_out = [text]
else:
- chat = lms.Chat(system_prompt) if system_prompt else lms.Chat()
- chat.add_user_message(user_prompt or "")
params: Dict[str, Any] = {
"temperature": temperature,
"max_tokens": max_tokens,
@@ -599,8 +775,8 @@ class WASLMStudioQuery:
params[k] = v
except Exception:
pass
- pred = model_handle.respond(chat, config=map_lms_params(params))
- responses_out = [str(pred)]
+ text = predict_text(model, system_prompt, user_prompt or "", params=params, model_handle=model_handle)
+ responses_out = [text]
except Exception as e:
print(f"Error in run_query: {e}")
@@ -642,7 +818,7 @@ class WASLMStudioOptions:
),
"max_tokens": (
"INT",
- {"default": 512, "min": 1, "max": 32768, "tooltip": "Maximum number of new tokens to generate for the response."},
+ {"default": 512, "min": -1, "max": 32768, "tooltip": "Maximum number of new tokens to generate for the response. Use -1 for unlimited (recommended for thinking models)."},
),
"seed": (
"INT",
@@ -899,9 +1075,6 @@ class WASLMStudioChat:
for idx, img in enumerate(imgs):
path = image_tensor_to_temp_png_path(img, max_edge=max_edge)
tmp_img_paths.append(path)
- img_handle = lms.prepare_image(path)
- chat = lms.Chat.from_history({"messages": msgs}) if msgs else (lms.Chat(system_prompt) if system_prompt else lms.Chat())
- chat.add_user_message(user_prompt or "", images=[img_handle])
params: Dict[str, Any] = {
"temperature": temperature,
"max_tokens": max_tokens,
@@ -914,20 +1087,18 @@ class WASLMStudioChat:
params[k] = v
except Exception:
pass
- pred = model_handle.respond(chat, config=map_lms_params(params))
- out.append(str(pred))
+ resp_text = predict_text(model, system_prompt, user_prompt or "", image_paths=[path], params=params, model_handle=model_handle, history=msgs if msgs else None)
+ out.append(resp_text)
msgs.append({"role": "user", "content": user_prompt or ""})
- msgs.append({"role": "assistant", "content": str(pred)})
+ msgs.append({"role": "assistant", "content": resp_text})
pbar.update_absolute(idx + 1)
responses_out = out
else:
- img_handles = []
+ all_paths = []
for i in imgs:
path = image_tensor_to_temp_png_path(i, max_edge=max_edge)
tmp_img_paths.append(path)
- img_handles.append(lms.prepare_image(path))
- chat = lms.Chat.from_history({"messages": msgs}) if msgs else (lms.Chat(system_prompt) if system_prompt else lms.Chat())
- chat.add_user_message(user_prompt or "", images=img_handles)
+ all_paths.append(path)
params: Dict[str, Any] = {
"temperature": temperature,
"max_tokens": max_tokens,
@@ -940,13 +1111,11 @@ class WASLMStudioChat:
params[k] = v
except Exception:
pass
- pred = model_handle.respond(chat, config=map_lms_params(params))
- responses_out = [str(pred)]
+ resp_text = predict_text(model, system_prompt, user_prompt or "", image_paths=all_paths, params=params, model_handle=model_handle, history=msgs if msgs else None)
+ responses_out = [resp_text]
msgs.append({"role": "user", "content": user_prompt or ""})
- msgs.append({"role": "assistant", "content": str(pred)})
+ msgs.append({"role": "assistant", "content": resp_text})
else:
- chat = lms.Chat.from_history({"messages": msgs}) if msgs else (lms.Chat(system_prompt) if system_prompt else lms.Chat())
- chat.add_user_message(user_prompt or "")
params: Dict[str, Any] = {
"temperature": temperature,
"max_tokens": max_tokens,
@@ -959,10 +1128,10 @@ class WASLMStudioChat:
params[k] = v
except Exception:
pass
- pred = model_handle.respond(chat, config=map_lms_params(params))
- responses_out = [str(pred)]
+ resp_text = predict_text(model, system_prompt, user_prompt or "", params=params, model_handle=model_handle, history=msgs if msgs else None)
+ responses_out = [resp_text]
msgs.append({"role": "user", "content": user_prompt or ""})
- msgs.append({"role": "assistant", "content": str(pred)})
+ msgs.append({"role": "assistant", "content": resp_text})
except Exception as e:
print(f"Error in chat: {e}")
responses_out = [f"Error: {str(e)}"]
@@ -1101,9 +1270,6 @@ class WASLMStudioCaption:
for idx, img in enumerate(imgs):
path = image_tensor_to_temp_png_path(img, max_edge=max_edge)
tmp_img_paths.append(path)
- img_handle = lms.prepare_image(path)
- chat = lms.Chat(system_text) if system_text else lms.Chat()
- chat.add_user_message(user_prompt or "", images=[img_handle])
params: Dict[str, Any] = {
"temperature": temperature,
"max_tokens": max_tokens,
@@ -1116,19 +1282,17 @@ class WASLMStudioCaption:
params[k] = v
except Exception:
pass
- pred = model_handle.respond(chat, config=map_lms_params(params))
- out.append(str(pred))
+ text = predict_text(model, system_text, user_prompt or "", image_paths=[path], params=params, model_handle=model_handle)
+ out.append(text)
pbar.update_absolute(idx + 1)
result = string_list(out)
else:
- img_handles = []
tmp_img_paths: List[str] = []
+ all_paths: List[str] = []
for i in imgs:
path = image_tensor_to_temp_png_path(i, max_edge=max_edge)
tmp_img_paths.append(path)
- img_handles.append(lms.prepare_image(path))
- chat = lms.Chat(system_text) if system_text else lms.Chat()
- chat.add_user_message(user_prompt or "", images=img_handles)
+ all_paths.append(path)
params: Dict[str, Any] = {
"temperature": temperature,
"max_tokens": max_tokens,
@@ -1141,8 +1305,8 @@ class WASLMStudioCaption:
params[k] = v
except Exception:
pass
- pred = model_handle.respond(chat, config=map_lms_params(params))
- result = string_list([str(pred)])
+ text = predict_text(model, system_text, user_prompt or "", image_paths=all_paths, params=params, model_handle=model_handle)
+ result = string_list([text])
except Exception as e:
print(f"Error in generate_captions: {e}")
@@ -1349,6 +1513,7 @@ class WASLMStudioCaptionDataset:
temperature = float(model.get("temperature", 0.2))
max_tokens = int(model.get("max_tokens", 512))
seed = model.get("seed", None)
+ max_edge = int(model.get("image_max_size", 1024))
unload = model.get("unload", False)
# Apply options overrides if provided
@@ -1398,9 +1563,7 @@ class WASLMStudioCaptionDataset:
skipped_count += 1
pbar.update_absolute(idx + 1)
continue
- img_handle = lms.prepare_image(p)
- chat = lms.Chat(system_text) if system_text else lms.Chat()
- chat.add_user_message(user_prompt or "", images=[img_handle])
+ resized_path = resize_disk_image_to_temp(p, max_edge=max_edge)
params: Dict[str, Any] = {
"temperature": temperature,
"max_tokens": max_tokens,
@@ -1413,12 +1576,19 @@ class WASLMStudioCaptionDataset:
params[k] = v
except Exception:
pass
- pred = model_handle.respond(chat, config=map_lms_params(params))
- text = str(pred).strip()
+ text = predict_text(model, system_text, user_prompt or "", image_paths=[resized_path], params=params, model_handle=model_handle)
+ print(f"[LMStudio DEBUG] after predict_text: len={len(text)} preview={text[:80]!r}")
# Apply thinking tag filtering if requested
if options and options.get("strip_thinking_tags", False):
text = strip_thinking_tags(text)
+ print(f"[LMStudio DEBUG] after strip_thinking: len={len(text)}")
+
+ if not text.strip():
+ print(f"[LMStudio WARNING] Empty result for {os.path.basename(p)} — model thinking may have filled the entire context. Increase context_length in LM Studio.")
+ failed_count += 1
+ pbar.update_absolute(idx + 1)
+ continue
if trigger_word_or_phrase.strip():
trigger = trigger_word_or_phrase.strip()
@@ -1575,6 +1745,7 @@ class WASLMStudioCaptionDatasetCustom:
temperature = float(model.get("temperature", 0.2))
max_tokens = int(model.get("max_tokens", 512))
seed = model.get("seed", None)
+ max_edge = int(model.get("image_max_size", 1024))
unload = model.get("unload", False)
# Apply options overrides if provided
@@ -1624,9 +1795,7 @@ class WASLMStudioCaptionDatasetCustom:
skipped_count += 1
pbar.update_absolute(idx + 1)
continue
- img_handle = lms.prepare_image(p)
- chat = lms.Chat(system_text) if system_text else lms.Chat()
- chat.add_user_message(user_prompt or "", images=[img_handle])
+ resized_path = resize_disk_image_to_temp(p, max_edge=max_edge)
params: Dict[str, Any] = {
"temperature": temperature,
"max_tokens": max_tokens,
@@ -1639,12 +1808,19 @@ class WASLMStudioCaptionDatasetCustom:
params[k] = v
except Exception:
pass
- pred = model_handle.respond(chat, config=map_lms_params(params))
- text = str(pred).strip()
+ text = predict_text(model, system_text, user_prompt or "", image_paths=[resized_path], params=params, model_handle=model_handle)
+ print(f"[LMStudio DEBUG] after predict_text: len={len(text)} preview={text[:80]!r}")
# Apply thinking tag filtering if requested
if options and options.get("strip_thinking_tags", False):
text = strip_thinking_tags(text)
+ print(f"[LMStudio DEBUG] after strip_thinking: len={len(text)}")
+
+ if not text.strip():
+ print(f"[LMStudio WARNING] Empty result for {os.path.basename(p)} — model thinking may have filled the entire context. Increase context_length in LM Studio.")
+ failed_count += 1
+ pbar.update_absolute(idx + 1)
+ continue
if trigger_word_or_phrase.strip():
trigger = trigger_word_or_phrase.strip()
diff --git a/pyproject.toml b/pyproject.toml
index 9f51f8f..0cbe7c6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
[project]
name = "was-lmstudio-easyquery"
-version = "0.2.0"
+version = "0.1.1"
description = "LM Studio calls with image support and easy task directions."
requires-python = ">=3.10"
license = { file = "LICENSE" }
@@ -22,4 +22,4 @@ PublisherId = "was"
DisplayName = "WAS LMStudio Easy-Query"
# Icon = "https://raw.githubusercontent.com/WASasquatch/WAS_LMStudio_EasyQuery/main/icon.png"
# Banner = "https://raw.githubusercontent.com/WASasquatch/WAS_LMStudio_EasyQuery/main/banner.png"
-requires-comfyui = ">=1.0.0"
+requires-comfyui = ">=1.0.0"
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index a701d80..45a704f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
lmstudio>=1.5.0
pillow>=10.0.0
-numpy>=1.23.0
\ No newline at end of file
+numpy>=1.23.0
+requests>=2.28.0
\ No newline at end of file