Update translator.py

translator.py

@@ -309,134 +309,75 @@ def handle_tts_request(request, output_dir):
         logger.debug(f"Stack trace: {traceback.format_exc()}")
         return jsonify({"error": f"Internal server error: {str(e)}"}), 500
 
-def handle_translation_request(request):
-    """Handle translation requests"""
+def handle_tts_request(request, output_dir):
+    """Handle TTS (Text-to-Speech) requests"""
     try:
         data = request.get_json()
         if not data:
-            logger.warning("⚠️ Translation endpoint called with no JSON data")
+            logger.warning("⚠️ TTS endpoint called with no JSON data")
             return jsonify({"error": "No JSON data provided"}), 400
 
-        source_text = data.get("text", "").strip()
-        source_language = data.get("source_language", "").lower()
-        target_language = data.get("target_language", "").lower()
+        text_input = data.get("text", "").strip()
+        language = data.get("language", "kapampangan").lower()
 
-        if not source_text:
-            logger.warning("⚠️ Translation request with empty text")
+        if not text_input:
+            logger.warning("⚠️ TTS request with empty text")
             return jsonify({"error": "No text provided"}), 400
 
-        # Map language names to codes
-        source_code = LANGUAGE_CODES.get(source_language, source_language)
-        target_code = LANGUAGE_CODES.get(target_language, target_language)
-
-        logger.info(f"🔄 Translating from {source_language} to {target_language}: '{source_text}'")
-
-        # Special handling for pam-fil, fil-pam, pam-tgl and tgl-pam using the phi model
-        use_phi_model = False
-        actual_source_code = source_code
-        actual_target_code = target_code
-
-        # Check if we need to use the phi model with fil replacement
-        if (source_code == "pam" and target_code == "fil") or (source_code == "fil" and target_code == "pam"):
-            use_phi_model = True
-        elif (source_code == "pam" and target_code == "tgl"):
-            use_phi_model = True
-            actual_target_code = "fil"  # Replace tgl with fil for the phi model
-        elif (source_code == "tgl" and target_code == "pam"):
-            use_phi_model = True
-            actual_source_code = "fil"  # Replace tgl with fil for the phi model
-
-        if use_phi_model:
-            model_key = "phi"
-
-            # Check if we have the phi model
-            if model_key not in translation_models or translation_models[model_key] is None:
-                logger.error(f"❌ Translation model for {model_key} not loaded")
-                return jsonify({"error": f"Translation model not available"}), 503
-
-            try:
-                # Get the phi model and tokenizer
-                model = translation_models[model_key]
-                tokenizer = translation_tokenizers[model_key]
-
-                # Prepend target language token to input
-                input_text = f">>{actual_target_code}<< {source_text}"
-
-                logger.info(f"🔄 Using phi model with input: '{input_text}'")
-
-                # Tokenize the text
-                tokenized = tokenizer(input_text, return_tensors="pt", padding=True)
-                tokenized = {k: v.to(model.device) for k, v in tokenized.items()}
-
-                with torch.no_grad():
-                    translated = model.generate(
-                        **tokenized,
-                        max_length=100,              # Reasonable output length
-                        num_beams=4,                 # Same as in training
-                        length_penalty=0.6,          # Same as in training
-                        early_stopping=True,         # Same as in training
-                        repetition_penalty=1.5,      # Add this to prevent repetition
-                        no_repeat_ngram_size=3       # Add this to prevent repetition
-                    )
-
-                # Decode the translation
-                result = tokenizer.decode(translated[0], skip_special_tokens=True)
-
-                logger.info(f"✅ Translation result: '{result}'")
-
-                return jsonify({
-                    "translated_text": result,
-                    "source_language": source_language,
-                    "target_language": target_language
-                })
-            except Exception as e:
-                logger.error(f"❌ Translation processing failed: {str(e)}")
-                logger.debug(f"Stack trace: {traceback.format_exc()}")
-                return jsonify({"error": f"Translation processing failed: {str(e)}"}), 500
-        else:
-            # Create the regular language pair key for other language pairs
-            lang_pair = f"{source_code}-{target_code}"
-
-            # Check if we have a model for this language pair
-            if lang_pair not in translation_models:
-                logger.warning(f"⚠️ No translation model available for {lang_pair}")
-                return jsonify(
-                    {"error": f"Translation from {source_language} to {target_language} is not supported yet"}), 400
-
-            if translation_models[lang_pair] is None or translation_tokenizers[lang_pair] is None:
-                logger.error(f"❌ Translation model for {lang_pair} not loaded")
-                return jsonify({"error": f"Translation model not available"}), 503
-
-            try:
-                # Regular translation process for other language pairs
-                model = translation_models[lang_pair]
-                tokenizer = translation_tokenizers[lang_pair]
+        if language not in TTS_MODELS:
+            logger.warning(f"⚠️ TTS requested for unsupported language: {language}")
+            return jsonify({"error": f"Invalid language. Available options: {list(TTS_MODELS.keys())}"}), 400
 
-                # Tokenize the text
-                tokenized = tokenizer(source_text, return_tensors="pt", padding=True)
-                tokenized = {k: v.to(model.device) for k, v in tokenized.items()}
+        if tts_models[language] is None:
+            logger.error(f"❌ TTS model for {language} not loaded")
+            return jsonify({"error": f"TTS model for {language} not available"}), 503
 
-                # Generate translation
-                with torch.no_grad():
-                    translated = model.generate(**tokenized)
+        logger.info(f"🔄 Generating TTS for language: {language}, text: '{text_input}'")
 
-                # Decode the translation
-                result = tokenizer.decode(translated[0], skip_special_tokens=True)
+        try:
+            processor = tts_processors[language]
+            model = tts_models[language]
+            inputs = processor(text_input, return_tensors="pt")
+            inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        except Exception as e:
+            logger.error(f"❌ TTS preprocessing failed: {str(e)}")
+            return jsonify({"error": f"TTS preprocessing failed: {str(e)}"}), 500
 
-                logger.info(f"✅ Translation result: '{result}'")
+        # Generate speech
+        try:
+            with torch.no_grad():
+                output = model(**inputs).waveform
+                waveform = output.squeeze().cpu().numpy()
+        except Exception as e:
+            logger.error(f"❌ TTS inference failed: {str(e)}")
+            logger.debug(f"Stack trace: {traceback.format_exc()}")
+            return jsonify({"error": f"TTS inference failed: {str(e)}"}), 500
 
-                return jsonify({
-                    "translated_text": result,
-                    "source_language": source_language,
-                    "target_language": target_language
-                })
-            except Exception as e:
-                logger.error(f"❌ Translation processing failed: {str(e)}")
-                logger.debug(f"Stack trace: {traceback.format_exc()}")
-                return jsonify({"error": f"Translation processing failed: {str(e)}"}), 500
+        # Save to file with a unique name to prevent overwriting
+        try:
+            # Create a unique filename using timestamp and text hash
+            import hashlib
+            import time
+            text_hash = hashlib.md5(text_input.encode()).hexdigest()[:8]
+            timestamp = int(time.time())
+            
+            output_filename = os.path.join(output_dir, f"{language}_{text_hash}_{timestamp}.wav")
+            sampling_rate = model.config.sampling_rate
+            sf.write(output_filename, waveform, sampling_rate)
+            logger.info(f"✅ Speech generated! File saved: {output_filename}")
+        except Exception as e:
+            logger.error(f"❌ Failed to save audio file: {str(e)}")
+            return jsonify({"error": f"Failed to save audio file: {str(e)}"}), 500
 
+        # Add cache-busting parameter to URL
+        return jsonify({
+            "message": "TTS audio generated",
+            "file_url": f"/download/{os.path.basename(output_filename)}?t={timestamp}",
+            "language": language,
+            "text_length": len(text_input)
+        })
     except Exception as e:
-        logger.error(f"❌ Unhandled exception in translation endpoint: {str(e)}")
+        logger.error(f"❌ Unhandled exception in TTS endpoint: {str(e)}")
         logger.debug(f"Stack trace: {traceback.format_exc()}")
         return jsonify({"error": f"Internal server error: {str(e)}"}), 500