diff --git a/AILab_QwenVL_GGUF.py b/AILab_QwenVL_GGUF.py index 3e2b32f..1a9a303 100644 --- a/AILab_QwenVL_GGUF.py +++ b/AILab_QwenVL_GGUF.py @@ -44,6 +44,7 @@ from AILab_LlamaCppInstaller import ( get_last_llama_cpp_backend_info, relax_windows_dll_directory_for_long_paths, ) +import comfy.model_management as comfy_model_management from comfy.model_management import throw_exception_if_processing_interrupted # Import cache functions from main module @@ -78,9 +79,9 @@ LAST_SAVED_PROMPT = None # kept only to avoid ImportError in legacy importers # Load per-node state at startup for per-node prompt and trace access load_node_prompt_state() -_GEMMA4_CPU_MULTIMODAL_MAX_CTX = 32768 -_GEMMA4_CPU_MULTIMODAL_MAX_BATCH = 256 -_GEMMA4_CPU_MULTIMODAL_MAX_IMAGE_TOKENS = 2048 +_GEMMA4_MULTIMODAL_MAX_CTX = 32768 +_GEMMA4_MULTIMODAL_IMAGE_TOKENS = 1120 +_GEMMA4_MULTIMODAL_MAX_BATCH = 2048 def read_gguf_architecture(filepath: Path) -> str | None: @@ -882,6 +883,12 @@ def _current_cuda_device_index() -> int | None: try: if not torch.cuda.is_available(): return None + try: + comfy_device = comfy_model_management.get_torch_device() + if getattr(comfy_device, "type", None) == "cuda" and getattr(comfy_device, "index", None) is not None: + return int(comfy_device.index) + except Exception: + pass return int(torch.cuda.current_device()) except Exception: return None @@ -1125,6 +1132,7 @@ class QwenVLGGUFBase: self.uses_qwen_template_thinking = False self.current_backend_gpu_offload = None self.current_batch_size = None + self.current_ubatch_size = None self.current_image_token_budget = None self.current_context_length = None self.gguf_arch = None @@ -1365,16 +1373,19 @@ class QwenVLGGUFBase: self.supports_qwen_soft_think = arch == "qwen3" if arch else "qwen3-" in model_name.lower() self.uses_qwen_template_thinking = bool(is_qwen35 or arch == "qwen3" or arch == "qwen") - if arch == "gemma4" and has_mmproj and gpu_offload is False: + if arch == "gemma4" and has_mmproj: original_settings = (n_ctx, n_batch_val, n_ubatch_val, img_max) - n_ctx = min(n_ctx, _GEMMA4_CPU_MULTIMODAL_MAX_CTX) - n_batch_val = min(n_batch_val, _GEMMA4_CPU_MULTIMODAL_MAX_BATCH) - n_ubatch_val = min(n_ubatch_val, n_batch_val) - img_max = min(img_max, _GEMMA4_CPU_MULTIMODAL_MAX_IMAGE_TOKENS) + n_ctx = min(n_ctx, _GEMMA4_MULTIMODAL_MAX_CTX) + img_max = min(img_max, _GEMMA4_MULTIMODAL_IMAGE_TOKENS) + required_batch = max(img_max, 1) + n_batch_val = max(n_batch_val, required_batch) + n_batch_val = min(n_batch_val, _GEMMA4_MULTIMODAL_MAX_BATCH) + n_ubatch_val = max(n_ubatch_val, required_batch) + n_ubatch_val = min(n_ubatch_val, n_batch_val, _GEMMA4_MULTIMODAL_MAX_BATCH) downgraded_settings = (n_ctx, n_batch_val, n_ubatch_val, img_max) if downgraded_settings != original_settings: load_warnings.append( - "downgraded unsafe Gemma 4 multimodal settings for CPU-only llama.cpp backend: " + "applied Gemma 4 multimodal compatibility settings: " f"ctx {original_settings[0]}->{n_ctx}, " f"n_batch {original_settings[1]}->{n_batch_val}, " f"n_ubatch {original_settings[2]}->{n_ubatch_val}, " @@ -1468,6 +1479,7 @@ class QwenVLGGUFBase: self.current_context_length = n_ctx self.current_image_token_budget = img_max self.current_batch_size = n_batch_val + self.current_ubatch_size = n_ubatch_val print( f"[QwenVL] Loading GGUF: {model_path.name} " @@ -1541,24 +1553,32 @@ class QwenVLGGUFBase: arch = str(getattr(self, "gguf_arch", "") or "").lower() if arch != "gemma4": return - if getattr(self, "current_backend_gpu_offload", None) is not False: - return + if audio_b64: + raise RuntimeError( + "Gemma 4 GGUF audio input is disabled because the installed llama.cpp backend can abort the " + "ComfyUI process for Gemma 4 audio. Use image/text input with Gemma 4 GGUF, or use a backend " + "that explicitly supports Gemma 4 audio." + ) context_length = int(getattr(self, "current_context_length", 0) or 0) batch_size = int(getattr(self, "current_batch_size", 0) or 0) + ubatch_size = int(getattr(self, "current_ubatch_size", 0) or 0) image_tokens = int(getattr(self, "current_image_token_budget", 0) or 0) if ( - context_length <= _GEMMA4_CPU_MULTIMODAL_MAX_CTX - and batch_size <= _GEMMA4_CPU_MULTIMODAL_MAX_BATCH - and image_tokens <= _GEMMA4_CPU_MULTIMODAL_MAX_IMAGE_TOKENS + context_length <= _GEMMA4_MULTIMODAL_MAX_CTX + and image_tokens <= _GEMMA4_MULTIMODAL_IMAGE_TOKENS + and batch_size >= max(image_tokens, 1) + and ubatch_size >= max(image_tokens, 1) + and batch_size <= _GEMMA4_MULTIMODAL_MAX_BATCH + and ubatch_size <= _GEMMA4_MULTIMODAL_MAX_BATCH ): return raise RuntimeError( "Gemma 4 GGUF multimodal generation was blocked before llama.cpp could abort ComfyUI. " - "The installed llama-cpp-python backend reports no GPU offload support, while this run uses " - f"ctx={context_length}, batch={batch_size}, image_max_tokens={image_tokens}. " - "Install a CUDA-enabled llama-cpp-python vision backend, or retry with device=cpu, " - f"ctx<={_GEMMA4_CPU_MULTIMODAL_MAX_CTX}, n_batch<={_GEMMA4_CPU_MULTIMODAL_MAX_BATCH}, " - f"and image_max_tokens<={_GEMMA4_CPU_MULTIMODAL_MAX_IMAGE_TOKENS}." + "The installed llama-cpp-python backend can abort with unsafe Gemma 4 image settings; this run uses " + f"ctx={context_length}, batch={batch_size}, ubatch={ubatch_size}, image_max_tokens={image_tokens}. " + f"Retry with ctx<={_GEMMA4_MULTIMODAL_MAX_CTX}, " + f"image_max_tokens<={_GEMMA4_MULTIMODAL_IMAGE_TOKENS}, and n_batch/n_ubatch large enough " + "to hold the image token chunk." ) def _invoke( diff --git a/tests/test_new_features.py b/tests/test_new_features.py index e165e42..dcdb81a 100644 --- a/tests/test_new_features.py +++ b/tests/test_new_features.py @@ -79,6 +79,7 @@ def build_loader_test_stubs() -> dict[str, types.ModuleType]: pil_package = build_stub_module("PIL", Image=pil_image) comfy_model_management = build_stub_module( "comfy.model_management", + get_torch_device=lambda: types.SimpleNamespace(type="cuda", index=0), throw_exception_if_processing_interrupted=lambda: None, ) comfy_package = build_stub_module("comfy", model_management=comfy_model_management) @@ -2228,6 +2229,7 @@ class TestGGUFWindllRuntimeFallback(unittest.TestCase): ) stub_modules = build_loader_test_stubs() stub_modules["AILab_LlamaCppInstaller"] = installer_module + stub_modules["comfy.model_management"].get_torch_device = lambda: types.SimpleNamespace(type="cuda", index=0) with tempfile.TemporaryDirectory() as temp_dir: model_path = Path(temp_dir) / "model.gguf" @@ -2282,10 +2284,10 @@ class TestGGUFWindllRuntimeFallback(unittest.TestCase): unique_id=None, ) - self.assertEqual(captured_llama_kwargs[0]["main_gpu"], 1) - self.assertIn("main_gpu=1", loader.last_backend_trace) + self.assertEqual(captured_llama_kwargs[0]["main_gpu"], 0) + self.assertIn("main_gpu=0", loader.last_backend_trace) - def test_gemma4_cpu_only_multimodal_load_downgrades_unsafe_settings(self): + def test_gemma4_multimodal_load_uses_compatible_image_ubatch_settings(self): captured_chat_kwargs = [] captured_llama_kwargs = [] @@ -2369,15 +2371,105 @@ class TestGGUFWindllRuntimeFallback(unittest.TestCase): unique_id=None, ) - self.assertEqual(captured_chat_kwargs[0]["image_max_tokens"], 2048) + self.assertEqual(captured_chat_kwargs[0]["image_max_tokens"], 1120) self.assertEqual(captured_llama_kwargs[0]["n_ctx"], 32768) - self.assertEqual(captured_llama_kwargs[0]["n_batch"], 256) - self.assertEqual(captured_llama_kwargs[0]["n_ubatch"], 256) - self.assertEqual(captured_llama_kwargs[0]["image_max_tokens"], 2048) + self.assertEqual(captured_llama_kwargs[0]["n_batch"], 1120) + self.assertEqual(captured_llama_kwargs[0]["n_ubatch"], 1120) + self.assertEqual(captured_llama_kwargs[0]["image_max_tokens"], 1120) self.assertEqual(loader.current_context_length, 32768) - self.assertEqual(loader.current_batch_size, 256) - self.assertEqual(loader.current_image_token_budget, 2048) - self.assertIn("downgraded unsafe Gemma 4 multimodal settings", loader.last_backend_trace) + self.assertEqual(loader.current_batch_size, 1120) + self.assertEqual(loader.current_image_token_budget, 1120) + self.assertIn("applied Gemma 4 multimodal compatibility settings", loader.last_backend_trace) + + def test_gemma4_gpu_multimodal_load_clamps_crash_log_settings(self): + captured_chat_kwargs = [] + captured_llama_kwargs = [] + + Gemma4ChatHandler = type( + "Gemma4ChatHandler", + (), + {"__init__": lambda self, **kwargs: captured_chat_kwargs.append(dict(kwargs))}, + ) + + class DummyLlama: + def __init__(self, **kwargs): + captured_llama_kwargs.append(dict(kwargs)) + + chat_format_module = build_stub_module( + "llama_cpp.llama_chat_format", + Gemma4ChatHandler=Gemma4ChatHandler, + ) + installer_module = build_stub_module( + "AILab_LlamaCppInstaller", + ensure_llama_cpp_backend=lambda *args, **kwargs: DummyLlama, + format_llama_cpp_backend_info=lambda info=None: "stub backend", + get_last_llama_cpp_backend_info=lambda: {"gpu_offload": True, "vision_handlers": ["Gemma4ChatHandler"]}, + relax_windows_dll_directory_for_long_paths=contextlib.nullcontext, + ) + stub_modules = build_loader_test_stubs() + stub_modules.update( + { + "AILab_LlamaCppInstaller": installer_module, + "llama_cpp": build_stub_module("llama_cpp", llama_chat_format=chat_format_module), + "llama_cpp.llama_chat_format": chat_format_module, + } + ) + + with tempfile.TemporaryDirectory() as temp_dir: + model_path = Path(temp_dir) / "google_gemma-4-E4B-it-Q8_0.gguf" + mmproj_path = Path(temp_dir) / "mmproj-F16.gguf" + model_path.write_bytes(b"GGUF") + mmproj_path.write_bytes(b"mmproj") + + with mock.patch.dict(sys.modules, stub_modules, clear=False): + module = load_module_from_file( + "AILab_QwenVL_GGUF.py", + "thinkingllm_gemma4_gpu_load_downgrade_test", + ) + loader = module.QwenVLGGUFBase() + resolved = module.GGUFVLResolved( + display_name="gemma", + repo_id=None, + alt_repo_ids=[], + author=None, + repo_dirname="gemma", + model_filename=str(model_path), + mmproj_filename=str(mmproj_path), + context_length=131072, + image_max_tokens=4096, + n_batch=512, + gpu_layers=-1, + top_k=64, + pool_size=4194304, + ) + + with mock.patch.object(module, "_resolve_model_entry", return_value=resolved): + with mock.patch.object(module, "_pick_device", return_value="cuda"): + with mock.patch.object(module, "read_gguf_architecture", return_value="gemma4"): + loader._load_model( + model_name="google_gemma-4-E4B-it-Q8_0.gguf", + device="cuda", + ctx=None, + n_batch=None, + n_ubatch=None, + gpu_layers=-1, + image_max_tokens=None, + top_k=None, + pool_size=None, + n_threads=None, + n_threads_batch=None, + flash_attn=True, + offload_kqv=True, + ctx_checkpoints=0, + enable_thinking=True, + unique_id=None, + ) + + self.assertEqual(captured_chat_kwargs[0]["image_max_tokens"], 1120) + self.assertEqual(captured_llama_kwargs[0]["n_ctx"], 32768) + self.assertEqual(captured_llama_kwargs[0]["n_batch"], 1120) + self.assertEqual(captured_llama_kwargs[0]["n_ubatch"], 1120) + self.assertEqual(captured_llama_kwargs[0]["image_max_tokens"], 1120) def test_gemma4_cpu_only_multimodal_guard_blocks_native_abort_risk(self): with mock.patch.dict(sys.modules, build_loader_test_stubs(), clear=False): @@ -2412,6 +2504,40 @@ class TestGGUFWindllRuntimeFallback(unittest.TestCase): top_k=64, ) + def test_gemma4_audio_guard_blocks_unsupported_native_audio_path(self): + with mock.patch.dict(sys.modules, build_loader_test_stubs(), clear=False): + module = load_module_from_file( + "AILab_QwenVL_GGUF.py", + "thinkingllm_gemma4_audio_abort_guard_test", + ) + + loader = module.QwenVLGGUFBase() + loader.chat_handler = object() + loader.gguf_arch = "gemma4" + loader.current_backend_gpu_offload = True + loader.current_context_length = 32768 + loader.current_batch_size = 1120 + loader.current_ubatch_size = 1120 + loader.current_image_token_budget = 1120 + + with self.assertRaisesRegex(RuntimeError, "Gemma 4 GGUF audio input is disabled"): + loader._invoke( + system_prompt="You are helpful.", + user_prompt="Transcribe this audio.", + images_b64=[], + audio_b64=["BASE64WAV"], + max_tokens=256, + temperature=0.0, + top_p=1.0, + repetition_penalty=1.0, + seed=1, + model_name="google_gemma-4-E4B-it-Q8_0.gguf", + stream_to_terminal=False, + enable_thinking=False, + auto_finalization_retry=False, + top_k=64, + ) + def test_gguf_thinking_toggle_updates_without_reload_signature(self): with mock.patch.dict(sys.modules, build_loader_test_stubs(), clear=False): module = load_module_from_file(