Update main.py

main.py

@@ -4,7 +4,9 @@ from fastapi.responses import StreamingResponse
 import io
 import requests
 from PIL import Image
-import subprocess
+import ffmpeg
+import tempfile
+import os
 
 app = FastAPI()
 
@@ -40,45 +42,32 @@ async def generate_video_api(
             # Add the generated image to frames
             frames.append(Image.open(io.BytesIO(response.content)))
 
-        # Create a video from frames using PIL
-        video_frames = []
-        for frame in frames:
-            img_byte_arr = io.BytesIO()
-            frame.save(img_byte_arr, format='PNG')
-            video_frames.append(img_byte_arr.getvalue())
+        # Create a temporary directory
+        with tempfile.TemporaryDirectory() as tmpdir:
+            # Save frames as temporary files
+            frame_files = []
+            for i, frame in enumerate(frames):
+                frame_file = os.path.join(tmpdir, f"frame_{i:03d}.png")
+                frame.save(frame_file)
+                frame_files.append(frame_file)
 
-        # Use FFmpeg to combine images into a video
-        ffmpeg_cmd = [
-            'ffmpeg',
-            '-f', 'image2pipe',
-            '-framerate', str(fps),
-            '-i', '-',
-            '-c:v', 'libx264',
-            '-pix_fmt', 'yuv420p',
-            '-movflags', '+faststart',
-            '-f', 'mp4',
-            '-'
-        ]
-
-        process = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-        
-        for frame in video_frames:
-            process.stdin.write(frame)
-        
-        process.stdin.close()
-        video_content, stderr = process.communicate()
-
-        if process.returncode != 0:
-            raise Exception(f"FFmpeg error: {stderr.decode()}")
+            # Use ffmpeg-python to combine images into a video
+            input_stream = ffmpeg.input(os.path.join(tmpdir, 'frame_%03d.png'), framerate=fps)
+            output_stream = ffmpeg.output(input_stream, 'pipe:', vcodec='libx264', pix_fmt='yuv420p', format='mp4')
+            out, err = ffmpeg.run(output_stream, capture_stdout=True, capture_stderr=True)
 
         # Return the video as a streaming response
-        return StreamingResponse(io.BytesIO(video_content), media_type="video/mp4")
+        return StreamingResponse(io.BytesIO(out), media_type="video/mp4")
 
     except requests.exceptions.HTTPError as http_err:
         # Handle HTTP errors
         error_detail = f"HTTP error occurred: {http_err}"
         raise HTTPException(status_code=500, detail=error_detail)
 
+    except ffmpeg.Error as e:
+        # Handle FFmpeg errors
+        raise HTTPException(status_code=500, detail=f"FFmpeg error occurred: {e.stderr.decode()}")
+
     except Exception as err:
         # Handle any other errors
         raise HTTPException(status_code=500, detail=f"An error occurred: {err}")