Upload app_parallel.py

app_parallel.py

@@ -0,0 +1,359 @@
+from flask import Flask, request, jsonify, stream_with_context
+import torch
+import shutil
+import os
+import sys
+from time import strftime
+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 elevenlabs
+from elevenlabs import set_api_key, generate, play, clone, Voice, VoiceSettings
+import uuid
+import time
+from PIL import Image
+import moviepy.editor as mp
+import requests
+import json
+import pickle
+from dotenv import load_dotenv
+from concurrent.futures import ProcessPoolExecutor, as_completed
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Initialize ProcessPoolExecutor for parallel processing
+executor = ProcessPoolExecutor(max_workers=3)
+
+
+class AnimationConfig:
+    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
+        self.ref_pose = ref_pose_video_path
+        self.checkpoint_dir = './checkpoints'
+        self.result_dir = result_folder
+        self.pose_style = pose_style
+        self.batch_size = 2
+        self.expression_scale = expression_scale
+        self.input_yaw = None
+        self.input_pitch = None
+        self.input_roll = None
+        self.enhancer = enhancer
+        self.background_enhancer = None
+        self.cpu = False
+        self.face3dvis = False
+        self.still = still  
+        self.preprocess = preprocess
+        self.verbose = False
+        self.old_version = False
+        self.net_recon = 'resnet50'
+        self.init_path = None
+        self.use_last_fc = False
+        self.bfm_folder = './checkpoints/BFM_Fitting/'
+        self.bfm_model = 'BFM_model_front.mat'
+        self.focal = 1015.
+        self.center = 112.
+        self.camera_d = 10.
+        self.z_near = 5.
+        self.z_far = 15.
+        self.device = 'cpu'
+        self.image_hardcoded = image_hardcoded
+
+
+app = Flask(__name__)
+# CORS(app)
+
+TEMP_DIR = None
+start_time = None
+
+app.config['temp_response'] = None
+app.config['generation_thread'] = None
+app.config['text_prompt'] = None
+app.config['final_video_path'] = None
+app.config['final_video_duration'] = None
+
+# 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, 'checkpoints', 'mapping_00229-model.pth.tar')
+        facerender_yaml_path = os.path.join(current_root_path, 'src', 'config', 'facerender.yaml')
+
+    first_coeff_path = preprocessed_data["first_coeff_path"]
+    crop_pic_path = preprocessed_data["crop_pic_path"]
+    crop_info = preprocessed_data["crop_info"]
+
+    print("first_coeff_path",first_coeff_path)
+    print("crop_pic_path",crop_pic_path)
+    print("crop_info",crop_info)
+
+    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, args.device)
+    coeff_path = audio_to_coeff.generate(batch, args.result_dir, args.pose_style, ref_pose_coeff_path=None)
+    
+    # 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_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)
+
+    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)
+
+    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
+
+
+def create_temp_dir():
+    return tempfile.TemporaryDirectory()
+
+def save_uploaded_file(file, filename,TEMP_DIR):
+    print("Entered save_uploaded_file")
+    unique_filename = str(uuid.uuid4()) + "_" + filename
+    file_path = os.path.join(TEMP_DIR.name, unique_filename)
+    file.save(file_path)
+    return file_path
+
+
+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}")
+
+
+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(os.getenv('ELEVENLABS_API_KEY'))
+
+            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(os.getenv('ELEVENLABS_API_KEY'))
+        # 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)
+
+    return driven_audio_path
+
+# 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 = "C:/Users/fd01076/Downloads/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":
+        with open(preprocessed_data_path, "rb") as f:
+            preprocessed_data = pickle.load(f)
+        print("Loaded existing preprocessed data from:", preprocessed_data_path)
+    else:
+        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)
+
+        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)
+
+    return preprocessed_data
+
+def split_audio(audio_path, chunk_duration=5):
+    audio_clip = mp.AudioFileClip(audio_path)
+    total_duration = audio_clip.duration
+
+    audio_chunks = []
+    for start_time in range(0, int(total_duration), chunk_duration):
+        end_time = min(start_time + chunk_duration, total_duration)
+        chunk = audio_clip.subclip(start_time, end_time)
+        with tempfile.NamedTemporaryFile(suffix=f"_chunk_{start_time}-{end_time}.wav", prefix="audio_chunk_", dir=TEMP_DIR.name, delete=False) as temp_file:
+            chunk_path = temp_file.name
+            chunk.write_audiofile(chunk_path)
+            audio_chunks.append(chunk_path)
+    
+    return audio_chunks
+
+# Generator function to yield chunk results as they are processed
+def generate_chunks(audio_chunks, preprocessed_data, args):
+    future_to_chunk = {executor.submit(process_chunk, chunk, preprocessed_data, args): chunk for chunk in audio_chunks}
+    
+    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
+            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"
+
+@app.route("/run", methods=['POST'])
+def parallel_processing():
+    global start_time
+    start_time = time.time()
+    global TEMP_DIR
+    global audio_chunks
+    TEMP_DIR = create_temp_dir()
+    print('request:',request.method)
+    try:
+        if request.method == 'POST':
+            # source_image = request.files['source_image']
+            image_path = 'C:/Users/fd01076/Downloads/marc_smile_image_videos/marc_smile_enhanced.png'
+            source_image = Image.open(image_path)
+            text_prompt = request.form['text_prompt']
+            
+            print('Input text prompt: ',text_prompt)
+            text_prompt = text_prompt.strip()
+            if not text_prompt:
+                return jsonify({'error': 'Input text prompt cannot be blank'}), 400
+                
+            voice_cloning = request.form.get('voice_cloning', 'yes')
+            image_hardcoded = request.form.get('image_hardcoded', 'no')
+            chat_model_used = request.form.get('chat_model_used', 'openai')
+            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 = float(request.form.get('expression_scale', 1))
+            enhancer = request.form.get('enhancer', None)
+            voice_gender = request.form.get('voice_gender', 'male')
+            still_str = request.form.get('still', 'False')
+            still = still_str.lower() == 'false'
+            print('still', still)
+            preprocess = request.form.get('preprocess', 'crop')
+            print('preprocess selected: ',preprocess)
+            ref_pose_video = request.files.get('ref_pose', None)
+
+            app.config['text_prompt'] = text_prompt
+            print('Final output text prompt using openai: ',text_prompt)
+    
+            source_image_path = save_uploaded_file(source_image, 'source_image.png',TEMP_DIR)
+            print(source_image_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)
+    
+            ref_pose_video_path = None
+            if ref_pose_video:
+                with tempfile.NamedTemporaryFile(suffix=".mp4", prefix="ref_pose_",dir=TEMP_DIR.name, delete=False) as temp_file:
+                    ref_pose_video_path = temp_file.name
+                    ref_pose_video.save(ref_pose_video_path)
+                    print('ref_pose_video_path',ref_pose_video_path)
+                    
+    except Exception as e:
+        app.logger.error(f"An error occurred: {e}")
+        return "An error occurred", 500
+    
+    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)
+    
+    preprocessed_data = run_preprocessing(args)
+    chunk_duration = 5
+    print(f"Splitting the audio into {chunk_duration}-second chunks...")
+    audio_chunks = split_audio(driven_audio_path, chunk_duration=chunk_duration)
+    print(f"Audio has been split into {len(audio_chunks)} chunks: {audio_chunks}")
+  
+    try:
+        return stream_with_context(generate_chunks(audio_chunks, preprocessed_data, args))
+        # base64_video, temp_file_path, duration = process_chunk(driven_audio_path, preprocessed_data, args)
+    except Exception as e:
+        return jsonify({'status': 'error', 'message': str(e)}), 500
+
+
+@app.route("/health", methods=["GET"])
+def health_status():
+    response = {"online": "true"}
+    return jsonify(response)
+
+if __name__ == '__main__':
+    app.run(debug=True)