diff --git a/qwen_inference.py b/qwen_inference.py
index 0172726..5c1bb7e 100644
--- a/qwen_inference.py
+++ b/qwen_inference.py
@@ -1,615 +1,561 @@
-import os
-import torch
-import folder_paths
-import requests
-import time
-import re
-import gc
-from typing import Optional, Tuple, Dict
-
-class QwenGPUInference:
- """
- Qwen3-4B GPU Inference Node with intelligent memory management
- Auto-downloads required config files and performs GPU inference
- Includes GPU memory checking and cleanup to avoid conflicts with ComfyUI's CLIP models
- """
-
- def __init__(self):
- self.model = None
- self.tokenizer = None
- self.current_model_path = None
- self.config_dir = None
- self.current_quantization = None
-
- @classmethod
- def _get_safetensors_files(cls):
- """Get all safetensors files from text_encoders and clip folders"""
- safetensors_files = []
-
- # Search in text_encoders folder
- try:
- text_encoder_paths = folder_paths.get_folder_paths("text_encoders")
- for path in text_encoder_paths:
- if os.path.exists(path):
- for file in os.listdir(path):
- if file.lower().endswith('.safetensors'):
- full_path = os.path.join(path, file)
- if full_path not in safetensors_files:
- safetensors_files.append(full_path)
- except:
- pass
-
- # Search in clip folder
- try:
- clip_paths = folder_paths.get_folder_paths("clip")
- for path in clip_paths:
- if os.path.exists(path):
- for file in os.listdir(path):
- if file.lower().endswith('.safetensors'):
- full_path = os.path.join(path, file)
- if full_path not in safetensors_files:
- safetensors_files.append(full_path)
- except:
- pass
-
- if not safetensors_files:
- return ["No safetensors files found"]
-
- return sorted(safetensors_files)
-
- @classmethod
- def _get_prompt_templates(cls):
- """Get all .md template files from Prompt folder"""
- current_dir = os.path.dirname(os.path.abspath(__file__))
- prompt_dir = os.path.join(current_dir, "Prompt")
-
- templates = []
-
- if os.path.exists(prompt_dir):
- for file in os.listdir(prompt_dir):
- if file.lower().endswith('.md'):
- templates.append(file)
-
- if not templates:
- return ["No Template"]
-
- return sorted(templates)
-
- @classmethod
- def _load_template_content(cls, template_name):
- """Load template content"""
- if template_name == "No Template" or template_name == "Custom":
- return ""
-
- current_dir = os.path.dirname(os.path.abspath(__file__))
- template_path = os.path.join(current_dir, "Prompt", template_name)
-
- if os.path.exists(template_path):
- try:
- with open(template_path, 'r', encoding='utf-8') as f:
- return f.read()
- except:
- return ""
-
- return ""
-
- @classmethod
- def INPUT_TYPES(cls):
- templates = cls._get_prompt_templates()
- # Add "Custom" option at the beginning of the list
- template_options = ["Custom"] + templates
-
- return {
- "required": {
- "user_prompt": ("STRING", {
- "multiline": True,
- "default": "A girl in a coffee shop"
- }),
- "prompt_template": (template_options, {
- "default": template_options[0] if template_options else "Custom"
- }),
- "system_prompt": ("STRING", {
- "multiline": True,
- "default": ""
- }),
- "max_new_tokens": ("INT", {
- "default": 2048,
- "min": 1,
- "max": 4096,
- "step": 1
- }),
- "temperature": ("FLOAT", {
- "default": 0.7,
- "min": 0.0,
- "max": 2.0,
- "step": 0.1
- }),
- },
- "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"
- }),
- "top_p": ("FLOAT", {
- "default": 0.9,
- "min": 0.0,
- "max": 1.0,
- "step": 0.05
- }),
- "top_k": ("INT", {
- "default": 50,
- "min": 0,
- "max": 100,
- "step": 1
- }),
- }
- }
-
- RETURN_TYPES = ("STRING",)
- RETURN_NAMES = ("text",)
- FUNCTION = "inference"
- CATEGORY = "ListHelper"
-
- def _find_qwen_model(self) -> Optional[str]:
- """Auto-find qwen_3_4b.safetensors model"""
- safetensors_files = self._get_safetensors_files()
-
- # Prioritize finding qwen_3_4b.safetensors
- for path in safetensors_files:
- if path != "No safetensors files found":
- basename = os.path.basename(path).lower()
- if "qwen" in basename and "3" in basename and "4b" in basename:
- return path
-
- # If specific model not found, return first safetensors file
- if safetensors_files and safetensors_files[0] != "No safetensors files found":
- return safetensors_files[0]
-
- return None
-
- def _remove_thinking_tags(self, text: str) -> str:
- """Remove ... tags and their content"""
- # Use regex to remove all ... blocks
- cleaned_text = re.sub(r'.*?', '', text, flags=re.DOTALL)
- # Remove extra blank lines
- cleaned_text = re.sub(r'\n\s*\n', '\n', cleaned_text)
- return cleaned_text.strip()
-
- def _check_gpu_memory(self, required_gb: float = 8.0) -> Tuple[bool, str]:
- """
- Check if GPU memory is sufficient
-
- Args:
- required_gb: Required GPU memory size (GB)
-
- Returns:
- (is_sufficient, detailed_message)
- """
- if not torch.cuda.is_available():
- return True, "GPU Memory: Using CPU mode"
-
- try:
- total_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
- allocated_memory = torch.cuda.memory_allocated(0) / 1024**3
- reserved_memory = torch.cuda.memory_reserved(0) / 1024**3
- free_memory = total_memory - reserved_memory
-
- info = f"GPU Memory: Total {total_memory:.2f}GB | Free {free_memory:.2f}GB | Required {required_gb:.2f}GB"
-
- if free_memory < required_gb:
- return False, info + f" | Insufficient: need {required_gb - free_memory:.2f}GB more"
-
- return True, info + " | Sufficient"
-
- except Exception as e:
- return True, f"GPU Memory: Cannot check - {e}"
-
- def _free_gpu_memory(self) -> None:
- """
- Free GPU memory
- Clear PyTorch cache and run garbage collection
- """
- try:
- if torch.cuda.is_available():
- # Record memory before cleanup
- before_reserved = torch.cuda.memory_reserved(0) / 1024**3
-
- # Clear CUDA cache
- torch.cuda.empty_cache()
- torch.cuda.synchronize()
-
- # Force garbage collection
- gc.collect()
-
- # Clear again
- torch.cuda.empty_cache()
-
- # Record memory after cleanup
- after_reserved = torch.cuda.memory_reserved(0) / 1024**3
-
- freed_reserved = before_reserved - after_reserved
-
- print(f"GPU Memory: Freed {freed_reserved:.2f}GB | Current reserved {after_reserved:.2f}GB")
- else:
- gc.collect()
- print("Memory cleanup: CPU mode")
-
- except Exception as e:
- print(f"Memory cleanup error: {e}")
- # Attempt garbage collection even if error occurs
- gc.collect()
-
- def _download_config_files(self, model_path: str, repo_id: str) -> bool:
- """Auto-download HuggingFace config files"""
- try:
- model_dir = os.path.dirname(model_path)
- model_basename = os.path.splitext(os.path.basename(model_path))[0]
- self.config_dir = os.path.join(model_dir, f"{model_basename}_config")
-
- config_files = [
- "config.json",
- "generation_config.json",
- "merges.txt",
- "tokenizer.json",
- "tokenizer_config.json",
- "vocab.json"
- ]
-
- all_exist = all(os.path.exists(os.path.join(self.config_dir, f)) for f in config_files)
-
- if all_exist:
- print(f"Config files exist: {self.config_dir}")
- return True
-
- os.makedirs(self.config_dir, exist_ok=True)
- print(f"Downloading config files to: {self.config_dir}")
-
- base_url = f"https://huggingface.co/{repo_id}/resolve/main/"
-
- for filename in config_files:
- filepath = os.path.join(self.config_dir, filename)
-
- if os.path.exists(filepath):
- print(f" {filename} exists")
- continue
-
- url = base_url + filename
- print(f" Downloading {filename}...")
-
- try:
- response = requests.get(url, timeout=30)
- response.raise_for_status()
-
- with open(filepath, 'wb') as f:
- f.write(response.content)
- print(f" {filename} downloaded")
-
- except Exception as e:
- print(f" Failed to download {filename}: {str(e)}")
- continue
-
- return True
-
- except Exception as e:
- print(f"Failed to download config files: {e}")
- import traceback
- traceback.print_exc()
- return False
-
- 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:
- # 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
- except ImportError as e:
- print(f"Missing required packages: {e}")
- print("Please run: pip install transformers safetensors")
- return False
-
- if not self._download_config_files(model_path, repo_id):
- return False
-
- print(f"\nLoading model: {os.path.basename(model_path)}")
-
- # 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
- if not is_enough:
- print("Memory insufficient, cleaning up...")
- self._free_gpu_memory()
-
- # Check again
- is_enough, memory_info = self._check_gpu_memory(required_gb=6.0)
- print(memory_info)
-
- if not is_enough:
- print("Still insufficient, will use CPU Offload strategy (slower inference)")
-
- overall_start = time.time()
-
- device = "cuda:0" if torch.cuda.is_available() else "cpu"
- dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
- # 1. Load tokenizer
- self.tokenizer = AutoTokenizer.from_pretrained(
- self.config_dir,
- trust_remote_code=True
- )
-
- # 2. Load config
- config = AutoConfig.from_pretrained(self.config_dir, trust_remote_code=True)
-
- # 3. Choose loading strategy based on memory
- import shutil
- temp_model_dir = os.path.join(self.config_dir, "temp_model")
- os.makedirs(temp_model_dir, exist_ok=True)
-
- # Determine loading strategy
- if torch.cuda.is_available():
- # 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
-
- # 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
- offload_folder = None
- else:
- # Insufficient memory: Use CPU Offload
- available_gpu = max(3.0, free_memory - 1.0) # Reserve at least 1GB
- print(f"Strategy: CPU Offload (Free: {free_memory:.2f}GB, Allocate: {available_gpu:.2f}GB)")
- max_memory_config = {
- 0: f"{available_gpu:.1f}GB",
- "cpu": "16GB"
- }
- # Create offload folder
- offload_folder = os.path.join(self.config_dir, "offload")
- os.makedirs(offload_folder, exist_ok=True)
- else:
- print("Strategy: CPU mode")
- max_memory_config = None
- offload_folder = None
-
- load_start = time.time()
-
- try:
- # Copy config files
- for file in ["config.json", "generation_config.json"]:
- src = os.path.join(self.config_dir, file)
- if os.path.exists(src):
- shutil.copy(src, temp_model_dir)
-
- # Create link or copy safetensors
- safetensors_target = os.path.join(temp_model_dir, "model.safetensors")
- if os.path.exists(safetensors_target):
- os.remove(safetensors_target)
-
- # Windows uses hard links instead of symbolic links
- try:
- os.link(model_path, safetensors_target)
- except:
- shutil.copy(model_path, safetensors_target)
-
- # Use appropriate loading strategy
- load_kwargs = {
- "pretrained_model_name_or_path": temp_model_dir,
- "trust_remote_code": True,
- "device_map": "auto",
- "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
-
- if offload_folder is not None:
- load_kwargs["offload_folder"] = offload_folder
-
- self.model = AutoModelForCausalLM.from_pretrained(**load_kwargs)
-
- print(f"Model loaded (Time: {time.time() - load_start:.2f}s)")
-
- finally:
- # Clean up temp directory
- try:
- if os.path.exists(temp_model_dir):
- shutil.rmtree(temp_model_dir)
- except:
- 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():
- current_allocated = torch.cuda.memory_allocated(0) / 1024**3
- current_reserved = torch.cuda.memory_reserved(0) / 1024**3
- print(f"GPU Memory: Allocated {current_allocated:.2f}GB | Reserved {current_reserved:.2f}GB")
-
- print(f"Model loaded successfully (Total time: {total_time:.2f}s)")
- return True
-
- except Exception as e:
- print(f"Failed to load model: {e}")
- import traceback
- traceback.print_exc()
- self.model = None
- self.tokenizer = None
- self.current_model_path = None
- return False
-
- def inference(
- self,
- user_prompt: str,
- prompt_template: str,
- 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 with optional FlashAttention-2 and INT8 quantization"""
-
- # Auto-find Qwen model
- model_path = self._find_qwen_model()
- if model_path is None:
- error_msg = "Error: Qwen model file not found.\nPlease place the correct model file (e.g., qwen_3_4b.safetensors) in ComfyUI's models/text_encoders or models/clip folder."
- print(error_msg)
- return (error_msg,)
-
- if not os.path.exists(model_path):
- error_msg = f"Error: Model file does not exist: {model_path}\nPlease place the correct model file in the text_encoders or clip folder."
- print(error_msg)
- return (error_msg,)
-
- # Use fixed repo_id
- repo_id = "Qwen/Qwen3-4B"
- 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,)
-
- try:
- # Load and apply template
- template_content = ""
- if prompt_template != "Custom":
- template_content = self._load_template_content(prompt_template)
-
- # If template content exists, replace system_prompt with template
- if template_content:
- system_prompt = template_content
-
- messages = []
- if system_prompt and system_prompt.strip():
- messages.append({"role": "system", "content": system_prompt})
- messages.append({"role": "user", "content": user_prompt})
-
- text = self.tokenizer.apply_chat_template(
- messages,
- tokenize=False,
- add_generation_prompt=True
- )
-
- inputs = self.tokenizer(text, return_tensors="pt")
-
- device = next(self.model.parameters()).device
- inputs = {k: v.to(device) for k, v in inputs.items()}
-
- inference_start = time.time()
- print(f"Inference starting...")
-
- with torch.no_grad():
- outputs = self.model.generate(
- input_ids=inputs['input_ids'],
- attention_mask=inputs['attention_mask'],
- max_new_tokens=max_new_tokens,
- do_sample=do_sample,
- temperature=temperature,
- top_p=top_p,
- top_k=top_k,
- pad_token_id=self.tokenizer.eos_token_id
- )
-
- response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-
- # Remove assistant markers
- if "assistant" in response:
- for separator in ["<|im_start|>assistant\n", "assistant\n", "Assistant:", "assistant:"]:
- if separator in response:
- response = response.split(separator)[-1].strip()
- break
-
- # Remove tags
- response = self._remove_thinking_tags(response)
-
- inference_time = time.time() - inference_start
-
- tokens_generated = len(outputs[0]) - len(inputs['input_ids'][0])
- tokens_per_sec = tokens_generated / inference_time if inference_time > 0 else 0
-
- print(f"Inference completed (Time: {inference_time:.2f}s | Tokens: {tokens_generated} | Speed: {tokens_per_sec:.1f} tokens/s)")
-
- if torch.cuda.is_available():
- print(f"GPU Memory: Used {torch.cuda.memory_allocated(0) / 1024**3:.2f}GB | Peak {torch.cuda.max_memory_allocated(0) / 1024**3:.2f}GB")
-
- return (response,)
-
- except Exception as e:
- import traceback
- error_msg = f"Inference failed: {str(e)}\n{traceback.format_exc()}"
- print(error_msg)
- return (error_msg,)
+import os
+import torch
+import folder_paths
+import requests
+import time
+import re
+import gc
+from typing import Optional, Tuple, Dict
+
+class QwenGPUInference:
+ """
+ Qwen3-4B GPU Inference Node with intelligent memory management
+ Auto-downloads required config files and performs GPU inference
+ Includes GPU memory checking and cleanup to avoid conflicts with ComfyUI's CLIP models
+ """
+
+ def __init__(self):
+ self.model = None
+ self.tokenizer = None
+ self.current_model_path = None
+ self.config_dir = None
+
+ @classmethod
+ def _get_safetensors_files(cls):
+ """Get all safetensors files from text_encoders and clip folders"""
+ safetensors_files = []
+
+ # Search in text_encoders folder
+ try:
+ text_encoder_paths = folder_paths.get_folder_paths("text_encoders")
+ for path in text_encoder_paths:
+ if os.path.exists(path):
+ for file in os.listdir(path):
+ if file.lower().endswith('.safetensors'):
+ full_path = os.path.join(path, file)
+ if full_path not in safetensors_files:
+ safetensors_files.append(full_path)
+ except:
+ pass
+
+ # Search in clip folder
+ try:
+ clip_paths = folder_paths.get_folder_paths("clip")
+ for path in clip_paths:
+ if os.path.exists(path):
+ for file in os.listdir(path):
+ if file.lower().endswith('.safetensors'):
+ full_path = os.path.join(path, file)
+ if full_path not in safetensors_files:
+ safetensors_files.append(full_path)
+ except:
+ pass
+
+ if not safetensors_files:
+ return ["No safetensors files found"]
+
+ return sorted(safetensors_files)
+
+ @classmethod
+ def _get_prompt_templates(cls):
+ """Get all .md template files from Prompt folder"""
+ current_dir = os.path.dirname(os.path.abspath(__file__))
+ prompt_dir = os.path.join(current_dir, "Prompt")
+
+ templates = []
+
+ if os.path.exists(prompt_dir):
+ for file in os.listdir(prompt_dir):
+ if file.lower().endswith('.md'):
+ templates.append(file)
+
+ if not templates:
+ return ["No Template"]
+
+ return sorted(templates)
+
+ @classmethod
+ def _load_template_content(cls, template_name):
+ """Load template content"""
+ if template_name == "No Template" or template_name == "Custom":
+ return ""
+
+ current_dir = os.path.dirname(os.path.abspath(__file__))
+ template_path = os.path.join(current_dir, "Prompt", template_name)
+
+ if os.path.exists(template_path):
+ try:
+ with open(template_path, 'r', encoding='utf-8') as f:
+ return f.read()
+ except:
+ return ""
+
+ return ""
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ templates = cls._get_prompt_templates()
+ # Add "Custom" option at the beginning of the list
+ template_options = ["Custom"] + templates
+
+ return {
+ "required": {
+ "user_prompt": ("STRING", {
+ "multiline": True,
+ "default": "A girl in a coffee shop"
+ }),
+ "prompt_template": (template_options, {
+ "default": template_options[0] if template_options else "Custom"
+ }),
+ "system_prompt": ("STRING", {
+ "multiline": True,
+ "default": ""
+ }),
+ "max_new_tokens": ("INT", {
+ "default": 2048,
+ "min": 1,
+ "max": 4096,
+ "step": 1
+ }),
+ "temperature": ("FLOAT", {
+ "default": 0.7,
+ "min": 0.0,
+ "max": 2.0,
+ "step": 0.1
+ }),
+ },
+ "optional": {
+ "keep_model_loaded": ("BOOLEAN", {
+ "default": True,
+ "tooltip": "WARNING: Set to False to unload model after generation. Required for low VRAM workflows."
+ }),
+ "do_sample": ("BOOLEAN", {
+ "default": True,
+ "tooltip": "Enable sampling"
+ }),
+ "top_p": ("FLOAT", {
+ "default": 0.9,
+ "min": 0.0,
+ "max": 1.0,
+ "step": 0.05
+ }),
+ "top_k": ("INT", {
+ "default": 50,
+ "min": 0,
+ "max": 100,
+ "step": 1
+ }),
+ }
+ }
+
+ RETURN_TYPES = ("STRING",)
+ RETURN_NAMES = ("text",)
+ FUNCTION = "inference"
+ CATEGORY = "ListHelper"
+
+ def _find_qwen_model(self) -> Optional[str]:
+ """Auto-find qwen_3_4b.safetensors model"""
+ safetensors_files = self._get_safetensors_files()
+
+ # Prioritize finding qwen_3_4b.safetensors
+ for path in safetensors_files:
+ if path != "No safetensors files found":
+ basename = os.path.basename(path).lower()
+ if "qwen" in basename and "3" in basename and "4b" in basename:
+ return path
+
+ # If specific model not found, return first safetensors file
+ if safetensors_files and safetensors_files[0] != "No safetensors files found":
+ return safetensors_files[0]
+
+ return None
+
+ def _remove_thinking_tags(self, text: str) -> str:
+ """Remove ... tags and their content"""
+ # Use regex to remove all ... blocks
+ cleaned_text = re.sub(r'.*?', '', text, flags=re.DOTALL)
+ # Remove extra blank lines
+ cleaned_text = re.sub(r'\n\s*\n', '\n', cleaned_text)
+ return cleaned_text.strip()
+
+ def _check_gpu_memory(self, required_gb: float = 8.0) -> Tuple[bool, str]:
+ """
+ Check if GPU memory is sufficient
+
+ Args:
+ required_gb: Required GPU memory size (GB)
+
+ Returns:
+ (is_sufficient, detailed_message)
+ """
+ if not torch.cuda.is_available():
+ return True, "GPU Memory: Using CPU mode"
+
+ try:
+ total_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
+ allocated_memory = torch.cuda.memory_allocated(0) / 1024**3
+ reserved_memory = torch.cuda.memory_reserved(0) / 1024**3
+ free_memory = total_memory - reserved_memory
+
+ info = f"GPU Memory: Total {total_memory:.2f}GB | Free {free_memory:.2f}GB | Required {required_gb:.2f}GB"
+
+ if free_memory < required_gb:
+ return False, info + f" | Insufficient: need {required_gb - free_memory:.2f}GB more"
+
+ return True, info + " | Sufficient"
+
+ except Exception as e:
+ return True, f"GPU Memory: Cannot check - {e}"
+
+ def _free_gpu_memory(self) -> None:
+ """
+ Free GPU memory
+ Clear PyTorch cache and run garbage collection
+ """
+ try:
+ if torch.cuda.is_available():
+ # Record memory before cleanup
+ before_reserved = torch.cuda.memory_reserved(0) / 1024**3
+
+ # Clear CUDA cache
+ torch.cuda.empty_cache()
+ torch.cuda.synchronize()
+
+ # Force garbage collection
+ gc.collect()
+
+ # Clear again
+ torch.cuda.empty_cache()
+
+ # Record memory after cleanup
+ after_reserved = torch.cuda.memory_reserved(0) / 1024**3
+
+ freed_reserved = before_reserved - after_reserved
+
+ print(f"GPU Memory: Freed {freed_reserved:.2f}GB | Current reserved {after_reserved:.2f}GB")
+ else:
+ gc.collect()
+ print("Memory cleanup: CPU mode")
+
+ except Exception as e:
+ print(f"Memory cleanup error: {e}")
+ # Attempt garbage collection even if error occurs
+ gc.collect()
+
+ def _download_config_files(self, model_path: str, repo_id: str) -> bool:
+ """Auto-download HuggingFace config files"""
+ try:
+ model_dir = os.path.dirname(model_path)
+ model_basename = os.path.splitext(os.path.basename(model_path))[0]
+ self.config_dir = os.path.join(model_dir, f"{model_basename}_config")
+
+ config_files = [
+ "config.json",
+ "generation_config.json",
+ "merges.txt",
+ "tokenizer.json",
+ "tokenizer_config.json",
+ "vocab.json"
+ ]
+
+ all_exist = all(os.path.exists(os.path.join(self.config_dir, f)) for f in config_files)
+
+ if all_exist:
+ print(f"Config files exist: {self.config_dir}")
+ return True
+
+ os.makedirs(self.config_dir, exist_ok=True)
+ print(f"Downloading config files to: {self.config_dir}")
+
+ base_url = f"https://huggingface.co/{repo_id}/resolve/main/"
+
+ for filename in config_files:
+ filepath = os.path.join(self.config_dir, filename)
+
+ if os.path.exists(filepath):
+ print(f" {filename} exists")
+ continue
+
+ url = base_url + filename
+ print(f" Downloading {filename}...")
+
+ try:
+ response = requests.get(url, timeout=30)
+ response.raise_for_status()
+
+ with open(filepath, 'wb') as f:
+ f.write(response.content)
+ print(f" {filename} downloaded")
+
+ except Exception as e:
+ print(f" Failed to download {filename}: {str(e)}")
+ continue
+
+ return True
+
+ except Exception as e:
+ print(f"Failed to download config files: {e}")
+ import traceback
+ 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)"""
+ try:
+ if self.model is not None and self.current_model_path == model_path:
+ print(f"Model already loaded: {os.path.basename(model_path)}")
+ return True
+
+ try:
+ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
+ from safetensors.torch import load_file
+ except ImportError as e:
+ print(f"Missing required packages: {e}")
+ print("Please run: pip install transformers safetensors")
+ return False
+
+ if not self._download_config_files(model_path, repo_id):
+ return False
+
+ 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)
+ print(memory_info)
+
+ # Step 2: If insufficient memory, try to free up
+ if not is_enough:
+ print("Memory insufficient, cleaning up...")
+ self._free_gpu_memory()
+
+ # Check again
+ is_enough, memory_info = self._check_gpu_memory(required_gb=7.5)
+ print(memory_info)
+
+ if not is_enough:
+ print("Still insufficient, will use CPU Offload strategy (slower inference)")
+
+ overall_start = time.time()
+
+ device = "cuda:0" if torch.cuda.is_available() else "cpu"
+ dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+ # 1. Load tokenizer
+ self.tokenizer = AutoTokenizer.from_pretrained(
+ self.config_dir,
+ trust_remote_code=True
+ )
+
+ # 2. Load config
+ config = AutoConfig.from_pretrained(self.config_dir, trust_remote_code=True)
+
+ # 3. Choose loading strategy based on memory
+ import shutil
+ temp_model_dir = os.path.join(self.config_dir, "temp_model")
+ os.makedirs(temp_model_dir, exist_ok=True)
+
+ # 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
+
+ if free_memory >= 7.5:
+ # Sufficient memory: Full GPU loading
+ print(f"Strategy: Full GPU loading (Free: {free_memory:.2f}GB)")
+ max_memory_config = None
+ offload_folder = None
+ else:
+ # Insufficient memory: Use CPU Offload
+ available_gpu = max(3.0, free_memory - 1.0) # Reserve at least 1GB
+ print(f"Strategy: CPU Offload (Free: {free_memory:.2f}GB, Allocate: {available_gpu:.2f}GB)")
+ max_memory_config = {
+ 0: f"{available_gpu:.1f}GB",
+ "cpu": "16GB"
+ }
+ # Create offload folder
+ offload_folder = os.path.join(self.config_dir, "offload")
+ os.makedirs(offload_folder, exist_ok=True)
+ else:
+ print("Strategy: CPU mode")
+ max_memory_config = None
+ offload_folder = None
+
+ load_start = time.time()
+
+ try:
+ # Copy config files
+ for file in ["config.json", "generation_config.json"]:
+ src = os.path.join(self.config_dir, file)
+ if os.path.exists(src):
+ shutil.copy(src, temp_model_dir)
+
+ # Create link or copy safetensors
+ safetensors_target = os.path.join(temp_model_dir, "model.safetensors")
+ if os.path.exists(safetensors_target):
+ os.remove(safetensors_target)
+
+ # Windows uses hard links instead of symbolic links
+ try:
+ os.link(model_path, safetensors_target)
+ except:
+ shutil.copy(model_path, safetensors_target)
+
+ # Use appropriate loading strategy
+ load_kwargs = {
+ "pretrained_model_name_or_path": temp_model_dir,
+ "trust_remote_code": True,
+ "device_map": "auto",
+ "torch_dtype": dtype,
+ "low_cpu_mem_usage": True
+ }
+
+ if max_memory_config is not None:
+ load_kwargs["max_memory"] = max_memory_config
+
+ if offload_folder is not None:
+ load_kwargs["offload_folder"] = offload_folder
+
+ self.model = AutoModelForCausalLM.from_pretrained(**load_kwargs)
+
+ print(f"Model loaded (Time: {time.time() - load_start:.2f}s)")
+
+ finally:
+ # Clean up temp directory
+ try:
+ if os.path.exists(temp_model_dir):
+ shutil.rmtree(temp_model_dir)
+ except:
+ pass
+
+ self.current_model_path = model_path
+ total_time = time.time() - overall_start
+
+ if torch.cuda.is_available():
+ current_allocated = torch.cuda.memory_allocated(0) / 1024**3
+ current_reserved = torch.cuda.memory_reserved(0) / 1024**3
+ print(f"GPU Memory: Allocated {current_allocated:.2f}GB | Reserved {current_reserved:.2f}GB")
+
+ print(f"Model loaded successfully (Total time: {total_time:.2f}s)")
+ return True
+
+ except Exception as e:
+ print(f"Failed to load model: {e}")
+ import traceback
+ traceback.print_exc()
+ self.model = None
+ self.tokenizer = None
+ self.current_model_path = None
+ return False
+
+ def inference(
+ self,
+ user_prompt: str,
+ prompt_template: str,
+ system_prompt: str,
+ max_new_tokens: int,
+ temperature: float,
+ keep_model_loaded: bool = True,
+ do_sample: bool = True,
+ top_p: float = 0.9,
+ top_k: int = 50,
+ ) -> Tuple[str]:
+ """Execute inference"""
+
+ # Auto-find Qwen model
+ model_path = self._find_qwen_model()
+ if model_path is None:
+ error_msg = "Error: Qwen model file not found.\nPlease place the correct model file (e.g., qwen_3_4b.safetensors) in ComfyUI's models/text_encoders or models/clip folder."
+ print(error_msg)
+ return (error_msg,)
+
+ if not os.path.exists(model_path):
+ error_msg = f"Error: Model file does not exist: {model_path}\nPlease place the correct model file in the text_encoders or clip folder."
+ print(error_msg)
+ return (error_msg,)
+
+ # Use fixed repo_id
+ repo_id = "Qwen/Qwen3-4B"
+ if not self._load_model(model_path, repo_id):
+ 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,)
+
+ try:
+ # Load and apply template
+ template_content = ""
+ if prompt_template != "Custom":
+ template_content = self._load_template_content(prompt_template)
+
+ # If template content exists, replace system_prompt with template
+ if template_content:
+ system_prompt = template_content
+
+ messages = []
+ if system_prompt and system_prompt.strip():
+ messages.append({"role": "system", "content": system_prompt})
+ messages.append({"role": "user", "content": user_prompt})
+
+ text = self.tokenizer.apply_chat_template(
+ messages,
+ tokenize=False,
+ add_generation_prompt=True
+ )
+
+ inputs = self.tokenizer(text, return_tensors="pt")
+
+ device = next(self.model.parameters()).device
+ inputs = {k: v.to(device) for k, v in inputs.items()}
+
+ inference_start = time.time()
+ print(f"Inference starting...")
+
+ with torch.no_grad():
+ outputs = self.model.generate(
+ input_ids=inputs['input_ids'],
+ attention_mask=inputs['attention_mask'],
+ max_new_tokens=max_new_tokens,
+ do_sample=do_sample,
+ temperature=temperature,
+ top_p=top_p,
+ top_k=top_k,
+ pad_token_id=self.tokenizer.eos_token_id
+ )
+
+ response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+ # Remove assistant markers
+ if "assistant" in response:
+ for separator in ["<|im_start|>assistant\n", "assistant\n", "Assistant:", "assistant:"]:
+ if separator in response:
+ response = response.split(separator)[-1].strip()
+ break
+
+ # Remove tags
+ response = self._remove_thinking_tags(response)
+
+ inference_time = time.time() - inference_start
+
+ tokens_generated = len(outputs[0]) - len(inputs['input_ids'][0])
+ tokens_per_sec = tokens_generated / inference_time if inference_time > 0 else 0
+
+ print(f"Inference completed (Time: {inference_time:.2f}s | Tokens: {tokens_generated} | Speed: {tokens_per_sec:.1f} tokens/s)")
+
+ if torch.cuda.is_available():
+ print(f"GPU Memory: Used {torch.cuda.memory_allocated(0) / 1024**3:.2f}GB | Peak {torch.cuda.max_memory_allocated(0) / 1024**3:.2f}GB")
+
+ # Check if we should unload the model to free VRAM
+ if not keep_model_loaded:
+ print("Explicitly unloading model to free VRAM...")
+ del self.model
+ del self.tokenizer
+ self.model = None
+ self.tokenizer = None
+ self.current_model_path = None
+ self._free_gpu_memory()
+
+ return (response,)
+
+ except Exception as e:
+ import traceback
+ error_msg = f"Inference failed: {str(e)}\n{traceback.format_exc()}"
+ print(error_msg)
+ return (error_msg,)