Feature parity with the other space

app.py

@@ -55,7 +55,15 @@ class Predictor:
         config = json.loads(open("index/cosine_infos.json").read())["index_param"]
         faiss.ParameterSpace().set_index_parameters(self.knn_index, config)
 
-    def predict(self, positive_tags, negative_tags, n_neighbours=5):
+    def predict(
+        self,
+        positive_tags,
+        negative_tags,
+        selected_ratings,
+        n_neighbours,
+        api_username,
+        api_key,
+    ):
         tags_df = self.tags_df
         model = self.model
 
@@ -99,12 +107,9 @@ class Predictor:
         for image_id, dist in zip(neighbours_ids, dists[0]):
             current_url = danbooru_id_to_url(
                 image_id,
-                [
-                    "General",
-                    "Sensitive",
-                    "Questionable",
-                    "Explicit",
-                ],
+                selected_ratings,
+                api_username,
+                api_key,
             )
             if current_url is not None:
                 image_urls.append(current_url)
@@ -117,16 +122,39 @@ def main():
 
     with gr.Blocks() as demo:
         with gr.Row():
-            positive_tags = gr.Textbox(label="Positive tags")
-            negative_tags = gr.Textbox(label="Negative tags")
-
+            with gr.Column():
+                positive_tags = gr.Textbox(label="Positive tags")
+                negative_tags = gr.Textbox(label="Negative tags")
+                n_neighbours = gr.Slider(
+                    minimum=1,
+                    maximum=20,
+                    value=5,
+                    step=1,
+                    label="# of images",
+                )
+
+            with gr.Column():
+                api_username = gr.Textbox(label="Danbooru API Username")
+                api_key = gr.Textbox(label="Danbooru API Key")
+                selected_ratings = gr.CheckboxGroup(
+                    choices=["General", "Sensitive", "Questionable", "Explicit"],
+                    value=["General", "Sensitive"],
+                    label="Ratings",
+                )
         find_btn = gr.Button("Find similar images")
 
         similar_images = gr.Gallery(label="Similar images", columns=[5])
 
         find_btn.click(
             fn=predictor.predict,
-            inputs=[positive_tags, negative_tags],
+            inputs=[
+                positive_tags,
+                negative_tags,
+                selected_ratings,
+                n_neighbours,
+                api_username,
+                api_key,
+            ],
             outputs=[similar_images],
         )