Decorate models that don't use infer

Author: grzegorz-roboflow

inference/models/clip/clip_model.py

@@ -24,12 +24,12 @@
     CLIP_MODEL_ID,
     ONNXRUNTIME_EXECUTION_PROVIDERS,
     REQUIRED_ONNX_PROVIDERS,
-    TENSORRT_CACHE_PATH,
 )
 from inference.core.exceptions import OnnxProviderNotAvailable
 from inference.core.models.roboflow import OnnxRoboflowCoreModel
 from inference.core.models.types import PreprocessReturnMetadata
 from inference.core.models.utils.batching import create_batches
+from inference.usage_tracking.collector import usage_collector
 from inference.core.utils.image_utils import load_image_rgb
 from inference.core.utils.onnx import get_onnxruntime_execution_providers
 from inference.core.utils.postprocess import cosine_similarity
@@ -304,6 +304,7 @@ def get_infer_bucket_file_list(self) -> List[str]:
         """
         return ["textual.onnx", "visual.onnx"]
 
+    @usage_collector("model")
     def infer_from_request(
         self, request: ClipInferenceRequest
     ) -> ClipEmbeddingResponse:

inference/models/depth_estimation/depthestimation.py

@@ -21,6 +21,7 @@
     WorkflowImageData,
 )
 from inference.models.transformers import TransformerModel
+from inference.usage_tracking.collector import usage_collector
 
 
 class DepthEstimator(TransformerModel):
@@ -50,6 +51,7 @@ def __init__(self, *args, **kwargs):
                 "Running DepthPro on CPU. This may be very slow. Consider using GPU or MPS if available."
             )
 
+    @usage_collector("model")
     def predict(self, image_in: Image.Image, prompt="", history=None, **kwargs):
         try:
             # Process input image

inference/models/sam3/segment_anything3.py

@@ -65,6 +65,7 @@
 )
 from inference.core.utils.image_utils import load_image_rgb
 from inference.core.utils.postprocess import masks2multipoly
+from inference.usage_tracking.collector import usage_collector
 
 
 def _to_numpy_masks(masks_any) -> np.ndarray:
@@ -372,6 +373,7 @@ def infer_from_request(self, request: Sam3InferenceRequest):
         else:
             raise ValueError(f"Invalid request type {type(request)}")
 
+    @usage_collector("model")
     def segment_image(
         self,
         image: Optional[InferenceRequestImage],

inference/usage_tracking/decorator_helpers.py

@@ -109,6 +109,10 @@ def get_request_resource_id_from_kwargs(func_kwargs: Dict[str, Any]) -> Optional
             return str(dataset_id)
         if hasattr(inference_request, "model_id"):
             return str(inference_request.model_id)
+    if "request" in func_kwargs:
+        request = func_kwargs["request"]
+        if hasattr(request, "model_id"):
+            return str(request.model_id)
     if "dataset_id" in func_kwargs and "version_id" in func_kwargs:
         dataset_id = func_kwargs["dataset_id"]
         version_id = func_kwargs["version_id"]
@@ -121,6 +125,12 @@ def get_request_resource_id_from_kwargs(func_kwargs: Dict[str, Any]) -> Optional
         workflow_request = func_kwargs["workflow_request"]
         if hasattr(workflow_request, "workflow_id"):
             return str(workflow_request.workflow_id)
+    if "self" in func_kwargs:
+        _self = func_kwargs["self"]
+        if hasattr(_self, "model_id"):
+            return str(_self.model_id)
+        if hasattr(_self, "endpoint"):
+            return str(_self.endpoint)
     return None