From 2deb431ea1c5a67a274dfb9681ddf0f47049a79d Mon Sep 17 00:00:00 2001 From: A Latent Place Date: Thu, 20 Mar 2025 19:45:40 +0100 Subject: [PATCH] Added reasoning output --- LMStudio.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/LMStudio.py b/LMStudio.py index 3365746..eacad88 100644 --- a/LMStudio.py +++ b/LMStudio.py @@ -18,6 +18,7 @@ class YANCLMSTUDIO: "multiline": True, "default": "You are an AI assistant specialized in generating detailed and creative image prompts for AI image generation. Your task is to expand a given user prompt into a well-structured, vivid, and highly descriptive prompt while ensuring that all terms from the original prompt are included. Enhance the visual quality and artistic impact by adding relevant details, but do not omit or alter any key elements provided by the user. Follow the given instructions or guidelines and respond only with the refined prompt." }), + "reasoning_tag": ("STRING", {"default": "think"}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "ip": ("STRING", {"default": "localhost"}), "port": ("INT", {"default": 1234}), @@ -29,8 +30,8 @@ class YANCLMSTUDIO: }, } - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("Extended Prompt",) + RETURN_TYPES = ("STRING", "STRING") + RETURN_NAMES = ("Extended Prompt", "Reasoning") OUTPUT_NODE = False @@ -63,7 +64,7 @@ class YANCLMSTUDIO: return response - def do_it(self, prompt, model_identifier, system_message, seed, ip, port, temperature, max_tokens, unload_llm, unload_comfy_models): + def do_it(self, prompt, model_identifier, system_message, reasoning_tag, seed, ip, port, temperature, max_tokens, unload_llm, unload_comfy_models): url = self.format_url(ip, port) @@ -80,11 +81,13 @@ class YANCLMSTUDIO: url, model_identifier, messages, temperature, max_tokens) if response.status_code == 200: - result = response.json()["choices"][0]["message"]["content"] + content = response.json()["choices"][0]["message"]["content"] # Removing the Deepseek block. - result = re.sub(r".*?", "", - result, flags=re.DOTALL).strip() + result = re.sub(rf"<{reasoning_tag}>.*?", "", + content, flags=re.DOTALL).strip() + reasoning = re.sub(rf".*<{reasoning_tag}>(.*?).*", r"\1", content, flags=re.DOTALL).strip() + else: raise Exception(f"Failed to get response: {response.status_code}") @@ -94,7 +97,7 @@ class YANCLMSTUDIO: text=True, check=True) - return (result,) + return (result, reasoning,) NODE_CLASS_MAPPINGS = {