[pull] master from Comfy-Org:master

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

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 + "<audio|>"
-                # Non-thinking mode primes an empty thought channel so the model answers directly.
-                model_open = "" if thinking else "<|channel>thought\n<channel|>"
+                # 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<channel|>" if self.prime_empty_thought and not thinking else ""
                 llama_text = f"{system}<|turn>user\n{text}{media}<turn|>\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).

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:

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

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<channel|>"
+
+# 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<turn|>\n<|turn>user\n{PROMPT}<turn|>\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}<turn|>\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}<turn|>\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