Update app_celery.py

app_celery.py

@@ -522,7 +522,7 @@ def generate_video():
     #     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}")
@@ -540,28 +540,28 @@ def generate_video():
 def stream_video_chunks():
     global chunk_tasks
     print("chunk_tasks:",chunk_tasks)
-    for index, task in enumerate(chunk_tasks):
-        print("task:",task)
     @stream_with_context
     def generate_chunks():
         video_chunk_paths = []
-        for index, task in enumerate(chunk_tasks):
-            print("task:",task)
-            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'
-                print(f"Chunk {index} generated and sent: {video_chunk_path}")
-                os.remove(video_chunk_path)
-                print(f"Deleted video chunk: {video_chunk_path}")
-            except Exception as e:
-                print(f"Error while fetching task result: {str(e)}")
-                yield f'data: error\n\n'
+        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)
-        print("Temporary files cleaned up, but preprocess_data is retained.")
+        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')