commit ae094ea5d13c96a9c9de90a4ad7db5850221b88e Author: pk Date: Mon May 19 12:37:10 2025 +0200 init diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..8cc22cf --- /dev/null +++ b/__init__.py @@ -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'] diff --git a/nodes/api/AQ_Gemini.py b/nodes/api/AQ_Gemini.py new file mode 100644 index 0000000..b9fd8d3 --- /dev/null +++ b/nodes/api/AQ_Gemini.py @@ -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 \ No newline at end of file diff --git a/nodes/pulid/AQ_MasksAndImagesAsList.py b/nodes/pulid/AQ_MasksAndImagesAsList.py new file mode 100644 index 0000000..34b3ca5 --- /dev/null +++ b/nodes/pulid/AQ_MasksAndImagesAsList.py @@ -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) \ No newline at end of file diff --git a/nodes/pulid/AQ_multiface_ApplyPulidFlux.py b/nodes/pulid/AQ_multiface_ApplyPulidFlux.py new file mode 100644 index 0000000..5880328 --- /dev/null +++ b/nodes/pulid/AQ_multiface_ApplyPulidFlux.py @@ -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,) \ No newline at end of file diff --git a/nodes/transformers/AQ_Qwen.py b/nodes/transformers/AQ_Qwen.py new file mode 100644 index 0000000..3e96c61 --- /dev/null +++ b/nodes/transformers/AQ_Qwen.py @@ -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() \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..4836bf4 --- /dev/null +++ b/requirements.txt @@ -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 \ No newline at end of file