Stabilize Gemma4 GGUF multimodal settings

This commit is contained in:
Marcel Hofer
2026-06-20 17:30:37 +02:00
parent 2b9ca76d41
commit 7d1b9c028b
2 changed files with 175 additions and 29 deletions
+39 -19
View File
@@ -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(
+136 -10
View File
@@ -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(