Here is our current code that is part of our integration for "FlorenceKit.py", "IFChatPromptNode.py" and "transformers_api.py" that we want to modify to achieve a seamless florence/comfy integration

#FlorenceKit.py
import torch
import logging
from typing import List, Union, Dict, Any, Optional
import numpy as np
from PIL import Image
import re
import supervision as sv
from transformers import set_seed
from .utils import hash_seed

logger = logging.getLogger(__name__)

class FlorenceModule:
    def __init__(self, model_manager):
        self.model_manager = model_manager
        self.colormap = sv.ColorLookup.INDEX  # Use supervision's color mapping

    async def run_florence(
        self,
        images: Union[torch.Tensor, List[torch.Tensor]],
        task_prompt: str,
        task: str,
        llm_model: str = "microsoft/Florence-2-large",
        precision: str = "fp16",
        fill_mask: bool = False,
        output_mask_select: str = "",
        keep_alive: bool = False,
        num_beams: int = 3,
        max_new_tokens: int = 1024,
        do_sample: bool = True,
        seed: Optional[int] = None,
        text_input: Optional[str] = None
    ) -> Dict[str, Any]:
        try:
            logger.info(f"Starting Florence task: {task}")
            logger.debug(f"Task prompt: {task_prompt}")
            logger.debug(f"Model: {llm_model}, Precision: {precision}")

            # Load Florence model
            florence_model = self.model_manager.load_model(
                model=llm_model,
                precision=precision,
                attention='sdpa',
                lora=None
            )
            if florence_model is None:
                raise ValueError(f"Failed to load model '{llm_model}'.")

            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            logger.info(f"Using device: {device}")
            model = florence_model['model'].to(device)
            processor = florence_model['processor']
            dtype = florence_model['dtype']
            logger.debug(f"Model dtype: {dtype}")

            if seed is not None:
                logger.info(f"Setting seed: {seed}")
                set_seed(hash_seed(seed))

            # Ensure images is a list of tensors
            if isinstance(images, torch.Tensor):
                if images.dim() == 4:
                    images = [img for img in images]
                elif images.dim() == 3:
                    images = [images]
                else:
                    raise ValueError(f"Invalid image tensor shape. Expected 3D or 4D, got {images.dim()}D")

            logger.debug(f"Number of images to process: {len(images)}")

            out_tensors = []
            out_mask_tensors = []
            out_results = []
            out_data = []

            for idx, img in enumerate(images):
                logger.debug(f"Processing image {idx+1}/{len(images)} with shape: {img.shape}")

                # Convert tensor to PIL Image
                image_pil = self.tensor_to_pil(img)
                W, H = image_pil.size
                logger.debug(f"PIL Image size: {W}x{H}")

                ## Determine the correct prompt based on the task (Same logic as before)
                if task_prompt in ["<REGION_TO_DESCRIPTION>", "<REGION_TO_CATEGORY>", "<REFERRING_EXPRESSION_SEGMENTATION>", "<CAPTION_TO_PHRASE_GROUNDING>"]:
                    prompt = f"{task_prompt} {text_input}".strip()
                elif task_prompt in ["<OD>", "<DENSE_REGION_CAPTION>", "<REGION_PROPOSAL>", "<OCR_WITH_REGION>"]:
                    prompt = text_input if text_input is not None else ""
                else:
                    prompt = task_prompt

                inputs = processor(text=prompt, images=image_pil, return_tensors="pt", do_rescale=False).to(dtype).to(device)
                logger.debug("Processed inputs with the processor.")
                logger.debug(f"Input keys: {inputs.keys()}")

                # Ensure input_ids is of type torch.long
                inputs["input_ids"] = inputs["input_ids"].long()

                # Generate outputs
                try:
                    with torch.no_grad():
                        generated_ids = model.generate(
                            input_ids=inputs["input_ids"],
                            pixel_values=inputs["pixel_values"],
                            max_new_tokens=max_new_tokens,
                            do_sample=do_sample,
                            num_beams=num_beams,
                        )
                    logger.debug("Generated IDs using the model.")
                    logger.debug(f"Generated IDs shape: {generated_ids.shape}")
                except Exception as e:
                    logger.error(f"Error during model.generate: {str(e)}", exc_info=True)
                    raise

                generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
                logger.debug(f"Generated text: {generated_text}")

                clean_results = self.clean_results(generated_text, task_prompt)
                logger.debug(f"Cleaned results: {clean_results}")

                out_results.append(clean_results)

                # Post-process the model's output
                try:
                    parsed_answer = processor.post_process_generation(generated_text, task=task_prompt, image_size=(W, H))
                    logger.debug(f"Parsed answer: {parsed_answer}")
                except Exception as e:
                    logger.error(f"Error during post-processing: {str(e)}", exc_info=True)
                    raise


                # Handle different tasks
                try:
                    if task in ["object_detection", "dense_region_caption", "caption_to_phrase_grounding", "region_proposal", "ocr_with_region"]:
                        # Access task-specific output using the task_prompt key
                        detections = sv.Detections.from_lmm(
                            lmm=sv.LMM.FLORENCE_2, result=parsed_answer[task_prompt], resolution_wh=(W, H)
                        )
                        annotated_image = self.annotate_with_boxes(image_pil, detections, output_mask_select)
                        tensor = self.pil_to_tensor(annotated_image).to(dtype)
                        out_tensors.append(tensor)
                        out_data.append(detections.xyxy.tolist())

                        if fill_mask:
                            mask_tensor = torch.zeros((1, H, W), dtype=dtype, device=device)
                            if detections.mask is not None:
                                for i, mask in enumerate(detections.mask):
                                    if str(i) in output_mask_select.split(","):
                                        mask_tensor[0] = torch.logical_or(mask_tensor[0], torch.from_numpy(mask.astype(np.uint8)))
                            out_mask_tensors.append(mask_tensor)
                        else:
                            out_mask_tensors.append(torch.zeros((1, H, W), dtype=dtype, device=device))

                        logger.debug(f"Handled detection-based task: {task}")

                    elif task == "referring_expression_segmentation" or task == "region_to_segmentation":
                        # Access task-specific output using the task_prompt key
                        detections = sv.Detections.from_lmm(
                            lmm=sv.LMM.FLORENCE_2, result=parsed_answer[task_prompt], resolution_wh=(W, H)
                        )
                        annotated_image = image_pil.copy()
                        mask_tensor = torch.zeros((1, H, W), dtype=dtype, device=device)

                        if detections.mask is not None:
                            for mask in detections.mask:
                                if fill_mask:
                                    # Assuming only 1 polygon per mask for now
                                    polygon = sv.mask_to_polygon(mask)[0].astype(np.int32)
                                    polygon_image = Image.fromarray(polygon).convert("RGBA")
                                    polygon_image = polygon_image.resize(image_pil.size)
                                    annotated_image = Image.alpha_composite(image_pil.convert("RGBA"), polygon_image)
                                mask_tensor[0] = torch.logical_or(mask_tensor[0], torch.from_numpy(mask.astype(np.uint8)))
                        
                        tensor = self.pil_to_tensor(annotated_image).to(dtype)
                        out_tensors.append(tensor)
                        out_mask_tensors.append(mask_tensor)
                        out_data.append(detections.xyxy.tolist())
                        logger.debug(f"Handled segmentation task: {task}")

                    elif task in ["caption", "detailed_caption", "more_detailed_caption", "ocr", "region_to_category", "region_to_description"]:
                        out_tensors.append(img.unsqueeze(0).to(dtype))
                        out_mask_tensors.append(torch.zeros((1, H, W), dtype=dtype, device=device))
                        # Access task-specific output using the task_prompt key
                        out_data.append(parsed_answer[task_prompt])
                        logger.debug(f"Handled text-based task: {task}")

                    else:
                        raise ValueError(f"Unsupported Florence task: {task}")

                except Exception as e:
                    logger.error(f"Error processing image {idx+1} for task: {task}: {str(e)}", exc_info=True)
                    # Use original image and empty mask if processing fails
                    out_tensors.append(img.unsqueeze(0).to(dtype))
                    out_mask_tensors.append(torch.zeros((1, H, W), dtype=dtype, device=device))
                    out_data.append({"error": str(e)})  # Store the error message

            # Prepare output tensors
            if out_tensors:
                out_tensors = torch.cat(out_tensors, dim=0)
                logger.debug(f"Concatenated out_tensors shape: {out_tensors.shape}")
            else:
                # Return the first image tensor if no output tensors were generated
                out_tensors = images[0].unsqueeze(0).to(dtype) if isinstance(images[0], torch.Tensor) else images[0]
                logger.debug("No output tensors generated; using original image.")

            if out_mask_tensors:
                out_mask_tensors = torch.cat(out_mask_tensors, dim=0)
                logger.debug(f"Concatenated out_mask_tensors shape: {out_mask_tensors.shape}")
            else:
                out_mask_tensors = torch.zeros_like(out_tensors)
                logger.debug("No mask tensors generated; using zero mask.")

            if not keep_alive:
                logger.info("Offloading model...")
                model.to(torch.device("cpu"))
                torch.cuda.empty_cache()

            logger.info("Florence task completed successfully.")
            
            result = {
                "out_tensors": out_tensors,
                "out_mask_tensors": out_mask_tensors,
                "out_results": out_results if out_results else ["No results generated"],
                "out_data": out_data if out_data else ["No data generated"]
            }
            
            logger.debug(f"Returning result: {result}")
            return result

        except Exception as e:
            logger.error(f"Error in run_florence method: {str(e)}", exc_info=True)
            # Return the original image if an error occurs
            return {
                "out_tensors": images[0].unsqueeze(0).to(dtype) if isinstance(images[0], torch.Tensor) else images[0],
                "out_mask_tensors": torch.zeros_like(images[0].unsqueeze(0)).to(dtype) if isinstance(images[0], torch.Tensor) else torch.zeros_like(images[0]),
                "out_results": [str(e)],
                "out_data": [str(e)]
            }

    def clean_results(self, results: str, task: str) -> str:
        # Clean up the model's output
        if task == '<OCR_WITH_REGION>':
            cleaned_string = re.sub(r'</?s>|<[^>]*>', '\n', results)
            return re.sub(r'\n+', '\n', cleaned_string)
        else:
            return results.replace(' { ', '').replace(' ', '')

    def tensor_to_pil(self, tensor: torch.Tensor) -> Image.Image:
        # Convert a PyTorch tensor to a PIL Image
        if tensor.dim() == 4:
            tensor = tensor.squeeze(0)  # Remove batch dimension if present
        
        if tensor.dim() == 3:
            if tensor.shape[0] in [1, 3]:
                # Channel-first format, convert to channel-last
                tensor = tensor.permute(1, 2, 0)
            elif tensor.shape[-1] not in [1, 3]:
                raise ValueError(f"Unexpected tensor shape: {tensor.shape}")
        elif tensor.dim() == 2:
            # Add channel dimension for grayscale
            tensor = tensor.unsqueeze(-1)
        else:
            raise ValueError(f"Unexpected tensor dimensions: {tensor.dim()}")

        # Ensure the tensor is on CPU and convert to numpy array
        numpy_image = tensor.cpu().numpy()

        # Normalize to 0-255 range and convert to uint8
        if numpy_image.dtype != np.uint8:
            numpy_image = (numpy_image * 255).clip(0, 255).astype(np.uint8)

        # Create PIL Image based on the number of channels
        if numpy_image.shape[-1] == 1:
            return Image.fromarray(numpy_image.squeeze(-1), mode='L')
        elif numpy_image.shape[-1] == 3:
            return Image.fromarray(numpy_image, mode='RGB')
        elif numpy_image.shape[-1] == 4:
            return Image.fromarray(numpy_image, mode='RGBA')
        else:
            raise ValueError(f"Unexpected number of channels: {numpy_image.shape[-1]}")

    def pil_to_tensor(self, image: Image.Image) -> torch.Tensor:
        # Convert a PIL Image to a PyTorch tensor
        tensor = torch.from_numpy(np.array(image).astype(np.float32) / 255.0)
        tensor = tensor.permute(2, 0, 1).unsqueeze(0)  # Add batch dimension
        return tensor

    def annotate_with_boxes(self, image: Image.Image, detections: sv.Detections, output_mask_select="") -> Image.Image:
        # Annotate the image with bounding boxes and labels
        annotated_image = image.copy()
        thickness = sv.calculate_optimal_line_thickness(resolution_wh=image.size)
        text_scale = sv.calculate_optimal_text_scale(resolution_wh=image.size)
        
        # Handle cases with invalid or empty output_mask_select
        mask_indexes = []
        if output_mask_select:
            try:
                mask_indexes = [int(n) for n in output_mask_select.split(",")]
            except ValueError:
                logger.warning(f"Invalid output_mask_select value: {output_mask_select}. Annotating all detections.")
                mask_indexes = list(range(len(detections.xyxy)))
        else:
            logger.debug("No output_mask_select provided. Annotating all detections.")
            mask_indexes = list(range(len(detections.xyxy)))

        bounding_box_annotator = sv.BoxAnnotator(color_lookup=self.colormap, thickness=thickness)
        label_annotator = sv.LabelAnnotator(
            color_lookup=self.colormap,
            text_scale=text_scale,
            text_thickness=thickness
        )

        # Iterate through mask indices, ensuring they're within bounds
        for i in mask_indexes:
            if i < len(detections.xyxy):
                annotated_image = bounding_box_annotator.annotate(annotated_image, detections[i])
                annotated_image = label_annotator.annotate(annotated_image, detections[i])
            else:
                logger.warning(f"Mask index {i} is out of bounds for the detections.")

        return annotated_image
cat << EOF

#IFChatPromptNode.py
import os
import sys
import json
import torch
import shutil
import base64
import platform
import importlib
import subprocess
import numpy as np
import folder_paths
from PIL import Image
import yaml
from io import BytesIO
from typing import List, Union, Dict, Any, Tuple, Optional
from .agent_tool import AgentTool
from .send_request import send_request
from .transformers_api import TransformersModelManager  # Updated import
from .utils import convert_image_to_base64, get_api_key, get_models, clean_text, validate_models, convert_image_for_comfy, tensor_to_pil, pil_to_tensor
import torchvision.transforms.functional as F
import asyncio
import tempfile
import threading
from aiohttp import web
from .graphRAG_module import GraphRAGapp
from .colpaliRAG_module import colpaliRAGapp
from .FlorenceKit import FlorenceModule 
# Set up logging
import logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
# Add the ComfyUI directory to the Python path
comfy_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
sys.path.insert(0, comfy_path)

ifchat_prompt_node = None

try:
    from server import PromptServer

    @PromptServer.instance.routes.post("/IF_ChatPrompt/get_llm_models")
    async def get_llm_models_endpoint(request):
        data = await request.json()
        llm_provider = data.get("llm_provider")
        engine = llm_provider
        base_ip = data.get("base_ip")
        port = data.get("port")
        models = get_models(engine, base_ip, port)
        return web.json_response(models)

    @PromptServer.instance.routes.post("/IF_ChatPrompt/get_embedding_models")
    async def get_embedding_models_endpoint(request):
        data = await request.json()
        embedding_provider = data.get("embedding_provider")
        engine = embedding_provider
        base_ip = data.get("base_ip")
        port = data.get("port")
        models = get_models(engine, base_ip, port)
        return web.json_response(models)

    @PromptServer.instance.routes.post("/IF_ChatPrompt/upload_file")
    async def upload_file_route(request):
        try:
            reader = await request.multipart()

            rag_folder_name = None
            file_content = None
            filename = None

            # Process all parts of the multipart request
            while True:
                part = await reader.next()
                if part is None:
                    break
                if part.name == "rag_root_dir":
                    rag_folder_name = await part.text()
                elif part.filename:
                    filename = part.filename
                    file_content = await part.read()

            if not filename or not file_content or not rag_folder_name:
                return web.json_response({"status": "error", "message": "Missing file, filename, or RAG folder name"})

            node = IFChatPrompt()
            input_dir = os.path.join(node.rag_dir, rag_folder_name, "input")

            if not os.path.exists(input_dir):
                os.makedirs(input_dir, exist_ok=True)

            file_path = os.path.join(input_dir, filename)

            with open(file_path, 'wb') as f:
                f.write(file_content)

            logger.info(f"File uploaded to: {file_path}")
            return web.json_response({"status": "success", "message": f"File uploaded to: {file_path}"})

        except Exception as e:
            logger.error(f"Error in upload_file_route: {str(e)}")
            return web.json_response({"status": "error", "message": f"Error uploading file: {str(e)}"})


    @PromptServer.instance.routes.post("/IF_ChatPrompt/setup_and_initialize")
    async def setup_and_initialize(request):
        global ifchat_prompt_node
        
        data = await request.json()
        folder_name = data.get('folder_name', 'rag_data')
        
        if ifchat_prompt_node is None:
            ifchat_prompt_node = IFChatPrompt()
        
        init_result = await ifchat_prompt_node.graphrag_app.setup_and_initialize_folder(folder_name, data)
        
        ifchat_prompt_node.rag_folder_name = folder_name
        ifchat_prompt_node.colpali_app.set_rag_root_dir(folder_name)   
        
        return web.json_response(init_result)

    @PromptServer.instance.routes.post("/IF_ChatPrompt/run_indexer")
    async def run_indexer_endpoint(request):
        try:
            data = await request.json()
            logger.debug(f"Received indexing request with data: {data}")

            global ifchat_prompt_node  # Access the global instance

            # Set the rag_root_dir in both modules using the global instance
            ifchat_prompt_node.graphrag_app.set_rag_root_dir(data.get('rag_folder_name'))
            ifchat_prompt_node.colpali_app.set_rag_root_dir(data.get('rag_folder_name'))

            query_type = data.get('mode_type')
            logger.debug(f"Query type: {query_type}")

            logger.debug(f"Starting indexing process for query type: {query_type}")

            # Initialize the colpali_model before calling insert, using the global instance
            if query_type == 'colpali':
                _ = ifchat_prompt_node.colpali_app.get_colpali_model()  # This will load or retrieve the cached model
                result = await ifchat_prompt_node.colpali_app.insert()
            else:
                result = await ifchat_prompt_node.graphrag_app.insert()

            logger.debug(f"Indexing process completed with result: {result}")

            if result:
                return web.json_response({"status": "success", "message": f"Indexing complete for {query_type}"})
            else:
                return web.json_response({"status": "error", "message": "Indexing failed. Check server logs."}, status=500)

        except Exception as e:
            logger.error(f"Error in run_indexer_endpoint: {str(e)}")
            return web.json_response({"status": "error", "message": f"Error during indexing: {str(e)}"}, status=500)
        
    @PromptServer.instance.routes.post("/IF_ChatPrompt/process_chat")
    async def process_chat_endpoint(request):
        data = await request.json()

        global ifchat_prompt_node 
        result = await ifchat_prompt_node.process_chat(**data)  

        return web.json_response(result)
        
except AttributeError:
    print("PromptServer.instance not available. Skipping route decoration for IF_ChatPrompt.")

class IFChatPrompt:

    def __init__(self):
        self.base_ip = "localhost"
        self.port = "11434"
        self.llm_provider = "ollama"
        self.embedding_provider = "sentence_transformers"
        self.llm_model = ""
        self.embedding_model = ""
        self.assistant = "Cortana"
        self.random = False
        self.comfy_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
        self.success_image = os.path.join(self.comfy_dir, "input", "IF_AI", "success.png")
        self.presets_dir = os.path.join(self.comfy_dir, "input", "IF_AI", "presets")
        self.rag_dir = os.path.join(self.comfy_dir, "input", "IF_AI", "rag")
        self.stop_file = os.path.join(self.presets_dir, "stop_strings.json")
        self.assistants_file = os.path.join(self.presets_dir, "assistants.json")
        self.neg_prompts_file = os.path.join(self.presets_dir, "neg_prompts.json")
        self.embellish_prompts_file = os.path.join(self.presets_dir, "embellishments.json")
        self.style_prompts_file = os.path.join(self.presets_dir, "style_prompts.json")
        self.tasks_file = os.path.join(self.presets_dir, "florence_prompts.json")
        self.agents_dir = os.path.join(self.presets_dir, "agents")
        self.agent_tools = self.load_agent_tools()
        self.stop_strings = self.load_presets(self.stop_file)
        self.assistants = self.load_presets(self.assistants_file)
        self.neg_prompts = self.load_presets(self.neg_prompts_file)
        self.embellish_prompts = self.load_presets(self.embellish_prompts_file)
        self.style_prompts = self.load_presets(self.style_prompts_file)
        self.florence_prompts = self.load_presets(self.tasks_file)
        self.keep_alive = False
        self.seed = 94687328150
        self.messages = []
        self.history_steps = 10
        self.external_llm_api_key = ""
        self.external_embedding_api_key = ""
        self.tool_input = ""
        self.prime_directives = None
        self.rag_folder_name = "rag_data"
        self.graphrag_app = GraphRAGapp()
        self.colpali_app = colpaliRAGapp()
        self.fix_json = True
        self.cached_colpali_model = None

        self.transformers_manager = TransformersModelManager(models_dir=folder_paths.models_dir)

        self.florence_app = FlorenceModule(model_manager=self.transformers_manager)

        self.florence_models = {}
        self.query_type = "global"  
        self.enable_RAG = False
        self.clear_history = False
        self.mode = False
        self.tool = "None"
        self.preset = "Default"
        self.precision = "fp16"
        self.task = None  # Added attribute to fix 'no attribute task' warning

    def load_presets(self, file_path):
        with open(file_path, 'r') as f:
            presets = json.load(f)
        return presets

    def load_agent_tools(self):
        agent_tools = {}
        for filename in os.listdir(self.agents_dir):
            if filename.endswith('.json'):
                full_path = os.path.join(self.agents_dir, filename)
                with open(full_path, 'r') as f:
                    try:
                        data = json.load(f)
                        if 'output_type' not in data:
                            data['output_type'] = None
                        agent_tool = AgentTool(**data)
                        agent_tool.load()
                        if agent_tool._class_instance is not None:
                            if agent_tool.python_function:
                                agent_tools[agent_tool.name] = agent_tool
                            else:
                                print(f"Warning: Agent tool {agent_tool.name} in {filename} does not have a python_function defined.")
                        else:
                            print(f"Failed to create class instance for {filename}")
                    except json.JSONDecodeError:
                        print(f"Error: Invalid JSON in {filename}")
                    except Exception as e:
                        print(f"Error loading {filename}: {str(e)}")
        return agent_tools
    
    async def process_chat(self, prompt, llm_provider, llm_model, base_ip, port, assistant, neg_prompt, embellish_prompt, style_prompt, 
                     embedding_model, embedding_provider, external_llm_api_key="", external_embedding_api_key="",
                     temperature=0.7, max_tokens=2048, seed=0, random=False, history_steps=10, keep_alive=False, top_k=40, top_p=0.2,
                     repeat_penalty=1.1, stop_string=None, images=None, mode=True, clear_history=False, text_cleanup=True,
                     tool=None, tool_input=None, prime_directives=None, enable_RAG=False, query_type='global', 
                     preset="Default", rag_folder_name=None, fix_json=True, task=None, fill_mask=False, output_mask_select="", precision = "fp16"):
    
        # Handle API key
        if external_llm_api_key != "":
            llm_api_key = external_llm_api_key
        else:
            llm_api_key = get_api_key(f"{llm_provider.upper()}_API_KEY", llm_provider)

        if external_embedding_api_key != "":
            embedding_api_key = external_embedding_api_key
        else:
            embedding_api_key = get_api_key(f"{embedding_provider.upper()}_API_KEY", embedding_provider)

        if prime_directives is not None:
            system_message_str = prime_directives
        else:
            system_message = self.assistants.get(assistant, "")
            system_message_str = json.dumps(system_message)

        base64_images = []
        if images is not None:
            # Ensure images is always a list
            if not isinstance(images, list):
                images = [images]
            
            for img in images:
                if isinstance(img, torch.Tensor):
                    if img.dim() == 4:  # [B, C, H, W] or [B, H, W, C]
                        if img.shape[1] in [1, 3, 4]:  # Channels likely in the second dimension
                            img = img.squeeze(0)  # [C, H, W]
                        else:  # Channels likely in the last dimension
                            img = img.squeeze(0).permute(2, 0, 1)  # [C, H, W]
                    elif img.dim() == 3:
                        if img.shape[0] not in [1, 3, 4]:  # If channels not in first dimension
                            img = img.permute(2, 0, 1)  # [C, H, W]
                    else:
                        raise ValueError(f"Unexpected image tensor shape: {img.shape}")
                    
                    # Convert PIL Image to Tensor and squeeze to [C, H, W]
                    img_tensor = pil_to_tensor(F.to_pil_image(img)).squeeze(0)
                elif isinstance(img, Image.Image):
                    img_tensor = pil_to_tensor(img).squeeze(0)
                else:
                    raise ValueError(f"Unsupported image type: {type(img)}")
                
                base64_images.append(convert_image_to_base64(img_tensor))

            
        print(f"Number of images being sent with request: {len(base64_images)}")

        # Validate LLM model
        validate_models(llm_model, llm_provider, "LLM", base_ip, port)

        # Validate embedding model
        validate_models(embedding_model, embedding_provider, "embedding", base_ip, port)

        # Handle history
        if clear_history:
            self.messages = []
        elif history_steps > 0:
            self.messages = self.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

        # Handle tools
        try:
            if tool and tool != "None":
                selected_tool = self.agent_tools.get(tool)
                if not selected_tool:
                    raise ValueError(f"Invalid agent tool selected: {tool}")

                # Prepare tool execution message
                tool_message = f"Execute the {tool} tool with the following input: {prompt}"
                system_prompt = json.dumps(selected_tool.system_prompt)

                # Send request to LLM for tool execution
                generated_text = await self.graphrag_app.unified_model_if_cache(
                    tool_message, system_prompt=system_prompt,
                    llm_provider=llm_provider,
                    llm_model=llm_model,
                    base_ip=base_ip,
                    port=port,
                    temperature=temperature,
                    max_tokens=max_tokens,
                    seed=seed,
                    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
                )

                # Parse the generated text for function calls
                function_call = None
                try:
                    response_data = json.loads(generated_text)
                    if 'function_call' in response_data:
                        function_call = response_data['function_call']
                        generated_text = response_data['content']
                except json.JSONDecodeError:
                    pass

                # Execute the tool with the LLM's response
                tool_args = {
                    "input": prompt,
                    "llm_response": generated_text,
                    "function_call": function_call,
                    "omni_input": tool_input,
                    "name": selected_tool.name,
                    "description": selected_tool.description,
                    "system_prompt": selected_tool.system_prompt
                }
                tool_result = selected_tool.execute(tool_args)

                # Update messages
                self.messages.append({"role": "user", "content": prompt})
                self.messages.append({
                    "role": "assistant",
                    "content": json.dumps(tool_result) if isinstance(tool_result, dict) else str(tool_result)
                })

                # Process the tool output
                if isinstance(tool_result, dict):
                    if "error" in tool_result:
                        generated_text = f"Error in {tool}: {tool_result['error']}"
                        tool_output = None
                    elif selected_tool.output_type and selected_tool.output_type in tool_result:
                        tool_output = tool_result[selected_tool.output_type]
                        generated_text = f"Agent {tool} executed successfully. Output generated."
                    else:
                        tool_output = tool_result
                        generated_text = str(tool_output)
                else:
                    tool_output = tool_result
                    generated_text = str(tool_output)

                return {
                    "Question": prompt,
                    "Response": generated_text,
                    "Negative": self.neg_prompts.get(neg_prompt, ""),
                    "Tool_Output": tool_output,
                    "Retrieved_Image": None  # No image retrieved in tool execution
                }
            else:
                # Check if the task is a Florence task
                florence_tasks = list(self.florence_prompts.keys())
                if florence_tasks and task.lower() != 'none':
                    return await self.process_florence_task(
                        prompt, llm_provider, llm_model, base_ip, port, assistant, neg_prompt, 
                        embellish_prompt, style_prompt, embedding_model, embedding_provider, 
                        external_llm_api_key, external_embedding_api_key, temperature, max_tokens, 
                        seed, random, history_steps, keep_alive, top_k, top_p, repeat_penalty, 
                        stop_string, images, mode, clear_history, text_cleanup, tool, tool_input, 
                        prime_directives, enable_RAG, query_type, preset, rag_folder_name, fix_json, 
                        task, fill_mask, output_mask_select, precision
                    )
                else:
                    response = await self.generate_response(
                        enable_RAG, query_type, prompt, preset, llm_provider, base_ip, port,
                        base64_images, llm_model, system_message_str, temperature, max_tokens,
                        random, top_k, top_p, repeat_penalty, stop, keep_alive, llm_api_key
                    )
                    
                    generated_text = response["text"]
                    retrieved_image = response.get("image")
                    graph_data = response.get("graph")

                    # Update messages
                    self.messages.append({"role": "user", "content": prompt})
                    self.messages.append({"role": "assistant", "content": generated_text})

                    # Clean up the generated text
                    if text_cleanup and isinstance(generated_text, str):
                        generated_text = clean_text(generated_text)
                    elif not isinstance(generated_text, str):
                        logger.warning(f"Generated text is not a string. Type: {type(generated_text)}")
                        generated_text = str(generated_text)  # Convert to string if possible

                    description = str(generated_text).strip()
                    print(f"Generated text: {description}")

                    if mode:
                        embellish_content = self.embellish_prompts.get(embellish_prompt, "").strip()
                        style_content = self.style_prompts.get(style_prompt, "").strip()
                        combined_prompt = f"{embellish_content} {description} {style_content}".strip()
                    else:
                        combined_prompt = description
                    
                    if retrieved_image is None:
                        retrieved_image = convert_image_for_comfy(self.success_image)
                    # Always return the same structure, but with different content based on the response type
                    return {
                        "Question": prompt,
                        "Response": combined_prompt,
                        "Negative": self.neg_prompts.get(neg_prompt, ""),
                        "Tool_Output": graph_data,
                        "Retrieved_Image": retrieved_image,
                        "Mask": None
                    }

        except Exception as e:
            logger.error(f"Exception occurred in process_chat: {str(e)}", exc_info=True)
            return {
                "Question": prompt,
                "Response": f"Error: {str(e)}",
                "Negative": "",
                "Tool_Output": None,
                "Retrieved_Image": None,
                "Mask": None
            }
 
    
    async def generate_response(self, enable_RAG, query_type, prompt, preset, llm_provider, base_ip, port,
                                base64_images, llm_model, system_message_str, temperature, max_tokens,
                                random, top_k, top_p, repeat_penalty, stop, keep_alive, llm_api_key):
        response_strategies = {
            "graphrag": self.graphrag_app.query,
            "colpali": self.colpali_app.query,
            "normal": lambda: send_request(
                llm_provider, base_ip, port, base64_images, llm_model, system_message_str, prompt, 
                self.messages, self.seed, temperature, max_tokens, random, top_k, top_p, 
                repeat_penalty, stop, keep_alive, llm_api_key
            )
        }

        if enable_RAG:
            if query_type == "colpali":
                strategy = "colpali"
            else:  # For "global", "local", and "naive" query types
                strategy = "graphrag"
        else:
            strategy = "normal"

        try:
            if strategy == "colpali":
                # Ensure the model is loaded before querying
                if self.cached_colpali_model is None:
                    self.cached_colpali_model = self.colpali_app.get_colpali_model()
                generated_text, retrieved_image = await response_strategies[strategy](prompt=prompt, system_message_str=system_message_str)
                return {"text": generated_text, "image": retrieved_image, "graph": None}
            elif strategy == "graphrag":
                result = await response_strategies[strategy](prompt=prompt, query_type=query_type, preset=preset)
                # GraphRAG returns a dict with 'result' and 'graphml_path'
                return {"text": result['result'], "image": None, "graph": result['graphml_path']}
            else:
                normal_response = response_strategies[strategy]()
                return {"text": normal_response, "image": None, "graph": None}
        except Exception as e:
            logger.error(f"Error in {strategy} query: {str(e)}")
            return {"text": f"Error in {strategy} query: {str(e)}", "image": None, "graph": None}

    async def process_florence_task(self, prompt, llm_provider, llm_model, base_ip, port, assistant, neg_prompt, embellish_prompt, style_prompt, embedding_model, embedding_provider, external_llm_api_key, external_embedding_api_key, temperature, max_tokens, seed, random, history_steps, keep_alive, top_k, top_p, repeat_penalty, stop_string, images, mode, clear_history, text_cleanup, tool, tool_input, prime_directives, enable_RAG, query_type, preset, rag_folder_name, fix_json, task, fill_mask, output_mask_select, precision):
        try:
            logger.debug(f"Processing Florence task: {task}")
            logger.debug(f"User prompt: {prompt}")

            # Retrieve task prompt (Florence tag) from florence_prompts.json
            task_prompt = self.florence_prompts.get(task, "").strip()
            logger.debug(f"Task prompt (Florence tag): {task_prompt}")
            
            # 'prompt' variable now holds the user's text input
            text_input = prompt  # Rename for clarity

            # Ensure images are in a list of tensors
            if isinstance(images, torch.Tensor):
                if images.dim() == 4:
                    images = [img for img in images]
                elif images.dim() == 3:
                    images = [images]
                else:
                    raise ValueError(f"Invalid image tensor shape. Expected 3D or 4D, got {images.dim()}D")
            
            logger.debug(f"Number of images: {len(images)}")
            logger.debug(f"Image shape: {images[0].shape}")

            # Call the FlorenceModule's run_florence method
            encode_result = await self.florence_app.run_florence(
                images=images,
                task_prompt=task_prompt,
                task=task,
                llm_model=llm_model,
                precision=precision,
                fill_mask=fill_mask,
                output_mask_select=output_mask_select,
                keep_alive=keep_alive,
                num_beams=top_k,
                max_new_tokens=max_tokens,
                do_sample=(not random),
                seed=seed,
                text_input=text_input
            )

            if encode_result is None:
                logger.error("Encoding failed. No result returned from run_florence.")
                return {
                    "Question": text_input,
                    "Response": "Encoding failed. No result returned from run_florence.",
                    "Negative": self.neg_prompts.get(neg_prompt, ""),
                    "Tool_Output": "Encoding failed",
                    "Retrieved_Image": images[0] if images else None,
                    "Mask": torch.zeros_like(images[0]) if images else None
                }

            logger.debug(f"Encode result keys: {encode_result.keys()}")

            out_tensors = encode_result.get("out_tensors")
            out_mask_tensors = encode_result.get("out_mask_tensors")
            out_results = encode_result.get("out_results", ["No results"])
            out_data = encode_result.get("out_data", ["No data"])

            logger.debug(f"Out tensors shape: {out_tensors.shape if out_tensors is not None else None}")
            logger.debug(f"Out mask tensors shape: {out_mask_tensors.shape if out_mask_tensors is not None else None}")
            logger.debug(f"Out results: {out_results}")
            logger.debug(f"Out data: {out_data}")

            return {
                "Question": text_input,
                "Response": out_results[0] if out_results else "No response",
                "Negative": self.neg_prompts.get(neg_prompt, ""),
                "Tool_Output": out_data[0] if out_data else "No data",
                # Check if images is not None and has at least one element
                "Retrieved_Image": out_tensors if out_tensors is not None and out_tensors.numel() > 0 else images[0] if images is not None and len(images) > 0 else None,
                "Mask": out_mask_tensors if out_mask_tensors is not None else torch.zeros_like(images[0]) if images is not None and len(images) > 0 else None
            }

        except Exception as e:
            logger.error(f"Error processing Florence task '{task}': {str(e)}", exc_info=True)
            return {
                "Question": text_input,
                "Response": f"Error processing Florence task: {str(e)}",
                "Negative": self.neg_prompts.get(neg_prompt, ""),
                "Tool_Output": f"Error: {str(e)}",
                # Check if images is not None and has at least one element
                "Retrieved_Image": images[0] if images is not None and len(images) > 0 else None,
                # Check if images is not None and has at least one element
                "Mask": torch.zeros_like(images[0]) if images is not None and len(images) > 0 else None
        }

       
    def process_chat_wrapper(self, *args, **kwargs):
        try:
            loop = asyncio.get_event_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
        
        return loop.run_until_complete(self.process_chat(*args, **kwargs))
    
    @classmethod
    def INPUT_TYPES(cls):
        node = cls()
        return {
            "required": {
                "prompt": ("STRING", {"multiline": True, "default": "", "tooltip": "The main text input for the chat or query."}),
                "llm_provider": (["llamacpp", "ollama", "kobold", "lmstudio", "textgen", "groq", "gemini", "openai", "anthropic", "mistral", "transformers"], {"default": node.llm_provider, "tooltip": "The provider of the language model to be used."}),
                "llm_model": ((), {"tooltip": "The specific language model to be used for processing."}),
                "base_ip": ("STRING", {"default": node.base_ip, "tooltip": "IP address of the LLM server."}),
                "port": ("STRING", {"default": node.port, "tooltip": "Port number for the LLM server connection."}),               
                "assistant": ([name for name in node.assistants.keys()], {"default": node.assistant, "tooltip": "The pre-defined assistant personality to use for responses."}),
            },
            "optional": {
                "images": ("IMAGE", {"list": True, "tooltip": "Input image(s) for visual processing or context."}),
                "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0, "step": 0.1, "tooltip": "Controls randomness in output generation. Higher values increase creativity but may reduce coherence."}),
                "max_tokens": ("INT", {"default": 2048, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Maximum number of tokens to generate in the response."}),
                "top_k": ("INT", {"default": 40, "min": 0, "max": 100, "tooltip": "Limits the next token selection to the K most likely tokens."}),
                "top_p": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.1, "tooltip": "Cumulative probability cutoff for token selection."}),
                "repeat_penalty": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.1, "tooltip": "Penalizes repetition in generated text."}),
                "stop_string": ([name for name in node.stop_strings.keys()], {"tooltip": "Specifies a string at which text generation should stop."}),
                "seed": ("INT", {"default": 94687328150, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Random seed for reproducible outputs."}),
                "random": ("BOOLEAN", {"default": False, "label_on": "Seed", "label_off": "Temperature", "tooltip": "Toggles between using a fixed seed or temperature-based randomness."}),
                "history_steps": ("INT", {"default": 10, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Number of previous conversation turns to consider for context."}),
                "clear_history": ("BOOLEAN", {"default": False, "label_on": "Clear History", "label_off": "Keep History", "tooltip": "Option to clear or retain conversation history."}),
                "mode": ("BOOLEAN", {"default": False, "label_on": "Using SD Mode", "label_off": "Using Chat Mode", "tooltip": "Switches between Stable Diffusion prompt generation and standard chat mode."}),
                "keep_alive": ("BOOLEAN", {"default": False, "label_on": "Keeps Model on Memory", "label_off": "Unloads Model from Memory", "tooltip": "Determines whether to keep the model loaded in memory between calls."}),
                "text_cleanup": ("BOOLEAN", {"default": True, "label_on": "Clean Response", "label_off": "Raw Text", "tooltip": "Applies text cleaning to the generated output."}),
                "precision": (['fp16','bf16','fp32'],{"default": 'bf16', "tooltip": "Select preccision on Transformer models."}),
                "fill_mask": ("BOOLEAN", {"default": False, "label_on": "Fill Mask", "label_off": "No Fill", "tooltip": "Option to fill masks for Florence tasks."}),
                "output_mask_select": ("STRING", {"default": ""}),
                "task": ([name for name in node.florence_prompts.keys()], {"default": "None", "tooltip": "Select a Florence task."}),
                "embellish_prompt": ([name for name in node.embellish_prompts.keys()], {"tooltip": "Adds pre-defined embellishments to the prompt."}),
                "style_prompt": ([name for name in node.style_prompts.keys()], {"tooltip": "Applies a pre-defined style to the prompt."}),
                "neg_prompt": ([name for name in node.neg_prompts.keys()], {"tooltip": "Adds a negative prompt to guide what should be avoided in generation."}),              
                "tool": (["None"] + [name for name in node.agent_tools.keys()], {"default": "None", "tooltip": "Selects a specific tool or agent for task execution."}),
                "embedding_provider": (["llamacpp", "ollama", "kobold", "lmstudio", "textgen", "groq", "gemini", "openai", "anthropic", "mistral", "sentence_transformers"], {"default": node.embedding_provider, "tooltip": "Provider for text embedding model."}),
                "embedding_model": ((), {"tooltip": "Specific embedding model to use."}),
                "tool_input": ("OMNI", {"default": None, "tooltip": "Additional input for the selected tool."}),
                "prime_directives": ("STRING", {"forceInput": True, "tooltip": "System message or prime directive for the AI assistant."}),
                "external_llm_api_key": ("STRING", {"forceInput": True, "tooltip": "API key for external LLM services."}),
                "external_embedding_api_key": ("STRING", {"forceInput": True, "tooltip": "API key for external embedding services."}),
                "enable_RAG": ("BOOLEAN", {"default": False, "label_on": "RAG is Enabled", "label_off": "RAG is Disabled", "tooltip": "Enables Retrieval-Augmented Generation for enhanced context."}),
                "query_type": (["global", "local", "naive", "colpali"], {"default": "global", "tooltip": "Selects the type of query strategy for RAG."}),
                "preset": (["Default", "Detailed", "Quick", "Bullet", "Comprehensive", "High-Level", "Focused"], {"default": "Default"}),
                "fix_json": ("BOOLEAN", {"default": True, "label_on": "Fix Response", "label_off": "No Fix"}),
            },
            "hidden": {
                "model": ("STRING", {"default": ""}),
                "rag_root_dir": ("STRING", {"default": "rag_data"})
            }
        }

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        node = cls()
        
        llm_provider = kwargs.get('llm_provider', node.llm_provider)
        embedding_provider = kwargs.get('embedding_provider', node.embedding_provider)
        base_ip = kwargs.get('base_ip', node.base_ip)
        port = kwargs.get('port', node.port)
        query_type = kwargs.get('query_type', node.query_type)
        external_llm_api_key = kwargs.get('external_llm_api_key', node.external_llm_api_key)
        external_llm_api_key = kwargs.get('external_llm_api_key', node.external_llm_api_key)
        task = kwargs.get('task', node.task)
        
        llm_provider_changed = llm_provider != node.llm_provider
        embedding_provider_changed = embedding_provider != node.embedding_provider
        llm_model_changed = node.llm_model != get_models(llm_provider, base_ip, port)
        embedding_model_changed = node.embedding_model != get_models(embedding_provider, base_ip, port)
        query_type_changed = kwargs.get('query_type') != node.query_type
        external_llm_api_key = kwargs.get('external_llm_api_key') != node.external_llm_api_key
        external_embedding_api_key = kwargs.get('external_embedding_api_key') != node.external_embedding_api_key
        task_changed = task != node.task
        
        if llm_provider_changed or embedding_provider_changed or llm_model_changed or embedding_model_changed or query_type_changed or task_changed:
            node.llm_provider = llm_provider
            node.embedding_provider = embedding_provider
            node.base_ip = base_ip
            node.port = port
            node.llm_model = get_models(llm_provider, base_ip, port)
            node.embedding_model = get_models(embedding_provider, base_ip, port)
            node.query_type = query_type
            node.external_llm_api_key = external_llm_api_key
            node.external_embedding_api_key = external_embedding_api_key
            # Update other attributes
            for attr in ['seed', 'random', 'history_steps', 'clear_history', 'mode', 
                         'keep_alive', 'tool', 'enable_RAG', 'preset']:
                setattr(node, attr, kwargs.get(attr, getattr(node, attr)))
            
            return True

        return False

    RETURN_TYPES = ("STRING", "STRING", "STRING", "OMNI", "IMAGE", "MASK")
    RETURN_NAMES = ("Question", "Response", "Negative", "Tool_Output", "Retrieved_Image", "Mask")
    OUTPUT_TOOLTIPS = (
        "The original input question or prompt.",
        "The generated response from the language model.",
        "The negative prompt used (if applicable) for guiding image generation.",
        "Output from the selected tool, which can be code or any other data type.",
        "An image retrieved by the RAG system, if applicable.",
        "Mask image generated by Florence tasks."
    )
    FUNCTION = "process_chat_wrapper"
    OUTPUT_NODE = True
    CATEGORY = "ImpactFrames💥🎞️"
    DESCRIPTION = "ComfyUI, Support API and Local LLM providers and RAG capabilities. Processes text prompts, handles image inputs, and integrates with different language models and indexing strategies."
cat << EOF

#transformers_api.py
from transformers import AutoModelForCausalLM, AutoTokenizer, Qwen2VLForConditionalGeneration, AutoProcessor, set_seed
from sentence_transformers import SentenceTransformer
from qwen_vl_utils import process_vision_info
from typing import List, Union, Optional, Dict, Any
from PIL import Image
from io import BytesIO
import base64
import torch
import logging
import os
import folder_paths
from unittest.mock import patch
from huggingface_hub.utils._errors import RepositoryNotFoundError
from pathlib import Path
from unittest.mock import patch
from transformers.dynamic_module_utils import get_imports


logger = logging.getLogger(__name__)

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)


class TransformersModelManager:
    def __init__(self, models_dir: str):
        self.models_dir = models_dir
        self.models = {}
        self.processors = {}
        self.offload_device = torch.device("cpu")  # Adjust as needed

    def load_model(self, model: str, precision: str, attention: str, lora: str = None) -> Dict[str, Any]:
        if model not in self.models:
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision]
            model_path = self.download_model_if_not_exists(model)
            if model_path is None:
                logger.error(f"Model path for '{model}' could not be determined.")
                return None

            if 'florence' in model.lower():
                logger.info(f"Using {attention} for attention")
                with patch("transformers.dynamic_module_utils.get_imports", self.fixed_get_imports):
                    loaded_model = AutoModelForCausalLM.from_pretrained(
                        model_path,
                        attn_implementation=attention,
                        device_map=device,
                        torch_dtype=dtype,
                        trust_remote_code=True
                    )
                processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
            elif 'qwen2-vl' in model.lower():
                loaded_model = Qwen2VLForConditionalGeneration.from_pretrained(
                    model_path, 
                    attn_implementation=attention, 
                    device_map=device, 
                    torch_dtype=dtype, 
                    trust_remote_code=True
                )
                processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
            else:
                loaded_model = AutoModelForCausalLM.from_pretrained(
                    model_path, 
                    attn_implementation=attention, 
                    device_map=device, 
                    torch_dtype=dtype, 
                    trust_remote_code=True
                )
                processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)

            # Store the loaded model and processor
            self.models[model] = loaded_model
            self.processors[model] = processor

            logger.info(f"Model '{model}' loaded successfully.")

        return {
            'model': self.models.get(model),
            'processor': self.processors.get(model),
            'dtype': self.models.get(model).dtype
        }

    def download_model_if_not_exists(self, model_name):
        from huggingface_hub import snapshot_download

        model_dir = model_name.rsplit('/', 1)[-1]
        model_path = os.path.join(self.models_dir, "LLM", model_dir)
        if not os.path.exists(model_path):
            logger.info(f"Downloading model '{model_name}' to: {model_path}")
            try:
                snapshot_download(
                    repo_id=model_name,
                    local_dir=model_path,
                    local_dir_use_symlinks=False,
                    token=os.getenv("HUGGINGFACE_TOKEN")  # Ensure your token is set in environment variables
                )
                logger.info(f"Model '{model_name}' downloaded successfully.")
            except RepositoryNotFoundError:
                logger.error(f"Repository '{model_name}' not found. Please check the model name and your access permissions.")
                return None
            except Exception as e:
                logger.error(f"An error occurred while downloading the model '{model_name}': {e}")
                return None
        else:
            logger.info(f"Model '{model_name}' already exists at: {model_path}")
        return model_path
        
    def send_transformers_request(self, model_name, system_message, user_message, messages, max_tokens, base64_images):
        """Sends a request to a Hugging Face Transformers model."""
        try:
            model_data = self.load_model(model_name)
            model = model_data['model']
            processor = model_data['processor']

            # Convert base64 images to PIL Images
            pil_images = [Image.open(BytesIO(base64.b64decode(img_str))) for img_str in base64_images]

            # Construct standardized messages
            formatted_messages = self.construct_messages(system_message, user_message, messages, pil_images)

            if 'qwen2-vl' in model_name.lower():
                # Process input for Qwen2-VL
                text = processor.apply_chat_template(formatted_messages, tokenize=False, add_generation_prompt=True)
                image_inputs, video_inputs = process_vision_info(formatted_messages)
                inputs = processor(
                    text=[text],
                    images=image_inputs,
                    videos=video_inputs,
                    padding=True,
                    return_tensors="pt"
                ).to(model.device)
                
                # Generate output
                generated_ids = model.generate(**inputs, max_new_tokens=max_tokens)
                generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
                output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)
                generated_text = output_text[0]
                
                result = generated_text
            else:
                # Handle other transformers models
                inputs = processor(formatted_messages, return_tensors="pt", padding=True).to(model.device)
                
                with torch.no_grad():
                    outputs = model.generate(**inputs, max_new_tokens=max_tokens)
                
                generated_text = processor.batch_decode(outputs, skip_special_tokens=True)[0]
                
                result = generated_text

            return result
        except Exception as e:
            print(f"Error in Transformers API request: {e}")
            return str(e)

    def construct_messages(self, system_message, user_message, messages, pil_images):
        """Constructs a standardized message format for transformer models."""
        formatted_messages = []
        if system_message:
            formatted_messages.append({"role": "system", "content": system_message})
        
        for msg in messages:
            formatted_messages.append({"role": msg['role'], "content": msg['content']})
        
        if user_message:
            formatted_messages.append({
                "role": "user",
                "content": [
                    {"type": "text", "text": user_message},
                    *[{"type": "image", "image": img} for img in pil_images]
                ]
            })
        
        return formatted_messages
    
    def create_sentence_transformers_embedding(self, model: str, input: Union[str, List[str]]) -> List[float]:
        st_model = SentenceTransformer(model, device="cpu")
        embeddings = st_model.encode(input, normalize_embeddings=True)
        return embeddings.tolist()   

    def unload_model(self, model: str):
        if model in self.models:
            del self.models[model]
            del self.processors[model]
            torch.cuda.empty_cache()
            logger.info(f"Unloaded Florence2 model: {model}")
        else:
            logger.warning(f"Attempted to unload a non-loaded model: {model}")

    def get_device(self):
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    @classmethod
    def fixed_get_imports(cls, filename: Union[str, os.PathLike], *args, **kwargs) -> List[str]:
        """Remove 'flash_attn' from imports if present."""
        try:
            if not str(filename).endswith("modeling_florence2.py"):
                return get_imports(filename)
            imports = get_imports(filename)
            if "flash_attn" in imports:
                imports.remove("flash_attn")
            return imports
        except Exception as e:
            print(f"No flash_attn import to remove: {e}")
            return get_imports(filename)

# Initialize a global manager instance
_transformers_manager = TransformersModelManager(models_dir=folder_paths.models_dir)

def send_transformers_request(model_name, system_message, user_message, messages, max_tokens, base64_images):
    return _transformers_manager.send_transformers_request(model_name, system_message, user_message, messages, max_tokens, base64_images)

def create_sentence_transformers_embedding(model, input):
    return _transformers_manager.create_sentence_transformers_embedding(model, input)

def load_model(model, precision, attention, lora=None):
    return _transformers_manager.load_model(model, precision, attention, lora)
cat << EOF


We want to create an script that handles the florence functionality intereating with the IFChat node and that harvest the supervision dependencies capabilities without overcomplicating things the "lmm.py" is part of the supervision dependencies library package installed on the environment of comfyui

#llm.py
import re
from enum import Enum
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np

from supervision.detection.utils import polygon_to_mask, polygon_to_xyxy


class LMM(Enum):
    PALIGEMMA = "paligemma"
    FLORENCE_2 = "florence_2"


RESULT_TYPES: Dict[LMM, type] = {LMM.PALIGEMMA: str, LMM.FLORENCE_2: dict}

REQUIRED_ARGUMENTS: Dict[LMM, List[str]] = {
    LMM.PALIGEMMA: ["resolution_wh"],
    LMM.FLORENCE_2: ["resolution_wh"],
}

ALLOWED_ARGUMENTS: Dict[LMM, List[str]] = {
    LMM.PALIGEMMA: ["resolution_wh", "classes"],
    LMM.FLORENCE_2: ["resolution_wh"],
}

SUPPORTED_TASKS_FLORENCE_2 = [
    "<OD>",
    "<CAPTION_TO_PHRASE_GROUNDING>",
    "<DENSE_REGION_CAPTION>",
    "<REGION_PROPOSAL>",
    "<OCR_WITH_REGION>",
    "<REFERRING_EXPRESSION_SEGMENTATION>",
    "<REGION_TO_SEGMENTATION>",
    "<OPEN_VOCABULARY_DETECTION>",
    "<REGION_TO_CATEGORY>",
    "<REGION_TO_DESCRIPTION>",
]


def validate_lmm_parameters(
    lmm: Union[LMM, str], result: Any, kwargs: Dict[str, Any]
) -> LMM:
    if isinstance(lmm, str):
        try:
            lmm = LMM(lmm.lower())
        except ValueError:
            raise ValueError(
                f"Invalid lmm value: {lmm}. Must be one of {[e.value for e in LMM]}"
            )

    if not isinstance(result, RESULT_TYPES[lmm]):
        raise ValueError(
            f"Invalid LMM result type: {type(result)}. Must be {RESULT_TYPES[lmm]}"
        )

    required_args = REQUIRED_ARGUMENTS.get(lmm, [])
    for arg in required_args:
        if arg not in kwargs:
            raise ValueError(f"Missing required argument: {arg}")

    allowed_args = ALLOWED_ARGUMENTS.get(lmm, [])
    for arg in kwargs:
        if arg not in allowed_args:
            raise ValueError(f"Argument {arg} is not allowed for {lmm.name}")

    return lmm


def from_paligemma(
    result: str, resolution_wh: Tuple[int, int], classes: Optional[List[str]] = None
) -> Tuple[np.ndarray, Optional[np.ndarray], np.ndarray]:
    w, h = resolution_wh
    pattern = re.compile(
        r"(?<!<loc\d{4}>)<loc(\d{4})><loc(\d{4})><loc(\d{4})><loc(\d{4})> ([\w\s\-]+)"
    )
    matches = pattern.findall(result)
    matches = np.array(matches) if matches else np.empty((0, 5))

    xyxy, class_name = matches[:, [1, 0, 3, 2]], matches[:, 4]
    xyxy = xyxy.astype(int) / 1024 * np.array([w, h, w, h])
    class_name = np.char.strip(class_name.astype(str))
    class_id = None

    if classes is not None:
        mask = np.array([name in classes for name in class_name]).astype(bool)
        xyxy, class_name = xyxy[mask], class_name[mask]
        class_id = np.array([classes.index(name) for name in class_name])

    return xyxy, class_id, class_name


def from_florence_2(
    result: dict, resolution_wh: Tuple[int, int]
) -> Tuple[
    np.ndarray, Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]
]:
    """
    Parse results from the Florence 2 multi-model model.
    https://huggingface.co/microsoft/Florence-2-large

    Parameters:
        result: dict containing the model output

    Returns:
        xyxy (np.ndarray): An array of shape `(n, 4)` containing
            the bounding boxes coordinates in format `[x1, y1, x2, y2]`
        labels: (Optional[np.ndarray]): An array of shape `(n,)` containing
            the class labels for each bounding box
        masks: (Optional[np.ndarray]): An array of shape `(n, h, w)` containing
            the segmentation masks for each bounding box
        obb_boxes: (Optional[np.ndarray]): An array of shape `(n, 4, 2)` containing
            oriented bounding boxes.
    """
    assert len(result) == 1, f"Expected result with a single element. Got: {result}"
    task = list(result.keys())[0]
    if task not in SUPPORTED_TASKS_FLORENCE_2:
        raise ValueError(
            f"{task} not supported. Supported tasks are: {SUPPORTED_TASKS_FLORENCE_2}"
        )
    result = result[task]

    if task in ["<OD>", "<CAPTION_TO_PHRASE_GROUNDING>", "<DENSE_REGION_CAPTION>"]:
        xyxy = np.array(result["bboxes"], dtype=np.float32)
        labels = np.array(result["labels"])
        return xyxy, labels, None, None

    if task == "<REGION_PROPOSAL>":
        xyxy = np.array(result["bboxes"], dtype=np.float32)
        # provides labels, but they are ["", "", "", ...]
        return xyxy, None, None, None

    if task == "<OCR_WITH_REGION>":
        xyxyxyxy = np.array(result["quad_boxes"], dtype=np.float32)
        xyxyxyxy = xyxyxyxy.reshape(-1, 4, 2)
        xyxy = np.array([polygon_to_xyxy(polygon) for polygon in xyxyxyxy])
        labels = np.array(result["labels"])
        return xyxy, labels, None, xyxyxyxy

    if task in ["<REFERRING_EXPRESSION_SEGMENTATION>", "<REGION_TO_SEGMENTATION>"]:
        xyxy_list = []
        masks_list = []
        for polygons_of_same_class in result["polygons"]:
            for polygon in polygons_of_same_class:
                polygon = np.reshape(polygon, (-1, 2)).astype(np.int32)
                mask = polygon_to_mask(polygon, resolution_wh).astype(bool)
                masks_list.append(mask)
                xyxy = polygon_to_xyxy(polygon)
                xyxy_list.append(xyxy)
            # per-class labels also provided, but they are ["", "", "", ...]
            # when we figure out how to set class names, we can do
            # zip(result["labels"], result["polygons"])
        xyxy = np.array(xyxy_list, dtype=np.float32)
        masks = np.array(masks_list)
        return xyxy, None, masks, None

    if task == "<OPEN_VOCABULARY_DETECTION>":
        xyxy = np.array(result["bboxes"], dtype=np.float32)
        labels = np.array(result["bboxes_labels"])
        # Also has "polygons" and "polygons_labels", but they don't seem to be used
        return xyxy, labels, None, None

    if task in ["<REGION_TO_CATEGORY>", "<REGION_TO_DESCRIPTION>"]:
        assert isinstance(
            result, str
        ), f"Expected string as <REGION_TO_CATEGORY> result, got {type(result)}"

        if result == "No object detected.":
            return np.empty((0, 4), dtype=np.float32), np.array([]), None, None

        pattern = re.compile(r"<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>")
        match = pattern.search(result)
        assert (
            match is not None
        ), f"Expected string to end in location tags, but got {result}"

        w, h = resolution_wh
        xyxy = np.array([match.groups()], dtype=np.float32)
        xyxy *= np.array([w, h, w, h]) / 1000
        result_string = result[: match.start()]
        labels = np.array([result_string])
        return xyxy, labels, None, None

    assert False, f"Unimplemented task: {task}"
cat << EOF

Our specific requirements are seamless integrating with comfy by respecting the input and output formats expected by the following core comfy nodes found on "comfy_core_image.py" our integration needs to process either a single individual image or a list of images and output the same a singel or list of individual images together with the text and/or data generated such as json files conatining the anotations and locators tags

#comfy_core_image.py
class PreviewImage(SaveImage):
    def __init__(self):
        self.output_dir = folder_paths.get_temp_directory()
        self.type = "temp"
        self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
        self.compress_level = 1

    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"images": ("IMAGE", ), },
                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                }

class LoadImage:
    @classmethod
    def INPUT_TYPES(s):
        input_dir = folder_paths.get_input_directory()
        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
        return {"required":
                    {"image": (sorted(files), {"image_upload": True})},
                }

    CATEGORY = "image"

    RETURN_TYPES = ("IMAGE", "MASK")
    FUNCTION = "load_image"
    def load_image(self, image):
        image_path = folder_paths.get_annotated_filepath(image)
        
        img = node_helpers.pillow(Image.open, image_path)
        
        output_images = []
        output_masks = []
        w, h = None, None

        excluded_formats = ['MPO']
        
        for i in ImageSequence.Iterator(img):
            i = node_helpers.pillow(ImageOps.exif_transpose, i)

            if i.mode == 'I':
                i = i.point(lambda i: i * (1 / 255))
            image = i.convert("RGB")

            if len(output_images) == 0:
                w = image.size[0]
                h = image.size[1]
            
            if image.size[0] != w or image.size[1] != h:
                continue
            
            image = np.array(image).astype(np.float32) / 255.0
            image = torch.from_numpy(image)[None,]
            if 'A' in i.getbands():
                mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
                mask = 1. - torch.from_numpy(mask)
            else:
                mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
            output_images.append(image)
            output_masks.append(mask.unsqueeze(0))

        if len(output_images) > 1 and img.format not in excluded_formats:
            output_image = torch.cat(output_images, dim=0)
            output_mask = torch.cat(output_masks, dim=0)
        else:
            output_image = output_images[0]
            output_mask = output_masks[0]

        return (output_image, output_mask)

    @classmethod
    def IS_CHANGED(s, image):
        image_path = folder_paths.get_annotated_filepath(image)
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()

    @classmethod
    def VALIDATE_INPUTS(s, image):
        if not folder_paths.exists_annotated_filepath(image):
            return "Invalid image file: {}".format(image)

        return True
cat << EOF

We are closely following other implementations to achieve the same integrations but following our requirements which were previously stated including using IFChat to manage the setup loading and tasks outputs from Florence the following is code that might be useful or we need to replicate for our implemntation

#florencecomfyexample.py
import torch
import torchvision.transforms.functional as F
import io
import os
import matplotlib
matplotlib.use('Agg')   
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from PIL import Image, ImageDraw, ImageColor, ImageFont
import random
import numpy as np
import re
from pathlib import Path

#workaround for unnecessary flash_attn requirement
from unittest.mock import patch
from transformers.dynamic_module_utils import get_imports

def fixed_get_imports(filename: str | os.PathLike) -> list[str]:
    try:
        if not str(filename).endswith("modeling_florence2.py"):
            return get_imports(filename)
        imports = get_imports(filename)
        imports.remove("flash_attn")
    except:
        print(f"No flash_attn import to remove")
        pass
    return imports


import comfy.model_management as mm
from comfy.utils import ProgressBar
import folder_paths

script_directory = os.path.dirname(os.path.abspath(__file__))

from transformers import AutoModelForCausalLM, AutoProcessor, set_seed

class DownloadAndLoadFlorence2Model:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "model": (
                    [ 
                    'microsoft/Florence-2-base',
                    'microsoft/Florence-2-base-ft',
                    'microsoft/Florence-2-large',
                    'microsoft/Florence-2-large-ft',
                    'HuggingFaceM4/Florence-2-DocVQA',
                    'thwri/CogFlorence-2.1-Large',
                    'thwri/CogFlorence-2.2-Large',
                    'gokaygokay/Florence-2-SD3-Captioner',
                    'gokaygokay/Florence-2-Flux-Large',
                    'MiaoshouAI/Florence-2-base-PromptGen-v1.5',
                    'MiaoshouAI/Florence-2-large-PromptGen-v1.5'
                    ],
                    {
                    "default": 'microsoft/Florence-2-base'
                    }),
            "precision": ([ 'fp16','bf16','fp32'],
                    {
                    "default": 'fp16'
                    }),
            "attention": (
                    [ 'flash_attention_2', 'sdpa', 'eager'],
                    {
                    "default": 'sdpa'
                    }),
            },
            "optional": {
                "lora": ("PEFTLORA",),
            }
        }

    RETURN_TYPES = ("FL2MODEL",)
    RETURN_NAMES = ("florence2_model",)
    FUNCTION = "loadmodel"
    CATEGORY = "Florence2"

    def loadmodel(self, model, precision, attention, lora=None):
        device = mm.get_torch_device()
        offload_device = mm.unet_offload_device()
        dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision]

        model_name = model.rsplit('/', 1)[-1]
        model_path = os.path.join(folder_paths.models_dir, "LLM", model_name)
        
        if not os.path.exists(model_path):
            print(f"Downloading Florence2 model to: {model_path}")
            from huggingface_hub import snapshot_download
            snapshot_download(repo_id=model,
                            local_dir=model_path,
                            local_dir_use_symlinks=False)
            
        print(f"using {attention} for attention")
        with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports): #workaround for unnecessary flash_attn requirement
            model = AutoModelForCausalLM.from_pretrained(model_path, attn_implementation=attention, device_map=device, torch_dtype=dtype,trust_remote_code=True)
        processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)

        if lora is not None:
            from peft import PeftModel
            adapter_name = lora
            model = PeftModel.from_pretrained(model, adapter_name, trust_remote_code=True)
        
        florence2_model = {
            'model': model, 
            'processor': processor,
            'dtype': dtype
            }

        return (florence2_model,)
    
class DownloadAndLoadFlorence2Lora:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "model": (
                    [ 
                    'NikshepShetty/Florence-2-pixelprose',
                    ],
                  ),            
            },
          
        }

    RETURN_TYPES = ("PEFTLORA",)
    RETURN_NAMES = ("lora",)
    FUNCTION = "loadmodel"
    CATEGORY = "Florence2"

    def loadmodel(self, model):
        model_name = model.rsplit('/', 1)[-1]
        model_path = os.path.join(folder_paths.models_dir, "LLM", model_name)
        
        if not os.path.exists(model_path):
            print(f"Downloading Florence2 lora model to: {model_path}")
            from huggingface_hub import snapshot_download
            snapshot_download(repo_id=model,
                            local_dir=model_path,
                            local_dir_use_symlinks=False)
        return (model_path,)
    
class Florence2ModelLoader:

    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "model": ([item.name for item in Path(folder_paths.models_dir, "LLM").iterdir() if item.is_dir()], {"tooltip": "models are expected to be in Comfyui/models/LLM folder"}),
            "precision": (['fp16','bf16','fp32'],),
            "attention": (
                    [ 'flash_attention_2', 'sdpa', 'eager'],
                    {
                    "default": 'sdpa'
                    }),
            },
            "optional": {
                "lora": ("PEFTLORA",),
            }
        }

    RETURN_TYPES = ("FL2MODEL",)
    RETURN_NAMES = ("florence2_model",)
    FUNCTION = "loadmodel"
    CATEGORY = "Florence2"

    def loadmodel(self, model, precision, attention, lora=None):
        device = mm.get_torch_device()
        dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision]
        model_path = Path(folder_paths.models_dir, "LLM", model)
        print(f"Loading model from {model_path}")
        print(f"using {attention} for attention")
        with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports): #workaround for unnecessary flash_attn requirement
            model = AutoModelForCausalLM.from_pretrained(model_path, attn_implementation=attention, device_map=device, torch_dtype=dtype,trust_remote_code=True)
        processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)

        if lora is not None:
            from peft import PeftModel
            adapter_name = lora
            model = PeftModel.from_pretrained(model, adapter_name, trust_remote_code=True)
        
        florence2_model = {
            'model': model, 
            'processor': processor,
            'dtype': dtype
            }
   
        return (florence2_model,)
    
class Florence2Run:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE", ),
                "florence2_model": ("FL2MODEL", ),
                "text_input": ("STRING", {"default": "", "multiline": True}),
                "task": (
                    [ 
                    'region_caption',
                    'dense_region_caption',
                    'region_proposal',
                    'caption',
                    'detailed_caption',
                    'more_detailed_caption',
                    'caption_to_phrase_grounding',
                    'referring_expression_segmentation',
                    'ocr',
                    'ocr_with_region',
                    'docvqa',
                    'prompt_gen_tags',
                    'prompt_gen_mixed_caption'
                    ],
                   ),
                "fill_mask": ("BOOLEAN", {"default": True}),
            },
            "optional": {
                "keep_model_loaded": ("BOOLEAN", {"default": False}),
                "max_new_tokens": ("INT", {"default": 1024, "min": 1, "max": 4096}),
                "num_beams": ("INT", {"default": 3, "min": 1, "max": 64}),
                "do_sample": ("BOOLEAN", {"default": True}),
                "output_mask_select": ("STRING", {"default": ""}),
                "seed": ("INT", {"default": 1, "min": 1, "max": 0xffffffffffffffff}),
            }
        }
    
    RETURN_TYPES = ("IMAGE", "MASK", "STRING", "JSON")
    RETURN_NAMES =("image", "mask", "caption", "data") 
    FUNCTION = "encode"
    CATEGORY = "Florence2"

    def hash_seed(self, seed):
        import hashlib
        # Convert the seed to a string and then to bytes
        seed_bytes = str(seed).encode('utf-8')
        # Create a SHA-256 hash of the seed bytes
        hash_object = hashlib.sha256(seed_bytes)
        # Convert the hash to an integer
        hashed_seed = int(hash_object.hexdigest(), 16)
        # Ensure the hashed seed is within the acceptable range for set_seed
        return hashed_seed % (2**32)

    def encode(self, image, text_input, florence2_model, task, fill_mask, keep_model_loaded=False, 
            num_beams=3, max_new_tokens=1024, do_sample=True, output_mask_select="", seed=None):
        device = mm.get_torch_device()
        _, height, width, _ = image.shape
        offload_device = mm.unet_offload_device()
        annotated_image_tensor = None
        mask_tensor = None
        processor = florence2_model['processor']
        model = florence2_model['model']
        dtype = florence2_model['dtype']
        model.to(device)
        
        if seed:
            set_seed(self.hash_seed(seed))

        colormap = ['blue','orange','green','purple','brown','pink','olive','cyan','red',
                    'lime','indigo','violet','aqua','magenta','gold','tan','skyblue']

        prompts = {
            'region_caption': '<OD>',
            'dense_region_caption': '<DENSE_REGION_CAPTION>',
            'region_proposal': '<REGION_PROPOSAL>',
            'caption': '<CAPTION>',
            'detailed_caption': '<DETAILED_CAPTION>',
            'more_detailed_caption': '<MORE_DETAILED_CAPTION>',
            'caption_to_phrase_grounding': '<CAPTION_TO_PHRASE_GROUNDING>',
            'referring_expression_segmentation': '<REFERRING_EXPRESSION_SEGMENTATION>',
            'ocr': '<OCR>',
            'ocr_with_region': '<OCR_WITH_REGION>',
            'docvqa': '<DocVQA>',
            'prompt_gen_tags': '<GENERATE_TAGS>',
            'prompt_gen_mixed_caption': '<MIXED_CAPTION>'
        }
        task_prompt = prompts.get(task, '<OD>')

        if (task not in ['referring_expression_segmentation', 'caption_to_phrase_grounding', 'docvqa']) and text_input:
            raise ValueError("Text input (prompt) is only supported for 'referring_expression_segmentation', 'caption_to_phrase_grounding', and 'docvqa'")

        if text_input != "":
            prompt = task_prompt + " " + text_input
        else:
            prompt = task_prompt

        image = image.permute(0, 3, 1, 2)
        
        out = []
        out_masks = []
        out_results = []
        out_data = []
        pbar = ProgressBar(len(image))
        for img in image:
            image_pil = F.to_pil_image(img)
            inputs = processor(text=prompt, images=image_pil, return_tensors="pt", do_rescale=False).to(dtype).to(device)

            generated_ids = model.generate(
                input_ids=inputs["input_ids"],
                pixel_values=inputs["pixel_values"],
                max_new_tokens=max_new_tokens,
                do_sample=do_sample,
                num_beams=num_beams,
            )

            results = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
            print(results)
            # cleanup the special tokens from the final list
            if task == 'ocr_with_region':
                clean_results = str(results)       
                cleaned_string = re.sub(r'</?s>|<[^>]*>', '\n',  clean_results)
                clean_results = re.sub(r'\n+', '\n', cleaned_string)
            else:
                clean_results = str(results)       
                clean_results = clean_results.replace('
', '')
                clean_results = clean_results.replace('
', '')

             #return single string if only one image for compatibility with nodes that can't handle string lists
            if len(image) == 1:
                out_results = clean_results
            else:
                out_results.append(clean_results)

            W, H = image_pil.size
            
            parsed_answer = processor.post_process_generation(results, task=task_prompt, image_size=(W, H))

            if task == 'region_caption' or task == 'dense_region_caption' or task == 'caption_to_phrase_grounding' or task == 'region_proposal':           
                fig, ax = plt.subplots(figsize=(W / 100, H / 100), dpi=100)
                fig.subplots_adjust(left=0, right=1, top=1, bottom=0)
                ax.imshow(image_pil)
                bboxes = parsed_answer[task_prompt]['bboxes']
                labels = parsed_answer[task_prompt]['labels']

                mask_indexes = []
                # Determine mask indexes outside the loop
                if output_mask_select != "":
                    mask_indexes = [n for n in output_mask_select.split(",")]
                    print(mask_indexes)
                else:
                    mask_indexes = [str(i) for i in range(len(bboxes))]

                # Initialize mask_layer only if needed
                if fill_mask:
                    mask_layer = Image.new('RGB', image_pil.size, (0, 0, 0))
                    mask_draw = ImageDraw.Draw(mask_layer)

                for index, (bbox, label) in enumerate(zip(bboxes, labels)):
                    # Modify the label to include the index
                    indexed_label = f"{index}.{label}"
                    
                    if fill_mask:
                        if str(index) in mask_indexes:
                            print("match index:", str(index), "in mask_indexes:", mask_indexes)
                            mask_draw.rectangle([bbox[0], bbox[1], bbox[2], bbox[3]], fill=(255, 255, 255))
                        if label in mask_indexes:
                            print("match label")
                            mask_draw.rectangle([bbox[0], bbox[1], bbox[2], bbox[3]], fill=(255, 255, 255))

                    # Create a Rectangle patch
                    rect = patches.Rectangle(
                        (bbox[0], bbox[1]),  # (x,y) - lower left corner
                        bbox[2] - bbox[0],   # Width
                        bbox[3] - bbox[1],   # Height
                        linewidth=1,
                        edgecolor='r',
                        facecolor='none',
                        label=indexed_label
                    )
                     # Calculate text width with a rough estimation
                    text_width = len(label) * 6  # Adjust multiplier based on your font size
                    text_height = 12  # Adjust based on your font size

                    # Initial text position
                    text_x = bbox[0]
                    text_y = bbox[1] - text_height  # Position text above the top-left of the bbox

                    # Adjust text_x if text is going off the left or right edge
                    if text_x < 0:
                        text_x = 0
                    elif text_x + text_width > W:
                        text_x = W - text_width

                    # Adjust text_y if text is going off the top edge
                    if text_y < 0:
                        text_y = bbox[3]  # Move text below the bottom-left of the bbox if it doesn't overlap with bbox

                    # Add the rectangle to the plot
                    ax.add_patch(rect)
                    facecolor = random.choice(colormap) if len(image) == 1 else 'red'
                    # Add the label
                    plt.text(
                        text_x,
                        text_y,
                        indexed_label,
                        color='white',
                        fontsize=12,
                        bbox=dict(facecolor=facecolor, alpha=0.5)
                    )
                if fill_mask:             
                    mask_tensor = F.to_tensor(mask_layer)
                    mask_tensor = mask_tensor.unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
                    mask_tensor = mask_tensor.mean(dim=0, keepdim=True)
                    mask_tensor = mask_tensor.repeat(1, 1, 1, 3)
                    mask_tensor = mask_tensor[:, :, :, 0]
                    out_masks.append(mask_tensor)           

                # Remove axis and padding around the image
                ax.axis('off')
                ax.margins(0,0)
                ax.get_xaxis().set_major_locator(plt.NullLocator())
                ax.get_yaxis().set_major_locator(plt.NullLocator())
                fig.canvas.draw() 
                buf = io.BytesIO()
                plt.savefig(buf, format='png', pad_inches=0)
                buf.seek(0)
                annotated_image_pil = Image.open(buf)

                annotated_image_tensor = F.to_tensor(annotated_image_pil)
                out_tensor = annotated_image_tensor[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
                out.append(out_tensor)
               
                out_data.append(bboxes)

                
                pbar.update(1)
    
                plt.close(fig)

            elif task == 'referring_expression_segmentation':
                # Create a new black image
                mask_image = Image.new('RGB', (W, H), 'black')
                mask_draw = ImageDraw.Draw(mask_image)
  
                predictions = parsed_answer[task_prompt]
    
                # Iterate over polygons and labels  
                for polygons, label in zip(predictions['polygons'], predictions['labels']):
                    color = random.choice(colormap)
                    for _polygon in polygons:  
                        _polygon = np.array(_polygon).reshape(-1, 2)
                        # Clamp polygon points to image boundaries
                        _polygon = np.clip(_polygon, [0, 0], [W - 1, H - 1])
                        if len(_polygon) < 3:  
                            print('Invalid polygon:', _polygon)
                            continue  
                        
                        _polygon = _polygon.reshape(-1).tolist()
                        
                        # Draw the polygon
                        if fill_mask:
                            overlay = Image.new('RGBA', image_pil.size, (255, 255, 255, 0))
                            image_pil = image_pil.convert('RGBA')
                            draw = ImageDraw.Draw(overlay)
                            color_with_opacity = ImageColor.getrgb(color) + (180,)
                            draw.polygon(_polygon, outline=color, fill=color_with_opacity, width=3)
                            image_pil = Image.alpha_composite(image_pil, overlay)
                        else:
                            draw = ImageDraw.Draw(image_pil)
                            draw.polygon(_polygon, outline=color, width=3)

                        #draw mask
                        mask_draw.polygon(_polygon, outline="white", fill="white")
                        
                image_tensor = F.to_tensor(image_pil)
                image_tensor = image_tensor[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float() 
                out.append(image_tensor)

                mask_tensor = F.to_tensor(mask_image)
                mask_tensor = mask_tensor.unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
                mask_tensor = mask_tensor.mean(dim=0, keepdim=True)
                mask_tensor = mask_tensor.repeat(1, 1, 1, 3)
                mask_tensor = mask_tensor[:, :, :, 0]
                out_masks.append(mask_tensor)
                pbar.update(1)

            elif task == 'ocr_with_region':
                try:
                    font = ImageFont.load_default().font_variant(size=24)
                except:
                    font = ImageFont.load_default()
                predictions = parsed_answer[task_prompt]
                scale = 1
                image_pil = image_pil.convert('RGBA')
                overlay = Image.new('RGBA', image_pil.size, (255, 255, 255, 0))
                draw = ImageDraw.Draw(overlay)
                bboxes, labels = predictions['quad_boxes'], predictions['labels']
                
                # Create a new black image for the mask
                mask_image = Image.new('RGB', (W, H), 'black')
                mask_draw = ImageDraw.Draw(mask_image)
                
                for box, label in zip(bboxes, labels):
                    scaled_box = [v / (width if idx % 2 == 0 else height) for idx, v in enumerate(box)]
                    out_data.append({"label": label, "box": scaled_box})
                    
                    color = random.choice(colormap)
                    new_box = (np.array(box) * scale).tolist()
                    
                    if fill_mask:
                        color_with_opacity = ImageColor.getrgb(color) + (180,)
                        draw.polygon(new_box, outline=color, fill=color_with_opacity, width=3)
                    else:
                        draw.polygon(new_box, outline=color, width=3)
                    
                    draw.text((new_box[0]+8, new_box[1]+2),
                              "{}".format(label),
                              align="right",
                              font=font,
                              fill=color)
                    
                    # Draw the mask
                    mask_draw.polygon(new_box, outline="white", fill="white")
                
                image_pil = Image.alpha_composite(image_pil, overlay)
                image_pil = image_pil.convert('RGB')
                
                image_tensor = F.to_tensor(image_pil)
                image_tensor = image_tensor[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
                out.append(image_tensor)

                # Process the mask
                mask_tensor = F.to_tensor(mask_image)
                mask_tensor = mask_tensor.unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
                mask_tensor = mask_tensor.mean(dim=0, keepdim=True)
                mask_tensor = mask_tensor.repeat(1, 1, 1, 3)
                mask_tensor = mask_tensor[:, :, :, 0]
                out_masks.append(mask_tensor)

                pbar.update(1)
            
            elif task == 'docvqa':
                if text_input == "":
                    raise ValueError("Text input (prompt) is required for 'docvqa'")
                prompt = "<DocVQA> " + text_input

                inputs = processor(text=prompt, images=image_pil, return_tensors="pt", do_rescale=False).to(dtype).to(device)
                generated_ids = model.generate(
                    input_ids=inputs["input_ids"],
                    pixel_values=inputs["pixel_values"],
                    max_new_tokens=max_new_tokens,
                    do_sample=do_sample,
                    num_beams=num_beams,
                )

                results = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
                clean_results = results.replace('
', '').replace('
', '')
                
                if len(image) == 1:
                    out_results = clean_results
                else:
                    out_results.append(clean_results)
                    
                out.append(F.to_tensor(image_pil).unsqueeze(0).permute(0, 2, 3, 1).cpu().float())

                pbar.update(1)
            
        if len(out) > 0:
            out_tensor = torch.cat(out, dim=0)
        else:
            out_tensor = torch.zeros((1, 64,64, 3), dtype=torch.float32, device="cpu")
        if len(out_masks) > 0:
            out_mask_tensor = torch.cat(out_masks, dim=0)
        else:
            out_mask_tensor = torch.zeros((1,64,64), dtype=torch.float32, device="cpu")

        if not keep_model_loaded:
            print("Offloading model...")
            model.to(offload_device)
            mm.soft_empty_cache()
        
        return (out_tensor, out_mask_tensor, out_results, out_data)
     
NODE_CLASS_MAPPINGS = {
    "DownloadAndLoadFlorence2Model": DownloadAndLoadFlorence2Model,
    "DownloadAndLoadFlorence2Lora": DownloadAndLoadFlorence2Lora,
    "Florence2ModelLoader": Florence2ModelLoader,
    "Florence2Run": Florence2Run,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "DownloadAndLoadFlorence2Model": "DownloadAndLoadFlorence2Model",
    "DownloadAndLoadFlorence2Lora": "DownloadAndLoadFlorence2Lora",
    "Florence2ModelLoader": "Florence2ModelLoader",
    "Florence2Run": "Florence2Run",
}
cat << EOF

finally within our code we are using a "florence_prompts.json" file to store task:prompt valuses where the task is the task name and the value is a tag that triggers the behaviour and the florence type of response

#florence_prompts.json
{
  "None": "",
  "object_detection": "<OD>",
  "region_proposal": "<REGION_PROPOSAL>",
  "dense_region_caption": "<DENSE_REGION_CAPTION>",
  "caption": "<CAPTION>",
  "detailed_caption": "<DETAILED_CAPTION>",
  "more_detailed_caption": "<MORE_DETAILED_CAPTION>",
  "ocr": "<OCR>",
  "ocr_with_region": "<OCR_WITH_REGION>",
  "region_category": "<REGION_TO_CATEGORY>",
  "region_to_description": "<REGION_TO_DESCRIPTION>"
} 
cat << EOF