fix pt2

app.py

@@ -236,13 +236,29 @@ def create_tags(threshold, sorted_tag_score: dict):
     return text_no_impl, filtered_tag_score
 
 def clear_image():
-    return "", {}, None, {}, None
+    return "", {}, None, {}, None, None
+
+def extract_selected_tag(evt: gr.SelectData):
+    # evt is a gr.SelectData; keep it out of GPU calls
+    try:
+        return evt.value
+    except Exception:
+        return None
 
 @spaces.GPU(duration=5)
 def cam_inference(img, threshold, alpha, selected_tag: str):
-    target_tag_index = tags[selected_tag] 
-    tensor = transform(img).unsqueeze(0)
+    if img is None or not selected_tag:
+        return img, None
 
+    # Map to index
+    if selected_tag not in tags:
+        key = selected_tag.replace("_", " ")
+        if key not in tags:
+            return img, None
+        selected_tag = key
+
+    target_tag_index = tags[selected_tag]
+    tensor = transform(img).unsqueeze(0)
 
     if torch.cuda.is_available():
         tensor = tensor.to(device, dtype=torch.float16)
@@ -438,6 +454,7 @@ with gr.Blocks(css=custom_css) as demo:
             original_image_state = gr.State() # stash a copy of the input image
             sorted_tag_score_state = gr.State(value={}) # stash a copy of the input image
             cam_state = gr.State()
+            selected_tag_state = gr.State(value=None)
             with gr.Row():
                 with gr.Column():
                     image = gr.Image(label="Source", sources=['upload', 'clipboard'], type='pil', show_label=False, elem_id="image_container")
@@ -458,7 +475,7 @@ with gr.Blocks(css=custom_css) as demo:
             image.clear(
                 fn=clear_image,
                 inputs=[],
-                outputs=[tag_string, label_box, original_image_state, sorted_tag_score_state, cam_state]
+                outputs=[tag_string, label_box, original_image_state, sorted_tag_score_state, cam_state, selected_tag_state]
             )
 
             threshold_slider.input(
@@ -469,8 +486,14 @@ with gr.Blocks(css=custom_css) as demo:
             )
 
             label_box.select(
+                fn=extract_selected_tag,
+                inputs=None,
+                outputs=selected_tag_state,
+                show_progress='hidden',
+                queue=False # This should be a very fast operation
+            ).then(
                 fn=cam_inference,
-                inputs=[original_image_state, cam_slider, alpha_slider, label_box],
+                inputs=[original_image_state, cam_slider, alpha_slider, selected_tag_state],
                 outputs=[image, cam_state],
                 show_progress='minimal'
             )