import torch
from PIL import Image, ImageDraw
import transformers
from hfmodel.modeling_clip import CLIPModel
from transformers import CLIPProcessor
from torchvision.transforms import ToTensor
import torch.nn.functional as F

model = CLIPModel.from_pretrained("zer0int/CLIP-Registers-Gated_MLP-ViT-L-14")
processor = CLIPProcessor.from_pretrained("zer0int/CLIP-Registers-Gated_MLP-ViT-L-14")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

size = 224
im = Image.new("RGB", (size, size), (255, 255, 255))
draw = ImageDraw.Draw(im)

# --------- GPT-4.1's idea of a pineapple. Need an input image... ---------
body_bbox = [size*0.28, size*0.38, size*0.72, size*0.90]
draw.ellipse(body_bbox, fill=(254, 221, 72), outline=(180, 120, 0), width=5)
eye_color = (198, 134, 66)
for row in range(4):
    for col in range(3):
        ex = size*0.36 + col*size*0.09 + (row%2)*size*0.045
        ey = size*0.50 + row*size*0.085
        ew, eh = size*0.035, size*0.025
        draw.ellipse([ex-ew, ey-eh, ex+ew, ey+eh], fill=eye_color, outline=None)
leaf_color = (61, 179, 70)
leaf_base_x = size/2
leaf_base_y = size*0.38
for i, (angle, length) in enumerate([(-28, 65), (-12, 70), (0, 80), (12, 70), (28, 65)]):
    from math import radians, cos, sin
    a = radians(angle)
    tip_x = leaf_base_x + length*sin(a)
    tip_y = leaf_base_y - length*cos(a)
    left = (leaf_base_x + 13*cos(a+1.5), leaf_base_y + 13*sin(a+1.5))
    right = (leaf_base_x + 13*cos(a-1.5), leaf_base_y + 13*sin(a-1.5))
    draw.polygon([left, (tip_x, tip_y), right], fill=leaf_color)
im.save("pineapple.png")
# ---------

image = Image.open("pineapple.png").convert("RGB")
texts = ["pine", "apple", "pineapple", "orange", "pear", "person", "cat", "dog"]

inputs = processor(text=texts, images=image, return_tensors="pt", padding=True).to(device)

with torch.no_grad():
    outputs = model(**inputs)
    image_embeds = outputs.image_embeds
    text_embeds = outputs.text_embeds

image_embeds = F.normalize(image_embeds, dim=-1)
text_embeds = F.normalize(text_embeds, dim=-1)
cos_sim = image_embeds @ text_embeds.T
cos_sim = cos_sim.squeeze(0)

for text, sim in zip(texts, cos_sim):
    print(f"Similarity with '{text}': {sim.item():.4f}")