nodes/dseditor-ComfyUI-ListHelper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Qwen_Node Speed Fix

Browse Source
This commit is contained in:
dseditor
2025-12-04 19:56:59 +08:00
parent 75abbbf4f2
commit e8df65c8f2
2 changed files with 83 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pyproject.toml
+2 -2
View File
@@ -1,7 +1,7 @@
[project]
name = "Listhelper"
description = "The ListHelper collection is a comprehensive set of custom nodes for ComfyUI that provides powerful list manipulation capabilities. This collection includes audio processing, text splitting, and number generation tools for enhanced workflow automation.And OpenRouter API Function"
version = "1.1.0"
description = "The ListHelper collection is a comprehensive set of custom nodes for ComfyUI that provides powerful list manipulation capabilities. This collection includes audio processing, text splitting, and number generation tools for enhanced workflow automation.Fix Qwen Node For LLM function"
version = "1.1.1"
license = {file = "LICENSE"}
dependencies = ["regex", "accelerate"]
qwen_inference.py
+81 -12
View File
@@ -19,6 +19,7 @@ class QwenGPUInference:
self.tokenizer = None
self.current_model_path = None
self.config_dir = None
self.current_quantization = None
@classmethod
def _get_safetensors_files(cls):
@@ -125,6 +126,14 @@ class QwenGPUInference:
}),
},
"optional": {
"use_flash_attention": ("BOOLEAN", {
"default": False,
"tooltip": "Enable FlashAttention-2 (may not improve speed for small batch inference)"
}),
"use_quantization": ("BOOLEAN", {
"default": False,
"tooltip": "Enable INT8 quantization for lower memory usage (slower inference)"
}),
"do_sample": ("BOOLEAN", {
"default": True,
"tooltip": "Enable sampling"
@@ -295,13 +304,24 @@ class QwenGPUInference:
traceback.print_exc()
return False
def _load_model(self, model_path: str, repo_id: str) -> bool:
"""Load model and tokenizer (optimized v3 - with memory management)"""
def _load_model(self, model_path: str, repo_id: str, use_quantization: bool = False, use_flash_attention: bool = False) -> bool:
"""Load model and tokenizer (optimized v4 - with memory management, quantization, and FlashAttention support)"""
try:
if self.model is not None and self.current_model_path == model_path:
print(f"Model already loaded: {os.path.basename(model_path)}")
# Check if model is already loaded with same settings
model_config_key = (model_path, use_quantization, use_flash_attention)
current_config_key = (self.current_model_path, self.current_quantization, getattr(self, 'current_flash_attention', None))
if self.model is not None and model_config_key == current_config_key:
print(f"Model already loaded: {os.path.basename(model_path)} (Quantized: {use_quantization}, FlashAttn: {use_flash_attention})")
return True
# If any setting changed, need to reload
if self.model is not None and model_config_key != current_config_key:
print(f"Model settings changed, reloading...")
self.model = None
self.tokenizer = None
self._free_gpu_memory()
try:
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
from safetensors.torch import load_file
@@ -315,8 +335,8 @@ class QwenGPUInference:
print(f"\nLoading model: {os.path.basename(model_path)}")
# Step 1: Check GPU memory
is_enough, memory_info = self._check_gpu_memory(required_gb=7.5)
# Step 1: Check GPU memory (reduced threshold from 7.5GB to 6.0GB)
is_enough, memory_info = self._check_gpu_memory(required_gb=6.0)
print(memory_info)
# Step 2: If insufficient memory, try to free up
@@ -325,7 +345,7 @@ class QwenGPUInference:
self._free_gpu_memory()
# Check again
is_enough, memory_info = self._check_gpu_memory(required_gb=7.5)
is_enough, memory_info = self._check_gpu_memory(required_gb=6.0)
print(memory_info)
if not is_enough:
@@ -352,9 +372,12 @@ class QwenGPUInference:
# Determine loading strategy
if torch.cuda.is_available():
free_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3 - torch.cuda.memory_reserved(0) / 1024**3
# Use allocated memory instead of reserved for more accurate free memory calculation
free_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3 - torch.cuda.memory_allocated(0) / 1024**3
if free_memory >= 7.5:
# Lower threshold to 6.0GB - model actually needs ~7.5GB but can work with less free space
# This prevents unnecessary CPU offload when other models are loaded
if free_memory >= 6.0:
# Sufficient memory: Full GPU loading
print(f"Strategy: Full GPU loading (Free: {free_memory:.2f}GB)")
max_memory_config = None
@@ -400,10 +423,52 @@ class QwenGPUInference:
"pretrained_model_name_or_path": temp_model_dir,
"trust_remote_code": True,
"device_map": "auto",
"torch_dtype": dtype,
"low_cpu_mem_usage": True
}
# Add FlashAttention configuration (only if enabled by user)
if use_flash_attention:
if torch.cuda.is_available():
try:
import flash_attn
load_kwargs["attn_implementation"] = "flash_attention_2"
print(f"FlashAttention: Enabled (v{flash_attn.__version__})")
except ImportError:
print("=" * 70)
print("WARNING: FlashAttention-2 is enabled but 'flash-attn' is not installed")
print("Install with: pip install flash-attn")
print("Falling back to standard attention (no performance impact)")
print("=" * 70)
else:
print("FlashAttention: Disabled (GPU required)")
# Add quantization configuration (only if enabled by user)
if use_quantization:
if torch.cuda.is_available():
try:
from transformers import BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(
load_in_8bit=True,
llm_int8_threshold=6.0,
llm_int8_has_fp16_weight=False,
)
load_kwargs["quantization_config"] = quantization_config
print("Quantization: Enabled (INT8) - ~45% memory reduction")
except ImportError:
print("=" * 70)
print("ERROR: Quantization is enabled but 'bitsandbytes' is not installed")
print("Install with: pip install bitsandbytes")
print("Falling back to FP16 (using more memory)")
print("=" * 70)
load_kwargs["torch_dtype"] = dtype
else:
print("Quantization: Disabled (GPU required)")
load_kwargs["torch_dtype"] = dtype
else:
# Default: use FP16 on GPU, FP32 on CPU
load_kwargs["torch_dtype"] = dtype
if max_memory_config is not None:
load_kwargs["max_memory"] = max_memory_config
@@ -423,6 +488,8 @@ class QwenGPUInference:
pass
self.current_model_path = model_path
self.current_quantization = use_quantization
self.current_flash_attention = use_flash_attention
total_time = time.time() - overall_start
if torch.cuda.is_available():
@@ -449,11 +516,13 @@ class QwenGPUInference:
system_prompt: str,
max_new_tokens: int,
temperature: float,
use_flash_attention: bool = False,
use_quantization: bool = False,
do_sample: bool = True,
top_p: float = 0.9,
top_k: int = 50,
) -> Tuple[str]:
"""Execute inference"""
"""Execute inference with optional FlashAttention-2 and INT8 quantization"""
# Auto-find Qwen model
model_path = self._find_qwen_model()
@@ -469,7 +538,7 @@ class QwenGPUInference:
# Use fixed repo_id
repo_id = "Qwen/Qwen3-4B"
if not self._load_model(model_path, repo_id):
if not self._load_model(model_path, repo_id, use_quantization, use_flash_attention):
error_msg = "Error: Model loading failed. Please check the model file and ensure it's properly placed in the text_encoders or clip folder."
print(error_msg)
return (error_msg,)