Update hf_model.py

hf_model.py

@@ -71,7 +71,16 @@ class UniBioTransferModel(LatentDiffusion, PyTorchModelHubMixin):
     """
 
     def __getstate__(self):
-        return {k: v for k, v in self.__dict__.items() if k != 'ptsM_Generator'}
+        state = self.__dict__.copy()
+        if 'ptsM_Generator' in state:
+            del state['ptsM_Generator']
+        return state
+
+    def __getattr__(self, name):
+        if name == 'ptsM_Generator':
+            self.ptsM_Generator = LandmarkExtractor(include_visualizer=True, img_256_mode=False)
+            return self.ptsM_Generator
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
 
     def __init__(self, config=None, task="face", **kwargs):
         self._task_name = task if isinstance(task, str) else TASK_ID2NAME.get(task, "face")
@@ -232,10 +241,7 @@ class UniBioTransferModel(LatentDiffusion, PyTorchModelHubMixin):
 
         # offload_unused_tasks__LD(model, task_id, method="cpu")
         
-        def lazy_pts():
-            return LandmarkExtractor(include_visualizer=True, img_256_mode=False)
-        model.get_pts_generator = lazy_pts
-        model.ptsM_Generator = None
+        model.ptsM_Generator = LandmarkExtractor(include_visualizer=True, img_256_mode=False)
         cleanup_gpu_memory()
         
         # ZeroGPU 兼容：只在 device 不是 "cpu" 且 CUDA 可用时才移动到 GPU
@@ -250,11 +256,4 @@ class UniBioTransferModel(LatentDiffusion, PyTorchModelHubMixin):
         model._task_name = task_name
         model._hf_config = {"task": task_name, "task_id": task_id}
         
-        _orig_call = model.__call__
-        def _wrapped_call(*args, **kwargs):
-            if model.ptsM_Generator is None:
-                model.ptsM_Generator = model.get_pts_generator()
-            return _orig_call(*args, **kwargs)
-        model.__call__ = _wrapped_call
-        
         return model