Update app.py

app.py

@@ -1,11 +1,13 @@
 import gradio as gr
 from PIL import Image
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, WhisperProcessor, WhisperForConditionalGeneration, ViTFeatureExtractor, ViTForImageClassification 
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, Wav2Vec2Processor, Wav2Vec2ForCTC, ViTFeatureExtractor, ViTForImageClassification
 
 import soundfile as sf
 import torch
 import numpy as np
+import time
 
+# Initialize the transformers and the models
 class_names = {
     0: "al qarawiyyin",
     1: "bab mansour el aleuj",
@@ -16,94 +18,79 @@ class_names = {
     6: "madrasa ben youssef",
     7: "majorel gardens",
     8: "menara"
-  }
+}
 
 model_name_or_path = "microsoft/DialoGPT-large"
-
 tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, padding_side="left", use_fast=False)
 tokenizer.pad_token = tokenizer.eos_token
+model = AutoModelForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float32, trust_remote_code=True)
 
-model = AutoModelForCausalLM.from_pretrained(
-    model_name_or_path,
-    torch_dtype=torch.float32,
-    device_map="auto",
-    trust_remote_code=True,
-)
-
-# Initialize the Wav2Vec2 model and processor
-wav2vec2_processor = WhisperProcessor.from_pretrained("openai/whisper-large")
-wav2vec2_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-large")
-wav2vec2_model.config.forced_decoder_ids = None
+wav2vec2_processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
+wav2vec2_model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
 
 vit_model = ViTForImageClassification.from_pretrained('ohidaoui/monuments-morocco-v1')
 vit_feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
 
-
 # Function to handle text input
 def handle_text(text):
-    new_user_input_ids = tokenizer.encode(text + tokenizer.eos_token, return_tensors='pt')
-    bot_input_ids = new_user_input_ids
-    chat_history_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
-    chat_output = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
-    return chat_output
+    chat_output = chat({"question": text})
+    return chat_output["answer"]
 
 # Function to handle image input
 def get_class_name(class_idx):
     return class_names[class_idx]
 
-
 def handle_image(img):
-    # Convert PIL image to numpy array
     img = np.array(img)
-
-    # Apply transformations and prepare image for the model
     inputs = vit_feature_extractor(images=img, return_tensors="pt")
-
-    # Pass through the Vision Transformer model
     outputs = vit_model(**inputs)
-
-    # Get the predicted class
     predicted_class_idx = torch.argmax(outputs.logits, dim=1).item()
-
-   
     predicted_class_name = get_class_name(predicted_class_idx)
+    chat_output = chat({"question": "what is " + predicted_class_name})
+    return chat_output["answer"]
 
-    return predicted_class_name
-    
 # Function to handle audio input
 def handle_audio(audio):
-    # gradio Audio returns a tuple (sample_rate, audio_np_array)
-    # we only need the audio data, hence accessing the second element
     audio = audio[1]
-    input_values = wav2vec2_processor(audio, sampling_rate=16000, return_tensors="pt").input_values
-    # Convert to the expected tensor type
+    input_values = wav2vec2_processor(audio, sampling_rate=16_000, return_tensors="pt").input_values
     input_values = input_values.to(torch.float32)
     logits = wav2vec2_model(input_values).logits
     predicted_ids = torch.argmax(logits, dim=-1)
     transcriptions = wav2vec2_processor.decode(predicted_ids[0])
-    return transcriptions
-
-
-
+    chat_output = chat({"question": transcriptions})
+    return chat_output["answer"]
 
-def chatbot(text, img, audio):
+# Main function to handle the inputs
+def chatbot(history, text=None, img=None, audio=None):
     text_output = handle_text(text) if text is not None else ''
     img_output = handle_image(img) if img is not None else ''
     audio_output = handle_audio(audio) if audio is not None else ''
-
     outputs = [o for o in [text_output, img_output, audio_output] if o]
-    return "\n".join(outputs)
-
-iface = gr.Interface(
-    fn=chatbot,
-    inputs=[
-        gr.inputs.Textbox(lines=2, placeholder="Input Text here..."),
-        gr.inputs.Image(label="Upload Image"),
-        gr.inputs.Audio(source="microphone", label="Audio Input"),
-    ],
-    outputs=gr.outputs.Textbox(label="Output"),
-    title="Multimodal Chatbot",
-    description="This chatbot can handle text, image, and audio inputs. Try it out!",
-)
-
-iface.launch()
+    output = "\n".join(outputs)
+    
+    history[-1][1] = output
+    for character in output:
+        history[-1][1] += character
+        time.sleep(0.05)
+        yield history
+
+with gr.Blocks() as demo:
+    chat_interface = gr.Chatbot([], elem_id="chatbot", height=750)
+
+    with gr.Row():
+        with gr.Column(scale=0.85):
+            text_input = gr.Textbox(
+                show_label=False,
+                placeholder="Input Text here...",
+                container=False
+            )
+        with gr.Column(scale=0.15, min_width=0):
+            img_input = gr.Image()
+            audio_input = gr.Audio(source="microphone", label="Audio Input")
+
+    text_msg = text_input.submit(chatbot, [chat_interface, text_input], [chat_interface, text_input], queue=False)
+    img_msg = img_input.upload(chatbot, [chat_interface, img_input], [chat_interface, img_input], queue=False)
+    audio_msg = audio_input.upload(chatbot, [chat_interface, audio_input], [chat_interface, audio_input], queue=False)
+
+demo.queue()
+demo.launch(share=True)