Update app.py

app.py

@@ -11,10 +11,18 @@ from threading import Thread
 import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, TextIteratorStreamer
+from moviepy.editor import VideoFileClip
+from PIL import Image
 
 model_name = 'AIDC-AI/Ovis2-16B'
+
 use_thread = False
 
+IMAGE_MAX_PARTITION = 16
+
+VIDEO_FRAME_NUMS = 32
+VIDEO_MAX_PARTITION = 1
+
 # load model
 model = AutoModelForCausalLM.from_pretrained(model_name,
                                              torch_dtype=torch.bfloat16,
@@ -49,8 +57,8 @@ def submit_chat(chatbot, text_input):
     return chatbot ,''
 
 @spaces.GPU
-def ovis_chat(chatbot: List[List[str]], image_input: Any):
-    conversations, model_inputs = prepare_inputs(chatbot, image_input)
+def ovis_chat(chatbot: List[List[str]], image_input: Any, video_input: Any):
+    conversations, model_inputs = prepare_inputs(chatbot, image_input, video_input)
     gen_kwargs = initialize_gen_kwargs()
 
     with torch.inference_mode():
@@ -74,7 +82,7 @@ def ovis_chat(chatbot: List[List[str]], image_input: Any):
     log_conversation(chatbot)
 
     
-def prepare_inputs(chatbot: List[List[str]], image_input: Any):
+def prepare_inputs(chatbot: List[List[str]], image_input: Any,  video_input: Any):
     # conversations = [{
     #     "from": "system",
     #     "value": "You are a helpful assistant, and your task is to provide reliable and structured responses to users."
@@ -95,10 +103,30 @@ def prepare_inputs(chatbot: List[List[str]], image_input: Any):
             if conv["from"] == "human":
                 conv["value"] = f'{image_placeholder}\n{conv["value"]}'
                 break
+        max_partition = IMAGE_MAX_PARTITION
+        image_input = [image_input]
+    
+    if video_input is not None:
+        for conv in conversations:
+            if conv["from"] == "human":
+                conv["value"] = f'{image_placeholder}\n' * VIDEO_FRAME_NUMS + f'{conv["value"]}'
+                break
+        # extract video frames here
+        with VideoFileClip(video_input) as clip:
+            total_frames = int(clip.fps * clip.duration)
+            if total_frames <= VIDEO_FRAME_NUMS:
+                sampled_indices = range(total_frames)
+            else:
+                stride = total_frames / VIDEO_FRAME_NUMS
+                sampled_indices = [min(total_frames - 1, int((stride * i + stride * (i + 1)) / 2)) for i in range(VIDEO_FRAME_NUMS)]
+            frames = [clip.get_frame(index / clip.fps) for index in sampled_indices]
+            frames = [Image.fromarray(frame, mode='RGB') for frame in frames]
+        image_input = frames
+        max_partition = VIDEO_MAX_PARTITION
 
     logger.info(conversations)
     
-    prompt, input_ids, pixel_values = model.preprocess_inputs(conversations, [image_input], max_partition=16)
+    prompt, input_ids, pixel_values = model.preprocess_inputs(conversations, image_input, max_partition=max_partition)
     attention_mask = torch.ne(input_ids, text_tokenizer.pad_token_id)
     
     model_inputs = {
@@ -115,7 +143,7 @@ def log_conversation(chatbot):
     logger.info("[OVIS_CONV_END]")
 
 def clear_chat():
-    return [], None, ""
+    return [], None, "", None
 
 with open(f"{cur_dir}/resource/logo.svg", "r", encoding="utf-8") as svg_file:
     svg_content = svg_file.read()
@@ -164,26 +192,58 @@ with gr.Blocks(title=model_name.split('/')[-1], theme=gr.themes.Ocean()) as demo
     gr.HTML(html)
     with gr.Row():
         with gr.Column(scale=3):
-            image_input = gr.Image(label="image", height=350, type="pil")
-            gr.Examples(
-                examples=[
-                    [f"{cur_dir}/examples/ovis2_math0.jpg", "Each face of the polyhedron shown is either a triangle or a square. Each square borders 4 triangles, and each triangle borders 3 squares. The polyhedron has 6 squares. How many triangles does it have?\n\nProvide a step-by-step solution to the problem, and conclude with 'the answer is' followed by the final solution."],
-                    [f"{cur_dir}/examples/ovis2_math1.jpg", "A large square touches another two squares, as shown in the picture. The numbers inside the smaller squares indicate their areas. What is the area of the largest square?\n\nProvide a step-by-step solution to the problem, and conclude with 'the answer is' followed by the final solution."],
-                    [f"{cur_dir}/examples/ovis2_figure0.png", "Explain this model."],
-                    [f"{cur_dir}/examples/ovis2_figure1.png", "Organize the notes about GRPO in the figure."],
-                    [f"{cur_dir}/examples/ovis2_multi0.jpg", "Posso avere un frappuccino e un caffè americano di taglia M? Quanto costa in totale?"],
-                ],
-                inputs=[image_input, text_input]
-            )
+            input_type = gr.Radio(choices=["image + prompt", "video + prompt"], label="Select input type:", value="image + prompt", elem_classes="my_radio")
+
+            image_input = gr.Image(label="image", height=350, type="pil", visible=True)
+            video_input = gr.Video(label="video", height=350, format='mp4', visible=False)
+            with gr.Column(visible=True) as image_examples_col:
+                image_examples = gr.Examples(
+                    examples=[
+                        [f"{cur_dir}/examples/ovis2_math0.jpg", "Each face of the polyhedron shown is either a triangle or a square. Each square borders 4 triangles, and each triangle borders 3 squares. The polyhedron has 6 squares. How many triangles does it have?\n\nProvide a step-by-step solution to the problem, and conclude with 'the answer is' followed by the final solution."],
+                        [f"{cur_dir}/examples/ovis2_math1.jpg", "A large square touches another two squares, as shown in the picture. The numbers inside the smaller squares indicate their areas. What is the area of the largest square?\n\nProvide a step-by-step solution to the problem, and conclude with 'the answer is' followed by the final solution."],
+                        [f"{cur_dir}/examples/ovis2_figure0.png", "Explain this model."],
+                        [f"{cur_dir}/examples/ovis2_figure1.png", "Organize the notes about GRPO in the figure."],
+                        [f"{cur_dir}/examples/ovis2_multi0.jpg", "Posso avere un frappuccino e un caffè americano di taglia M? Quanto costa in totale?"],
+                    ],
+                    inputs=[image_input, text_input]
+                )
+
+            def update_visibility_on_example(video_input, text_input):
+                return (gr.update(visible=True), text_input)
+            
+            with gr.Column(visible=False) as video_examples_col:
+                video_examples = gr.Examples(
+                    examples=[
+                        [f"{cur_dir}/examples/video_demo_1.mp4", "Describe the video."]
+                    ],
+                    inputs=[video_input, text_input],
+                    fn = update_visibility_on_example,
+                    run_on_click = True,
+                    outputs=[video_input, text_input]
+                )
+
         with gr.Column(scale=7):
             chatbot = gr.Chatbot(label="Ovis", layout="panel", height=600, show_copy_button=True, latex_delimiters=latex_delimiters_set)
             text_input.render()
             with gr.Row():
                 send_btn = gr.Button("Send", variant="primary")
                 clear_btn = gr.Button("Clear", variant="secondary")
-
-    send_click_event = send_btn.click(submit_chat, [chatbot, text_input], [chatbot, text_input]).then(ovis_chat,[chatbot, image_input],chatbot)
-    submit_event = text_input.submit(submit_chat, [chatbot, text_input], [chatbot, text_input]).then(ovis_chat,[chatbot, image_input],chatbot)
-    clear_btn.click(clear_chat, outputs=[chatbot, image_input, text_input])
+                
+        def update_input_and_clear(selected):
+            if selected == "image + prompt":
+                visibility_updates = (gr.update(visible=True), gr.update(visible=False), 
+                                      gr.update(visible=True), gr.update(visible=False))
+            else:
+                visibility_updates = (gr.update(visible=False), gr.update(visible=True), 
+                                      gr.update(visible=False), gr.update(visible=True))
+            clear_chat_outputs = clear_chat()
+            return visibility_updates + clear_chat_outputs
+
+        input_type.change(fn=update_input_and_clear, inputs=input_type, 
+                        outputs=[image_input, video_input, image_examples_col, video_examples_col, chatbot, image_input, text_input, video_input])
+
+    send_click_event = send_btn.click(submit_chat, [chatbot, text_input], [chatbot, text_input]).then(ovis_chat,[chatbot, image_input, video_input],chatbot)
+    submit_event = text_input.submit(submit_chat, [chatbot, text_input], [chatbot, text_input]).then(ovis_chat,[chatbot, image_input, video_input],chatbot)
+    clear_btn.click(clear_chat, outputs=[chatbot, image_input, text_input, video_input])
 
 demo.launch()