Update app.py

app.py

@@ -17,7 +17,7 @@ print(f"[SYSTEM] | Using {DEVICE} type compute device.")
 DEFAULT_INPUT = "Describe in one paragraph."
 MAX_FRAMES = 64
 
-repo_name = "openbmb/MiniCPM-o-2_6-int4" # "openbmb/MiniCPM-V-2_6-int4" # "openbmb/MiniCPM-V-2_6"
+repo_name = "openbmb/MiniCPM-o-2_6" # "openbmb/MiniCPM-V-2_6-int4" # "openbmb/MiniCPM-V-2_6"
 repo = AutoModel.from_pretrained(repo_name, trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
 
@@ -48,19 +48,28 @@ def encode_video(video_path):
     return frames
     
 @spaces.GPU(duration=60)
-def generate(image, video, instruction=DEFAULT_INPUT, sampling=False, temperature=0.7, top_p=0.8, top_k=100, repetition_penalty=1.05, max_tokens=512):
+def generate(image, video, audio, instruction=DEFAULT_INPUT, sampling=False, temperature=0.7, top_p=0.8, top_k=100, repetition_penalty=1.05, max_tokens=512):
     # repo.to(DEVICE)
 
     print(image)
     print(video)
+    print(audio)
     print(instruction)
     
-    if not video:
+    if image is not None:
         image_data = Image.fromarray(image.astype('uint8'), 'RGB')
         inputs = [{"role": "user", "content": [image_data, instruction]}]
-    else:
+    elif video is not None:
         video_data = encode_video(video)
-        inputs = [{"role": "user", "content": video_data + [instruction]}]
+        inputs = [{"role": "user", "content": [video_data, instruction]}]
+    elif audio is not None:
+        if isinstance(audio, str):
+            audio_data, _ = librosa.load(audio, sr=16000, mono=True)
+        else:
+            audio_data = audio
+        inputs = [{"role": "user", "content": [audio_data, instruction]}]
+    else:
+        return "No input provided."
     
     parameters = {
         "sampling": sampling,
@@ -69,8 +78,6 @@ def generate(image, video, instruction=DEFAULT_INPUT, sampling=False, temperatur
         "top_k": top_k,
         "repetition_penalty": repetition_penalty,
         "max_new_tokens": max_tokens,
-        "use_image_id": False,
-        "max_slice_nums": 2,
     }
     
     output = repo.chat(image=None, msgs=inputs, tokenizer=tokenizer, **parameters)
@@ -90,6 +97,7 @@ with gr.Blocks(css=css) as main:
     with gr.Column():
         input = gr.Image(label="Image")
         input_2 = gr.Video(label="Video")
+        input_3 = gr.Audio(label="Audio"),
         instruction = gr.Textbox(lines=1, value=DEFAULT_INPUT, label="Instruction")
         sampling = gr.Checkbox(value=False, label="Sampling")
         temperature = gr.Slider(minimum=0.01, maximum=1.99, step=0.01, value=0.7, label="Temperature")
@@ -103,7 +111,7 @@ with gr.Blocks(css=css) as main:
     with gr.Column():
         output = gr.Textbox(lines=1, value="", label="Output")
 
-    submit.click(fn=generate, inputs=[input, input_2, instruction, sampling, temperature, top_p, top_k, repetition_penalty, max_tokens], outputs=[output], queue=False)
+    submit.click(fn=generate, inputs=[input, input_2, input_3, instruction, sampling, temperature, top_p, top_k, repetition_penalty, max_tokens], outputs=[output], queue=False)
     maintain.click(cloud, inputs=[], outputs=[], queue=False)
 
 main.launch(show_api=True)