Update src/facerender/animate.py

src/facerender/animate.py

@@ -128,29 +128,39 @@ 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'):
-
-        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'])
+    def load_cpk_mapping(self, checkpoint_path, mapping=None, discriminator=None,
+                     optimizer_mapping=None, optimizer_discriminator=None, device='cpu'):
+
+    # Eğer dosya değil de klasörse ve içinde archive/data.pkl varsa onu yükle
+    if os.path.isdir(checkpoint_path):
+        archive_path = os.path.join(checkpoint_path, "archive", "data.pkl")
+        if os.path.isfile(archive_path):
+            print(f"📦 data.pkl bulundu: {archive_path}")
+            checkpoint_path = archive_path
+        else:
+            raise FileNotFoundError("archive/data.pkl bulunamadı!")
+
+    # Eğer tar dosyasıysa sadece kontrol amaçlı bakıyoruz
+    elif checkpoint_path.endswith(".tar"):
+        try:
+            with tarfile.open(checkpoint_path, "r") as tar:
+                members = tar.getnames()
+                if not any("storages" in name for name in members):
+                    print("⚠️ 'storages' klasörü .tar içinde yok. Devam ediliyor...")
+        except Exception as e:
+            print(f"Tar kontrol hatası: {e}")
+
+    # Yükleme işlemi
+    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)
 
-        return checkpoint.get('epoch', 0)