Update app.py

app.py

@@ -1,10 +1,18 @@
 import gradio as gr
+from torchvision.transforms import Resize
+import torch
+from upstash_vector import Index
+import os
+
+index = Index.from_env()
+
+resize_transform = Resize((250,250))
 
 with gr.Blocks() as demo:
     gr.Markdown(
         """
         # Find Your Twins
-        
+
         Upload your face and find the most similar people from the X dataset. Powered by [Upstash Vector](https://upstash.com) 🚀
         """
     )
@@ -12,28 +20,39 @@ with gr.Blocks() as demo:
     with gr.Tab("Basic"):
         with gr.Row():
             with gr.Column(scale=1):
-                input_image = gr.Image()
+                input_image = gr.Image(type="pil")
             with gr.Column(scale=3):
                 output_image = gr.Gallery()
 
-
-        @input_image.change(inputs=input_image, outputs=output_image)
-        def find_similar_faces(image):
-            return [image] * 3
-
+        
+        @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]
+  
     with gr.Tab("Advanced"):
         with gr.Row():
             with gr.Column(scale=1):
-                adv_input_image = gr.Image()
+                adv_input_image = gr.Image(type="pil")
                 adv_image_count = gr.Number(9, label="Image Count")
 
             with gr.Column(scale=3):
                 adv_output_image = gr.Gallery(height=1000)
 
 
-        @adv_input_image.change(inputs=[adv_input_image, adv_image_count], outputs=[adv_output_image])
+        @adv_input_image.upload(inputs=[adv_input_image, adv_image_count], outputs=[adv_output_image])
         def find_similar_faces(image, count):
-            return [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]
+  
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=True)