Update app.py

app.py

@@ -20,6 +20,9 @@ from elevenlabs import set_api_key, generate, play, clone
 from flask_cors import CORS, cross_origin
 from flask_swagger_ui import get_swaggerui_blueprint
 import uuid
+import time
+
+start_time = time.time()
 
 class AnimationConfig:
     def __init__(self, driven_audio_path, source_image_path, result_folder,pose_style,expression_scale,enhancer):
@@ -69,23 +72,22 @@ swagger_ui_blueprint = get_swaggerui_blueprint(
 )
 
 app = Flask(__name__)
+
+TEMP_DIR = None
 CORS(app)
 app.register_blueprint(swagger_ui_blueprint, url_prefix=SWAGGER_URL)
 
 app.config['temp_response'] = None
 app.config['generation_thread'] = None
 app.config['text_prompt'] = None
+app.config['final_video_path'] = None
 
-TEMP_DIR = tempfile.TemporaryDirectory()
 
 
 def main(args):
     pic_path = args.source_image
     audio_path = args.driven_audio
     save_dir = args.result_dir
-    # save_dir = os.path.join(args.result_folder, strftime("%Y_%m_%d_%H.%M.%S"))
-    # os.makedirs(save_dir, exist_ok=True)
-    print('save_dir',save_dir)
     pose_style = args.pose_style
     device = args.device
     batch_size = args.batch_size
@@ -100,7 +102,6 @@ def main(args):
     print('current_root_path ',current_root_path)
 
     sadtalker_paths = init_path(args.checkpoint_dir, os.path.join(current_root_path, 'src/config'), args.size, args.old_version, args.preprocess)
-    print('sadtalker_paths ',sadtalker_paths)
 
     
 
@@ -144,7 +145,6 @@ def main(args):
         print('ref_eyeblink_coeff_path',ref_pose_coeff_path)
 
     batch = get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path, still=args.still)
-    print('batch',batch)
     coeff_path = audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path)
 
     if args.face3dvis:
@@ -154,19 +154,15 @@ def main(args):
                                 batch_size, input_yaw_list, input_pitch_list, input_roll_list,
                                 expression_scale=args.expression_scale, still_mode=args.still, preprocess=args.preprocess, size=args.size)
 
-    print('data ',data)
-    print('save_dir ', save_dir)
-    print('pic_path ',pic_path)
-    print('crop ',crop_info)
 
-    result, base64_video = animate_from_coeff.generate(data, save_dir, pic_path, crop_info, \
+    result, base64_video,temp_file_path= animate_from_coeff.generate(data, save_dir, pic_path, crop_info, \
                                 enhancer=args.enhancer, background_enhancer=args.background_enhancer, preprocess=args.preprocess, img_size=args.size)
 
 
     print('The generated video is named:')
     app.config['temp_response'] = base64_video
-
-    return base64_video
+    app.config['final_video_path'] = temp_file_path
+    return base64_video, temp_file_path
 
     # shutil.move(result, save_dir+'.mp4')
 
@@ -174,7 +170,10 @@ def main(args):
     if not args.verbose:
         shutil.rmtree(save_dir)
 
-def save_uploaded_file(file, filename):
+def create_temp_dir():
+    return tempfile.TemporaryDirectory()
+
+def save_uploaded_file(file, filename,TEMP_DIR):
     unique_filename = str(uuid.uuid4()) + "_" + filename
     file_path = os.path.join(TEMP_DIR.name, unique_filename)
     file.save(file_path)
@@ -197,23 +196,28 @@ def translate_text(text, target_language):
 
 @app.route("/run", methods=['POST'])
 def generate_video():
+    global TEMP_DIR
+    TEMP_DIR = create_temp_dir()
     if request.method == 'POST':
         source_image = request.files['source_image']
         text_prompt = request.form['text_prompt']
+        print('Input text prompt: ',text_prompt)
         voice_cloning = request.form.get('voice_cloning', 'no')
-        target_language = request.form.get('target_language', None)
+        target_language = request.form.get('target_language', 'original_text')
+        print('target_language',target_language)
         pose_style = int(request.form.get('pose_style', 1))
         expression_scale = int(request.form.get('expression_scale', 1))
         enhancer = request.form.get('enhancer', None)
         voice_gender = request.form.get('voice_gender', 'male')
 
-        if target_language is not None:
+        if target_language != 'original_text':
             response = translate_text(text_prompt, target_language)
             text_prompt = response.choices[0].message.content.strip()
-            print('text_prompt',text_prompt)
 
         app.config['text_prompt'] = text_prompt
-        source_image_path = save_uploaded_file(source_image, 'source_image.png')
+        print('Final text prompt: ',text_prompt)
+
+        source_image_path = save_uploaded_file(source_image, 'source_image.png',TEMP_DIR)
         print(source_image_path)
 
         if voice_cloning == 'no':
@@ -226,7 +230,7 @@ def generate_video():
                                                 voice=voice,
                                                 input = text_prompt)
 
-            with tempfile.NamedTemporaryFile(suffix=".wav", prefix="text_to_speech_", delete=False) as temp_file:
+            with tempfile.NamedTemporaryFile(suffix=".wav", prefix="text_to_speech_",dir=TEMP_DIR.name, delete=False) as temp_file:
                 driven_audio_path = temp_file.name
                 
             response.write_to_file(driven_audio_path)
@@ -234,7 +238,7 @@ def generate_video():
         elif voice_cloning == 'yes':
             user_voice = request.files['user_voice']
 
-            with tempfile.NamedTemporaryFile(suffix=".wav", prefix="user_voice_", delete=False) as temp_file:
+            with tempfile.NamedTemporaryFile(suffix=".wav", prefix="user_voice_",dir=TEMP_DIR.name, delete=False) as temp_file:
                 user_voice_path = temp_file.name
                 user_voice.save(user_voice_path)
                 print('user_voice_path',user_voice_path)
@@ -244,11 +248,11 @@ def generate_video():
                         files = [user_voice_path] )
 
             audio = generate(text = text_prompt, voice = voice, model = "eleven_multilingual_v2")
-            with tempfile.NamedTemporaryFile(suffix=".mp3", prefix="cloned_audio_", delete=False) as temp_file:
+            with tempfile.NamedTemporaryFile(suffix=".mp3", prefix="cloned_audio_",dir=TEMP_DIR.name, delete=False) as temp_file:
                 driven_audio_path = temp_file.name
                 elevenlabs.save(audio, driven_audio_path)
 
-        save_dir = tempfile.mkdtemp()
+        save_dir = tempfile.mkdtemp(dir=TEMP_DIR.name)
         result_folder = os.path.join(save_dir, "results")
         os.makedirs(result_folder, exist_ok=True)
 
@@ -275,6 +279,7 @@ def generate_video():
 
 @app.route("/status", methods=["GET"])
 def check_generation_status():
+    global TEMP_DIR
     response = {"base64_video": "","text_prompt":"", "status": ""}
     process_id = request.args.get('process_id', None)
 
@@ -289,6 +294,26 @@ def check_generation_status():
             response["base64_video"] = final_response
             response["text_prompt"] = app.config.get('text_prompt')
             response["status"] = "completed"
+
+            final_video_path = app.config['final_video_path']
+            print('final_video_path',final_video_path)
+
+
+            if final_video_path and os.path.exists(final_video_path):
+                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.")
+            end_time = time.time()
+            total_time = round(end_time - start_time, 2)
+            print("Total time taken for execution:", total_time, " seconds")
             return jsonify(response)
     return jsonify({"error":"No process id provided"})