From 07b066c5108f033c983cc656741d298513c968e8 Mon Sep 17 00:00:00 2001 From: loxotron Date: Fri, 16 May 2025 19:45:50 +0300 Subject: [PATCH] fixes for directml, use_pytorch_cross_attention and channels_last args workaround for torch.count_nonzero on DirectML --- comfy/model_management.py | 137 ++++++++++++++++++++++++++------------ comfy/samplers.py | 19 +++++- fast_sampler.py | 37 +++++++--- 3 files changed, 140 insertions(+), 53 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 7bc3e5075..dac5985dc 100644 --- a/comfy/model_management.py +++ b/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: - float8_types.append(dtype) - except: + # 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 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 -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: - 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): - ENABLE_PYTORCH_ATTENTION = True -if ENABLE_PYTORCH_ATTENTION: - 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: - logging.warning("Could not set allow_fp16_bf16_reduction_math_sdp") + +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: # 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", "gfx1030", "gfx1031", "gfx1032"]) and not (args.use_split_cross_attention or args.use_quad_cross_attention): + ENABLE_PYTORCH_ATTENTION = True +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) + elif DEBUG_ENABLED: + 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,21 +1395,37 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma if is_ascend_npu(): return False if is_amd(): - 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) - if is_nvidia(): - # 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) + 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 + 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 any(platform.win32_ver()) and props.major <= 7: - return manual_cast and torch.cuda.is_bf16_supported() - 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)) + 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) + return False + if any(platform.win32_ver()) and props.major <= 7: + return manual_cast and torch.cuda.is_bf16_supported() + 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.""" diff --git a/comfy/samplers.py b/comfy/samplers.py index 67ae09a25..a672d870a 100644 --- a/comfy/samplers.py +++ b/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. - latent_image = self.inner_model.process_latent_in(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, diff --git a/fast_sampler.py b/fast_sampler.py index 47a8efc6d..a821d7895 100644 --- a/fast_sampler.py +++ b/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: - logging.debug(f"Latent samples device: {latent_samples.device}, dtype: {latent_samples.dtype}") + elif DEBUG_ENABLED: + 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: