app_celery.py

from flask import Flask, request, jsonify, Response, stream_with_context
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, AsyncOpenAI
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
import uuid
import time
from PIL import Image
import moviepy.editor as mp
import requests
import json
import pickle
from celery import Celery
# from gevent import monkey
# monkey.patch_all()
import torch.multiprocessing as t
import multiprocessing
multiprocessing.set_start_method('spawn', force=True)

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 = 'cuda'
        self.image_hardcoded = image_hardcoded


app = Flask(__name__)
# CORS(app)
app.config['broker_url'] = 'redis://localhost:6379/0'
app.config['result_backend'] = 'redis://localhost:6379/0'

celery = Celery(app.name, broker=app.config['broker_url'])
celery.conf.update(app.config)

TEMP_DIR = None
start_time = None
chunk_tasks = []

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




def main(args):
    print("Entered main function")
    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')
        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')
        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)
    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}")

    first_coeff_path, crop_pic_path, crop_info =  preprocess_model.generate(pic_path, first_frame_dir, args.preprocess, source_image_flag=True)

    
    print('first_coeff_path ',first_coeff_path)
    print('crop_pic_path ',crop_pic_path)
    print('crop_info ',crop_info)

    if first_coeff_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_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_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)


    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)

    # 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)

    multiprocessing.set_start_method('spawn', force=True)
    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)


    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

# client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))

# def openai_chat_avatar(text_prompt):
#     response = client.chat.completions.create(
#         model="gpt-4o-mini",
#         messages=[{"role": "system", "content": "Answer 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}"},
#         ],
#         max_tokens = len(text_prompt) + 300 # Use the length of the input text
#         # temperature=0.3,
#         # stop=["Translate:", "Text:"]
#     )
#     return response

def ryzedb_chat_avatar(question):
    url = "https://inference.dev.ryzeai.ai/chat/stream"
    question = question + ". Summarize and Answer using the minimum words you can ever use."
    payload = json.dumps({
    "input": {
    "chat_history": [],
    "app_id": os.getenv('RYZE_APP_ID'),
    "question": question
    },
    "config": {}
    })
    headers = {
        'Content-Type': 'application/json'
    }
    
    try:
        # Send the POST request
        response = requests.request("POST", url, headers=headers, data=payload)
        
        # Check for successful request
        response.raise_for_status()
        
        # Return the response JSON
        return response.text
    
    except requests.exceptions.RequestException as e:
        print(f"An error occurred: {e}")
        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}")


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)

    return driven_audio_path

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

@celery.task
def process_video_for_chunk(audio_chunk_path, args_dict, chunk_index):
    print("Entered process_video_for_chunk")
    args = AnimationConfig(
        driven_audio_path=args_dict['driven_audio_path'],
        source_image_path=args_dict['source_image_path'],
        result_folder=args_dict['result_folder'],
        pose_style=args_dict['pose_style'],
        expression_scale=args_dict['expression_scale'],
        enhancer=args_dict['enhancer'],
        still=args_dict['still'],
        preprocess=args_dict['preprocess'],
        ref_pose_video_path=args_dict['ref_pose_video_path'],
        image_hardcoded=args_dict['image_hardcoded']
    )
    args.driven_audio = audio_chunk_path
    chunk_save_dir = os.path.join(args.result_dir, f"chunk_{chunk_index}")
    os.makedirs(chunk_save_dir, exist_ok=True)
    print("args",args)
    try:
        base64_video, video_chunk_path, duration = main(args)
        print(f"Main function returned: {video_chunk_path}, {duration}")
        return video_chunk_path
    except Exception as e:
        print(f"Error in process_video_for_chunk: {str(e)}")
        raise
    # base64_video, video_chunk_path, duration = main(args)
    # return video_chunk_path


@app.route("/run", methods=['POST'])
def generate_video():
    global start_time
    global chunk_tasks
    start_time = time.time()
    global TEMP_DIR
    TEMP_DIR = create_temp_dir()
    print('request:',request.method)
    try:
        if request.method == 'POST':
            # source_image = request.files['source_image']
            image_path = '/home/user/app/images/out.jpg'
            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', '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)
            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)

            if chat_model_used == 'ryzedb':
                response = ryzedb_chat_avatar(text_prompt)
                events  = response.split('\r\n\r\n')
                content = None
                for event in events:
                # Split each event block by "\r\n" to get the lines
                    lines = event.split('\r\n')
                    if len(lines) > 1 and lines[0] == 'event: data':
                        # Extract the JSON part from the second line and parse it
                        json_data = lines[1].replace('data: ', '')
                        try:
                            data = json.loads(json_data)
                            text_prompt = data.get('content')
                            app.config['text_prompt'] = text_prompt
                            print('Final output text prompt using ryzedb: ',text_prompt)
                            break  # Exit the loop once content is found
                        except json.JSONDecodeError:
                            continue

            else:
                # 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)
    
            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)
            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}")
        
            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)
    args_dict = {
    '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,
    'device': 'cuda' if torch.cuda.is_available() else 'cpu'}

    # if torch.cuda.is_available() and not args.cpu:
    #     args.device = "cuda"
    # else:
    #     args.device = "cpu"
    print("audio_chunks:",audio_chunks)   
    try:
        for index, audio_chunk in enumerate(audio_chunks):
            print(f"Submitting chunk {index} with audio_chunk: {audio_chunk}")
            task = process_video_for_chunk.apply_async(args=[audio_chunk, args_dict, index])
            print(f"Task {task.id} submitted for chunk {index}")
            chunk_tasks.append(task)
            print("chunk_tasks",chunk_tasks)
        return jsonify({'status': 'Video generation started'}), 200

        
    except Exception as e:
        return jsonify({'status': 'error', 'message': str(e)}), 500

@app.route("/stream", methods=['GET'])
def stream_video_chunks():
    global chunk_tasks
    print("chunk_tasks:",chunk_tasks)
    @stream_with_context
    def generate_chunks():
        video_chunk_paths = []
        unfinished_tasks = chunk_tasks[:]
        while unfinished_tasks:  # Keep running until all tasks are finished
            for task in unfinished_tasks[:]:  # Iterate over a copy of the list
                if task.ready():  # Check if the task is finished
                    try:
                        video_chunk_path = task.get()  # Get the result (chunk path)
                        video_chunk_paths.append(video_chunk_path)
                        yield f'data: {video_chunk_path}\n\n'  # Stream the chunk path to frontend
                        app.logger.info(f"Chunk generated and sent: {video_chunk_path}")
                        os.remove(video_chunk_path)  # Optionally delete the chunk after sending
                        unfinished_tasks.remove(task)  # Remove the finished task
                    except Exception as e:
                        app.logger.error(f"Error while fetching task result: {str(e)}")
                        yield f'data: error\n\n'
            time.sleep(1)  # Avoid busy waiting, check every second

        preprocess_dir = os.path.join("/tmp", "preprocess_data")
        custom_cleanup(TEMP_DIR.name, preprocess_dir)
        app.logger.info("Temporary files cleaned up, but preprocess_data is retained.")

    # Return the generator that streams the data as it becomes available
    return Response(generate_chunks(), content_type='text/event-stream')

@app.route("/health", methods=["GET"])
def health_status():
    response = {"online": "true"}
    return jsonify(response)
if __name__ == '__main__':
    app.run(debug=True)