Flask API & YoutubeTask class

Former-commit-id: 209a92dc549e6b55c68d8cf3fde376d6b53d149a

src/task.py

@@ -53,9 +53,12 @@ class TaskStatus(Enum):
     OUTPUT_MODULE = (auto(), None)
 
 
+
 class Task:
     def __init__(self, task_id, task_local_dir, launch_info):
+        self.status = TaskStatus.CREATED
         openai.api_key = getenv("OPENAI_API_KEY")
+        self.launch_info = launch_info
         self.task_local_dir = task_local_dir
         self.model = launch_info["model"]
         self.gpu_status = 0
@@ -63,34 +66,14 @@ class Task:
         self.task_id = task_id
         self.progress = NotImplemented
         self.SRT_Script = None
-    
+
 
     @staticmethod
     def fromYoutubeLink(youtube_url, task_id, launch_info):
         # convert to audio
-        local_dump = Path(launch_info['local_dump'])# should get from launch config
-        yt = YouTube(youtube_url)
-        video = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().first()
-        
-        if video:
-            video.download(str(local_dump.joinpath(f"task_{task_id}")), filename=f"task_{task_id}.mp4")
-            logging.info(f'Video download completed to {local_dump.joinpath(f"task_{task_id}")}!')
-        else:
-            raise FileNotFoundError(f"Video stream not found for link {youtube_url}")
-
-        audio = yt.streams.filter(only_audio=True).first()
-        if audio:
-            audio.download(str(local_dump.joinpath(f"task_{task_id}")), filename=f"task_{task_id}.mp3")
-            logging.info(f'Audio download completed to {local_dump.joinpath(f"task_{task_id}")}!')
-        else:
-            logging.info("download audio failed, using ffmpeg to extract audio")
-            subprocess.run(['ffmpeg', '-i', local_dump.joinpath(f"task_{task_id}").joinpath(f"task_{task_id}.mp4"), '-f', 'mp3', '-ab', '192000', '-vn', local_dump.joinpath(f"task_{task_id}").joinpath(f"task_{task_id}.mp3")])
-            logging.info("audio extraction finished")
-        
-        logging.info("Task Creation Complete.")
         logging.info("Task Creation method: Youtube Link")
-
-        return Task(task_id, local_dump.joinpath(f"task_{task_id}"), launch_info)
+        local_dump = Path(launch_info['local_dump'])  # should get from launch config
+        return YoutubeTask(task_id, local_dump.joinpath(f"task_{task_id}"), launch_info, youtube_url)
 
     @staticmethod
     def fromAudioFile():
@@ -133,4 +116,36 @@ class Task:
         self.translation()
         self.postprocess()
         out = self.output_render()
-        return out
+        return out
+
+class YoutubeTask(Task):
+    def __init__(self, task_id, task_local_dir, launch_info, youtube_url):
+        super().__init__(task_id, task_local_dir, launch_info)
+        self.youtube_url = youtube_url
+
+    def run(self):
+        yt = YouTube(self.youtube_url)
+        local_dump = self.task_local_dir
+        video = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().first()
+
+        if video:
+            video.download(str(local_dump.joinpath(f"task_{self.task_id}")), filename=f"task_{self.task_id}.mp4")
+            logging.info(f'Video download completed to {local_dump.joinpath(f"task_{self.task_id}")}!')
+        else:
+            raise FileNotFoundError(f"Video stream not found for link {self.youtube_url}")
+
+        audio = yt.streams.filter(only_audio=True).first()
+        if audio:
+            audio.download(str(local_dump.joinpath(f"task_{self.task_id}")), filename=f"task_{self.task_id}.mp3")
+            logging.info(f'Audio download completed to {local_dump.joinpath(f"task_{self.task_id}")}!')
+        else:
+            logging.info("download audio failed, using ffmpeg to extract audio")
+            subprocess.run(
+                ['ffmpeg', '-i', local_dump.joinpath(f"task_{self.task_id}").joinpath(f"task_{self.task_id}.mp4"), '-f', 'mp3',
+                 '-ab', '192000', '-vn', local_dump.joinpath(f"task_{self.task_id}").joinpath(f"task_{self.task_id}.mp3")])
+            logging.info("audio extraction finished")
+
+        logging.info("Task Creation Complete.")
+        logging.info("Task Creation method: Youtube Link")
+
+        super().run_pipeline()

src/web/web.py

@@ -0,0 +1,28 @@
+import yaml
+from flask import Flask, request, jsonify
+from concurrent.futures import ThreadPoolExecutor
+from src.task import Task
+
+app = Flask(__name__)
+
+# Global thread pool
+executor = ThreadPoolExecutor(max_workers=4)  # Adjust max_workers as per your requirement
+
+@app.route('/api/task', methods=['POST'])
+def create_task():
+    data = request.get_json()
+
+    if not data or 'youtube_link' not in data:
+        return jsonify({'error': 'YouTube link not provided'}), 400
+    youtube_link = data['youtube_link']
+    launch_config = yaml.load(open("./configs/local_launch.yaml"), Loader=yaml.Loader)
+    task = Task.fromYoutubeLink(youtube_link, 1, launch_config)
+
+    # Submit task to thread pool
+    executor.submit(task.run)
+
+    return jsonify({'task_id': task.task_id})
+
+
+if __name__ == '__main__':
+    app.run(debug=True)