358 lines
14 KiB
Python
358 lines
14 KiB
Python
import hashlib
|
|
import io
|
|
import os
|
|
import random
|
|
from typing import Any, Dict, List, Tuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torchvision.transforms.functional as TF
|
|
from PIL import Image, ImageColor, ImageDraw, ImageFont
|
|
|
|
import comfy.model_management as mm
|
|
from comfy.utils import ProgressBar, load_torch_file
|
|
import comfy.ops
|
|
import folder_paths
|
|
|
|
from models.florence2.config import Florence2Config
|
|
from models.florence2.model import Florence2
|
|
from models.florence2.processing import Processor
|
|
|
|
MODEL_DIR = os.path.join(folder_paths.models_dir, "LLM")
|
|
os.makedirs(MODEL_DIR, exist_ok=True)
|
|
folder_paths.add_model_folder_path("LLM", MODEL_DIR)
|
|
|
|
MODE_CONFIGS = {
|
|
"Polygon Mask (prompt)": {
|
|
"token": "<REFERRING_EXPRESSION_SEGMENTATION>",
|
|
"type": "polygon",
|
|
"allows_text": True,
|
|
},
|
|
"Phrase Grounding (boxes)": {
|
|
"token": "<CAPTION_TO_PHRASE_GROUNDING>",
|
|
"type": "bbox",
|
|
"allows_text": True,
|
|
},
|
|
"Region Proposals (boxes only)": {
|
|
"token": "<REGION_PROPOSAL>",
|
|
"type": "bbox",
|
|
"allows_text": False,
|
|
},
|
|
}
|
|
|
|
MODE_CHOICES = tuple(MODE_CONFIGS.keys())
|
|
TASK_CONFIGS = MODE_CONFIGS
|
|
TASK_CHOICES = MODE_CHOICES
|
|
|
|
|
|
MODEL_CHOICES = (
|
|
"microsoft/Florence-2-base",
|
|
"microsoft/Florence-2-base-ft",
|
|
"microsoft/Florence-2-large",
|
|
"microsoft/Florence-2-large-ft",
|
|
"thwri/CogFlorence-2.1-Large",
|
|
"thwri/CogFlorence-2.2-Large",
|
|
)
|
|
|
|
COLOR_BANK = ["blue", "orange", "green", "purple", "pink", "cyan"]
|
|
|
|
|
|
|
|
|
|
|
|
class AILab_Florence2:
|
|
CATEGORY = "🧪AILab/🧽RMBG"
|
|
RETURN_TYPES = ("IMAGE", "MASK", "JSON", "BBOX")
|
|
RETURN_NAMES = ("IMAGE", "MASK", "DATA", "BBOXES")
|
|
FUNCTION = "analyze"
|
|
MODEL_CACHE: Dict[Tuple[str, str, str], Dict[str, Any]] = {}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"prompt": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True,
|
|
"placeholder": "Prompt: e.g. face, a person wearing red coat",
|
|
"tooltip": "Used for polygon masks or phrase grounding; ignored for region proposals.",
|
|
},
|
|
),
|
|
"model": (
|
|
MODEL_CHOICES,
|
|
{
|
|
"default": MODEL_CHOICES[0],
|
|
"tooltip": "Base = stable, +ft = fine-tuned captions, CogFlorence = sharper phrase alignment.",
|
|
},
|
|
),
|
|
"mode": (
|
|
MODE_CHOICES,
|
|
{
|
|
"default": MODE_CHOICES[0],
|
|
"tooltip": "Polygon masks use prompts; phrase grounding/region proposals return boxes.",
|
|
},
|
|
),
|
|
"precision": (
|
|
("fp16", "bf16", "fp32"),
|
|
{"default": "fp16", "tooltip": "Lower precision saves VRAM; fp32 is safest if you hit NaNs."},
|
|
),
|
|
"attention": (
|
|
("flash_attention_2", "sdpa", "eager"),
|
|
{"default": "sdpa", "tooltip": "flash_attn2 needs PyTorch 2.1+; use eager if kernels fail."},
|
|
),
|
|
"fill_mask": (
|
|
"BOOLEAN",
|
|
{"default": True, "tooltip": "When true, bbox tasks also output filled mask tensors."},
|
|
),
|
|
},
|
|
"optional": {
|
|
"output_mask_select": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"tooltip": "Comma-separated indices or labels (e.g. 0,2,person) to limit masks.",
|
|
},
|
|
),
|
|
"max_new_tokens": (
|
|
"INT",
|
|
{"default": 1024, "min": 256, "max": 4096, "step": 32},
|
|
),
|
|
"unload_model": (
|
|
"BOOLEAN",
|
|
{"default": False, "tooltip": "Unload model from VRAM after execution to free memory."},
|
|
),
|
|
},
|
|
}
|
|
|
|
def _ensure_weights(self, model_name: str) -> str:
|
|
target = os.path.join(MODEL_DIR, model_name.split("/")[-1])
|
|
if os.path.exists(target):
|
|
return target
|
|
from huggingface_hub import snapshot_download
|
|
|
|
snapshot_download(model_name, local_dir=target, local_dir_use_symlinks=False)
|
|
return target
|
|
|
|
def _get_model(self, model_name: str, precision: str, attention: str) -> Dict[str, Any]:
|
|
key = (model_name, precision)
|
|
if key in self.MODEL_CACHE:
|
|
return self.MODEL_CACHE[key]
|
|
|
|
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[precision]
|
|
weights = self._ensure_weights(model_name)
|
|
offload = mm.unet_offload_device()
|
|
|
|
config_path = os.path.join(weights, "config.json")
|
|
checkpoint_path = os.path.join(weights, "model.safetensors")
|
|
if not os.path.exists(checkpoint_path):
|
|
checkpoint_path = os.path.join(weights, "pytorch_model.bin")
|
|
if not os.path.exists(checkpoint_path):
|
|
raise FileNotFoundError(f"No model weights found at {weights}")
|
|
|
|
config = Florence2Config.from_json(config_path)
|
|
model = Florence2(config, dtype=dtype, device="cpu", operations=comfy.ops.manual_cast)
|
|
|
|
sd = load_torch_file(checkpoint_path)
|
|
for key_name in [
|
|
"language_model.model.encoder.embed_tokens.weight",
|
|
"language_model.model.decoder.embed_tokens.weight",
|
|
]:
|
|
if key_name in sd and "language_model.model.shared.weight" in sd:
|
|
sd.pop(key_name, None)
|
|
|
|
model.load_state_dict(sd, strict=False)
|
|
del sd
|
|
model.language_model.tie_weights()
|
|
model = model.eval().to(offload)
|
|
|
|
processor = Processor(model_path=weights)
|
|
bundle = {"model": model, "processor": processor, "dtype": dtype}
|
|
self.MODEL_CACHE[key] = bundle
|
|
return bundle
|
|
|
|
@staticmethod
|
|
def _prepare_prompt(mode: str, text_prompt: str) -> str:
|
|
config = MODE_CONFIGS.get(mode, MODE_CONFIGS[MODE_CHOICES[0]])
|
|
text = text_prompt.strip()
|
|
if text and not config["allows_text"]:
|
|
text = ""
|
|
base = config["token"]
|
|
return f"{base} {text}" if text else base
|
|
|
|
def _draw_regions(
|
|
self,
|
|
image_pil: Image.Image,
|
|
predictions: Dict[str, Any],
|
|
fill_mask: bool,
|
|
select_filter: List[str],
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
annotated = image_pil.copy()
|
|
draw = ImageDraw.Draw(annotated)
|
|
|
|
mask_layer = Image.new("RGB", image_pil.size, (0, 0, 0)) if fill_mask else None
|
|
mask_draw = ImageDraw.Draw(mask_layer) if mask_layer else None
|
|
|
|
try:
|
|
font = ImageFont.load_default(size=14)
|
|
except Exception:
|
|
font = ImageFont.load_default()
|
|
|
|
for index, (bbox, label) in enumerate(zip(predictions["bboxes"], predictions["labels"])):
|
|
x0, y0, x1, y1 = bbox
|
|
if y1 < y0:
|
|
y0, y1 = y1, y0
|
|
if x1 < x0:
|
|
x0, x1 = x1, x0
|
|
filter_hit = not select_filter or str(index) in select_filter or label in select_filter
|
|
if fill_mask and filter_hit:
|
|
mask_draw.rectangle([x0, y0, x1, y1], fill=(255, 255, 255))
|
|
|
|
draw.rectangle([x0, y0, x1, y1], outline="red", width=2)
|
|
label_text = f"{index}.{label}"
|
|
color = random.choice(COLOR_BANK)
|
|
text_pos = (x0, max(0, y0 - 16))
|
|
text_bbox = draw.textbbox(text_pos, label_text, font=font)
|
|
draw.rectangle(text_bbox, fill=color)
|
|
draw.text(text_pos, label_text, fill="white", font=font)
|
|
|
|
annotated_tensor = TF.to_tensor(annotated)[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
|
|
if mask_layer is not None:
|
|
mask_tensor = TF.to_tensor(mask_layer).unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
|
|
mask_tensor = mask_tensor.mean(dim=0, keepdim=True)[:, :, :, 0]
|
|
else:
|
|
mask_tensor = torch.zeros((1, annotated_tensor.shape[1], annotated_tensor.shape[2]), dtype=torch.float32)
|
|
|
|
return annotated_tensor, mask_tensor
|
|
|
|
def _segment_polygons(
|
|
self,
|
|
image_pil: Image.Image,
|
|
predictions: Dict[str, Any],
|
|
fill_mask: bool,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
width, height = image_pil.size
|
|
mask_image = Image.new("RGB", (width, height), "black")
|
|
mask_draw = ImageDraw.Draw(mask_image)
|
|
|
|
for polygons, label in zip(predictions["polygons"], predictions["labels"]):
|
|
color = random.choice(COLOR_BANK)
|
|
for polygon in polygons:
|
|
polygon = np.array(polygon).reshape(-1, 2)
|
|
polygon = np.clip(polygon, [0, 0], [width - 1, height - 1])
|
|
if len(polygon) < 3:
|
|
continue
|
|
pts = polygon.reshape(-1).tolist()
|
|
if fill_mask:
|
|
overlay = Image.new("RGBA", image_pil.size, (255, 255, 255, 0))
|
|
draw = ImageDraw.Draw(overlay)
|
|
rgba = ImageColor.getrgb(color) + (180,)
|
|
draw.polygon(pts, outline=color, fill=rgba, width=3)
|
|
image_pil = Image.alpha_composite(image_pil.convert("RGBA"), overlay).convert("RGB")
|
|
else:
|
|
draw = ImageDraw.Draw(image_pil)
|
|
draw.polygon(pts, outline=color, width=3)
|
|
mask_draw.polygon(pts, outline="white", fill="white")
|
|
|
|
image_tensor = TF.to_tensor(image_pil)[:3, :, :].unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
|
|
mask_tensor = TF.to_tensor(mask_image).unsqueeze(0).permute(0, 2, 3, 1).cpu().float()
|
|
mask_tensor = mask_tensor.mean(dim=0, keepdim=True)[:, :, :, 0]
|
|
return image_tensor, mask_tensor
|
|
|
|
def analyze(
|
|
self,
|
|
image: torch.Tensor,
|
|
prompt: str = "",
|
|
model: str = MODEL_CHOICES[0],
|
|
mode: str = MODE_CHOICES[0],
|
|
precision: str = "fp16",
|
|
attention: str = "sdpa",
|
|
fill_mask: bool = True,
|
|
output_mask_select: str = "",
|
|
max_new_tokens: int = 1024,
|
|
unload_model: bool = False,
|
|
model_name: str = None,
|
|
task: str = None,
|
|
text_prompt: str = None,
|
|
keep_model_loaded: bool = None,
|
|
**kwargs,
|
|
):
|
|
if model_name is not None and (model is None or model == MODEL_CHOICES[0]):
|
|
model = model_name
|
|
if task is not None and not mode:
|
|
mode = task
|
|
if text_prompt is not None and not prompt:
|
|
prompt = text_prompt
|
|
if keep_model_loaded is not None:
|
|
unload_model = not keep_model_loaded
|
|
|
|
bundle = self._get_model(model, precision, attention)
|
|
florence_model = bundle["model"]
|
|
processor = bundle["processor"]
|
|
dtype = bundle["dtype"]
|
|
|
|
device = mm.get_torch_device()
|
|
offload = mm.unet_offload_device()
|
|
florence_model.to(device)
|
|
images = image.permute(0, 3, 1, 2)
|
|
|
|
mode_config = MODE_CONFIGS.get(mode, MODE_CONFIGS[MODE_CHOICES[0]])
|
|
prepared_prompt = self._prepare_prompt(mode, prompt)
|
|
out_images, out_masks, out_data = [], [], []
|
|
pbar = ProgressBar(len(images))
|
|
select_filter = [s.strip() for s in output_mask_select.split(",") if s.strip()]
|
|
|
|
for tensor in images:
|
|
image_pil = TF.to_pil_image(tensor)
|
|
inputs = processor(
|
|
text=prepared_prompt,
|
|
images=tensor.unsqueeze(0),
|
|
)
|
|
|
|
generated_ids = florence_model.generate(
|
|
input_ids=inputs["input_ids"].to(device),
|
|
pixel_values=inputs["pixel_values"].to(dtype=dtype, device=device),
|
|
max_new_tokens=max_new_tokens,
|
|
do_sample=False,
|
|
num_beams=3,
|
|
)
|
|
raw = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
|
|
parsed = processor.post_process_generation(raw, task=mode_config["token"], image_size=image_pil.size)
|
|
predictions = parsed[mode_config["token"]]
|
|
|
|
if mode_config["type"] == "bbox":
|
|
annotated, mask = self._draw_regions(image_pil, predictions, fill_mask, select_filter)
|
|
out_images.append(annotated)
|
|
out_masks.append(mask)
|
|
else:
|
|
image_tensor, mask_tensor = self._segment_polygons(image_pil, predictions, fill_mask)
|
|
out_images.append(image_tensor)
|
|
out_masks.append(mask_tensor)
|
|
|
|
out_data.append(predictions)
|
|
pbar.update(1)
|
|
|
|
image_out = torch.cat(out_images, dim=0)
|
|
mask_out = torch.cat(out_masks, dim=0)
|
|
|
|
if unload_model:
|
|
florence_model.to(offload)
|
|
mm.soft_empty_cache()
|
|
|
|
all_bboxes = []
|
|
for entry in out_data:
|
|
if isinstance(entry, dict) and "bboxes" in entry:
|
|
all_bboxes.extend(entry["bboxes"])
|
|
|
|
return image_out, mask_out, out_data, all_bboxes
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AILab_Florence2": AILab_Florence2,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AILab_Florence2": "Florence2 (RMBG)",
|
|
}
|