#@title Setup
#!pip3 install ftfy regex tqdm transformers==4.15.0 timm==0.4.12 fairscale==0.4.4
#!pip3 install git+https://github.com/openai/CLIP.git
#!git clone https://github.com/pharmapsychotic/clip-interrogator.git
#!git clone https://github.com/salesforce/BLIP
#%cd /content/BLIP

import clip
import gc
#import numpy as np
import os
import pandas as pd
import requests
import torch
#import torchvision.transforms as T
#import torchvision.transforms.functional as TF

from IPython.display import display
from PIL import Image
#from torch import nn
#from torch.nn import functional as F
from torchvision import transforms
from torchvision.transforms.functional import InterpolationMode
from ldm.models.blip import blip_decoder

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

blip_image_eval_size = 384
blip_model_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model*_base_caption.pth'        
blip_model = blip_decoder(pretrained=blip_model_url, image_size=blip_image_eval_size, vit='base')
blip_model.eval()
blip_model = blip_model.to(device)

def generate_caption(pil_image):
	gpu_image = transforms.Compose([
	    transforms.Resize((blip_image_eval_size, blip_image_eval_size), interpolation=InterpolationMode.BICUBIC),
	    transforms.ToTensor(),
	    transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
	    ])(image).unsqueeze(0).to(device)

	with torch.no_grad():
		caption = blip_model.generate(gpu_image, sample=False, num_beams=3, max_length=20, min_length=5)
	return caption[0]

def load_list(filename):
	with open(filename, 'r', encoding='utf-8', errors='replace') as f:
		items = [line.strip() for line in f.readlines()]
	return items

def rank(model, image_features, text_array, top_count=1):
	top_count = min(top_count, len(text_array))
	text_tokens = clip.tokenize([text for text in text_array]).cuda()
	with torch.no_grad():
		text_features = model.encode_text(text_tokens).float()
	text_features /= text_features.norm(dim=-1, keepdim=True)

	similarity = torch.zeros((1, len(text_array))).to(device)
	for i in range(image_features.shape[0]):
		similarity += (100.0 * image_features[i].unsqueeze(0) @ text_features.T).softmax(dim=-1)
	similarity /= image_features.shape[0]

	top_probs, top_labels = similarity.cpu().topk(top_count, dim=-1)  
	return [(text_array[top_labels[0][i].numpy()], (top_probs[0][i].numpy()*100)) for i in range(top_count)]

def interrogate(image, models):
	caption = generate_caption(image)
	if len(models) == 0:
		print(f"\n\n{caption}")
		return

	table = []
	bests = [[('',0)]]*5
	for model_name in models:
		print(f"Interrogating with {model_name}...")
		model, preprocess = clip.load(model_name)
		model.cuda().eval()

		images = preprocess(image).unsqueeze(0).cuda()
		with torch.no_grad():
			image_features = model.encode_image(images).float()
		image_features /= image_features.norm(dim=-1, keepdim=True)

		ranks = [
		    rank(model, image_features, mediums),
		    rank(model, image_features, ["by "+artist for artist in artists]),
		    rank(model, image_features, trending_list),
		    rank(model, image_features, movements),
		    rank(model, image_features, flavors, top_count=3)
		]

		for i in range(len(ranks)):
			confidence_sum = 0
			for ci in range(len(ranks[i])):
				confidence_sum += ranks[i][ci][1]
			if confidence_sum > sum(bests[i][t][1] for t in range(len(bests[i]))):
				bests[i] = ranks[i]

		row = [model_name]
		for r in ranks:
			row.append(', '.join([f"{x[0]} ({x[1]:0.1f}%)" for x in r]))

		table.append(row)

		del model
		gc.collect()
	display(pd.DataFrame(table, columns=["Model", "Medium", "Artist", "Trending", "Movement", "Flavors"]))

	flaves = ', '.join([f"{x[0]}" for x in bests[4]])
	medium = bests[0][0][0]
	if caption.startswith(medium):
		print(f"\n\n{caption} {bests[1][0][0]}, {bests[2][0][0]}, {bests[3][0][0]}, {flaves}")
	else:
		print(f"\n\n{caption}, {medium} {bests[1][0][0]}, {bests[2][0][0]}, {bests[3][0][0]}, {flaves}")

data_path = "../clip-interrogator/data/"

artists = load_list(os.path.join(data_path, 'artists.txt'))
flavors = load_list(os.path.join(data_path, 'flavors.txt'))
mediums = load_list(os.path.join(data_path, 'mediums.txt'))
movements = load_list(os.path.join(data_path, 'movements.txt'))

sites = ['Artstation', 'behance', 'cg society', 'cgsociety', 'deviantart', 'dribble', 'flickr', 'instagram', 'pexels', 'pinterest', 'pixabay', 'pixiv', 'polycount', 'reddit', 'shutterstock', 'tumblr', 'unsplash', 'zbrush central']
trending_list = [site for site in sites]
trending_list.extend(["trending on "+site for site in sites])
trending_list.extend(["featured on "+site for site in sites])
trending_list.extend([site+" contest winner" for site in sites])

#@title Interrogate!

#@markdown 

#@markdown #####**Image:**

image_path_or_url = "https://i.redd.it/e2e8gimigjq91.jpg" #@param {type:"string"}

#@markdown 

#@markdown #####**CLIP models:**

#@markdown For [StableDiffusion](https://stability.ai/blog/stable-diffusion-announcement) you can just use ViTL14<br>
#@markdown For [DiscoDiffusion](https://colab.research.google.com/github/alembics/disco-diffusion/blob/main/Disco_Diffusion.ipynb) and 
#@markdown [JAX](https://colab.research.google.com/github/huemin-art/jax-guided-diffusion/blob/v2.7/Huemin_Jax_Diffusion_2_7.ipynb) enable all the same models here as you intend to use when generating your images

ViTB32 = True #@param{type:"boolean"}
ViTB16 = True #@param{type:"boolean"}
ViTL14 = False #@param{type:"boolean"}
ViTL14_336px = False #@param{type:"boolean"}
RN101 = False #@param{type:"boolean"}
RN50 = True #@param{type:"boolean"}
RN50x4 = False #@param{type:"boolean"}
RN50x16 = False #@param{type:"boolean"}
RN50x64 = False #@param{type:"boolean"}

models = []
if ViTB32: models.append('ViT-B/32')
if ViTB16: models.append('ViT-B/16')
if ViTL14: models.append('ViT-L/14')
if ViTL14_336px: models.append('ViT-L/14@336px')
if RN101: models.append('RN101')
if RN50: models.append('RN50')
if RN50x4: models.append('RN50x4')
if RN50x16: models.append('RN50x16')
if RN50x64: models.append('RN50x64')

if str(image_path_or_url).startswith('http://') or str(image_path_or_url).startswith('https://'):
	image = Image.open(requests.get(image_path_or_url, stream=True).raw).convert('RGB')
else:
	image = Image.open(image_path_or_url).convert('RGB')

thumb = image.copy()
thumb.thumbnail([blip_image_eval_size, blip_image_eval_size])
display(thumb)

interrogate(image, models=models)