fixes for directml, use_pytorch_cross_attention and channels_last args

workaround for torch.count_nonzero on DirectML

comfy/model_management.py

@@ -114,13 +114,34 @@ def is_directml_enabled():
     return False
 
 def get_supported_float8_types():
-    """Get supported float8 data types."""
+    """Get supported float8 data types available in the current PyTorch version."""
     float8_types = []
-    for dtype in [torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz, torch.float8_e8m0fnu]:
+    # List of potential float8 type names to check
+    float8_type_names = [
+        'float8_e4m3fn',
+        'float8_e4m3fnuz',
+        'float8_e5m2',
+        'float8_e5m2fnuz',
+        'float8_e8m0fnu',
+    ]
+    
+    for dtype_name in float8_type_names:
         try:
+            # Check if the dtype exists in torch module
+            if hasattr(torch, dtype_name):
+                dtype = getattr(torch, dtype_name)
+                # Verify that the dtype is a valid torch.dtype
+                if isinstance(dtype, torch.dtype):
                     float8_types.append(dtype)
-        except:
+        except Exception as e:
+            # Log the error only in debug mode to avoid clutter
+            if DEBUG_ENABLED:
+                logging.debug(f"Failed to access torch.{dtype_name}: {str(e)}")
             pass
+    
+    if DEBUG_ENABLED:
+        logging.debug(f"Supported float8 types: {[str(dtype) for dtype in float8_types]}")
+    
     return float8_types
 
 def get_directml_vram(dev):
@@ -191,7 +212,11 @@ def get_directml_vram(dev):
 FLOAT8_TYPES = get_supported_float8_types()
 XFORMERS_IS_AVAILABLE = False
 XFORMERS_ENABLED_VAE = True
-ENABLE_PYTORCH_ATTENTION = True  # Enable PyTorch attention for better performance
+ENABLE_PYTORCH_ATTENTION = False
+if args.use_pytorch_cross_attention:
+    ENABLE_PYTORCH_ATTENTION = True
+    XFORMERS_IS_AVAILABLE = False
+
 FORCE_FP32 = args.force_fp32
 DISABLE_SMART_MEMORY = args.disable_smart_memory
 
@@ -437,7 +462,7 @@ def flash_attention_enabled():
 def pytorch_attention_enabled():
     """Check if PyTorch attention is enabled."""
     global ENABLE_PYTORCH_ATTENTION
-    return ENABLE_PYTORCH_ATTENTION or not (xformers_enabled() or sage_attention_enabled() or flash_attention_enabled())
+    return ENABLE_PYTORCH_ATTENTION
 
 def pytorch_attention_enabled_vae():
     """Check if PyTorch attention is enabled for VAE."""
@@ -502,29 +527,37 @@ class OOM_EXCEPTION(Exception):
     """Exception raised for out-of-memory errors."""
     pass
 
-if args.use_pytorch_cross_attention:
-    ENABLE_PYTORCH_ATTENTION = True
-    XFORMERS_IS_AVAILABLE = False
 MIN_WEIGHT_MEMORY_RATIO = 0.4 if is_nvidia() else 0.0
+
+try:
     if is_nvidia() and torch_version_numeric[0] >= 2:
         if not (ENABLE_PYTORCH_ATTENTION or args.use_split_cross_attention or args.use_quad_cross_attention):
             ENABLE_PYTORCH_ATTENTION = True
     elif is_intel_xpu() or is_ascend_npu() or is_mlu():
         if not (args.use_split_cross_attention or args.use_quad_cross_attention):
             ENABLE_PYTORCH_ATTENTION = True
-elif is_amd() and torch_version_numeric[0] >= 2 and torch_version_numeric[1] >= 7:
+    elif is_amd() and torch_version_numeric[0] >= 2 and torch_version_numeric[1] >= 7:  # works on 2.6 but doesn't actually seem to improve much
         arch = torch.cuda.get_device_properties(get_torch_device()).gcnArchName
         logging.info(f"AMD arch: {arch}")
-    if any(a in arch for a in ["gfx1100", "gfx1101"]) and not (args.use_split_cross_attention or args.use_quad_cross_attention):
+        if any(a in arch for a in ["gfx1100", "gfx1101", "gfx1030", "gfx1031", "gfx1032"]) and not (args.use_split_cross_attention or args.use_quad_cross_attention):
             ENABLE_PYTORCH_ATTENTION = True
-if ENABLE_PYTORCH_ATTENTION:
+except:
+    pass
+    
+if ENABLE_PYTORCH_ATTENTION and not directml_enabled:
+    try:
         torch.backends.cuda.enable_math_sdp(True)
         torch.backends.cuda.enable_flash_sdp(True)
         torch.backends.cuda.enable_mem_efficient_sdp(True)
         if torch_version_numeric[0] == 2 and torch_version_numeric[1] >= 5:
             torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(True)
-else:
+        elif DEBUG_ENABLED:
-    logging.warning("Could not set allow_fp16_bf16_reduction_math_sdp")
+            logging.debug("Could not set allow_fp16_bf16_reduction_math_sdp due to PyTorch version < 2.5")
+    except Exception as e:
+        if DEBUG_ENABLED:
+            logging.debug(f"Failed to enable CUDA SDP optimizations: {str(e)}")
+elif directml_enabled and DEBUG_ENABLED:
+    logging.debug("Skipped CUDA-specific SDP optimizations (math_sdp, flash_sdp, mem_efficient_sdp, allow_fp16_bf16_reduction_math_sdp) for DirectML")
 
 def get_free_memory(dev=None, torch_free_too=False):
     """
@@ -1350,6 +1383,10 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
         return False
     if args.force_fp16:
         return supports_cast(torch.float16, device)
+    if directml_enabled:
+        if DEBUG_ENABLED:
+            logging.debug("should_use_fp16: DirectML detected, disabling FP16 due to potential instability")
+        return False
     if is_intel_xpu():
         return True
     if is_mlu():
@@ -1358,12 +1395,19 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
     if is_ascend_npu():
         return False
     if is_amd():
+        try:
             arch = torch.cuda.get_device_properties(device).gcnArchName
             if any(a in arch for a in ["gfx1030", "gfx1031", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]):
                 return manual_cast
             return True
-    props = torch.cuda.get_device_properties(device)
+        except AssertionError:
+            # Fallback for non-CUDA AMD GPUs (e.g., via DirectML)
+            if DEBUG_ENABLED:
+                logging.debug("should_use_fp16: Fallback to False for AMD GPU without CUDA")
+            return False
     if is_nvidia():
+        try:
+            props = torch.cuda.get_device_properties(device)
             # Prefer FP32 for low VRAM or older GPUs
             total_vram = get_total_memory(device) / (1024**3)
             if total_vram < 5.9 or props.major <= 7:  # Turing (7.5) or Pascal (6.x)
@@ -1373,6 +1417,15 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
             if props.major >= 8:
                 return True
             return torch.cuda.is_bf16_supported() and manual_cast and (not prioritize_performance or model_params * 4 > get_total_memory(device))
+        except AssertionError:
+            # Fallback for non-CUDA NVIDIA GPUs
+            if DEBUG_ENABLED:
+                logging.debug("should_use_fp16: Fallback to False for NVIDIA GPU without CUDA")
+            return False
+    # Fallback for other devices
+    if DEBUG_ENABLED:
+        logging.debug("should_use_fp16: Fallback to False for unknown device")
+    return False
 
 def should_use_bf16(device=None, model_params=0, prioritize_performance=True, manual_cast=False):
     """Determine if BF16 should be used for the device."""

comfy/samplers.py

@@ -10,6 +10,7 @@ import torch
 from functools import partial
 import collections
 from comfy import model_management
+from comfy.cli_args import args
 import math
 import logging
 import comfy.sampler_helpers
@@ -19,6 +20,7 @@ import comfy.hooks
 import scipy.stats
 import numpy
 
+DEBUG_ENABLED = args.debug
 
 def add_area_dims(area, num_dims):
     while (len(area) // 2) < num_dims:
@@ -942,15 +944,28 @@ class CFGGuider:
         return sampling_function(self.inner_model, x, timestep, self.conds.get("negative", None), self.conds.get("positive", None), self.cfg, model_options=model_options, seed=seed)
 
     def inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed):
-        if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image.
+        # Workaround for torch.count_nonzero on DirectML
+        if latent_image is not None:
+            if model_management.is_directml_enabled():
+                nonzero_count = torch.sum(latent_image != 0).item()
+                if DEBUG_ENABLED:
+                    logging.debug(f"inner_sample: DirectML count_nonzero replacement: nonzero_count={nonzero_count}")
+            else:
+                nonzero_count = torch.count_nonzero(latent_image).item()
+            if nonzero_count > 0:  # Don't shift the empty latent image
                 latent_image = self.inner_model.process_latent_in(latent_image)
+        else:
+            nonzero_count = 0
 
+        # Process conditions
         self.conds = process_conds(self.inner_model, noise, self.conds, device, latent_image, denoise_mask, seed)
 
+        # Clone model options and add sample sigmas
         extra_model_options = comfy.model_patcher.create_model_options_clone(self.model_options)
         extra_model_options.setdefault("transformer_options", {})["sample_sigmas"] = sigmas
         extra_args = {"model_options": extra_model_options, "seed": seed}
 
+        # Execute sampler with wrappers
         executor = comfy.patcher_extension.WrapperExecutor.new_class_executor(
             sampler.sample,
             sampler,

fast_sampler.py

@@ -4,7 +4,7 @@ import gc
 import time
 from torch.amp import autocast
 from comfy.cli_args import args
-from comfy.model_management import get_torch_device, vae_dtype, soft_empty_cache, free_memory, force_channels_last, estimate_vae_decode_memory, device_supports_non_blocking
+from comfy.model_management import get_torch_device, vae_dtype, soft_empty_cache, free_memory, force_channels_last, estimate_vae_decode_memory, device_supports_non_blocking, directml_enabled
 from contextlib import contextmanager
 import latent_preview
 import logging
@@ -150,7 +150,12 @@ def finalize_images(images, device):
     """Process and finalize output images."""
     if len(images.shape) == 5:  # Combine batches
         images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1])
-    return images.to(device=device, memory_format=torch.channels_last)
+    # Apply channels_last only for CUDA devices if force_channels_last is enabled
+    is_gpu = device.type == 'cuda' and torch.cuda.is_available()
+    memory_format = torch.channels_last if (is_gpu and not directml_enabled and force_channels_last()) else torch.contiguous_format
+    if DEBUG_ENABLED:
+        logging.debug(f"finalize_images: Using memory_format={memory_format} for device={device}, directml_enabled={directml_enabled}")
+    return images.to(device=device, memory_format=memory_format)
 
 def fast_sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
                 denoise, disable_noise, start_step, last_step, force_full_denoise, noise_mask, callback, seed, device, dtype, is_gpu):
@@ -170,7 +175,11 @@ def fast_sample(model, noise, steps, cfg, sampler_name, scheduler, positive, neg
                 force_full_denoise=force_full_denoise,
                 noise_mask=noise_mask, callback=callback, seed=seed
             )
-            samples = samples.to(device=device, dtype=dtype, memory_format=torch.channels_last)
+            # Apply channels_last only for CUDA devices if force_channels_last is enabled
+            memory_format = torch.channels_last if (is_gpu and not directml_enabled and force_channels_last()) else torch.contiguous_format
+            if DEBUG_ENABLED:
+                logging.debug(f"fast_sample: Using memory_format={memory_format} for device={device}, directml_enabled={directml_enabled}")
+            samples = samples.to(device=device, dtype=dtype, memory_format=memory_format)
     
     if PROFILING_ENABLED:
         logging.debug(f"Sampling completed, took {time.time() - start_time:.3f} s")
@@ -330,15 +339,19 @@ def fast_vae_decode(vae, samples):
         # Preload VAE to device
         preload_model(vae, device, is_vae=True)
 
-        # Transfer latents with channels_last
+        # Transfer latents with appropriate memory format
         with profile_section("VAE latent transfer"):
             non_blocking = is_gpu and device_supports_non_blocking(device)
             latent_samples = samples["samples"].to(device, dtype=vae_dtype_val, non_blocking=non_blocking)
-            if is_gpu and force_channels_last():
+            # Apply channels_last only for CUDA devices if force_channels_last is enabled
+            memory_format = torch.channels_last if (is_gpu and not directml_enabled and force_channels_last()) else torch.contiguous_format
+            if is_gpu and memory_format == torch.channels_last:
+                if DEBUG_ENABLED:
+                    logging.debug(f"fast_vae_decode: Using memory_format={memory_format} for device={device}, directml_enabled={directml_enabled}")
                 latent_samples = latent_samples.to(memory_format=torch.channels_last)
                 vae.first_stage_model.to(memory_format=torch.channels_last)
-            if PROFILING_ENABLED:
+            elif DEBUG_ENABLED:
-                logging.debug(f"Latent samples device: {latent_samples.device}, dtype: {latent_samples.dtype}")
+                logging.debug(f"fast_vae_decode: Using memory_format={memory_format} for device={device}, directml_enabled={directml_enabled}")
 
         # Decode latents
         with torch.no_grad():
@@ -408,15 +421,21 @@ def fast_vae_tiled_decode(vae, samples, tile_size=512, overlap=64, temporal_size
             logging.debug(f"VAE preload took {time.time() - preload_start:.3f} s")
             logging.debug(f"Post-preload checkpoint: {time.time()}")
 
-        # Transfer latents
+        # Transfer latents with appropriate memory format
         with profile_section("VAE latent transfer"):
             latent_samples = samples["samples"]
             if PROFILING_ENABLED:
                 logging.debug(f"Latent samples device: {latent_samples.device}, dtype: {latent_samples.dtype}")
             latent_samples = optimized_transfer(latent_samples, device, vae_dtype_val)
-            if is_gpu and force_channels_last():
+            # Apply channels_last only for CUDA devices if force_channels_last is enabled
+            memory_format = torch.channels_last if (is_gpu and not directml_enabled and force_channels_last()) else torch.contiguous_format
+            if is_gpu and memory_format == torch.channels_last:
+                if DEBUG_ENABLED:
+                    logging.debug(f"fast_vae_tiled_decode: Using memory_format={memory_format} for device={device}, directml_enabled={directml_enabled}")
                 latent_samples = latent_samples.to(memory_format=torch.channels_last)
                 vae.first_stage_model.to(memory_format=torch.channels_last)
+            elif DEBUG_ENABLED:
+                logging.debug(f"fast_vae_tiled_decode: Using memory_format={memory_format} for device={device}, directml_enabled={directml_enabled}")
 
         # Log before decoding
         if PROFILING_ENABLED: