91 lines
3.0 KiB
Python
91 lines
3.0 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 true CLIP cosine similarity between prompt and image."""
|
|
|
|
@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 (cached after first run)
|
|
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] # First from batch
|
|
else:
|
|
img = image
|
|
|
|
# [H,W,C] float32 [0,1] → numpy uint8 → PIL RGB
|
|
img_np = (img.cpu().numpy() * 255).astype(np.uint8)
|
|
if img_np.shape[-1] == 4: # RGBA → RGB
|
|
img_np = img_np[:, :, :3]
|
|
pil_image = Image.fromarray(img_np, mode='RGB')
|
|
|
|
# GPU inference
|
|
device = next(model.parameters()).device
|
|
inputs = processor(text=[prompt], images=pil_image, return_tensors="pt", padding=True)
|
|
inputs = {k: v.to(device) for k, v in inputs.items()}
|
|
model.to(device)
|
|
|
|
# Raw CLIP score: cosine similarity of normalized embeddings
|
|
with torch.no_grad():
|
|
outputs = model(**inputs)
|
|
image_embeds = outputs.image_embeds
|
|
text_embeds = outputs.text_embeds
|
|
|
|
# L2 normalize + cosine similarity (standard CLIP score)
|
|
image_embeds /= image_embeds.norm(p=2, dim=-1, keepdim=True) + 1e-8
|
|
text_embeds /= text_embeds.norm(p=2, dim=-1, keepdim=True) + 1e-8
|
|
score = (image_embeds @ text_embeds.T)[0, 0].item()
|
|
|
|
# Cleanup
|
|
model.cpu()
|
|
del model, inputs
|
|
|
|
except Exception as e:
|
|
score = 0.0
|
|
print(f"CLIP error: {str(e)}")
|
|
|
|
score_str = f"CLIP Score: {score:.4f}"
|
|
print(score_str)
|
|
|
|
return (score, score_str)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"CLIPScoreNode": CLIPScoreNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"CLIPScoreNode": "CLIP Score Calculator",
|
|
} |