365 lines
12 KiB
Python
365 lines
12 KiB
Python
import os
|
|
import ast
|
|
import json
|
|
from dataclasses import dataclass
|
|
from typing import List, Dict, Any, Tuple
|
|
|
|
import torch
|
|
from PIL import Image
|
|
from huggingface_hub import snapshot_download
|
|
from transformers import (
|
|
AutoProcessor,
|
|
Qwen2_5_VLForConditionalGeneration,
|
|
BitsAndBytesConfig,
|
|
)
|
|
|
|
import folder_paths
|
|
|
|
|
|
def parse_json(json_output: str) -> str:
|
|
"""Extract the JSON payload from a model response string."""
|
|
if "```json" in json_output:
|
|
json_output = json_output.split("```json", 1)[1]
|
|
json_output = json_output.split("```", 1)[0]
|
|
|
|
try:
|
|
parsed = json.loads(json_output)
|
|
if isinstance(parsed, dict) and "content" in parsed:
|
|
inner = parsed["content"]
|
|
if isinstance(inner, str):
|
|
json_output = inner
|
|
except Exception:
|
|
pass
|
|
return json_output
|
|
|
|
|
|
def parse_boxes(
|
|
text: str,
|
|
img_width: int,
|
|
img_height: int,
|
|
input_w: int,
|
|
input_h: int,
|
|
score_threshold: float = 0.0,
|
|
) -> List[Dict[str, Any]]:
|
|
"""Return bounding boxes parsed from the model's raw JSON output."""
|
|
text = parse_json(text)
|
|
try:
|
|
data = json.loads(text)
|
|
except Exception:
|
|
try:
|
|
data = ast.literal_eval(text)
|
|
except Exception:
|
|
end_idx = text.rfind('"}') + len('"}')
|
|
truncated = text[:end_idx] + "]"
|
|
data = ast.literal_eval(truncated)
|
|
if isinstance(data, dict):
|
|
inner = data.get("content")
|
|
if isinstance(inner, str):
|
|
try:
|
|
data = ast.literal_eval(inner)
|
|
except Exception:
|
|
data = []
|
|
else:
|
|
data = []
|
|
items: List[DetectedBox] = []
|
|
x_scale = img_width / input_w
|
|
y_scale = img_height / input_h
|
|
|
|
for item in data:
|
|
box = item.get("bbox_2d") or item.get("bbox") or item
|
|
label = item.get("label", "")
|
|
score = float(item.get("score", 1.0))
|
|
y1, x1, y2, x2 = box[1], box[0], box[3], box[2]
|
|
abs_y1 = int(y1 * y_scale)
|
|
abs_x1 = int(x1 * x_scale)
|
|
abs_y2 = int(y2 * y_scale)
|
|
abs_x2 = int(x2 * x_scale)
|
|
if abs_x1 > abs_x2:
|
|
abs_x1, abs_x2 = abs_x2, abs_x1
|
|
if abs_y1 > abs_y2:
|
|
abs_y1, abs_y2 = abs_y2, abs_y1
|
|
if score >= score_threshold:
|
|
items.append(DetectedBox([abs_x1, abs_y1, abs_x2, abs_y2], score, label))
|
|
items.sort(key=lambda x: x.score, reverse=True)
|
|
return [
|
|
{"score": b.score, "bbox": b.bbox, "label": b.label}
|
|
for b in items
|
|
]
|
|
|
|
|
|
@dataclass
|
|
class DetectedBox:
|
|
bbox: List[int]
|
|
score: float
|
|
label: str = ""
|
|
|
|
|
|
@dataclass
|
|
class QwenModel:
|
|
model: Any
|
|
processor: Any
|
|
device: str
|
|
|
|
|
|
class DownloadAndLoadQwenModel:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_name": ([
|
|
"Qwen/Qwen2.5-VL-3B-Instruct",
|
|
"Qwen/Qwen2.5-VL-7B-Instruct",
|
|
"Qwen/Qwen2.5-VL-32B-Instruct",
|
|
"Qwen/Qwen2.5-VL-72B-Instruct",
|
|
], ),
|
|
"device": ([
|
|
"auto",
|
|
"cuda:0",
|
|
"cuda:1",
|
|
"cpu",
|
|
], ),
|
|
"precision": ([
|
|
"INT4",
|
|
"INT8",
|
|
"BF16",
|
|
"FP16",
|
|
"FP32",
|
|
], ),
|
|
"attention": ([
|
|
"flash_attention_2",
|
|
"sdpa",
|
|
], ),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("QWEN_MODEL",)
|
|
RETURN_NAMES = ("qwen_model",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "Qwen2.5-VL"
|
|
|
|
def load(self, model_name: str, device: str, precision: str, attention: str):
|
|
model_dir = os.path.join(folder_paths.models_dir, "Qwen", model_name.replace("/", "_"))
|
|
# Always attempt download with resume enabled so an interrupted download
|
|
# can be continued when the node is executed again.
|
|
snapshot_download(
|
|
repo_id=model_name,
|
|
local_dir=model_dir,
|
|
local_dir_use_symlinks=False,
|
|
resume_download=True,
|
|
)
|
|
if device == "auto":
|
|
device_map = "auto"
|
|
elif device == "cpu":
|
|
device_map = {"": "cpu"}
|
|
else:
|
|
device_map = {"": device}
|
|
|
|
precision = precision.upper()
|
|
dtype_map = {
|
|
"BF16": torch.bfloat16,
|
|
"FP16": torch.float16,
|
|
"FP32": torch.float32,
|
|
}
|
|
torch_dtype = dtype_map.get(precision, torch.bfloat16)
|
|
quant_config = None
|
|
if precision == "INT4":
|
|
quant_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True)
|
|
elif precision == "INT8":
|
|
quant_config = BitsAndBytesConfig(load_in_8bit=True)
|
|
|
|
attn_impl = attention
|
|
if precision == "FP32" and attn_impl == "flash_attention_2":
|
|
# FlashAttention doesn't support fp32. Fall back to SDPA.
|
|
attn_impl = "sdpa"
|
|
|
|
try:
|
|
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
|
model_dir,
|
|
torch_dtype=torch_dtype,
|
|
quantization_config=quant_config,
|
|
device_map=device_map,
|
|
attn_implementation=attn_impl,
|
|
)
|
|
except OSError:
|
|
# If loading fails due to missing or corrupt files, force a
|
|
# re-download and try again.
|
|
snapshot_download(
|
|
repo_id=model_name,
|
|
local_dir=model_dir,
|
|
local_dir_use_symlinks=False,
|
|
resume_download=True,
|
|
force_download=True,
|
|
)
|
|
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
|
model_dir,
|
|
torch_dtype=torch_dtype,
|
|
quantization_config=quant_config,
|
|
device_map=device_map,
|
|
attn_implementation=attn_impl,
|
|
)
|
|
except Exception:
|
|
# Surface any other errors (e.g. missing flash_attn)
|
|
raise
|
|
processor = AutoProcessor.from_pretrained(model_dir)
|
|
return (QwenModel(model=model, processor=processor, device=device),)
|
|
|
|
|
|
class QwenVLDetection:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"qwen_model": ("QWEN_MODEL",),
|
|
"image": ("IMAGE",),
|
|
"target": ("STRING", {"default": "object"}),
|
|
"bbox_selection": ("STRING", {"default": "all"}),
|
|
"score_threshold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"merge_boxes": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("JSON", "BBOX")
|
|
RETURN_NAMES = ("text", "bboxes")
|
|
FUNCTION = "detect"
|
|
CATEGORY = "Qwen2.5-VL"
|
|
|
|
def detect(
|
|
self,
|
|
qwen_model: QwenModel,
|
|
image,
|
|
target: str,
|
|
bbox_selection: str = "all",
|
|
score_threshold: float = 0.0,
|
|
merge_boxes: bool = False,
|
|
):
|
|
"""Generate bounding boxes for ``target`` within ``image``."""
|
|
model = qwen_model.model
|
|
processor = qwen_model.processor
|
|
device = qwen_model.device
|
|
if device == "auto":
|
|
device = str(next(model.parameters()).device)
|
|
if device.startswith("cuda") and torch.cuda.is_available():
|
|
try:
|
|
torch.cuda.set_device(int(device.split(":")[1]))
|
|
except Exception:
|
|
pass
|
|
|
|
prompt = f"Locate the {target} and output bbox in JSON"
|
|
|
|
if isinstance(image, torch.Tensor):
|
|
image = (image.squeeze().clamp(0, 1) * 255).to(torch.uint8).cpu().numpy()
|
|
image = Image.fromarray(image)
|
|
elif not isinstance(image, Image.Image):
|
|
raise ValueError("Unsupported image type")
|
|
|
|
messages = [
|
|
{"role": "system", "content": "You are a helpful assistant"},
|
|
{"role": "user", "content": [{"type": "text", "text": prompt}, {"image": image}]},
|
|
]
|
|
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
|
inputs = processor(text=[text], images=[image], return_tensors="pt", padding=True)
|
|
# Avoid keeping metadata tensors such as ``image_grid_thw`` on the GPU to
|
|
# prevent CUDA assertion failures on newer PyTorch builds when converting
|
|
# them to Python scalars for box rescaling. Fall back to pixel values or
|
|
# raw image dimensions if the grid is missing or produces invalid sizes.
|
|
input_h = None
|
|
input_w = None
|
|
image_grid = inputs.get("image_grid_thw")
|
|
if image_grid is not None:
|
|
image_grid = image_grid.detach().cpu()
|
|
if image_grid.numel() >= 3:
|
|
h_val = int(image_grid[0][1].item() * 14)
|
|
w_val = int(image_grid[0][2].item() * 14)
|
|
if h_val > 0 and w_val > 0:
|
|
input_h, input_w = h_val, w_val
|
|
|
|
if input_h is None or input_w is None:
|
|
pixel_values = inputs.get("pixel_values")
|
|
if pixel_values is not None:
|
|
input_h = int(pixel_values.shape[-2])
|
|
input_w = int(pixel_values.shape[-1])
|
|
else:
|
|
input_w, input_h = image.size
|
|
|
|
input_h = max(1, input_h)
|
|
input_w = max(1, input_w)
|
|
|
|
inputs = inputs.to(device)
|
|
with torch.no_grad():
|
|
output_ids = model.generate(**inputs, max_new_tokens=1024)
|
|
gen_ids = [output_ids[len(inp):] for inp, output_ids in zip(inputs.input_ids, output_ids)]
|
|
output_text = processor.batch_decode(
|
|
gen_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
|
|
)[0]
|
|
items = parse_boxes(
|
|
output_text,
|
|
image.width,
|
|
image.height,
|
|
input_w,
|
|
input_h,
|
|
score_threshold,
|
|
)
|
|
|
|
selection = bbox_selection.strip().lower()
|
|
boxes = items
|
|
if selection != "all" and selection:
|
|
idxs = []
|
|
for part in selection.replace(",", " ").split():
|
|
try:
|
|
idxs.append(int(part))
|
|
except Exception:
|
|
continue
|
|
boxes = [boxes[i] for i in idxs if 0 <= i < len(boxes)]
|
|
|
|
if merge_boxes and boxes:
|
|
x1 = min(b["bbox"][0] for b in boxes)
|
|
y1 = min(b["bbox"][1] for b in boxes)
|
|
x2 = max(b["bbox"][2] for b in boxes)
|
|
y2 = max(b["bbox"][3] for b in boxes)
|
|
score = max(b["score"] for b in boxes)
|
|
label = boxes[0].get("label", target)
|
|
boxes = [{"bbox": [x1, y1, x2, y2], "score": score, "label": label}]
|
|
|
|
json_boxes = [
|
|
{"bbox_2d": b["bbox"], "label": b.get("label", target)} for b in boxes
|
|
]
|
|
json_output = json.dumps(json_boxes, ensure_ascii=False)
|
|
bboxes_only = [b["bbox"] for b in boxes]
|
|
return (json_output, bboxes_only)
|
|
|
|
|
|
class BBoxesToSAM2:
|
|
"""Convert a list of bounding boxes to the format expected by SAM2 nodes."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"bboxes": ("BBOX",)}}
|
|
|
|
RETURN_TYPES = ("BBOXES",)
|
|
RETURN_NAMES = ("sam2_bboxes",)
|
|
FUNCTION = "convert"
|
|
CATEGORY = "Qwen2.5-VL"
|
|
|
|
def convert(self, bboxes):
|
|
if not isinstance(bboxes, list):
|
|
raise ValueError("bboxes must be a list")
|
|
|
|
# If already batched, return as-is
|
|
if bboxes and isinstance(bboxes[0], (list, tuple)) and bboxes[0] and isinstance(bboxes[0][0], (list, tuple)):
|
|
return (bboxes,)
|
|
|
|
return ([bboxes],)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DownloadAndLoadQwenModel": DownloadAndLoadQwenModel,
|
|
"QwenVLDetection": QwenVLDetection,
|
|
"BBoxesToSAM2": BBoxesToSAM2,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DownloadAndLoadQwenModel": "Download and Load Qwen2.5-VL Model",
|
|
"QwenVLDetection": "Qwen2.5-VL Object Detection",
|
|
"BBoxesToSAM2": "Prepare BBoxes for SAM2",
|
|
}
|