Update src/facerender/animate.py

src/facerender/animate.py

@@ -68,7 +68,7 @@ class AnimateFromCoeff():
             raise AttributeError("Checkpoint should be specified for video head pose estimator.")
 
         if sadtalker_path['mappingnet_checkpoint'] is not None:
-            self.load_cpk_mapping(sadtalker_path['mappingnet_checkpoint'], mapping=mapping)
+            self._load_cpk_mapping(sadtalker_path['mappingnet_checkpoint'], mapping=mapping)
         else:
             raise AttributeError("Checkpoint should be specified for video head pose estimator.")
 
@@ -128,33 +128,29 @@ class AnimateFromCoeff():
 
         return checkpoint.get('epoch', 0)
 
-    def load_cpk_mapping(self, checkpoint_path, mapping=None, discriminator=None,
-                     optimizer_mapping=None, optimizer_discriminator=None, device='cpu'):
-
-    # Eğer .tar ise aç ve klasör haline getir
-    if checkpoint_path.endswith(".tar"):
-        try:
-            print(f"Extracting tar file: {checkpoint_path}")
-            with tarfile.open(checkpoint_path, "r") as tar:
-                extract_dir = checkpoint_path.replace(".tar", "")
-                os.makedirs(extract_dir, exist_ok=True)
-                tar.extractall(path=extract_dir)
-            checkpoint_path = os.path.join(extract_dir, "data.pkl")
-        except Exception as e:
-            print(f"Tar açma hatası: {e}")
-            return 0  # fail gracefully
-
-    # Checkpoint (data.pkl) yükle
-    checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
-
-    if mapping is not None and 'mapping' in checkpoint:
-        mapping.load_state_dict(checkpoint['mapping'])
-    if discriminator is not None and 'discriminator' in checkpoint:
-        discriminator.load_state_dict(checkpoint['discriminator'])
-    if optimizer_mapping is not None and 'optimizer_mapping' in checkpoint:
-        optimizer_mapping.load_state_dict(checkpoint['optimizer_mapping'])
-    if optimizer_discriminator is not None and 'optimizer_discriminator' in checkpoint:
-        optimizer_discriminator.load_state_dict(checkpoint['optimizer_discriminator'])
-
-    return checkpoint.get('epoch', 0)
+    def _load_cpk_mapping(self, checkpoint_path, mapping=None, discriminator=None,
+                         optimizer_mapping=None, optimizer_discriminator=None, device='cpu'):
 
+        if checkpoint_path.endswith(".tar"):
+            try:
+                with tarfile.open(checkpoint_path, "r") as tar:
+                    extract_dir = checkpoint_path.replace(".tar", "")
+                    os.makedirs(extract_dir, exist_ok=True)
+                    tar.extractall(path=extract_dir)
+                    checkpoint_path = os.path.join(extract_dir, "data.pkl")
+            except Exception as e:
+                print(f"Tar açma hatası: {e}")
+                return 0
+
+        checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
+
+        if mapping is not None and 'mapping' in checkpoint:
+            mapping.load_state_dict(checkpoint['mapping'])
+        if discriminator is not None and 'discriminator' in checkpoint:
+            discriminator.load_state_dict(checkpoint['discriminator'])
+        if optimizer_mapping is not None and 'optimizer_mapping' in checkpoint:
+            optimizer_mapping.load_state_dict(checkpoint['optimizer_mapping'])
+        if optimizer_discriminator is not None and 'optimizer_discriminator' in checkpoint:
+            optimizer_discriminator.load_state_dict(checkpoint['optimizer_discriminator'])
+
+        return checkpoint.get('epoch', 0)