Clean and add model link

app.py

@@ -37,9 +37,11 @@ examples = [
     ["examples/1811.jpg", "Bags"],
 ]
 
+
 @torch.inference_mode()
 def retrieval(image, category):
-    if image is None or category is None: return
+    if image is None or category is None:
+        return
 
     q_emb = m(tfs(image).unsqueeze(0), torch.tensor([category]))
 
@@ -48,42 +50,61 @@ def retrieval(image, category):
     imgs = [process_img(idx, gal_imgs) for idx in r[1][0]]
 
     html = [make_img_html(i) for i in imgs]
-    html += ["<p></p>"] # Avoid Gradio's last-child{margin-bottom:0!important;}
+    html += ["<p></p>"]  # Avoid Gradio's last-child{margin-bottom:0!important;}
 
     return "\n".join(html)
 
 
-
 JavaScriptLoader("src/custom_functions.js")
 with gr.Blocks(css="src/style.css") as demo:
     with gr.Column():
-        gr.Markdown("""
+        gr.Markdown(
+            """
         # Conditional ViT Demo 
         [[`Paper`](https://arxiv.org/abs/2306.02928)] 
         [[`Code`](https://github.com/Simon-Lepage/CondViT-LRVSF)] 
         [[`Dataset`](https://huggingface.co/datasets/Slep/LAION-RVS-Fashion)] 
+        [[`Model`](https://huggingface.co/Slep/CondViT-B16-cat)] 
+
 
         *Running on 2 vCPU, 16Go RAM.*
 
         - **Model :** Categorical CondViT-B/16
         - **Gallery :** 93K images.
-        """)
+        """
+        )
 
         # Input section
         with gr.Row():
             img = gr.Image(label="Query Image", type="pil", elem_id="query_img")
             with gr.Column():
-                cat = gr.Dropdown(choices = categories, label="Category", value="Upper Body", type='index', elem_id="dropdown")
+                cat = gr.Dropdown(
+                    choices=categories,
+                    label="Category",
+                    value="Upper Body",
+                    type="index",
+                    elem_id="dropdown",
+                )
                 submit = gr.Button("Submit")
-        
+
         # Examples
-        gr.Examples(examples, inputs=[img, cat], fn=retrieval, elem_id = "preset_examples", examples_per_page=100)
-        gr.HTML(value=ExamplesHandler(examples).to_html(), label = "examples", elem_id = "html_examples")
-        
+        gr.Examples(
+            examples,
+            inputs=[img, cat],
+            fn=retrieval,
+            elem_id="preset_examples",
+            examples_per_page=100,
+        )
+        gr.HTML(
+            value=ExamplesHandler(examples).to_html(),
+            label="examples",
+            elem_id="html_examples",
+        )
+
         # Outputs
         gr.Markdown("# Retrieved Items")
-        out = gr.HTML(label="Results", elem_id = "html_output")
+        out = gr.HTML(label="Results", elem_id="html_output")
 
         submit.click(fn=retrieval, inputs=[img, cat], outputs=out)
 
-demo.launch()
+demo.launch()

src/custom_js.py

@@ -1,27 +0,0 @@
-import gradio
-
-# Adapted from https://github.com/gradio-app/gradio/discussions/2932
-
-class JavaScriptLoader:
-    def __init__(self, target):
-        #Copy the template response
-        self.original_template = gradio.routes.templates.TemplateResponse
-        #Prep the js files
-        self.load_js(target)
-        #reassign the template response to your method, so gradio calls your method instead
-        gradio.routes.templates.TemplateResponse = self.template_response
-
-    def load_js(self, target):
-        with open(target, 'r', encoding="utf-8") as file:
-                self.loaded_script = f"<script>\n{file.read()}\n</script>"
-
-    def template_response(self, *args, **kwargs):
-        """Once gradio calls your method, you call the original, you modify it to include
-        your scripts and you return the modified version
-        """
-        response = self.original_template(*args, **kwargs)
-        response.body = response.body.replace(
-            '</head>'.encode('utf-8'), self.loaded_script + "\n</head>".encode("utf-8")
-        )
-        response.init_headers()
-        return response