Compare commits

...

26 Commits

Author SHA1 Message Date
Maifee Ul Asad
c4db212dc7
Merge 6fa13d9836 into 7cf4e78335 2026-07-07 17:30:15 +06:00
Maifee Ul Asad
6fa13d9836
Merge branch 'master' into offloader-maifee 2026-07-07 17:30:11 +06:00
comfyanonymous
7cf4e78335
Delete symlink that breaks our updates. (#14803)
Some checks are pending
Detect Unreviewed Merge / detect (push) Waiting to run
Python Linting / Run Ruff (push) Waiting to run
Python Linting / Run Pylint (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Waiting to run
Execution Tests / test (macos-latest) (push) Waiting to run
Execution Tests / test (ubuntu-latest) (push) Waiting to run
Execution Tests / test (windows-latest) (push) Waiting to run
Test server launches without errors / test (push) Waiting to run
Unit Tests / test (macos-latest) (push) Waiting to run
Unit Tests / test (ubuntu-latest) (push) Waiting to run
Unit Tests / test (windows-2022) (push) Waiting to run
2026-07-06 22:24:05 -04:00
Alexis Rolland
7747c342d4
ci: add CLA Assistant workflow (#14582)
Some checks are pending
Detect Unreviewed Merge / detect (push) Waiting to run
Python Linting / Run Ruff (push) Waiting to run
Python Linting / Run Pylint (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Waiting to run
Execution Tests / test (macos-latest) (push) Waiting to run
Execution Tests / test (ubuntu-latest) (push) Waiting to run
Execution Tests / test (windows-latest) (push) Waiting to run
Test server launches without errors / test (push) Waiting to run
Unit Tests / test (macos-latest) (push) Waiting to run
Unit Tests / test (ubuntu-latest) (push) Waiting to run
Unit Tests / test (windows-2022) (push) Waiting to run
2026-07-07 06:44:19 +08:00
comfyanonymous
439bd807f8
Skip unloading dynamic model patchers in current workflow. (#14799) 2026-07-06 14:35:12 -07:00
Maifee Ul Asad
5efe0d14bf
Merge branch 'Comfy-Org:master' into offloader-maifee 2026-06-19 14:03:31 +06:00
Maifee Ul Asad
9dd4d35e80
[cleanup]: removed duplicate deps from merging; 2026-03-29 18:31:41 +06:00
Maifee Ul Asad
e9d082a228
Merge branch 'master' into offloader-maifee 2026-03-28 13:05:32 +06:00
Maifee Ul Asad
7602203696
Merge branch 'Comfy-Org:master' into offloader-maifee 2026-01-11 12:59:53 +06:00
Maifee Ul Asad
ffa7a369ba
Merge branch 'Comfy-Org:master' into offloader-maifee 2026-01-07 21:26:04 +06:00
Maifee Ul Asad
7ec1656735
Merge branch 'comfyanonymous:master' into offloader-maifee 2025-12-07 15:32:22 +06:00
Maifee Ul Asad
cee75f301a
Merge branch 'comfyanonymous:master' into offloader-maifee 2025-11-27 08:47:41 +06:00
Maifee Ul Asad
1a59686ca8
Merge branch 'comfyanonymous:master' into offloader-maifee 2025-11-25 22:09:53 +06:00
Maifee Ul Asad
6d96d26795
Merge branch 'comfyanonymous:master' into offloader-maifee 2025-11-25 22:08:51 +06:00
Maifee Ul Asad
e07a32c9b8
Merge branch 'master' into offloader-maifee 2025-11-09 17:25:46 +06:00
Maifee Ul Asad
a19f0a88e4 refactor(gds): add show_stats option; moved GDS initialization to dedicated file; 2025-10-12 00:54:04 +06:00
Maifee Ul Asad
64811809a0 docs: updated GDS environment limitations;
- works only on linux and nvidia
2025-10-11 21:51:24 +06:00
Maifee Ul Asad
529109083a docs: added GDS docs and deps documentation; 2025-10-10 22:23:13 +06:00
Maifee Ul Asad
a7be9f6fc3 review: remove GDS-related dependencies from requirements.txt 2025-10-10 22:08:58 +06:00
Maifee Ul Asad
6075c44ec8 feat(gds): add GDS-related dependencies to requirements.txt 2025-10-08 14:43:55 +06:00
Maifee Ul Asad
154b73835a feat(gds): implement GPUDirect Storage initialization based on CLI arguments 2025-10-08 14:40:59 +06:00
Maifee Ul Asad
862e7784f4 feat(gds): add nodes_gds.py to built-in extra nodes initialization 2025-10-08 14:40:41 +06:00
Maifee Ul Asad
f6b6636bf3 feat(gds): implement GDS loading fallback in load_torch_file function;
- need to work with tensorflow and other formats
 - afaik, almost all models shared now is in torch format
 - converting types should not be that big of a deal
2025-10-08 14:40:29 +06:00
Maifee Ul Asad
83b00df3f0 feat(gds): add GPUDirect Storage options for SSD-to-GPU model loading; 2025-10-08 14:39:22 +06:00
Maifee Ul Asad
5f24eb699c feat(gds): implement GPUDirect Storage for efficient model loading 2025-10-08 14:38:49 +06:00
Maifee Ul Asad
fab0954077 feat(gds): add GPUDirect Storage support for model loading and prefetching
- limited to NVIDIA GPUs only
2025-10-08 14:38:38 +06:00
11 changed files with 963 additions and 2 deletions

91
.github/workflows/cla.yml vendored Normal file
View File

@ -0,0 +1,91 @@
name: CLA Assistant
on:
issue_comment:
types: [created]
pull_request_target:
types: [opened, synchronize, closed]
permissions:
actions: write
contents: read # 'read' is enough because signatures live in a REMOTE repo
pull-requests: write
statuses: write
jobs:
cla-assistant:
runs-on: ubuntu-latest
steps:
# The CLA action normally requires every commit author in a PR to sign.
# We only want the PR author to sign, so we allowlist all other committers
# by computing them from the PR's commits and excluding the PR author.
- name: Build author-only allowlist
id: allowlist
if: >
github.event_name == 'pull_request_target' ||
(github.event_name == 'issue_comment' && github.event.issue.pull_request && (
github.event.comment.body == 'recheck' ||
github.event.comment.body == 'I have read and agree to the Contributor License Agreement'
))
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }}
PR_AUTHOR: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
BASE_ALLOWLIST: action@github.com,actions-user,ampagent,claude,comfy-pr-bot,GitHub Action,github-actions,github-actions[bot],Glary Bot,Glary-Bot,*[bot]
run: |
others=$(gh api "repos/${{ github.repository }}/pulls/${PR_NUMBER}/commits" --paginate \
--jq '.[] | (.author.login // empty), (.committer.login // empty)' \
| sort -u | grep -vix "${PR_AUTHOR}" | paste -sd, -)
if [ -n "$others" ]; then
echo "allowlist=${BASE_ALLOWLIST},${others}" >> "$GITHUB_OUTPUT"
else
echo "allowlist=${BASE_ALLOWLIST}" >> "$GITHUB_OUTPUT"
fi
- name: CLA Assistant
# Run on PR events, on "recheck" comment, or when someone posts the exact signing phrase.
# IMPORTANT: this phrase must match `custom-pr-sign-comment` below.
if: >
github.event_name == 'pull_request_target' ||
(github.event_name == 'issue_comment' && github.event.issue.pull_request && (
github.event.comment.body == 'recheck' ||
github.event.comment.body == 'I have read and agree to the Contributor License Agreement'
))
uses: contributor-assistant/github-action@ca4a40a7d1004f18d9960b404b97e5f30a505a08 # v2.6.1
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# PAT required to write to the centralized signatures repo.
PERSONAL_ACCESS_TOKEN: ${{ secrets.PERSONAL_ACCESS_TOKEN }}
with:
# Where the CLA document lives (shown to contributors)
path-to-document: https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md
# Centralized signature storage
remote-organization-name: comfy-org
remote-repository-name: comfy-cla
path-to-signatures: signatures/cla.json
branch: main
# Only the PR author must sign: bots plus every non-author committer
# are allowlisted via the "Build author-only allowlist" step above.
# *[bot] is a catch-all for any GitHub App bot account.
allowlist: ${{ steps.allowlist.outputs.allowlist }}
# Custom PR comment messages
custom-notsigned-prcomment: |
🎉 Thank you for your contribution, we really appreciate it! 🎉
Like many open source projects, we require contributors to sign our [Contributor License Agreement (CLA)](https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md). A CLA makes the ownership of contributions explicit, so contributors and the project share a clear understanding of how the code can be used. By signing, you:
- Confirm that you own your contribution.
- Keep the right to reuse your own code.
- Grant us a copyright license to include and share it within our projects.
CLAs are standard practice across major open source projects including those under the Apache Software Foundation and the Linux Foundation. Ours is based on the Apache Software Foundation's CLA. Most importantly, it would enable us to relicense the project under a more permissive license in the future, giving the project and its community greater flexibility.
✍ **To sign, please post a new comment on this PR with exactly the following text:** ✍
custom-pr-sign-comment: I have read and agree to the Contributor License Agreement
custom-allsigned-prcomment: |
✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.

View File

@ -1 +0,0 @@
AGENTS.md

View File

@ -417,6 +417,14 @@ Use `--tls-keyfile key.pem --tls-certfile cert.pem` to enable TLS/SSL, the app w
> Note: Windows users can use [alexisrolland/docker-openssl](https://github.com/alexisrolland/docker-openssl) or one of the [3rd party binary distributions](https://wiki.openssl.org/index.php/Binaries) to run the command example above.
<br/><br/>If you use a container, note that the volume mount `-v` can be a relative path so `... -v ".\:/openssl-certs" ...` would create the key & cert files in the current directory of your command prompt or powershell terminal.
## How to run heavy workflow on mid range GPU (NVIDIA-Linux)?
Use the `--enable-gds` flag to activate NVIDIA [GPUDirect Storage](https://docs.nvidia.com/gpudirect-storage/) (GDS), which allows data to be transferred directly between SSDs and GPUs. This eliminates traditional CPU-mediated data paths, significantly reducing I/O latency and CPU overhead. System RAM will still be utilized for caching to further optimize performance, along with SSD.
This feature is tested on NVIDIA GPUs on Linux based system only.
Requires: `cupy-cuda12x>=12.0.0`, `pynvml>=11.4.1`, `cudf>=23.0.0`, `numba>=0.57.0`, `nvidia-ml-py>=12.0.0`.
## Support and dev channel
[Discord](https://comfy.org/discord): Try the #help or #feedback channels.

View File

@ -160,6 +160,17 @@ parser.add_argument("--default-hashing-function", type=str, choices=['md5', 'sha
parser.add_argument("--disable-smart-memory", action="store_true", help="Force ComfyUI to agressively offload to regular ram instead of keeping models in vram when it can.")
parser.add_argument("--deterministic", action="store_true", help="Make pytorch use slower deterministic algorithms when it can. Note that this might not make images deterministic in all cases.")
# GPUDirect Storage (GDS) arguments
gds_group = parser.add_argument_group('gds', 'GPUDirect Storage options for direct SSD-to-GPU model loading')
gds_group.add_argument("--enable-gds", action="store_true", help="Enable GPUDirect Storage for direct SSD-to-GPU model loading (requires CUDA 11.4+, cuFile).")
gds_group.add_argument("--disable-gds", action="store_true", help="Explicitly disable GPUDirect Storage.")
gds_group.add_argument("--gds-min-file-size", type=int, default=100, help="Minimum file size in MB to use GDS (default: 100MB).")
gds_group.add_argument("--gds-chunk-size", type=int, default=64, help="GDS transfer chunk size in MB (default: 64MB).")
gds_group.add_argument("--gds-streams", type=int, default=4, help="Number of CUDA streams for GDS operations (default: 4).")
gds_group.add_argument("--gds-prefetch", action="store_true", help="Enable GDS prefetching for better performance.")
gds_group.add_argument("--gds-no-fallback", action="store_true", help="Disable fallback to CPU loading if GDS fails.")
gds_group.add_argument("--gds-stats", action="store_true", help="Print GDS statistics on exit.")
class PerformanceFeature(enum.Enum):
Fp16Accumulation = "fp16_accumulation"
Fp8MatrixMultiplication = "fp8_matrix_mult"

494
comfy/gds_loader.py Normal file
View File

@ -0,0 +1,494 @@
# copyright 2025 Maifee Ul Asad @ github.com/maifeeulasad
# copyright under GNU GENERAL PUBLIC LICENSE, Version 3, 29 June 2007
"""
GPUDirect Storage (GDS) Integration for ComfyUI
Direct SSD-to-GPU model loading without RAM/CPU bottlenecks
Still there will be some CPU/RAM usage, mostly for safetensors parsing and small buffers.
This module provides GPUDirect Storage functionality to load models directly
from NVMe SSDs to GPU memory, bypassing system RAM and CPU.
"""
import os
import logging
import torch
import time
from typing import Optional, Dict, Any, Union
from pathlib import Path
import safetensors
import gc
import mmap
from dataclasses import dataclass
try:
import cupy
import cupy.cuda.runtime as cuda_runtime
CUPY_AVAILABLE = True
except ImportError:
CUPY_AVAILABLE = False
logging.warning("CuPy not available. GDS will use fallback mode.")
try:
import cudf # RAPIDS for GPU dataframes
RAPIDS_AVAILABLE = True
except ImportError:
RAPIDS_AVAILABLE = False
try:
import pynvml
pynvml.nvmlInit()
NVML_AVAILABLE = True
except ImportError:
NVML_AVAILABLE = False
logging.warning("NVIDIA-ML-Py not available. GPU monitoring disabled.")
@dataclass
class GDSConfig:
"""Configuration for GPUDirect Storage"""
enabled: bool = True
min_file_size_mb: int = 100 # Only use GDS for files larger than this
chunk_size_mb: int = 64 # Size of chunks to transfer
use_pinned_memory: bool = True
prefetch_enabled: bool = True
compression_aware: bool = True
max_concurrent_streams: int = 4
fallback_to_cpu: bool = True
show_stats: bool = False # Whether to show stats on exit
class GDSError(Exception):
"""GDS-specific errors"""
pass
class GPUDirectStorage:
"""
GPUDirect Storage implementation for ComfyUI
Enables direct SSD-to-GPU transfers for model loading
"""
def __init__(self, config: Optional[GDSConfig] = None):
self.config = config or GDSConfig()
self.device = torch.cuda.current_device() if torch.cuda.is_available() else None
self.cuda_streams = []
self.pinned_buffers = {}
self.stats = {
'gds_loads': 0,
'fallback_loads': 0,
'total_bytes_gds': 0,
'total_time_gds': 0.0,
'avg_bandwidth_gbps': 0.0
}
# Initialize GDS if available
self._gds_available = self._check_gds_availability()
if self._gds_available:
self._init_gds()
else:
logging.warning("GDS not available, using fallback methods")
def _check_gds_availability(self) -> bool:
"""Check if GDS is available on the system"""
if not torch.cuda.is_available():
return False
if not CUPY_AVAILABLE:
return False
# Check for GPUDirect Storage support
try:
# Check CUDA version (GDS requires CUDA 11.4+)
cuda_version = torch.version.cuda
if cuda_version:
major, minor = map(int, cuda_version.split('.')[:2])
if major < 11 or (major == 11 and minor < 4):
logging.warning(f"CUDA {cuda_version} detected. GDS requires CUDA 11.4+")
return False
# Check if cuFile is available (part of CUDA toolkit)
try:
import cupy.cuda.cufile as cufile
# Try to initialize cuFile
cufile.initialize()
return True
except (ImportError, RuntimeError) as e:
logging.warning(f"cuFile not available: {e}")
return False
except Exception as e:
logging.warning(f"GDS availability check failed: {e}")
return False
def _init_gds(self):
"""Initialize GDS resources"""
try:
# Create CUDA streams for async operations
for i in range(self.config.max_concurrent_streams):
stream = torch.cuda.Stream()
self.cuda_streams.append(stream)
# Pre-allocate pinned memory buffers
if self.config.use_pinned_memory:
self._allocate_pinned_buffers()
logging.info(f"GDS initialized with {len(self.cuda_streams)} streams")
except Exception as e:
logging.error(f"Failed to initialize GDS: {e}")
self._gds_available = False
def _allocate_pinned_buffers(self):
"""Pre-allocate pinned memory buffers for staging"""
try:
# Allocate buffers of different sizes
buffer_sizes = [16, 32, 64, 128, 256] # MB
for size_mb in buffer_sizes:
size_bytes = size_mb * 1024 * 1024
# Allocate pinned memory using CuPy
if CUPY_AVAILABLE:
buffer = cupy.cuda.alloc_pinned_memory(size_bytes)
self.pinned_buffers[size_mb] = buffer
except Exception as e:
logging.warning(f"Failed to allocate pinned buffers: {e}")
def _get_file_size(self, file_path: str) -> int:
"""Get file size in bytes"""
return os.path.getsize(file_path)
def _should_use_gds(self, file_path: str) -> bool:
"""Determine if GDS should be used for this file"""
if not self._gds_available or not self.config.enabled:
return False
file_size_mb = self._get_file_size(file_path) / (1024 * 1024)
return file_size_mb >= self.config.min_file_size_mb
def _load_with_gds(self, file_path: str) -> Dict[str, torch.Tensor]:
"""Load model using GPUDirect Storage"""
start_time = time.time()
try:
if file_path.lower().endswith(('.safetensors', '.sft')):
return self._load_safetensors_gds(file_path)
else:
return self._load_pytorch_gds(file_path)
except Exception as e:
logging.error(f"GDS loading failed for {file_path}: {e}")
if self.config.fallback_to_cpu:
logging.info("Falling back to CPU loading")
self.stats['fallback_loads'] += 1
return self._load_fallback(file_path)
else:
raise GDSError(f"GDS loading failed: {e}")
finally:
load_time = time.time() - start_time
self.stats['total_time_gds'] += load_time
def _load_safetensors_gds(self, file_path: str) -> Dict[str, torch.Tensor]:
"""Load safetensors file using GDS"""
try:
import cupy.cuda.cufile as cufile
# Open file with cuFile for direct GPU loading
with cufile.CuFileManager() as manager:
# Memory-map the file for efficient access
with open(file_path, 'rb') as f:
# Use mmap for large files
with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mmapped_file:
# Parse safetensors header
header_size = int.from_bytes(mmapped_file[:8], 'little')
header_bytes = mmapped_file[8:8+header_size]
import json
header = json.loads(header_bytes.decode('utf-8'))
# Load tensors directly to GPU
tensors = {}
data_offset = 8 + header_size
for name, info in header.items():
if name == "__metadata__":
continue
dtype_map = {
'F32': torch.float32,
'F16': torch.float16,
'BF16': torch.bfloat16,
'I8': torch.int8,
'I16': torch.int16,
'I32': torch.int32,
'I64': torch.int64,
'U8': torch.uint8,
}
dtype = dtype_map.get(info['dtype'], torch.float32)
shape = info['shape']
start_offset = data_offset + info['data_offsets'][0]
end_offset = data_offset + info['data_offsets'][1]
# Direct GPU allocation
tensor = torch.empty(shape, dtype=dtype, device=f'cuda:{self.device}')
# Use cuFile for direct transfer
tensor_bytes = end_offset - start_offset
# Get GPU memory pointer
gpu_ptr = tensor.data_ptr()
# Direct file-to-GPU transfer
cufile.copy_from_file(
gpu_ptr,
mmapped_file[start_offset:end_offset],
tensor_bytes
)
tensors[name] = tensor
self.stats['gds_loads'] += 1
self.stats['total_bytes_gds'] += self._get_file_size(file_path)
return tensors
except Exception as e:
logging.error(f"GDS safetensors loading failed: {e}")
raise
def _load_pytorch_gds(self, file_path: str) -> Dict[str, torch.Tensor]:
"""Load PyTorch file using GDS with staging"""
try:
# For PyTorch files, we need to use a staging approach
# since torch.load doesn't support direct GPU loading
# Load to pinned memory first
with open(file_path, 'rb') as f:
file_size = self._get_file_size(file_path)
# Choose appropriate buffer or allocate new one
buffer_size_mb = min(256, max(64, file_size // (1024 * 1024)))
if buffer_size_mb in self.pinned_buffers:
pinned_buffer = self.pinned_buffers[buffer_size_mb]
else:
# Allocate temporary pinned buffer
pinned_buffer = cupy.cuda.alloc_pinned_memory(file_size)
# Read file to pinned memory
f.readinto(pinned_buffer)
# Use torch.load with map_location to specific GPU
# This will be faster due to pinned memory
state_dict = torch.load(
f,
map_location=f'cuda:{self.device}',
weights_only=True
)
self.stats['gds_loads'] += 1
self.stats['total_bytes_gds'] += file_size
return state_dict
except Exception as e:
logging.error(f"GDS PyTorch loading failed: {e}")
raise
def _load_fallback(self, file_path: str) -> Dict[str, torch.Tensor]:
"""Fallback loading method using standard approaches"""
if file_path.lower().endswith(('.safetensors', '.sft')):
# Use safetensors with device parameter
with safetensors.safe_open(file_path, framework="pt", device=f'cuda:{self.device}') as f:
return {k: f.get_tensor(k) for k in f.keys()}
else:
# Standard PyTorch loading
return torch.load(file_path, map_location=f'cuda:{self.device}', weights_only=True)
def load_model(self, file_path: str, device: Optional[torch.device] = None) -> Dict[str, torch.Tensor]:
"""
Main entry point for loading models with GDS
Args:
file_path: Path to the model file
device: Target device (if None, uses current CUDA device)
Returns:
Dictionary of tensors loaded directly to GPU
"""
if device is not None and device.type == 'cuda':
self.device = device.index or 0
if self._should_use_gds(file_path):
logging.info(f"Loading {file_path} with GDS")
return self._load_with_gds(file_path)
else:
logging.info(f"Loading {file_path} with standard method")
self.stats['fallback_loads'] += 1
return self._load_fallback(file_path)
def prefetch_model(self, file_path: str) -> bool:
"""
Prefetch model to GPU memory cache (if supported)
Args:
file_path: Path to the model file
Returns:
True if prefetch was successful
"""
if not self.config.prefetch_enabled or not self._gds_available:
return False
try:
# Basic prefetch implementation
# This would ideally use NVIDIA's GPUDirect Storage API
# to warm up the storage cache
file_size = self._get_file_size(file_path)
logging.info(f"Prefetching {file_path} ({file_size // (1024*1024)} MB)")
# Read file metadata to warm caches
with open(file_path, 'rb') as f:
# Read first and last chunks to trigger prefetch
f.read(1024 * 1024) # First 1MB
f.seek(-min(1024 * 1024, file_size), 2) # Last 1MB
f.read()
return True
except Exception as e:
logging.warning(f"Prefetch failed for {file_path}: {e}")
return False
def get_stats(self) -> Dict[str, Any]:
"""Get loading statistics"""
total_loads = self.stats['gds_loads'] + self.stats['fallback_loads']
if self.stats['total_time_gds'] > 0 and self.stats['total_bytes_gds'] > 0:
bandwidth_gbps = (self.stats['total_bytes_gds'] / (1024**3)) / self.stats['total_time_gds']
self.stats['avg_bandwidth_gbps'] = bandwidth_gbps
return {
**self.stats,
'total_loads': total_loads,
'gds_usage_percent': (self.stats['gds_loads'] / max(1, total_loads)) * 100,
'gds_available': self._gds_available,
'config': self.config.__dict__
}
def cleanup(self):
"""Clean up GDS resources"""
try:
# Clear CUDA streams
for stream in self.cuda_streams:
stream.synchronize()
self.cuda_streams.clear()
# Free pinned buffers
for buffer in self.pinned_buffers.values():
if CUPY_AVAILABLE:
cupy.cuda.free_pinned_memory(buffer)
self.pinned_buffers.clear()
# Force garbage collection
gc.collect()
torch.cuda.empty_cache()
except Exception as e:
logging.warning(f"GDS cleanup failed: {e}")
def __del__(self):
"""Destructor to ensure cleanup"""
self.cleanup()
# Global GDS instance
_gds_instance: Optional[GPUDirectStorage] = None
def get_gds_instance(config: Optional[GDSConfig] = None) -> GPUDirectStorage:
"""Get or create the global GDS instance"""
global _gds_instance
if _gds_instance is None:
_gds_instance = GPUDirectStorage(config)
return _gds_instance
def load_torch_file_gds(ckpt: str, safe_load: bool = False, device: Optional[torch.device] = None) -> Dict[str, torch.Tensor]:
"""
GDS-enabled replacement for comfy.utils.load_torch_file
Args:
ckpt: Path to checkpoint file
safe_load: Whether to use safe loading (for compatibility)
device: Target device
Returns:
Dictionary of loaded tensors
"""
gds = get_gds_instance()
try:
# Load with GDS
return gds.load_model(ckpt, device)
except Exception as e:
logging.error(f"GDS loading failed, falling back to standard method: {e}")
# Fallback to original method
import comfy.utils
return comfy.utils.load_torch_file(ckpt, safe_load=safe_load, device=device)
def prefetch_model_gds(file_path: str) -> bool:
"""Prefetch model for faster loading"""
gds = get_gds_instance()
return gds.prefetch_model(file_path)
def get_gds_stats() -> Dict[str, Any]:
"""Get GDS statistics"""
gds = get_gds_instance()
return gds.get_stats()
def configure_gds(config: GDSConfig):
"""Configure GDS settings"""
global _gds_instance
_gds_instance = GPUDirectStorage(config)
def init_gds(config: GDSConfig):
"""
Initialize GPUDirect Storage with the provided configuration
Args:
config: GDSConfig object with initialization parameters
"""
try:
# Configure GDS
configure_gds(config)
logging.info(f"GDS initialized: enabled={config.enabled}, min_size={config.min_file_size_mb}MB, streams={config.max_concurrent_streams}")
# Set up exit handler for stats if requested
if hasattr(config, 'show_stats') and config.show_stats:
import atexit
def print_gds_stats():
stats = get_gds_stats()
logging.info("=== GDS Statistics ===")
logging.info(f"Total loads: {stats['total_loads']}")
logging.info(f"GDS loads: {stats['gds_loads']} ({stats['gds_usage_percent']:.1f}%)")
logging.info(f"Fallback loads: {stats['fallback_loads']}")
logging.info(f"Total bytes via GDS: {stats['total_bytes_gds'] / (1024**3):.2f} GB")
logging.info(f"Average bandwidth: {stats['avg_bandwidth_gbps']:.2f} GB/s")
logging.info("===================")
atexit.register(print_gds_stats)
except ImportError as e:
logging.warning(f"GDS initialization failed - missing dependencies: {e}")
except Exception as e:
logging.error(f"GDS initialization failed: {e}")

View File

@ -468,6 +468,9 @@ class CLIP:
def decode(self, token_ids, skip_special_tokens=True):
return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
def is_dynamic(self):
return self.patcher.is_dynamic()
class VAE:
def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None):
if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
@ -1251,6 +1254,8 @@ class VAE:
except:
return None
def is_dynamic(self):
return self.patcher.is_dynamic()
class StyleModel:
def __init__(self, model, device="cpu"):

View File

@ -120,6 +120,18 @@ def load_safetensors(ckpt):
def load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False):
# Try GDS loading first if available and device is GPU
if device is not None and device.type == 'cuda':
try:
from . import gds_loader
gds_result = gds_loader.load_torch_file_gds(ckpt, safe_load=safe_load, device=device)
if return_metadata:
# For GDS, we return empty metadata for now (can be enhanced)
return (gds_result, {})
return gds_result
except Exception as e:
logging.debug(f"GDS loading failed, using fallback: {e}")
if device is None:
device = torch.device("cpu")
metadata = None

View File

@ -503,6 +503,21 @@ RAM_CACHE_DEFAULT_RAM_USAGE = 0.05
RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3
def all_outputs_dynamic(outputs):
if outputs is None:
return False
for output in outputs:
if isinstance(output, (list, tuple)):
if not all_outputs_dynamic(output):
return False
elif not hasattr(output, "is_dynamic") or not output.is_dynamic():
return False
return True
class RAMPressureCache(LRUCache):
def __init__(self, key_class, enable_providers=False):
@ -533,7 +548,11 @@ class RAMPressureCache(LRUCache):
for key, cache_entry in self.cache.items():
if not free_active and self.used_generation[key] == self.generation:
continue
oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
if all_outputs_dynamic(cache_entry.outputs) and self.used_generation[key] == self.generation:
continue
oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE
def scan_list_for_ram_usage(outputs):

293
comfy_extras/nodes_gds.py Normal file
View File

@ -0,0 +1,293 @@
# copyright 2025 Maifee Ul Asad @ github.com/maifeeulasad
# copyright under GNU GENERAL PUBLIC LICENSE, Version 3, 29 June 2007
"""
Enhanced model loading nodes with GPUDirect Storage support
"""
import logging
import time
import asyncio
from typing import Optional, Dict, Any
import torch
import folder_paths
import comfy.sd
import comfy.utils
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
class CheckpointLoaderGDS(ComfyNodeABC):
"""
Enhanced checkpoint loader with GPUDirect Storage support
Provides direct SSD-to-GPU loading and prefetching capabilities
"""
@classmethod
def INPUT_TYPES(s) -> InputTypeDict:
return {
"required": {
"ckpt_name": (folder_paths.get_filename_list("checkpoints"), {
"tooltip": "The name of the checkpoint (model) to load with GDS optimization."
}),
},
"optional": {
"prefetch": ("BOOLEAN", {
"default": False,
"tooltip": "Prefetch model to GPU cache for faster loading."
}),
"use_gds": ("BOOLEAN", {
"default": True,
"tooltip": "Use GPUDirect Storage if available."
}),
"target_device": (["auto", "cuda:0", "cuda:1", "cuda:2", "cuda:3", "cpu"], {
"default": "auto",
"tooltip": "Target device for model loading."
})
}
}
RETURN_TYPES = ("MODEL", "CLIP", "VAE", "STRING")
RETURN_NAMES = ("model", "clip", "vae", "load_info")
OUTPUT_TOOLTIPS = (
"The model used for denoising latents.",
"The CLIP model used for encoding text prompts.",
"The VAE model used for encoding and decoding images to and from latent space.",
"Loading information and statistics."
)
FUNCTION = "load_checkpoint_gds"
CATEGORY = "loaders/advanced"
DESCRIPTION = "Enhanced checkpoint loader with GPUDirect Storage support for direct SSD-to-GPU loading."
EXPERIMENTAL = True
def load_checkpoint_gds(self, ckpt_name: str, prefetch: bool = False, use_gds: bool = True, target_device: str = "auto"):
start_time = time.time()
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
# Determine target device
if target_device == "auto":
device = None # Let the system decide
elif target_device == "cpu":
device = torch.device("cpu")
else:
device = torch.device(target_device)
load_info = {
"file": ckpt_name,
"path": ckpt_path,
"target_device": str(device) if device else "auto",
"gds_enabled": use_gds,
"prefetch_used": prefetch
}
try:
# Prefetch if requested
if prefetch and use_gds:
try:
from comfy.gds_loader import prefetch_model_gds
prefetch_success = prefetch_model_gds(ckpt_path)
load_info["prefetch_success"] = prefetch_success
if prefetch_success:
logging.info(f"Prefetched {ckpt_name} to GPU cache")
except Exception as e:
logging.warning(f"Prefetch failed for {ckpt_name}: {e}")
load_info["prefetch_error"] = str(e)
# Load checkpoint with potential GDS optimization
if use_gds and device and device.type == 'cuda':
try:
from comfy.gds_loader import get_gds_instance
gds = get_gds_instance()
# Check if GDS should be used for this file
if gds._should_use_gds(ckpt_path):
load_info["loader_used"] = "GDS"
logging.info(f"Loading {ckpt_name} with GDS")
else:
load_info["loader_used"] = "Standard"
logging.info(f"Loading {ckpt_name} with standard method (file too small for GDS)")
except Exception as e:
logging.warning(f"GDS check failed, using standard loading: {e}")
load_info["loader_used"] = "Standard (GDS failed)"
else:
load_info["loader_used"] = "Standard"
# Load the actual checkpoint
out = comfy.sd.load_checkpoint_guess_config(
ckpt_path,
output_vae=True,
output_clip=True,
embedding_directory=folder_paths.get_folder_paths("embeddings")
)
load_time = time.time() - start_time
load_info["load_time_seconds"] = round(load_time, 3)
load_info["load_success"] = True
# Format load info as string
info_str = f"Loaded: {ckpt_name}\n"
info_str += f"Method: {load_info['loader_used']}\n"
info_str += f"Time: {load_info['load_time_seconds']}s\n"
info_str += f"Device: {load_info['target_device']}"
if "prefetch_success" in load_info:
info_str += f"\nPrefetch: {'' if load_info['prefetch_success'] else ''}"
logging.info(f"Checkpoint loaded: {ckpt_name} in {load_time:.3f}s using {load_info['loader_used']}")
return (*out[:3], info_str)
except Exception as e:
load_info["load_success"] = False
load_info["error"] = str(e)
error_str = f"Failed to load: {ckpt_name}\nError: {str(e)}"
logging.error(f"Checkpoint loading failed: {e}")
raise RuntimeError(error_str)
class ModelPrefetcher(ComfyNodeABC):
"""
Node for prefetching models to GPU cache
"""
@classmethod
def INPUT_TYPES(s) -> InputTypeDict:
return {
"required": {
"checkpoint_names": ("STRING", {
"multiline": True,
"default": "",
"tooltip": "List of checkpoint names to prefetch (one per line)."
}),
"prefetch_enabled": ("BOOLEAN", {
"default": True,
"tooltip": "Enable/disable prefetching."
})
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("prefetch_report",)
OUTPUT_TOOLTIPS = ("Report of prefetch operations.",)
FUNCTION = "prefetch_models"
CATEGORY = "loaders/advanced"
DESCRIPTION = "Prefetch multiple models to GPU cache for faster loading."
OUTPUT_NODE = True
def prefetch_models(self, checkpoint_names: str, prefetch_enabled: bool = True):
if not prefetch_enabled:
return ("Prefetching disabled",)
# Parse checkpoint names
names = [name.strip() for name in checkpoint_names.split('\n') if name.strip()]
if not names:
return ("No checkpoints specified for prefetching",)
try:
from comfy.gds_loader import prefetch_model_gds
except ImportError:
return ("GDS not available for prefetching",)
results = []
successful_prefetches = 0
for name in names:
try:
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", name)
success = prefetch_model_gds(ckpt_path)
if success:
results.append(f"{name}")
successful_prefetches += 1
else:
results.append(f"{name} (prefetch failed)")
except Exception as e:
results.append(f"{name} (error: {str(e)[:50]})")
report = f"Prefetch Report ({successful_prefetches}/{len(names)} successful):\n"
report += "\n".join(results)
return (report,)
class GDSStats(ComfyNodeABC):
"""
Node for displaying GDS statistics
"""
@classmethod
def INPUT_TYPES(s) -> InputTypeDict:
return {
"required": {
"refresh": ("BOOLEAN", {
"default": False,
"tooltip": "Refresh statistics."
})
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("stats_report",)
OUTPUT_TOOLTIPS = ("GDS statistics and performance report.",)
FUNCTION = "get_stats"
CATEGORY = "utils/advanced"
DESCRIPTION = "Display GPUDirect Storage statistics and performance metrics."
OUTPUT_NODE = True
def get_stats(self, refresh: bool = False):
try:
from comfy.gds_loader import get_gds_stats
stats = get_gds_stats()
report = "=== GPUDirect Storage Statistics ===\n\n"
# Availability
report += f"GDS Available: {'' if stats['gds_available'] else ''}\n"
# Usage statistics
report += f"Total Loads: {stats['total_loads']}\n"
report += f"GDS Loads: {stats['gds_loads']} ({stats['gds_usage_percent']:.1f}%)\n"
report += f"Fallback Loads: {stats['fallback_loads']}\n\n"
# Performance metrics
if stats['total_bytes_gds'] > 0:
gb_transferred = stats['total_bytes_gds'] / (1024**3)
report += f"Data Transferred: {gb_transferred:.2f} GB\n"
report += f"Average Bandwidth: {stats['avg_bandwidth_gbps']:.2f} GB/s\n"
report += f"Total GDS Time: {stats['total_time_gds']:.2f}s\n\n"
# Configuration
config = stats.get('config', {})
if config:
report += "Configuration:\n"
report += f"- Enabled: {config.get('enabled', 'Unknown')}\n"
report += f"- Min File Size: {config.get('min_file_size_mb', 'Unknown')} MB\n"
report += f"- Chunk Size: {config.get('chunk_size_mb', 'Unknown')} MB\n"
report += f"- Max Streams: {config.get('max_concurrent_streams', 'Unknown')}\n"
report += f"- Prefetch: {config.get('prefetch_enabled', 'Unknown')}\n"
report += f"- Fallback: {config.get('fallback_to_cpu', 'Unknown')}\n"
return (report,)
except ImportError:
return ("GDS module not available",)
except Exception as e:
return (f"Error retrieving GDS stats: {str(e)}",)
# Node mappings
NODE_CLASS_MAPPINGS = {
"CheckpointLoaderGDS": CheckpointLoaderGDS,
"ModelPrefetcher": ModelPrefetcher,
"GDSStats": GDSStats,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"CheckpointLoaderGDS": "Load Checkpoint (GDS)",
"ModelPrefetcher": "Model Prefetcher",
"GDSStats": "GDS Statistics",
}

28
main.py
View File

@ -233,6 +233,34 @@ import comfyui_version
import app.logger
import hook_breaker_ac10a0
# Initialize GPUDirect Storage if enabled
def init_gds():
"""Initialize GPUDirect Storage based on CLI arguments"""
if hasattr(args, 'disable_gds') and args.disable_gds:
logging.info("GDS explicitly disabled via --disable-gds")
return
if not hasattr(args, 'enable_gds') and not hasattr(args, 'gds_prefetch') and not hasattr(args, 'gds_stats'):
# GDS not explicitly requested, use auto-detection
return
if hasattr(args, 'enable_gds') and args.enable_gds:
from comfy.gds_loader import GDSConfig, init_gds as gds_init
config = GDSConfig(
enabled=getattr(args, 'enable_gds', False) or getattr(args, 'gds_prefetch', False),
min_file_size_mb=getattr(args, 'gds_min_file_size', 100),
chunk_size_mb=getattr(args, 'gds_chunk_size', 64),
max_concurrent_streams=getattr(args, 'gds_streams', 4),
prefetch_enabled=getattr(args, 'gds_prefetch', True),
fallback_to_cpu=not getattr(args, 'gds_no_fallback', False),
show_stats=getattr(args, 'gds_stats', False)
)
gds_init(config)
# Initialize GDS
init_gds()
import comfy.memory_management
import comfy.model_patcher

View File

@ -2466,6 +2466,7 @@ async def init_builtin_extra_nodes():
"nodes_model_patch.py",
"nodes_easycache.py",
"nodes_audio_encoder.py",
"nodes_gds.py",
"nodes_rope.py",
"nodes_logic.py",
"nodes_resolution.py",