Update app_parallel.py

app_parallel.py

@@ -113,6 +113,7 @@ def process_chunk(audio_chunk, preprocessed_data, args):
     with open(crop_info_path , "rb") as f:
             crop_info = pickle.load(f)
 
+    print(f"Loaded existing preprocessed data")
     print("first_coeff_path",first_coeff_path)
     print("crop_pic_path",crop_pic_path)
     print("crop_info",crop_info)
@@ -157,20 +158,19 @@ def save_uploaded_file(file, filename,TEMP_DIR):
     return file_path
 
 
-def custom_cleanup(temp_dir, exclude_dir):
+def custom_cleanup(temp_dir):
     # Iterate over the files and directories in TEMP_DIR
     for filename in os.listdir(temp_dir):
         file_path = os.path.join(temp_dir, filename)
-        # Skip the directory we want to exclude
-        if file_path != exclude_dir:
-            try:
-                if os.path.isdir(file_path):
-                    shutil.rmtree(file_path)
-                else:
-                    os.remove(file_path)
-                print(f"Deleted: {file_path}")
-            except Exception as e:
-                print(f"Failed to delete {file_path}. Reason: {e}")
+        if os.path.isdir(file_path):
+            shutil.rmtree(file_path)
+        else:
+            os.remove(file_path)
+        print(f"Deleted: {file_path}")
+    
+    torch.cuda.empty_cache()
+    import gc
+    gc.collect()
 
 
 def generate_audio(voice_cloning, voice_gender, text_prompt):
@@ -225,8 +225,8 @@ 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)
-    fixed_temp_dir = "/tmp/preprocess_data"
-    os.makedirs(fixed_temp_dir, exist_ok=True)
+    fixed_temp_dir = "/home/user/app/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":
@@ -234,29 +234,22 @@ def run_preprocessing(args):
         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}
-
-
-        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
+
+client = OpenAI(api_key="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")
+
+def openai_chat_avatar(text_prompt):
+    response = client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[{"role": "system", "content": "Ensure answers are concise, human-like, and clear while maintaining quality. Use the fewest possible words, avoiding unnecessary articles, prepositions, and adjectives. Responses should be short but still address the question thoroughly without being verbose.Keep them to one sentence only"},
+            {"role": "user", "content": f"Hi! I need help with something. {text_prompt}"},
+        ],
+        max_tokens = len(text_prompt) + 300 # Use the length of the input text
+        # temperature=0.3,
+        # stop=["Translate:", "Text:"]
+    )
+    return response
     
 def split_audio(audio_path, chunk_duration):
     audio_clip = mp.AudioFileClip(audio_path)
@@ -277,13 +270,17 @@ def split_audio(audio_path, chunk_duration):
 def generate_chunks(audio_chunks, preprocessed_data, args):
     future_to_chunk = {executor.submit(process_chunk, chunk[1], preprocessed_data, args): chunk[0] for chunk in audio_chunks}
     
-    for future in as_completed(future_to_chunk):
-        idx = future_to_chunk[future]  # Get the original chunk that was processed
-        try:
-            base64_video, temp_file_path = future.result()  # Get the result of the completed task
-            yield json.dumps({'start_time': idx, 'path': temp_file_path}).encode('utf-8')
-        except Exception as e:
-            yield f"Task for chunk {idx} failed: {e}\n"
+    try:
+        for future in as_completed(future_to_chunk):
+            idx = future_to_chunk[future]  # Get the original chunk that was processed
+            try:
+                base64_video, temp_file_path = future.result()  # Get the result of the completed task
+                yield json.dumps({'start_time': idx, 'base64_video': base64_video}).encode('utf-8')
+            except Exception as e:
+                yield f"Task for chunk {idx} failed: {e}\n"
+    finally:
+        if TEMP_DIR:
+            custom_cleanup(TEMP_DIR.name)
 
 @app.route("/run", methods=['POST'])
 def parallel_processing():
@@ -321,6 +318,8 @@ def parallel_processing():
             print('preprocess selected: ',preprocess)
             # ref_pose_video = request.files.get('ref_pose', None)
 
+            # response = openai_chat_avatar(text_prompt)
+            # text_prompt = response.choices[0].message.content.strip()
             app.config['text_prompt'] = text_prompt
             print('Final output text prompt using openai: ',text_prompt)