Update app_parallel.py

app_parallel.py

@@ -1,4 +1,5 @@
-from flask import Flask, request, jsonify, stream_with_context
+from flask import Flask, request, jsonify, stream_with_context, send_file, send_from_directory
+import asyncio
 import torch
 import shutil
 import os
@@ -23,6 +24,7 @@ import json
 import pickle
 # from dotenv import load_dotenv
 from concurrent.futures import ProcessPoolExecutor, as_completed, ThreadPoolExecutor
+from stream_server import add_video,HLS_DIR, generate_m3u8
 
 # Load environment variables from .env file
 # load_dotenv()
@@ -174,6 +176,12 @@ def custom_cleanup(temp_dir):
     import gc
     gc.collect()
 
+def get_audio_duration(audio_path):
+    audio_clip = mp.AudioFileClip(audio_path)
+    duration_in_seconds = audio_clip.duration
+    audio_clip.close()  # Don't forget to close the clip
+    return duration_in_seconds
+
 
 def generate_audio(voice_cloning, voice_gender, text_prompt):
     print("generate_audio")
@@ -219,8 +227,10 @@ def generate_audio(voice_cloning, voice_gender, text_prompt):
                 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
+    return driven_audio_path, audio_duration
 
 def run_preprocessing(args):
     global path_of_lm_croper, path_of_net_recon_model, dir_of_BFM_fitting
@@ -273,17 +283,22 @@ def generate_chunks(audio_chunks, preprocessed_data, args):
     future_to_chunk = {executor.submit(process_chunk, chunk[1], preprocessed_data, args): chunk[0] for chunk in audio_chunks}
     
     try:
-        for future in as_completed(future_to_chunk):
+        for chunk_idx, future in enumerate(as_completed(future_to_chunk)):
             idx = future_to_chunk[future]  # Get the original chunk that was processed
             try:
                 base64_video, temp_file_path = future.result()  # Get the result of the completed task
-                yield json.dumps({'start_time': idx, 'base64_video': base64_video}).encode('utf-8')
+                # add video for streaming
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+                loop.run_until_complete(add_video(temp_file_path))
+                loop.close()
+                yield json.dumps({'start_time': idx, 'video_index': chunk_idx}).encode('utf-8')
             except Exception as e:
                 yield f"Task for chunk {idx} failed: {e}\n"
     finally:
         if TEMP_DIR:
             custom_cleanup(TEMP_DIR.name)
-
+            
 @app.route("/run", methods=['POST'])
 def parallel_processing():
     global start_time
@@ -291,6 +306,7 @@ def parallel_processing():
     start_time = time.time()
     global TEMP_DIR
     TEMP_DIR = create_temp_dir()
+    unique_id = str(uuid.uuid4())
     print('request:',request.method)
     try:
         if request.method == 'POST':
@@ -328,37 +344,30 @@ def parallel_processing():
             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)
+            driven_audio_path, audio_duration = 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)
+            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 = 3
+            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}")
+
+            os.makedirs('lives', exist_ok=True)
+            print("Entering generate m3u8")
+            generate_m3u8(audio_duration, f'lives/{unique_id}.m3u8')
+            
+            return jsonify({'video_url': f'{unique_id}.m3u8'}), 200
                     
     except Exception as e:
         app.logger.error(f"An error occurred: {e}")
         return jsonify({'status': 'error', 'message': str(e)}), 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 = 3
-    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}")
-    return jsonify({"status": "processing started, use /stream to get video chunks."}), 200
-
-    # 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("/stream", methods=["GET"])
 def stream_results():
@@ -371,6 +380,24 @@ def stream_results():
     except Exception as e:
         return jsonify({'status': 'error', 'message': str(e)}), 500
 
+@app.route("/live_stream/<string:playlist>", methods=['GET'])
+async def get_concatenated_playlist(playlist: str):
+    """
+    Endpoint to serve the concatenated HLS playlist.
+
+    Returns:
+        FileResponse: The concatenated playlist file.
+    """
+    playlist_path = os.path.join('lives', playlist)
+    if not os.path.exists(playlist_path):
+        return jsonify({'status': 'error', "msg":"Playlist not found"}), 404
+
+    return send_file(playlist_path, mimetype='application/vnd.apple.mpegurl')
+
+@app.route("/live_stream/<string:filename>", methods=["GET"])
+def live_stream(filename):
+    return send_from_directory(directory="hls_videos", filename=filename)
+
 @app.route("/health", methods=["GET"])
 def health_status():
     response = {"online": "true"}