diff --git a/comfy/model_management.py b/comfy/model_management.py index 65599424b21..15c03dc771d 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -490,28 +490,36 @@ def amd_min_version(device=None, min_rdna_version=0): except: rocm_version = (6, -1) - def aotriton_supported(gpu_arch): - path = torch.__path__[0] - path = os.path.join(os.path.join(path, "lib"), "aotriton.images") - gfx = set(map(lambda a: a[4:], filter(lambda a: a.startswith("amd-gfx"), os.listdir(path)))) - if gpu_arch in gfx: - return True - if "{}x".format(gpu_arch[:-1]) in gfx: - return True - if "{}xx".format(gpu_arch[:-2]) in gfx: - return True - return False + def aotriton_supported(): + """Whether pytorch reports flash attention as usable on this gpu. + + can_use_flash_attention() evaluates runtime eligibility for the given + parameters; on a ROCm build that includes checking the gpu arch against the + kernel images AOTriton was compiled for. Querying it avoids assuming where + those images live inside the torch install. The probe tensor is shaped and + typed to pass the unrelated SDPA checks, so False means no hardware support + rather than a rejected shape. + """ + try: + if not torch.backends.cuda.is_flash_attention_available(): # not built with flash attention + return False + q = torch.empty((1, 1, 8, 64), dtype=torch.float16, device=get_torch_device()) + params = torch.backends.cuda.SDPAParams(q, q, q, None, 0.0, False, False) + return torch.backends.cuda.can_use_flash_attention(params, False) + except (AttributeError, RuntimeError, TypeError) as e: + logging.warning("Could not query aotriton support: {}".format(e)) + return False logging.info("AMD arch: {}".format(arch)) logging.info("ROCm version: {}".format(rocm_version)) if args.use_split_cross_attention == False and args.use_quad_cross_attention == False: - if aotriton_supported(arch): # AMD efficient attention implementation depends on aotriton. + if aotriton_supported(): # AMD efficient attention implementation depends on aotriton. if torch_version_numeric >= (2, 7): # works on 2.6 but doesn't actually seem to improve much if any((a in arch) for a in ["gfx90a", "gfx942", "gfx950", "gfx1100", "gfx1101", "gfx1150", "gfx1151"]): # TODO: more arches, TODO: gfx950 ENABLE_PYTORCH_ATTENTION = True if rocm_version >= (7, 0): - if any((a in arch) for a in ["gfx1200", "gfx1201"]): - ENABLE_PYTORCH_ATTENTION = True + if any((a in arch) for a in ["gfx1200", "gfx1201"]): + ENABLE_PYTORCH_ATTENTION = True if torch_version_numeric >= (2, 7) and rocm_version >= (6, 4): if any((a in arch) for a in ["gfx1200", "gfx1201", "gfx950"]): # TODO: more arches, "gfx942" gives error on pytorch nightly 2.10 1013 rocm7.0 SUPPORT_FP8_OPS = True diff --git a/comfy/text_encoders/gemma4.py b/comfy/text_encoders/gemma4.py index 5b0b968c93e..606f8993e5c 100644 --- a/comfy/text_encoders/gemma4.py +++ b/comfy/text_encoders/gemma4.py @@ -1183,6 +1183,7 @@ def _get_aspect_ratio_preserving_size(height, width, patch_size, max_patches, po class Gemma4_Tokenizer(): tokenizer_json_data = None + prime_empty_thought = False def state_dict(self): if self.tokenizer_json_data is not None: @@ -1333,8 +1334,8 @@ def tokenize_with_weights(self, text, return_word_ids=False, image=None, audio=N num_samples = int(waveform.shape[-1] * 16000 / sample_rate) if sample_rate != 16000 else waveform.shape[-1] n_audio_tokens = self._audio_token_count(num_samples) media += "<|audio>" + "<|audio|>" * n_audio_tokens + "" - # Non-thinking mode primes an empty thought channel so the model answers directly. - model_open = "" if thinking else "<|channel>thought\n" + # 12B/31B prime a closed thought block for non-thinking mode, E2B/E4B must not: it cues them into reasoning inline. + model_open = "<|channel>thought\n" if self.prime_empty_thought and not thinking else "" llama_text = f"{system}<|turn>user\n{text}{media}\n<|turn>model\n{model_open}" text_tokens = super().tokenize_with_weights(llama_text, return_word_ids) @@ -1418,6 +1419,7 @@ def __init__(self, embedding_directory=None, tokenizer_data={}): class Gemma4UnifiedSDTokenizer(Gemma4SDTokenizer): """Encoder-free (gemma4_unified) audio: raw 16kHz waveform frames instead of mel spectrogram.""" embedding_size = 3840 + prime_empty_thought = True def _extract_audio_features(self, waveform, sample_rate): audio = self._resample_16k(waveform, sample_rate) @@ -1500,7 +1502,7 @@ def __init__(self, device="cpu", dtype=None, model_options={}): # Variants -def _make_variant(config_cls): +def _make_variant(config_cls, prime_empty_thought=False): audio = config_cls.audio_config is not None bases = (Gemma4AudioMixin, Gemma4Base) if audio else (Gemma4Base,) class Variant(*bases): @@ -1510,8 +1512,8 @@ def __init__(self, config_dict, dtype, device, operations): if audio: self._init_audio(self.model.config, dtype, device, operations) embedding_size = config_cls.hidden_size - if embedding_size != Gemma4SDTokenizer.embedding_size: - tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size}) + if embedding_size != Gemma4SDTokenizer.embedding_size or prime_empty_thought: + tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size, 'prime_empty_thought': prime_empty_thought}) class Tokenizer(Gemma4Tokenizer): tokenizer_class = tok_cls Variant.tokenizer = Tokenizer @@ -1521,7 +1523,7 @@ class Tokenizer(Gemma4Tokenizer): Gemma4_E4B = _make_variant(Gemma4Config) Gemma4_E2B = _make_variant(Gemma4_E2B_Config) -Gemma4_31B = _make_variant(Gemma4_31B_Config) +Gemma4_31B = _make_variant(Gemma4_31B_Config, prime_empty_thought=True) # Gemma4 12B Unified: encoder-free multimodal, distinct base/tokenizer (not via _make_variant). diff --git a/main.py b/main.py index 361b1fc899a..eb64726ea4e 100644 --- a/main.py +++ b/main.py @@ -248,7 +248,7 @@ def execute_script(script_path): import comfy.memory_management import comfy.model_patcher -if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_management.is_nvidia() and not comfy.model_management.is_wsl()): +if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_management.is_nvidia()): if (not args.enable_dynamic_vram) and (comfy.model_management.torch_version_numeric < (2, 8)): logging.warning("Unsupported Pytorch detected. DynamicVRAM support requires Pytorch version 2.8 or later. Falling back to legacy ModelPatcher. VRAM estimates may be unreliable especially on Windows") else: diff --git a/requirements.txt b/requirements.txt index f461e3b76a2..4f505cc9c2f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -22,7 +22,7 @@ alembic SQLAlchemy>=2.0.0 filelock av>=16.0.0 -comfy-kitchen==0.2.30 +comfy-kitchen==0.2.31 comfy-aimdo==0.4.13 requests simpleeval>=1.0.0 diff --git a/tests-unit/comfy_test/gemma4_template_test.py b/tests-unit/comfy_test/gemma4_template_test.py new file mode 100644 index 00000000000..77e274653f7 --- /dev/null +++ b/tests-unit/comfy_test/gemma4_template_test.py @@ -0,0 +1,61 @@ +"""Gemma4 chat template regression tests.""" + +import pytest +import torch + +from comfy.cli_args import args + +if not torch.cuda.is_available(): + args.cpu = True + +import comfy.text_encoders.gemma4 as gemma4 # noqa: E402 + +PROMPT = "describe a cute anime girl with fennec ears" +THOUGHT_BLOCK = "<|channel>thought\n" + +# E2B/E4B and 12B/31B ship different canonical chat templates: only the latter prime a +# closed thought block when thinking is off. +NO_PRIMING = [gemma4.Gemma4_E2B, gemma4.Gemma4_E4B] +PRIMING = [gemma4.Gemma4_31B, gemma4.Gemma4_12B] + + +class _CaptureTemplate: + """Stands in for SDTokenizer.tokenize_with_weights so the built template is checked without model files.""" + llama_text = "" + + def tokenize_with_weights(self, text, return_word_ids=False, **kwargs): + self.llama_text = text + return {} + + +def build_template(variant, **kwargs): + prime = variant.tokenizer.tokenizer_class.prime_empty_thought + probe = type("Probe", (gemma4.Gemma4_Tokenizer, _CaptureTemplate), {"prime_empty_thought": prime})() + probe.tokenize_with_weights(PROMPT, **kwargs) + return probe.llama_text + + +@pytest.mark.parametrize("variant", NO_PRIMING + PRIMING) +def test_thinking_enabled_only_asks_via_the_system_turn(variant): + template = build_template(variant, skip_template=False, thinking=True) + assert template == f"<|turn>system\n<|think|>\n\n<|turn>user\n{PROMPT}\n<|turn>model\n" + + +@pytest.mark.parametrize("variant", NO_PRIMING) +def test_thinking_disabled_does_not_prime_a_thought_channel(variant): + template = build_template(variant, skip_template=False, thinking=False) + assert template == f"<|turn>user\n{PROMPT}\n<|turn>model\n" + assert "channel" not in template + assert "<|think|>" not in template + + +@pytest.mark.parametrize("variant", PRIMING) +def test_thinking_disabled_primes_a_thought_channel(variant): + template = build_template(variant, skip_template=False, thinking=False) + assert template == f"<|turn>user\n{PROMPT}\n<|turn>model\n{THOUGHT_BLOCK}" + + +@pytest.mark.parametrize("variant", NO_PRIMING + PRIMING) +@pytest.mark.parametrize("thinking", [False, True]) +def test_skip_template_passes_text_through_unchanged(variant, thinking): + assert build_template(variant, skip_template=True, thinking=thinking) == PROMPT