app.py

import gradio as gr
from datasets import load_dataset
import random
import numpy as np
from transformers import CLIPProcessor, CLIPModel
from os import environ
import clip
import pickle
import requests
import torch
import os
from huggingface_hub import hf_hub_download
from torch import nn
import torch.nn.functional as nnf
import sys
from typing import Tuple, List, Union, Optional
from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup


N = type(None)
V = np.array
ARRAY = np.ndarray
ARRAYS = Union[Tuple[ARRAY, ...], List[ARRAY]]
VS = Union[Tuple[V, ...], List[V]]
VN = Union[V, N]
VNS = Union[VS, N]
T = torch.Tensor
TS = Union[Tuple[T, ...], List[T]]
TN = Optional[T]
TNS = Union[Tuple[TN, ...], List[TN]]
TSN = Optional[TS]
TA = Union[T, ARRAY]

D = torch.device
CPU = torch.device('cpu')

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

# # Load the pre-trained model and processor
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

clip_model = clip_model.to(device)
#orig_clip_model, orig_clip_processor = clip.load("ViT-B/32", device=device, jit=False)


# Load the Unsplash dataset
dataset = load_dataset("jamescalam/unsplash-25k-photos", split="train")  # all 25K images are in train split
dataset_size = len(dataset)

# Load gpt and modifed weights for captions
gpt = GPT2LMHeadModel.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
conceptual_weight = hf_hub_download(repo_id="akhaliq/CLIP-prefix-captioning-conceptual-weights", filename="conceptual_weights.pt")
coco_weight = hf_hub_download(repo_id="akhaliq/CLIP-prefix-captioning-COCO-weights", filename="coco_weights.pt")


height = 256   # height for resizing images

def predict(image, labels):
    with torch.no_grad():
        inputs = clip_processor(text=[f"a photo of {c}" for c in labels], images=image, return_tensors="pt", padding=True).to(device)
        outputs = clip_model(**inputs)
        logits_per_image = outputs.logits_per_image # this is the image-text similarity score
        probs = logits_per_image.softmax(dim=1).cpu().numpy() # we can take the softmax to get the label probabilities
    return {k: float(v) for k, v in zip(labels, probs[0])}


# def predict2(image, labels):
#     image = orig_clip_processor(image).unsqueeze(0).to(device)     
#     text = clip.tokenize(labels).to(device)
#     with torch.no_grad():
#         image_features = orig_clip_model.encode_image(image)
#         text_features = orig_clip_model.encode_text(text)
#         logits_per_image, logits_per_text = orig_clip_model(image, text)
#         probs = logits_per_image.softmax(dim=-1).cpu().numpy()
#         return {k: float(v) for k, v in zip(labels, probs[0])}

def rand_image():
    n = dataset.num_rows
    r = random.randrange(0,n)
    return dataset[r]["photo_image_url"] + f"?h={height}"  # Unsplash allows dynamic requests, including size of image

def set_labels(text):
    return text.split(",")
    
# get_caption = gr.load("ryaalbr/caption", src="spaces", hf_token=environ["api_key"])
# def generate_text(image, model_name):
#     return get_caption(image, model_name)


class MLP(nn.Module):

    def forward(self, x: T) -> T:
        return self.model(x)

    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
        super(MLP, self).__init__()
        layers = []
        for i in range(len(sizes) -1):
            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
            if i < len(sizes) - 2:
                layers.append(act())
        self.model = nn.Sequential(*layers)


class ClipCaptionModel(nn.Module):

    def get_dummy_token(self, batch_size: int, device: D) -> T:
        return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64, device=device)

    def forward(self, tokens: T, prefix: T, mask: Optional[T] = None, labels: Optional[T] = None):
        embedding_text = self.gpt.transformer.wte(tokens)
        prefix_projections = self.clip_project(prefix).view(-1, self.prefix_length, self.gpt_embedding_size)
        embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1)
        if labels is not None:
            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)
            labels = torch.cat((dummy_token, tokens), dim=1)
        out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask)
        return out

    def __init__(self, prefix_length: int, prefix_size: int = 512):
        super(ClipCaptionModel, self).__init__()
        self.prefix_length = prefix_length
        self.gpt = gpt
        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
        if prefix_length > 10:  # not enough memory
            self.clip_project = nn.Linear(prefix_size, self.gpt_embedding_size * prefix_length)
        else:
            self.clip_project = MLP((prefix_size, (self.gpt_embedding_size * prefix_length) // 2, self.gpt_embedding_size * prefix_length))

#clip_model, preprocess = clip.load("ViT-B/32", device=device, jit=False)


def get_caption(img,model_name):   
    prefix_length = 10
    
    model = ClipCaptionModel(prefix_length)
    
    if model_name == "COCO":
      model_path = coco_weight
    else:
      model_path = conceptual_weight
    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
    model = model.eval() 
    model = model.to(device)

    

    input = clip_processor(images=img, return_tensors="pt").to(device)
    with torch.no_grad():
        prefix = clip_model.get_image_features(**input)    

    # image = preprocess(img).unsqueeze(0).to(device)
    # with torch.no_grad():
    #     prefix = clip_model.encode_image(image).to(device, dtype=torch.float32)
        prefix_embed = model.clip_project(prefix).reshape(1, prefix_length, -1)
        output = model.gpt.generate(inputs_embeds=prefix_embed,
                                  num_beams=1,
                                  do_sample=False,
                                  num_return_sequences=1,
                                  no_repeat_ngram_size=1,
                                  max_new_tokens = 67,
                                  pad_token_id = tokenizer.eos_token_id,
                                  eos_token_id = tokenizer.encode('.')[0],
                                  renormalize_logits = True)
        generated_text_prefix = tokenizer.decode(output[0], skip_special_tokens=True)
    return generated_text_prefix[:-1] if generated_text_prefix[-1] == "." else generated_text_prefix  #remove period at end if present

    
    # get_images = gr.load("ryaalbr/ImageSearch", src="spaces", hf_token=environ["api_key"])
# def search_images(text):
#     return get_images(text, api_name="images")

emb_filename = 'unsplash-25k-photos-embeddings-indexes.pkl'
with open(emb_filename, 'rb') as emb:
        id2url, img_names, img_emb = pickle.load(emb)


def search(search_query):

    with torch.no_grad():

        # Encode and normalize the description using CLIP (HF CLIP)
        inputs = clip_processor(text=search_query, images=None, return_tensors="pt", padding=True).to(device)
        text_encoded = clip_model.get_text_features(**inputs)
        
        # # Encode and normalize the description using CLIP (original CLIP)
        # text_encoded = orig_clip_model.encode_text(clip.tokenize(search_query))
        # text_encoded /= text_encoded.norm(dim=-1, keepdim=True)


    # Retrieve the description vector
    text_features = text_encoded.cpu().numpy()

    # Compute the similarity between the descrption and each photo using the Cosine similarity
    similarities = (text_features @ img_emb.T).squeeze(0)

    # Sort the photos by their similarity score
    best_photos = similarities.argsort()[::-1]
    best_photos = best_photos[:15]
    #best_photos = sorted(zip(similarities, range(img_emb.shape[0])), key=lambda x: x[0], reverse=True)

    best_photo_ids = img_names[best_photos]

    imgs = []

    # Iterate over the top 5 results
    for id in best_photo_ids:

        id, _ = id.split('.')
        url = id2url.get(id, "")
        if url == "": continue

        img = url  + "?h=512"
       # r = requests.get(url + "?w=512", stream=True)
       # img = Image.open(r.raw)
        #credits = f'Photo by <a href="https://unsplash.com/@{photo["photographer_username"]}?utm_source=NaturalLanguageImageSearch&utm_medium=referral">{photo["photographer_first_name"]} {photo["photographer_last_name"]}</a> on <a href="https://unsplash.com/?utm_source=NaturalLanguageImageSearch&utm_medium=referral">Unsplash</a>'
        imgs.append(img)
        #display(HTML(f'Photo by <a href="https://unsplash.com/@{photo["photographer_username"]}?utm_source=NaturalLanguageImageSearch&utm_medium=referral">{photo["photographer_first_name"]} {photo["photographer_last_name"]}</a> on <a href="https://unsplash.com/?utm_source=NaturalLanguageImageSearch&utm_medium=referral">Unsplash</a>'))

        if len(imgs) == 5: break

    return imgs



with gr.Blocks() as demo:

    with gr.Tab("Classification"):
        labels = gr.State([])   # creates hidden component that can store a value and can be used as input/output; here, initial value is an empty list
        instructions = """## Instructions:
        1. Enter list of labels separated by commas (or select one of the examples below)
        2. Click **Get Random Image** to grab a random image from dataset
        3. Click **Classify Image** to analyze current image against the labels (including after changing labels)
        """
        gr.Markdown(instructions)
        with gr.Row(variant="compact"):
            label_text = gr.Textbox(show_label=False, placeholder="Enter classification labels").style(container=False)
            #submit_btn = gr.Button("Submit").style(full_width=False)
        gr.Examples(["spring, summer, fall, winter",
                     "mountain, city, beach, ocean, desert, forest, valley", 
                     "red, blue, green, white, black, purple, brown", 
                     "person, animal, landscape, something else",
                     "day, night, dawn, dusk"], inputs=label_text)
        with gr.Row():
            with gr.Column(variant="panel"):
                im = gr.Image(interactive=False).style(height=height)
                with gr.Row():
                    get_btn = gr.Button("Get Random Image").style(full_width=False)
                    class_btn = gr.Button("Classify Image").style(full_width=False)
            cf = gr.Label()
        #submit_btn.click(fn=set_labels, inputs=label_text)
        label_text.change(fn=set_labels, inputs=label_text, outputs=labels)  # parse list if changed
        label_text.blur(fn=set_labels, inputs=label_text, outputs=labels)  # parse list if focus is moved elsewhere; ensures that list is fully parsed before classification
        label_text.submit(fn=set_labels, inputs=label_text, outputs=labels)  # parse list if user hits enter; ensures that list is fully parsed before classification
        get_btn.click(fn=rand_image, outputs=im)
        #im.change(predict, inputs=[im, labels], outputs=cf)
        class_btn.click(predict, inputs=[im, labels], outputs=cf)
        gr.HTML(f"Dataset: <a href='https://github.com/unsplash/datasets' target='_blank'>Unsplash Lite</a>; Number of Images: {dataset_size}")
    

    with gr.Tab("Captioning"):
        instructions = """## Instructions:
        1. Click **Get Random Image** to grab a random image from dataset
        1. Click **Create Caption** to generate a caption for the image (usually takes 5-10s but could be over 60s)
        1. Different models can be selected:
          * **COCO** generally produces more straight-forward captions, but it is a smaller dataset and therefore struggles to recognize certain objects
          * **Conceptual Captions** is a much larger dataset but sometimes produces results that resemble social media posts rather than captions
        """
        gr.Markdown(instructions)
        with gr.Row():
            with gr.Column(variant="panel"):
                im_cap = gr.Image(interactive=False).style(height=height)
                model_name = gr.Radio(choices=["COCO","Conceptual Captions"], type="value", value="COCO", label="Model").style(container=True, item_container = False)
                with gr.Row():
                    get_btn_cap = gr.Button("Get Random Image").style(full_width=False)
                    caption_btn = gr.Button("Create Caption").style(full_width=False)
            caption = gr.Textbox(label='Caption', elem_classes="caption-text")
        get_btn_cap.click(fn=rand_image, outputs=im_cap)
        #im_cap.change(generate_text, inputs=im_cap, outputs=caption)
        caption_btn.click(get_caption, inputs=[im_cap, model_name], outputs=caption)
        gr.HTML(f"Dataset: <a href='https://github.com/unsplash/datasets' target='_blank'>Unsplash Lite</a>; Number of Images: {dataset_size}")

    with gr.Tab("Search"):
        instructions = """## Instructions:
        1. Enter a search query (or select one of the examples below)
        2. Click **Find Images** to find images that match the query (top 5 are shown in order from left to right)
        3. Keep in mind that the dataset contains mostly nature-focused images"""
        gr.Markdown(instructions)
        with gr.Column(variant="panel"):
            desc = gr.Textbox(show_label=False, placeholder="Enter description").style(container=False)
            gr.Examples(["someone holding flowers",
                         "someone holding pink flowers",
                         "red fruit in a person's hands",
                         "an aerial view of forest",
                         "a waterfall in Iceland with a rainbow"
                        ], inputs=desc)
            search_btn = gr.Button("Find Images").style(full_width=False)
        gallery = gr.Gallery(show_label=False).style(grid=(2,2,3,5))
        search_btn.click(search,inputs=desc, outputs=gallery, postprocess=False)
        gr.HTML(f"Dataset: <a href='https://github.com/unsplash/datasets' target='_blank'>Unsplash Lite</a>; Number of Images: {dataset_size}")

demo.queue(concurrency_count=3)
demo.launch()