Update main.py

main.py

@@ -334,7 +334,7 @@ def generate_story_endpoint():
         return jsonify({'error': str(e)}), 500
 
 # ---------- Video Generation Endpoint ----------
-@app.route('/api/video/generate', methods=['POST'])
+app.route('/api/video/generate', methods=['POST'])
 def generate_video_endpoint():
     try:
         # Verify user token
@@ -361,53 +361,33 @@ def generate_video_endpoint():
         if not sections:
             return jsonify({'error': 'No sections found in the story'}), 404
 
-        # Initialize storage for images and audio
+        # Download the image and audio files from their URLs.
         image_files = []
         audio_files = []
-        temp_files = []  # Keep track of all files for cleanup
-
         for section in sections:
-            try:
-                image_url = section.get("image_url")
-                audio_url = section.get("audio_url")
-
-                # Download and validate image
-                if image_url:
-                    img_resp = requests.get(image_url, timeout=10)
-                    if img_resp.status_code == 200:
-                        img = Image.open(BytesIO(img_resp.content))
-                        if img.mode != "RGB":
-                            img = img.convert("RGB")
-
-                        img_path = f"/tmp/{uuid.uuid4().hex}.jpg"
-                        img.save(img_path)
-                        image_files.append(img_path)
-                        temp_files.append(img_path)
-                    else:
-                        print(f"Warning: Failed to download image from {image_url}")
-
-                # Download and validate audio
-                if audio_url:
-                    aud_resp = requests.get(audio_url, timeout=10)
-                    if aud_resp.status_code == 200:
-                        aud_path = f"/tmp/{uuid.uuid4().hex}.mp3"
-                        with open(aud_path, "wb") as f:
-                            f.write(aud_resp.content)
-                        audio_files.append(aud_path)
-                        temp_files.append(aud_path)
-                    else:
-                        print(f"Warning: Failed to download audio from {audio_url}")
-
-            except Exception as e:
-                print(f"Error processing section: {e}")
+            image_url = section.get("image_url")
+            audio_url = section.get("audio_url")
+            # Download image
+            img_resp = requests.get(image_url)
+            if img_resp.status_code == 200:
+                img_path = f"/tmp/{uuid.uuid4().hex}.jpg"
+                with open(img_path, "wb") as f:
+                    f.write(img_resp.content)
+                image_files.append(img_path)
+            # Download audio
+            aud_resp = requests.get(audio_url)
+            if aud_resp.status_code == 200:
+                aud_path = f"/tmp/{uuid.uuid4().hex}.mp3"
+                with open(aud_path, "wb") as f:
+                    f.write(aud_resp.content)
+                audio_files.append(aud_path)
 
         if not image_files:
-            return jsonify({'error': 'No valid images available for video generation'}), 500
+            return jsonify({'error': 'No images available for video generation'}), 500
 
-        # Generate video
+        # Generate video using the video generation function (from video_gen.py)
+        from video_gen import create_video  # from your video_gen.py
         video_output_path = f"/tmp/{uuid.uuid4().hex}.mp4"
-        temp_files.append(video_output_path)
-
         video_file = create_video(image_files, audio_files, output_path=video_output_path)
         if not video_file:
             return jsonify({'error': 'Video generation failed'}), 500
@@ -416,31 +396,27 @@ def generate_video_endpoint():
         video_blob_name = f"stories/{uid}/{uuid.uuid4().hex}.mp4"
         video_url = upload_to_storage(video_file, video_blob_name)
 
-        # Update the story record with the video URL.
+        # Optionally, update the story record with the video URL.
         story_ref.update({"video_url": video_url})
 
-        # Subtract 5 Credits from the User
+          # --- Subtract 5 Credits from the User ---
         user_ref = db.reference(f"users/{uid}")
         user_data = user_ref.get() or {}
         current_credits = user_data.get("credits", 0)
         new_credits = max(0, current_credits - 5)
         user_ref.update({"credits": new_credits})
 
-        return jsonify({"video_url": video_url, "new_credits": new_credits})
+        # Clean up temporary local files.
+        for f in image_files:
+            os.remove(f)
+        for f in audio_files:
+            os.remove(f)
+        os.remove(video_file)
 
+        return jsonify({"video_url": video_url,"new_credits": new_credits})
     except Exception as e:
-        print("Error generating video:", traceback.format_exc())
         return jsonify({'error': str(e)}), 500
 
-    finally:
-        # Cleanup all temporary files
-        for file_path in temp_files:
-            try:
-                if os.path.exists(file_path):
-                    os.remove(file_path)
-            except Exception as cleanup_error:
-                print(f"Failed to delete {file_path}: {cleanup_error}")
-
 
 # ---------- Audio Generation Endpoint ----------
 @app.route('/api/audio/generate', methods=['POST'])