Improve Windows ROCm inference handling

comfy/model_management.py

@@ -395,6 +395,7 @@ def raise_non_oom(e):
 
 XFORMERS_VERSION = ""
 XFORMERS_ENABLED_VAE = True
+ENABLE_PYTORCH_VAE_ON_AMD = "COMFYUI_ENABLE_PYTORCH_VAE_ON_AMD"
 if args.disable_xformers:
     XFORMERS_IS_AVAILABLE = False
 else:
@@ -1628,9 +1629,14 @@ def pytorch_attention_enabled():
 
 def pytorch_attention_enabled_vae():
     if is_amd():
-        return False  # enabling pytorch attention on AMD currently causes crash when doing high res
+        if os.getenv(ENABLE_PYTORCH_VAE_ON_AMD) == "1":
+            return hasattr(torch.nn.functional, "scaled_dot_product_attention")
+        return False  # enabling pytorch attention on AMD can corrupt high-res VAE decode
     return pytorch_attention_enabled()
 
+def pytorch_attention_vae_single_batch():
+    return sys.platform == "win32" and is_amd() and pytorch_attention_enabled_vae()
+
 def pytorch_attention_flash_attention():
     global ENABLE_PYTORCH_ATTENTION
     if ENABLE_PYTORCH_ATTENTION:

comfy/ops.py

@@ -41,7 +41,7 @@ def scaled_dot_product_attention(q, k, v, *args, **kwargs):
 
 
 try:
-    if torch.cuda.is_available() and comfy.model_management.WINDOWS:
+    if torch.cuda.is_available() and comfy.model_management.WINDOWS and comfy.model_management.is_nvidia():
         from torch.nn.attention import SDPBackend, sdpa_kernel
         import inspect
         if "set_priority" in inspect.signature(sdpa_kernel).parameters:

comfy/sd.py

@@ -1105,6 +1105,8 @@ class VAE:
                 free_memory = self.patcher.get_free_memory(self.device)
                 batch_number = int(free_memory / memory_used)
                 batch_number = max(1, batch_number)
+                if model_management.pytorch_attention_vae_single_batch():
+                    batch_number = 1
 
                 # Pre-allocate output for VAEs that support direct buffer writes
                 preallocated = False
@@ -1958,10 +1960,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
     if unet_dtype is None:
         unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=unet_weight_dtype, weight_dtype=weight_dtype)
 
-    if model_config.quant_config is not None:
+    manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
-        manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
-    else:
-        manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
     model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device)
 
     if model_config.clip_vision_prefix is not None:
@@ -2099,10 +2098,7 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None, disable
     else:
         unet_dtype = dtype
 
-    if model_config.quant_config is not None:
+    manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
-        manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
-    else:
-        manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
     model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device)
 
     if custom_operations is not None:

comfy_api/torch_helpers/torch_compile.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
+import logging
+import sys
 import torch
 
 import comfy.utils
 from comfy.patcher_extension import WrappersMP
-from typing import TYPE_CHECKING, Callable, Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional
 if TYPE_CHECKING:
     from comfy.model_patcher import ModelPatcher
     from comfy.patcher_extension import WrapperExecutor
@@ -11,6 +13,35 @@ if TYPE_CHECKING:
 
 COMPILE_KEY = "torch.compile"
 TORCH_COMPILE_KWARGS = "torch_compile_kwargs"
+WINDOWS_ROCM_INDUCTOR_OPTIONS = {
+    "triton.cudagraphs": False,
+    "triton.cudagraph_trees": False,
+}
+
+
+def _is_windows_rocm_inductor(backend: Optional[str]) -> bool:
+    return backend == "inductor" and sys.platform == "win32" and getattr(torch.version, "hip", None) is not None
+
+
+def normalize_torch_compile_kwargs(compile_kwargs: dict[str, Any]) -> dict[str, Any]:
+    compile_kwargs = dict(compile_kwargs)
+    if _is_windows_rocm_inductor(compile_kwargs.get("backend")) and compile_kwargs.get("mode") in (None, "", "default"):
+        options = dict(compile_kwargs.get("options") or {})
+        if set(options) <= {"guard_filter_fn"}:
+            compile_kwargs["mode"] = None
+            compile_kwargs["options"] = None
+            logging.info("torch.compile: using default mode for Windows ROCm inductor.")
+        else:
+            changed = False
+            for key, value in WINDOWS_ROCM_INDUCTOR_OPTIONS.items():
+                if options.get(key) is not value:
+                    options[key] = value
+                    changed = True
+            compile_kwargs["options"] = options
+            compile_kwargs["mode"] = None
+            if changed:
+                logging.info("torch.compile: disabled inductor cudagraphs for Windows ROCm.")
+    return compile_kwargs
 
 
 def apply_torch_compile_factory(compiled_module_dict: dict[str, Callable]) -> Callable:
@@ -30,7 +61,7 @@ def apply_torch_compile_factory(compiled_module_dict: dict[str, Callable]) -> Ca
     return apply_torch_compile_wrapper
 
 
-def set_torch_compile_wrapper(model: ModelPatcher, backend: str, options: Optional[dict[str,str]]=None,
+def set_torch_compile_wrapper(model: ModelPatcher, backend: str, options: Optional[dict[str, Any]]=None,
                               mode: Optional[str]=None, fullgraph=False, dynamic: Optional[bool]=None,
                               keys: list[str]=["diffusion_model"], *args, **kwargs):
     '''
@@ -52,6 +83,7 @@ def set_torch_compile_wrapper(model: ModelPatcher, backend: str, options: Option
         "fullgraph": fullgraph,
         "dynamic": dynamic,
     }
+    compile_kwargs = normalize_torch_compile_kwargs(compile_kwargs)
     # get a dict of compiled keys
     compiled_modules = {}
     for key in keys:

tests-unit/comfy_api_test/torch_compile_test.py

@@ -0,0 +1,50 @@
+import torch
+
+from comfy_api.torch_helpers import torch_compile
+
+
+def test_windows_rocm_default_mode_drops_injected_guard_options(monkeypatch):
+    monkeypatch.setattr(torch_compile.sys, "platform", "win32")
+    monkeypatch.setattr(torch.version, "hip", "7.15", raising=False)
+
+    result = torch_compile.normalize_torch_compile_kwargs(
+        {
+            "backend": "inductor",
+            "mode": "default",
+            "options": {"guard_filter_fn": object()},
+        }
+    )
+
+    assert result["mode"] is None
+    assert result["options"] is None
+
+
+def test_windows_rocm_custom_options_disable_cudagraphs(monkeypatch):
+    monkeypatch.setattr(torch_compile.sys, "platform", "win32")
+    monkeypatch.setattr(torch.version, "hip", "7.15", raising=False)
+
+    result = torch_compile.normalize_torch_compile_kwargs(
+        {
+            "backend": "inductor",
+            "mode": "default",
+            "options": {"max_autotune": True},
+        }
+    )
+
+    assert result["mode"] is None
+    assert result["options"] == {
+        "max_autotune": True,
+        "triton.cudagraphs": False,
+        "triton.cudagraph_trees": False,
+    }
+
+
+def test_non_rocm_compile_options_are_unchanged(monkeypatch):
+    monkeypatch.setattr(torch_compile.sys, "platform", "linux")
+    compile_kwargs = {
+        "backend": "inductor",
+        "mode": "default",
+        "options": {"max_autotune": True},
+    }
+
+    assert torch_compile.normalize_torch_compile_kwargs(compile_kwargs) == compile_kwargs