Update app_celery.py

app_celery.py

@@ -77,6 +77,7 @@ celery.conf.update(app.config)
 
 TEMP_DIR = None
 start_time = None
+chunk_tasks = []
 
 app.config['temp_response'] = None
 app.config['generation_thread'] = None
@@ -419,6 +420,7 @@ def process_video_for_chunk(audio_chunk_path, args_dict, chunk_index):
 @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()
@@ -520,64 +522,42 @@ def generate_video():
     #     args.device = "cpu"
         
     try:
-        # base64_video, temp_file_path, duration = main(args)
-        # final_video_path = app.config['final_video_path']
-        # print('final_video_path',final_video_path)
-        chunk_tasks = []
         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
 
-        # # video_chunk_paths = [task.get() for task in chunk_tasks]
-        # # print(f"Video chunks generated: {video_chunk_paths}")
-        # video_chunk_paths = []
-        # for task in chunk_tasks:
-        #     try:
-        #         video_chunk_path = task.get()  # Wait for the task to complete
-        #         video_chunk_paths.append(video_chunk_path)
-        #     except Exception as e:
-        #         print(f"Error while fetching task result: {str(e)}")
-        #         return jsonify({'status': 'error', 'message': str(e)}), 500
-
-        # print(f"Video chunks generated: {video_chunk_paths}")
-        # preprocess_dir = os.path.join("/tmp", "preprocess_data")
-        # custom_cleanup(TEMP_DIR.name, preprocess_dir)
-
-        # print("Temporary files cleaned up, but preprocess_data is retained.")
-        # return jsonify({
-        #     'status': 'completed',
-        #     'video_chunk_paths': video_chunk_paths
-        # })
-        
-
-        @stream_with_context
-        def generate_chunks():
-            video_chunk_paths = []
-            for index, task in enumerate(chunk_tasks):
-                try:
-                    video_chunk_path = task.get()  # Wait for each task to complete
-                    video_chunk_paths.append(video_chunk_path)
-                    print("video_chunk_paths",video_chunk_paths)
-                    yield f'data: {video_chunk_path}\n\n'  # Send the chunk path to the frontend
-                    print(f"Chunk {index} generated and sent: {video_chunk_path}")
-                except Exception as e:
-                    print(f"Error while fetching task result: {str(e)}")
-                    yield f'data: error\n\n'
-
-            preprocess_dir = os.path.join("/tmp", "preprocess_data")
-            custom_cleanup(TEMP_DIR.name, preprocess_dir)
-            print("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')
         
     except Exception as e:
         return jsonify({'status': 'error', 'message': str(e)}), 500
 
+@app.route("/stream", methods=['GET'])
+def stream_video_chunks():
+    global chunk_tasks
+    
+    @stream_with_context
+    def generate_chunks():
+        video_chunk_paths = []
+        for index, task in enumerate(chunk_tasks):
+            try:
+                video_chunk_path = task.get()  # Wait for each task to complete
+                video_chunk_paths.append(video_chunk_path)
+                print("video_chunk_paths",video_chunk_paths)
+                yield f'data: {video_chunk_path}\n\n'  # Send the chunk path to the frontend
+                print(f"Chunk {index} generated and sent: {video_chunk_path}")
+            except Exception as e:
+                print(f"Error while fetching task result: {str(e)}")
+                yield f'data: error\n\n'
+
+        preprocess_dir = os.path.join("/tmp", "preprocess_data")
+        custom_cleanup(TEMP_DIR.name, preprocess_dir)
+        print("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():