Added GPU Models

main.py

@@ -14,13 +14,10 @@ from s3_handler import S3Handler
 from media_download import YoutubeDownloader
 from transcription import StableWhisper
 from summarizer import Extract_Summary, AudioBookNarration
-from audiobook import AudioBook
+# from audiobook import AudioBook
 
 
 
-# For Storing Models
-global MODELS
-
 ### API Configurations
 
 # Context Manager for FastAPI Start/Shutdown
@@ -31,7 +28,11 @@ async def lifespan(app: FastAPI):
     
     # Loading ML models
     print('Loading ML Models..')
+
+    # For Storing Models
+    global MODELS
     MODELS = load_models()
+    
     print('ML Models Loaded!')
     
     yield
@@ -228,7 +229,7 @@ async def get_translation(user_id: str, target_language: str = 'en'):
 
     # NLLB based Translation
     user_folder_path = os.path.join(output_folder, user_id)
-    nllb_translator = Translation(model=MODELS['transcription'], 
+    nllb_translator = Translation(model=MODELS['translation'], 
                                   transcript_dict=transcript, 
                                   source_lang=transcript['language'], 
                                   target_lang=target_language, 

transcription.py

@@ -77,9 +77,8 @@ class YouTubeTranscriptAPI(Transcription):
 
 
 class Whisper(Transcription):
-    def __init__(self, model, media_path, output_path, subtitle_format, word_level):
+    def __init__(self, media_path, output_path, subtitle_format, word_level):
         super().__init__(media_path, output_path, subtitle_format)
-        self.model = model
         self.word_level = word_level
         self.supported_subtitle_formats = ['ass', 'srt', 'vtt']
         assert(self.subtitle_format.lower() in self.supported_subtitle_formats)