Update app.py

app.py

@@ -8,151 +8,80 @@ from gtts import gTTS
 import soundfile as sf
 from transformers import VitsTokenizer, VitsModel, set_seed
 
-# Clone and Install IndicTransToolkit repository
-if not os.path.exists('IndicTransToolkit'):
-    os.system('git clone https://github.com/VarunGumma/IndicTransToolkit')
-    os.system('cd IndicTransToolkit && python3 -m pip install --editable ./')
+# Set Hugging Face token (via environment or user input)
+hf_token = st.text_input("Enter your Hugging Face API token", type="password")
 
-# Initialize BLIP for image captioning
-blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to("cuda" if torch.cuda.is_available() else "cpu")
+# Ensure token is provided
+if hf_token:
 
-# Function to generate captions
-def generate_caption(image_path):
-    image = Image.open(image_path).convert("RGB")
-    inputs = blip_processor(image, "image of", return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
-    with torch.no_grad():
-        generated_ids = blip_model.generate(**inputs)
-    caption = blip_processor.decode(generated_ids[0], skip_special_tokens=True)
-    return caption
+    # Initialize BLIP for image captioning
+    blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+    blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Function for translation using IndicTrans2
-def translate_caption(caption, target_languages):
-    # Load model and tokenizer
-    model_name = "ai4bharat/indictrans2-en-indic-1B"
-    tokenizer_IT2 = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-    model_IT2 = AutoModelForSeq2SeqLM.from_pretrained(model_name, trust_remote_code=True)
-    model_IT2 = torch.quantization.quantize_dynamic(
-        model_IT2, {torch.nn.Linear}, dtype=torch.qint8
-    )
-
-    ip = IndicProcessor(inference=True)
-
-    # Source language (English)
-    src_lang = "eng_Latn"
-    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-    model_IT2.to(DEVICE)  # Move model to the device
-
-    # Integrating with workflow now
-    input_sentences = [caption]
-    translations = {}
-
-    for tgt_lang in target_languages:
-        # Preprocess input sentences
-        batch = ip.preprocess_batch(input_sentences, src_lang=src_lang, tgt_lang=tgt_lang)
-
-        # Tokenize the sentences and generate input encodings
-        inputs = tokenizer_IT2(batch, truncation=True, padding="longest", return_tensors="pt").to(DEVICE)
-
-        # Generate translations using the model
+    # Function to generate captions
+    def generate_caption(image_path):
+        image = Image.open(image_path).convert("RGB")
+        inputs = blip_processor(image, "image of", return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
         with torch.no_grad():
-            generated_tokens = model_IT2.generate(
-                **inputs,
-                use_cache=True,
-                min_length=0,
-                max_length=256,
-                num_beams=5,
-                num_return_sequences=1,
-            )
-
-        # Decode the generated tokens into text
-        with tokenizer_IT2.as_target_tokenizer():
-            generated_tokens = tokenizer_IT2.batch_decode(generated_tokens.detach().cpu().tolist(), skip_special_tokens=True, clean_up_tokenization_spaces=True)
-
-        # Postprocess the translations
-        translated_texts = ip.postprocess_batch(generated_tokens, lang=tgt_lang)
-        translations[tgt_lang] = translated_texts[0]
-
-    return translations
-
-# Function to generate audio using gTTS
-def generate_audio_gtts(text, lang_code, output_file):
-    tts = gTTS(text=text, lang=lang_code)
-    tts.save(output_file)
-    return output_file
-
-# Function to generate audio using Facebook MMS-TTS
-def generate_audio_fbmms(text, model_name, output_file):
-    tokenizer = VitsTokenizer.from_pretrained(model_name)
-    model = VitsModel.from_pretrained(model_name)
-    inputs = tokenizer(text=text, return_tensors="pt")
-    set_seed(555)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    waveform = outputs.waveform[0].cpu().numpy()
-    sf.write(output_file, waveform, samplerate=model.config.sampling_rate)
-    return output_file
-
-# Streamlit UI
-st.title("Multilingual Assistive Model")
-
-uploaded_image = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
-
-if uploaded_image is not None:
-    # Display the uploaded image
-    image = Image.open(uploaded_image)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
-
-    # Generate Caption
-    st.write("Generating Caption...")
-    caption = generate_caption(uploaded_image)
-    st.write(f"Caption: {caption}")
+            generated_ids = blip_model.generate(**inputs)
+        caption = blip_processor.decode(generated_ids[0], skip_special_tokens=True)
+        return caption
+
+    # Function to load FB MMS TTS model with Hugging Face token
+    def load_fbmms_model(model_name, hf_token):
+        tokenizer = VitsTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
+        model = VitsModel.from_pretrained(model_name, use_auth_token=hf_token)
+        return tokenizer, model
+
+    # Function to generate audio using Facebook MMS-TTS
+    def generate_audio_fbmms(text, model_name, hf_token, output_file):
+        tokenizer, model = load_fbmms_model(model_name, hf_token)
+        inputs = tokenizer(text=text, return_tensors="pt")
+        set_seed(555)
+        with torch.no_grad():
+            outputs = model(**inputs)
+        waveform = outputs.waveform[0].cpu().numpy()
+        sf.write(output_file, waveform, samplerate=model.config.sampling_rate)
+        return output_file
+
+    # Streamlit UI for TTS method
+    tts_method = st.selectbox(
+        "Choose Text-to-Speech Method",
+        options=["gTTS (Google)", "Facebook MMS TTS"],
+        index=0  # Default to gTTS
+    )
 
-    # Select target languages for translation
+    # Select target languages with human-readable names
     language_options = {
         "hin_Deva": "Hindi (Devanagari)",
         "mar_Deva": "Marathi (Devanagari)",
-        "guj_Gujr": "Gujarati (Gujrati)",
-        "urd_Arab": "Urdu (Arabic)",
+        "guj_Gujr": "Gujarati (Gujarati)",
+        "urd_Arab": "Urdu (Arabic)"
     }
-    
+
     target_languages = st.multiselect(
-        "Select target languages for translation", 
-        list(language_options.keys()), 
-        ["hin_Deva", "mar_Deva"]
+        "Select target languages for translation",
+        options=list(language_options.keys()),
+        format_func=lambda x: language_options[x]
     )
 
-    # Generate Translations
-    if target_languages:
-        st.write("Translating Caption...")
-        translations = translate_caption(caption, target_languages)
-        st.write("Translations:")
+    if uploaded_image is not None and target_languages:
+        caption = generate_caption(uploaded_image)
         for lang in target_languages:
-            st.write(f"{language_options[lang]}: {translations[lang]}")
-        
-        # Select audio generation method
-        audio_method = st.radio("Choose Audio Generation Method", ("gTTS (Default)", "Facebook MMS-TTS"))
-
-        # Generate audio for each target language
-        for lang in target_languages:
-            st.write(f"Generating audio for {language_options[lang]}...")
-
-            lang_code = {
-                "hin_Deva": "hi",  # Hindi
-                "mar_Deva": "mr",  # Marathi
-                "guj_Gujr": "gu",  # Gujarati
-                "urd_Arab": "ur"   # Urdu
-            }.get(lang, "en")
-            
+            st.write(f"Generating audio for {lang}...")
             output_file = f"{lang}_audio.mp3"
             
-            if audio_method == "gTTS (Default)":
+            if tts_method == "gTTS (Google)":
+                lang_code = {
+                    "hin_Deva": "hi",
+                    "mar_Deva": "mr",
+                    "guj_Gujr": "gu",
+                    "urd_Arab": "ur"
+                }.get(lang, "en")
                 audio_file = generate_audio_gtts(translations[lang], lang_code, output_file)
             else:
-                model_name = "your_facebook_mms_model_name"  # Update this to the correct model name
-                audio_file = generate_audio_fbmms(translations[lang], model_name, output_file)
-
-            st.write(f"Playing {language_options[lang]} audio:")
-            st.audio(audio_file)
-else:
-    st.write("Upload an image to start.")
+                model_name = f"facebook/mms-tts-{lang}"
+                audio_file = generate_audio_fbmms(translations[lang], model_name, hf_token, output_file)
+            
+            if audio_file:
+                st.audio(audio_file)