Update src/facerender/animate.py

src/facerender/animate.py

@@ -153,42 +153,24 @@ class AnimateFromCoeff:
                          optimizer_mapping=None, optimizer_discriminator=None,
                          device='cpu'):
 
-        # Eğer .tar ile bitiyorsa, önce aç ve içinden .pth veya .pkl ara
+        # 1) Eğer bir .tar checkpoint ise, direkt torch.load ile oku
         if checkpoint_path.endswith('.tar'):
-            tmpdir = tempfile.mkdtemp()
-            with tarfile.open(checkpoint_path, 'r') as tar:
-                tar.extractall(path=tmpdir)
-
-            found_pth = None
-            found_pkl = None
-            for root, _, files in os.walk(tmpdir):
-                for fname in files:
-                    if fname.endswith('.pth') and found_pth is None:
-                        found_pth = os.path.join(root, fname)
-                    if fname.endswith('.pkl') and found_pkl is None:
-                        found_pkl = os.path.join(root, fname)
-                if found_pth:
-                    break
-
-            if found_pth:
-                checkpoint_path = found_pth
-            elif found_pkl:
-                checkpoint_path = found_pkl
-            else:
-                raise FileNotFoundError(
-                    f"{checkpoint_path} içinden ne .pth ne de .pkl dosyası bulunabildi."
-                )
-
-        # Eğer bir klasör yoluna geldi ise (nadiren kullanılır), archive altındaki 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
-
-        # Artık checkpoint_path kesin .pth veya .pkl
-        checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
+            # Burada torch, içindeki storages klasörünü de kullanarak yükleyecek
+            checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
 
-        # State dict’leri ilgili modellere yükle
+        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
         if mapping is not None and 'mapping' in checkpoint:
             mapping.load_state_dict(checkpoint['mapping'])
         if discriminator is not None and 'discriminator' in checkpoint:
@@ -198,5 +180,5 @@ class AnimateFromCoeff:
         if optimizer_discriminator is not None and 'optimizer_discriminator' in checkpoint:
             optimizer_discriminator.load_state_dict(checkpoint['optimizer_discriminator'])
 
-        # Epoch bilgisi varsa döndür, yoksa 0
+        # 5) Epoch bilgisi varsa dön
         return checkpoint.get('epoch', 0)