Update app.py

app.py

@@ -310,11 +310,13 @@ g_current_model = None
 
 # --- Global ONNX session ---
 g_session = None
+g_use_openvino = False
+g_execution_provider = None
 
 # --- Initialization Function ---
 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
-    global g_session
+    global g_session, g_use_openvino, g_execution_provider
     
     if not model_choice in MODEL_OPTIONS:
         print(f"Invalid model choice: {model_choice}, falling back to default: {DEFAULT_MODEL}")
@@ -325,7 +327,7 @@ def initialize_onnx_paths(model_choice=DEFAULT_MODEL):
     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}...")
+    print(f"Initializing paths and labels for model: {model_choice}...")
     hf_token = os.environ.get("HF_TOKEN")
     
     try:
@@ -353,13 +355,24 @@ def initialize_onnx_paths(model_choice=DEFAULT_MODEL):
         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)}")
 
-        # Load ONNX session ONCE here
-        print("Creating ONNX Runtime session (CPUExecutionProvider)...")
-        g_session = ort.InferenceSession(
-            g_onnx_model_path,
-            providers=["CPUExecutionProvider"]
-        )
-        print("ONNX Runtime session ready.")
+        # Try OpenVINO first, then fall back to ONNX Runtime
+        print("Attempting to initialize inference runtime...")
+        try:
+            import openvino as ov
+            
+            print("OpenVINO available, attempting to load model...")
+            core = ov.Core()
+            model = core.read_model(g_onnx_model_path)
+            g_session = core.compile_model(model, "CPU")
+            g_use_openvino = True
+            g_execution_provider = "CPU – OpenVINO™"
+            print("Successfully initialized with OpenVINO runtime")
+        except ImportError:
+            print("OpenVINO not available, falling back to ONNX Runtime CPU")
+            _init_onnx_runtime()
+        except Exception as e:
+            print(f"OpenVINO initialization failed: {e}, falling back to ONNX Runtime CPU")
+            _init_onnx_runtime()
 
         return True
 
@@ -373,9 +386,29 @@ def initialize_onnx_paths(model_choice=DEFAULT_MODEL):
         g_idx_to_tag = None
         g_tag_to_category = None
         g_current_model = None
-        # Raise Gradio error to make it visible in the UI
+        g_session = None
+        g_use_openvino = False
+        g_execution_provider = None
         raise gr.Error(f"Initialization failed: {e}. Check logs and HF_TOKEN.")
 
+def _init_onnx_runtime():
+    """Initialize ONNX Runtime with CPU as fallback"""
+    global g_session, g_use_openvino, g_execution_provider
+    
+    sess_options = ort.SessionOptions()
+    sess_options.log_severity_level = 3  # Only show errors
+    
+    providers = ["CPUExecutionProvider"]
+    g_session = ort.InferenceSession(
+        g_onnx_model_path,
+        sess_options=sess_options,
+        providers=providers
+    )
+    g_use_openvino = False
+    g_execution_provider = g_session.get_providers()[0]
+    print(f"ONNX Runtime session ready with {g_execution_provider}")
+
+
 # Function to handle model change
 def change_model(model_choice):
     try:
@@ -388,8 +421,10 @@ def change_model(model_choice):
         return f"Error changing model: {str(e)}"
         
 # --- Main Prediction Function (ONNX) ---
+# --- Main Prediction Function (ONNX/OpenVINO) ---
 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} ---")
+    print(f"--- predict_onnx function started with model {model_choice} ---")
+    print(f"Using runtime: {g_execution_provider}")
     
     # Ensure current model matches selected model
     global g_current_model
@@ -404,25 +439,12 @@ def predict_onnx(image_input, model_choice, gen_threshold, char_threshold, outpu
     if g_onnx_model_path is None or g_labels_data is None:
         message = "Error: Paths or labels not initialized. Check startup logs."
         print(message)
-        # Return error message and None for the image output
         return message, None
 
-    # --- 2. Load ONNX Session (inside worker) ---
-    session = None
-    try:
-        print(f"Loading ONNX session from: {g_onnx_model_path}")
-        available_providers = ort.get_available_providers()
-        providers = []
-        if 'CUDAExecutionProvider' in available_providers:
-            providers.append('CUDAExecutionProvider')
-        providers.append('CPUExecutionProvider')
-        print(f"Attempting to load session with providers: {providers}")
-        session = g_session
-        print(f"ONNX session loaded using: {session.get_providers()[0]}")
-    except Exception as e:
-        message = f"Error loading ONNX session in worker: {e}"
+    # --- 2. Check session is available ---
+    if g_session is None:
+        message = "Error: Inference session not initialized."
         print(message)
-        import traceback; traceback.print_exc()
         return message, None
 
     # --- 3. Process Input Image ---
@@ -433,26 +455,23 @@ def predict_onnx(image_input, model_choice, gen_threshold, char_threshold, outpu
     try:
         # Handle different input types (PIL, numpy, URL, file path)
         if isinstance(image_input, str):
-            if image_input.startswith("http"): # URL
+            if image_input.startswith("http"):
                 response = requests.get(image_input, timeout=10)
                 response.raise_for_status()
                 image = Image.open(io.BytesIO(response.content))
-            elif os.path.exists(image_input): # File path
+            elif os.path.exists(image_input):
                 image = Image.open(image_input)
             else:
-                 raise ValueError(f"Invalid image input string: {image_input}")
+                raise ValueError(f"Invalid image input string: {image_input}")
         elif isinstance(image_input, np.ndarray):
-             image = Image.fromarray(image_input)
+            image = Image.fromarray(image_input)
         elif isinstance(image_input, Image.Image):
-             image = image_input # Already a PIL image
+            image = image_input
         else:
-             raise TypeError(f"Unsupported image input type: {type(image_input)}")
+            raise TypeError(f"Unsupported image input type: {type(image_input)}")
 
         # Preprocess the PIL image
         original_pil_image, input_tensor = preprocess_image(image)
-
-        # Ensure input tensor is float32, as expected by most ONNX models
-        # (even if the model internally uses float16)
         input_tensor = input_tensor.astype(np.float32)
 
     except Exception as e:
@@ -462,49 +481,51 @@ def predict_onnx(image_input, model_choice, gen_threshold, char_threshold, outpu
 
     # --- 4. Run Inference ---
     try:
-        input_name = session.get_inputs()[0].name
-        output_name = session.get_outputs()[0].name
-        print(f"Running inference with input '{input_name}', output '{output_name}'")
+        print(f"Running inference with {'OpenVINO' if g_use_openvino else 'ONNX Runtime'}")
         start_time = time.time()
-        outputs = session.run([output_name], {input_name: input_tensor})[0]
+        
+        if g_use_openvino:
+            # OpenVINO inference
+            results = g_session(input_tensor)
+            outputs = list(results.values())[0]
+        else:
+            # ONNX Runtime inference
+            input_name = g_session.get_inputs()[0].name
+            output_name = g_session.get_outputs()[0].name
+            outputs = g_session.run([output_name], {input_name: input_tensor})[0]
+        
         inference_time = time.time() - start_time
-        print(f"Inference completed in {inference_time:.3f} seconds")
+        print(f"Inference completed in {inference_time:.3f} seconds using {g_execution_provider}")
 
         # Check for NaN/Inf in outputs
         if np.isnan(outputs).any() or np.isinf(outputs).any():
             print("Warning: NaN or Inf detected in model output. Clamping...")
-            outputs = np.nan_to_num(outputs, nan=0.0, posinf=1.0, neginf=0.0) # Clamp to 0-1 range
+            outputs = np.nan_to_num(outputs, nan=0.0, posinf=1.0, neginf=0.0)
 
-        # Apply sigmoid (outputs are likely logits)
-        # Use a stable sigmoid implementation
+        # Apply sigmoid
         def stable_sigmoid(x):
-            return 1 / (1 + np.exp(-np.clip(x, -30, 30))) # Clip to avoid overflow
-        probs = stable_sigmoid(outputs[0]) # Assuming batch size 1
+            return 1 / (1 + np.exp(-np.clip(x, -30, 30)))
+        probs = stable_sigmoid(outputs[0])
 
     except Exception as e:
-        message = f"Error during ONNX inference: {e}"
+        message = f"Error during inference: {e}"
         print(message)
         import traceback; traceback.print_exc()
         return message, None
-    finally:
-        # Clean up session if needed (might reduce memory usage between clicks)
-        del session
 
     # --- 5. Post-process and Format Output ---
     try:
         print("Post-processing results...")
-        # Use the correct global variable for labels
         predictions = get_tags(probs, g_labels_data, gen_threshold, char_threshold)
 
         # Format output text string
         output_tags = []
         if predictions.get("rating"): output_tags.append(predictions["rating"][0][0].replace("_", " "))
         if predictions.get("quality"): output_tags.append(predictions["quality"][0][0].replace("_", " "))
-        # Add other categories, respecting order and filtering meta if needed
+        
         for category in ["artist", "character", "copyright", "general", "meta", "model"]:
             tags_in_category = predictions.get(category, [])
             for tag, prob in tags_in_category:
-                # Basic meta tag filtering for text output
                 if category == "meta" and any(p in tag.lower() for p in ['id', 'commentary', 'request', 'mismatch']):
                     continue
                 output_tags.append(tag.replace("_", " "))
@@ -514,12 +535,8 @@ def predict_onnx(image_input, model_choice, gen_threshold, char_threshold, outpu
         viz_image = None
         if output_mode == "Tags + Visualization":
             print("Generating visualization...")
-            # Pass the correct threshold for display title (can pass both if needed)
-            # For simplicity, passing gen_threshold as a representative value
             viz_image = visualize_predictions(original_pil_image, predictions, gen_threshold)
             print("Visualization generated.")
-        else:
-            print("Visualization skipped.")
 
         print("Prediction complete.")
         return output_text, viz_image
@@ -540,6 +557,7 @@ footer { display: none !important; }
 
 with gr.Blocks(css=css) as demo:
     gr.Markdown("# CL EVA02 ONNX Tagger (CPU)")
+    gr.Markdown("OpenVINO™ is used for accelerated CPU inference when available, with ONNX Runtime as fallback.")
     gr.Markdown("This space is a duplicate of https://huggingface.co/spaces/cella110n/cl_tagger running on CPU and uses the [non-gated releases](https://huggingface.co/cella110n/cl_tagger) of cl-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).")