update app.py

.gitignore

@@ -0,0 +1,4 @@
+__pycache__
+*.swp
+hf_models/
+pretrained_models/

app.py

@@ -6,11 +6,12 @@ CC BY-NC-4.0 (https://creativecommons.org/licenses/by-nc/4.0/)
 import os
 import time
 from argparse import ArgumentParser
+import json
 
 import numpy as np
 import torch
 import gradio as gr
-from clip_retrieval.clip_client import ClipClient
+import faiss
 
 from encode_with_pseudo_tokens import encode_with_pseudo_tokens_HF
 from models import build_text_encoder, Phi, PIC2WORD
@@ -19,6 +20,7 @@ import transformers
 from huggingface_hub import hf_hub_url, cached_download
 
 
+
 def parse_args():
     parser = ArgumentParser()
     parser.add_argument("--lincir_ckpt_path", default=None, type=str,
@@ -100,6 +102,7 @@ def load_models(args):
             }
 
 
+@torch.no_grad()
 def predict(images, input_text, model_name):
     start_time = time.time()
     input_images = model_dict['clip_preprocess'](images, return_tensors='pt')['pixel_values'].to(model_dict['device'])
@@ -125,18 +128,15 @@ def predict(images, input_text, model_name):
     clip_text_time = time.time() - start_time
 
     start_time = time.time()
-    try:
-        results = client.query(embedding_input=text_embeddings[0].tolist())
-        output = ''
-    except:
-        results = []
-        output = 'The server for image retrieval is not working. Please try again later.'
-    retrieval_time = time.time() - start_time
 
+    _, results = faiss_index.search(text_embeddings.cpu().numpy(), k=10)
+
+    retrieval_time = time.time() - start_time
 
+    output = ''
 
-    for idx, result in enumerate(results):
-        image_url = result['url']
+    for idx, retrieved_idx in enumerate(results[0]):
+        image_url = image_urls[retrieved_idx]
         output += f'![image]({image_url})\n'
 
     time_output = {'CLIP visual extractor': clip_image_time,
@@ -180,7 +180,7 @@ def test_fps(batch_size=1):
 if __name__ == '__main__':
     args = parse_args()
 
-    global model_dict, client
+    global model_dict, faiss_index, image_urls
 
     model_dict = load_models(args)
 
@@ -189,19 +189,18 @@ if __name__ == '__main__':
         test_fps(1)
         exit()
 
+    faiss_index = faiss.read_index('./clip_large.index', faiss.IO_FLAG_MMAP | faiss.IO_FLAG_READ_ONLY)
 
-    client = ClipClient(url="https://knn.laion.ai/knn-service",
-                        indice_name="laion5B-H-14" if args.clip_model_name == "huge" else "laion5B-L-14",
-                        )
+    image_urls = json.load(open('./image_urls.json'))
 
-    title = 'Zeroshot CIR demo'
+    title = 'Zeroshot CIR demo to search high-quality AI images'
 
     md_title = f'''# {title}
     [LinCIR](https://arxiv.org/abs/2312.01998): Language-only Training of Zero-shot Composed Image Retrieval  
     [SEARLE](https://arxiv.org/abs/2303.15247): Zero-shot Composed Image Retrieval with Textual Inversion  
     [Pic2Word](https://arxiv.org/abs/2302.03084): Mapping Pictures to Words for Zero-shot Composed Image Retrieval  
 
-    K-NN index for the retrieval results are entirely trained using the entire Laion-5B imageset. This is made possible thanks to the great work of [rom1504](https://github.com/rom1504/clip-retrieval).
+    K-NN index for the retrieval results are entirely trained using [the upscaled midjourney v5 images (444,901)](https://huggingface.co/datasets/wanng/midjourney-v5-202304-clean).
     '''
 
     with gr.Blocks(title=title) as demo:

requirements.txt

@@ -6,3 +6,4 @@ accelerate
 datasets
 spacy
 git+https://github.com/rom1504/clip-retrieval
+faiss