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
@@ -133,24 +133,38 @@ class STFPipeline:
         pivot = 0
         results = []
 
+        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(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
+        
+        # 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
 
         print('STFPipeline execute 3')
         images2video(results, save_path)