Update sadtalker_utils.py

sadtalker_utils.py

@@ -17,7 +17,7 @@ from scipy import signal
 from scipy.io import loadmat, savemat, wavfile
 import glob
 import tempfile
-from tqdm import tqdm
+import tqdm
 import math
 import torchaudio
 import urllib.request
@@ -64,6 +64,34 @@ def save_wav_util(wav, path, sr):
     wavfile.write(path, sr, wav.astype(np.int16))
 
 
+def load_state_dict_robust(model, checkpoint_path, device, model_name="model"):
+    if not os.path.exists(checkpoint_path):
+        raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+    if checkpoint_path.endswith('safetensors'):
+        checkpoint = safetensors.torch.load_file(checkpoint_path, device=device)
+    else:
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+
+    state_dict = checkpoint.get(model_name, checkpoint)
+    try:
+        model.load_state_dict(state_dict)
+    except RuntimeError as e:
+        print(f"Error loading {model_name} state_dict: {e}")
+        print(f"Trying to load state_dict with key mapping for {model_name}.")
+        model_state_dict = model.state_dict()
+        mapped_state_dict = {}
+        for key, value in state_dict.items():
+            if key in model_state_dict and model_state_dict[key].shape == value.shape:
+                mapped_state_dict[key] = value
+            else:
+                print(f"Skipping key {key} due to shape mismatch or missing in model.")
+        missing_keys, unexpected_keys = model.load_state_dict(mapped_state_dict, strict=False)
+        if missing_keys or unexpected_keys:
+            print(f"Missing keys: {missing_keys}")
+            print(f"Unexpected keys: {unexpected_keys}")
+        print(f"Successfully loaded {model_name} state_dict with key mapping.")
+
+
 class OcclusionAwareKPDetector(nn.Module):
 
     def __init__(self, kp_channels, num_kp, num_dilation_blocks, dropout_rate):
@@ -175,12 +203,6 @@ class Face3DHelper:
         return [x_min, y_min, x_max, y_max]
 
 
-class Face3DHelperOld(Face3DHelper):
-
-    def __init__(self, local_pca_path, device):
-        super(Face3DHelperOld, self).__init__(local_pca_path, device)
-
-
 class MouthDetector:
 
     def __init__(self):
@@ -258,14 +280,12 @@ class SadTalker:
         self.cfg['INPUT_IMAGE']['SIZE'] = size
         self.cfg['INPUT_IMAGE']['OLD_VERSION'] = old_version
 
-        download_model(kp_url, kp_file, checkpoint_path)
-        download_model(aud_url, aud_file, checkpoint_path)
-        download_model(wav_url, wav_file, checkpoint_path)
-        download_model(gen_url, gen_file, checkpoint_path)
-        download_model(mapx_url, mapx_file, checkpoint_path)
-        download_model(den_url, den_file, checkpoint_path)
-        download_model(GFPGAN_URL, 'GFPGANv1.4.pth', checkpoint_path)
-        download_model(REALESRGAN_URL, 'RealESRGAN_x2plus.pth', checkpoint_path)
+        for filename, url in [
+            (kp_file, kp_url), (aud_file, aud_url), (wav_file, wav_url), (gen_file, gen_url),
+            (mapx_file, mapx_url), (den_file, den_url), ('GFPGANv1.4.pth', GFPGAN_URL),
+            ('RealESRGAN_x2plus.pth', REALESRGAN_URL)
+        ]:
+            download_model(url, filename, checkpoint_path)
 
         self.sadtalker_model = SadTalkerModel(self.cfg, device_id=[0])
 
@@ -383,9 +403,6 @@ class SadTalkerInner:
         if self.still_mode or self.use_idle_mode:
             ref_pose_coeff = proc.generate_still_pose(self.pose_style)
             ref_expression_coeff = proc.generate_still_expression(self.exp_scale)
-        elif self.use_idle_mode:
-            ref_pose_coeff = proc.generate_idles_pose(self.length_of_audio, self.pose_style)
-            ref_expression_coeff = proc.generate_idles_expression(self.length_of_audio)
         else:
             ref_pose_coeff = None
             ref_expression_coeff = None
@@ -414,7 +431,7 @@ class SadTalkerInner:
         he_estimator = self.sadtalker_model.he_estimator
         audio_to_coeff = self.sadtalker_model.audio_to_coeff
         animate_from_coeff = self.sadtalker_model.animate_from_coeff
-        proc = self.sadtalker_model.preprocesser
+        face_enhancer = self.sadtalker_model.face_enhancer if self.use_enhancer else None
         with torch.no_grad():
             kp_source = kp_extractor(batch['source_image'])
             if self.still_mode or self.use_idle_mode:
@@ -444,7 +461,6 @@ class SadTalkerInner:
             kp_norm = animate_from_coeff.normalize_kp(kp_driving)
             coeff['kp_driving'] = kp_norm
             coeff['jacobian'] = [torch.eye(2).unsqueeze(0).unsqueeze(0).to(self.device)] * 4
-            face_enhancer = self.sadtalker_model.face_enhancer if self.use_enhancer else None
             output_video = animate_from_coeff.generate(batch['source_image'], kp_source, coeff, generator, mapping,
                                                         he_estimator, batch['audio'], batch['source_image_crop'],
                                                         face_enhancer=face_enhancer)
@@ -507,10 +523,7 @@ class Preprocesser:
     def __init__(self, sadtalker_cfg, device):
         self.cfg = sadtalker_cfg
         self.device = device
-        if self.cfg['INPUT_IMAGE'].get('OLD_VERSION', False):
-            self.face3d_helper = Face3DHelperOld(self.cfg['INPUT_IMAGE'].get('LOCAL_PCA_PATH', ''), device)
-        else:
-            self.face3d_helper = Face3DHelper(self.cfg['INPUT_IMAGE'].get('LOCAL_PCA_PATH', ''), device)
+        self.face3d_helper = Face3DHelper(self.cfg['INPUT_IMAGE'].get('LOCAL_PCA_PATH', ''), device)
         self.mouth_detector = MouthDetector()
 
     def crop(self, source_image_pil, preprocess_type, size=256):
@@ -607,23 +620,7 @@ class KeyPointExtractor(nn.Module):
                                                       num_dilation_blocks=2,
                                                       dropout_rate=0.1).to(device)
         checkpoint_path = os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'kp_detector.safetensors')
-        self.load_kp_detector(checkpoint_path, device)
-
-    def load_kp_detector(self, checkpoint_path, device):
-        if os.path.exists(checkpoint_path):
-            if checkpoint_path.endswith('safetensors'):
-                checkpoint = safetensors.torch.load_file(checkpoint_path, device=device)
-            else:
-                checkpoint = torch.load(checkpoint_path, map_location=device)
-            try:
-                self.kp_extractor.load_state_dict(checkpoint.get('kp_detector', {}))
-            except RuntimeError as e:
-                print(f"Error loading kp_detector state_dict: {e}")
-                print("Trying to load state_dict without prefix 'kp_detector.'")
-                self.kp_extractor.load_state_dict(checkpoint, strict=False)
-
-        else:
-            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+        load_state_dict_robust(self.kp_extractor, checkpoint_path, device, model_name='kp_detector')
 
     def forward(self, x):
         kp = self.kp_extractor(x)
@@ -636,34 +633,12 @@ class Audio2Coeff(nn.Module):
         super(Audio2Coeff, self).__init__()
         self.audio_model = Wav2Vec2Model().to(device)
         checkpoint_path = os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'wav2vec2.pth')
-        self.load_audio_model(checkpoint_path, device)
+        load_state_dict_robust(self.audio_model, checkpoint_path, device, model_name='wav2vec2')
         self.pose_mapper = AudioCoeffsPredictor(2048, 64).to(device)
         self.exp_mapper = AudioCoeffsPredictor(2048, 64).to(device)
         self.blink_mapper = AudioCoeffsPredictor(2048, 1).to(device)
-        mapping_checkpoint = os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'audio2pose_00140-model.pth')
-        self.load_mapping_model(mapping_checkpoint, device)
-
-    def load_audio_model(self, checkpoint_path, device):
-        if os.path.exists(checkpoint_path):
-            if checkpoint_path.endswith('safetensors'):
-                checkpoint = safetensors.torch.load_file(checkpoint_path, device=device)
-            else:
-                checkpoint = torch.load(checkpoint_path, map_location=device)
-            self.audio_model.load_state_dict(checkpoint.get("wav2vec2", {}))
-        else:
-            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
-
-    def load_mapping_model(self, checkpoint_path, device):
-        if os.path.exists(checkpoint_path):
-            if checkpoint_path.endswith('safetensors'):
-                checkpoint = safetensors.torch.load_file(checkpoint_path, device=device)
-            else:
-                checkpoint = torch.load(checkpoint_path, map_location=device)
-            self.pose_mapper.load_state_dict(checkpoint.get("pose_predictor", {}))
-            self.exp_mapper.load_state_dict(checkpoint.get("exp_predictor", {}))
-            self.blink_mapper.load_state_dict(checkpoint.get("blink_predictor", {}))
-        else:
-            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+        mapping_checkpoint = os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'auido2pose_00140-model.pth')
+        load_state_dict_robust(self, mapping_checkpoint, device)
 
     def get_pose_coeff(self, audio_tensor, ref_pose_coeff=None, kp_ref=None, use_ref_info=''):
         audio_embedding = self.audio_model(audio_tensor)
@@ -718,31 +693,19 @@ class AnimateFromCoeff(nn.Module):
         pose_coeff = coeff['pose_coeff']
         expression_coeff = coeff['expression_coeff']
         blink_coeff = coeff['blink_coeff']
-        with torch.no_grad():
-            if blink_coeff is not None:
-                sparse_motion = he_estimator(kp_source, kp_driving, jacobian)
-                dense_motion = sparse_motion['dense_motion']
-                video_deocclusion = generator(source_image, dense_motion, bg_param={'mask': None, 'color': None})
-                face_3d = mapping(expression_coeff, pose_coeff, blink_coeff)
-                video_3d = generator(source_image, dense_motion, bg_param={'mask': None, 'color': None},
-                                      face_3d_param=face_3d)
-                video_output = video_deocclusion['video_no_reocclusion'] + video_3d['video_3d']
-                video_output = self.make_animation(video_output)
-            else:
-                sparse_motion = he_estimator(kp_source, kp_driving, jacobian)
-                dense_motion = sparse_motion['dense_motion']
-                face_3d = mapping(expression_coeff, pose_coeff)
-                video_3d = generator(source_image, dense_motion, bg_param={'mask': None, 'color': None},
-                                      face_3d_param=face_3d)
-                video_output = video_3d['video_3d']
-                video_output = self.make_animation(video_output)
-            if face_enhancer is not None:
-                video_output_enhanced = []
-                for frame in tqdm(video_output, 'Face enhancer running'):
-                    pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
-                    enhanced_image = face_enhancer.enhance(np.array(pil_image))[0]
-                    video_output_enhanced.append(cv2.cvtColor(enhanced_image, cv2.COLOR_BGR2RGB))
-                video_output = video_output_enhanced
+        face_3d = mapping(expression_coeff, pose_coeff, blink_coeff) if blink_coeff is not None else mapping(expression_coeff, pose_coeff)
+        sparse_motion = he_estimator(kp_source, kp_driving, jacobian)
+        dense_motion = sparse_motion['dense_motion']
+        video_deocclusion = generator(source_image, dense_motion, bg_param={'mask': None, 'color': None})
+        video_3d = generator(source_image, dense_motion, bg_param={'mask': None, 'color': None}, face_3d_param=face_3d)
+        video_output = video_deocclusion['video_no_reocclusion'] + video_3d['video_3d']
+        if face_enhancer is not None:
+            video_output_enhanced = []
+            for frame in tqdm(video_output, 'Face enhancer running'):
+                pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+                enhanced_image = face_enhancer.forward(np.array(pil_image))
+                video_output_enhanced.append(cv2.cvtColor(enhanced_image, cv2.COLOR_BGR2RGB))
+            video_output = video_output_enhanced
         return video_output
 
     def make_animation(self, video_array):
@@ -767,24 +730,10 @@ class Generator(nn.Module):
                                      kp_size=10,
                                      num_deform_blocks=sadtalker_cfg['MODEL']['NUM_MOTION_FRAMES']).to(device)
         checkpoint_path = os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'generator.pth')
-        self.load_generator(checkpoint_path, device)
-
-    def load_generator(self, checkpoint_path, device):
-        if os.path.exists(checkpoint_path):
-            if checkpoint_path.endswith('safetensors'):
-                checkpoint = safetensors.torch.load_file(checkpoint_path, device=device)
-            else:
-                checkpoint = torch.load(checkpoint_path, map_location=device)
-            self.generator.load_state_dict(checkpoint.get('generator', {}))
-        else:
-            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+        load_state_dict_robust(self.generator, checkpoint_path, device, model_name='generator')
 
     def forward(self, source_image, dense_motion, bg_param, face_3d_param=None):
-        if face_3d_param is not None:
-            video_3d = self.generator(source_image, kp_driving=dense_motion, bg_param=bg_param,
-                                      face_3d_param=face_3d_param)
-        else:
-            video_3d = self.generator(source_image, kp_driving=dense_motion, bg_param=bg_param)
+        video_3d = self.generator(source_image, kp_driving=dense_motion, bg_param=bg_param, face_3d_param=face_3d_param)
         return {'video_3d': video_3d, 'video_no_reocclusion': video_3d}
 
 
@@ -794,19 +743,9 @@ class Mapping(nn.Module):
         super(Mapping, self).__init__()
         self.mapping_net = MappingNet(num_coeffs=64, num_layers=3, hidden_dim=128).to(device)
         checkpoint_path = os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'mapping.pth')
-        self.load_mapping_net(checkpoint_path, device)
+        load_state_dict_robust(self.mapping_net, checkpoint_path, device, model_name='mapping')
         self.f_3d_mean = torch.zeros(1, 64, device=device)
 
-    def load_mapping_net(self, checkpoint_path, device):
-        if os.path.exists(checkpoint_path):
-            if checkpoint_path.endswith('safetensors'):
-                checkpoint = safetensors.torch.load_file(checkpoint_path, device=device)
-            else:
-                checkpoint = torch.load(checkpoint_path, map_location=device)
-            self.mapping_net.load_state_dict(checkpoint.get('mapping', {}))
-        else:
-            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
-
     def forward(self, expression_coeff, pose_coeff, blink_coeff=None):
         coeff = torch.cat([expression_coeff, pose_coeff], dim=1)
         face_3d = self.mapping_net(coeff) + self.f_3d_mean
@@ -825,17 +764,7 @@ class OcclusionAwareDenseMotion(nn.Module):
                                                         num_blocks=sadtalker_cfg['MODEL']['NUM_MOTION_FRAMES'] - 1,
                                                         max_features=sadtalker_cfg['MODEL']['MAX_FEATURES']).to(device)
         checkpoint_path = os.path.join(sadtalker_cfg['MODEL']['CHECKPOINTS_DIR'], 'dense_motion.pth')
-        self.load_dense_motion_network(checkpoint_path, device)
-
-    def load_dense_motion_network(self, checkpoint_path, device):
-        if os.path.exists(checkpoint_path):
-            if checkpoint_path.endswith('safetensors'):
-                checkpoint = safetensors.torch.load_file(checkpoint_path, device=device)
-            else:
-                checkpoint = torch.load(checkpoint_path, map_location=device)
-            self.dense_motion_network.load_state_dict(checkpoint.get('dense_motion', {}))
-        else:
-            raise FileNotFoundError(f"Checkpoint file not found: {checkpoint_path}")
+        load_state_dict_robust(self.dense_motion_network, checkpoint_path, device, model_name='dense_motion')
 
     def forward(self, kp_source, kp_driving, jacobian):
         sparse_motion = self.dense_motion_network(kp_source, kp_driving, jacobian)
@@ -870,7 +799,10 @@ class FaceEnhancer(nn.Module):
             self.face_enhancer = None
 
     def forward(self, x):
-        return self.face_enhancer.enhance(x, outscale=1)[0]
+        if self.face_enhancer:
+            return self.face_enhancer.enhance(x, outscale=1)[0]
+        return x
+
 
 def load_models():
     checkpoint_path = './checkpoints'
@@ -883,5 +815,6 @@ def load_models():
     print("SadTalker models loaded successfully!")
     return sadtalker_instance
 
+
 if __name__ == '__main__':
     sadtalker_instance = load_models()