Update app.py

app.py

@@ -137,7 +137,7 @@ class GatedHead(torch.nn.Module):
 
 model.head = GatedHead(min(model.head.weight.shape), 9083)
 
-safetensors.torch.load_model(model, "JTP_PILOT2-e3-vit_so400m_patch14_siglip_384.safetensors")
+safetensors.torch.load_model(model, "JTP_PILOT2-2-e3-vit_so400m_patch14_siglip_384.safetensors")
 model.eval()
 
 with open("tagger_tags.json", "r") as file:
@@ -147,26 +147,36 @@ allowed_tags = list(tags.keys())
 for idx, tag in enumerate(allowed_tags):
     allowed_tags[idx] = tag.replace("_", " ")
 
+sorted_tag_score = {}
+
 @spaces.GPU(duration=5)
-def create_tags(image, threshold):
+def run_classifier(image, threshold):
+    global sorted_tag_score
     img = image.convert('RGB')
     tensor = transform(img).unsqueeze(0)
 
     with torch.no_grad():
-        probits = model(tensor).squeeze()
-        indices = torch.where(probits > threshold)[0]
-        values = probits[indices]
+        values, indices = torch.topk(model(tensor)[0], 250)
 
-    temp = []
     tag_score = dict()
     for i in range(indices.size(0)):
-        temp.append([allowed_tags[indices[i]], values[i].item()])
         tag_score[allowed_tags[indices[i]]] = values[i].item()
-    temp = [t[0] for t in temp]
-    text_no_impl = ", ".join(temp)
-    return text_no_impl, tag_score
+    sorted_tag_score = dict(sorted(tag_score.items(), key=lambda item: item[1], reverse=True))
+
+    return create_tags(threshold)
+
+def create_tags(threshold):
+    global sorted_tag_score
+    filtered_tag_score = {key: value for key, value in sorted_tag_score.items() if value > threshold}
+    text_no_impl = ", ".join(filtered_tag_score.keys())
+    return text_no_impl, filtered_tag_score
+
+def clear_image():
+    global sorted_tag_score
+    sorted_tag_score = {}
+    return "", {}
 
-with gr.Blocks() as demo:
+with gr.Blocks(css=".output-class { display: none; }") as demo:
     gr.Markdown("""
     ## Joint Tagger Project: JTP-PILOT² Demo **BETA**
     This tagger is designed for use on furry images (though may very well work on out-of-distribution images, potentially with funny results).  A threshold of 0.2 is recommended.  Lower thresholds often turn up more valid tags, but can also result in some amount of hallucinated tags.
@@ -175,14 +185,30 @@ with gr.Blocks() as demo:
 
     Special thanks to Minotoro at frosting.ai for providing the compute power for this project.
     """)
-    gr.Interface(
-        create_tags,
-        inputs=[gr.Image(label="Source", sources=['upload', 'webcam'], type='pil'), gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")],
-        outputs=[
-            gr.Textbox(label="Tag String"),
-            gr.Label(label="Tag Predictions", num_top_classes=200),
-        ],
-        allow_flagging="never",
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(label="Source", sources=['upload'], type='pil', height=512, show_label=False)
+            threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")
+        with gr.Column():
+            tag_string = gr.Textbox(label="Tag String")
+            label_box = gr.Label(label="Tag Predictions", num_top_classes=250, show_label=False)
+
+    image_input.upload(
+        fn=run_classifier,
+        inputs=[image_input, threshold_slider],
+        outputs=[tag_string, label_box]
+    )
+
+    image_input.clear(
+        fn=clear_image,
+        inputs=[],
+        outputs=[tag_string, label_box]
+    )
+
+    threshold_slider.input(
+        fn=create_tags,
+        inputs=[threshold_slider],
+        outputs=[tag_string, label_box]
     )
 
 if __name__ == "__main__":