Update app.py

app.py

@@ -32,7 +32,7 @@ import json
 
 
 class AnimationConfig:
-    def __init__(self, driven_audio_path, source_image_path, result_folder,pose_style,expression_scale,enhancer,still,preprocess,ref_pose_video_path):
+    def __init__(self, driven_audio_path, source_image_path, result_folder,pose_style,expression_scale,enhancer,still,preprocess,ref_pose_video_path, image_hardcoded):
         self.driven_audio = driven_audio_path
         self.source_image = source_image_path
         self.ref_eyeblink = None
@@ -64,6 +64,7 @@ class AnimationConfig:
         self.z_near = 5.
         self.z_far = 15.
         self.device = 'cpu'
+        self.image_hardcoded = image_hardcoded
 
 
 app = Flask(__name__)
@@ -93,6 +94,7 @@ def main(args):
     ref_eyeblink = args.ref_eyeblink
     ref_pose = args.ref_pose
     preprocess = args.preprocess
+    image_hardcoded = args.image_hardcoded
 
     dir_path = os.path.dirname(os.path.realpath(__file__))
     current_root_path = dir_path
@@ -139,8 +141,30 @@ def main(args):
     # first_coeff_path, crop_pic_path, crop_info =  preprocess_model.generate(pic_path, first_frame_dir, args.preprocess,\
                                                                             #  source_image_flag=True, pic_size=args.size)
 
-   
-    first_coeff_path, crop_pic_path, crop_info =  preprocess_model.generate(pic_path, first_frame_dir, args.preprocess, source_image_flag=True)
+    preprocess_dir = os.path.join(TEMP_DIR.name, "preprocess_data")
+    os.makedirs(preprocess_dir, exist_ok=True)
+    preprocessed_data_path = os.path.join(preprocess_dir, "preprocessed_data.pkl")
+
+    if os.path.exists(preprocessed_data_path) and image_hardcoded == "yes":
+        print("Loading preprocessed data...")
+         with open(preprocessed_data_path, "rb") as f:
+            preprocessed_data = pickle.load(f)
+        first_coeff_path = preprocessed_data["first_coeff_path"]
+        crop_pic_path = preprocessed_data["crop_pic_path"]
+        crop_info = preprocessed_data["crop_info"]
+        print(f"Loaded existing preprocessed data from: {preprocessed_data_path}")
+    
+    else:
+        print("Running preprocessing...")
+        first_coeff_path, crop_pic_path, crop_info =  preprocess_model.generate(pic_path, first_frame_dir, args.preprocess, source_image_flag=True)
+        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}")
+    
     print('first_coeff_path ',first_coeff_path)
     print('crop_pic_path ',crop_pic_path)
 
@@ -316,6 +340,20 @@ def ryzedb_chat_avatar(question):
         return None
 
 
+def custom_cleanup(temp_dir, exclude_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}")
 
 @app.route("/run", methods=['POST'])
 def generate_video():
@@ -337,6 +375,7 @@ def generate_video():
                 return jsonify({'error': 'Input text prompt cannot be blank'}), 400
                 
             voice_cloning = request.form.get('voice_cloning', 'no')
+            image_hardcoded = request.form.get('image_hardcoded', 'yes')
             chat_model_used = request.form.get('chat_model_used', 'openai')
             target_language = request.form.get('target_language', 'original_text')
             print('target_language',target_language)
@@ -458,7 +497,7 @@ def generate_video():
         return "An error occurred", 500
     
     # Example of using the class with some hypothetical paths
-    args = AnimationConfig(driven_audio_path=driven_audio_path, source_image_path=source_image_path, result_folder=result_folder, pose_style=pose_style, expression_scale=expression_scale,enhancer=enhancer,still=still,preprocess=preprocess,ref_pose_video_path=ref_pose_video_path)
+    args = AnimationConfig(driven_audio_path=driven_audio_path, source_image_path=source_image_path, result_folder=result_folder, pose_style=pose_style, expression_scale=expression_scale,enhancer=enhancer,still=still,preprocess=preprocess,ref_pose_video_path=ref_pose_video_path, image_hardcoded=image_hardcoded)
         
     if torch.cuda.is_available() and not args.cpu:
         args.device = "cuda"
@@ -506,14 +545,12 @@ def check_generation_status():
                 os.remove(final_video_path)
                 print("Deleted video file:", final_video_path)
 
-            TEMP_DIR.cleanup()
-            # print("Temporary Directory:", TEMP_DIR.name)
-            # if TEMP_DIR:
-            #     print("Contents of Temporary Directory:")
-            #     for filename in os.listdir(TEMP_DIR.name):
-            #         print(filename)
-            # else:
-            #     print("Temporary Directory is None or already cleaned up.")
+            # TEMP_DIR.cleanup()
+            preprocess_dir = os.path.join(TEMP_DIR.name, "preprocess_data")
+            custom_cleanup(TEMP_DIR.name, preprocess_dir)
+
+            print("Temporary files cleaned up, but preprocess_data is retained.")
+            
             end_time = time.time()
             total_time = round(end_time - start_time, 2)
             print("Total time taken for execution:", total_time, " seconds")