xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-04 11:47:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

fixes for directml, use_pytorch_cross_attention and channels_last args

Browse Source 
workaround for torch.count_nonzero on DirectML
This commit is contained in:
loxotron 2025-05-16 19:45:50 +03:00
parent aaed282c3a
commit 07b066c510
3 changed files with 140 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

137
comfy/model_management.py
View File

@ -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."""

19
comfy/samplers.py
View File

@ -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,

37
fast_sampler.py
View File

@ -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: