Update stf_utils.py

stf_utils.py

@@ -2,7 +2,7 @@ import torch
 import os
 from concurrent.futures import ThreadPoolExecutor
 from pydub import AudioSegment
-import cv2
+import cv2; cv2.setNumThreads(0); cv2.ocl.setUseOpenCL(False)
 from pathlib import Path
 import subprocess
 from pathlib import Path
@@ -14,7 +14,6 @@ from tqdm import tqdm
 
 import stf_alternative
 
-import spaces
 
 
 def exec_cmd(cmd):
@@ -69,138 +68,126 @@ def merge_audio_video(video_fp, audio_fp, wfp):
 
 
 class STFPipeline:
-    def __init__(
-        self,
-        stf_path: str = "/home/user/app/stf/",
-        template_video_path: str = "templates/front_one_piece_dress_nodded_cut.webm",
-        config_path: str = "front_config.json",
-        checkpoint_path: str = "089.pth",
-        root_path: str = "works",
-        wavlm_path: str = "microsoft/wavlm-large",
-        device: str = "cuda:0"
+    def __init__(self,
+                 stf_path: str = "/home/user/app/stf/",
+                 device: str = "cuda:0",
+                 template_video_path: str = "templates/front_one_piece_dress_nodded_cut.webm",
+                 config_path: str = "front_config.json",
+                 checkpoint_path: str = "089.pth",
+                 #root_path: str = "works"
+                 root_path: str = "/tmp/works",
+                 male : bool = False
+                 
     ):
-        self.device = device
-        self.stf_path = stf_path
+        #os.makedirs(root_path, exist_ok=True)
+        import shutil; shutil.copytree('/home/user/app/stf/works', '/tmp/works', dirs_exist_ok=True)
+        
+        import zipfile
+
+        if not male:
+            dir_zip='/tmp/works/preprocess/nasilhong_f_v1_front/crop_video_front_one_piece_dress_nodded_cut.zip'
+            dir_target='/tmp/works/preprocess/nasilhong_f_v1_front/'
+            zipfile.ZipFile(dir_zip, 'r').extractall(dir_target)
+            
+            dir_zip='/tmp/works/preprocess/nasilhong_f_v1_front/front_one_piece_dress_nodded_cut.zip'
+            dir_target='/tmp/works/preprocess/nasilhong_f_v1_front/'
+            zipfile.ZipFile(dir_zip, 'r').extractall(dir_target)
+        else:
+            dir_zip='/tmp/works/preprocess/Ian_v3_front/crop_video_Cam2_2309071202_0012_Natural_Looped.zip'
+            dir_target='/tmp/works/preprocess/Ian_v3_front/'
+            zipfile.ZipFile(dir_zip, 'r').extractall(dir_target)
+            
+            dir_zip='/tmp/works/preprocess/Ian_v3_front/Cam2_2309071202_0012_Natural_Looped.zip'
+            dir_target='/tmp/works/preprocess/Ian_v3_front/'
+            zipfile.ZipFile(dir_zip, 'r').extractall(dir_target)
+        
+
         self.config_path = os.path.join(stf_path, config_path)
         self.checkpoint_path = os.path.join(stf_path, checkpoint_path)
-        self.work_root_path = os.path.join(stf_path, root_path)
-        self.wavlm_path = wavlm_path
-        self.template_video_path = template_video_path
+        #self.work_root_path = os.path.join(stf_path, root_path)
+        self.work_root_path = os.path.join(root_path)
+        self.device = device
+        self.template_video_path=os.path.join(stf_path, template_video_path)
+        
+        # model = stf_alternative.create_model(
+        # config_path=config_path,
+        # checkpoint_path=checkpoint_path,
+        # work_root_path=work_root_path,
+        # device=device,
+        # wavlm_path="microsoft/wavlm-large",
+        # )
+        # self.template = stf_alternative.Template(
+        # model=model,
+        # config_path=config_path,
+        # template_video_path=template_video_path,
+        # )
 
-        # 비동기적으로 모델 로딩
-        self.model = self.load_model()
-        self.template = self.create_template()
+    
 
-    @spaces.GPU(duration=120)
-    def load_model(self):
-        """모델을 생성하고 GPU에 할당."""
+    def execute(self, audio: str):
+
+        
         model = stf_alternative.create_model(
             config_path=self.config_path,
             checkpoint_path=self.checkpoint_path,
             work_root_path=self.work_root_path,
             device=self.device,
-            wavlm_path=self.wavlm_path
+            wavlm_path="microsoft/wavlm-large",
         )
-        return model
 
-    @spaces.GPU(duration=120)
-    def create_template(self):
-        """템플릿 생성."""
-        template = stf_alternative.Template(
-            model=self.model,
+
+        self.template = stf_alternative.Template(
+            model=model,
             config_path=self.config_path,
-            template_video_path=self.template_video_path
+            template_video_path=self.template_video_path,
         )
-        return template
 
-    def execute(self, audio: str) -> str:
-        """오디오를 입력 받아 비디오를 생성."""
-        # 폴더 생성
-        Path("dubbing").mkdir(exist_ok=True)
-        save_path = os.path.join("dubbing", Path(audio).stem + "--lip.mp4")
 
+        
+        # Path("dubbing").mkdir(exist_ok=True)
+        # save_path = os.path.join("dubbing", Path(audio).stem+"--lip.mp4")
+        Path("/tmp/dubbing").mkdir(exist_ok=True)
+        save_path = os.path.join("/tmp/dubbing", Path(audio).stem+"--lip.mp4")
+        
         reader = iter(self.template._get_reader(num_skip_frames=0))
+
         audio_segment = AudioSegment.from_file(audio)
+        pivot = 0
         results = []
 
-        # 비동기 프레임 생성
-        with ThreadPoolExecutor(max_workers=4) as executor:
+        # try:
+
+        #     gen_infer = self.template.gen_infer(
+        #         audio_segment,
+        #         pivot,
+        #     )
+        #     for idx, (it, chunk) in enumerate(gen_infer, pivot):
+        #         frame = next(reader)
+        #         composed = self.template.compose(idx, frame, it)
+        #         frame_name = f"{idx}".zfill(5)+".jpg"
+        #         results.append(it['pred'])
+        #     pivot = idx + 1
+        # except StopIteration as e:
+        #     pass
+
+        
+        with ThreadPoolExecutor(1) as p:
             try:
+
                 gen_infer = self.template.gen_infer_concurrent(
-                    executor, audio_segment, 0
+                    p,
+                    audio_segment,
+                    pivot,
                 )
-                for idx, (it, _) in enumerate(gen_infer):
+                for idx, (it, chunk) in enumerate(gen_infer, pivot):
                     frame = next(reader)
                     composed = self.template.compose(idx, frame, it)
-                    results.append(it["pred"])
-            except StopIteration:
+                    frame_name = f"{idx}".zfill(5)+".jpg"
+                    results.append(it['pred'])
+                pivot = idx + 1
+            except StopIteration as e:
                 pass
 
-        self.images_to_video(results, save_path)
-        return save_path
-
-    @staticmethod
-    def images_to_video(images, output_path, fps=24):
-        """이미지 배열을 비디오로 변환."""
-        writer = imageio.get_writer(output_path, fps=fps, format="mp4", codec="libx264")
-        for i in track(range(len(images)), description="비디오 생성 중"):
-            writer.append_data(images[i])
-        writer.close()
-        print(f"비디오 저장 완료: {output_path}")
-        
-# class STFPipeline:
-#     def __init__(self,
-#                  stf_path: str = "/home/user/app/stf/",
-#                  device: str = "cuda:0",
-#                  template_video_path: str = "templates/front_one_piece_dress_nodded_cut.webm",
-#                  config_path: str = "front_config.json",
-#                  checkpoint_path: str = "089.pth",
-#                  root_path: str = "works"
-                 
-#     ):
-        
-#         config_path = os.path.join(stf_path, config_path)
-#         checkpoint_path = os.path.join(stf_path, checkpoint_path)
-#         work_root_path = os.path.join(stf_path, root_path)
-        
-#         model = stf_alternative.create_model(
-#         config_path=config_path,
-#         checkpoint_path=checkpoint_path,
-#         work_root_path=work_root_path,
-#         device=device,
-#         wavlm_path="microsoft/wavlm-large",
-#         )
-#         self.template = stf_alternative.Template(
-#         model=model,
-#         config_path=config_path,
-#         template_video_path=template_video_path,
-#         )
-    
-
-#     def execute(self, audio: str):
-#         Path("dubbing").mkdir(exist_ok=True)
-#         save_path = os.path.join("dubbing", Path(audio).stem+"--lip.mp4")
-#         reader = iter(self.template._get_reader(num_skip_frames=0))
-#         audio_segment = AudioSegment.from_file(audio)
-#         pivot = 0
-#         results = []
-#         with ThreadPoolExecutor(4) as p:
-#             try:
-
-#                 gen_infer = self.template.gen_infer_concurrent(
-#                     p,
-#                     audio_segment,
-#                     pivot,
-#                 )
-#                 for idx, (it, chunk) in enumerate(gen_infer, pivot):
-#                     frame = next(reader)
-#                     composed = self.template.compose(idx, frame, it)
-#                     frame_name = f"{idx}".zfill(5)+".jpg"
-#                     results.append(it['pred'])
-#                 pivot = idx + 1
-#             except StopIteration as e:
-#                 pass
-            
-#         images2video(results, save_path)
+        images2video(results, save_path)
                                 
-#         return save_path
+        return save_path