step02

pipeline/pipeline.py

@@ -1,28 +1,29 @@
 # coding=utf-8
 
-from .step0 import Step0
+# from .step0 import Step0
 from .step1 import Step1
-from .step2 import Step2
+# from .step2 import Step2
 from .step3 import Step3
 from .step4 import Step4
+from .step02 import Step02
 import logging
 import re
 import os
 
 class Pipeline:
     def __init__(self, step0_model_dir, step1_mode, step2_model_dir, step2_mode, step3_mode):
-        self.step0 = Step0(step0_model_dir)
+        # self.step0 = Step0(step0_model_dir)
+        self.step02 = Step02(step0_model_dir, step2_mode)
         self.step1 = Step1(step1_mode)
-        self.step2 = Step2(step2_model_dir, step2_mode)
+        # self.step2 = Step2(step2_model_dir, step2_mode)
         self.step3 = Step3(model_type=step3_mode)
         self.step4 = Step4()
-        self.step_processors = [self.step1, self.step2, self.step3, self.step4]
         self.log = logging.getLogger(self.__class__.__name__)
         self.log.setLevel(logging.INFO)
         
 
     def run(self, video_input, output_dir, mode='s4', postp_mode='rep', prompt='', negative_prompt='', duration=10, seed=42):
-        step0_resp = self.step0.run(video_input)
+        step0_resp = self.step02.run_step0(video_input)
         step0_resp_list = re.findall(r'(Step\d:.*?)(?=Step\d:|$)', step0_resp, re.DOTALL)
         step_infos = [step_info.strip().split("\n")[0] for step_info in step0_resp_list]
         step3_temp_dir = os.path.join(output_dir, "remove_vo")
@@ -36,7 +37,7 @@ class Pipeline:
                 step_results["step1_video_path"] = step1_video_path
 
             elif step_info == 'Step2: Given a video and its generated audio, determine whether the audio contains voice-over.':
-                is_vo = self.step2.run(str(step_results["step1_video_path"]))
+                is_vo = self.step02.run_step2(str(step_results["step1_video_path"]))
                 step_results["is_vo"] = is_vo
                 if not step_results["is_vo"]: # not voice-over
                     step_results["temp_final_audio_path"] = step_results["step1_audio_path"]

pipeline/step02.py

@@ -0,0 +1,73 @@
+# coding=utf-8
+# judge voice-over
+
+from third_party.VideoLLaMA2.videollama2 import model_init, mm_infer
+import logging
+
+    
+class Step02:
+    def __init__(self, model_path, step2_mode):
+        self.modal = "video"
+        self.log = logging.getLogger(self.__class__.__name__)
+        self.log.setLevel(logging.INFO)
+        self.model, self.processor, self.tokenizer = model_init(model_path)
+        self.preprocess = self.processor[self.modal]
+        self.step2_mode = step2_mode
+
+    def run_step0(self, video_path, modal_type='v'):
+        question = f"Generate high-quality audio from video step-by-step."
+        if modal_type == "a":
+            self.model.model.vision_tower = None
+        elif modal_type == "v":
+            self.model.model.audio_tower = None
+        elif modal_type == "av":
+            pass
+        else:
+            raise NotImplementedError
+
+        self.log.info("######################################################################################################")
+        self.log.info("Generate high-quality audio from video step-by-step...")
+        audio_video_tensor = self.preprocess(video_path, va=False)
+        output = mm_infer(
+            audio_video_tensor,
+            question,
+            model=self.model,
+            tokenizer=self.tokenizer,
+            modal=self.modal,
+            do_sample=False,
+        )
+
+        return output
+
+
+    def run_step2(self, video_audio_path, modal_type='av'):
+        question = f"Given a video and its corresponding audio, determine whether the audio contains voice-over? Options: A. Yes, B. No. Choose A or B."
+        if modal_type == "a":
+            self.model.model.vision_tower = None
+        elif modal_type == "v":
+            self.model.model.audio_tower = None
+        elif modal_type == "av":
+            pass
+        else:
+            raise NotImplementedError
+        audio_video_tensor = self.preprocess(video_audio_path, va=True)
+        output = mm_infer(
+            audio_video_tensor,
+            question,
+            model=self.model,
+            tokenizer=self.tokenizer,
+            modal=self.modal,
+            do_sample=False,
+        )
+        
+        if self.step2_mode == "cot":
+            output = output.split("<CONCLUSION>")[-1][1]
+        print("1111111111111111111111111: ", output)
+        output = (output == "A")
+
+        if output:
+            self.log.info(f"The video generated by Step1 ({video_audio_path}) contains voice-over.")
+        else:
+            self.log.info(f"The video generated by Step1 ({video_audio_path}) does not contain voice-over.")
+        self.log.info("Finish Step2 successfully.\n")
+        return output