app.py

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

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize ZeroGPU
spaces.init()

# Function to safely load pipeline
@spaces.GPU
def load_pipeline(model_name, **kwargs):
    try:
        return pipeline(model=model_name, device=0, **kwargs)
    except Exception as e:
        logger.error(f"Error loading {model_name} pipeline: {e}")
        return None

# Load Whisper model for speech recognition
@spaces.GPU
def load_whisper():
    try:
        processor = WhisperProcessor.from_pretrained("openai/whisper-small")
        model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").cuda()
        return processor, model
    except Exception as e:
        logger.error(f"Error loading Whisper model: {e}")
        return None, None

# Load sarvam-2b for text generation
@spaces.GPU
def load_sarvam():
    return load_pipeline('sarvamai/sarvam-2b-v0.5')

# Load vision model
@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="auto", attn_implementation="flash_attention_2").cuda().eval()
        processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
        return model, processor
    except Exception as e:
        logger.error(f"Error loading vision model: {e}")
        return None, None

# Process audio input
@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.cuda()
        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:
        logger.error(f"Error processing audio: {e}")
        return f"Error processing audio. Please type your message instead."

# Generate response
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)
        
        output_path = "/tmp/response.mp3"
        tts.save(output_path)
        return output_path
    except Exception as e:
        logger.error(f"Error in text-to-speech: {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:
        logger.error(f"Error generating response: {e}")
        return f"Error generating response. Please try again."

@spaces.GPU
def process_image(image, text_input, vision_model, vision_processor):
    if vision_model is None or vision_processor is None:
        return "Error: Vision model is not available."
    
    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")
        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:
        logger.error(f"Error processing image: {e}")
        return f"Error processing image. Please try again."

@spaces.GPU
def multimodal_assistant(input_type, audio_input, text_input, image_input):
    try:
        # Load models
        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:
        logger.error(f"An error occurred in multimodal_assistant: {e}")
        return f"An error occurred. Please try again.", 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()