Create app.py

app.py

@@ -0,0 +1,309 @@
+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 os
+import io
+import subprocess
+from langdetect import detect
+
+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 Phi-3.5-vision-instruct 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=torch.float16, use_flash_attention_2=False
+        )
+        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
+
+# 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."
+
+# Process image input
+@spaces.GPU
+def process_image_input(image, vision_model, vision_processor):
+    if vision_model is None or vision_processor is None:
+        return "Error: Vision model is not available."
+    
+    try:
+        inputs = vision_processor(images=image, return_tensors="pt")
+        inputs = {k: v.to(vision_model.device) for k, v in inputs.items()}
+        
+        with torch.no_grad():
+            outputs = vision_model.generate(**inputs, max_new_tokens=512, do_sample=True, top_k=50, top_p=0.95)
+        
+        generated_text = vision_processor.batch_decode(outputs, skip_special_tokens=True)[0]
+        return generated_text
+    except Exception as e:
+        return f"Error processing image: {str(e)}"
+
+# Generate response within a GPU-decorated function
+@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)}"
+
+# Text-to-speech function
+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']:
+            # You might want to use a different TTS library here
+            # For example, you could use the Google Cloud Text-to-Speech API
+            # or a specialized Indic language TTS library
+            
+            # This is a placeholder for a better Indic TTS solution
+            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
+
+# Improved language detection function
+def detect_language(text):
+    lang_codes = {
+        'bn': 'Bengali', 'gu': 'Gujarati', 'hi': 'Hindi', 'kn': 'Kannada',
+        'ml': 'Malayalam', 'mr': 'Marathi', 'or': 'Oriya', 'pa': 'Punjabi',
+        'ta': 'Tamil', 'te': 'Telugu', 'en': 'English'
+    }
+    
+    try:
+        detected_lang = detect(text)
+        return detected_lang if detected_lang in lang_codes else 'en'
+    except:
+        # Fallback to simple script-based detection
+        for code, lang in lang_codes.items():
+            if any(ord(char) >= 0x0900 and ord(char) <= 0x097F for char in text):  # Devanagari script
+                return 'hi'
+        return 'en'  # Default to English if no Indic script is detected
+
+@spaces.GPU
+def indic_vision_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:
+            transcription = process_image_input(image_input, vision_model, vision_processor)
+        else:
+            return "Please provide either audio, text, or image input.", "No input provided.", None
+
+        response = generate_response(transcription, sarvam_pipe)
+        lang = detect_language(response)
+        audio_response = text_to_speech(response, lang)
+        
+        return transcription, response, audio_response
+    except Exception as e:
+        error_message = f"An error occurred: {str(e)}"
+        return error_message, error_message, None
+
+# Custom CSS
+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
+custom_header = """
+<div id="custom-header">
+    <h1>
+        <span class="blue">Hello,</span>
+        <span class="pink">User</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("### Indic Vision 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 (if text input selected)")
+    image_input = gr.Image(type="pil", 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,
+        inputs=[input_type, audio_input, text_input, image_input],
+        outputs=[output_transcription, output_response, output_audio]
+    )
+    gr.HTML("<footer>Powered by Indic Language AI with Vision Capabilities</footer>")
+
+# Launch the app
+iface.launch()