Upload app.py

app.py

@@ -268,10 +268,12 @@ def visualize_predictions(image: Image.Image, predictions: Dict, threshold: floa
 
 # --- Constants ---
 REPO_ID = "cella110n/cl_tagger"
-# Use the specified ONNX model filename
-ONNX_FILENAME = "cl_eva02_tagger_v1_250426/model.onnx"
-# Correct the tag mapping path to match the ONNX model's directory
-TAG_MAPPING_FILENAME = "cl_eva02_tagger_v1_250426/tag_mapping.json"
+# Model options
+MODEL_OPTIONS = {
+    "cl_eva02_tagger_v1_250426": "cl_eva02_tagger_v1_250426/model.onnx",
+    "cl_eva02_tagger_v1_250427": "cl_eva02_tagger_v1_250427/model.onnx"
+}
+DEFAULT_MODEL = "cl_eva02_tagger_v1_250426"
 CACHE_DIR = "./model_cache"
 
 # --- Global variables for paths (initialized at startup) ---
@@ -280,35 +282,76 @@ g_tag_mapping_path = None
 g_labels_data = None
 g_idx_to_tag = None
 g_tag_to_category = None
+g_current_model = None
 
 # --- Initialization Function ---
-def initialize_onnx_paths():
-    global g_onnx_model_path, g_tag_mapping_path, g_labels_data, g_idx_to_tag, g_tag_to_category
-    print("Initializing ONNX paths and labels...")
+def initialize_onnx_paths(model_choice=DEFAULT_MODEL):
+    global g_onnx_model_path, g_tag_mapping_path, g_labels_data, g_idx_to_tag, g_tag_to_category, g_current_model
+    
+    if not model_choice in MODEL_OPTIONS:
+        print(f"Invalid model choice: {model_choice}, falling back to default: {DEFAULT_MODEL}")
+        model_choice = DEFAULT_MODEL
+    
+    g_current_model = model_choice
+    model_dir = model_choice
+    onnx_filename = MODEL_OPTIONS[model_choice]
+    tag_mapping_filename = f"{model_dir}/tag_mapping.json"
+    
+    print(f"Initializing ONNX paths and labels for model: {model_choice}...")
     hf_token = os.environ.get("HF_TOKEN")
     try:
-        print(f"Attempting to download ONNX model: {ONNX_FILENAME}")
-        g_onnx_model_path = hf_hub_download(repo_id=REPO_ID, filename=ONNX_FILENAME, cache_dir=CACHE_DIR, token=hf_token, force_download=False)
+        print(f"Attempting to download ONNX model: {onnx_filename}")
+        g_onnx_model_path = hf_hub_download(repo_id=REPO_ID, filename=onnx_filename, cache_dir=CACHE_DIR, token=hf_token, force_download=False)
         print(f"ONNX model path: {g_onnx_model_path}")
 
-        print(f"Attempting to download Tag mapping: {TAG_MAPPING_FILENAME}")
-        g_tag_mapping_path = hf_hub_download(repo_id=REPO_ID, filename=TAG_MAPPING_FILENAME, cache_dir=CACHE_DIR, token=hf_token, force_download=False)
+        print(f"Attempting to download Tag mapping: {tag_mapping_filename}")
+        g_tag_mapping_path = hf_hub_download(repo_id=REPO_ID, filename=tag_mapping_filename, cache_dir=CACHE_DIR, token=hf_token, force_download=False)
         print(f"Tag mapping path: {g_tag_mapping_path}")
 
         print("Loading labels from mapping...")
         g_labels_data, g_idx_to_tag, g_tag_to_category = load_tag_mapping(g_tag_mapping_path)
         print(f"Labels loaded. Count: {len(g_labels_data.names)}")
+        
+        return True
 
     except Exception as e:
         print(f"Error during initialization: {e}")
         import traceback; traceback.print_exc()
+        # Reset globals to force reinitialization
+        g_onnx_model_path = None
+        g_tag_mapping_path = None
+        g_labels_data = None
+        g_idx_to_tag = None
+        g_tag_to_category = None
+        g_current_model = None
         # Raise Gradio error to make it visible in the UI
         raise gr.Error(f"Initialization failed: {e}. Check logs and HF_TOKEN.")
 
+# Function to handle model change
+def change_model(model_choice):
+    try:
+        success = initialize_onnx_paths(model_choice)
+        if success:
+            return f"Model changed to: {model_choice}"
+        else:
+            return "Failed to change model. See logs for details."
+    except Exception as e:
+        return f"Error changing model: {str(e)}"
+        
 # --- Main Prediction Function (ONNX) ---
 @spaces.GPU()
-def predict_onnx(image_input, gen_threshold, char_threshold, output_mode):
-    print("--- predict_onnx function started (GPU worker) ---")
+def predict_onnx(image_input, model_choice, gen_threshold, char_threshold, output_mode):
+    print(f"--- predict_onnx function started (GPU worker) with model {model_choice} ---")
+    
+    # Ensure current model matches selected model
+    global g_current_model
+    if g_current_model != model_choice:
+        print(f"Model mismatch! Current: {g_current_model}, Selected: {model_choice}. Reinitializing...")
+        try:
+            initialize_onnx_paths(model_choice)
+        except Exception as e:
+            return f"Error initializing model '{model_choice}': {str(e)}", None
+    
     # --- 1. Ensure paths and labels are loaded ---
     if g_onnx_model_path is None or g_labels_data is None:
         message = "Error: Paths or labels not initialized. Check startup logs."
@@ -450,11 +493,16 @@ footer { display: none !important; }
 with gr.Blocks(css=css) as demo:
     gr.Markdown("# CL EVA02 ONNX Tagger")
     gr.Markdown("Upload an image or paste an image URL to predict tags using the CL EVA02 Tagger model (ONNX), fine-tuned from [SmilingWolf/wd-eva02-large-tagger-v3](https://huggingface.co/SmilingWolf/wd-eva02-large-tagger-v3).")
+    
     with gr.Row():
         with gr.Column(scale=1):
             image_input = gr.Image(type="pil", label="Input Image", elem_id="input-image")
-            # Add back URL input capability if desired (needs JS or separate component)
-            # gr.HTML("<div id='url-input-container'></div>")
+            model_choice = gr.Dropdown(
+                choices=list(MODEL_OPTIONS.keys()), 
+                value=DEFAULT_MODEL, 
+                label="Model Version",
+                interactive=True
+            )
             gen_threshold = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.55, label="General/Meta Tag Threshold")
             char_threshold = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.60, label="Character/Copyright/Artist Tag Threshold")
             output_mode = gr.Radio(choices=["Tags Only", "Tags + Visualization"], value="Tags + Visualization", label="Output Mode")
@@ -462,28 +510,37 @@ with gr.Blocks(css=css) as demo:
         with gr.Column(scale=1):
             output_tags = gr.Textbox(label="Predicted Tags", lines=10, interactive=False)
             output_visualization = gr.Image(type="pil", label="Prediction Visualization", interactive=False)
+    
+    # Handle model change
+    model_status = gr.Textbox(label="Model Status", interactive=False, visible=False)
+    model_choice.change(
+        fn=change_model,
+        inputs=[model_choice],
+        outputs=[model_status]
+    )
+    
     gr.Examples(
         examples=[
-            ["https://pbs.twimg.com/media/GXBXsRvbQAAg1kp.jpg", 0.55, 0.70, "Tags + Visualization"],
-            ["https://pbs.twimg.com/media/GjlX0gibcAA4EJ4.jpg", 0.55, 0.70, "Tags Only"],
-            ["https://pbs.twimg.com/media/Gj4nQbjbEAATeoH.jpg", 0.55, 0.70, "Tags + Visualization"],
-            ["https://pbs.twimg.com/media/GkbtX0GaoAMlUZt.jpg", 0.55, 0.70, "Tags + Visualization"]
+            ["https://pbs.twimg.com/media/GXBXsRvbQAAg1kp.jpg", DEFAULT_MODEL, 0.55, 0.70, "Tags + Visualization"],
+            ["https://pbs.twimg.com/media/GjlX0gibcAA4EJ4.jpg", DEFAULT_MODEL, 0.55, 0.70, "Tags Only"],
+            ["https://pbs.twimg.com/media/Gj4nQbjbEAATeoH.jpg", DEFAULT_MODEL, 0.55, 0.70, "Tags + Visualization"],
+            ["https://pbs.twimg.com/media/GkbtX0GaoAMlUZt.jpg", DEFAULT_MODEL, 0.55, 0.70, "Tags + Visualization"]
         ],
-        inputs=[image_input, gen_threshold, char_threshold, output_mode],
+        inputs=[image_input, model_choice, gen_threshold, char_threshold, output_mode],
         outputs=[output_tags, output_visualization],
         fn=predict_onnx, # Use the ONNX prediction function
         cache_examples=False # Disable caching for examples during testing
     )
     predict_button.click(
         fn=predict_onnx, # Use the ONNX prediction function
-        inputs=[image_input, gen_threshold, char_threshold, output_mode],
+        inputs=[image_input, model_choice, gen_threshold, char_threshold, output_mode],
         outputs=[output_tags, output_visualization]
     )
 
 # --- Main Block ---
 if __name__ == "__main__":
     if not os.environ.get("HF_TOKEN"): print("Warning: HF_TOKEN environment variable not set.")
-    # Initialize paths and labels at startup
-    initialize_onnx_paths()
+    # Initialize paths and labels at startup (with default model)
+    initialize_onnx_paths(DEFAULT_MODEL)
     # Launch Gradio app
     demo.launch(share=True)