Update app_parallel.py

app_parallel.py

@@ -109,6 +109,8 @@ def process_chunk(audio_chunk, preprocessed_data, args):
     first_coeff_path = preprocessed_data["first_coeff_path"]
     crop_pic_path = preprocessed_data["crop_pic_path"]
     crop_info = preprocessed_data["crop_info"]
+    with open(crop_info_path, "rb") as f:
+            crop_info = pickle.load(f)
 
     print("first_coeff_path",first_coeff_path)
     print("crop_pic_path",crop_pic_path)
@@ -132,14 +134,14 @@ def process_chunk(audio_chunk, preprocessed_data, args):
     result, base64_video, temp_file_path, _ = animate_from_coeff.generate(data, args.result_dir, args.source_image, crop_info, 
                                     enhancer=args.enhancer, background_enhancer=args.background_enhancer, preprocess=args.preprocess)
 
-    video_clip = mp.VideoFileClip(temp_file_path)
-    duration = video_clip.duration
+    # video_clip = mp.VideoFileClip(temp_file_path)
+    # duration = video_clip.duration
     
     app.config['temp_response'] = base64_video
     app.config['final_video_path'] = temp_file_path
-    app.config['final_video_duration'] = duration
+    # app.config['final_video_duration'] = duration
     
-    return base64_video, temp_file_path, duration
+    return base64_video, temp_file_path
 
 
 def create_temp_dir():
@@ -219,36 +221,39 @@ def generate_audio(voice_cloning, voice_gender, text_prompt):
 # Preprocessing step that runs only once
 def run_preprocessing(args):
     global path_of_lm_croper, path_of_net_recon_model, dir_of_BFM_fitting
-    first_frame_dir = os.path.join(args.result_dir, 'first_frame_dir')
-    os.makedirs(first_frame_dir, exist_ok=True)
-
-    # Check if preprocessed data already exists
-    fixed_temp_dir = 'tmp/preprocessed_data'
+    fixed_temp_dir = "/tmp/preprocess_data"
     os.makedirs(fixed_temp_dir, exist_ok=True)
     preprocessed_data_path = os.path.join(fixed_temp_dir, "preprocessed_data.pkl")
 
     if os.path.exists(preprocessed_data_path) and args.image_hardcoded == "yes":
+        print("Loading preprocessed data...")
         with open(preprocessed_data_path, "rb") as f:
             preprocessed_data = pickle.load(f)
         print("Loaded existing preprocessed data from:", preprocessed_data_path)
     else:
+        print("Running preprocessing...")
         preprocess_model = CropAndExtract(path_of_lm_croper, path_of_net_recon_model, dir_of_BFM_fitting, args.device)
         first_coeff_path, crop_pic_path, crop_info = preprocess_model.generate(args.source_image, first_frame_dir, args.preprocess, source_image_flag=True)
+        first_coeff_new_path = os.path.join(fixed_temp_dir, os.path.basename(first_coeff_path))
+        crop_pic_new_path = os.path.join(fixed_temp_dir, os.path.basename(crop_pic_path))
+        crop_info_new_path = os.path.join(fixed_temp_dir, "crop_info.pkl")
+        shutil.move(first_coeff_path, first_coeff_new_path)
+        shutil.move(crop_pic_path, crop_pic_new_path)
+
+        with open(crop_info_new_path, "wb") as f:
+            pickle.dump(crop_info, f)
+
+        preprocessed_data = {"first_coeff_path": first_coeff_new_path,
+                            "crop_pic_path": crop_pic_new_path,
+                            "crop_info": crop_info_new_path}
 
-        if not first_coeff_path:
-            raise Exception("Failed to get coefficients")
 
-        # Save the preprocessed data
-        preprocessed_data = {
-            "first_coeff_path": first_coeff_path,
-            "crop_pic_path": crop_pic_path,
-            "crop_info": crop_info
-        }
         with open(preprocessed_data_path, "wb") as f:
             pickle.dump(preprocessed_data, f)
+        print(f"Preprocessed data saved to: {preprocessed_data_path}")
 
     return preprocessed_data
-
+    
 def split_audio(audio_path, chunk_duration=5):
     audio_clip = mp.AudioFileClip(audio_path)
     total_duration = audio_clip.duration
@@ -271,7 +276,7 @@ def generate_chunks(audio_chunks, preprocessed_data, args):
     for future in as_completed(future_to_chunk):
         chunk = future_to_chunk[future]  # Get the original chunk that was processed
         try:
-            base64_video, temp_file_path, duration = future.result()  # Get the result of the completed task
+            base64_video, temp_file_path = future.result()  # Get the result of the completed task
             yield f"Task for chunk {chunk} completed with video path: {temp_file_path}\n"
         except Exception as e:
             yield f"Task for chunk {chunk} failed: {e}\n"