Update app_celery.py

app_celery.py

@@ -1,4 +1,4 @@
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, Response, stream_with_context
 import torch
 import shutil
 import os
@@ -532,33 +532,47 @@ def generate_video():
             chunk_tasks.append(task)
             print("chunk_tasks",chunk_tasks)
 
-        # video_chunk_paths = [task.get() for task in chunk_tasks]
+        # # 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}")
-        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
-        })
-        
+        # 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({
-        #         'base64_video': base64_video,
-        #         'text_prompt': text_prompt,
-        #         'duration': duration,
-        #         'status': 'completed'
-        #     })
+        #     '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)
+                    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