debug

src/create_modules.py

@@ -33,14 +33,11 @@ from src.utils.crop_face_single import crop_face
 
 class Processer():
     def __init__(self):
-        self.create_models()
-
+        self.a2m_model, self.pipe = self.create_models()
+    
     @spaces.GPU
     def create_models(self):
 
-        self.lmk_extractor = LMKExtractor()
-        self.vis = FaceMeshVisualizer(forehead_edge=False)
-
         config = OmegaConf.load('./configs/prompts/animation_audio.yaml')
 
         if config.weight_dtype == "fp16":
@@ -50,64 +47,69 @@ class Processer():
             
         audio_infer_config = OmegaConf.load(config.audio_inference_config)
         # prepare model
-        self.a2m_model = Audio2MeshModel(audio_infer_config['a2m_model'])
-        self.a2m_model.load_state_dict(torch.load(audio_infer_config['pretrained_model']['a2m_ckpt'], map_location="cpu"), strict=False)
-        self.a2m_model.to("cuda").eval()
+        a2m_model = Audio2MeshModel(audio_infer_config['a2m_model'])
+        a2m_model.load_state_dict(torch.load(audio_infer_config['pretrained_model']['a2m_ckpt'], map_location="cpu"), strict=False)
+        a2m_model.to("cuda").eval()
 
-        self.vae = AutoencoderKL.from_pretrained(
+        vae = AutoencoderKL.from_pretrained(
             config.pretrained_vae_path,
         ).to("cuda", dtype=weight_dtype)
 
-        self.reference_unet = UNet2DConditionModel.from_pretrained(
+        reference_unet = UNet2DConditionModel.from_pretrained(
             config.pretrained_base_model_path,
             subfolder="unet",
         ).to(dtype=weight_dtype, device="cuda")
 
         inference_config_path = config.inference_config
         infer_config = OmegaConf.load(inference_config_path)
-        self.denoising_unet = UNet3DConditionModel.from_pretrained_2d(
+        denoising_unet = UNet3DConditionModel.from_pretrained_2d(
             config.pretrained_base_model_path,
             config.motion_module_path,
             subfolder="unet",
             unet_additional_kwargs=infer_config.unet_additional_kwargs,
         ).to(dtype=weight_dtype, device="cuda")
 
-        self.pose_guider = PoseGuider(noise_latent_channels=320, use_ca=True).to(device="cuda", dtype=weight_dtype) # not use cross attention
+        pose_guider = PoseGuider(noise_latent_channels=320, use_ca=True).to(device="cuda", dtype=weight_dtype) # not use cross attention
 
-        self.image_enc = CLIPVisionModelWithProjection.from_pretrained(
+        image_enc = CLIPVisionModelWithProjection.from_pretrained(
             config.image_encoder_path
         ).to(dtype=weight_dtype, device="cuda")
 
         sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)
-        self.scheduler = DDIMScheduler(**sched_kwargs)
+        scheduler = DDIMScheduler(**sched_kwargs)
 
         # load pretrained weights
-        self.denoising_unet.load_state_dict(
+        denoising_unet.load_state_dict(
             torch.load(config.denoising_unet_path, map_location="cpu"),
             strict=False,
         )
-        self.reference_unet.load_state_dict(
+        reference_unet.load_state_dict(
             torch.load(config.reference_unet_path, map_location="cpu"),
         )
-        self.pose_guider.load_state_dict(
+        pose_guider.load_state_dict(
             torch.load(config.pose_guider_path, map_location="cpu"),
         )
 
-        self.pipe = Pose2VideoPipeline(
-            vae=self.vae,
-            image_encoder=self.image_enc,
-            reference_unet=self.reference_unet,
-            denoising_unet=self.denoising_unet,
-            pose_guider=self.pose_guider,
-            scheduler=self.scheduler,
+        pipe = Pose2VideoPipeline(
+            vae=vae,
+            image_encoder=image_enc,
+            reference_unet=reference_unet,
+            denoising_unet=denoising_unet,
+            pose_guider=pose_guider,
+            scheduler=scheduler,
         )
-        self.pipe = self.pipe.to("cuda", dtype=weight_dtype)
+        pipe = pipe.to("cuda", dtype=weight_dtype)
+        
+        return a2m_model, pipe 
         
     
     @spaces.GPU
     def audio2video(self, input_audio, ref_img, headpose_video=None, size=512, steps=25, length=150, seed=42):
         fps = 30
         cfg = 3.5
+        
+        lmk_extractor = LMKExtractor()
+        vis = FaceMeshVisualizer()
 
         config = OmegaConf.load('./configs/prompts/animation_audio.yaml')
         audio_infer_config = OmegaConf.load(config.audio_inference_config)
@@ -123,19 +125,19 @@ class Processer():
         save_dir.mkdir(exist_ok=True, parents=True)
 
         ref_image_np = cv2.cvtColor(ref_img, cv2.COLOR_RGB2BGR)
-        ref_image_np = crop_face(ref_image_np, self.lmk_extractor)
+        ref_image_np = crop_face(ref_image_np, lmk_extractor)
         if ref_image_np is None:
             return None, Image.fromarray(ref_img)
         
         ref_image_np = cv2.resize(ref_image_np, (size, size))
         ref_image_pil = Image.fromarray(cv2.cvtColor(ref_image_np, cv2.COLOR_BGR2RGB))
         
-        face_result = self.lmk_extractor(ref_image_np)
+        face_result = lmk_extractor(ref_image_np)
         if face_result is None: 
             return None, ref_image_pil
         
         lmks = face_result['lmks'].astype(np.float32)
-        ref_pose = self.vis.draw_landmarks((ref_image_np.shape[1], ref_image_np.shape[0]), lmks, normed=True)
+        ref_pose = vis.draw_landmarks((ref_image_np.shape[1], ref_image_np.shape[0]), lmks, normed=True)
         
         sample = prepare_audio_feature(input_audio, wav2vec_model_path=audio_infer_config['a2m_model']['model_path'])
         sample['audio_feature'] = torch.from_numpy(sample['audio_feature']).float().cuda()
@@ -148,7 +150,7 @@ class Processer():
         pred = pred + face_result['lmks3d']
         
         if headpose_video is not None:
-            pose_seq = get_headpose_temp(headpose_video, self.lmk_extractor)
+            pose_seq = get_headpose_temp(headpose_video, lmk_extractor)
         else:
             pose_seq = np.load(config['pose_temp'])
         mirrored_pose_seq = np.concatenate((pose_seq, pose_seq[-2:0:-1]), axis=0)
@@ -159,7 +161,7 @@ class Processer():
 
         pose_images = []
         for i, verts in enumerate(projected_vertices):
-            lmk_img = self.vis.draw_landmarks((width, height), verts, normed=False)
+            lmk_img = vis.draw_landmarks((width, height), verts, normed=False)
             pose_images.append(lmk_img)
 
         pose_list = []
@@ -210,6 +212,9 @@ class Processer():
     @spaces.GPU
     def video2video(self, ref_img, source_video, size=512, steps=25, length=150, seed=42):
         cfg = 3.5
+        
+        lmk_extractor = LMKExtractor()
+        vis = FaceMeshVisualizer()
 
         generator = torch.manual_seed(seed)
         width, height = size, size
@@ -222,19 +227,19 @@ class Processer():
         save_dir.mkdir(exist_ok=True, parents=True)
 
         ref_image_np = cv2.cvtColor(ref_img, cv2.COLOR_RGB2BGR)
-        ref_image_np = crop_face(ref_image_np, self.lmk_extractor)
+        ref_image_np = crop_face(ref_image_np, lmk_extractor)
         if ref_image_np is None:
             return None, Image.fromarray(ref_img)
         
         ref_image_np = cv2.resize(ref_image_np, (size, size))
         ref_image_pil = Image.fromarray(cv2.cvtColor(ref_image_np, cv2.COLOR_BGR2RGB))
         
-        face_result = self.lmk_extractor(ref_image_np)
+        face_result = lmk_extractor(ref_image_np)
         if face_result is None: 
             return None, ref_image_pil
         
         lmks = face_result['lmks'].astype(np.float32)
-        ref_pose = self.vis.draw_landmarks((ref_image_np.shape[1], ref_image_np.shape[0]), lmks, normed=True)
+        ref_pose = vis.draw_landmarks((ref_image_np.shape[1], ref_image_np.shape[0]), lmks, normed=True)
 
         source_images = read_frames(source_video)
         src_fps = get_fps(source_video)
@@ -257,7 +262,7 @@ class Processer():
             src_tensor_list.append(pose_transform(src_image_pil))
             src_img_np = cv2.cvtColor(np.array(src_image_pil), cv2.COLOR_RGB2BGR)
             frame_height, frame_width, _ = src_img_np.shape
-            src_img_result = self.lmk_extractor(src_img_np)
+            src_img_result = lmk_extractor(src_img_np)
             if src_img_result is None:
                 break
             pose_trans_list.append(src_img_result['trans_mat'])
@@ -291,7 +296,7 @@ class Processer():
         
         pose_list = []
         for i, verts in enumerate(projected_vertices):
-            lmk_img = self.vis.draw_landmarks((frame_width, frame_height), verts, normed=False)
+            lmk_img = vis.draw_landmarks((frame_width, frame_height), verts, normed=False)
             pose_image_np = cv2.resize(lmk_img,  (width, height))
             pose_list.append(pose_image_np)