Spaces:
Running
Running
| # -*- encoding: utf-8 -*- | |
| # File: app.py | |
| # Description: None | |
| from copy import deepcopy | |
| from typing import Dict, List | |
| from PIL import Image | |
| import io | |
| import subprocess | |
| import requests | |
| import json | |
| import base64 | |
| import gradio as gr | |
| import librosa | |
| IMAGE_EXTENSIONS = (".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp") | |
| VIDEO_EXTENSIONS = (".mp4", ".mkv", ".mov", ".avi", ".flv", ".wmv", ".webm", ".m4v") | |
| AUDIO_EXTENSIONS = (".mp3", ".wav", "flac", ".m4a") | |
| DEFAULT_SAMPLING_PARAMS = { | |
| "top_p": 0.8, | |
| "top_k": 100, | |
| "temperature": 0.7, | |
| "do_sample": True, | |
| "num_beams": 1, | |
| "repetition_penalty": 1.2, | |
| } | |
| MAX_NEW_TOKENS = 1024 | |
| def load_image_to_base64(image_path): | |
| """Load image and convert to base64 string""" | |
| with Image.open(image_path) as img: | |
| if img.mode != 'RGB': | |
| img = img.convert('RGB') | |
| img_byte_arr = io.BytesIO() | |
| img.save(img_byte_arr, format='PNG') | |
| img_byte_arr = img_byte_arr.getvalue() | |
| return base64.b64encode(img_byte_arr).decode('utf-8') | |
| def wav_to_bytes_with_ffmpeg(wav_file_path): | |
| process = subprocess.Popen( | |
| ['ffmpeg', '-i', wav_file_path, '-f', 'wav', '-'], | |
| stdout=subprocess.PIPE, | |
| stderr=subprocess.PIPE | |
| ) | |
| out, _ = process.communicate() | |
| return base64.b64encode(out).decode('utf-8') | |
| def parse_sse_response(response): | |
| for line in response.iter_lines(): | |
| if line: | |
| line = line.decode('utf-8') | |
| if line.startswith('data: '): | |
| data = line[6:] # Remove 'data: ' prefix | |
| if data == '[DONE]': | |
| break | |
| try: | |
| json_data = json.loads(data) | |
| yield json_data['text'] | |
| except json.JSONDecodeError: | |
| raise gr.Error(f"Failed to parse JSON: {data}") | |
| def history2messages(history: List[Dict]) -> List[Dict]: | |
| """ | |
| Transform gradio history to chat messages. | |
| """ | |
| messages = [] | |
| cur_message = dict() | |
| for item in history: | |
| if item["role"] == "assistant": | |
| if len(cur_message) > 0: | |
| messages.append(deepcopy(cur_message)) | |
| cur_message = dict() | |
| messages.append(deepcopy(item)) | |
| continue | |
| if "role" not in cur_message: | |
| cur_message["role"] = "user" | |
| if "content" not in cur_message: | |
| cur_message["content"] = dict() | |
| if "metadata" not in item: | |
| item["metadata"] = {"title": None} | |
| if item["metadata"]["title"] is None: | |
| cur_message["content"]["text"] = item["content"] | |
| elif item["metadata"]["title"] == "image": | |
| cur_message["content"]["image"] = load_image_to_base64(item["content"][0]) | |
| elif item["metadata"]["title"] == "audio": | |
| cur_message["content"]["audio"] = wav_to_bytes_with_ffmpeg(item["content"][0]) | |
| if len(cur_message) > 0: | |
| messages.append(cur_message) | |
| return messages | |
| def check_messages(history, message, audio): | |
| has_text = message["text"] and message["text"].strip() | |
| has_files = len(message["files"]) > 0 | |
| has_audio = audio is not None | |
| if not (has_text or has_files or has_audio): | |
| raise gr.Error("请输入文字或上传音频/图片后再发送。") | |
| audios = [] | |
| images = [] | |
| for file_msg in message["files"]: | |
| if file_msg.endswith(AUDIO_EXTENSIONS) or file_msg.endswith(VIDEO_EXTENSIONS): | |
| duration = librosa.get_duration(filename=file_msg) | |
| if duration > 30: | |
| raise gr.Error("音频时长不能超过30秒。") | |
| if duration == 0: | |
| raise gr.Error("音频时长不能为0秒。") | |
| audios.append(file_msg) | |
| elif file_msg.endswith(IMAGE_EXTENSIONS): | |
| images.append(file_msg) | |
| else: | |
| filename = file_msg.split("/")[-1] | |
| raise gr.Error(f"Unsupported file type: {filename}. It should be an image or audio file.") | |
| if len(audios) > 1: | |
| raise gr.Error("Please upload only one audio file.") | |
| if len(images) > 1: | |
| raise gr.Error("Please upload only one image file.") | |
| if audio is not None: | |
| if len(audios) > 0: | |
| raise gr.Error("Please upload only one audio file or record audio.") | |
| audios.append(audio) | |
| # Append the message to the history | |
| for image in images: | |
| history.append({"role": "user", "content": (image,), "metadata": {"title": "image"}}) | |
| for audio in audios: | |
| history.append({"role": "user", "content": (audio,), "metadata": {"title": "audio"}}) | |
| if message["text"]: | |
| history.append({"role": "user", "content": message["text"]}) | |
| return history, gr.MultimodalTextbox(value=None, interactive=False), None | |
| def bot( | |
| history: list, | |
| top_p: float, | |
| top_k: int, | |
| temperature: float, | |
| repetition_penalty: float, | |
| max_new_tokens: int = MAX_NEW_TOKENS, | |
| regenerate: bool = False, | |
| ): | |
| if history and regenerate: | |
| history = history[:-1] | |
| if not history: | |
| return history | |
| msgs = history2messages(history) | |
| API_URL = "http://8.152.0.142:8000/v1/chat" | |
| payload = { | |
| "messages": msgs, | |
| "sampling_params": { | |
| "top_p": top_p, | |
| "top_k": top_k, | |
| "temperature": temperature, | |
| "repetition_penalty": repetition_penalty, | |
| "max_new_tokens": max_new_tokens, | |
| "num_beams": 3, | |
| } | |
| } | |
| response = requests.get( | |
| API_URL, | |
| json=payload, | |
| headers={'Accept': 'text/event-stream'}, | |
| stream=True | |
| ) | |
| response_text = "" | |
| for text in parse_sse_response(response): | |
| response_text += text | |
| yield history + [{"role": "assistant", "content": response_text}] | |
| return response_text | |
| def change_state(state): | |
| return gr.update(visible=not state), not state | |
| def reset_user_input(): | |
| return gr.update(value="") | |
| if __name__ == "__main__": | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown( | |
| f""" | |
| # 🪐 Chat with <a href="https://github.com/infinigence/Infini-Megrez-Omni">Megrez-3B-Omni</a> | |
| """ | |
| ) | |
| chatbot = gr.Chatbot(elem_id="chatbot", bubble_full_width=False, type="messages", height='48vh') | |
| sampling_params_group_hidden_state = gr.State(False) | |
| with gr.Row(equal_height=True): | |
| chat_input = gr.MultimodalTextbox( | |
| file_count="multiple", | |
| placeholder="Enter your prompt or upload image/audio here, then press ENTER...", | |
| show_label=False, | |
| scale=8, | |
| file_types=["image", "audio"], | |
| interactive=True, | |
| # stop_btn=True, | |
| ) | |
| with gr.Row(equal_height=True): | |
| audio_input = gr.Audio( | |
| sources=["microphone", "upload"], | |
| type="filepath", | |
| scale=1, | |
| max_length=30 | |
| ) | |
| with gr.Row(equal_height=True): | |
| with gr.Column(scale=1, min_width=150): | |
| with gr.Row(equal_height=True): | |
| regenerate_btn = gr.Button("Regenerate", variant="primary") | |
| clear_btn = gr.ClearButton( | |
| [chat_input, audio_input, chatbot], | |
| ) | |
| with gr.Row(): | |
| sampling_params_toggle_btn = gr.Button("Sampling Parameters") | |
| with gr.Group(visible=False) as sampling_params_group: | |
| with gr.Row(): | |
| temperature = gr.Slider( | |
| minimum=0, maximum=1.2, value=DEFAULT_SAMPLING_PARAMS["temperature"], label="Temperature" | |
| ) | |
| repetition_penalty = gr.Slider( | |
| minimum=0, | |
| maximum=2, | |
| value=DEFAULT_SAMPLING_PARAMS["repetition_penalty"], | |
| label="Repetition Penalty", | |
| ) | |
| with gr.Row(): | |
| top_p = gr.Slider(minimum=0, maximum=1, value=DEFAULT_SAMPLING_PARAMS["top_p"], label="Top-p") | |
| top_k = gr.Slider(minimum=0, maximum=1000, value=DEFAULT_SAMPLING_PARAMS["top_k"], label="Top-k") | |
| with gr.Row(): | |
| max_new_tokens = gr.Slider( | |
| minimum=1, | |
| maximum=MAX_NEW_TOKENS, | |
| value=MAX_NEW_TOKENS, | |
| label="Max New Tokens", | |
| interactive=True, | |
| ) | |
| sampling_params_toggle_btn.click( | |
| change_state, | |
| sampling_params_group_hidden_state, | |
| [sampling_params_group, sampling_params_group_hidden_state], | |
| ) | |
| chat_msg = chat_input.submit( | |
| check_messages, | |
| [chatbot, chat_input, audio_input], | |
| [chatbot, chat_input, audio_input], | |
| ) | |
| bot_msg = chat_msg.then( | |
| bot, | |
| inputs=[chatbot, top_p, top_k, temperature, repetition_penalty, max_new_tokens], | |
| outputs=chatbot, | |
| api_name="bot_response", | |
| ) | |
| bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input]) | |
| regenerate_btn.click( | |
| bot, | |
| inputs=[chatbot, top_p, top_k, temperature, repetition_penalty, max_new_tokens, gr.State(True)], | |
| outputs=chatbot, | |
| ) | |
| demo.launch() |