Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # must come before torch | |
| import tempfile | |
| import time | |
| import av # PyAV: used to check whether a video has an audio stream | |
| import numpy as np | |
| import torch | |
| import soundfile as sf | |
| import gradio as gr | |
| from transformers import AutoModelForMultimodalLM, AutoProcessor | |
| MODEL_ID = "LiquidAI/LFM3-Omni-2B" | |
| DEVICE = "cuda" | |
| DTYPE = torch.bfloat16 | |
| print("Loading model...") | |
| model = ( | |
| AutoModelForMultimodalLM.from_pretrained( | |
| MODEL_ID, dtype=DTYPE, trust_remote_code=True | |
| ) | |
| .to(DEVICE) | |
| .eval() | |
| ) | |
| print("Loading processor...") | |
| processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True) | |
| codec = processor.audio_tokenizer.to(device=DEVICE, dtype=DTYPE).eval() | |
| print("Ready.") | |
| SYSTEM_PROMPTS = { | |
| "Text": "Respond with text.", | |
| "Speech": "Respond in audio only.", | |
| "Text + Speech": "Respond with interleaved text and audio.", | |
| } | |
| GENERATION_MODES = { | |
| "Text": "text", | |
| "Speech": "audio", | |
| "Text + Speech": "interleaved", | |
| } | |
| def _video_has_audio(path: str) -> bool: | |
| """Return True if the video file has at least one audio stream.""" | |
| try: | |
| with av.open(path) as container: | |
| return any(s.type == "audio" for s in container.streams) | |
| except Exception: | |
| return False | |
| def _save_wav(waveform) -> str | None: | |
| """Save a 24 kHz waveform tensor to a temp wav file for Gradio.""" | |
| if waveform is None: | |
| return None | |
| arr = waveform.detach().squeeze(0).float().cpu().numpy() | |
| if arr.ndim == 0 or arr.size == 0: | |
| return None | |
| path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name | |
| sf.write(path, arr, 24000) | |
| return path | |
| def respond( | |
| message_text: str, | |
| audio_input, | |
| image_input, | |
| video_input, | |
| output_modality: str = "Text", | |
| max_new_tokens: int = 384, | |
| use_video_soundtrack: bool = False, | |
| history_state: list | None = None, | |
| ): | |
| """Answer a question about text, an image, a video, or a recording. | |
| Args: | |
| message_text: The user's text message or question. | |
| audio_input: Optional audio clip (microphone or upload). | |
| image_input: Optional image to talk about. | |
| video_input: Optional video clip to talk about. | |
| output_modality: "Text", "Speech", or "Text + Speech". | |
| max_new_tokens: Maximum number of tokens to generate. | |
| use_video_soundtrack: Let the model hear the video's audio track too. | |
| history_state: Chat history (list of {role, content} dicts). | |
| Returns: | |
| (chat_history, reply_text, reply_audio_path, timing_note) | |
| """ | |
| t0 = time.perf_counter() | |
| modality = output_modality if output_modality in GENERATION_MODES else "Text" | |
| system = SYSTEM_PROMPTS[modality] | |
| generation_mode = GENERATION_MODES[modality] | |
| content = [] | |
| if image_input is not None: | |
| content.append({"type": "image", "image": image_input}) | |
| if video_input is not None: | |
| content.append({"type": "video", "video": video_input}) | |
| if audio_input is not None: | |
| content.append({"type": "audio", "audio": audio_input}) | |
| text = (message_text or "").strip() | |
| if text or not content: | |
| content.append({"type": "text", "text": text if text else "Hello!"}) | |
| history = list(history_state or []) | |
| conversation = [{"role": "system", "content": system}] | |
| conversation += history | |
| conversation.append({"role": "user", "content": content}) | |
| load_audio_from_video = bool( | |
| use_video_soundtrack | |
| and video_input is not None | |
| and _video_has_audio(video_input) | |
| ) | |
| processor_kwargs = {} | |
| if load_audio_from_video: | |
| processor_kwargs["processor_kwargs"] = {"use_audio_in_video": True} | |
| inputs = processor.apply_chat_template( | |
| conversation, | |
| add_generation_prompt=True, | |
| tokenize=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| load_audio_from_video=load_audio_from_video, | |
| **processor_kwargs, | |
| ).to(DEVICE) | |
| speech = generation_mode in ("audio", "interleaved") | |
| gen_kwargs = dict( | |
| max_new_tokens=int(max_new_tokens), | |
| generation_mode=generation_mode, | |
| ) | |
| if speech: | |
| gen_kwargs.update(audio_temperature=0.4, audio_repetition_penalty=1.2) | |
| with torch.no_grad(): | |
| outputs = model.generate(**inputs, **gen_kwargs) | |
| reply_text = processor.decode( | |
| outputs.sequences[0, inputs["input_ids"].shape[1] : outputs.sequence_lengths[0]], | |
| skip_special_tokens=True, | |
| ).strip() | |
| reply_audio_path = None | |
| if speech and outputs.audio_codes is not None and int(outputs.audio_lengths[0]) > 0: | |
| waveform = codec.decode( | |
| outputs.audio_codes[:1, :, : outputs.audio_lengths[0]] | |
| ).audio_values[0] | |
| reply_audio_path = _save_wav(waveform) | |
| history.append({"role": "user", "content": text or "(media)"}) | |
| history.append({"role": "assistant", "content": reply_text}) | |
| elapsed = time.perf_counter() - t0 | |
| note = f"Generated in {elapsed:.1f}s on GPU" | |
| return history, "", reply_audio_path, note | |
| CSS = """ | |
| #col-container { max-width: 1100px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="LFM3-Omni-2B") as demo: | |
| gr.Markdown( | |
| """ | |
| # 🎙️ LFM3-Omni-2B — any input, text or speech out | |
| Liquid AI's **omni-modal** 2.2B model: one network that understands **text, images, video and audio**, | |
| and answers in **text, 24 kHz speech, or both**. Ask a question by voice, upload a photo, drop in a | |
| video (with its soundtrack), and pick how it should answer. | |
| """ | |
| ) | |
| with gr.Column(elem_id="col-container"): | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| chatbot = gr.Chatbot(label="Conversation", height=340) | |
| message_text = gr.Textbox( | |
| label="Your message", | |
| placeholder="Ask about the image, video, or audio… or just chat.", | |
| lines=2, | |
| ) | |
| with gr.Row(): | |
| audio_input = gr.Audio( | |
| label="🎤 Speak or upload audio", type="filepath", sources=["microphone", "upload"] | |
| ) | |
| with gr.Row(): | |
| image_input = gr.Image(label="🖼️ Image (optional)", type="filepath") | |
| video_input = gr.Video(label="🎬 Video (optional)") | |
| with gr.Row(): | |
| clear_btn = gr.Button("🧹 New conversation") | |
| with gr.Accordion("Advanced", open=False): | |
| output_modality = gr.Radio( | |
| ["Text", "Speech", "Text + Speech"], | |
| value="Text", | |
| label="Answer modality", | |
| info="Speech output uses audio_temperature=0.4, audio_repetition_penalty=1.2 (recommended).", | |
| ) | |
| max_new_tokens = gr.Slider( | |
| 64, 1024, value=384, step=64, label="Max new tokens" | |
| ) | |
| use_video_soundtrack = gr.Checkbox( | |
| value=True, | |
| label="Use video soundtrack", | |
| info="When the attached video has an audio track, let the model hear it too (see and hear together). Videos without audio fall back to frames only.", | |
| ) | |
| with gr.Column(scale=2): | |
| reply_audio = gr.Audio(label="🔊 Spoken answer (24 kHz)", type="filepath") | |
| timing = gr.Markdown("") | |
| history_state = gr.State([]) | |
| btn = gr.Button("Send", variant="primary") | |
| for trigger in (btn.click, message_text.submit): | |
| trigger( | |
| fn=respond, | |
| inputs=[ | |
| message_text, | |
| audio_input, | |
| image_input, | |
| video_input, | |
| output_modality, | |
| max_new_tokens, | |
| use_video_soundtrack, | |
| history_state, | |
| ], | |
| outputs=[chatbot, message_text, reply_audio, timing], | |
| ) | |
| clear_btn.click( | |
| lambda: ([], "", None, "", []), | |
| outputs=[chatbot, message_text, reply_audio, timing, history_state], | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["Name a famous bridge in San Francisco.", None, None, None, "Text", 256, False], | |
| [ | |
| "What does the sign say, and what objects are around it?", | |
| None, | |
| "assets/sign_beach.jpg", | |
| None, | |
| "Text", | |
| 384, | |
| False, | |
| ], | |
| [ | |
| "Transcribe this recording.", | |
| "assets/librispeech_question.wav", | |
| None, | |
| None, | |
| "Text + Speech", | |
| 384, | |
| False, | |
| ], | |
| [ | |
| "What is happening in this video? Answer in one sentence.", | |
| None, | |
| None, | |
| "assets/chef_wok_flames.mp4", | |
| "Text", | |
| 384, | |
| True, | |
| ], | |
| [ | |
| "Read this out loud: The quick brown fox jumps over the lazy dog.", | |
| None, | |
| None, | |
| None, | |
| "Speech", | |
| 256, | |
| False, | |
| ], | |
| ], | |
| inputs=[ | |
| message_text, | |
| audio_input, | |
| image_input, | |
| video_input, | |
| output_modality, | |
| max_new_tokens, | |
| use_video_soundtrack, | |
| ], | |
| outputs=[chatbot, message_text, reply_audio, timing], | |
| fn=respond, | |
| cache_examples=False, | |
| run_on_click=True, | |
| label="Try an example", | |
| ) | |
| demo.launch(mcp_server=True) | |