Update app.py

app.py

@@ -8,6 +8,8 @@ import subprocess
 import logging
 from typing import Optional, Tuple
 import re
+import warnings
+warnings.filterwarnings("ignore")
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -30,12 +32,19 @@ class SubtitleTranslator:
         if self.translator is None:
             logger.info("Loading translation model...")
             # Use a lightweight translation model
-            self.translator = pipeline(
-                "translation", 
-                model="Helsinki-NLP/opus-mt-mul-en",
-                device=0 if self.device == "cuda" else -1,
-                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
-            )
+            try:
+                self.translator = pipeline(
+                    "translation", 
+                    model="Helsinki-NLP/opus-mt-mul-en",
+                    device=0 if self.device == "cuda" else -1
+                )
+            except Exception as e:
+                logger.warning(f"Failed to load Helsinki model, using Facebook model: {e}")
+                self.translator = pipeline(
+                    "translation",
+                    model="facebook/m2m100_418M",
+                    device=0 if self.device == "cuda" else -1
+                )
     
     def extract_audio(self, video_path: str) -> str:
         """Extract audio from video file"""
@@ -80,32 +89,16 @@ class SubtitleTranslator:
             if source_lang == "en":
                 return text
             
-            # Chunk long text to avoid memory issues
-            max_length = 500
-            if len(text) <= max_length:
-                result = self.translator(text, max_length=512)
-                return result[0]['translation_text']
-            
-            # Process in chunks
-            sentences = re.split(r'[.!?]+', text)
-            translated_chunks = []
-            current_chunk = ""
+            # For Helsinki model, use direct translation
+            if "Helsinki" in str(type(self.translator.model)):
+                result = self.translator(text)
+                return result[0]['translation_text'] if result else text
             
-            for sentence in sentences:
-                if len(current_chunk + sentence) <= max_length:
-                    current_chunk += sentence + ". "
-                else:
-                    if current_chunk:
-                        result = self.translator(current_chunk.strip(), max_length=512)
-                        translated_chunks.append(result[0]['translation_text'])
-                    current_chunk = sentence + ". "
-            
-            if current_chunk:
-                result = self.translator(current_chunk.strip(), max_length=512)
-                translated_chunks.append(result[0]['translation_text'])
-            
-            return " ".join(translated_chunks)
-        
+            # For M2M100 model, specify target language
+            else:
+                result = self.translator(text, forced_bos_token_id=self.translator.tokenizer.get_lang_id("en"))
+                return result[0]['translation_text'] if result else text
+                
         except Exception as e:
             logger.error(f"Translation failed: {e}")
             return text  # Return original if translation fails
@@ -294,8 +287,4 @@ def create_interface():
 # Launch the app
 if __name__ == "__main__":
     demo = create_interface()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=true 
-    )
+    demo.launch(share=True)