tweaks

app.py

@@ -11,7 +11,9 @@ model = AutoModel.from_pretrained(model_ckpt)
 hidden_dim = model.config.hidden_size
 dataset = load_dataset("BounharAbdelaziz/Face-Aging-Dataset")
 
-MAX_K = 50
+TOP_K = 1000
+BASE_COUNT=4
+MAX_COUNT = 30
 
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -36,8 +38,8 @@ with gr.Blocks() as demo:
             inputs = extractor(images=image, return_tensors="pt")
             outputs = model(**inputs)
             embed = outputs.last_hidden_state[0][0]
-            result = await index.query(vector=embed.tolist(), top_k=4)
-            return [dataset["train"][int(vector.id)]["image"] for vector in result]
+            result = await index.query(vector=embed.tolist(), top_k=TOP_K)
+            return [dataset["train"][int(vector.id)]["image"] for vector in result[:BASE_COUNT]]
 
         gr.Examples(
             examples=[
@@ -55,31 +57,46 @@ with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column(scale=1):
                 adv_input_image = gr.Image(type="pil")
-                adv_image_count = gr.Slider(1, MAX_K, 10, label="Image Count")
+                adv_image_count = gr.Slider(1, MAX_COUNT, 10, label="Image Count")
                 adv_button = gr.Button("Submit")
 
             with gr.Column(scale=2):
-                adv_output_image = gr.Gallery()
+                adv_output_images = gr.Gallery()
 
         async def find_similar_faces(image, count):
+            if image is None:
+                return None
             inputs = extractor(images=image, return_tensors="pt")
             outputs = model(**inputs)
             embed = outputs.last_hidden_state[0][0]
             result = await index.query(
-                vector=embed.tolist(), top_k=max(1, min(MAX_K, int(count)))
+                vector=embed.tolist(), top_k=TOP_K
             )
-            return [dataset["train"][int(vector.id)]["image"] for vector in result]
+            return [dataset["train"][int(vector.id)]["image"] for vector in result[:int(count)]]
 
         adv_button.click(
             fn=find_similar_faces,
             inputs=[adv_input_image, adv_image_count],
-            outputs=[adv_output_image],
+            outputs=[adv_output_images],
         )
-        adv_input_image.upload(
+        adv_input_image.change(
             fn=find_similar_faces,
             inputs=[adv_input_image, adv_image_count],
-            outputs=[adv_output_image],
+            outputs=[adv_output_images],
         )
+        gr.Examples(
+            examples=[
+                [dataset["train"][6]["image"], MAX_COUNT],
+                [dataset["train"][7]["image"], MAX_COUNT],
+                [dataset["train"][8]["image"], MAX_COUNT],
+            ],
+            inputs=[adv_input_image, adv_image_count],
+            outputs=adv_output_images,
+            fn=find_similar_faces,
+            cache_examples=False,
+        )
+
+
 
 if __name__ == "__main__":
     demo.queue(default_concurrency_limit=40)