randomseed

This commit is contained in:
dseditor
2025-12-09 23:23:02 +08:00
parent c210234b5b
commit 39cae890c1
3 changed files with 48 additions and 116 deletions
+31 -21
View File
@@ -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
View File
@@ -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)
-80
View File
@@ -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()