tweaks

app.py

@@ -2,27 +2,22 @@ import gradio as gr
 import os
 from torchvision.transforms import Resize
 from upstash_vector import Index
-
+from datasets import load_dataset
+from transformers import AutoFeatureExtractor, AutoModel
 
 index = Index.from_env()
-print(os.environ("UPSTASH_VECTOR_REST_URL"))
-print(os.environ("UPSTASH_VECTOR_REST_TOKEN"))
+print(os.environ["UPSTASH_VECTOR_REST_URL"])
+print(os.environ["UPSTASH_VECTOR_REST_TOKEN"])
 
-resize_transform = Resize((250,250))
-
-
-from transformers import AutoFeatureExtractor, AutoModel
+resize_transform = Resize((250, 250))
 
 model_ckpt = "google/vit-base-patch16-224-in21k"
 extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
 model = AutoModel.from_pretrained(model_ckpt)
 hidden_dim = model.config.hidden_size
 
-from datasets import load_dataset
-
 dataset = load_dataset("HengJi/human_faces")
 
-
 with gr.Blocks() as demo:
     gr.Markdown(
         """
@@ -39,16 +34,16 @@ with gr.Blocks() as demo:
             with gr.Column(scale=3):
                 output_image = gr.Gallery()
 
-        
+
         @input_image.upload(inputs=input_image, outputs=output_image)
-        def find_similar_faces(image): 
-          resized_image = resize_transform(image)
-          inputs = extractor(images=image, return_tensors="pt")
-          outputs = model(**inputs)
-          embed = outputs.last_hidden_state[0][0]
-          result = index.query(vector=embed.tolist(), top_k=3)
-          return[dataset["train"][int(vector.id[3:])]["image"] for vector in result]
-  
+        def find_similar_faces(image):
+            resized_image = resize_transform(image)
+            inputs = extractor(images=resized_image, return_tensors="pt")
+            outputs = model(**inputs)
+            embed = outputs.last_hidden_state[0][0]
+            result = index.query(vector=embed.tolist(), top_k=3)
+            return [dataset["train"][int(vector.id[3:])]["image"] for vector in result]
+
     with gr.Tab("Advanced"):
         with gr.Row():
             with gr.Column(scale=1):
@@ -61,13 +56,12 @@ with gr.Blocks() as demo:
 
         @adv_input_image.upload(inputs=[adv_input_image, adv_image_count], outputs=[adv_output_image])
         def find_similar_faces(image, count):
-          resized_image = resize_transform(image)
-          inputs = extractor(images=image, return_tensors="pt")
-          outputs = model(**inputs)
-          embed = outputs.last_hidden_state[0][0]
-          result = index.query(vector=embed.tolist(), top_k=min(count, 9))
-          return[dataset["train"][int(vector.id[3:])]["image"] for vector in result]
-  
+            resized_image = resize_transform(image)
+            inputs = extractor(images=resized_image, return_tensors="pt")
+            outputs = model(**inputs)
+            embed = outputs.last_hidden_state[0][0]
+            result = index.query(vector=embed.tolist(), top_k=min(count, 9))
+            return [dataset["train"][int(vector.id[3:])]["image"] for vector in result]
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    demo.launch(debug=True)

requirements.txt

@@ -1,4 +1,5 @@
 torchvision
 transformers
 datasets
-upstash-vector
+upstash-vector
+gradio