refactor: consolidate model registration logic for ONNX and Gradio API

app.py

@@ -80,8 +80,6 @@ CLASS_NAMES = {
     "model_8": ['Fake', 'Real'],
 }
 
-# Register all models (ONNX, HuggingFace, Gradio API)
-register_all_models(MODEL_PATHS, CLASS_NAMES, device, infer_onnx_model, preprocess_onnx_input, postprocess_onnx_output)
 
 
 
@@ -124,167 +122,14 @@ def infer_onnx_model(hf_model_id, preprocessed_image_np, model_config: dict):
         # Return a structure consistent with other model errors
         return {"logits": np.array([]), "probabilities": np.array([])}
 
+# Register all models (ONNX, HuggingFace, Gradio API)
+register_all_models(MODEL_PATHS, CLASS_NAMES, device, infer_onnx_model, preprocess_onnx_input, postprocess_onnx_output)
 
 # Register the ONNX quantized model
 # Dummy entry for ONNX model to be loaded dynamically
 # We will now register a 'wrapper' that handles dynamic loading
 
-class ONNXModelWrapper:
-    def __init__(self, hf_model_id):
-        self.hf_model_id = hf_model_id
-        self._session = None
-        self._preprocessor_config = None
-        self._model_config = None
-
-    def load(self):
-        if self._session is None:
-            self._session, self._preprocessor_config, self._model_config = get_onnx_model_from_cache(
-                self.hf_model_id, _onnx_model_cache, load_onnx_model_and_preprocessor
-            )
-            logger.info(f"ONNX model {self.hf_model_id} loaded into wrapper.")
-
-    def __call__(self, image_np):
-        self.load() # Ensure model is loaded on first call
-        # Pass model_config to infer_onnx_model
-        return infer_onnx_model(self.hf_model_id, image_np, self._model_config)
-
-    def preprocess(self, image: Image.Image):
-        self.load()
-        return preprocess_onnx_input(image, self._preprocessor_config)
-
-    def postprocess(self, onnx_output: dict, class_names_from_registry: list): # class_names_from_registry is ignored
-        self.load()
-        return postprocess_onnx_output(onnx_output, self._model_config)
-
-# Consolidate all model loading and registration
-for model_key, hf_model_path in MODEL_PATHS.items():
-    logger.debug(f"Attempting to register model: {model_key} with path: {hf_model_path}")
-    model_num = model_key.replace("model_", "").upper()
-    contributor = "Unknown"
-    architecture = "Unknown"
-    dataset = "TBA"
-
-    current_class_names = CLASS_NAMES.get(model_key, [])
-
-    # Logic for ONNX models (1, 2, 3, 5, 6, 7)
-    if "ONNX" in hf_model_path:
-        logger.debug(f"Model {model_key} identified as ONNX.")
-        logger.info(f"Registering ONNX model: {model_key} from {hf_model_path}")
-        onnx_wrapper_instance = ONNXModelWrapper(hf_model_path)
-
-        # Attempt to derive contributor, architecture, dataset based on model_key
-        if model_key == "model_1":
-            contributor = "haywoodsloan"
-            architecture = "SwinV2"
-            dataset = "DeepFakeDetection"
-        elif model_key == "model_2":
-            contributor = "Heem2"
-            architecture = "ViT"
-            dataset = "DeepFakeDetection"
-        elif model_key == "model_3":
-            contributor = "Organika"
-            architecture = "VIT"
-            dataset = "SDXL"
-        elif model_key == "model_5":
-            contributor = "prithivMLmods"
-            architecture = "VIT"
-        elif model_key == "model_6":
-            contributor = "ideepankarsharma2003"
-            architecture = "SWINv1"
-            dataset = "SDXL, Midjourney"
-        elif model_key == "model_7":
-            contributor = "date3k2"
-            architecture = "VIT"
-
-        display_name_parts = [model_num]
-        if architecture and architecture not in ["Unknown"]:
-            display_name_parts.append(architecture)
-        if dataset and dataset not in ["TBA"]:
-            display_name_parts.append(dataset)
-        display_name = "-".join(display_name_parts)
-        display_name += "_ONNX" # Always append _ONNX for ONNX models
-
-        register_model_with_metadata(
-            model_id=model_key,
-            model=onnx_wrapper_instance, # The callable wrapper for the ONNX model
-            preprocess=onnx_wrapper_instance.preprocess,
-            postprocess=onnx_wrapper_instance.postprocess,
-            class_names=current_class_names, # Initial class names; will be overridden by model_config if available
-            display_name=display_name,
-            contributor=contributor,
-            model_path=hf_model_path,
-            architecture=architecture,
-            dataset=dataset
-        )
-    # Logic for Gradio API model (model_8)
-    elif model_key == "model_8":
-        logger.debug(f"Model {model_key} identified as Gradio API.")
-        logger.info(f"Registering Gradio API model: {model_key} from {hf_model_path}")
-        contributor = "aiwithoutborders-xyz"
-        architecture = "ViT"
-        dataset = "DeepfakeDetection"
-
-        display_name_parts = [model_num]
-        if architecture and architecture not in ["Unknown"]:
-            display_name_parts.append(architecture)
-        if dataset and dataset not in ["TBA"]:
-            display_name_parts.append(dataset)
-        display_name = "-".join(display_name_parts)
-
-        register_model_with_metadata(
-            model_id=model_key,
-            model=infer_gradio_api,
-            preprocess=preprocess_gradio_api,
-            postprocess=postprocess_gradio_api,
-            class_names=current_class_names,
-            display_name=display_name,
-            contributor=contributor,
-            model_path=hf_model_path,
-            architecture=architecture,
-            dataset=dataset
-        )
-    # Logic for PyTorch/Hugging Face pipeline models (currently only model_4)
-    elif model_key == "model_4": # Explicitly handle model_4
-        logger.debug(f"Model {model_key} identified as PyTorch/HuggingFace pipeline.")
-        logger.info(f"Registering HuggingFace pipeline/AutoModel: {model_key} from {hf_model_path}")
-        contributor = "cmckinle"
-        architecture = "VIT"
-        dataset = "SDXL, FLUX"
-
-        display_name_parts = [model_num]
-        if architecture and architecture not in ["Unknown"]:
-            display_name_parts.append(architecture)
-        if dataset and dataset not in ["TBA"]:
-            display_name_parts.append(dataset)
-        display_name = "-".join(display_name_parts)
-
-        current_processor = AutoFeatureExtractor.from_pretrained(hf_model_path, device=device)
-        model_instance = AutoModelForImageClassification.from_pretrained(hf_model_path).to(device)
-
-        preprocess_func = preprocess_resize_256
-        postprocess_func = postprocess_logits
-
-        def custom_infer(image, processor_local=current_processor, model_local=model_instance):
-            inputs = processor_local(image, return_tensors="pt").to(device)
-            with torch.no_grad():
-                outputs = model_local(**inputs)
-            return outputs
-        model_instance = custom_infer
-
-        register_model_with_metadata(
-            model_id=model_key,
-            model=model_instance,
-            preprocess=preprocess_func,
-            postprocess=postprocess_func,
-            class_names=current_class_names,
-            display_name=display_name,
-            contributor=contributor,
-            model_path=hf_model_path,
-            architecture=architecture,
-            dataset=dataset
-        )
-    else: # Fallback for any unhandled models (shouldn't happen if MODEL_PATHS is fully covered)
-        logger.warning(f"Could not automatically load and register model: {model_key} from {hf_model_path}. No matching registration logic found.")
+
 
 
 def infer(image: Image.Image, model_id: str, confidence_threshold: float = 0.75) -> dict: