|
|
|
import gradio.components as gc |
|
import gradio as gr |
|
|
|
import numpy as np |
|
import pandas as pd |
|
import torch |
|
from PIL import Image |
|
from transformers import CLIPModel, CLIPProcessor |
|
device = 'cpu' |
|
torch.no_grad().__enter__() |
|
torch.autocast('cuda').__enter__() |
|
|
|
|
|
|
|
t = pd.read_pickle("clip_texts_1_fp16.pkl") |
|
words = t.reset_index().word |
|
wordsv = torch.tensor(t.values).to(device) |
|
|
|
|
|
|
|
|
|
model_name = "openai/clip-vit-large-patch14" |
|
mmm = CLIPModel.from_pretrained(model_name) |
|
mmm.eval() |
|
mmm.to(device) |
|
|
|
processor = CLIPProcessor.from_pretrained(model_name) |
|
|
|
|
|
|
|
|
|
def slerp(t, v0, v1, DOT_THRESHOLD=0.9995): |
|
""" helper function to spherically interpolate two arrays v1 v2 """ |
|
inputs_are_torch = False |
|
if not isinstance(v0, np.ndarray): |
|
inputs_are_torch = True |
|
input_device = v0.device |
|
v0 = v0.cpu().numpy() |
|
v1 = v1.cpu().numpy() |
|
|
|
dot = np.sum(v0 * v1 / (np.linalg.norm(v0) * np.linalg.norm(v1))) |
|
if np.abs(dot) > DOT_THRESHOLD: |
|
v2 = (1 - t) * v0 + t * v1 |
|
else: |
|
theta_0 = np.arccos(dot) |
|
sin_theta_0 = np.sin(theta_0) |
|
theta_t = theta_0 * t |
|
sin_theta_t = np.sin(theta_t) |
|
s0 = np.sin(theta_0 - theta_t) / sin_theta_0 |
|
s1 = sin_theta_t / sin_theta_0 |
|
v2 = s0 * v0 + s1 * v1 |
|
|
|
if inputs_are_torch: |
|
v2 = torch.from_numpy(v2).to(input_device) |
|
|
|
return v2 |
|
|
|
|
|
def query(text: str, img: Image.Image, limit: int, score_threshold: float, slerp_degree: float): |
|
if text != '': |
|
inp = processor(text=text, return_tensors='pt').to(device) |
|
rout = mmm.get_text_features(**inp) |
|
tout = rout.detach().cpu().numpy()[0] |
|
out = tout |
|
|
|
if img is not None: |
|
inp = processor(images=[img], return_tensors="pt",).to(device) |
|
rout = mmm.get_image_features(**inp) |
|
iout = rout.detach().cpu().numpy()[0] |
|
out = iout |
|
|
|
if text != '' and img is not None: |
|
out = slerp(slerp_degree, tout, iout) |
|
|
|
if out is not None: |
|
|
|
scores = np.dot(out, wordsv.T) |
|
|
|
topk = ( |
|
pd.concat( |
|
[words, pd.Series(scores, name='score')], |
|
axis=1 |
|
) |
|
.sort_values('score', ascending=False) |
|
.query(f'score > {score_threshold}') |
|
.head(limit) |
|
) |
|
|
|
topwords = "\n".join( |
|
f'{word}: {score:.2f} ' |
|
for _, word, score in topk.itertuples() |
|
) |
|
|
|
return topwords |
|
|
|
|
|
searchtext = gc.Textbox(lines=2, placeholder="Search text") |
|
searchimage = gc.Image(shape=(224, 224), label="Search image", type='pil') |
|
inp_limit = gc.Slider(1, 50, 10, step=1, label='Limit') |
|
score_threshold = gc.Slider(0, 30, 0, step=.5, label='Score threshold') |
|
slerp_degree = gc.Slider( |
|
0, 1, 0.5, step=.01, label='Slerp degree (if both text and image are provided)\nFinds a midpoint between image and text embeddings') |
|
|
|
|
|
dsurl = 'https://www.kaggle.com/datasets/yk1598/479k-english-words' |
|
gr.Interface( |
|
query, |
|
[searchtext, searchimage, inp_limit, score_threshold, slerp_degree], |
|
[gc.Textbox(label='Top words')], |
|
title="Initial Token Finder for Textual Inversion", |
|
description=f"find the closest single token word for a given text and/or image.\nbased on {model_name}.\n\nData: {dsurl}", |
|
analytics_enabled=False, |
|
allow_flagging='never', |
|
).launch() |
|
|
