dk - edit

person_attr.py

@@ -184,14 +184,30 @@ class PersonAttributeModel(nn.Module):
 # ---------------------------------------------------------------------------
 
 class PersonAttribute:
-    """image_url + xywh(top-left) → Market person attributes."""
+    """image_url + xywh(top-left) → Market person attributes.
 
-    def __init__(self, checkpoint_path: str | None = None, device: str | None = None):
+    ClearML Serving: subclass with auto_load=False, then call load(model_filepath).
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str | None = None,
+        device: str | None = None,
+        auto_load: bool = True,
+    ):
         self.checkpoint_path = checkpoint_path or DEFAULT_CHECKPOINT
         self.device = torch.device(
             device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
         )
+        self.model: PersonAttributeModel | None = None
+        if auto_load:
+            self.model = self.load_model()
+
+    def load(self, model_filepath: str) -> PersonAttribute:
+        """ClearML Serving hook — load weights from a downloaded checkpoint path."""
+        self.checkpoint_path = model_filepath
         self.model = self.load_model()
+        return self
 
     def load_model(self) -> PersonAttributeModel:
         if not os.path.isfile(self.checkpoint_path):
@@ -244,6 +260,8 @@ class PersonAttribute:
         return results
 
     def predict(self, image: Image.Image) -> dict[str, str]:
+        if self.model is None:
+            raise RuntimeError('Model is not loaded. Call load(model_filepath) first.')
         src = self.preprocess(image).to(self.device)
         with torch.no_grad():
             out = self.model.forward(src)