Update src/facerender/animate.py

src/facerender/animate.py

@@ -2,6 +2,7 @@ import os
 import cv2
 import yaml
 import tarfile
+import tempfile
 import numpy as np
 import warnings
 from skimage import img_as_ubyte
@@ -128,36 +129,45 @@ 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'):
-
-        # .tar içeriği kontrolü
-        if checkpoint_path.endswith(".tar"):
-            try:
-                with tarfile.open(checkpoint_path, "r") as tar:
-                    members = tar.getnames()
-                    if not any(name.startswith("storages") for name in members):
-                        print("⚠️ 'storages' klasörü .tar dosyasında bulunamadı. Devam ediliyor...")
-                    else:
-                        print("✔️ 'storages' bulundu.")
-            except Exception as e:
-                print(f"Tar kontrol hatası: {e}")
-
-        # Eğer klasörse ve archive/data.pkl varsa
-        if os.path.isdir(checkpoint_path):
-            archive_path = os.path.join(checkpoint_path, "archive", "data.pkl")
-            if os.path.isfile(archive_path):
-                checkpoint_path = archive_path
-
-        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'):
+
+    # 1) Eğer .tar ile bitiyorsa, önce geçici bir klasöre aç
+    if checkpoint_path.endswith(".tar"):
+        tmpdir = tempfile.mkdtemp()
+        with tarfile.open(checkpoint_path, "r") as tar:
+            tar.extractall(path=tmpdir)
+        # içinden ilk .pth dosyasını bul ve path'i ona yönlendir
+        found = False
+        for root, _, files in os.walk(tmpdir):
+            for fname in files:
+                if fname.endswith(".pth"):
+                    checkpoint_path = os.path.join(root, fname)
+                    found = True
+                    break
+            if found: 
+                break
+        if not found:
+            raise FileNotFoundError(f"{checkpoint_path} içinde .pth dosyası bulunamadı.")
+
+    # 2) Eğer doğrudan klasör olarak yüklediysen de içindeki 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) Artık checkpoint_path kesin .pth veya .safetensors uzantılı bir dosya:
+    checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
+
+    # 4) State dict’leri yüklüyoruz
+    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)