From 33dec1631b59238bb08fb5cb845a4fed4dda0a5f Mon Sep 17 00:00:00 2001 From: Marcel Hofer <49859266+goodguy1963@users.noreply.github.com> Date: Thu, 6 Aug 2026 16:55:16 +0200 Subject: [PATCH] fix: reuse ComfyUI text encoders --- AILab_QwenVL.py | 95 ++++++++++++++++++++++++++++++++++++-- README.md | 9 ++++ tests/test_new_features.py | 61 ++++++++++++++++++++++++ 3 files changed, 162 insertions(+), 3 deletions(-) diff --git a/AILab_QwenVL.py b/AILab_QwenVL.py index e114bf1..b7236fc 100644 --- a/AILab_QwenVL.py +++ b/AILab_QwenVL.py @@ -709,7 +709,7 @@ NO_PRESET_PROMPT = "🚫 No preset (image-only)" PRESET_PROMPTS: list[str] = [NO_PRESET_PROMPT, "Describe this image in detail."] TOOLTIPS = { - "model_name": "Pick the Qwen-VL checkpoint. First run downloads weights into models/LLM/Qwen-VL, so leave disk space.", + "model_name": "Pick the checkpoint. [ComfyUI] entries reuse compatible models/text_encoders files; other entries download into models/LLM/Qwen-VL on first use.", "quantization": "Precision vs VRAM. FP16 gives the best quality if memory allows; 8-bit suits 8–16 GB GPUs; 4-bit fits 6 GB or lower but is slower.", "attention_mode": "auto tries SageAttention → FlashAttention 2 → SDPA in order. SDPA is stable and recommended. Only override when debugging attention backends.", "preset_prompt": "Built-in instruction describing how Qwen-VL should analyze the media input.", @@ -893,7 +893,7 @@ def load_model_configs(): def _scan_local_hf_models(): - """Scan all LLM/Qwen-VL directories for locally available HF models not already in the config.""" + """Scan local HF directories and compatible ComfyUI text encoders.""" global HF_VL_MODELS, HF_ALL_MODELS # Collect all Qwen-VL directories to scan from ComfyUI's multi-path system @@ -946,8 +946,33 @@ def _scan_local_hf_models(): except PermissionError: pass + try: + text_encoders = folder_paths.get_filename_list("text_encoders") + except (AttributeError, KeyError): + text_encoders = [] + for filename in text_encoders: + normalized = Path(filename).name.lower().replace("-", "").replace("_", "") + family = None + if "gemma312b" in normalized: + family = "gemma3" + elif "qwen3vl" in normalized and ("4b" in normalized or "8b" in normalized): + family = "qwen3vl" + if not family or not normalized.endswith(".safetensors"): + continue + display = f"[ComfyUI] {filename}" + model_info = { + "comfy_text_encoder": filename, + "repo_id": None, + "is_local": True, + "quantized": True, + "native_family": family, + } + HF_VL_MODELS[display] = model_info + HF_ALL_MODELS[display] = model_info + count += 1 + if count: - print(f"[QwenVL] Discovered {count} local HF model(s) on disk") + print(f"[QwenVL] Discovered {count} local model(s) on disk") if not HF_ALL_MODELS: @@ -1522,6 +1547,23 @@ class QwenVLBase: unique_id=None, hf_token: str | None = None, ): + model_info = HF_ALL_MODELS.get(model_name, {}) + comfy_filename = model_info.get("comfy_text_encoder") + if comfy_filename: + model_path = folder_paths.get_full_path_or_raise("text_encoders", comfy_filename) + signature = (model_name, "comfy", model_path) + if keep_model_loaded and self.model is not None and self.current_signature == signature: + return + self.clear() + from comfy import sd as comfy_sd + self.model = comfy_sd.load_clip( + ckpt_paths=[model_path], + embedding_directory=folder_paths.get_folder_paths("embeddings"), + ) + self.current_signature = signature + print(f"[QwenVL] Loaded shared ComfyUI text encoder: {comfy_filename}") + return + quant = Quantization.from_value(quant_value) # Skip enforce_memory for now # Safety check: ensure quant is not None @@ -1679,7 +1721,53 @@ class QwenVLBase: model_name="", stream_to_terminal=False, enable_thinking=True, + seed=0, ): + model_info = HF_ALL_MODELS.get(model_name, {}) + native_family = model_info.get("native_family") + if native_family: + if num_beams != 1: + raise ValueError("[QwenVL] ComfyUI text encoders support num_beams=1 only") + if native_family == "gemma3" and video is not None: + raise ValueError("[QwenVL] Gemma 3 shared text encoders support text and image input, not video") + + tokenize_kwargs = {"skip_template": False, "thinking": bool(enable_thinking)} + if native_family == "gemma3": + tokenize_kwargs["image"] = image[:1] if image is not None else None + else: + images = [] + if image is not None: + if image.dim() == 4 and image.shape[0] > 1: + print(f"[QwenVL] IMAGE input contains {image.shape[0]} items; using the first item only. Use the video input for multi-frame analysis.") + images.append(image[:1] if image.dim() == 4 else image) + if video is not None: + frame_indexes = list(range(video.shape[0])) + if len(frame_indexes) > frame_count: + frame_indexes = np.linspace(0, len(frame_indexes) - 1, frame_count, dtype=int).tolist() + images.extend(video[index:index + 1] for index in frame_indexes) + tokenize_kwargs["images"] = images + + throw_exception_if_processing_interrupted() + tokens = self.model.tokenize(prompt_text, **tokenize_kwargs) + output_tokens = self.model.generate( + tokens, + do_sample=True, + max_length=max_tokens, + temperature=temperature, + top_k=50, + top_p=top_p, + min_p=0.0, + repetition_penalty=repetition_penalty, + seed=seed, + presence_penalty=0.0, + ) + throw_exception_if_processing_interrupted() + raw_text = self.model.decode(output_tokens).strip() + if stream_to_terminal: + print("[QwenVL] ComfyUI native generation completed (live token streaming is unavailable):") + print(raw_text) + return raw_text, raw_text + # Memory optimization: clear cache before generation ensure_cuda_vram_headroom("QwenVL", min_free_gb=1.0, min_free_ratio=0.08) supports_soft_think = getattr(self, "supports_qwen_soft_think", False) @@ -1931,6 +2019,7 @@ class QwenVLBase: model_name=model_name, stream_to_terminal=stream_to_terminal, enable_thinking=enable_thinking, + seed=seed, ) # Cache the generated text diff --git a/README.md b/README.md index e3ab3c9..4fb29df 100644 --- a/README.md +++ b/README.md @@ -112,6 +112,15 @@ The HF path is the most straightforward option for Windows and Linux when you wa HF tokenizer note: some Gemma and Qwen variants need `sentencepiece` or `tiktoken` available in the ComfyUI Python environment. ThinkingLLM now declares those dependencies and will attempt to install them automatically if a tokenizer backend is missing at first load. +#### Shared ComfyUI text encoders + +Compatible single-file checkpoints already in `ComfyUI/models/text_encoders` appear in the model selector with a `[ComfyUI]` prefix. ThinkingLLM loads them through ComfyUI, so it does not download or duplicate a Transformers snapshot. + +- Qwen3-VL 4B and 8B `.safetensors`, including `qwen3vl_8b_fp8_scaled.safetensors` +- Gemma 3 12B `.safetensors`, including `gemma-3-12b-it-heretic-v2_int8.safetensors` + +The MiniMax H3/Qwen3-VL 32B conditioning checkpoint is intentionally excluded because it is not a complete text-generation model. The valid SafeTensors extension is `.safetensors` (plural). ComfyUI-native generation supports one-beam sampling; terminal output is printed after generation rather than streamed token by token. + ### GGUF / llama.cpp nodes - `ThinkingLLM (GGUF)` diff --git a/tests/test_new_features.py b/tests/test_new_features.py index 6c4d105..4b633c2 100644 --- a/tests/test_new_features.py +++ b/tests/test_new_features.py @@ -116,6 +116,8 @@ def build_loader_test_stubs() -> dict[str, types.ModuleType]: models_dir=str(PKG), folder_names_and_paths={}, get_folder_paths=lambda *args, **kwargs: [], + get_filename_list=lambda *args, **kwargs: [], + get_full_path_or_raise=lambda category, filename: str(PKG / filename), ), "AILab_StreamDisplay": build_stub_module( "AILab_StreamDisplay", @@ -768,6 +770,65 @@ class TestPromptEnhancerMetadata(unittest.TestCase): self.assertIn(marker, read_source(filename)) +class TestSharedComfyUITextEncoders(unittest.TestCase): + def test_discovers_and_runs_supported_native_text_encoders(self): + filenames = [ + "qwen3vl_4b_convrot.safetensors", + "qwen3vl_8b_fp8_scaled.safetensors", + "gemma-3-12b-it-heretic-v2_int8.safetensors", + "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", + "qwen3vl_32b_minimax_h3_nvfp4_awq.safetensor", + "clip_l.safetensors", + ] + calls = {} + + class FakeClip: + def tokenize(self, text, **kwargs): + calls["tokenize"] = (text, kwargs) + return "tokens" + + def generate(self, tokens, **kwargs): + calls["generate"] = (tokens, kwargs) + return [1, 2, 3] + + def decode(self, tokens): + return "native answer" + + def load_clip(**kwargs): + calls["load"] = kwargs + return FakeClip() + + stub_modules = build_loader_test_stubs() + stub_modules.pop("AILab_QwenVL", None) + stub_modules["folder_paths"].get_filename_list = lambda category: filenames if category == "text_encoders" else [] + stub_modules["folder_paths"].get_full_path_or_raise = lambda category, filename: str(PKG / filename) + comfy_sd = build_stub_module("comfy.sd", load_clip=load_clip) + stub_modules["comfy"].sd = comfy_sd + stub_modules["comfy.sd"] = comfy_sd + + with mock.patch.dict(sys.modules, stub_modules, clear=False): + module = load_module_from_file("AILab_QwenVL.py", "thinkingllm_native_text_encoder_test") + qwen_name = "[ComfyUI] qwen3vl_8b_fp8_scaled.safetensors" + gemma_name = "[ComfyUI] gemma-3-12b-it-heretic-v2_int8.safetensors" + + self.assertIn("[ComfyUI] qwen3vl_4b_convrot.safetensors", module.HF_VL_MODELS) + self.assertIn(qwen_name, module.HF_VL_MODELS) + self.assertIn(gemma_name, module.HF_VL_MODELS) + self.assertNotIn("[ComfyUI] qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", module.HF_VL_MODELS) + self.assertNotIn("[ComfyUI] qwen3vl_32b_minimax_h3_nvfp4_awq.safetensor", module.HF_VL_MODELS) + + node = module.QwenVLBase() + node.load_model(qwen_name, "unused", "auto", False, "auto", True) + text, raw = node.generate("hello", None, None, 4, 32, 0.7, 0.9, 1, 1.1, model_name=qwen_name, seed=7) + + self.assertEqual(("native answer", "native answer"), (text, raw)) + self.assertEqual([str(PKG / "qwen3vl_8b_fp8_scaled.safetensors")], calls["load"]["ckpt_paths"]) + self.assertTrue(calls["tokenize"][1]["thinking"]) + self.assertEqual(7, calls["generate"][1]["seed"]) + with self.assertRaisesRegex(ValueError, "not video"): + node.generate("hello", None, object(), 4, 32, 0.7, 0.9, 1, 1.1, model_name=gemma_name) + + class TestHuggingFaceTokenSupport(unittest.TestCase): TOKEN_NODE_FILES = [ "AILab_QwenVL.py",