Update app.py

app.py

@@ -3,19 +3,13 @@ import torch
 import shutil
 import os
 import sys
-from argparse import ArgumentParser
-from time import strftime
-from argparse import Namespace
 from src.utils.preprocess import CropAndExtract
 from src.test_audio2coeff import Audio2Coeff
 from src.facerender.animate import AnimateFromCoeff
 from src.generate_batch import get_data
 from src.generate_facerender_batch import get_facerender_data
-# from src.utils.init_path import init_path
 import tempfile
 from openai import OpenAI
-import threading
-import elevenlabs
 from elevenlabs import set_api_key, generate, play, clone, Voice, VoiceSettings
 from flask_cors import CORS, cross_origin
 # from flask_swagger_ui import get_swaggerui_blueprint
@@ -29,6 +23,12 @@ import pickle
 # from videoretalking import inference_function
 # import base64
 # import gfpgan_enhancer
+# import threading
+# import elevenlabs
+# from argparse import Namespace
+# from argparse import ArgumentParser
+# from time import strftime
+# from src.utils.init_path import init_path
 
 
 
@@ -80,216 +80,71 @@ app.config['text_prompt'] = None
 app.config['final_video_path'] = None
 app.config['final_video_duration'] = None
 
-
-
-def main(args):
-    pic_path = args.source_image
-    audio_path = args.driven_audio
-    save_dir = args.result_dir
-    pose_style = args.pose_style
-    device = args.device
-    batch_size = args.batch_size
-    input_yaw_list = args.input_yaw
-    input_pitch_list = args.input_pitch
-    input_roll_list = args.input_roll
-    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
-    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)
-
-    path_of_lm_croper = os.path.join(current_root_path, args.checkpoint_dir, 'shape_predictor_68_face_landmarks.dat')
-    path_of_net_recon_model = os.path.join(current_root_path, args.checkpoint_dir, 'epoch_20.pth')
-    dir_of_BFM_fitting = os.path.join(current_root_path, args.checkpoint_dir, 'BFM_Fitting')
-    wav2lip_checkpoint = os.path.join(current_root_path, args.checkpoint_dir, 'wav2lip.pth')
-
-    audio2pose_checkpoint = os.path.join(current_root_path, args.checkpoint_dir, 'auido2pose_00140-model.pth')
-    audio2pose_yaml_path = os.path.join(current_root_path, 'src', 'config', 'auido2pose.yaml')
-    
-    audio2exp_checkpoint = os.path.join(current_root_path, args.checkpoint_dir, 'auido2exp_00300-model.pth')
-    audio2exp_yaml_path = os.path.join(current_root_path, 'src', 'config', 'auido2exp.yaml')
-
-    free_view_checkpoint = os.path.join(current_root_path, args.checkpoint_dir, 'facevid2vid_00189-model.pth.tar')
-
-    if preprocess == 'full':
-        mapping_checkpoint = os.path.join(current_root_path, args.checkpoint_dir, 'mapping_00109-model.pth.tar')
+# Global paths
+dir_path = os.path.dirname(os.path.realpath(__file__))
+current_root_path = dir_path
+
+path_of_lm_croper = os.path.join(current_root_path, 'checkpoints', 'shape_predictor_68_face_landmarks.dat')
+path_of_net_recon_model = os.path.join(current_root_path, 'checkpoints', 'epoch_20.pth')
+dir_of_BFM_fitting = os.path.join(current_root_path, 'checkpoints', 'BFM_Fitting')
+wav2lip_checkpoint = os.path.join(current_root_path, 'checkpoints', 'wav2lip.pth')
+audio2pose_checkpoint = os.path.join(current_root_path, 'checkpoints', 'auido2pose_00140-model.pth')
+audio2pose_yaml_path = os.path.join(current_root_path, 'src', 'config', 'auido2pose.yaml')
+audio2exp_checkpoint = os.path.join(current_root_path, 'checkpoints', 'auido2exp_00300-model.pth')
+audio2exp_yaml_path = os.path.join(current_root_path, 'src', 'config', 'auido2exp.yaml')
+free_view_checkpoint = os.path.join(current_root_path, 'checkpoints', 'facevid2vid_00189-model.pth.tar')
+
+# Function for running the actual task (using preprocessed data)
+def process_chunk(audio_chunk, preprocessed_data, args):
+    print("Entered Process Chunk Function")
+    global audio2pose_checkpoint, audio2pose_yaml_path, audio2exp_checkpoint, audio2exp_yaml_path, wav2lip_checkpoint
+    global free_view_checkpoint
+    if args.preprocess == 'full':
+        mapping_checkpoint = os.path.join(current_root_path, 'checkpoints', 'mapping_00109-model.pth.tar')
         facerender_yaml_path = os.path.join(current_root_path, 'src', 'config', 'facerender_still.yaml')
     else:
-        mapping_checkpoint = os.path.join(current_root_path, args.checkpoint_dir, 'mapping_00229-model.pth.tar')
+        mapping_checkpoint = os.path.join(current_root_path, 'checkpoints', 'mapping_00229-model.pth.tar')
         facerender_yaml_path = os.path.join(current_root_path, 'src', 'config', 'facerender.yaml')
 
-
-    # preprocess_model = CropAndExtract(sadtalker_paths, device)
-    #init model
-    print(path_of_net_recon_model)
-    preprocess_model = CropAndExtract(path_of_lm_croper, path_of_net_recon_model, dir_of_BFM_fitting, device)
-
-    # audio_to_coeff = Audio2Coeff(sadtalker_paths,  device)
+    first_coeff_path = preprocessed_data["first_coeff_path"]
+    crop_pic_path = preprocessed_data["crop_pic_path"]
+    crop_info_path = "/home/user/app/preprocess_data/crop_info.json"
+    with open(crop_info_path , "rb") as f:
+            crop_info = json.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)
+    torch.cuda.empty_cache()
+    batch = get_data(first_coeff_path, audio_chunk, args.device, ref_eyeblink_coeff_path=None, still=args.still)
     audio_to_coeff = Audio2Coeff(audio2pose_checkpoint, audio2pose_yaml_path, 
                                 audio2exp_checkpoint, audio2exp_yaml_path, 
-                                wav2lip_checkpoint, device)
-    # animate_from_coeff = AnimateFromCoeff(sadtalker_paths, device)
-    animate_from_coeff = AnimateFromCoeff(free_view_checkpoint, mapping_checkpoint, 
-                                            facerender_yaml_path, device)
-
-    first_frame_dir = os.path.join(save_dir, 'first_frame_dir')
-    os.makedirs(first_frame_dir, exist_ok=True)
-    # 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)
-
-
-    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 image_hardcoded == "yes":
-        print("Loading preprocessed data...")
-        with open(preprocessed_data_path, "rb") as f:
-            preprocessed_data = pickle.load(f)
-        first_coeff_new_path = preprocessed_data["first_coeff_path"]
-        crop_pic_new_path = preprocessed_data["crop_pic_path"]
-        crop_info_path = preprocessed_data["crop_info_path"]
-        with open(crop_info_path, "rb") as f:
-            crop_info = pickle.load(f)
-            
-        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)
-        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_path": 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}")
+                                wav2lip_checkpoint, args.device)
+    coeff_path = audio_to_coeff.generate(batch, args.result_dir, args.pose_style, ref_pose_coeff_path=None)
     
-    print('first_coeff_path ',first_coeff_new_path)
-    print('crop_pic_path ',crop_pic_new_path)
-    print('crop_info ',crop_info)
-
-    if first_coeff_new_path is None:
-        print("Can't get the coeffs of the input")
-        return
-
-    if ref_eyeblink is not None:
-        ref_eyeblink_videoname = os.path.splitext(os.path.split(ref_eyeblink)[-1])[0]
-        ref_eyeblink_frame_dir = os.path.join(save_dir, ref_eyeblink_videoname)
-        os.makedirs(ref_eyeblink_frame_dir, exist_ok=True)
-        # ref_eyeblink_coeff_path, _, _ =  preprocess_model.generate(ref_eyeblink, ref_eyeblink_frame_dir, args.preprocess, source_image_flag=False)
-        ref_eyeblink_coeff_path, _, _ =  preprocess_model.generate(ref_eyeblink, ref_eyeblink_frame_dir)
-    else:
-        ref_eyeblink_coeff_path=None
-        print('ref_eyeblink_coeff_path',ref_eyeblink_coeff_path)
-
-    if ref_pose is not None:
-        if ref_pose == ref_eyeblink:
-            ref_pose_coeff_path = ref_eyeblink_coeff_path
-        else:
-            ref_pose_videoname = os.path.splitext(os.path.split(ref_pose)[-1])[0]
-            ref_pose_frame_dir = os.path.join(save_dir, ref_pose_videoname)
-            os.makedirs(ref_pose_frame_dir, exist_ok=True)
-            # ref_pose_coeff_path, _, _ =  preprocess_model.generate(ref_pose, ref_pose_frame_dir, args.preprocess, source_image_flag=False)
-            ref_pose_coeff_path, _, _ =  preprocess_model.generate(ref_pose, ref_pose_frame_dir)
-    else:
-        ref_pose_coeff_path=None
-        print('ref_eyeblink_coeff_path',ref_pose_coeff_path)
-
-    batch = get_data(first_coeff_new_path, audio_path, device, ref_eyeblink_coeff_path, still=args.still)
-    coeff_path = audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path)
-
-
-    if args.face3dvis:
-        from src.face3d.visualize import gen_composed_video
-        gen_composed_video(args, device, first_coeff_new_path, coeff_path, audio_path, os.path.join(save_dir, '3dface.mp4'))
-  
-    # data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path,
-                                # 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)
-
+    # Further processing with animate_from_coeff using the coeff_path
+    animate_from_coeff = AnimateFromCoeff(free_view_checkpoint, mapping_checkpoint, 
+                                            facerender_yaml_path, args.device)
 
-    data = get_facerender_data(coeff_path, crop_pic_new_path, first_coeff_new_path, audio_path, 
-                                batch_size, input_yaw_list, input_pitch_list, input_roll_list,
+    torch.cuda.empty_cache()
+    data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_chunk, 
+                                args.batch_size, args.input_yaw, args.input_pitch, args.input_roll, 
                                 expression_scale=args.expression_scale, still_mode=args.still, preprocess=args.preprocess)
+    torch.cuda.empty_cache()
+    print("Will Enter Animation")
+    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)
 
-    # 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)
-
-
-    result, base64_video,temp_file_path,new_audio_path = animate_from_coeff.generate(data, save_dir, pic_path, crop_info, \
-                                enhancer=args.enhancer, background_enhancer=args.background_enhancer, preprocess=args.preprocess)
-
-    # face_path = temp_file_path
-    # audio_path = new_audio_path
-    # temp_file = tempfile.NamedTemporaryFile(delete=False, dir=TEMP_DIR.name, suffix='.mp4')
-    # video_lipsync_file_path = temp_file.name
-    # output_path = video_lipsync_file_path
-
-    # # Call the function
-    # inference_function.video_lipsync_correctness(
-    #     face=face_path,
-    #     audio_path=audio_path,
-    #     face3d_net_path = path_of_net_recon_model,
-    #     outfile=output_path,
-    #     tmp_dir="temp",
-    #     crop=[0, -1, 0, -1],
-    #     re_preprocess=True,  # Set to True if you want to reprocess; False otherwise
-    #     exp_img="neutral",  # Can be 'smile', 'neutral', or path to an expression image
-    #     one_shot=False,
-    #     up_face="original",  # Options: 'original', 'sad', 'angry', 'surprise'
-    #     LNet_batch_size=16,
-    #     without_rl1=False
-    # )
-    
-    # print('The video with lip sync is generated')
-    # print("GFPGAN Activated")
-    
-    # gfpgan_enhancer.process_video_with_gfpgan(output_path, output_path)
-    # audio_clip = mp.AudioFileClip(new_audio_path) 
-    # video_clip = mp.VideoFileClip(output_path)
-    # # Combine audio and video
-    # final_clip = video_clip.set_audio(audio_clip)
-
-    # temp_file = tempfile.NamedTemporaryFile(suffix='.mp4', dir=TEMP_DIR.name, delete=False)
-    # temp_file.close()
-    # final_video_path = temp_file.name
-    # final_clip.write_videofile(final_video_path)
-
-    # with open(final_video_path, 'rb') as f:
-    #     video_content = f.read()
-
-    # base64_lipsync_video = base64.b64encode(video_content).decode('utf-8')
-
-    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
-    
-    return base64_video, temp_file_path, duration
-
-    # shutil.move(result, save_dir+'.mp4')
-
+    # app.config['final_video_duration'] = duration
+    torch.cuda.empty_cache()
+    return base64_video, temp_file_path
 
-    if not args.verbose:
-        shutil.rmtree(save_dir)
 
 def create_temp_dir():
     return tempfile.TemporaryDirectory()
@@ -305,8 +160,8 @@ client = OpenAI(api_key="sk-proj-04146TPzEmvdV6DzSxsvNM7jxOnzys5TnB7iZB0tp59B-jM
 def openai_chat_avatar(text_prompt):
     response = client.chat.completions.create(
         model="gpt-4o-mini",
-        messages=[{"role": "system", "content": "Answer in English language always using the minimum words you can ever use."},
-            {"role": "user", "content": f"Hi! I need help with something. Can you assist me with the following: {text_prompt}"},
+        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,
@@ -359,6 +214,72 @@ def custom_cleanup(temp_dir, exclude_dir):
             except Exception as e:
                 print(f"Failed to delete {file_path}. Reason: {e}")
 
+
+def generate_audio(voice_cloning, voice_gender, text_prompt):
+    print("generate_audio")
+    if voice_cloning == 'no':
+        if voice_gender == 'male':
+            voice = 'echo'
+            print('Entering Audio creation using elevenlabs')
+            set_api_key('92e149985ea2732b4359c74346c3daee')
+
+            audio = generate(text = text_prompt, voice = "Daniel", model = "eleven_multilingual_v2",stream=True, latency=4)
+            with tempfile.NamedTemporaryFile(suffix=".mp3", prefix="text_to_speech_",dir=TEMP_DIR.name, delete=False) as temp_file:
+                for chunk in audio:
+                    temp_file.write(chunk)
+                driven_audio_path = temp_file.name
+                print('driven_audio_path',driven_audio_path)
+                print('Audio file saved using elevenlabs')
+                    
+        else:
+            voice = 'nova'
+
+            print('Entering Audio creation using whisper')
+            response = client.audio.speech.create(model="tts-1-hd",
+                                            voice=voice,
+                                            input = text_prompt)
+
+            print('Audio created using whisper')
+            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)
+            print('Audio file saved using whisper')
+    
+    elif voice_cloning == 'yes':
+        set_api_key('92e149985ea2732b4359c74346c3daee')
+        # voice = clone(name = "User Cloned Voice",
+        #             files = [user_voice_path] )
+        voice = Voice(voice_id="CEii8R8RxmB0zhAiloZg",name="Marc",settings=VoiceSettings(
+                        stability=0.71, similarity_boost=0.5, style=0.0, use_speaker_boost=True),)
+
+        audio = generate(text = text_prompt, voice = voice, model = "eleven_multilingual_v2",stream=True, latency=4)
+        with tempfile.NamedTemporaryFile(suffix=".mp3", prefix="cloned_audio_",dir=TEMP_DIR.name, delete=False) as temp_file:
+            for chunk in audio:
+                temp_file.write(chunk)
+            driven_audio_path = temp_file.name
+            print('driven_audio_path',driven_audio_path)
+            # audio_duration = get_audio_duration(driven_audio_path)
+            # print('Total Audio Duration in seconds',audio_duration)
+
+    return driven_audio_path
+
+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 = "/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":
+        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)
+
+    return preprocessed_data
+
 @app.route("/run", methods=['POST'])
 def generate_video():
     global start_time
@@ -430,64 +351,8 @@ def generate_video():
             source_image_path = save_uploaded_file(source_image, 'source_image.png',TEMP_DIR)
             print(source_image_path)
         
-            if voice_cloning == 'no':
-                if voice_gender == 'male':
-                    voice = 'echo'
-                    print('Entering Audio creation using elevenlabs')
-                    set_api_key("92e149985ea2732b4359c74346c3daee")
-        
-                    audio = generate(text = text_prompt, voice = "George", model = "eleven_multilingual_v2",stream=True, latency=4)
-                    with tempfile.NamedTemporaryFile(suffix=".mp3", prefix="text_to_speech_",dir=TEMP_DIR.name, delete=False) as temp_file:
-                        for chunk in audio:
-                            temp_file.write(chunk)
-                        driven_audio_path = temp_file.name
-                        print('driven_audio_path',driven_audio_path)
-                        print('Audio file saved using elevenlabs')
-                    
-                else:
-                    voice = 'nova'
-    
-                    print('Entering Audio creation using whisper')
-                    response = client.audio.speech.create(model="tts-1-hd",
-                                                    voice=voice,
-                                                    input = text_prompt)
-    
-                    print('Audio created using whisper')
-                    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)
-                    print('Audio file saved using whisper')
-    
-            elif voice_cloning == 'yes':
-                # user_voice = request.files['user_voice']
-                # user_voice = '/home/user/app/images/marc_voice.mp3'
-    
-                # with tempfile.NamedTemporaryFile(suffix=".wav", prefix="user_voice_",dir=TEMP_DIR.name, delete=False) as temp_file:
-                #     with open(user_voice, 'rb') as source_file:
-                #         file_contents = source_file.read()
-                #         temp_file.write(file_contents)
-
-                    # temp_file.flush()
-                    # user_voice_path = temp_file.name
-                    # user_voice.save(user_voice_path)
-                    # print('user_voice_path',user_voice_path)
-    
-                set_api_key("92e149985ea2732b4359c74346c3daee")
-                # voice = clone(name = "User Cloned Voice",
-                #             files = [user_voice_path] )
-                voice = Voice(voice_id="DeZH4ash9IU9gUcNjVXh",name="Marc",settings=VoiceSettings(
-                                stability=0.71, similarity_boost=0.5, style=0.0, use_speaker_boost=True),)
-    
-                audio = generate(text = text_prompt, voice = voice, model = "eleven_multilingual_v2",stream=True, latency=4)
-                with tempfile.NamedTemporaryFile(suffix=".mp3", prefix="cloned_audio_",dir=TEMP_DIR.name, delete=False) as temp_file:
-                    for chunk in audio:
-                        temp_file.write(chunk)
-                    driven_audio_path = temp_file.name
-                    print('driven_audio_path',driven_audio_path)
-                    
-                #     elevenlabs.save(audio, driven_audio_path)
-    
+            driven_audio_path = generate_audio(voice_cloning, voice_gender, text_prompt)
+            
             save_dir = tempfile.mkdtemp(dir=TEMP_DIR.name)
             result_folder = os.path.join(save_dir, "results")
             os.makedirs(result_folder, exist_ok=True)
@@ -503,7 +368,6 @@ def generate_video():
         app.logger.error(f"An error occurred: {e}")
         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, image_hardcoded=image_hardcoded)
         
     if torch.cuda.is_available() and not args.cpu:
@@ -518,7 +382,7 @@ def generate_video():
     #                 "process_id": generation_thread.ident}
 
     try:
-        base64_video, temp_file_path, duration = main(args)
+        base64_video, temp_file_path, duration = process_chunk(driven_audio_path, preprocessed_data, args)
         final_video_path = app.config['final_video_path']
         print('final_video_path',final_video_path)