IFLLMNode.py

# 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