diff --git a/gguf_inference.py b/gguf_inference.py index d668424..aeeea3e 100644 --- a/gguf_inference.py +++ b/gguf_inference.py @@ -7,6 +7,7 @@ import subprocess import platform import base64 import io +import random import numpy as np import folder_paths import urllib.request @@ -244,6 +245,7 @@ class GGUFInference: }), }, "optional": { + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}), "keep_model_loaded": ("BOOLEAN", { "default": False, "tooltip": "Keep model in memory after inference" @@ -266,8 +268,8 @@ class GGUFInference: } } - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("text",) + RETURN_TYPES = ("STRING", "INT",) + RETURN_NAMES = ("text", "used_seed",) FUNCTION = "inference" CATEGORY = "ListHelper" @@ -471,12 +473,13 @@ class GGUFInference: temperature: float, top_p: float, top_k: int, + seed: int = 0, keep_model_loaded: bool = False, enable_vision: bool = False, mmproj_file: str = "No mmproj files", image = None, auto_install_llama_cpp: bool = False, - ) -> Tuple[str]: + ) -> Tuple[str, int]: """Execute GGUF model inference""" # Check if llama-cpp-python is available @@ -490,7 +493,7 @@ class GGUFInference: else: error_msg = "ERROR: llama-cpp-python is not installed.\n\nPlease either:\n1. Enable 'auto_install_llama_cpp' option (Windows only)\n2. Install manually: pip install llama-cpp-python" print(error_msg) - return (error_msg,) + return (error_msg, seed) # Check if model is a suggested download if model in SUGGESTED_MODELS: @@ -505,11 +508,11 @@ class GGUFInference: else: error_msg = "Error: Cannot find clip folder for downloading model." print(error_msg) - return (error_msg,) + return (error_msg, seed) except: error_msg = "Error: Cannot access clip folder for downloading model." print(error_msg) - return (error_msg,) + return (error_msg, seed) model_path = os.path.join(download_dir, filename) @@ -521,7 +524,7 @@ class GGUFInference: if not self._download_file(download_url, model_path): error_msg = f"Error: Failed to download model from {download_url}" print(error_msg) - return (error_msg,) + return (error_msg, seed) else: # Get full path from model name gguf_files = self._get_gguf_files() @@ -534,12 +537,12 @@ class GGUFInference: if model_path is None: error_msg = f"Error: Model not found: {model}\nPlease place GGUF files in text_encoders or clip folder." print(error_msg) - return (error_msg,) + return (error_msg, seed) if not os.path.exists(model_path): error_msg = f"Error: Model file does not exist: {model_path}" print(error_msg) - return (error_msg,) + return (error_msg, seed) # Check if this is a vision model is_vision_model = self._is_vision_model(model_path) @@ -621,7 +624,7 @@ class GGUFInference: if not self._load_model(model_path, enable_vision, mmproj_path): error_msg = "Error: Model loading failed.\nPlease check if llama-cpp-python is properly installed." print(error_msg) - return (error_msg,) + return (error_msg, seed) try: # Load and apply template @@ -645,7 +648,7 @@ class GGUFInference: if image_url is None: error_msg = "Error: Failed to convert image to base64 format." print(error_msg) - return (error_msg,) + return (error_msg, seed) messages.append({ "role": "user", @@ -661,16 +664,23 @@ class GGUFInference: print("Note: Image input provided but vision mode is disabled, ignoring image") # Run inference - print(f"Starting inference...") + print(f"Starting inference (Seed: {seed})...") inference_start = time.time() - response = self.model.create_chat_completion( - messages=messages, - max_tokens=max_tokens, - temperature=temperature, - top_p=top_p, - top_k=top_k, - ) + # Prepare generation parameters + gen_params = { + "messages": messages, + "max_tokens": max_tokens, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + } + + # Only add seed if it's set (> 0) + if seed > 0: + gen_params["seed"] = seed + + response = self.model.create_chat_completion(**gen_params) # Extract response text response_text = response["choices"][0]["message"]["content"] @@ -689,10 +699,10 @@ class GGUFInference: print("Unloading model to free memory...") self._free_memory() - return (response_text,) + return (response_text, seed) except Exception as e: import traceback error_msg = f"Inference failed: {str(e)}\n{traceback.format_exc()}" print(error_msg) - return (error_msg,) + return (error_msg, seed) diff --git a/qwen_inference.py b/qwen_inference.py index 6ef1781..8ef0ff4 100644 --- a/qwen_inference.py +++ b/qwen_inference.py @@ -5,6 +5,7 @@ import requests import time import re import gc +import random from typing import Optional, Tuple, Dict class QwenGPUInference: @@ -126,6 +127,7 @@ class QwenGPUInference: }), }, "optional": { + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}), "keep_model_loaded": ("BOOLEAN", { "default": False, "tooltip": "WARNING: Set to False to unload model after generation. Required for low VRAM workflows." @@ -138,10 +140,6 @@ class QwenGPUInference: "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, @@ -157,8 +155,8 @@ class QwenGPUInference: } } - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("text",) + RETURN_TYPES = ("STRING", "INT",) + RETURN_NAMES = ("text", "used_seed",) FUNCTION = "inference" CATEGORY = "ListHelper" @@ -520,13 +518,13 @@ class QwenGPUInference: system_prompt: str, max_new_tokens: int, temperature: float, + seed: int = 0, keep_model_loaded: bool = True, use_flash_attention: bool = False, use_quantization: bool = False, - do_sample: bool = True, top_p: float = 0.9, top_k: int = 50, - ) -> Tuple[str]: + ) -> Tuple[str, int]: """Execute inference with optional FlashAttention-2 and INT8 quantization""" # Auto-find Qwen model @@ -534,21 +532,25 @@ class QwenGPUInference: 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,) + return (error_msg, seed) 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,) + return (error_msg, seed) # 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,) + return (error_msg, seed) try: + # Handle seed - ComfyUI's control_after_generate will handle randomization + if seed > 0: + torch.manual_seed(seed) + # Load and apply template template_content = "" if prompt_template != "Custom": @@ -575,14 +577,14 @@ class QwenGPUInference: inputs = {k: v.to(device) for k, v in inputs.items()} inference_start = time.time() - print(f"Inference starting...") + print(f"Inference starting (Seed: {seed})...") 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, + do_sample=True, temperature=temperature, top_p=top_p, top_k=top_k, @@ -621,10 +623,10 @@ class QwenGPUInference: self.current_model_path = None self._free_gpu_memory() - return (response,) + return (response, seed) except Exception as e: import traceback error_msg = f"Inference failed: {str(e)}\n{traceback.format_exc()}" print(error_msg) - return (error_msg,) + return (error_msg, seed) diff --git a/test_qwen_node.py b/test_qwen_node.py deleted file mode 100644 index ca51cf1..0000000 --- a/test_qwen_node.py +++ /dev/null @@ -1,80 +0,0 @@ -""" -測試 QwenGPUInference 節點 -""" - -import sys -import os - -# 添加當前目錄和 ComfyUI 路徑 -current_dir = os.path.dirname(os.path.abspath(__file__)) -comfy_path = os.path.abspath(os.path.join(current_dir, "..", "..")) -sys.path.insert(0, current_dir) -sys.path.insert(0, comfy_path) - -# 模擬 folder_paths -class MockFolderPaths: - @staticmethod - def get_folder_paths(folder_type): - if folder_type == "text_encoders": - return [os.path.join(comfy_path, "models", "text_encoders")] - return [] - -# 替換 folder_paths -import folder_paths as real_folder_paths -try: - # 嘗試使用真實的 folder_paths - pass -except: - # 如果失敗，使用模擬的 - sys.modules['folder_paths'] = MockFolderPaths() - -# 現在可以導入節點 -from qwen_inference import QwenGPUInference - -def test_node(): - """測試節點基本功能""" - - print("=" * 60) - print("測試 QwenGPUInference 節點") - print("=" * 60) - - # 創建節點實例 - node = QwenGPUInference() - - # 測試獲取檔案列表 - print("\n1. 測試獲取 safetensors 檔案列表...") - files = node._get_safetensors_files() - print(f" 找到 {len(files)} 個檔案:") - for f in files: - print(f" - {os.path.basename(f) if f != 'No safetensors files found' else f}") - - if files[0] == "No safetensors files found": - print("\n⚠️ 未找到 safetensors 檔案，無法繼續測試") - print("請確保在 models/text_encoders 目錄中有 .safetensors 檔案") - return - - # 測試推理 - print("\n2. 測試推理功能...") - model_file = os.path.basename(files[0]) - - result = node.inference( - model_file=model_file, - user_prompt="你好，請介紹一下你自己。", - system_prompt="你是一個有用的 AI 助手。", - max_new_tokens=100, - temperature=0.7, - do_sample=True, - top_p=0.9, - top_k=50, - repo_id="Qwen/Qwen3-4B" - ) - - print("\n推理結果:") - print("-" * 60) - print(result[0]) - print("-" * 60) - - print("\n✓ 測試完成!") - -if __name__ == "__main__": - test_node()