init
This commit is contained in:
pk
+24
@@ -0,0 +1,24 @@
|
||||
from .nodes.pulid.AQ_MasksAndImagesAsList import AQ_MasksAndImagesAsList
|
||||
from .nodes.transformers.AQ_Qwen import AQ_Qwen
|
||||
from .nodes.transformers.AQ_Qwen import AQ_QwenLoader
|
||||
from .nodes.pulid.AQ_multiface_ApplyPulidFlux import AQ_multiface_ApplyPulidFlux
|
||||
from .nodes.api.AQ_Gemini import AQ_Gemini
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AQ_MasksAndImagesAsList": AQ_MasksAndImagesAsList,
|
||||
"AQ_multiface_ApplyPulidFlux": AQ_multiface_ApplyPulidFlux,
|
||||
"AQ_Qwen": AQ_Qwen,
|
||||
"AQ_QwenLoader": AQ_QwenLoader,
|
||||
"AQ_Gemini": AQ_Gemini
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AQ_MasksAndImagesAsList": "AQ_MasksAndImagesAsList",
|
||||
"AQ_multiface_ApplyPulidFlux": "AQ_multiface_ApplyPulidFlux",
|
||||
"AQ_Qwen": "AQ_Qwen",
|
||||
"AQ_QwenLoader": "AQ_QwenLoader",
|
||||
"AQ_Gemini": "AQ_Gemini"
|
||||
}
|
||||
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -0,0 +1,276 @@
|
||||
import requests
|
||||
import base64
|
||||
from PIL import Image
|
||||
import io
|
||||
import json
|
||||
import numpy as np
|
||||
import torch
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
|
||||
class AQ_Gemini:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
default_schema = '''{
|
||||
"type": "object",
|
||||
"required": ["description", "style"],
|
||||
"properties": {
|
||||
"description": {
|
||||
"type": "string"
|
||||
},
|
||||
"style": {
|
||||
"type": "string"
|
||||
},
|
||||
"tags": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
}
|
||||
}'''
|
||||
|
||||
models = [
|
||||
"gemma-3-12b-it", # default
|
||||
"gemma-3-1b-it",
|
||||
"gemma-3-4b-it",
|
||||
"gemma-3-27b-it",
|
||||
"gemini-2.5-flash-preview-04-17",
|
||||
"gemini-2.5-pro-preview-05-06",
|
||||
"gemini-2.0-flash",
|
||||
"gemini-2.0-flash-preview-image-generation",
|
||||
"gemini-2.0-flash-lite",
|
||||
"gemini-1.5-flash",
|
||||
"gemini-1.5-flash-8b",
|
||||
"gemini-1.5-pro",
|
||||
"gemini-embedding-exp",
|
||||
"imagen-3.0-generate-002",
|
||||
"veo-2.0-generate-001",
|
||||
"gemini-2.0-flash-live-001",
|
||||
"custom" # Option for custom model
|
||||
]
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"gemini_api_key": ("STRING", {"default": ""}),
|
||||
"model_selection": (models, {"default": "gemma-3-12b-it"}),
|
||||
"custom_model": ("STRING", {"default": "", "multiline": False}),
|
||||
"prompt": ("STRING", {"default": "", "multiline": True}),
|
||||
"system_message": ("STRING", {"default": "You are a helpful assistant.", "multiline": True}),
|
||||
"temperature": ("FLOAT", {"default": 1, "min": 0.0, "max": 2.0, "step": 0.05}),
|
||||
"top_k": ("INT", {"default": 64, "min": 0, "max": 100}),
|
||||
"top_p": ("FLOAT", {"default": 0.95, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"enable_json": ("BOOLEAN", {"default": False}),
|
||||
"json_schema": ("STRING", {"default": default_schema, "multiline": True}),
|
||||
"result_template": ("STRING", {"default": "", "multiline": True}),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING", "STRING",)
|
||||
RETURN_NAMES = ("response", "formatted_response", "instruction")
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "Aquasite/LLM"
|
||||
|
||||
def generate(self, gemini_api_key, model_selection, custom_model, prompt, system_message, temperature=0.8,
|
||||
top_k=64, top_p=0.95, enable_json=False, json_schema="", result_template="", image=None):
|
||||
|
||||
instruction = """
|
||||
How to use AQ_Gemini node:
|
||||
1. Provide your Gemini API key in the 'gemini_api_key' field
|
||||
2. Select a model from the dropdown or choose 'custom' and enter your model name
|
||||
3. Enter your prompt in the 'prompt' field
|
||||
4. Customize the system message to control the assistant's behavior
|
||||
5. Optionally provide an image for image-based prompts
|
||||
6. For JSON output:
|
||||
- Enable 'enable_json'
|
||||
- Provide JSON schema in 'json_schema'
|
||||
- Use 'result_template' to format the response
|
||||
|
||||
Available Models:
|
||||
- Gemma Models: gemma-3-{1b,4b,12b,27b}-it
|
||||
- Gemini 2.5: flash-preview, pro-preview
|
||||
- Gemini 2.0: flash, flash-preview-image-generation, flash-lite, flash-live
|
||||
- Gemini 1.5: flash, flash-8b, pro
|
||||
- Special Models: gemini-embedding-exp, imagen-3.0, veo-2.0
|
||||
|
||||
Note: If any error occurs, empty strings will be returned instead of raising an error.
|
||||
"""
|
||||
|
||||
try:
|
||||
if not gemini_api_key:
|
||||
return ("", "", instruction)
|
||||
|
||||
# Initialize Gemini client
|
||||
client = genai.Client(api_key=gemini_api_key)
|
||||
|
||||
# Determine which model to use
|
||||
model_name = custom_model if model_selection == "custom" and custom_model else model_selection
|
||||
|
||||
# Prepare contents list
|
||||
contents = []
|
||||
|
||||
# Add system message if provided
|
||||
if system_message:
|
||||
contents.append(
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[types.Part.from_text(text=system_message)]
|
||||
)
|
||||
)
|
||||
|
||||
# Handle image if provided
|
||||
if image is not None:
|
||||
# Convert tensor to numpy array
|
||||
if isinstance(image, torch.Tensor):
|
||||
image = image.cpu().numpy()
|
||||
|
||||
# Handle ComfyUI image format (B, H, W, C)
|
||||
if len(image.shape) == 4:
|
||||
image = image[0]
|
||||
|
||||
# Ensure proper type and range for PIL
|
||||
if image.dtype != np.uint8:
|
||||
if image.max() <= 1:
|
||||
image = (image * 255).astype(np.uint8)
|
||||
else:
|
||||
image = image.astype(np.uint8)
|
||||
|
||||
# Convert numpy array to PIL Image and then to base64
|
||||
img = Image.fromarray(image)
|
||||
img_byte_arr = io.BytesIO()
|
||||
img.save(img_byte_arr, format="PNG")
|
||||
img_byte_arr.seek(0)
|
||||
img_base64 = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
|
||||
|
||||
# Create content with image and prompt
|
||||
contents.append(
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[{
|
||||
"inline_data": {
|
||||
"mime_type": "image/png",
|
||||
"data": img_base64
|
||||
}
|
||||
},
|
||||
{
|
||||
"text": prompt
|
||||
}]
|
||||
)
|
||||
)
|
||||
else:
|
||||
# Text-only content
|
||||
contents.append(
|
||||
types.Content(
|
||||
role="user",
|
||||
parts=[types.Part.from_text(text=prompt)]
|
||||
)
|
||||
)
|
||||
|
||||
# Configure generation settings
|
||||
generate_config = types.GenerateContentConfig(
|
||||
temperature=temperature,
|
||||
top_k=top_k,
|
||||
top_p=top_p
|
||||
)
|
||||
|
||||
# Add JSON schema if enabled and not a Gemma-3 model
|
||||
is_gemma_model = model_name.startswith("gemma-3")
|
||||
if enable_json and json_schema and not is_gemma_model:
|
||||
try:
|
||||
schema_dict = json.loads(json_schema)
|
||||
generate_config.response_mime_type = "application/json"
|
||||
generate_config.response_schema = self._convert_schema_dict_to_genai(schema_dict)
|
||||
except json.JSONDecodeError:
|
||||
print("Invalid JSON schema provided")
|
||||
|
||||
# Generate content
|
||||
response = ""
|
||||
for chunk in client.models.generate_content_stream(
|
||||
model=model_name,
|
||||
contents=contents,
|
||||
config=generate_config,
|
||||
):
|
||||
if chunk.text:
|
||||
response += chunk.text
|
||||
|
||||
formatted_response = response
|
||||
|
||||
# Handle template formatting if result_template is provided and response might be JSON
|
||||
if result_template and enable_json:
|
||||
success, json_response = self.try_parse_json(response)
|
||||
if success:
|
||||
try:
|
||||
formatted_response = result_template.format(json=json_response)
|
||||
except KeyError as e:
|
||||
print(f"Template formatting error: {str(e)}")
|
||||
formatted_response = response
|
||||
else:
|
||||
print("Response is not valid JSON")
|
||||
formatted_response = response
|
||||
|
||||
return (response, formatted_response, instruction)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in Gemini API: {str(e)}")
|
||||
return ("", "", instruction)
|
||||
|
||||
def _convert_schema_dict_to_genai(self, schema_dict):
|
||||
"""Convert a JSON schema dictionary to genai Schema object."""
|
||||
schema_type = schema_dict.get("type", "object")
|
||||
|
||||
if schema_type == "object":
|
||||
properties = {}
|
||||
for prop_name, prop_schema in schema_dict.get("properties", {}).items():
|
||||
properties[prop_name] = self._convert_schema_dict_to_genai(prop_schema)
|
||||
|
||||
return genai.types.Schema(
|
||||
type=genai.types.Type.OBJECT,
|
||||
required=schema_dict.get("required", []),
|
||||
properties=properties
|
||||
)
|
||||
|
||||
elif schema_type == "array":
|
||||
items_schema = self._convert_schema_dict_to_genai(schema_dict.get("items", {}))
|
||||
return genai.types.Schema(
|
||||
type=genai.types.Type.ARRAY,
|
||||
items=items_schema
|
||||
)
|
||||
|
||||
elif schema_type == "string":
|
||||
return genai.types.Schema(type=genai.types.Type.STRING)
|
||||
|
||||
elif schema_type == "number":
|
||||
return genai.types.Schema(type=genai.types.Type.NUMBER)
|
||||
|
||||
elif schema_type == "integer":
|
||||
return genai.types.Schema(type=genai.types.Type.INTEGER)
|
||||
|
||||
elif schema_type == "boolean":
|
||||
return genai.types.Schema(type=genai.types.Type.BOOLEAN)
|
||||
|
||||
else:
|
||||
return genai.types.Schema(type=genai.types.Type.STRING)
|
||||
|
||||
def prepare_json(self, text):
|
||||
"""Clean up text that might contain markdown JSON code blocks."""
|
||||
# Remove markdown JSON code block markers
|
||||
text = text.replace("```json", "").replace("```", "").strip()
|
||||
return text
|
||||
|
||||
def try_parse_json(self, text):
|
||||
"""Try to parse JSON from text, handling markdown code blocks.
|
||||
Returns tuple (success, result) where result is either parsed JSON or original text."""
|
||||
try:
|
||||
# First clean up any markdown formatting
|
||||
cleaned_text = self.prepare_json(text)
|
||||
# Try to parse as JSON
|
||||
return True, json.loads(cleaned_text)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"JSON parsing error: {str(e)}")
|
||||
return False, text
|
||||
except Exception as e:
|
||||
print(f"Unexpected error while parsing JSON: {str(e)}")
|
||||
return False, text
|
||||
@@ -0,0 +1,47 @@
|
||||
import torch
|
||||
|
||||
class AQ_MasksAndImagesAsList:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"masks": ("MASK", {"forceInput": True}), # From SEGSToMaskList
|
||||
"images": ("IMAGE", {"forceInput": True}), # From SEGSToImageList
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE_LIST", "IMAGE_LIST", "STRING")
|
||||
RETURN_NAMES = ("images", "masks", "count")
|
||||
FUNCTION = "convert_to_lists"
|
||||
CATEGORY = "AQ/Pulid"
|
||||
INPUT_IS_LIST = True
|
||||
|
||||
def convert_to_lists(self, masks, images):
|
||||
# Ensure we have the same number of masks and images
|
||||
if len(masks) != len(images):
|
||||
print(f"Warning: Number of masks ({len(masks)}) does not match number of images ({len(images)})")
|
||||
# Use the shorter length to avoid index errors
|
||||
length = min(len(masks), len(images))
|
||||
masks = masks[:length]
|
||||
images = images[:length]
|
||||
|
||||
# Process masks and images together to maintain order
|
||||
processed_images = []
|
||||
processed_masks = []
|
||||
|
||||
for mask, img in zip(masks, images):
|
||||
# Process mask
|
||||
if len(mask.shape) == 2:
|
||||
mask = mask.unsqueeze(-1)
|
||||
processed_masks.append(mask)
|
||||
|
||||
# Process image
|
||||
processed_images.append(img)
|
||||
|
||||
# Create count string
|
||||
count_str = f"{len(processed_images)}"
|
||||
|
||||
return (processed_images, processed_masks, count_str)
|
||||
@@ -0,0 +1,95 @@
|
||||
import torch
|
||||
import sys
|
||||
import os
|
||||
import importlib
|
||||
|
||||
class AQ_multiface_ApplyPulidFlux:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL", ),
|
||||
"pulid_flux": ("PULIDFLUX", ),
|
||||
"eva_clip": ("EVA_CLIP", ),
|
||||
"face_analysis": ("FACEANALYSIS", ),
|
||||
"images": ("IMAGE_LIST", ),
|
||||
"weight": ("FLOAT", {"default": 1.0, "min": -1.0, "max": 5.0, "step": 0.05 }),
|
||||
"start_at": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001 }),
|
||||
"end_at": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001 }),
|
||||
},
|
||||
"optional": {
|
||||
"attn_masks": ("IMAGE_LIST", ),
|
||||
"options": ("OPTIONS",),
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID"
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "apply_multi_pulid_flux"
|
||||
CATEGORY = "AQ/Pulid"
|
||||
|
||||
def apply_multi_pulid_flux(self, model, pulid_flux, eva_clip, face_analysis, images, weight, start_at, end_at, attn_masks=None, options={}, unique_id=None):
|
||||
# Rather than importing the class, we'll call the original node directly
|
||||
# We know the original node is successfully registered in ComfyUI
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
|
||||
if "ApplyPulidFlux" not in NODE_CLASS_MAPPINGS:
|
||||
# This fallback will help with debugging if the node isn't found
|
||||
print("Available nodes:", list(NODE_CLASS_MAPPINGS.keys()))
|
||||
|
||||
# Try an alternative approach - import from the PuLID_Flux module directly
|
||||
try:
|
||||
# Add the parent directory to sys.path to make absolute imports work
|
||||
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
if parent_dir not in sys.path:
|
||||
sys.path.insert(0, parent_dir)
|
||||
|
||||
# Import the module (not relatively)
|
||||
print(f"Trying to import from {parent_dir}")
|
||||
pulidflux_module = importlib.import_module("ComfyUI_PuLID_Flux_ll.pulidflux")
|
||||
ApplyPulidFlux = pulidflux_module.ApplyPulidFlux
|
||||
apply_pulid = ApplyPulidFlux()
|
||||
except Exception as e:
|
||||
print(f"Error importing directly from module: {e}")
|
||||
raise ValueError(f"Could not find or import ApplyPulidFlux. Is ComfyUI_PuLID_Flux_ll installed? Error: {e}")
|
||||
else:
|
||||
# Use the registered node class
|
||||
ApplyPulidFlux = NODE_CLASS_MAPPINGS["ApplyPulidFlux"]
|
||||
apply_pulid = ApplyPulidFlux()
|
||||
|
||||
# Work directly on the original model
|
||||
current_model = model
|
||||
|
||||
# Get number of images in the list
|
||||
num_images = len(images)
|
||||
|
||||
# Process each image-mask pair
|
||||
for i in range(num_images):
|
||||
# Get single image from the list
|
||||
single_image = images[i].unsqueeze(0) if len(images[i].shape) == 3 else images[i]
|
||||
|
||||
# Extract corresponding mask if provided
|
||||
attn_mask = None
|
||||
if attn_masks is not None:
|
||||
if i < len(attn_masks):
|
||||
attn_mask = attn_masks[i].unsqueeze(0) if len(attn_masks[i].shape) == 3 else attn_masks[i]
|
||||
|
||||
# Apply PuLID Flux to this image-mask pair
|
||||
print(f"Processing image {i+1}/{num_images}")
|
||||
current_model, = apply_pulid.apply_pulid_flux(
|
||||
model=current_model,
|
||||
pulid_flux=pulid_flux,
|
||||
eva_clip=eva_clip,
|
||||
face_analysis=face_analysis,
|
||||
image=single_image,
|
||||
weight=weight,
|
||||
start_at=start_at,
|
||||
end_at=end_at,
|
||||
attn_mask=attn_mask,
|
||||
options=options,
|
||||
unique_id=unique_id
|
||||
)
|
||||
|
||||
return (current_model,)
|
||||
@@ -0,0 +1,308 @@
|
||||
import torch
|
||||
import os
|
||||
import folder_paths
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from transformers import (
|
||||
Qwen2_5_VLForConditionalGeneration,
|
||||
BitsAndBytesConfig,
|
||||
AutoProcessor,
|
||||
)
|
||||
from pathlib import Path
|
||||
import json
|
||||
|
||||
model_directory = os.path.join(folder_paths.models_dir, "VLM")
|
||||
os.makedirs(model_directory, exist_ok=True)
|
||||
|
||||
class AQ_QwenLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": (
|
||||
[
|
||||
"unsloth/Qwen2.5-VL-3B-Instruct-bnb-4bit",
|
||||
"unsloth/Qwen2.5-VL-7B-Instruct-bnb-4bit",
|
||||
"unsloth/Qwen2.5-VL-32B-Instruct-bnb-4bit",
|
||||
],
|
||||
{"default": "unsloth/Qwen2.5-VL-3B-Instruct-bnb-4bit"},
|
||||
),
|
||||
"quantization": (
|
||||
["none", "4bit"],
|
||||
{"default": "4bit"},
|
||||
),
|
||||
"attention": (
|
||||
["flash_attention_2", "sdpa", "eager"],
|
||||
{"default": "sdpa"},
|
||||
),
|
||||
"device": (
|
||||
["auto", "cuda:0", "cuda:1"],
|
||||
{"default": "auto"},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("QWEN_MODEL",)
|
||||
RETURN_NAMES = ("qwen_model",)
|
||||
FUNCTION = "load_model"
|
||||
CATEGORY = "AQ/LLM"
|
||||
|
||||
def load_model(self, model, quantization, attention, device):
|
||||
qwen_model = {"model": None, "processor": None, "model_path": None}
|
||||
|
||||
# Setup model path
|
||||
model_name = model.rsplit("/", 1)[-1]
|
||||
model_path = os.path.join(model_directory, model_name)
|
||||
|
||||
# Download if needed
|
||||
if not os.path.exists(model_path):
|
||||
print(f"Downloading Qwen model to: {model_path}")
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download(
|
||||
repo_id=model,
|
||||
local_dir=model_path,
|
||||
local_dir_use_symlinks=False
|
||||
)
|
||||
|
||||
# Setup quantization
|
||||
if quantization == "4bit":
|
||||
quantization_config = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_compute_dtype=torch.float16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
bnb_4bit_quant_type="nf4"
|
||||
)
|
||||
elif quantization == "8bit":
|
||||
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
|
||||
else:
|
||||
quantization_config = None
|
||||
|
||||
# Load model and processor
|
||||
qwen_model["model"] = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype=torch.float16,
|
||||
device_map=device,
|
||||
attn_implementation=attention,
|
||||
quantization_config=quantization_config,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
qwen_model["processor"] = AutoProcessor.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
use_fast=True
|
||||
)
|
||||
qwen_model["model_path"] = model_path
|
||||
|
||||
return (qwen_model,)
|
||||
|
||||
class AQ_Qwen:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
default_format = '''{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"user_idea": {
|
||||
"type": "string"
|
||||
},
|
||||
"new_scene_description": {
|
||||
"type": "string"
|
||||
},
|
||||
"new_image_type":{
|
||||
"type": "string"
|
||||
},
|
||||
"new_style":{
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"user_idea",
|
||||
"new_scene_description",
|
||||
"new_image_type",
|
||||
"new_style"
|
||||
]
|
||||
}'''
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"qwen_model": ("QWEN_MODEL",),
|
||||
"prompt": ("STRING", {"default": "", "multiline": True}),
|
||||
"system_message": ("STRING", {"default": "You are a helpful assistant.", "multiline": True}),
|
||||
"temperature": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 2.0, "step": 0.05}),
|
||||
"max_tokens": ("INT", {"default": 256, "min": 1, "max": 2048}),
|
||||
"top_p": ("FLOAT", {"default": 0.001, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"repetition_penalty": ("FLOAT", {"default": 1.05, "min": 1.0, "max": 2.0, "step": 0.01}),
|
||||
"min_pixels": ("INT", {"default": 224, "min": 16, "max": 1280}),
|
||||
"max_pixels": ("INT", {"default": 1280, "min": 224, "max": 2048}),
|
||||
"enable_json": ("BOOLEAN", {"default": False}),
|
||||
"json_format": ("STRING", {"default": default_format, "multiline": True}),
|
||||
"result_template": ("STRING", {
|
||||
"default": "{json[new_scene_description]} in style {json[new_style]}, {json[new_image_type]}",
|
||||
"multiline": True
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING",)
|
||||
RETURN_NAMES = ("response", "formatted_response")
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "AQ/LLM"
|
||||
|
||||
def format_response(self, response, enable_json, json_format, result_template):
|
||||
if not enable_json:
|
||||
return response, response
|
||||
|
||||
try:
|
||||
# First try to parse the response as JSON
|
||||
try:
|
||||
json_response = json.loads(response)
|
||||
except json.JSONDecodeError:
|
||||
# If response is not JSON, it means the model didn't follow the format
|
||||
print("Response is not valid JSON, returning as is")
|
||||
return response, response
|
||||
|
||||
# If result_template is provided, format using it
|
||||
if result_template:
|
||||
try:
|
||||
formatted_response = result_template.format(json=json_response)
|
||||
return response, formatted_response
|
||||
except KeyError as e:
|
||||
error_msg = f"Error: Missing required field in JSON response: {str(e)}"
|
||||
print(error_msg)
|
||||
return response, error_msg
|
||||
except Exception as e:
|
||||
error_msg = f"Error formatting template: {str(e)}"
|
||||
print(error_msg)
|
||||
return response, error_msg
|
||||
|
||||
return response, response
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in JSON formatting: {str(e)}")
|
||||
return response, response
|
||||
|
||||
def generate(self, qwen_model, prompt, system_message, temperature=0.1,
|
||||
max_tokens=256, top_p=0.001, repetition_penalty=1.05,
|
||||
min_pixels=224, max_pixels=1280, enable_json=False,
|
||||
json_format="", result_template="", image=None):
|
||||
try:
|
||||
min_pixels = min_pixels * min_pixels
|
||||
max_pixels = max_pixels * max_pixels
|
||||
|
||||
content = []
|
||||
|
||||
# Handle system message
|
||||
if system_message:
|
||||
messages = [{"role": "system", "content": system_message}]
|
||||
else:
|
||||
messages = []
|
||||
|
||||
# Handle image if provided
|
||||
user_content = []
|
||||
if image is not None:
|
||||
# Convert tensor to numpy array
|
||||
if isinstance(image, torch.Tensor):
|
||||
image = image.cpu().numpy()
|
||||
|
||||
# Handle ComfyUI image format (B, H, W, C)
|
||||
if len(image.shape) == 4:
|
||||
image = image[0] # Take first image from batch
|
||||
|
||||
# Ensure proper type and range for PIL
|
||||
if image.dtype != np.uint8:
|
||||
if image.max() <= 1:
|
||||
image = (image * 255).astype(np.uint8)
|
||||
else:
|
||||
image = image.astype(np.uint8)
|
||||
|
||||
# Convert numpy array to PIL Image and save temporarily
|
||||
img = Image.fromarray(image)
|
||||
image_path = Path(folder_paths.temp_directory) / "temp_qwen_image.png"
|
||||
img.save(str(image_path))
|
||||
|
||||
# Add image to content
|
||||
user_content.append({
|
||||
"type": "image",
|
||||
"image": f"file://{image_path.resolve().as_posix()}",
|
||||
"min_pixels": min_pixels,
|
||||
"max_pixels": max_pixels,
|
||||
})
|
||||
|
||||
# Modify prompt if JSON is enabled
|
||||
if enable_json and json_format:
|
||||
try:
|
||||
schema = json.loads(json_format)
|
||||
prompt = f"""Please provide your response in the following JSON format:
|
||||
{json.dumps(schema, indent=2)}
|
||||
|
||||
{prompt}"""
|
||||
except json.JSONDecodeError:
|
||||
print("Invalid JSON format provided, using prompt as is")
|
||||
|
||||
# Add prompt
|
||||
user_content.append({"type": "text", "text": prompt})
|
||||
messages.append({"role": "user", "content": user_content})
|
||||
|
||||
# Process with Qwen's template
|
||||
model_text = qwen_model["processor"].apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)
|
||||
|
||||
# Process vision info
|
||||
from qwen_vl_utils import process_vision_info
|
||||
image_inputs, video_inputs, video_kwargs = process_vision_info(
|
||||
messages,
|
||||
return_video_kwargs=True
|
||||
)
|
||||
|
||||
# Prepare inputs
|
||||
inputs = qwen_model["processor"](
|
||||
text=[model_text],
|
||||
images=image_inputs,
|
||||
padding=True,
|
||||
return_tensors="pt",
|
||||
**video_kwargs,
|
||||
)
|
||||
|
||||
# Move to correct device
|
||||
inputs = inputs.to(qwen_model["model"].device)
|
||||
|
||||
# Generate response
|
||||
generated_ids = qwen_model["model"].generate(
|
||||
**inputs,
|
||||
max_new_tokens=max_tokens,
|
||||
do_sample=True,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
repetition_penalty=repetition_penalty,
|
||||
pad_token_id=qwen_model["processor"].tokenizer.pad_token_id,
|
||||
eos_token_id=qwen_model["processor"].tokenizer.eos_token_id,
|
||||
)
|
||||
|
||||
# Process output
|
||||
generated_ids_trimmed = [
|
||||
out_ids[len(in_ids):]
|
||||
for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
|
||||
]
|
||||
|
||||
response = qwen_model["processor"].batch_decode(
|
||||
generated_ids_trimmed,
|
||||
skip_special_tokens=True,
|
||||
clean_up_tokenization_spaces=False,
|
||||
)[0]
|
||||
|
||||
# Format response if needed
|
||||
return self.format_response(response, enable_json, json_format, result_template)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in Qwen generation: {str(e)}")
|
||||
return ("")
|
||||
|
||||
finally:
|
||||
# Clean up temp files
|
||||
if image is not None and 'image_path' in locals() and image_path.exists():
|
||||
image_path.unlink()
|
||||
@@ -0,0 +1,14 @@
|
||||
huggingface_hub
|
||||
accelerate
|
||||
insightface
|
||||
onnxruntime-gpu
|
||||
facexlib
|
||||
segment-anything
|
||||
dill
|
||||
qwen-vl-utils
|
||||
qwen-vl-utils[decord]
|
||||
|
||||
# for gemini nodes
|
||||
google-generativeai>=0.3.2
|
||||
google-genai
|
||||
requests>=2.31.0
|
||||
Reference in New Issue
Block a user