GXXyang/VideoLLaMA3

1import os2import os.path as osp3 4import gradio as gr5import spaces6import torch7from threading import Thread8from transformers import AutoModelForCausalLM, AutoProcessor, TextIteratorStreamer9 10 11HEADER = ("""12<div style="display: flex; justify-content: center; align-items: center; text-align: center;">13  <a href="https://github.com/DAMO-NLP-SG/VideoLLaMA3" style="margin-right: 20px; text-decoration: none; display: flex; align-items: center;">14    <img src="https://github.com/DAMO-NLP-SG/VideoLLaMA3/blob/main/assets/logo.png?raw=true" alt="VideoLLaMA 3 🔥🚀🔥" style="max-width: 120px; height: auto;">15  </a>16  <div>17    <h1>VideoLLaMA 3: Frontier Multimodal Foundation Models for Video Understanding</h1>18    <h5 style="margin: 0;">If this demo please you, please give us a star ⭐ on Github or 💖 on this space.</h5>19  </div>20</div>21 22<div style="display: flex; justify-content: center; margin-top: 10px;">23  <a href="https://github.com/DAMO-NLP-SG/VideoLLaMA3"><img src='https://img.shields.io/badge/Github-VideoLLaMA3-9C276A' style="margin-right: 5px;"></a>24  <a href="https://arxiv.org/pdf/2501.13106"><img src="https://img.shields.io/badge/Arxiv-2501.13106-AD1C18" style="margin-right: 5px;"></a>25  <a href="https://huggingface.co/collections/DAMO-NLP-SG/videollama3-678cdda9281a0e32fe79af15"><img src="https://img.shields.io/badge/🤗-Checkpoints-ED5A22.svg" style="margin-right: 5px;"></a>26  <a href="https://github.com/DAMO-NLP-SG/VideoLLaMA3/stargazers"><img src="https://img.shields.io/github/stars/DAMO-NLP-SG/VideoLLaMA3.svg?style=social"></a>27</div>28""")29 30device = "cuda"31model = AutoModelForCausalLM.from_pretrained(32    "DAMO-NLP-SG/VideoLLaMA3-7B",33    trust_remote_code=True,34    torch_dtype=torch.bfloat16,35    attn_implementation="flash_attention_2",36)37model.to(device)38processor = AutoProcessor.from_pretrained("DAMO-NLP-SG/VideoLLaMA3-7B", trust_remote_code=True)39 40 41example_dir = "./examples"42image_formats = ("png", "jpg", "jpeg")43video_formats = ("mp4",)44 45image_examples, video_examples = [], []46if example_dir is not None:47    example_files = [48        osp.join(example_dir, f) for f in os.listdir(example_dir)49    ]50    for example_file in example_files:51        if example_file.endswith(image_formats):52            image_examples.append([example_file])53        elif example_file.endswith(video_formats):54            video_examples.append([example_file])55 56 57def _on_video_upload(messages, video):58        if video is not None:59            # messages.append({"role": "user", "content": gr.Video(video)})60            messages.append({"role": "user", "content": {"path": video}})61        return messages, None62    63def _on_image_upload(messages, image):64    if image is not None:65        # messages.append({"role": "user", "content": gr.Image(image)})66        messages.append({"role": "user", "content": {"path": image}})67    return messages, None68 69def _on_text_submit(messages, text):70    messages.append({"role": "user", "content": text})71    return messages, ""72 73@spaces.GPU(duration=120)74def _predict(messages, input_text, do_sample, temperature, top_p, max_new_tokens,75             fps, max_frames):76    if len(input_text) > 0:77        messages.append({"role": "user", "content": input_text})78    new_messages = []79    contents = []80    for message in messages:81        if message["role"] == "assistant":82            if len(contents):83                new_messages.append({"role": "user", "content": contents})84                contents = []85            new_messages.append(message)86        elif message["role"] == "user":87            if isinstance(message["content"], str):88                contents.append(message["content"])89            else:90                media_path = message["content"][0]91                if media_path.endswith(video_formats):92                    contents.append({"type": "video", "video": {"video_path": media_path, "fps": fps, "max_frames": max_frames}})93                elif media_path.endswith(image_formats):94                    contents.append({"type": "image", "image": {"image_path": media_path}})95                else:96                    raise ValueError(f"Unsupported media type: {media_path}")97 98    if len(contents):99        new_messages.append({"role": "user", "content": contents})100 101    if len(new_messages) == 0 or new_messages[-1]["role"] != "user":102        return messages103 104    generation_config = {105        "do_sample": do_sample,106        "temperature": temperature,107        "top_p": top_p,108        "max_new_tokens": max_new_tokens109    }110 111    inputs = processor(112        conversation=new_messages,113        add_system_prompt=True,114        add_generation_prompt=True,115        return_tensors="pt"116    )117    inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}118    if "pixel_values" in inputs:119        inputs["pixel_values"] = inputs["pixel_values"].to(torch.bfloat16)120 121    streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)122    generation_kwargs = {123        **inputs,124        **generation_config,125        "streamer": streamer,126    }127 128    thread = Thread(target=model.generate, kwargs=generation_kwargs)129    thread.start()130 131    messages.append({"role": "assistant", "content": ""})132    for token in streamer:133        messages[-1]['content'] += token134        yield messages135 136 137with gr.Blocks() as interface:138    gr.HTML(HEADER)139    with gr.Row():140        chatbot = gr.Chatbot(type="messages", elem_id="chatbot", height=835)141 142        with gr.Column():143            with gr.Tab(label="Input"):144 145                with gr.Row():146                    input_video = gr.Video(sources=["upload"], label="Upload Video")147                    input_image = gr.Image(sources=["upload"], type="filepath", label="Upload Image")148                149                input_text = gr.Textbox(label="Input Text", placeholder="Type your message here and press enter to submit")150 151                submit_button = gr.Button("Generate")152 153                gr.Examples(examples=[154                    [f"examples/bear.mp4", "What is unusual in the video?"],155                    [f"examples/dog.mp4", "Please describe the video in detail."],156                    [f"examples/exercise.mp4", "What is the man doing in the video?"],157                ], inputs=[input_video, input_text], label="Video examples")158 159            with gr.Tab(label="Configure"):160                with gr.Accordion("Generation Config", open=True):161                    do_sample = gr.Checkbox(value=True, label="Do Sample")162                    temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, label="Temperature")163                    top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, label="Top P")164                    max_new_tokens = gr.Slider(minimum=0, maximum=4096, value=2048, step=1, label="Max New Tokens")165 166                with gr.Accordion("Video Config", open=True):167                    fps = gr.Slider(minimum=0.0, maximum=10.0, value=1, label="FPS")168                    max_frames = gr.Slider(minimum=0, maximum=256, value=180, step=1, label="Max Frames")169 170    input_video.change(_on_video_upload, [chatbot, input_video], [chatbot, input_video])171    input_image.change(_on_image_upload, [chatbot, input_image], [chatbot, input_image])172    input_text.submit(_on_text_submit, [chatbot, input_text], [chatbot, input_text])173    submit_button.click(174        _predict,175        [176            chatbot, input_text, do_sample, temperature, top_p, max_new_tokens,177            fps, max_frames178        ],179        [chatbot],180    )181 182 183if __name__ == "__main__":184    interface.launch()185