File size: 11,828 Bytes
7d5976d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
51911c0
 
7d5976d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8720372
 
7d5976d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
import gradio as gr
import torch
from transformers import Qwen2_5OmniModel, Qwen2_5OmniProcessor
from qwen_omni_utils import process_mm_info
import soundfile as sf
import tempfile
import spaces

# Initialize the model and processor
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float16

model = Qwen2_5OmniModel.from_pretrained(
    "Qwen/Qwen2.5-Omni-7B",
    torch_dtype=torch_dtype,
    device_map="auto",
    enable_audio_output=True,
    # attn_implementation="flash_attention_2" if torch.cuda.is_available() else None
)

processor = Qwen2_5OmniProcessor.from_pretrained("Qwen/Qwen2.5-Omni-7B")

# System prompt
SYSTEM_PROMPT = {
    "role": "system",
    "content": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
}

# Voice options
VOICE_OPTIONS = {
    "Chelsie (Female)": "Chelsie",
    "Ethan (Male)": "Ethan"
}

@spaces.GPU
def process_input(image, audio, video, text, chat_history, voice_type, enable_audio_output):
    # Combine multimodal inputs
    user_input = {
        "text": text,
        "image": image if image is not None else None,
        "audio": audio if audio is not None else None,
        "video": video if video is not None else None
    }
    
    # Prepare conversation history for model processing
    conversation = [SYSTEM_PROMPT]
    
    # Add previous chat history
    if isinstance(chat_history, list):
        for item in chat_history:
            if isinstance(item, tuple) and len(item) == 2:
                user_msg, bot_msg = item
                conversation.append({"role": "user", "content": user_input_to_content(user_msg)})
                conversation.append({"role": "assistant", "content": bot_msg})
    else:
        # Initialize chat history if it's not a list
        chat_history = []
    
    # Add current user input
    conversation.append({"role": "user", "content": user_input_to_content(user_input)})
    
    # Prepare for inference
    text = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
    audios, images, videos = process_mm_info(conversation, use_audio_in_video=True)
    
    inputs = processor(
        text=text, 
        audios=audios, 
        images=images, 
        videos=videos, 
        return_tensors="pt", 
        padding=True
    )
    inputs = inputs.to(model.device).to(model.dtype)
    
    # Generate response
    if enable_audio_output:
        voice_type_value = VOICE_OPTIONS.get(voice_type, "Chelsie")
        text_ids, audio = model.generate(
            **inputs, 
            use_audio_in_video=True,
            return_audio=True,
            spk=voice_type_value
        )
        
        # Save audio to temporary file
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file:
            sf.write(
                tmp_file.name,
                audio.reshape(-1).detach().cpu().numpy(),
                samplerate=24000,
            )
            audio_path = tmp_file.name
    else:
        text_ids = model.generate(
            **inputs, 
            use_audio_in_video=True,
            return_audio=False
        )
        audio_path = None
    
    # Decode text response
    text_response = processor.batch_decode(
        text_ids, 
        skip_special_tokens=True, 
        clean_up_tokenization_spaces=False
    )[0]
    
    # Clean up text response
    text_response = text_response.strip()
    
    # Format user message for chat history display
    user_message_for_display = str(text) if text is not None else ""
    if image is not None:
        user_message_for_display = (user_message_for_display or "Image uploaded") + " [Image]"
    if audio is not None:
        user_message_for_display = (user_message_for_display or "Audio uploaded") + " [Audio]"
    if video is not None:
        user_message_for_display = (user_message_for_display or "Video uploaded") + " [Video]"
    
    # If empty, provide a default message
    if not user_message_for_display.strip():
        user_message_for_display = "Multimodal input"
    
    # Update chat history with properly formatted entries
    if not isinstance(chat_history, list):
        chat_history = []
    chat_history.append((user_message_for_display, text_response))
    
    # Prepare output
    if enable_audio_output and audio_path:
        return chat_history, text_response, audio_path
    else:
        return chat_history, text_response, None

def user_input_to_content(user_input):
    if isinstance(user_input, str):
        return user_input
    elif isinstance(user_input, dict):
        # Handle file uploads
        content = []
        if "text" in user_input and user_input["text"]:
            content.append({"type": "text", "text": user_input["text"]})
        if "image" in user_input and user_input["image"]:
            content.append({"type": "image", "image": user_input["image"]})
        if "audio" in user_input and user_input["audio"]:
            content.append({"type": "audio", "audio": user_input["audio"]})
        if "video" in user_input and user_input["video"]:
            content.append({"type": "video", "video": user_input["video"]})
        return content
    return user_input

def create_demo():
    with gr.Blocks(title="Qwen2.5-Omni ChatBot", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# Qwen2.5-Omni Multimodal ChatBot")
        gr.Markdown("Experience the omni-modal capabilities of Qwen2.5-Omni through text, images, audio, and video interactions.")
        
        # Hidden placeholder components for text-only input
        placeholder_image = gr.Image(type="filepath", visible=False)
        placeholder_audio = gr.Audio(type="filepath", visible=False)
        placeholder_video = gr.Video(visible=False)
        
        # Chat interface
        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(height=600)
                with gr.Accordion("Advanced Options", open=False):
                    voice_type = gr.Dropdown(
                        choices=list(VOICE_OPTIONS.keys()),
                        value="Chelsie (Female)",
                        label="Voice Type"
                    )
                    enable_audio_output = gr.Checkbox(
                        value=True,
                        label="Enable Audio Output"
                    )
                
                # Multimodal input components
                with gr.Tabs():
                    with gr.TabItem("Text Input"):
                        text_input = gr.Textbox(
                            placeholder="Type your message here...",
                            label="Text Input"
                        )
                        text_submit = gr.Button("Send Text")
                    
                    with gr.TabItem("Multimodal Input"):
                        with gr.Row():
                            image_input = gr.Image(
                                type="filepath",
                                label="Upload Image"
                            )
                            audio_input = gr.Audio(
                                type="filepath",
                                label="Upload Audio"
                            )
                        with gr.Row():
                            video_input = gr.Video(
                                label="Upload Video"
                            )
                        additional_text = gr.Textbox(
                            placeholder="Additional text message...",
                            label="Additional Text"
                        )
                        multimodal_submit = gr.Button("Send Multimodal Input")
                
                clear_button = gr.Button("Clear Chat")
                
            with gr.Column(scale=1):
                gr.Markdown("## Model Capabilities")
                gr.Markdown("""
                **Qwen2.5-Omni can:**
                - Process and understand text
                - Analyze images and answer questions about them
                - Transcribe and understand audio
                - Analyze video content (with or without audio)
                - Generate natural speech responses
                """)
                
                gr.Markdown("### Example Prompts")
                gr.Examples(
                    examples=[
                        ["Describe what you see in this image", "image"],
                        ["What is being said in this audio clip?", "audio"],
                        ["What's happening in this video?", "video"],
                        ["Explain Artificial Intelligence in simple terms", "text"],
                        ["Generate a short story about a robot learning to play AlphaGo", "text"]
                    ],
                    inputs=[text_input, gr.Textbox(visible=False)],
                    label="Text Examples"
                )
                
                audio_output = gr.Audio(
                    label="Model Speech Output",
                    visible=True,
                    autoplay=True
                )
                text_output = gr.Textbox(
                    label="Model Text Response",
                    interactive=False
                )
        
        # Text input handling
        text_submit.click(
            fn=lambda text: str(text) if text is not None else "",
            inputs=text_input,
            outputs=[chatbot],
            queue=False
        ).then(
            fn=process_input,
            inputs=[placeholder_image, placeholder_audio, placeholder_video, text_input, chatbot, voice_type, enable_audio_output],
            outputs=[chatbot, text_output, audio_output]
        )
        
        # Multimodal input handling
        def prepare_multimodal_input(image, audio, video, text):
            # Create a display message that indicates what was uploaded
            display_message = str(text) if text is not None else ""
            if image is not None:
                display_message = (display_message + " " if display_message.strip() else "") + "[Image]"
            if audio is not None:
                display_message = (display_message + " " if display_message.strip() else "") + "[Audio]"
            if video is not None:
                display_message = (display_message + " " if display_message.strip() else "") + "[Video]"
            
            if not display_message.strip():
                display_message = "Multimodal content"
                
            return display_message
        
        multimodal_submit.click(
            fn=prepare_multimodal_input,
            inputs=[image_input, audio_input, video_input, additional_text],
            outputs=[chatbot],
            queue=False
        ).then(
            fn=process_input,
            inputs=[image_input, audio_input, video_input, additional_text, 
                   chatbot, voice_type, enable_audio_output],
            outputs=[chatbot, text_output, audio_output]
        )
        
        # Clear chat
        def clear_chat():
            return [], None, None
        
        clear_button.click(
            fn=clear_chat,
            outputs=[chatbot, text_output, audio_output]
        )
        
        # Update audio output visibility
        def toggle_audio_output(enable_audio):
            return gr.Audio(visible=enable_audio)
        
        enable_audio_output.change(
            fn=toggle_audio_output,
            inputs=enable_audio_output,
            outputs=audio_output
        )
    
    return demo

if __name__ == "__main__":
    demo = create_demo()
    demo.launch(server_name="0.0.0.0", server_port=7860)