137 lines
5.2 KiB
Python
137 lines
5.2 KiB
Python
import os
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from pycocotools import mask as mask_utils
|
|
import folder_paths
|
|
import comfy.model_management as mm
|
|
|
|
from .model_loader import load_falcon_perception_local
|
|
|
|
class LoadFalconPerception:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
# We can look for .safetensors files in checkpoints or let user specify path
|
|
return {"required": {
|
|
"model_name_or_path": ("STRING", {"default": "falcon_perception/model.safetensors"}),
|
|
"device": (["auto", "cuda", "cpu"], {"default": "auto"}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("FALCON_PERCEPTION_MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "FalconPerception"
|
|
|
|
def load_model(self, model_name_or_path, device):
|
|
# Resolve full path
|
|
full_path = os.path.join(folder_paths.models_dir, model_name_or_path)
|
|
if not os.path.exists(full_path):
|
|
if os.path.exists(model_name_or_path):
|
|
full_path = model_name_or_path
|
|
else:
|
|
raise FileNotFoundError(f"Model weights not found at {full_path} or {model_name_or_path}")
|
|
|
|
if device == "auto":
|
|
device_map = "auto"
|
|
elif device == "cuda":
|
|
device_map = {"": mm.get_torch_device()}
|
|
else:
|
|
device_map = {"": torch.device("cpu")}
|
|
|
|
print(f"[Falcon Perception] Loading model from {full_path} on {device_map}...")
|
|
|
|
# Use our custom loader to bypass HF cache and handle PyTorch 2.7.0 compatibility
|
|
model = load_falcon_perception_local(full_path, device_map)
|
|
|
|
return (model,)
|
|
|
|
class FalconPerceptionGenerate:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model": ("FALCON_PERCEPTION_MODEL",),
|
|
"image": ("IMAGE",),
|
|
"query": ("STRING", {"multiline": True, "default": "cat"}),
|
|
"max_new_tokens": ("INT", {"default": 2048, "min": 1, "max": 8192}),
|
|
"min_dimension": ("INT", {"default": 256, "min": 64, "max": 2048}),
|
|
"max_dimension": ("INT", {"default": 1024, "min": 256, "max": 4096}),
|
|
"compile_model": ("BOOLEAN", {"default": False}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("MASK", "MASK", "STRING")
|
|
RETURN_NAMES = ("combined_mask", "individual_masks", "info")
|
|
FUNCTION = "generate"
|
|
CATEGORY = "FalconPerception"
|
|
|
|
def generate(self, model, image, query, max_new_tokens, min_dimension, max_dimension, compile_model):
|
|
# Convert ComfyUI IMAGE [B, H, W, C] to list of PIL Images
|
|
batch_size = image.shape[0]
|
|
pil_images = []
|
|
for i in range(batch_size):
|
|
img_np = 255. * image[i].cpu().numpy()
|
|
pil_img = Image.fromarray(np.clip(img_np, 0, 255).astype(np.uint8))
|
|
pil_images.append(pil_img)
|
|
|
|
queries = [query] * batch_size
|
|
|
|
print(f"[Falcon Perception] Generating for batch size {batch_size} with query '{query}'...")
|
|
|
|
with torch.no_grad():
|
|
preds_batch = model.generate(
|
|
pil_images,
|
|
queries,
|
|
max_new_tokens=max_new_tokens,
|
|
min_dimension=min_dimension,
|
|
max_dimension=max_dimension,
|
|
compile=compile_model
|
|
)
|
|
|
|
combined_masks = []
|
|
individual_masks = []
|
|
info_strs = []
|
|
|
|
for b, preds in enumerate(preds_batch):
|
|
img_h, img_w = pil_images[b].height, pil_images[b].width
|
|
|
|
if len(preds) == 0:
|
|
combined_masks.append(torch.zeros((img_h, img_w), dtype=torch.float32))
|
|
info_strs.append(f"Image {b}: No instances found.")
|
|
continue
|
|
|
|
comb_m = np.zeros((img_h, img_w), dtype=bool)
|
|
info = [f"Image {b}: Found {len(preds)} instances"]
|
|
|
|
for i, p in enumerate(preds):
|
|
rle = p["mask_rle"]
|
|
m = {"size": rle["size"], "counts": rle["counts"].encode("utf-8")}
|
|
mask = mask_utils.decode(m).astype(bool)
|
|
|
|
comb_m = comb_m | mask
|
|
individual_masks.append(torch.from_numpy(mask.astype(np.float32)))
|
|
info.append(f" Instance {i}: xy={p['xy']}, hw={p['hw']}")
|
|
|
|
combined_masks.append(torch.from_numpy(comb_m.astype(np.float32)))
|
|
info_strs.append("\n".join(info))
|
|
|
|
combined_tensor = torch.stack(combined_masks, dim=0)
|
|
|
|
if len(individual_masks) > 0:
|
|
# individual masks are stacked into a batch [N, H, W] where N is total instances across batch
|
|
individual_tensor = torch.stack(individual_masks, dim=0)
|
|
else:
|
|
individual_tensor = torch.zeros((1, pil_images[0].height, pil_images[0].width), dtype=torch.float32)
|
|
|
|
final_info = "\n\n".join(info_strs)
|
|
|
|
return (combined_tensor, individual_tensor, final_info)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LoadFalconPerception": LoadFalconPerception,
|
|
"FalconPerceptionGenerate": FalconPerceptionGenerate
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LoadFalconPerception": "Load Falcon Perception Model",
|
|
"FalconPerceptionGenerate": "Falcon Perception Generate"
|
|
}
|