Update src/facerender/animate.py

src/facerender/animate.py

@@ -153,24 +153,53 @@ class AnimateFromCoeff:
                          optimizer_mapping=None, optimizer_discriminator=None,
                          device='cpu'):
 
-        # 1) Eğer bir .tar checkpoint ise, direkt torch.load ile oku
-        if checkpoint_path.endswith('.tar'):
-            # Burada torch, içindeki storages klasörünü de kullanarak yükleyecek
-            checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
-
-        else:
-            # 2) Eğer klasör verilmişse archive/data.pkl'e bak
-            if os.path.isdir(checkpoint_path):
-                candidate = os.path.join(checkpoint_path, 'archive', 'data.pkl')
-                if os.path.isfile(candidate):
-                    checkpoint_path = candidate
-                else:
-                    raise FileNotFoundError(f"{checkpoint_path} bir klasör ama archive/data.pkl bulunamadı.")
-
-            # 3) .pth veya .pkl ise yine torch.load
-            checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
-
-        # 4) State dict’leri ilgili objelere yükle
+           def load_cpk_mapping(self,
+                         checkpoint_path,
+                         mapping=None,
+                         discriminator=None,
+                         optimizer_mapping=None,
+                         optimizer_discriminator=None,
+                         device='cpu'):
+        import os, tarfile, tempfile, torch
+
+        # 1) Eğer .tar veya .pth.tar ile bitiyorsa:
+        if checkpoint_path.endswith('.tar') or checkpoint_path.endswith('.pth.tar'):
+            tmpdir = tempfile.mkdtemp()
+            with tarfile.open(checkpoint_path, 'r') as tar:
+                tar.extractall(path=tmpdir)
+
+            # 1.a) Önce .pth arıyoruz, bulamazsak .pkl
+            candidate_pth = None
+            candidate_pkl = None
+            for root, _, files in os.walk(tmpdir):
+                for f in files:
+                    if f.endswith('.pth') and candidate_pth is None:
+                        candidate_pth = os.path.join(root, f)
+                    if f.endswith('.pkl') and candidate_pkl is None:
+                        candidate_pkl = os.path.join(root, f)
+                if candidate_pth:
+                    break
+
+            if candidate_pth:
+                checkpoint_path = candidate_pth
+            elif candidate_pkl:
+                checkpoint_path = candidate_pkl
+            else:
+                raise FileNotFoundError(
+                    f"{checkpoint_path} içinden ne .pth ne de .pkl dosyası bulunabildi."
+                )
+
+        # 2) Eğer checkpoint_path bir klasörse, archive/data.pkl’e bak
+        if os.path.isdir(checkpoint_path):
+            possible = os.path.join(checkpoint_path, 'archive', 'data.pkl')
+            if os.path.isfile(possible):
+                checkpoint_path = possible
+
+        # 3) Torch ile gerçek dosyayı yükle
+        checkpoint = torch.load(checkpoint_path,
+                                map_location=torch.device(device))
+
+        # 4) State dict’leri ilgili modellere ata
         if mapping is not None and 'mapping' in checkpoint:
             mapping.load_state_dict(checkpoint['mapping'])
         if discriminator is not None and 'discriminator' in checkpoint:
@@ -180,5 +209,6 @@ class AnimateFromCoeff:
         if optimizer_discriminator is not None and 'optimizer_discriminator' in checkpoint:
             optimizer_discriminator.load_state_dict(checkpoint['optimizer_discriminator'])
 
-        # 5) Epoch bilgisi varsa dön
+        # 5) Epoch bilgisi varsa dön, yoksa 0
         return checkpoint.get('epoch', 0)
+