app.py

import torch
import librosa
from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, AutoModelForCausalLM, AutoProcessor
from gtts import gTTS
import gradio as gr
import spaces
from PIL import Image
import subprocess

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)

# 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").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_sarvam():
    return load_pipeline('sarvamai/sarvam-2b-v0.5')

# Load vision model
@spaces.GPU
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_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, 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)}. Please type your message instead."

# Generate response within a GPU-decorated function
@spaces.GPU
def text_to_speech(text, lang='hi'):
    try:
        # 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:
        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)}"

@spaces.GPU
def process_image(image, text_input, vision_model, vision_processor):
    try:
        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:
        return f"Error processing image: {str(e)}"

@spaces.GPU
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()
        sarvam_pipe = load_sarvam()
        vision_model, vision_processor = load_vision_model()

        if input_type == "audio" and audio_input is not None:
            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:
            return process_image(image_input, text_input, vision_model, vision_processor), None
        else:
            return "Please provide either audio, text, or image input.", None

        response = generate_response(transcription, sarvam_pipe)
        lang = detect_language(response)
        audio_response = text_to_speech(response, lang)
        
        return response, audio_response
    except Exception as e:
        error_message = f"An error occurred: {str(e)}"
        return error_message, None

# Custom CSS (you can keep your existing custom CSS here)
custom_css = """
body {
    background-color: #0b0f19;
    color: #e2e8f0;
    font-family: 'Arial', sans-serif;
}
#custom-header {
    text-align: center;
    padding: 20px 0;
    background-color: #1a202c;
    margin-bottom: 20px;
    border-radius: 10px;
}
#custom-header h1 {
    font-size: 2.5rem;
    margin-bottom: 0.5rem;
}
#custom-header h1 .blue {
    color: #60a5fa;
}
#custom-header h1 .pink {
    color: #f472b6;
}
#custom-header h2 {
    font-size: 1.5rem;
    color: #94a3b8;
}
.suggestions {
    display: flex;
    justify-content: center;
    flex-wrap: wrap;
    gap: 1rem;
    margin: 20px 0;
}
.suggestion {
    background-color: #1e293b;
    border-radius: 0.5rem;
    padding: 1rem;
    display: flex;
    align-items: center;
    transition: transform 0.3s ease;
    width: 200px;
}
.suggestion:hover {
    transform: translateY(-5px);
}
.suggestion-icon {
    font-size: 1.5rem;
    margin-right: 1rem;
    background-color: #2d3748;
    padding: 0.5rem;
    border-radius: 50%;
}
.gradio-container {
    max-width: 100% !important;
}
#component-0, #component-1, #component-2 {
    max-width: 100% !important;
}
footer {
    text-align: center;
    margin-top: 2rem;
    color: #64748b;
}
"""

# Custom HTML for the header (you can keep your existing custom header here)
custom_header = """
<div id="custom-header">
    <h1>
        <span class="blue">Multimodal</span>
        <span class="pink">Indic Assistant</span>
    </h1>
    <h2>How can I help you today?</h2>
</div>
"""

# Custom HTML for suggestions
custom_suggestions = """
<div class="suggestions">
    <div class="suggestion">
        <span class="suggestion-icon">🎤</span>
        <p>Speak in any Indic language</p>
    </div>
    <div class="suggestion">
        <span class="suggestion-icon">⌨️</span>
        <p>Type in any Indic language</p>
    </div>
    <div class="suggestion">
        <span class="suggestion-icon">📷</span>
        <p>Upload an image for analysis</p>
    </div>
    <div class="suggestion">
        <span class="suggestion-icon">🤖</span>
        <p>Get AI-generated responses</p>
    </div>
    <div class="suggestion">
        <span class="suggestion-icon">🔊</span>
        <p>Listen to audio responses</p>
    </div>
</div>
"""

# Create Gradio interface
with gr.Blocks(css=custom_css, theme=gr.themes.Base().set(
    body_background_fill="#0b0f19",
    body_text_color="#e2e8f0",
    button_primary_background_fill="#3b82f6",
    button_primary_background_fill_hover="#2563eb",
    button_primary_text_color="white",
    block_title_text_color="#94a3b8",
    block_label_text_color="#94a3b8",
)) as iface:
    gr.HTML(custom_header)
    gr.HTML(custom_suggestions)
    
    with gr.Row():
        with gr.Column(scale=1):
            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 question")
    image_input = gr.Image(label="Upload an image (if image input selected)")
    
    submit_btn = gr.Button("Submit")
    
    output_response = gr.Textbox(label="Generated Response")
    output_audio = gr.Audio(label="Audio Response")
    
    submit_btn.click(
        fn=multimodal_assistant,
        inputs=[input_type, audio_input, text_input, image_input],
        outputs=[output_response, output_audio]
    )
    gr.HTML("<footer>Powered by Multimodal Indic Language AI</footer>")

# Launch the app
iface.launch()