113 lines
3.3 KiB
Python
113 lines
3.3 KiB
Python
"""
|
|
ComfyUI Custom Node: CLIP Score Calculator
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from transformers import CLIPProcessor, CLIPModel
|
|
|
|
|
|
class CLIPScoreNode:
|
|
"""Calculates CLIP cosine similarity with simple prompt chunking."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"model_name": (
|
|
[
|
|
"openai/clip-vit-base-patch16",
|
|
"openai/clip-vit-base-patch32",
|
|
"openai/clip-vit-large-patch14",
|
|
],
|
|
{"default": "openai/clip-vit-base-patch16"},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT", "STRING")
|
|
RETURN_NAMES = ("clip_score", "score_text")
|
|
FUNCTION = "calculate_clip_score"
|
|
CATEGORY = "metrics"
|
|
OUTPUT_NODE = True
|
|
|
|
def calculate_clip_score(self, image, prompt, model_name):
|
|
try:
|
|
# Load CLIP
|
|
model = CLIPModel.from_pretrained(model_name)
|
|
processor = CLIPProcessor.from_pretrained(model_name)
|
|
|
|
# Handle ComfyUI IMAGE: [B,H,W,C] or [H,W,C]
|
|
if len(image.shape) == 4:
|
|
img = image[0]
|
|
else:
|
|
img = image
|
|
|
|
# Convert to PIL
|
|
img_np = (img.cpu().numpy() * 255).astype(np.uint8)
|
|
if img_np.shape[-1] == 4:
|
|
img_np = img_np[:, :, :3]
|
|
pil_image = Image.fromarray(img_np, mode="RGB")
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
model.to(device)
|
|
|
|
chunks = [prompt[i:i + 200] for i in range(0, len(prompt), 200)]
|
|
scores = []
|
|
|
|
with torch.no_grad():
|
|
for chunk in chunks:
|
|
inputs = processor(
|
|
text=[chunk],
|
|
images=pil_image,
|
|
return_tensors="pt",
|
|
padding=True,
|
|
truncation=True,
|
|
max_length=77,
|
|
)
|
|
|
|
inputs = {k: v.to(device) for k, v in inputs.items()}
|
|
|
|
outputs = model(**inputs)
|
|
|
|
image_embeds = outputs.image_embeds
|
|
text_embeds = outputs.text_embeds
|
|
|
|
# Normalize
|
|
image_embeds = image_embeds / (
|
|
image_embeds.norm(p=2, dim=-1, keepdim=True) + 1e-8
|
|
)
|
|
text_embeds = text_embeds / (
|
|
text_embeds.norm(p=2, dim=-1, keepdim=True) + 1e-8
|
|
)
|
|
|
|
score = (image_embeds @ text_embeds.T)[0, 0].item()
|
|
scores.append(score)
|
|
|
|
final_score = sum(scores) / len(scores) if scores else 0.0
|
|
|
|
# Cleanup
|
|
model.cpu()
|
|
del model
|
|
|
|
except Exception as e:
|
|
final_score = 0.0
|
|
print(f"CLIP error: {str(e)}")
|
|
|
|
score_str = f"CLIP Score: {final_score:.4f}"
|
|
print(score_str)
|
|
|
|
return (final_score, score_str)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"CLIPScoreNode": CLIPScoreNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"CLIPScoreNode": "CLIP Score Calculator",
|
|
}
|