1714 lines
76 KiB
Python
1714 lines
76 KiB
Python
# IFLLMNode.py
|
|
import os
|
|
import sys
|
|
import json
|
|
import torch
|
|
import asyncio
|
|
import requests
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
import tempfile
|
|
import time
|
|
from typing import List, Dict, Any, Optional, Union, Tuple
|
|
from pathlib import Path
|
|
from .send_request import send_request
|
|
from .utils import (
|
|
get_api_key,
|
|
get_models,
|
|
process_images_for_comfy,
|
|
clean_text,
|
|
load_placeholder_image,
|
|
validate_models,
|
|
save_combo_settings,
|
|
load_combo_settings,
|
|
create_settings_from_ui,
|
|
prepare_batch_images,
|
|
process_auto_mode_images,
|
|
tensor_to_pil,
|
|
gemini2_process_images,
|
|
gemini2_prepare_response,
|
|
gemini2_create_client,
|
|
validate_gemini_key
|
|
)
|
|
import base64
|
|
import numpy as np
|
|
import codecs
|
|
import random
|
|
import math
|
|
|
|
# Add Google Gemini SDK imports
|
|
try:
|
|
from google import genai
|
|
from google.genai import types
|
|
GEMINI_SDK_AVAILABLE = True
|
|
except ImportError:
|
|
GEMINI_SDK_AVAILABLE = False
|
|
print("Google Generative AI SDK not found. Install with: pip install google-generativeai")
|
|
|
|
# Add ComfyUI directory to path
|
|
comfy_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
|
|
if comfy_path not in sys.path:
|
|
sys.path.insert(0, comfy_path)
|
|
|
|
try:
|
|
import folder_paths
|
|
except ImportError:
|
|
print("Error: Could not import folder_paths. Make sure ComfyUI core is in your Python path.")
|
|
folder_paths = None
|
|
|
|
# Set up logging
|
|
import logging
|
|
logging.basicConfig(level=logging.DEBUG)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
try:
|
|
from server import PromptServer
|
|
from aiohttp import web
|
|
|
|
@PromptServer.instance.routes.post("/IF_LLM/get_llm_models")
|
|
async def get_llm_models_endpoint(request):
|
|
try:
|
|
data = await request.json()
|
|
llm_provider = data.get("llm_provider")
|
|
engine = llm_provider
|
|
base_ip = data.get("base_ip")
|
|
port = data.get("port")
|
|
external_api_key = data.get("external_api_key")
|
|
|
|
if external_api_key:
|
|
api_key = external_api_key
|
|
else:
|
|
api_key_name = f"{llm_provider.upper()}_API_KEY"
|
|
try:
|
|
api_key = get_api_key(api_key_name, engine)
|
|
except ValueError:
|
|
api_key = None
|
|
|
|
node = IFLLM()
|
|
models = node.get_models(engine, base_ip, port, api_key)
|
|
return web.json_response(models)
|
|
|
|
except Exception as e:
|
|
print(f"Error in get_llm_models_endpoint: {str(e)}")
|
|
return web.json_response([], status=500)
|
|
|
|
@PromptServer.instance.routes.post("/IF_LLM/add_routes")
|
|
async def add_routes_endpoint(request):
|
|
return web.json_response({"status": "success"})
|
|
|
|
@PromptServer.instance.routes.post("/IF_LLM/save_combo_settings")
|
|
async def save_combo_settings_endpoint(request):
|
|
try:
|
|
data = await request.json()
|
|
|
|
# Convert UI settings to proper format
|
|
settings = create_settings_from_ui(data)
|
|
|
|
# Get node instance
|
|
node = IFLLM()
|
|
|
|
# Save settings
|
|
saved_settings = save_combo_settings(settings, node.combo_presets_dir)
|
|
|
|
return web.json_response({
|
|
"status": "success",
|
|
"message": "Combo settings saved successfully",
|
|
"settings": saved_settings
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error saving combo settings: {str(e)}")
|
|
return web.json_response({
|
|
"status": "error",
|
|
"message": str(e)
|
|
}, status=500)
|
|
|
|
except AttributeError:
|
|
print("PromptServer.instance not available. Skipping route decoration for IF_LLM.")
|
|
|
|
class IFLLM:
|
|
def __init__(self):
|
|
self.strategies = "normal"
|
|
# Initialize paths and load presets
|
|
# Get the directory where the current script is located
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
# Build paths relative to the script location
|
|
self.presets_dir = os.path.join(current_dir, "IF_AI", "presets")
|
|
self.combo_presets_dir = os.path.join(self.presets_dir, "AutoCombo")
|
|
# Load preset configurations
|
|
self.profiles = self.load_presets(os.path.join(self.presets_dir, "profiles.json"))
|
|
self.neg_prompts = self.load_presets(os.path.join(self.presets_dir, "neg_prompts.json"))
|
|
self.embellish_prompts = self.load_presets(os.path.join(self.presets_dir, "embellishments.json"))
|
|
self.style_prompts = self.load_presets(os.path.join(self.presets_dir, "style_prompts.json"))
|
|
self.stop_strings = self.load_presets(os.path.join(self.presets_dir, "stop_strings.json"))
|
|
|
|
# Initialize placeholder image path
|
|
self.placeholder_image_path = os.path.join(self.presets_dir, "placeholder.png")
|
|
|
|
# Default values
|
|
|
|
self.base_ip = "localhost"
|
|
self.port = "11434"
|
|
self.engine = "transformers"
|
|
self.selected_model = "Qwen2.5-VL-3B-Instruct-AWQ"
|
|
self.profile = "IF_PromptMKR_IMG"
|
|
self.messages = []
|
|
self.keep_alive = False
|
|
self.seed = 94687328150
|
|
self.history_steps = 10
|
|
self.external_api_key = ""
|
|
self.preset = "Default"
|
|
self.precision = "fp16"
|
|
self.attention = "sdpa"
|
|
self.Omni = None
|
|
self.mask = None
|
|
self.aspect_ratio = "1:1"
|
|
self.keep_alive = False
|
|
self.clear_history = False
|
|
self.random = False
|
|
self.max_tokens = 2048
|
|
self.temperature = 0.8
|
|
self.top_k = 40
|
|
self.top_p = 0.9
|
|
self.repeat_penalty = 1.1
|
|
self.batch_count = 4
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
node = cls()
|
|
return {
|
|
"required": {
|
|
"llm_provider": (["transformers","llamacpp", "ollama", "kobold", "lmstudio", "textgen", "groq", "gemini", "openai", "anthropic", "mistral","deepseek","xai"], {"default": "transformers"}),
|
|
"llm_model": ((), {}),
|
|
"base_ip": ("STRING", {"default": "localhost"}),
|
|
"port": ("STRING", {"default": "11434"}),
|
|
"user_prompt": ("STRING", {"multiline": True}),
|
|
},
|
|
"optional": {
|
|
"images": ("IMAGE", {"list": True}),
|
|
"strategy": (["normal", "omost", "create", "edit", "variations", "gemini2_create"], {"default": "normal"}),
|
|
"mask": ("MASK", {}),
|
|
"prime_directives": ("STRING", {"forceInput": True, "tooltip": "The system prompt for the LLM."}),
|
|
"profiles": (["None"] + list(cls().profiles.keys()), {"default": "None", "tooltip": "The pre-defined system_prompt from the json profile file on the presets folder you can edit or make your own will be listed here."}),
|
|
"embellish_prompt": (list(cls().embellish_prompts.keys()), {"tooltip": "The pre-defined embellishment from the json embellishments file on the presets folder you can edit or make your own will be listed here."}),
|
|
"style_prompt": (list(cls().style_prompts.keys()), {"tooltip": "The pre-defined style from the json style_prompts file on the presets folder you can edit or make your own will be listed here."}),
|
|
"neg_prompt": (list(cls().neg_prompts.keys()), {"tooltip": "The pre-defined negative prompt from the json neg_prompts file on the presets folder you can edit or make your own will be listed here."}),
|
|
"stop_string": (list(cls().stop_strings.keys()), {"tooltip": "Specifies a string at which text generation should stop."}),
|
|
"max_tokens": ("INT", {"default": 2048, "min": 1, "max": 8192, "tooltip": "Maximum number of tokens to generate in the response."}),
|
|
"random": ("BOOLEAN", {"default": False, "label_on": "Seed", "label_off": "Temperature", "tooltip": "Toggles between using a fixed seed or temperature-based randomness."}),
|
|
"seed": ("INT", {"default": 0, "tooltip": "Random seed for reproducible outputs."}),
|
|
"keep_alive": ("BOOLEAN", {"default": True, "label_on": "Keeps Model on Memory", "label_off": "Unloads Model from Memory", "tooltip": "Determines whether to keep the model loaded in memory between calls."}),
|
|
"clear_history": ("BOOLEAN", {"default": True, "label_on": "Clear History", "label_off": "Keep History", "tooltip": "Determines whether to clear the history between calls."}),
|
|
"history_steps": ("INT", {"default": 10, "tooltip": "Number of steps to keep in history."}),
|
|
"aspect_ratio": (["1:1", "16:9", "4:5", "3:4", "5:4", "9:16"], {"default": "1:1", "tooltip": "Aspect ratio for the generated images."}),
|
|
"auto": ("BOOLEAN", {"default": False, "label_on": "Auto Is Enabled", "label_off": "Auto is Disabled", "tooltip": "If true, it generates auto promts based on the listed images click the save Auto settings to set the auto prompt generation file"}),
|
|
"batch_count": ("INT", {"default": 1, "tooltip": "Number of images to generate. only for create, edit and variations strategies."}),
|
|
"external_api_key": ("STRING", {"default": "", "tooltip": "If this is not empty, it will be used instead of the API key from the .env file. Make sure it is empty to use the .env file."}),
|
|
"Omni": ("OMNI", {"default": None, "tooltip": "Additional input for the selected tool."}),
|
|
"attention": (["sdpa", "flash_attention_2", "xformers"], {"default": "sdpa", "tooltip": "Select attention mechanism on Transformer models."}),
|
|
},
|
|
"hidden": {
|
|
"temperature": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 1.0, "tooltip": "Controls randomness in output generation. Higher values increase creativity but may reduce coherence."}),
|
|
"top_k": ("INT", {"default": 40, "tooltip": "Limits the next token selection to the K most likely tokens."}),
|
|
"top_p": ("FLOAT", {"default": 0.9, "tooltip": "Cumulative probability cutoff for token selection."}),
|
|
"repeat_penalty": ("FLOAT", {"default": 1.1, "tooltip": "Penalizes repetition in generated text."}),
|
|
"precision": (["fp16", "fp32", "bf16"], {"tooltip": "Select preccision on Transformer models."}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING", "OMNI", "IMAGE", "MASK")
|
|
RETURN_NAMES = ("question", "response", "negative", "omni", "generated_images", "mask")
|
|
|
|
FUNCTION = "process_image_wrapper"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "ImpactFrames💥🎞️/IF_LLM"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, llm_provider, llm_model, **kwargs):
|
|
# Only report a change when the model or provider has actually changed
|
|
# This prevents ComfyUI from resetting the model selection
|
|
# Using a unique identifier instead of random to maintain consistency
|
|
import hashlib
|
|
|
|
# Create a unique hash based on the provider and model
|
|
unique_id = f"{llm_provider}:{llm_model}"
|
|
hash_obj = hashlib.md5(unique_id.encode())
|
|
|
|
# Return a deterministic value based on the hash
|
|
# This ensures the same provider/model combo always returns the same value
|
|
# but different combos return different values
|
|
return int(hash_obj.hexdigest(), 16) / (2**128)
|
|
|
|
async def process_image(
|
|
self,
|
|
llm_provider: str,
|
|
llm_model: str,
|
|
base_ip: str,
|
|
port: str,
|
|
user_prompt: str,
|
|
strategy: str = "normal",
|
|
images=None,
|
|
messages=None,
|
|
prime_directives: Optional[str] = None,
|
|
profiles: Optional[str] = None,
|
|
embellish_prompt: Optional[str] = None,
|
|
style_prompt: Optional[str] = None,
|
|
neg_prompt: Optional[str] = None,
|
|
stop_string: Optional[str] = None,
|
|
max_tokens: int = 2048,
|
|
seed: int = 0,
|
|
random: bool = False,
|
|
temperature: float = 0.8,
|
|
top_k: int = 40,
|
|
top_p: float = 0.9,
|
|
repeat_penalty: float = 1.1,
|
|
keep_alive: bool = False,
|
|
clear_history: bool = False,
|
|
history_steps: int = 10,
|
|
external_api_key: str = "",
|
|
precision: str = "fp16",
|
|
attention: str = "sdpa",
|
|
Omni: Optional[str] = None,
|
|
aspect_ratio: str = "1:1",
|
|
mask: Optional[torch.Tensor] = None,
|
|
batch_count: int = 4,
|
|
auto: bool = False,
|
|
auto_mode: bool = False,
|
|
**kwargs
|
|
) -> Union[str, Dict[str, Any]]:
|
|
try:
|
|
# Initialize variables at the start
|
|
formatted_response = None
|
|
generated_images = None
|
|
generated_masks = None
|
|
tool_output = None
|
|
|
|
current_images = None
|
|
current_mask = None
|
|
|
|
if external_api_key != "":
|
|
llm_api_key = external_api_key
|
|
else:
|
|
llm_api_key = get_api_key(f"{llm_provider.upper()}_API_KEY", llm_provider)
|
|
print(f"LLM API key: {llm_api_key[:5]}...")
|
|
|
|
# Validate LLM model
|
|
validate_models(llm_model, llm_provider, "LLM", base_ip, port, llm_api_key)
|
|
|
|
# Handle history
|
|
messages = messages or []
|
|
if clear_history:
|
|
messages = []
|
|
elif history_steps > 0:
|
|
messages = messages[-history_steps:]
|
|
|
|
|
|
# Handle stop
|
|
if stop_string is None or stop_string == "None":
|
|
stop_content = None
|
|
else:
|
|
stop_content = self.stop_strings.get(stop_string, None)
|
|
stop = stop_content
|
|
|
|
if llm_provider not in ["ollama", "llamacpp", "vllm", "lmstudio", "gemeni"]:
|
|
if llm_provider == "kobold":
|
|
stop = stop_content + \
|
|
["\n\n\n\n\n"] if stop_content else ["\n\n\n\n\n"]
|
|
elif llm_provider == "mistral":
|
|
stop = stop_content + \
|
|
["\n\n"] if stop_content else ["\n\n"]
|
|
else:
|
|
stop = stop_content if stop_content else None
|
|
|
|
# Prepare embellishments and styles
|
|
embellish_content = self.embellish_prompts.get(embellish_prompt, "").strip() if embellish_prompt else ""
|
|
style_content = self.style_prompts.get(style_prompt, "").strip() if style_prompt else ""
|
|
neg_content = self.neg_prompts.get(neg_prompt, "").strip() if neg_prompt else ""
|
|
profile_content = self.profiles.get(profiles, "")
|
|
|
|
# Prepare system prompt
|
|
if prime_directives is not None:
|
|
system_message = prime_directives
|
|
else:
|
|
system_message= json.dumps(profile_content)
|
|
|
|
tool_type = Omni
|
|
strategy_name = strategy
|
|
|
|
kwargs = {
|
|
'batch_count': batch_count,
|
|
'llm_provider': llm_provider,
|
|
'base_ip': base_ip,
|
|
'port': port,
|
|
'llm_model': llm_model,
|
|
'system_message': system_message,
|
|
'seed': seed,
|
|
'temperature': temperature,
|
|
'max_tokens': max_tokens,
|
|
'random': random,
|
|
'top_k': top_k,
|
|
'top_p': top_p,
|
|
'repeat_penalty': repeat_penalty,
|
|
'stop': stop,
|
|
'keep_alive': keep_alive,
|
|
'llm_api_key': llm_api_key,
|
|
'precision': precision,
|
|
'attention': attention,
|
|
'aspect_ratio': aspect_ratio,
|
|
'neg_prompt': neg_prompt,
|
|
'neg_content': neg_content,
|
|
'formatted_response': formatted_response,
|
|
'generated_images': generated_images,
|
|
'generated_masks': generated_masks,
|
|
'tool_output': tool_output,
|
|
'omni': tool_type,
|
|
}
|
|
|
|
# If images is None or empty, skip "image-based" logic but still allow LLM tasks to proceed
|
|
if images is not None and len(images) > 0:
|
|
current_images = images
|
|
else:
|
|
print("No images connected; continuing with text-based tasks only.")
|
|
|
|
# If no mask is connected, load a placeholder or just skip
|
|
if mask is not None:
|
|
current_mask = mask
|
|
else:
|
|
current_mask = load_placeholder_image(self.placeholder_image_path)[1]
|
|
|
|
if auto:
|
|
try:
|
|
# Use the main auto mode processing function
|
|
result = await self.process_auto_mode(
|
|
images=current_images,
|
|
mask=current_mask,
|
|
messages=messages,
|
|
strategy=strategy,
|
|
auto_mode=auto_mode,
|
|
**kwargs
|
|
)
|
|
|
|
if result:
|
|
return result
|
|
else:
|
|
#self, images, masks, error_message, prompt=""
|
|
return self.create_error_response(
|
|
current_images,
|
|
current_mask,
|
|
"No results generated from auto mode processing.",
|
|
user_prompt
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in auto mode processing: {str(e)}")
|
|
return self.create_error_response(
|
|
current_images,
|
|
current_mask,
|
|
"No results generated from auto mode processing.",
|
|
user_prompt
|
|
)
|
|
|
|
else:
|
|
# Execute strategy-specific logic
|
|
if strategy_name == "normal":
|
|
return await self.execute_normal_strategy(
|
|
user_prompt, current_images, current_mask, messages, embellish_content, style_content, **kwargs)
|
|
elif strategy_name == "create":
|
|
return await self.execute_create_strategy(
|
|
user_prompt, current_mask, **kwargs)
|
|
elif strategy_name == "omost":
|
|
return await self.execute_omost_strategy(
|
|
user_prompt, current_images, current_mask, embellish_content, style_content, **kwargs)
|
|
elif strategy_name == "variations":
|
|
return await self.execute_variations_strategy(
|
|
user_prompt, current_images, **kwargs)
|
|
elif strategy_name == "edit":
|
|
return await self.execute_edit_strategy(
|
|
user_prompt, current_images, current_mask, **kwargs)
|
|
elif strategy_name == "gemini2_create":
|
|
return await self.execute_gemini2_create_strategy(
|
|
user_prompt, current_images, current_mask, **kwargs)
|
|
else:
|
|
raise ValueError(f"Unsupported strategy: {strategy_name}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in process_image: {str(e)}")
|
|
return self.create_error_response(
|
|
current_images,
|
|
current_mask,
|
|
"No results generated from auto mode processing.",
|
|
user_prompt
|
|
)
|
|
|
|
async def process_auto_mode(self, images, mask, messages, strategy, auto_mode=True, embellish_content="", style_content="", **kwargs):
|
|
"""
|
|
Main auto mode processing function that preserves batch handling.
|
|
"""
|
|
try:
|
|
# Determine batch size based on mode
|
|
batch_size = 4 if auto_mode else 1
|
|
|
|
# Process images into appropriate batches
|
|
image_batches, mask_batches = process_auto_mode_images(
|
|
images=images,
|
|
mask=mask,
|
|
batch_size=batch_size
|
|
)
|
|
|
|
all_results = []
|
|
user_prompt = kwargs.get('user_prompt', '')
|
|
batch_count = kwargs.get('batch_count', 1)
|
|
|
|
# Process each image/mask batch
|
|
for img_batch, mask_batch in zip(image_batches, mask_batches):
|
|
|
|
for i in range(img_batch.size(0)):
|
|
single_img = img_batch[i:i+1]
|
|
single_mask = mask_batch[i:i+1]
|
|
|
|
# Generate combo prompt once for this image
|
|
combo_prompt = await self.generate_combo_prompts(
|
|
images=single_img,
|
|
settings_dict=None
|
|
)
|
|
|
|
# Process batch_count iterations for this image
|
|
for iteration in range(batch_count):
|
|
batch_results = await self.process_auto_batch(
|
|
batch_images=single_img,
|
|
batch_mask=single_mask,
|
|
strategy=strategy,
|
|
prompt=combo_prompt,
|
|
messages=messages,
|
|
embellish_content=embellish_content,
|
|
style_content=style_content,
|
|
**{**kwargs,
|
|
'batch_count': 1, # Process single iteration here
|
|
'seed': kwargs.get('seed', 0) + iteration if kwargs.get('seed') is not None else None
|
|
}
|
|
)
|
|
|
|
if batch_results:
|
|
if isinstance(batch_results, list):
|
|
all_results.extend(batch_results)
|
|
else:
|
|
all_results.append(batch_results)
|
|
|
|
if not all_results:
|
|
return [{
|
|
"Question": user_prompt,
|
|
"Response": "No results generated",
|
|
"Negative": "",
|
|
"Tool_Output": None,
|
|
"Retrieved_Image": images,
|
|
"Mask": mask
|
|
}]
|
|
|
|
return all_results
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in process_auto_mode: {str(e)}")
|
|
return [{
|
|
"Question": kwargs.get('user_prompt', ''),
|
|
"Response": f"Error: {str(e)}",
|
|
"Negative": "",
|
|
"Tool_Output": None,
|
|
"Retrieved_Image": images,
|
|
"Mask": mask
|
|
}]
|
|
|
|
async def process_auto_batch(self, batch_images, batch_mask, strategy, prompt, messages,
|
|
embellish_content="", style_content="", **kwargs):
|
|
"""
|
|
Process single iteration of auto mode batch.
|
|
Batch count iterations are handled by process_auto_mode.
|
|
"""
|
|
try:
|
|
# Create clean kwargs without user_prompt
|
|
batch_kwargs = {
|
|
k: v for k, v in kwargs.items()
|
|
if k not in ['user_prompt']
|
|
}
|
|
|
|
# Execute strategy (should process just one iteration)
|
|
if strategy == "normal":
|
|
results = await self.execute_normal_strategy(
|
|
user_prompt=prompt,
|
|
current_images=batch_images,
|
|
current_mask=batch_mask,
|
|
messages=messages,
|
|
embellish_content=embellish_content,
|
|
style_content=style_content,
|
|
**batch_kwargs
|
|
)
|
|
elif strategy == "omost":
|
|
results = await self.execute_omost_strategy(
|
|
user_prompt=prompt,
|
|
current_images=batch_images,
|
|
current_mask=batch_mask,
|
|
embellish_content=embellish_content,
|
|
style_content=style_content,
|
|
**batch_kwargs
|
|
)
|
|
else:
|
|
raise ValueError(f"Unsupported strategy for auto mode: {strategy}")
|
|
|
|
return results
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing auto batch: {str(e)}")
|
|
return None
|
|
|
|
async def execute_normal_strategy(self, user_prompt, current_images, current_mask, messages, embellish_content, style_content, **kwargs):
|
|
"""
|
|
Execute normal strategy with robust error handling and response validation.
|
|
"""
|
|
try:
|
|
results = []
|
|
batch_count = kwargs.get('batch_count', 1)
|
|
|
|
# Process and validate images
|
|
images_to_send = current_images if (current_images is not None and current_images.nelement() > 0) else None
|
|
|
|
# Process batch_count times
|
|
for i in range(batch_count):
|
|
try:
|
|
# Update seed for each iteration if using random seeding
|
|
current_seed = kwargs['seed'] + i if kwargs.get('random', False) and kwargs.get('seed') is not None else kwargs.get('seed')
|
|
|
|
# Make the API request
|
|
response = await send_request(
|
|
llm_provider=kwargs.get('llm_provider'),
|
|
base_ip=kwargs.get('base_ip'),
|
|
port=kwargs.get('port'),
|
|
images=images_to_send,
|
|
llm_model=kwargs.get('llm_model'),
|
|
system_message=kwargs.get('system_message'),
|
|
user_message=user_prompt,
|
|
messages=messages or [], # Ensure messages is never None
|
|
seed=current_seed,
|
|
temperature=kwargs.get('temperature', 0.7),
|
|
max_tokens=kwargs.get('max_tokens', 2048),
|
|
random=kwargs.get('random', False),
|
|
top_k=kwargs.get('top_k', 40),
|
|
top_p=kwargs.get('top_p', 0.9),
|
|
repeat_penalty=kwargs.get('repeat_penalty', 1.1),
|
|
stop=kwargs.get('stop'),
|
|
keep_alive=kwargs.get('keep_alive', False),
|
|
llm_api_key=kwargs.get('llm_api_key'),
|
|
precision=kwargs.get('precision', 'fp16'),
|
|
attention=kwargs.get('attention', 'sdpa'),
|
|
aspect_ratio=kwargs.get('aspect_ratio', '1:1'),
|
|
strategy="normal",
|
|
mask=current_mask
|
|
)
|
|
|
|
# Validate and extract response content
|
|
response_content = ""
|
|
if response is None:
|
|
logger.error("Received a None response from the LLM API.")
|
|
continue # Skip to the next iteration
|
|
elif isinstance(response, dict):
|
|
if "choices" in response and response["choices"]:
|
|
message = response["choices"][0].get("message", {})
|
|
response_content = message.get("content", "")
|
|
|
|
# Additional validation for empty content
|
|
if not response_content:
|
|
logger.warning("Empty response content in choices")
|
|
continue
|
|
|
|
elif "response" in response:
|
|
response_content = response["response"]
|
|
else:
|
|
logger.warning(f"Unexpected response format: {response}")
|
|
continue
|
|
|
|
elif isinstance(response, str):
|
|
response_content = response
|
|
|
|
if not response_content:
|
|
logger.warning("Empty response content received")
|
|
continue
|
|
|
|
# Proceed with cleaning and formatting the response
|
|
cleaned_response = clean_text(response_content)
|
|
final_prompt = "\n".join(filter(None, [
|
|
embellish_content.strip() if embellish_content else "",
|
|
cleaned_response.strip(),
|
|
style_content.strip() if style_content else ""
|
|
]))
|
|
|
|
# Generate negative prompt if needed
|
|
if kwargs.get('neg_prompt') == "AI_Fill":
|
|
neg_prompt = await self.generate_negative_prompt(
|
|
cleaned_response,
|
|
images=current_images,
|
|
**kwargs
|
|
)
|
|
else:
|
|
neg_prompt = kwargs.get('neg_content', '')
|
|
|
|
# Add result to results list
|
|
results.append({
|
|
"Question": user_prompt,
|
|
"Response": final_prompt,
|
|
"Negative": neg_prompt,
|
|
"Tool_Output": None,
|
|
"Retrieved_Image": current_images,
|
|
"Mask": current_mask
|
|
})
|
|
|
|
except Exception as batch_error:
|
|
logger.error(f"Error in batch {i}: {str(batch_error)}")
|
|
continue
|
|
|
|
# Keep message history if enabled
|
|
if kwargs.get('keep_alive') and results:
|
|
messages.extend([
|
|
{"role": "user", "content": user_prompt},
|
|
{"role": "assistant", "content": results[-1]["Response"]}
|
|
])
|
|
|
|
# Return results or error response
|
|
if not results:
|
|
return [self.create_error_response(
|
|
current_images,
|
|
current_mask,
|
|
"No valid results generated from normal strategy.",
|
|
user_prompt
|
|
)]
|
|
|
|
return results
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in normal strategy: {str(e)}")
|
|
return [self.create_error_response(
|
|
current_images,
|
|
current_mask,
|
|
f"Error in normal strategy: {str(e)}",
|
|
user_prompt
|
|
)]
|
|
|
|
async def execute_omost_strategy(
|
|
self, user_prompt, current_images, current_mask,
|
|
embellish_content="", style_content="", **kwargs
|
|
):
|
|
"""Execute OMOST strategy with batch processing and proper negative prompt generation"""
|
|
omni = kwargs.get("omni", None)
|
|
|
|
# Make sure user_prompt is a string, in case it's a list
|
|
if isinstance(user_prompt, list):
|
|
user_prompt = " ".join(user_prompt)
|
|
|
|
try:
|
|
batch_count = kwargs.get('batch_count', 1)
|
|
messages = []
|
|
system_prompt = self.profiles.get("IF_Omost")
|
|
results = []
|
|
|
|
logger.debug(f"Processing {batch_count} batches in OMOST strategy")
|
|
|
|
# Process batch_count times
|
|
for batch_idx in range(batch_count):
|
|
try:
|
|
# Get LLM response (dict or str).
|
|
llm_response = await send_request(
|
|
llm_provider=kwargs.get('llm_provider'),
|
|
base_ip=kwargs.get('base_ip'),
|
|
port=kwargs.get('port'),
|
|
images=current_images,
|
|
llm_model=kwargs.get('llm_model'),
|
|
system_message=system_prompt,
|
|
user_message=user_prompt,
|
|
messages=messages,
|
|
seed=kwargs.get('seed', 0) + batch_idx if kwargs.get('seed', 0) != 0 else kwargs.get('seed', 0),
|
|
temperature=kwargs.get('temperature', 0.7),
|
|
max_tokens=kwargs.get('max_tokens', 2048),
|
|
random=kwargs.get('random', False),
|
|
top_k=kwargs.get('top_k', 40),
|
|
top_p=kwargs.get('top_p', 0.9),
|
|
repeat_penalty=kwargs.get('repeat_penalty', 1.1),
|
|
stop=kwargs.get('stop', None),
|
|
keep_alive=kwargs.get('keep_alive', False),
|
|
llm_api_key=kwargs.get('llm_api_key'),
|
|
precision=kwargs.get('precision', 'fp16'),
|
|
attention=kwargs.get('attention', 'sdpa'),
|
|
aspect_ratio=kwargs.get('aspect_ratio', '1:1'),
|
|
strategy="omost",
|
|
mask=current_mask
|
|
)
|
|
|
|
if not llm_response:
|
|
logger.warning(f"No response from LLM in batch {batch_idx}")
|
|
continue
|
|
|
|
# If llm_response is dict, extract text from "choices"
|
|
# or fallback to stringifying.
|
|
if isinstance(llm_response, dict):
|
|
if "choices" in llm_response and llm_response["choices"]:
|
|
choice = llm_response["choices"][0]
|
|
if "message" in choice and "content" in choice["message"]:
|
|
llm_response = choice["message"]["content"]
|
|
else:
|
|
llm_response = json.dumps(llm_response)
|
|
elif "response" in llm_response:
|
|
llm_response = llm_response["response"]
|
|
else:
|
|
llm_response = json.dumps(llm_response)
|
|
elif not isinstance(llm_response, str):
|
|
llm_response = str(llm_response)
|
|
|
|
# IMPORTANT: Avoid calling clean_text() so Canvas code remains intact.
|
|
final_prompt = "\n".join(
|
|
filter(None,
|
|
[
|
|
embellish_content.strip(),
|
|
llm_response.strip(),
|
|
style_content.strip()
|
|
]
|
|
)
|
|
)
|
|
|
|
# Lazy load omost_function
|
|
omost_function = get_omost_function()
|
|
tool_result = await omost_function({
|
|
"name": "omost_tool",
|
|
"description": "Analyzes images composition and generates a Canvas representation.",
|
|
"system_prompt": system_prompt,
|
|
"input": user_prompt,
|
|
"llm_response": llm_response,
|
|
"function_call": None,
|
|
"omni_input": omni
|
|
})
|
|
|
|
# Handle negative prompt if requested
|
|
if kwargs.get('neg_prompt') == "AI_Fill":
|
|
neg_prompt = await self.generate_negative_prompt(
|
|
llm_response, # pass raw text if you want the LLM to see code
|
|
images=current_images,
|
|
**kwargs
|
|
)
|
|
else:
|
|
neg_prompt = kwargs.get('neg_content', '')
|
|
|
|
if isinstance(tool_result, dict):
|
|
if "error" in tool_result:
|
|
logger.warning(
|
|
f"OMOST tool warning in batch {batch_idx}: {tool_result['error']}"
|
|
)
|
|
continue
|
|
|
|
canvas_cond = tool_result.get("canvas_conditioning")
|
|
if canvas_cond is not None:
|
|
# Ensure canvas_conditioning is a flat list of dicts
|
|
if (
|
|
isinstance(canvas_cond, list)
|
|
and len(canvas_cond) == 1
|
|
and isinstance(canvas_cond[0], list)
|
|
):
|
|
# Flatten once
|
|
canvas_cond = canvas_cond[0]
|
|
tool_result["canvas_conditioning"] = canvas_cond
|
|
|
|
results.append({
|
|
"Question": user_prompt,
|
|
"Response": final_prompt,
|
|
"Negative": neg_prompt,
|
|
"Tool_Output": canvas_cond,
|
|
"Retrieved_Image": current_images,
|
|
"Mask": current_mask
|
|
})
|
|
|
|
except Exception as batch_error:
|
|
logger.error(f"Error in OMOST batch {batch_idx}: {str(batch_error)}")
|
|
continue
|
|
|
|
# Keep message history if enabled
|
|
if kwargs.get('keep_alive') and results:
|
|
messages.append({"role": "user", "content": user_prompt})
|
|
messages.append({"role": "assistant", "content": results[-1]["Response"]})
|
|
|
|
logger.debug(f"Generated {len(results)} results in OMOST strategy")
|
|
|
|
if not results:
|
|
return [self.create_error_response(
|
|
current_images,
|
|
current_mask,
|
|
"No valid results generated",
|
|
user_prompt
|
|
)]
|
|
|
|
return results
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in OMOST strategy: {str(e)}")
|
|
return [self.create_error_response(
|
|
current_images,
|
|
current_mask,
|
|
"No valid results generated",
|
|
user_prompt
|
|
)]
|
|
|
|
async def execute_create_strategy(self, user_prompt, current_mask, **kwargs):
|
|
try:
|
|
# Create strategy - no input images needed
|
|
messages = []
|
|
api_response = await send_request(
|
|
llm_provider=kwargs.get('llm_provider'),
|
|
base_ip=kwargs.get('base_ip'),
|
|
port=kwargs.get('port'),
|
|
images=None, # No input images needed for create
|
|
llm_model=kwargs.get('llm_model'),
|
|
system_message=kwargs.get('system_message'),
|
|
user_message=user_prompt,
|
|
messages=messages,
|
|
seed=kwargs.get('seed', 0),
|
|
temperature=kwargs.get('temperature'),
|
|
max_tokens=kwargs.get('max_tokens'),
|
|
random=kwargs.get('random'),
|
|
top_k=kwargs.get('top_k'),
|
|
top_p=kwargs.get('top_p'),
|
|
repeat_penalty=kwargs.get('repeat_penalty'),
|
|
stop=kwargs.get('stop'),
|
|
keep_alive=kwargs.get('keep_alive'),
|
|
llm_api_key=kwargs.get('llm_api_key'),
|
|
precision=kwargs.get('precision'),
|
|
attention=kwargs.get('attention'),
|
|
aspect_ratio=kwargs.get('aspect_ratio'),
|
|
strategy="create",
|
|
batch_count= 1,
|
|
mask=current_mask
|
|
)
|
|
|
|
# Extract base64 images from response
|
|
all_base64_images = []
|
|
if isinstance(api_response, dict) and "images" in api_response:
|
|
base64_images = api_response.get("images", [])
|
|
all_base64_images.extend(base64_images if isinstance(base64_images, list) else [base64_images])
|
|
|
|
# Process the images if we have any
|
|
if all_base64_images:
|
|
# Prepare data for processing
|
|
image_data = {
|
|
"data": [{"b64_json": img} for img in all_base64_images]
|
|
}
|
|
|
|
# Process images
|
|
images_tensor, mask_tensor = process_images_for_comfy(
|
|
image_data,
|
|
placeholder_image_path=self.placeholder_image_path,
|
|
response_key="data",
|
|
field_name="b64_json"
|
|
)
|
|
|
|
logger.debug(f"Retrieved_Image tensor shape: {images_tensor.shape}")
|
|
|
|
return {
|
|
"Question": user_prompt,
|
|
"Response": f"Create image{'s' if len(all_base64_images) > 1 else ''} successfully generated.",
|
|
"Negative": kwargs.get('neg_content', ''),
|
|
"Tool_Output": all_base64_images,
|
|
"Retrieved_Image": images_tensor,
|
|
"Mask": mask_tensor
|
|
}
|
|
else:
|
|
# No images were generated
|
|
image_tensor, mask_tensor = load_placeholder_image(self.placeholder_image_path)
|
|
return self.create_error_response(
|
|
image_tensor,
|
|
mask_tensor,
|
|
"No images were generated in create strategy",
|
|
user_prompt
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in create strategy: {str(e)}")
|
|
image_tensor, mask_tensor = load_placeholder_image(self.placeholder_image_path)
|
|
return self.create_error_response(
|
|
image_tensor,
|
|
mask_tensor,
|
|
f"Error in create strategy: {str(e)}",
|
|
user_prompt
|
|
)
|
|
|
|
async def execute_variations_strategy(self, user_prompt, images, **kwargs):
|
|
"""Core implementation of variations strategy"""
|
|
try:
|
|
batch_count = kwargs.get('batch_count', 1)
|
|
messages = []
|
|
api_responses = []
|
|
|
|
# Prepare input images
|
|
input_images = prepare_batch_images(images)
|
|
|
|
# Process each input image
|
|
for img in input_images:
|
|
try:
|
|
# Send request for variations
|
|
api_response = await send_request(
|
|
images=img,
|
|
user_message=user_prompt,
|
|
messages=messages,
|
|
strategy="variations",
|
|
batch_count=batch_count,
|
|
mask=None, # Variations don't use masks
|
|
**kwargs
|
|
)
|
|
if api_response:
|
|
api_responses.append(api_response)
|
|
except Exception as e:
|
|
logger.error(f"Error processing image variation: {str(e)}")
|
|
continue
|
|
|
|
# Extract and process base64 images from responses
|
|
all_base64_images = []
|
|
for response in api_responses:
|
|
if isinstance(response, dict) and "images" in response:
|
|
base64_images = response.get("images", [])
|
|
if isinstance(base64_images, list):
|
|
all_base64_images.extend(base64_images)
|
|
else:
|
|
all_base64_images.append(base64_images)
|
|
|
|
# Process the generated images
|
|
if all_base64_images:
|
|
# Prepare data for processing
|
|
image_data = {
|
|
"data": [{"b64_json": img} for img in all_base64_images]
|
|
}
|
|
|
|
# Convert to tensors
|
|
images_tensor, mask_tensor = process_images_for_comfy(
|
|
image_data,
|
|
placeholder_image_path=self.placeholder_image_path,
|
|
response_key="data",
|
|
field_name="b64_json"
|
|
)
|
|
|
|
logger.debug(f"Variations image tensor shape: {images_tensor.shape}")
|
|
|
|
return {
|
|
"Question": user_prompt,
|
|
"Response": f"Generated {len(all_base64_images)} variations successfully.",
|
|
"Negative": kwargs.get('neg_content', ''),
|
|
"Tool_Output": all_base64_images,
|
|
"Retrieved_Image": images_tensor,
|
|
"Mask": mask_tensor
|
|
}
|
|
else:
|
|
# No variations were generated
|
|
image_tensor, mask_tensor = load_placeholder_image(self.placeholder_image_path)
|
|
return self.create_error_response(
|
|
image_tensor,
|
|
mask_tensor,
|
|
"No variations were generated",
|
|
user_prompt
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in variations strategy: {str(e)}")
|
|
image_tensor, mask_tensor = load_placeholder_image(self.placeholder_image_path)
|
|
return self.create_error_response(
|
|
image_tensor,
|
|
mask_tensor,
|
|
f"Error in variations strategy: {str(e)}",
|
|
user_prompt
|
|
)
|
|
|
|
async def execute_edit_strategy(self, user_prompt, images, mask, **kwargs):
|
|
"""Core implementation of edit strategy"""
|
|
try:
|
|
batch_count = kwargs.get('batch_count', 1)
|
|
messages = []
|
|
api_responses = []
|
|
|
|
# Prepare input images and masks
|
|
input_images = prepare_batch_images(images)
|
|
input_masks = prepare_batch_images(mask) if mask is not None else [None] * len(input_images)
|
|
|
|
# Process each image-mask pair
|
|
for img, msk in zip(input_images, input_masks):
|
|
try:
|
|
# Send request for edit
|
|
api_response = await send_request(
|
|
images=img,
|
|
user_message=user_prompt,
|
|
messages=messages,
|
|
strategy="edit",
|
|
batch_count=batch_count,
|
|
mask=msk,
|
|
**kwargs
|
|
)
|
|
if api_response:
|
|
api_responses.append(api_response)
|
|
except Exception as e:
|
|
logger.error(f"Error processing image-mask pair: {str(e)}")
|
|
continue
|
|
|
|
# Extract and process base64 images from responses
|
|
all_base64_images = []
|
|
for response in api_responses:
|
|
if isinstance(response, dict) and "images" in response:
|
|
base64_images = response.get("images", [])
|
|
if isinstance(base64_images, list):
|
|
all_base64_images.extend(base64_images)
|
|
else:
|
|
all_base64_images.append(base64_images)
|
|
|
|
# Process the edited images
|
|
if all_base64_images:
|
|
# Prepare data for processing
|
|
image_data = {
|
|
"data": [{"b64_json": img} for img in all_base64_images]
|
|
}
|
|
|
|
# Convert to tensors
|
|
images_tensor, mask_tensor = process_images_for_comfy(
|
|
image_data,
|
|
placeholder_image_path=self.placeholder_image_path,
|
|
response_key="data",
|
|
field_name="b64_json"
|
|
)
|
|
|
|
logger.debug(f"Edited image tensor shape: {images_tensor.shape}")
|
|
|
|
return {
|
|
"Question": user_prompt,
|
|
"Response": f"Generated {len(all_base64_images)} variations successfully.",
|
|
"Negative": kwargs.get('neg_content', ''),
|
|
"Tool_Output": all_base64_images,
|
|
"Retrieved_Image": images_tensor,
|
|
"Mask": mask_tensor
|
|
}
|
|
else:
|
|
# No edits were generated
|
|
image_tensor, mask_tensor = load_placeholder_image(self.placeholder_image_path)
|
|
return self.create_error_response(
|
|
image_tensor,
|
|
mask_tensor,
|
|
"No edited images were generated",
|
|
user_prompt
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in edit strategy: {str(e)}")
|
|
image_tensor, mask_tensor = load_placeholder_image(self.placeholder_image_path)
|
|
return self.create_error_response(
|
|
image_tensor,
|
|
mask_tensor,
|
|
f"Error in edit strategy: {str(e)}",
|
|
user_prompt
|
|
)
|
|
|
|
async def execute_gemini2_create_strategy(self, user_prompt, current_images, current_mask=None, **kwargs):
|
|
"""
|
|
Execute Gemini 2.0 create strategy using the Google Gemini API SDK.
|
|
Handles batches of images as input and returns generated images.
|
|
|
|
Args:
|
|
user_prompt (str): The prompt for image generation
|
|
current_images (torch.Tensor): Batch of input images [B,H,W,C]
|
|
current_mask (torch.Tensor, optional): Mask tensor
|
|
**kwargs: Additional arguments including API key, model settings, etc.
|
|
|
|
Returns:
|
|
dict: Response dictionary with generated images and other metadata
|
|
"""
|
|
try:
|
|
# Check if Gemini SDK is available
|
|
if not GEMINI_SDK_AVAILABLE:
|
|
error_msg = "Google Generative AI SDK not installed. Install with: pip install google-generativeai"
|
|
logger.error(error_msg)
|
|
return self.create_error_response(
|
|
current_images,
|
|
current_mask,
|
|
error_msg,
|
|
user_prompt
|
|
)
|
|
|
|
# Initialize variables for response
|
|
response_text = ""
|
|
temp_img_paths = []
|
|
|
|
# Get API key
|
|
if kwargs.get('external_api_key'):
|
|
api_key = kwargs.get('external_api_key')
|
|
else:
|
|
api_key = kwargs.get('llm_api_key')
|
|
|
|
if not api_key:
|
|
logger.error("No valid Gemini API key provided")
|
|
return self.create_error_response(
|
|
current_images,
|
|
current_mask,
|
|
"Error: No valid Gemini API key provided. Please set GEMINI_API_KEY in your environment or provide external_api_key.",
|
|
user_prompt
|
|
)
|
|
|
|
# Process parameters
|
|
temperature = kwargs.get('temperature', 0.8)
|
|
seed = kwargs.get('seed', 0)
|
|
batch_count = kwargs.get('batch_count', 1)
|
|
|
|
# Use random seed if seed is 0 or random is True
|
|
if seed == 0 or kwargs.get('random', False):
|
|
import random
|
|
seed = random.randint(1, 2**31 - 1)
|
|
|
|
logger.info(f"Using Gemini 2.0 Create strategy with seed: {seed}, temperature: {temperature}")
|
|
|
|
# Create Gemini client
|
|
client = genai.Client(api_key=api_key)
|
|
|
|
# Process input images
|
|
if current_images is not None and current_images.nelement() > 0:
|
|
# Prepare input images for Gemini API
|
|
input_images = prepare_batch_images(current_images)
|
|
logger.info(f"Processing {len(input_images)} input images for Gemini")
|
|
|
|
# Convert images to format required by Gemini
|
|
contents = []
|
|
|
|
# Add each image to the request
|
|
for idx, img in enumerate(input_images):
|
|
try:
|
|
# Convert tensor to PIL image
|
|
pil_image = tensor_to_pil(img)
|
|
|
|
# Save as temporary file
|
|
temp_img_path = os.path.join(tempfile.gettempdir(), f"gemini_input_{idx}_{int(time.time())}.png")
|
|
pil_image.save(temp_img_path)
|
|
temp_img_paths.append(temp_img_path)
|
|
|
|
# Read image data
|
|
with open(temp_img_path, "rb") as f:
|
|
image_bytes = f.read()
|
|
|
|
# Add image to content
|
|
contents.append({
|
|
"inline_data": {
|
|
"mime_type": "image/png",
|
|
"data": image_bytes
|
|
}
|
|
})
|
|
|
|
except Exception as img_error:
|
|
logger.error(f"Error processing input image {idx}: {str(img_error)}")
|
|
|
|
# Add the prompt after all images
|
|
contents.append({"text": user_prompt})
|
|
else:
|
|
# No input images, just use the prompt
|
|
contents = user_prompt
|
|
logger.info("No input images provided, using text prompt only")
|
|
|
|
# Configure generation parameters
|
|
gen_config = types.GenerateContentConfig(
|
|
temperature=temperature,
|
|
seed=seed,
|
|
response_modalities=['Text', 'Image']
|
|
# Request multiple images based on batch_count
|
|
# generation_parameters parameter is not supported and causing errors
|
|
# generation_parameters={
|
|
# "num_iterations": batch_count
|
|
# }
|
|
)
|
|
|
|
# Note: Gemini 2.0 API doesn't support the num_iterations parameter directly
|
|
# It can return multiple images for some prompts but doesn't guarantee batch_count
|
|
# The API will decide how many images to return based on the prompt
|
|
|
|
# Call Gemini API
|
|
logger.info(f"Calling Gemini API with {len(contents) if isinstance(contents, list) else 1} content parts")
|
|
response = client.models.generate_content(
|
|
model="models/gemini-2.0-flash-exp", # Using the latest image generation model
|
|
contents=contents,
|
|
config=gen_config
|
|
)
|
|
|
|
logger.info("Received response from Gemini API")
|
|
|
|
# Process the response to extract generated images
|
|
if not hasattr(response, 'candidates') or not response.candidates:
|
|
logger.error("API response contained no candidates")
|
|
return self.create_error_response(
|
|
current_images,
|
|
current_mask,
|
|
"Error: Gemini API returned no candidates in the response",
|
|
user_prompt
|
|
)
|
|
|
|
# Extract generated images and text
|
|
generated_images = []
|
|
|
|
for candidate_idx, candidate in enumerate(response.candidates):
|
|
if not hasattr(candidate, 'content') or not hasattr(candidate.content, 'parts'):
|
|
continue
|
|
|
|
for part in candidate.content.parts:
|
|
# Extract text content
|
|
if hasattr(part, 'text') and part.text:
|
|
response_text += part.text + "\n"
|
|
|
|
# Extract image content
|
|
if hasattr(part, 'inline_data') and part.inline_data:
|
|
try:
|
|
# Get binary image data
|
|
image_binary = part.inline_data.data
|
|
generated_images.append(image_binary)
|
|
logger.info(f"Extracted image {len(generated_images)} from response")
|
|
except Exception as img_error:
|
|
logger.error(f"Error extracting image from response: {str(img_error)}")
|
|
|
|
# Clean up temporary files
|
|
for temp_path in temp_img_paths:
|
|
try:
|
|
if os.path.exists(temp_path):
|
|
os.remove(temp_path)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to remove temporary file {temp_path}: {str(e)}")
|
|
|
|
# If no images were generated, return error
|
|
if not generated_images:
|
|
logger.warning("No images found in Gemini API response")
|
|
return self.create_error_response(
|
|
current_images,
|
|
current_mask,
|
|
f"No images generated. API response: {response_text[:500]}...",
|
|
user_prompt
|
|
)
|
|
|
|
# Process generated images for ComfyUI
|
|
image_data = {
|
|
"data": [{"b64_json": base64.b64encode(img).decode('utf-8')} for img in generated_images]
|
|
}
|
|
|
|
# Convert binary image data to tensors
|
|
images_tensor, mask_tensor = process_images_for_comfy(
|
|
image_data,
|
|
placeholder_image_path=self.placeholder_image_path,
|
|
response_key="data",
|
|
field_name="b64_json"
|
|
)
|
|
|
|
logger.info(f"Successfully processed {len(generated_images)} generated images")
|
|
|
|
return {
|
|
"Question": user_prompt,
|
|
"Response": f"Generated {len(generated_images)} images with Gemini 2.0.\n\n{response_text}",
|
|
"Negative": kwargs.get('neg_content', ''),
|
|
"Tool_Output": generated_images,
|
|
"Retrieved_Image": images_tensor,
|
|
"Mask": mask_tensor
|
|
}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in Gemini 2.0 create strategy: {str(e)}", exc_info=True)
|
|
return self.create_error_response(
|
|
current_images,
|
|
current_mask,
|
|
f"Error in Gemini 2.0 create strategy: {str(e)}",
|
|
user_prompt
|
|
)
|
|
|
|
def get_models(self, engine, base_ip, port, api_key=None):
|
|
return get_models(engine, base_ip, port, api_key)
|
|
|
|
def load_presets(self, file_path: str) -> Dict[str, Any]:
|
|
"""
|
|
Load JSON presets with support for multiple encodings and better error handling.
|
|
|
|
Args:
|
|
file_path (str): Path to the JSON preset file
|
|
|
|
Returns:
|
|
Dict[str, Any]: Loaded JSON data or empty dict if loading fails
|
|
"""
|
|
# List of encodings to try
|
|
encodings = ['utf-8', 'utf-8-sig', 'latin1', 'cp1252', 'gbk']
|
|
|
|
for encoding in encodings:
|
|
try:
|
|
with codecs.open(file_path, 'r', encoding=encoding) as f:
|
|
content = f.read()
|
|
|
|
# Debug: Print problematic content around error location
|
|
try:
|
|
data = json.loads(content)
|
|
except json.JSONDecodeError as je:
|
|
# Get context around the error
|
|
start = max(0, je.pos - 50)
|
|
end = min(len(content), je.pos + 50)
|
|
context = content[start:end]
|
|
|
|
print(f"\nError details for {file_path}:")
|
|
print(f"Error position: Line {je.lineno}, Column {je.colno}")
|
|
print(f"Context around error:\n{context}")
|
|
print(f"Error message: {str(je)}")
|
|
continue
|
|
|
|
# Only rewrite if encoding was NOT utf-8 or utf-8-sig
|
|
if encoding.lower() not in ('utf-8', 'utf-8-sig'):
|
|
try:
|
|
with codecs.open(file_path, 'w', encoding='utf-8') as out_f:
|
|
json.dump(data, out_f, ensure_ascii=False, indent=2)
|
|
except Exception as write_err:
|
|
print(f"Warning: Could not write back UTF-8 encoded file: {write_err}")
|
|
|
|
return data
|
|
|
|
except UnicodeDecodeError:
|
|
print(f"Unicode decode error with {encoding} encoding")
|
|
continue
|
|
except Exception as e:
|
|
print(f"Error loading presets from {file_path} with {encoding} encoding: {e}")
|
|
continue
|
|
|
|
# If all attempts fail, try to load a backup or create empty dict
|
|
try:
|
|
backup_path = file_path + '.backup'
|
|
if os.path.exists(backup_path):
|
|
print(f"Attempting to load backup file: {backup_path}")
|
|
with codecs.open(backup_path, 'r', encoding='utf-8') as f:
|
|
return json.load(f)
|
|
except Exception as e:
|
|
print(f"Error loading backup file: {e}")
|
|
|
|
print(f"Error: Failed to load {file_path} with any supported encoding")
|
|
return {}
|
|
|
|
def validate_outputs(self, outputs):
|
|
"""Helper to validate output types match expectations"""
|
|
if len(outputs) != len(self.RETURN_TYPES):
|
|
raise ValueError(
|
|
f"Expected {len(self.RETURN_TYPES)} outputs, got {len(outputs)}"
|
|
)
|
|
|
|
for i, (output, expected_type) in enumerate(zip(outputs, self.RETURN_TYPES)):
|
|
if output is None and expected_type in ["IMAGE", "MASK"]:
|
|
raise ValueError(
|
|
f"Output {i} ({self.RETURN_NAMES[i]}) cannot be None for type {expected_type}"
|
|
)
|
|
|
|
async def generate_combo_prompts(self, images, settings_dict=None, **kwargs):
|
|
try:
|
|
if settings_dict is None:
|
|
settings_dict = load_combo_settings(self.combo_presets_dir)
|
|
|
|
if not settings_dict:
|
|
raise ValueError("No combo settings available")
|
|
|
|
# Get the profile content
|
|
profile_name = settings_dict.get('profile', 'IF_PromptMKR')
|
|
profile_content = self.profiles.get(profile_name, {}).get('instruction', '')
|
|
|
|
if not settings_dict.get('prime_directives'):
|
|
settings_dict['prime_directives'] = profile_content
|
|
|
|
# Extract API key
|
|
llm_provider = settings_dict.get('llm_provider', '')
|
|
if settings_dict.get('external_api_key'):
|
|
llm_api_key = settings_dict['external_api_key']
|
|
else:
|
|
llm_api_key = get_api_key(f"{llm_provider.upper()}_API_KEY", llm_provider)
|
|
|
|
# Create request parameters with correct mappings
|
|
request_params = {
|
|
'llm_provider': settings_dict.get('llm_provider', ''),
|
|
'base_ip': settings_dict.get('base_ip', 'localhost'),
|
|
'port': settings_dict.get('port', '11434'),
|
|
'images': images,
|
|
'llm_model': settings_dict.get('llm_model', ''),
|
|
'system_message': settings_dict.get('prime_directives', ''), # Map prime_directives to system_message
|
|
'user_message': settings_dict.get('user_prompt', ''), # Map user_prompt to user_message
|
|
'messages': [],
|
|
'seed': settings_dict.get('seed', None),
|
|
'temperature': settings_dict.get('temperature', 0.7),
|
|
'max_tokens': settings_dict.get('max_tokens', 2048),
|
|
'random': settings_dict.get('random', False),
|
|
'top_k': settings_dict.get('top_k', 40),
|
|
'top_p': settings_dict.get('top_p', 0.9),
|
|
'repeat_penalty': settings_dict.get('repeat_penalty', 1.1),
|
|
'stop': settings_dict.get('stop_string', None), # Map stop_string to stop
|
|
'keep_alive': settings_dict.get('keep_alive', False),
|
|
'llm_api_key': llm_api_key,
|
|
'precision': settings_dict.get('precision', 'fp16'),
|
|
'attention': settings_dict.get('attention', 'sdpa'),
|
|
'aspect_ratio': settings_dict.get('aspect_ratio', '1:1'),
|
|
'strategy': 'normal',
|
|
'mask': None,
|
|
'batch_count': settings_dict.get('batch_count', 1)
|
|
}
|
|
|
|
response = await send_request(**request_params)
|
|
|
|
if isinstance(response, dict):
|
|
return response.get('response', '')
|
|
return response
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error generating combo prompts: {str(e)}")
|
|
return ""
|
|
|
|
def process_image_wrapper(self, **kwargs):
|
|
"""Wrapper to handle async execution of process_image"""
|
|
try:
|
|
# Attempt to get the current event loop
|
|
try:
|
|
loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
# Create a new event loop if one doesn't exist
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
|
|
# Validate required inputs
|
|
required_params = ['llm_provider', 'llm_model', 'base_ip', 'port', 'user_prompt']
|
|
missing_params = [p for p in required_params if p not in kwargs]
|
|
if missing_params:
|
|
raise ValueError(f"Missing required parameters: {', '.join(missing_params)}")
|
|
|
|
# Execute the asynchronous process_image method
|
|
result = loop.run_until_complete(self.process_image(**kwargs))
|
|
|
|
# Initialize aggregation lists
|
|
responses = []
|
|
prompts = []
|
|
negatives = []
|
|
omnis = []
|
|
retrieved_images = []
|
|
masks = []
|
|
|
|
# Aggregate results based on their type
|
|
if isinstance(result, list):
|
|
if not result:
|
|
raise ValueError("No results generated")
|
|
|
|
for result_item in result:
|
|
if isinstance(result_item, dict):
|
|
prompts.append(result_item.get("Response", ""))
|
|
responses.append(result_item.get("Question", ""))
|
|
negatives.append(result_item.get("Negative", ""))
|
|
omnis.append(result_item.get("Tool_Output"))
|
|
|
|
# Safely handle None for images/masks
|
|
img = result_item.get("Retrieved_Image")
|
|
msk = result_item.get("Mask")
|
|
if not isinstance(img, torch.Tensor):
|
|
# Replace None with a placeholder image
|
|
placeholder_img, placeholder_mask = load_placeholder_image(self.placeholder_image_path)
|
|
img = placeholder_img
|
|
# If we have a real mask, use it, else the same placeholder
|
|
if not isinstance(msk, torch.Tensor):
|
|
msk = placeholder_mask
|
|
elif not isinstance(msk, torch.Tensor):
|
|
# If the image is valid but mask isn't, load just the placeholder mask
|
|
_, msk = load_placeholder_image(self.placeholder_image_path)
|
|
|
|
# Convert 4-channel images to 3-channel if needed
|
|
if isinstance(img, torch.Tensor):
|
|
if img.dim() == 4: # Batch of images [B, C, H, W]
|
|
if img.shape[1] == 4: # Check if channels = 4
|
|
# Convert RGBA to RGB by removing alpha channel
|
|
img = img[:, :3, :, :]
|
|
logger.debug(f"Converted batch of 4-channel images to 3-channel, new shape: {img.shape}")
|
|
elif img.dim() == 3: # Single image [C, H, W]
|
|
if img.shape[0] == 4: # Check if channels = 4
|
|
# Convert RGBA to RGB by removing alpha channel
|
|
img = img[:3, :, :]
|
|
logger.debug(f"Converted single 4-channel image to 3-channel, new shape: {img.shape}")
|
|
|
|
retrieved_images.append(img)
|
|
masks.append(msk)
|
|
|
|
else:
|
|
raise ValueError(f"Unexpected result format: {type(result_item)}")
|
|
|
|
elif isinstance(result, dict):
|
|
prompts.append(result.get("Response", ""))
|
|
responses.append(result.get("Question", ""))
|
|
negatives.append(result.get("Negative", ""))
|
|
omnis.append(result.get("Tool_Output"))
|
|
|
|
# Same handling for the single dictionary return
|
|
img = result.get("Retrieved_Image")
|
|
msk = result.get("Mask")
|
|
if not isinstance(img, torch.Tensor):
|
|
placeholder_img, placeholder_mask = load_placeholder_image(self.placeholder_image_path)
|
|
img = placeholder_img
|
|
if not isinstance(msk, torch.Tensor):
|
|
msk = placeholder_mask
|
|
elif not isinstance(msk, torch.Tensor):
|
|
_, msk = load_placeholder_image(self.placeholder_image_path)
|
|
|
|
retrieved_images.append(img)
|
|
masks.append(msk)
|
|
|
|
else:
|
|
raise ValueError(f"Unexpected result type: {type(result)}")
|
|
|
|
# Concatenate image tensors
|
|
retrieved_images_tensor = torch.cat(retrieved_images, dim=0) if retrieved_images else load_placeholder_image(self.placeholder_image_path)[0]
|
|
|
|
# Concatenate mask tensors
|
|
masks_tensor = torch.cat(masks, dim=0) if masks else load_placeholder_image(self.placeholder_image_path)[1]
|
|
|
|
# Debug logging for verification
|
|
for idx in range(len(retrieved_images)):
|
|
logger.debug(f"Result {idx + 1}: Retrieved image type: {type(retrieved_images[idx])}")
|
|
if isinstance(retrieved_images[idx], torch.Tensor):
|
|
logger.debug(f"Result {idx + 1}: Retrieved image shape: {retrieved_images[idx].shape}")
|
|
logger.debug(f"Result {idx + 1}: Mask type: {type(masks[idx])}")
|
|
if isinstance(masks[idx], torch.Tensor):
|
|
logger.debug(f"Result {idx + 1}: Mask shape: {masks[idx].shape}")
|
|
|
|
# Ensure masks_tensor has the expected shape
|
|
# Expected: [batch_size, 1, H, W]
|
|
# If masks_tensor is not in the correct shape, adjust accordingly
|
|
if masks_tensor.dim() == 3:
|
|
masks_tensor = masks_tensor.unsqueeze(1) # Add channel dimension if missing
|
|
|
|
# Return the aggregated results
|
|
return (
|
|
responses, # List of STRING (questions/prompts)
|
|
prompts, # List of STRING (generated responses)
|
|
negatives, # List of STRING (negative prompts)
|
|
omnis, # List of OMNI
|
|
retrieved_images_tensor, # Concatenated IMAGE tensors [batch, 3, H, W]
|
|
masks_tensor # Concatenated MASK tensors [batch, 1, H, W]
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in process_image_wrapper: {str(e)}")
|
|
# Create fallback values as lists to match RETURN_TYPES
|
|
image_tensor, mask_tensor = load_placeholder_image(self.placeholder_image_path)
|
|
return (
|
|
[kwargs.get("user_prompt", "")], # List containing original prompt
|
|
[f"Error: {str(e)}"], # List containing error message as response
|
|
[""], # List containing empty negative prompt
|
|
[None], # List containing no OMNI data
|
|
image_tensor, # Single tensor
|
|
mask_tensor # Single tensor
|
|
)
|
|
|
|
def create_error_response(self, images, masks, error_message, prompt=""):
|
|
"""Create standardized error response"""
|
|
try:
|
|
if images is None:
|
|
image_tensor = load_placeholder_image(self.placeholder_image_path)[0]
|
|
else:
|
|
image_tensor = images
|
|
if masks is None:
|
|
mask_tensor = load_placeholder_image(self.placeholder_image_path)[1]
|
|
else:
|
|
mask_tensor = masks
|
|
return {
|
|
"Question": prompt,
|
|
"Response": f"Error: {error_message}",
|
|
"Negative": f"Error: {error_message}",
|
|
"Tool_Output": None,
|
|
"Retrieved_Image": image_tensor,
|
|
"Mask": mask_tensor
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"Error creating error response: {str(e)}")
|
|
# Fallback error response without images
|
|
return {
|
|
"Question": prompt,
|
|
"Response": f"Critical Error: {error_message}",
|
|
"Negative": f"Error: {error_message}",
|
|
"Tool_Output": None,
|
|
"Retrieved_Image": None,
|
|
"Mask": None
|
|
}
|
|
|
|
async def generate_negative_prompt(
|
|
self,
|
|
prompt: str,
|
|
images: List[Image.Image],
|
|
**kwargs
|
|
) -> List[str]:
|
|
"""
|
|
Generate negative prompts for the given input prompt.
|
|
|
|
Args:
|
|
prompt: Input prompt text
|
|
**kwargs: Generation parameters like seed, temperature etc
|
|
|
|
Returns:
|
|
List of generated negative prompts
|
|
"""
|
|
try:
|
|
if not prompt:
|
|
return []
|
|
|
|
# Get system message for negative prompts and ensure it's a string
|
|
neg_system_message = self.profiles.get("IF_NegativePromptEngineer_V2", "")
|
|
if isinstance(neg_system_message, dict):
|
|
neg_system_message = json.dumps(neg_system_message)
|
|
|
|
# Generate negative prompt using cleaned response
|
|
neg_prompt = await send_request(
|
|
llm_provider=kwargs.get('llm_provider'),
|
|
base_ip=kwargs.get('base_ip'),
|
|
port=kwargs.get('port'),
|
|
images=images,
|
|
llm_model=kwargs.get('llm_model'),
|
|
system_message=neg_system_message,
|
|
user_message=f"Generate negative prompts for:\n{prompt}",
|
|
messages=[], # Fresh context for negative generation
|
|
seed=kwargs.get('seed', 0),
|
|
temperature=kwargs.get('temperature'),
|
|
max_tokens=kwargs.get('max_tokens'),
|
|
random=kwargs.get('random'),
|
|
top_k=kwargs.get('top_k'),
|
|
top_p=kwargs.get('top_p'),
|
|
repeat_penalty=kwargs.get('repeat_penalty'),
|
|
stop=kwargs.get('stop'),
|
|
keep_alive=kwargs.get('keep_alive'),
|
|
llm_api_key=kwargs.get('llm_api_key'),
|
|
)
|
|
|
|
# If the response is a dict, extract the actual text before calling clean_text
|
|
if isinstance(neg_prompt, dict):
|
|
extracted = ""
|
|
if "choices" in neg_prompt and neg_prompt["choices"]:
|
|
extracted = neg_prompt["choices"][0].get("message", {}).get("content", "")
|
|
elif "response" in neg_prompt:
|
|
extracted = neg_prompt["response"]
|
|
else:
|
|
# Fallback: just serialize the dict
|
|
extracted = json.dumps(neg_prompt)
|
|
neg_prompt = extracted
|
|
|
|
if neg_prompt:
|
|
return clean_text(neg_prompt)
|
|
else:
|
|
return kwargs.get('neg_content', '')
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error generating negative prompts: {str(e)}")
|
|
return ["Error generating negative prompt"]
|
|
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"IF_LLM": IFLLM
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"IF_LLM": "IF LLM🎨"
|
|
}
|
|
|
|
def get_omost_function():
|
|
"""Lazily import omost_function only when needed"""
|
|
try:
|
|
if "omost" not in sys.modules:
|
|
from .omost import omost_function
|
|
else:
|
|
omost_function = sys.modules["omost"].omost_function
|
|
return omost_function
|
|
except ImportError as e:
|
|
print(f"Error importing omost_function: {e}")
|
|
raise |