Update app.py

app.py

@@ -1,124 +1,348 @@
-import os
+# Import spaces first to avoid CUDA initialization issues
+import spaces
+
+# Then import other libraries
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, TextIteratorStreamer, BitsAndBytesConfig
+import librosa
+from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, AutoModelForCausalLM, AutoProcessor
+from gtts import gTTS
 import gradio as gr
-from threading import Thread
 from PIL import Image
+import os
+import base64
+from io import BytesIO
+
+import io
 import subprocess
-import spaces  # Add this import
+from langdetect import detect
 
-# Install flash-attention
+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)
 
-# Constants
-TITLE = "<h1><center>Phi 3.5 Multimodal (Text + Vision)</center></h1>"
-DESCRIPTION = "# Phi-3.5 Multimodal Demo (Text + Vision)"
-
-# Model configurations
-TEXT_MODEL_ID = "microsoft/Phi-3.5-mini-instruct"
-VISION_MODEL_ID = "microsoft/Phi-3.5-vision-instruct"
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Quantization config for text model
-quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4"
-)
-
-# Load models and tokenizers
-text_tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_ID)
-text_model = AutoModelForCausalLM.from_pretrained(
-    TEXT_MODEL_ID,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-    quantization_config=quantization_config
-)
-
-vision_model = AutoModelForCausalLM.from_pretrained(
-    VISION_MODEL_ID, 
-    trust_remote_code=True, 
-    torch_dtype="auto", 
-    attn_implementation="flash_attention_2"
-).to(device).eval()
-
-vision_processor = AutoProcessor.from_pretrained(VISION_MODEL_ID, trust_remote_code=True)
-
-# Helper functions
+# Function to safely load pipeline within a GPU-decorated function
 @spaces.GPU
-def stream_text_chat(message, history, system_prompt, temperature=0.8, max_new_tokens=1024, top_p=1.0, top_k=20):
-    conversation = [{"role": "system", "content": system_prompt}]
-    for prompt, answer in history:
-        conversation.extend([
-            {"role": "user", "content": prompt},
-            {"role": "assistant", "content": answer},
-        ])
-    conversation.append({"role": "user", "content": message})
-
-    input_ids = text_tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(text_model.device)
-    streamer = TextIteratorStreamer(text_tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
-
-    generate_kwargs = dict(
-        input_ids=input_ids,
-        max_new_tokens=max_new_tokens,
-        do_sample=temperature > 0,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        eos_token_id=[128001, 128008, 128009],
-        streamer=streamer,
-    )
+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
 
-    with torch.no_grad():
-        thread = Thread(target=text_model.generate, kwargs=generate_kwargs)
-        thread.start()
+# Load sarvam-2b for text generation within a GPU-decorated function
+@spaces.GPU
+def load_sarvam():
+    return load_pipeline('sarvamai/sarvam-2b-v0.5')
 
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        yield history + [[message, buffer]]
+# 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
 
-@spaces.GPU  # Add this decorator
-def process_vision_query(image, text_input):
-    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(device)
-    
-    with torch.no_grad():
-        generate_ids = vision_model.generate(
-            **inputs, 
-            max_new_tokens=1000, 
-            eos_token_id=vision_processor.tokenizer.eos_token_id
+
+# 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:
+        print("Starting to load vision model...")
+        model_id = "microsoft/Phi-3.5-vision-instruct"
+        print(f"Loading model from {model_id}")
+        
+        # Check for CUDA availability
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {device}")
+        
+        # Load model with potential memory optimization
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id, 
+            trust_remote_code=True, 
+            torch_dtype=torch.float16, 
+            use_flash_attention_2=True,  # Enable if supported
+            device_map="auto",  # Automatically manage model placement
+            low_cpu_mem_usage=True
         )
+        print("Model loaded successfully")
+        
+        print("Loading processor...")
+        processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True, num_crops=16)
+        print("Processor loaded successfully")
+        
+        return model, processor
+    except ImportError as e:
+        print(f"Error importing required modules: {str(e)}")
+        print("Please ensure all required dependencies are installed.")
+    except RuntimeError as e:
+        print(f"Runtime error (possibly CUDA out of memory): {str(e)}")
+        print("Consider using a smaller model or enabling GPU offloading.")
+    except Exception as e:
+        print(f"Unexpected error in loading vision model: {str(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."
     
-    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
+    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."
+
+# Updated process_image_input function
+@spaces.GPU
+@spaces.GPU
+def process_image_input(image, text_prompt, vision_model, processor):
+    if vision_model is None or processor is None:
+        return "Error: Vision model is not available."
+    
+    try:
+        # Convert image to base64
+        if isinstance(image, Image.Image):
+            buffered = BytesIO()
+            image.save(buffered, format="PNG")
+            img_str = base64.b64encode(buffered.getvalue()).decode()
+        else:
+            # If it's not a PIL Image, assume it's a file path
+            with open(image, "rb") as image_file:
+                img_str = base64.b64encode(image_file.read()).decode()
+        
+        # Format the input with image tag
+        formatted_prompt = f"{text_prompt}\n<image>data:image/png;base64,{img_str}</image>"
+        
+        # Process the formatted prompt
+        inputs = processor(text=formatted_prompt, return_tensors="pt").to(vision_model.device)
+        
+        # Generate text
+        with torch.no_grad():
+            outputs = vision_model.generate(
+                **inputs,
+                max_new_tokens=100,
+                do_sample=True,
+                top_k=50,
+                top_p=0.95,
+                num_return_sequences=1
+            )
+        
+        generated_text = 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:
+        whisper_processor, whisper_model = load_whisper()
+        sarvam_pipe = load_sarvam()
+        vision_model, 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:
+            # Use a default prompt if no text input is provided
+            text_prompt = text_input if text_input else "Describe this image in detail."
+            transcription = process_image_input(image_input, text_prompt, vision_model, 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;}
+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 {@spaces.GPU
+def indic_vision_assistant(input_type, audio_input, text_input, image_input):
+    try:
+        whisper_processor, whisper_model = load_whisper()
+        sarvam_pipe = load_sarvam()
+        vision_model, 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:
+            # Use a default prompt if no text input is provided
+            text_prompt = text_input if text_input else "Describe this image in detail."
+            transcription = process_image_input(image_input, text_prompt, vision_model, 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
+
+    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">Phi 3.5</span> <span class="pink">Multimodal Assistant</span></h1>
-    <h2>Text and Vision AI at Your Service</h2>
+    <h1>
+        <span class="blue">Hello,</span>
+        <span class="pink">User</span>
+    </h1>
+    <h2>How can I help you today?</h2>
 </div>
 """
 
@@ -126,25 +350,28 @@ custom_header = """
 custom_suggestions = """
 <div class="suggestions">
     <div class="suggestion">
-        <span class="suggestion-icon">💬</span>
-        <p>Chat with the Text Model</p>
+        <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>Analyze Images with Vision Model</p>
+        <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>Explore advanced options</p>
+        <span class="suggestion-icon">🔊</span>
+        <p>Listen to audio responses</p>
     </div>
 </div>
 """
-
-# Gradio interface
+# Update the Gradio interface to allow text input for image processing
 with gr.Blocks(css=custom_css, theme=gr.themes.Base().set(
     body_background_fill="#0b0f19",
     body_text_color="#e2e8f0",
@@ -153,38 +380,30 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Base().set(
     button_primary_text_color="white",
     block_title_text_color="#94a3b8",
     block_label_text_color="#94a3b8",
-)) as demo:
+)) as iface:
     gr.HTML(custom_header)
     gr.HTML(custom_suggestions)
-
-    with gr.Tab("Text Model (Phi-3.5-mini)"):
-        chatbot = gr.Chatbot(height=400)
-        msg = gr.Textbox(label="Message", placeholder="Type your message here...")
-        with gr.Accordion("Advanced Options", open=False):
-            system_prompt = gr.Textbox(value="You are a helpful assistant", label="System Prompt")
-            temperature = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.8, label="Temperature")
-            max_new_tokens = gr.Slider(minimum=128, maximum=8192, step=1, value=1024, label="Max new tokens")
-            top_p = gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=1.0, label="top_p")
-            top_k = gr.Slider(minimum=1, maximum=20, step=1, value=20, label="top_k")
-        
-        submit_btn = gr.Button("Submit", variant="primary")
-        clear_btn = gr.Button("Clear Chat", variant="secondary")
-
-        submit_btn.click(stream_text_chat, [msg, chatbot, system_prompt, temperature, max_new_tokens, top_p, top_k], [chatbot])
-        clear_btn.click(lambda: None, None, chatbot, queue=False)
-
-    with gr.Tab("Vision Model (Phi-3.5-vision)"):
-        with gr.Row():
-            with gr.Column(scale=1):
-                vision_input_img = gr.Image(label="Upload an Image", type="pil")
-                vision_text_input = gr.Textbox(label="Ask a question about the image", placeholder="What do you see in this image?")
-                vision_submit_btn = gr.Button("Analyze Image", variant="primary")
-            with gr.Column(scale=1):
-                vision_output_text = gr.Textbox(label="AI Analysis", lines=10)
-        
-        vision_submit_btn.click(process_vision_query, [vision_input_img, vision_text_input], [vision_output_text])
-
-    gr.HTML("<footer>Powered by Phi 3.5 Multimodal AI</footer>")
-
-if __name__ == "__main__":
-    demo.launch()
+    
+    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 or image prompt")
+    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()