Update src/facerender/animate.py

src/facerender/animate.py

@@ -7,8 +7,7 @@ import numpy as np
 import warnings
 from skimage import img_as_ubyte
 import safetensors
-import safetensors.torch 
-
+import safetensors.torch
 warnings.filterwarnings('ignore')
 
 import imageio
@@ -18,9 +17,9 @@ import torchvision
 from src.facerender.modules.keypoint_detector import HEEstimator, KPDetector
 from src.facerender.modules.mapping import MappingNet
 from src.facerender.modules.generator import OcclusionAwareGenerator, OcclusionAwareSPADEGenerator
-from src.facerender.modules.make_animation import make_animation 
+from src.facerender.modules.make_animation import make_animation
 
-from pydub import AudioSegment 
+from pydub import AudioSegment
 from src.utils.face_enhancer import enhancer_generator_with_len, enhancer_list
 from src.utils.paste_pic import paste_pic
 from src.utils.videoio import save_video_with_watermark
@@ -28,11 +27,11 @@ from src.utils.videoio import save_video_with_watermark
 try:
     import webui  # in webui
     in_webui = True
-except:
+except ImportError:
     in_webui = False
 
 
-class AnimateFromCoeff():
+class AnimateFromCoeff:
 
     def __init__(self, sadtalker_path, device):
         with open(sadtalker_path['facerender_yaml']) as f:
@@ -60,53 +59,72 @@ class AnimateFromCoeff():
         for param in mapping.parameters():
             param.requires_grad = False
 
-        if sadtalker_path is not None:
-            if 'checkpoint' in sadtalker_path:
-                self.load_cpk_facevid2vid_safetensor(sadtalker_path['checkpoint'], kp_detector=kp_extractor, generator=generator, he_estimator=None)
-            else:
-                self.load_cpk_facevid2vid(sadtalker_path['free_view_checkpoint'], kp_detector=kp_extractor, generator=generator, he_estimator=he_estimator)
+        # FaceVid2Vid checkpoint yükleme
+        if 'checkpoint' in sadtalker_path:
+            self.load_cpk_facevid2vid_safetensor(
+                sadtalker_path['checkpoint'],
+                kp_detector=kp_extractor,
+                generator=generator,
+                he_estimator=None,
+                device=device
+            )
         else:
-            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_facevid2vid(
+                sadtalker_path['free_view_checkpoint'],
+                kp_detector=kp_extractor,
+                generator=generator,
+                he_estimator=he_estimator,
+                device=device
+            )
+
+        # MappingNet checkpoint yükleme
+        if sadtalker_path.get('mappingnet_checkpoint') is not None:
+            self.load_cpk_mapping(
+                sadtalker_path['mappingnet_checkpoint'],
+                mapping=mapping,
+                device=device
+            )
         else:
-            raise AttributeError("Checkpoint should be specified for video head pose estimator.")
+            raise AttributeError("mappingnet_checkpoint path belirtmelisiniz.")
 
         self.kp_extractor = kp_extractor
         self.generator = generator
         self.he_estimator = he_estimator
         self.mapping = mapping
+        self.device = device
 
         self.kp_extractor.eval()
         self.generator.eval()
         self.he_estimator.eval()
         self.mapping.eval()
 
-        self.device = device
-
-    def load_cpk_facevid2vid_safetensor(self, checkpoint_path, generator=None, 
-                                        kp_detector=None, he_estimator=None,  
-                                        device="cpu"):
+    def load_cpk_facevid2vid_safetensor(self, checkpoint_path,
+                                        generator=None, kp_detector=None,
+                                        he_estimator=None, device="cpu"):
 
         checkpoint = safetensors.torch.load_file(checkpoint_path)
 
         if generator is not None:
-            x_generator = {k.replace('generator.', ''): v for k, v in checkpoint.items() if 'generator' in k}
-            generator.load_state_dict(x_generator)
+            state = {k.replace('generator.', ''): v
+                     for k, v in checkpoint.items() if k.startswith('generator.')}
+            generator.load_state_dict(state)
         if kp_detector is not None:
-            x_generator = {k.replace('kp_extractor.', ''): v for k, v in checkpoint.items() if 'kp_extractor' in k}
-            kp_detector.load_state_dict(x_generator)
+            state = {k.replace('kp_extractor.', ''): v
+                     for k, v in checkpoint.items() if k.startswith('kp_extractor.')}
+            kp_detector.load_state_dict(state)
         if he_estimator is not None:
-            x_generator = {k.replace('he_estimator.', ''): v for k, v in checkpoint.items() if 'he_estimator' in k}
-            he_estimator.load_state_dict(x_generator)
+            state = {k.replace('he_estimator.', ''): v
+                     for k, v in checkpoint.items() if k.startswith('he_estimator.')}
+            he_estimator.load_state_dict(state)
 
         return None
 
-    def load_cpk_facevid2vid(self, checkpoint_path, generator=None, discriminator=None, 
-                              kp_detector=None, he_estimator=None, optimizer_generator=None, 
-                              optimizer_discriminator=None, optimizer_kp_detector=None, 
-                              optimizer_he_estimator=None, device="cpu"):
+    def load_cpk_facevid2vid(self, checkpoint_path,
+                              generator=None, discriminator=None,
+                              kp_detector=None, he_estimator=None,
+                              optimizer_generator=None, optimizer_discriminator=None,
+                              optimizer_kp_detector=None, optimizer_he_estimator=None,
+                              device="cpu"):
 
         checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
 
@@ -118,6 +136,7 @@ class AnimateFromCoeff():
             he_estimator.load_state_dict(checkpoint['he_estimator'])
         if discriminator is not None and 'discriminator' in checkpoint:
             discriminator.load_state_dict(checkpoint['discriminator'])
+        # Optimizeler varsa yükle
         if optimizer_generator is not None and 'optimizer_generator' in checkpoint:
             optimizer_generator.load_state_dict(checkpoint['optimizer_generator'])
         if optimizer_discriminator is not None and 'optimizer_discriminator' in checkpoint:
@@ -129,45 +148,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'):
-
-    # 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
+    def load_cpk_mapping(self, checkpoint_path,
+                         mapping=None, discriminator=None,
+                         optimizer_mapping=None, optimizer_discriminator=None,
+                         device='cpu'):
+
+        # 1) .tar ise içeriği aç ve .pth bul
+        if checkpoint_path.endswith(".tar"):
+            tmpdir = tempfile.mkdtemp()
+            with tarfile.open(checkpoint_path, "r") as tar:
+                tar.extractall(path=tmpdir)
+            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 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)
+            if not found:
+                raise FileNotFoundError(f"{checkpoint_path} içinde .pth dosyası bulunamadı.")
+
+        # 2) Klasör yüklendiyse 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) checkpoint’i yükle
+        checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
+
+        # 4) State dict’leri 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:
+            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)