diff --git a/comfy/model_management.py b/comfy/model_management.py index 222005b6f..e3b43636b 100644 --- a/comfy/model_management.py +++ b/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: diff --git a/comfy/ops.py b/comfy/ops.py index 13c2604fb..d48589f6a 100644 --- a/comfy/ops.py +++ b/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: diff --git a/comfy/sd.py b/comfy/sd.py index 4a0742e7a..b50fe6d83 100644 --- a/comfy/sd.py +++ b/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(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) + 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(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) + 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: diff --git a/comfy_api/torch_helpers/torch_compile.py b/comfy_api/torch_helpers/torch_compile.py index 9223f58db..a935a3e75 100644 --- a/comfy_api/torch_helpers/torch_compile.py +++ b/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: diff --git a/tests-unit/comfy_api_test/torch_compile_test.py b/tests-unit/comfy_api_test/torch_compile_test.py new file mode 100644 index 000000000..e9a17bb34 --- /dev/null +++ b/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