randomseed
This commit is contained in:
+31
-21
@@ -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)
|
||||
|
||||
+17
-15
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user