refactor: ONNX model loading and caching functions

app.py

@@ -17,6 +17,7 @@ import torch
 
 from utils.utils import softmax, augment_image, preprocess_resize_256, preprocess_resize_224, postprocess_pipeline, postprocess_logits, postprocess_binary_output, to_float_scalar, infer_gradio_api, preprocess_gradio_api, postprocess_gradio_api
 from utils.onnx_helpers import preprocess_onnx_input, postprocess_onnx_output
+from utils.onnx_model_loader import load_onnx_model_and_preprocessor, get_onnx_model_from_cache
 from forensics.gradient import gradient_processing
 from forensics.minmax import minmax_process
 from forensics.ela import ELA
@@ -97,46 +98,10 @@ def register_model_with_metadata(model_id, model, preprocess, postprocess, class
     MODEL_REGISTRY[model_id] = entry
 
 
-def load_onnx_model_and_preprocessor(hf_model_id):
-    # model_dir = snapshot_download(repo_id=hf_model_id, local_dir_use_symlinks=False)
-    
-    # Create a unique local directory for each ONNX model
-    model_specific_dir = os.path.join("./models", hf_model_id.replace('/', '_'))
-    os.makedirs(model_specific_dir, exist_ok=True)
-
-    # Use hf_hub_download to get specific files into the model-specific directory
-    onnx_model_path = hf_hub_download(repo_id=hf_model_id, filename="model_quantized.onnx", subfolder="onnx", local_dir=model_specific_dir, local_dir_use_symlinks=False)
-
-    # Load preprocessor config
-    preprocessor_config = {}
-    try:
-        preprocessor_config_path = hf_hub_download(repo_id=hf_model_id, filename="preprocessor_config.json", local_dir=model_specific_dir, local_dir_use_symlinks=False)
-        with open(preprocessor_config_path, 'r') as f:
-            preprocessor_config = json.load(f)
-    except Exception as e:
-        logger.warning(f"Could not download or load preprocessor_config.json for {hf_model_id}: {e}")
-    
-    # Load model config for class names if available
-    model_config = {}
-    try:
-        model_config_path = hf_hub_download(repo_id=hf_model_id, filename="config.json", local_dir=model_specific_dir, local_dir_use_symlinks=False)
-        with open(model_config_path, 'r') as f:
-            model_config = json.load(f)
-    except Exception as e:
-        logger.warning(f"Could not download or load config.json for {hf_model_id}: {e}")
-
-    return onnxruntime.InferenceSession(onnx_model_path), preprocessor_config, model_config
-
 
 # Cache for ONNX sessions and preprocessors
 _onnx_model_cache = {}
 
-def get_onnx_model_from_cache(hf_model_id):
-    if hf_model_id not in _onnx_model_cache:
-        logger.info(f"Loading ONNX model and preprocessor for {hf_model_id}...")
-        _onnx_model_cache[hf_model_id] = load_onnx_model_and_preprocessor(hf_model_id)
-    return _onnx_model_cache[hf_model_id]
-
 
 def infer_onnx_model(hf_model_id, preprocessed_image_np, model_config: dict):
     try:

utils/onnx_model_loader.py

@@ -0,0 +1,33 @@
+import os
+import json
+import onnxruntime
+from huggingface_hub import hf_hub_download
+import logging
+
+def load_onnx_model_and_preprocessor(hf_model_id):
+    model_specific_dir = os.path.join("./models", hf_model_id.replace('/', '_'))
+    os.makedirs(model_specific_dir, exist_ok=True)
+    onnx_model_path = hf_hub_download(repo_id=hf_model_id, filename="model_quantized.onnx", subfolder="onnx", local_dir=model_specific_dir, local_dir_use_symlinks=False)
+    preprocessor_config = {}
+    try:
+        preprocessor_config_path = hf_hub_download(repo_id=hf_model_id, filename="preprocessor_config.json", local_dir=model_specific_dir, local_dir_use_symlinks=False)
+        with open(preprocessor_config_path, 'r') as f:
+            preprocessor_config = json.load(f)
+    except Exception as e:
+        logging.getLogger(__name__).warning(f"Could not download or load preprocessor_config.json for {hf_model_id}: {e}")
+    model_config = {}
+    try:
+        model_config_path = hf_hub_download(repo_id=hf_model_id, filename="config.json", local_dir=model_specific_dir, local_dir_use_symlinks=False)
+        with open(model_config_path, 'r') as f:
+            model_config = json.load(f)
+    except Exception as e:
+        logging.getLogger(__name__).warning(f"Could not download or load config.json for {hf_model_id}: {e}")
+    return onnxruntime.InferenceSession(onnx_model_path), preprocessor_config, model_config
+
+def get_onnx_model_from_cache(hf_model_id, _onnx_model_cache, load_func):
+    import logging
+    logger = logging.getLogger(__name__)
+    if hf_model_id not in _onnx_model_cache:
+        logger.info(f"Loading ONNX model and preprocessor for {hf_model_id}...")
+        _onnx_model_cache[hf_model_id] = load_func(hf_model_id)
+    return _onnx_model_cache[hf_model_id]