Update app.py

app.py

@@ -1,122 +1,140 @@
-import spaces
-import gradio as gr
 import torch
-import os
+import librosa
 from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, AutoModelForCausalLM, AutoProcessor
 from gtts import gTTS
-from langdetect import detect
+import gradio as gr
+import spaces
+from PIL import Image
 import subprocess
-from io import BytesIO
+
+print("Using GPU for operations when available")
 
 # Install flash-attn
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-# Disable CUDA initialization at import
-os.environ['CUDA_VISIBLE_DEVICES'] = ''
-torch.set_grad_enabled(False)
-
-print("CUDA initialization disabled at import")
+# Function to safely load pipeline within a GPU-decorated function
+@spaces.GPU
+def load_pipeline(model_name, **kwargs):
+    try:
+        device = 0 if torch.cuda.is_available() else "cpu"
+        return pipeline(model=model_name, device=device, **kwargs)
+    except Exception as e:
+        print(f"Error loading {model_name} pipeline: {e}")
+        return None
 
+# Load Whisper model for speech recognition within a GPU-decorated function
 @spaces.GPU
 def load_whisper():
     try:
+        device = 0 if torch.cuda.is_available() else "cpu"
         processor = WhisperProcessor.from_pretrained("openai/whisper-small")
-        model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
+        model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").to(device)
         return processor, model
     except Exception as e:
         print(f"Error loading Whisper model: {e}")
         return None, None
 
+# Load sarvam-2b for text generation within a GPU-decorated function
 @spaces.GPU
-def load_vision_model():
-    try:
-        model_id = "microsoft/Phi-3.5-vision-instruct"
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id, trust_remote_code=True, torch_dtype=torch.float16
-        )
-        processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True, num_crops=16)
-        return model, processor
-    except Exception as e:
-        print(f"Error loading vision model: {e}")
-        return None, None
+def load_sarvam():
+    return load_pipeline('sarvamai/sarvam-2b-v0.5')
 
+# Load vision model
 @spaces.GPU
-def load_sarvam():
-    try:
-        return pipeline('sarvamai/sarvam-2b-v0.5')
-    except Exception as e:
-        print(f"Error loading Sarvam model: {e}")
-        return None
+def load_vision_model():
+    model_id = "microsoft/Phi-3.5-vision-instruct"
+    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype="auto", attn_implementation="flash_attention_2").cuda().eval()
+    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+    return model, processor
 
+# Process audio input within a GPU-decorated function
 @spaces.GPU
-def process_audio(audio_path, whisper_processor, whisper_model):
-    import librosa
+def process_audio_input(audio, whisper_processor, whisper_model):
+    if whisper_processor is None or whisper_model is None:
+        return "Error: Speech recognition model is not available. Please type your message instead."
+    
     try:
-        audio, sr = librosa.load(audio_path, sr=16000)
-        input_features = whisper_processor(audio, sampling_rate=sr, return_tensors="pt").input_features
+        audio, sr = librosa.load(audio, sr=16000)
+        input_features = whisper_processor(audio, sampling_rate=sr, return_tensors="pt").input_features.to(whisper_model.device)
         predicted_ids = whisper_model.generate(input_features)
         transcription = whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
         return transcription
     except Exception as e:
-        return f"Error processing audio: {str(e)}"
+        return f"Error processing audio: {str(e)}. Please type your message instead."
 
+# Generate response within a GPU-decorated function
 @spaces.GPU
-def process_image(image, text_prompt, vision_model, vision_processor):
+def text_to_speech(text, lang='hi'):
     try:
-        messages = [{"role": "user", "content": f"{text_prompt}\n<|image_1|>"}]
-        prompt = vision_processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = vision_processor(prompt, image, return_tensors="pt")
-        generate_ids = vision_model.generate(**inputs, max_new_tokens=1000, temperature=0.2, do_sample=True)
-        generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:]
-        response = vision_processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-        return response
+        # Use a better TTS engine for Indic languages
+        if lang in ['hi', 'bn', 'gu', 'kn', 'ml', 'mr', 'or', 'pa', 'ta', 'te']:
+            tts = gTTS(text=text, lang=lang, tld='co.in')  # Use Indian TLD
+        else:
+            tts = gTTS(text=text, lang=lang)
+        
+        tts.save("response.mp3")
+        return "response.mp3"
     except Exception as e:
-        return f"Error processing image: {str(e)}"
+        print(f"Error in text-to-speech: {str(e)}")
+        return None
+
+# Detect language (placeholder function, replace with actual implementation)
+def detect_language(text):
+    # Implement language detection logic here
+    return 'en'  # Default to English for now
 
 @spaces.GPU
 def generate_response(transcription, sarvam_pipe):
+    if sarvam_pipe is None:
+        return "Error: Text generation model is not available."
+    
     try:
+        # Generate response using the sarvam-2b model
         response = sarvam_pipe(transcription, max_length=100, num_return_sequences=1)[0]['generated_text']
         return response
     except Exception as e:
         return f"Error generating response: {str(e)}"
 
-def text_to_speech(text, lang='hi'):
+@spaces.GPU
+def process_image(image, text_input, vision_model, vision_processor):
     try:
-        tts = gTTS(text=text, lang=lang, tld='co.in')
-        tts.save("response.mp3")
-        return "response.mp3"
+        prompt = f"<|user|>\n<|image_1|>\n{text_input}<|end|>\n<|assistant|>\n"
+        image = Image.fromarray(image).convert("RGB")
+        inputs = vision_processor(prompt, image, return_tensors="pt").to("cuda:0")
+        generate_ids = vision_model.generate(**inputs, max_new_tokens=1000, eos_token_id=vision_processor.tokenizer.eos_token_id)
+        generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:]
+        response = vision_processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        return response
     except Exception as e:
-        print(f"Error in text-to-speech: {str(e)}")
-        return None
+        return f"Error processing image: {str(e)}"
 
 @spaces.GPU
-def indic_vision_assistant(input_type, audio_input, text_input, image_input):
+def multimodal_assistant(input_type, audio_input, text_input, image_input):
     try:
+        # Load models within the GPU-decorated function
         whisper_processor, whisper_model = load_whisper()
-        vision_model, vision_processor = load_vision_model()
         sarvam_pipe = load_sarvam()
+        vision_model, vision_processor = load_vision_model()
 
         if input_type == "audio" and audio_input is not None:
-            transcription = process_audio(audio_input, whisper_processor, whisper_model)
+            transcription = process_audio_input(audio_input, whisper_processor, whisper_model)
         elif input_type == "text" and text_input:
             transcription = text_input
         elif input_type == "image" and image_input is not None:
-            text_prompt = text_input if text_input else "Describe this image in detail."
-            transcription = process_image(image_input, text_prompt, vision_model, vision_processor)
+            return process_image(image_input, text_input, vision_model, vision_processor), None
         else:
-            return "Please provide either audio, text, or image input.", "No input provided.", None
+            return "Please provide either audio, text, or image input.", None
 
         response = generate_response(transcription, sarvam_pipe)
-        lang = detect(response)
+        lang = detect_language(response)
         audio_response = text_to_speech(response, lang)
         
-        return transcription, response, audio_response
+        return response, audio_response
     except Exception as e:
         error_message = f"An error occurred: {str(e)}"
-        return error_message, error_message, None
+        return error_message, None
 
-# Custom CSS
+# Custom CSS (you can keep your existing custom CSS here)
 custom_css = """
 body {
     background-color: #0b0f19;
@@ -183,12 +201,12 @@ footer {
 }
 """
 
-# Custom HTML for the header
+# Custom HTML for the header (you can keep your existing custom header here)
 custom_header = """
 <div id="custom-header">
     <h1>
-        <span class="blue">Hello,</span>
-        <span class="pink">User</span>
+        <span class="blue">Multimodal</span>
+        <span class="pink">Indic Assistant</span>
     </h1>
     <h2>How can I help you today?</h2>
 </div>
@@ -206,7 +224,7 @@ custom_suggestions = """
         <p>Type in any Indic language</p>
     </div>
     <div class="suggestion">
-        <span class="suggestion-icon">🖼️</span>
+        <span class="suggestion-icon">📷</span>
         <p>Upload an image for analysis</p>
     </div>
     <div class="suggestion">
@@ -220,7 +238,7 @@ custom_suggestions = """
 </div>
 """
 
-# Gradio interface
+# Create Gradio interface
 with gr.Blocks(css=custom_css, theme=gr.themes.Base().set(
     body_background_fill="#0b0f19",
     body_text_color="#e2e8f0",
@@ -235,25 +253,24 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Base().set(
     
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("### Indic Vision Assistant")
+            gr.Markdown("### Multimodal Indic Assistant")
     
     input_type = gr.Radio(["audio", "text", "image"], label="Input Type", value="audio")
     audio_input = gr.Audio(type="filepath", label="Speak (if audio input selected)")
-    text_input = gr.Textbox(label="Type your message or image prompt")
-    image_input = gr.Image(type="pil", label="Upload an image (if image input selected)")
+    text_input = gr.Textbox(label="Type your message or image question")
+    image_input = gr.Image(label="Upload an image (if image input selected)")
     
     submit_btn = gr.Button("Submit")
     
-    output_transcription = gr.Textbox(label="Transcription/Input")
     output_response = gr.Textbox(label="Generated Response")
     output_audio = gr.Audio(label="Audio Response")
     
     submit_btn.click(
-        fn=indic_vision_assistant,
+        fn=multimodal_assistant,
         inputs=[input_type, audio_input, text_input, image_input],
-        outputs=[output_transcription, output_response, output_audio]
+        outputs=[output_response, output_audio]
     )
-    gr.HTML("<footer>Powered by Indic Language AI with Vision Capabilities</footer>")
+    gr.HTML("<footer>Powered by Multimodal Indic Language AI</footer>")
 
 # Launch the app
 iface.launch()