Update app.py

app.py

@@ -1,279 +1,128 @@
+import os
 import torch
-import librosa
-from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration, AutoModelForCausalLM, AutoProcessor
-from gtts import gTTS
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor, TextIteratorStreamer, BitsAndBytesConfig
 import gradio as gr
+from threading import Thread
 from PIL import Image
-import logging
-import os
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+import subprocess
 
-# Check for GPU availability
-device = "cuda" if torch.cuda.is_available() else "cpu"
-logger.info(f"Using device: {device}")
-
-# Function to safely load pipeline
-def load_pipeline(model_name, **kwargs):
-    try:
-        return pipeline(model=model_name, device=device, **kwargs)
-    except Exception as e:
-        logger.error(f"Error loading {model_name} pipeline: {e}")
-        return None
+# Install flash-attention
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-# Load Whisper model for speech recognition
-def load_whisper():
-    try:
-        processor = WhisperProcessor.from_pretrained("openai/whisper-small")
-        model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").to(device)
-        return processor, model
-    except Exception as e:
-        logger.error(f"Error loading Whisper model: {e}")
-        return None, None
+# Constants
+TITLE = "<h1><center>Phi 3.5 Multimodal (Text + Vision)</center></h1>"
+DESCRIPTION = "# Phi-3.5 Multimodal Demo (Text + Vision)"
 
-# Load sarvam-2b for text generation
-def load_sarvam():
-    return load_pipeline('sarvamai/sarvam-2b-v0.5')
+# Model configurations
+TEXT_MODEL_ID = "microsoft/Phi-3.5-mini-instruct"
+VISION_MODEL_ID = "microsoft/Phi-3.5-vision-instruct"
 
-# Load vision model
-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").to(device).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
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Process audio input
-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(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:
-        logger.error(f"Error processing audio: {e}")
-        return f"Error processing audio. Please type your message instead."
+# 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
+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,
+    )
 
-# 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
+    with torch.no_grad():
+        thread = Thread(target=text_model.generate, kwargs=generate_kwargs)
+        thread.start()
 
-# Detect language (placeholder function, replace with actual implementation)
-def detect_language(text):
-    # Implement language detection logic here
-    return 'en'  # Default to English for now
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        yield buffer
 
-def generate_response(transcription, sarvam_pipe):
-    if sarvam_pipe is None:
-        return "Error: Text generation model is not available."
+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)
     
-    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."
-
-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."
+    with torch.no_grad():
+        generate_ids = vision_model.generate(
+            **inputs, 
+            max_new_tokens=1000, 
+            eos_token_id=vision_processor.tokenizer.eos_token_id
+        )
     
-    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(device)
-        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."
-
-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)
+    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
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.HTML(TITLE)
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Tab("Text Model (Phi-3.5-mini)"):
+        chatbot = gr.Chatbot(height=600)
+        gr.ChatInterface(
+            fn=stream_text_chat,
+            chatbot=chatbot,
+            additional_inputs=[
+                gr.Textbox(value="You are a helpful assistant", label="System Prompt"),
+                gr.Slider(minimum=0, maximum=1, step=0.1, value=0.8, label="Temperature"),
+                gr.Slider(minimum=128, maximum=8192, step=1, value=1024, label="Max new tokens"),
+                gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=1.0, label="top_p"),
+                gr.Slider(minimum=1, maximum=20, step=1, value=20, label="top_k"),
+            ],
+        )
+
+    with gr.Tab("Vision Model (Phi-3.5-vision)"):
+        with gr.Row():
+            with gr.Column():
+                vision_input_img = gr.Image(label="Input Picture")
+                vision_text_input = gr.Textbox(label="Question")
+                vision_submit_btn = gr.Button(value="Submit")
+            with gr.Column():
+                vision_output_text = gr.Textbox(label="Output Text")
         
-        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>")
+        vision_submit_btn.click(process_vision_query, [vision_input_img, vision_text_input], [vision_output_text])
 
-# Launch the app
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()