572 lines
22 KiB
Python
572 lines
22 KiB
Python
"""
|
|
Single Florence2 Detector Class - Replicating kijai's complete workflow
|
|
|
|
This module provides one clean class that internally does what kijai's
|
|
Florence2ModelLoader + Florence2Run do together, but as a single callable interface.
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
import os
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from typing import List, Tuple, Optional, Union, Dict, Any
|
|
from PIL import Image
|
|
|
|
# === DaViT Patch - Apply before any transformers imports ===
|
|
def patch_davit_missing_method():
|
|
"""Non-invasive DaViT patch that doesn't interfere with global imports."""
|
|
def _initialize_weights_dummy(self):
|
|
"""Dummy implementation for DaViT compatibility."""
|
|
pass
|
|
|
|
# Return a function that can patch objects after they're loaded
|
|
def patch_object_if_davit(obj):
|
|
"""Patch a specific object if it's a DaViT that needs fixing."""
|
|
try:
|
|
if (hasattr(obj, '__class__') and
|
|
'DaViT' in str(type(obj)) and
|
|
not hasattr(obj, '_initialize_weights')):
|
|
print(f"Patching DaViT object: {type(obj)}")
|
|
import types
|
|
obj._initialize_weights = types.MethodType(_initialize_weights_dummy, obj)
|
|
except Exception as e:
|
|
# Ignore patching errors
|
|
pass
|
|
|
|
return patch_object_if_davit
|
|
|
|
# Apply the patch immediately
|
|
davit_patcher = patch_davit_missing_method()
|
|
|
|
# DO NOT import transformers at module level - defer until needed
|
|
# This prevents interference with ComfyUI initialization
|
|
|
|
# ComfyUI memory management (with fallback)
|
|
try:
|
|
import comfy.model_management as mm
|
|
except ImportError:
|
|
print("ComfyUI not available, using fallback memory management")
|
|
class MockModelManagement:
|
|
@staticmethod
|
|
def get_torch_device():
|
|
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
@staticmethod
|
|
def unet_offload_device():
|
|
return torch.device("cpu")
|
|
|
|
@staticmethod
|
|
def soft_empty_cache():
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
mm = MockModelManagement()
|
|
|
|
|
|
def fixed_get_imports(filename):
|
|
"""Fix for unnecessary flash_attn requirement - exactly from kijai."""
|
|
# Import here to avoid module-level import issues
|
|
from transformers.dynamic_module_utils import get_imports
|
|
|
|
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
|
|
|
|
|
|
@dataclass
|
|
class BoundingBox:
|
|
"""Simple bounding box representation."""
|
|
x1: int
|
|
y1: int
|
|
x2: int
|
|
y2: int
|
|
confidence: Optional[float] = None
|
|
|
|
@property
|
|
def width(self) -> int:
|
|
return self.x2 - self.x1
|
|
|
|
@property
|
|
def height(self) -> int:
|
|
return self.y2 - self.y1
|
|
|
|
@property
|
|
def area(self) -> int:
|
|
return self.width * self.height
|
|
|
|
@property
|
|
def center(self) -> Tuple[int, int]:
|
|
return (self.x1 + self.width // 2, self.y1 + self.height // 2)
|
|
|
|
def to_tuple(self) -> Tuple[int, int, int, int]:
|
|
"""Returns (x1, y1, x2, y2) tuple."""
|
|
return (self.x1, self.y1, self.x2, self.y2)
|
|
|
|
def __repr__(self):
|
|
conf_str = f", conf={self.confidence:.3f}" if self.confidence else ""
|
|
return f"BoundingBox({self.x1}, {self.y1}, {self.x2}, {self.y2}, area={self.area}{conf_str})"
|
|
|
|
|
|
class ImageRectangleDetector:
|
|
"""
|
|
Single class that replicates kijai's complete Florence2ModelLoader + Florence2Run workflow.
|
|
|
|
Internally handles model loading with precision/attention settings AND inference,
|
|
but provides a simple callable interface for use in any ComfyUI node.
|
|
"""
|
|
|
|
def __init__(self,
|
|
model_name: str = "CogFlorence-2.2-Large",
|
|
comfyui_base_path: str = "A:\\Comfy_Dec\\ComfyUI",
|
|
precision: str = "fp16",
|
|
attention: str = "sdpa",
|
|
min_box_area: int = 1000,
|
|
convert_to_safetensors: bool = False):
|
|
"""
|
|
Initialize detector with kijai's exact parameters.
|
|
|
|
Args:
|
|
model_name: Model directory name in ComfyUI/models/LLM/
|
|
comfyui_base_path: Path to ComfyUI installation
|
|
precision: Model precision ('fp16', 'bf16', 'fp32') - kijai's Florence2ModelLoader param
|
|
attention: Attention mechanism ('sdpa', 'flash_attention_2', 'eager') - kijai's param
|
|
min_box_area: Minimum bounding box area to keep
|
|
convert_to_safetensors: Convert .bin to .safetensors - kijai's param
|
|
"""
|
|
# Store all kijai's parameters
|
|
self.model_name = model_name
|
|
self.comfyui_base_path = comfyui_base_path
|
|
self.precision = precision
|
|
self.attention = attention
|
|
self.min_box_area = min_box_area
|
|
self.convert_to_safetensors = convert_to_safetensors
|
|
|
|
# Model state - loaded lazily
|
|
self.model = None
|
|
self.processor = None
|
|
self.dtype = None
|
|
self.loaded_model_dict = None # kijai's format: {model, processor, dtype}
|
|
|
|
# Get model path
|
|
self.model_path = os.path.join(comfyui_base_path, "models", "LLM", model_name)
|
|
|
|
# Verify model exists
|
|
if not os.path.exists(self.model_path):
|
|
raise FileNotFoundError(f"Model not found at {self.model_path}")
|
|
|
|
def _load_model_kijai_style(self):
|
|
"""
|
|
Load model using kijai's exact Florence2ModelLoader logic.
|
|
This replicates the loadmodel() function from kijai's Florence2ModelLoader class.
|
|
"""
|
|
if self.loaded_model_dict is not None:
|
|
return self.loaded_model_dict
|
|
|
|
# Import transformers only when actually loading model
|
|
try:
|
|
from transformers import AutoProcessor, AutoModelForCausalLM, AutoModelForSeq2SeqLM
|
|
import transformers
|
|
except ImportError:
|
|
raise ImportError("transformers library required. Install with: pip install transformers")
|
|
|
|
print(f"Loading model from {self.model_path}")
|
|
print(f"Florence2 using {self.attention} for attention")
|
|
print(f"Transformers version: {transformers.__version__}")
|
|
|
|
# kijai's exact device and dtype setup
|
|
device = mm.get_torch_device()
|
|
offload_device = mm.unet_offload_device()
|
|
self.dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[self.precision]
|
|
|
|
# kijai's safetensors conversion logic
|
|
if self.convert_to_safetensors:
|
|
model_weight_path = os.path.join(self.model_path, 'pytorch_model.bin')
|
|
if os.path.exists(model_weight_path):
|
|
safetensors_weight_path = os.path.join(self.model_path, 'model.safetensors')
|
|
print(f"Converting {model_weight_path} to {safetensors_weight_path}")
|
|
if not os.path.exists(safetensors_weight_path):
|
|
from safetensors.torch import save_file
|
|
checkpoint = torch.load(model_weight_path, map_location="cpu")
|
|
save_file(checkpoint, safetensors_weight_path)
|
|
print(f"Deleting original file: {model_weight_path}")
|
|
os.remove(model_weight_path)
|
|
print(f"Original {model_weight_path} file deleted.")
|
|
|
|
# Load model with proper error handling
|
|
from unittest.mock import patch
|
|
|
|
print("Loading with trust_remote_code approach...")
|
|
try:
|
|
# Your config.json specifies AutoModelForSeq2SeqLM, not AutoModelForCausalLM!
|
|
with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
|
|
self.model = AutoModelForSeq2SeqLM.from_pretrained(
|
|
self.model_path,
|
|
attn_implementation=self.attention,
|
|
torch_dtype=self.dtype,
|
|
trust_remote_code=True
|
|
).to(offload_device)
|
|
print("AutoModelForSeq2SeqLM loading successful")
|
|
|
|
except Exception as e:
|
|
print(f"AutoModelForSeq2SeqLM failed: {e}")
|
|
print("Trying AutoModelForCausalLM as fallback...")
|
|
|
|
# Fallback to AutoModelForCausalLM for older transformers compatibility
|
|
try:
|
|
with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
|
|
self.model = AutoModelForCausalLM.from_pretrained(
|
|
self.model_path,
|
|
attn_implementation=self.attention,
|
|
torch_dtype=self.dtype,
|
|
trust_remote_code=True
|
|
).to(offload_device)
|
|
print("AutoModelForCausalLM fallback successful")
|
|
|
|
except Exception as e2:
|
|
print(f"Both loading methods failed!")
|
|
print(f"AutoModelForSeq2SeqLM error: {e}")
|
|
print(f"AutoModelForCausalLM error: {e2}")
|
|
raise Exception(f"Could not load model with either method. Check config.json auto_map.")
|
|
|
|
# Apply targeted DaViT patch after model is loaded
|
|
if hasattr(self, 'model') and self.model is not None:
|
|
try:
|
|
# Recursively patch any DaViT components in the loaded model
|
|
def patch_recursive(obj, visited=None):
|
|
if visited is None:
|
|
visited = set()
|
|
|
|
if id(obj) in visited:
|
|
return
|
|
visited.add(id(obj))
|
|
|
|
# Apply the patcher
|
|
davit_patcher(obj)
|
|
|
|
# Recursively check submodules
|
|
if hasattr(obj, '__dict__'):
|
|
for attr_value in obj.__dict__.values():
|
|
if hasattr(attr_value, '__dict__'):
|
|
patch_recursive(attr_value, visited)
|
|
|
|
# Check PyTorch module children
|
|
if hasattr(obj, 'children'):
|
|
try:
|
|
for child in obj.children():
|
|
patch_recursive(child, visited)
|
|
except Exception:
|
|
pass
|
|
|
|
patch_recursive(self.model)
|
|
print("DaViT patching complete")
|
|
|
|
except Exception as e:
|
|
print(f"DaViT patching skipped: {e}")
|
|
|
|
# Load processor - kijai's approach
|
|
self.processor = AutoProcessor.from_pretrained(self.model_path, trust_remote_code=True)
|
|
|
|
# Create kijai's exact return format
|
|
self.loaded_model_dict = {
|
|
'model': self.model,
|
|
'processor': self.processor,
|
|
'dtype': self.dtype
|
|
}
|
|
|
|
print(f"Model loaded successfully on {offload_device}")
|
|
return self.loaded_model_dict
|
|
|
|
|
|
|
|
def detect_objects(self,
|
|
image: Union[Image.Image, np.ndarray, torch.Tensor],
|
|
text_input: str = "rectangular images in page",
|
|
task: str = "caption_to_phrase_grounding",
|
|
fill_mask: bool = True,
|
|
keep_model_loaded: bool = True,
|
|
max_new_tokens: int = 1024,
|
|
num_beams: int = 3,
|
|
do_sample: bool = True) -> List[BoundingBox]:
|
|
"""
|
|
Detect objects using kijai's complete workflow.
|
|
This replicates what kijai's Florence2Run does after getting the model from Florence2ModelLoader.
|
|
|
|
Args:
|
|
image: Input image
|
|
text_input: Prompt text (kijai's parameter name)
|
|
task: Florence2 task type (kijai's parameter)
|
|
fill_mask: Whether to fill masks (kijai's parameter)
|
|
keep_model_loaded: Whether to keep model in GPU memory (kijai's parameter)
|
|
max_new_tokens: Maximum tokens to generate (kijai's parameter)
|
|
num_beams: Number of beams for generation (kijai's parameter)
|
|
do_sample: Whether to use sampling (kijai's parameter)
|
|
|
|
Returns:
|
|
List of BoundingBox objects for detected regions (filtered by min_box_area)
|
|
"""
|
|
# Load model using kijai's approach
|
|
florence2_model = self._load_model_kijai_style()
|
|
|
|
# Extract components (kijai's format)
|
|
model = florence2_model['model']
|
|
processor = florence2_model['processor']
|
|
dtype = florence2_model['dtype']
|
|
|
|
# kijai's device management
|
|
device = mm.get_torch_device()
|
|
offload_device = mm.unet_offload_device()
|
|
|
|
# Move model to GPU for inference (kijai's approach)
|
|
model.to(device)
|
|
|
|
# Convert image to PIL format
|
|
pil_image = self._preprocess_image(image)
|
|
|
|
# Prepare task prompt - kijai's exact format
|
|
task_prompt = f"<{task}>"
|
|
if text_input:
|
|
full_prompt = f"{task_prompt}{text_input}"
|
|
else:
|
|
full_prompt = task_prompt
|
|
|
|
print(f"Running Florence2 inference with prompt: '{full_prompt}'")
|
|
|
|
try:
|
|
# Process inputs - kijai's exact approach
|
|
inputs = processor(
|
|
text=full_prompt,
|
|
images=pil_image,
|
|
return_tensors="pt"
|
|
).to(device)
|
|
|
|
# Run generation - kijai's exact parameters from Florence2Run
|
|
with torch.no_grad():
|
|
generated_ids = model.generate(
|
|
input_ids=inputs["input_ids"],
|
|
pixel_values=inputs["pixel_values"],
|
|
max_new_tokens=max_new_tokens,
|
|
num_beams=num_beams,
|
|
do_sample=do_sample
|
|
)
|
|
|
|
# Decode results
|
|
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
|
|
|
|
print(f"Generated text: {generated_text}")
|
|
|
|
# Parse results using processor's post_process_generation
|
|
parsed_answer = processor.post_process_generation(
|
|
generated_text,
|
|
task=task,
|
|
image_size=(pil_image.width, pil_image.height)
|
|
)
|
|
|
|
print(f"Parsed answer: {parsed_answer}")
|
|
|
|
# Extract bounding boxes
|
|
boxes = self._extract_bounding_boxes(parsed_answer, task)
|
|
|
|
# Filter by minimum area (our addition)
|
|
filtered_boxes = self._filter_small_boxes(boxes)
|
|
|
|
print(f"Detected {len(filtered_boxes)} objects after filtering")
|
|
return filtered_boxes
|
|
|
|
finally:
|
|
# Handle model offloading - kijai's exact approach
|
|
if not keep_model_loaded:
|
|
print("Offloading model...")
|
|
model.to(offload_device)
|
|
mm.soft_empty_cache()
|
|
|
|
def _extract_bounding_boxes(self, parsed_answer: Dict[str, Any], task: str) -> List[BoundingBox]:
|
|
"""Extract bounding boxes from Florence2 parsed output."""
|
|
boxes = []
|
|
|
|
if task in parsed_answer:
|
|
result = parsed_answer[task]
|
|
|
|
if 'bboxes' in result:
|
|
bboxes = result['bboxes']
|
|
|
|
for bbox in bboxes:
|
|
if len(bbox) == 4:
|
|
x1, y1, x2, y2 = bbox
|
|
# Ensure valid coordinates
|
|
box = BoundingBox(
|
|
x1=int(max(0, x1)),
|
|
y1=int(max(0, y1)),
|
|
x2=int(x2),
|
|
y2=int(y2)
|
|
)
|
|
|
|
# Only add valid boxes
|
|
if box.width > 0 and box.height > 0:
|
|
boxes.append(box)
|
|
|
|
return boxes
|
|
|
|
def _filter_small_boxes(self, boxes: List[BoundingBox]) -> List[BoundingBox]:
|
|
"""Filter out boxes smaller than minimum area threshold."""
|
|
filtered = []
|
|
|
|
for box in boxes:
|
|
if box.area >= self.min_box_area:
|
|
filtered.append(box)
|
|
else:
|
|
print(f"Filtering small box: {box}")
|
|
|
|
return filtered
|
|
|
|
def _preprocess_image(self, image: Union[Image.Image, np.ndarray, torch.Tensor]) -> Image.Image:
|
|
"""Convert input image to PIL Image format."""
|
|
if isinstance(image, Image.Image):
|
|
return image
|
|
elif isinstance(image, np.ndarray):
|
|
# Handle numpy arrays
|
|
if image.dtype in [np.float32, np.float64]:
|
|
image = (image * 255).astype(np.uint8)
|
|
if len(image.shape) == 4:
|
|
image = image[0] # Remove batch dimension
|
|
if image.shape[0] == 3:
|
|
image = np.transpose(image, (1, 2, 0)) # CHW to HWC
|
|
return Image.fromarray(image)
|
|
elif isinstance(image, torch.Tensor):
|
|
# Handle torch tensors
|
|
if image.dim() == 4:
|
|
image = image[0] # Remove batch dimension
|
|
if image.shape[0] == 3:
|
|
image = image.permute(1, 2, 0) # CHW to HWC
|
|
if image.dtype in [torch.float32, torch.float16]:
|
|
image = (image * 255).clamp(0, 255).byte()
|
|
return Image.fromarray(image.cpu().numpy())
|
|
else:
|
|
raise ValueError(f"Unsupported image type: {type(image)}")
|
|
|
|
def get_model_info(self) -> Dict[str, Any]:
|
|
"""Get information about the detector configuration."""
|
|
# Import transformers only when needed
|
|
try:
|
|
import transformers
|
|
transformers_version = transformers.__version__
|
|
except ImportError:
|
|
transformers_version = "Not installed"
|
|
|
|
return {
|
|
"model_name": self.model_name,
|
|
"model_path": self.model_path,
|
|
"precision": self.precision,
|
|
"attention": self.attention,
|
|
"min_box_area": self.min_box_area,
|
|
"convert_to_safetensors": self.convert_to_safetensors,
|
|
"loaded": self.loaded_model_dict is not None,
|
|
"transformers_version": transformers_version
|
|
}
|
|
|
|
def update_settings(self, min_box_area: Optional[int] = None):
|
|
"""Update settings without reloading model."""
|
|
if min_box_area is not None:
|
|
self.min_box_area = min_box_area
|
|
|
|
|
|
# Utility functions
|
|
def visualize_detections(image: Union[Image.Image, np.ndarray],
|
|
boxes: List[BoundingBox],
|
|
show_labels: bool = True) -> Image.Image:
|
|
"""Draw bounding boxes on image."""
|
|
from PIL import ImageDraw, ImageFont
|
|
|
|
if isinstance(image, np.ndarray):
|
|
image = Image.fromarray(image)
|
|
|
|
vis_image = image.copy()
|
|
draw = ImageDraw.Draw(vis_image)
|
|
colors = ['red', 'blue', 'green', 'orange', 'purple', 'cyan']
|
|
|
|
for i, box in enumerate(boxes):
|
|
color = colors[i % len(colors)]
|
|
|
|
# Draw rectangle
|
|
draw.rectangle([box.x1, box.y1, box.x2, box.y2], outline=color, width=3)
|
|
|
|
# Draw label
|
|
if show_labels:
|
|
label = f"Area: {box.area}"
|
|
try:
|
|
font = ImageFont.truetype("arial.ttf", 16)
|
|
except:
|
|
font = ImageFont.load_default()
|
|
|
|
draw.text((box.x1, box.y1 - 20), label, fill=color, font=font)
|
|
|
|
return vis_image
|
|
|
|
|
|
def crop_images_from_boxes(image: Union[Image.Image, np.ndarray],
|
|
boxes: List[BoundingBox]) -> List[Image.Image]:
|
|
"""Crop individual images from bounding boxes."""
|
|
if isinstance(image, np.ndarray):
|
|
image = Image.fromarray(image)
|
|
|
|
cropped_images = []
|
|
for box in boxes:
|
|
cropped = image.crop((box.x1, box.y1, box.x2, box.y2))
|
|
cropped_images.append(cropped)
|
|
|
|
return cropped_images
|
|
|
|
|
|
# Example usage
|
|
if __name__ == "__main__":
|
|
# Initialize detector with your exact kijai settings
|
|
detector = ImageRectangleDetector(
|
|
model_name="CogFlorence-2.2-Large",
|
|
comfyui_base_path="A:\\Comfy_Dec\\ComfyUI",
|
|
precision="fp16", # kijai's Florence2ModelLoader param
|
|
attention="sdpa", # kijai's Florence2ModelLoader param
|
|
min_box_area=1000,
|
|
convert_to_safetensors=False # kijai's Florence2ModelLoader param
|
|
)
|
|
|
|
print("Detector configuration:")
|
|
model_info = detector.get_model_info()
|
|
for key, value in model_info.items():
|
|
print(f" {key}: {value}")
|
|
|
|
# Example usage (replace with your image)
|
|
image_path = "test_image.jpg"
|
|
if os.path.exists(image_path):
|
|
image = Image.open(image_path)
|
|
|
|
# Detect objects with your exact kijai settings
|
|
boxes = detector.detect_objects(
|
|
image=image,
|
|
text_input="rectangular images in page", # kijai's parameter name
|
|
task="caption_to_phrase_grounding", # kijai's parameter
|
|
fill_mask=True, # kijai's parameter
|
|
keep_model_loaded=True, # kijai's parameter
|
|
max_new_tokens=1024, # kijai's parameter (you use 2048)
|
|
num_beams=3, # kijai's parameter
|
|
do_sample=True # kijai's parameter
|
|
)
|
|
|
|
print(f"\nFound {len(boxes)} objects:")
|
|
for i, box in enumerate(boxes):
|
|
print(f" {i+1}: {box}")
|
|
|
|
# Visualize and save results
|
|
if boxes:
|
|
vis_image = visualize_detections(image, boxes)
|
|
vis_image.save("detection_results.jpg")
|
|
|
|
cropped_images = crop_images_from_boxes(image, boxes)
|
|
for i, cropped in enumerate(cropped_images):
|
|
cropped.save(f"cropped_{i+1}.jpg")
|
|
|
|
print("Results saved!")
|
|
else:
|
|
print("Please provide a test image to run the example") |