Update README.md

README.md

@@ -38,16 +38,12 @@ An example of using this model to run on your video. Please first clone [Otter](
 ```python
 import mimetypes
 import os
-from io import BytesIO
 from typing import Union
 import cv2
 import requests
 import torch
 import transformers
 from PIL import Image
-from torchvision.transforms import Compose, Resize, ToTensor
-from tqdm import tqdm
-import sys
 
 from otter.modeling_otter import OtterForConditionalGeneration
 
@@ -118,9 +114,7 @@ def get_formatted_prompt(prompt: str) -> str:
 
 def get_response(input_data, prompt: str, model=None, image_processor=None) -> str:
     if isinstance(input_data, Image.Image):
-        vision_x = (
-            image_processor.preprocess([input_data], return_tensors="pt")["pixel_values"].unsqueeze(1).unsqueeze(0)
-        )
+        vision_x = image_processor.preprocess([input_data], return_tensors="pt")["pixel_values"].unsqueeze(1).unsqueeze(0)
     elif isinstance(input_data, list):  # list of video frames
         vision_x = image_processor.preprocess(input_data, return_tensors="pt")["pixel_values"].unsqueeze(1).unsqueeze(0)
     else:
@@ -154,31 +148,36 @@ def get_response(input_data, prompt: str, model=None, image_processor=None) -> s
     )
     return parsed_output
 
-
-# ------------------- Main Function -------------------
-
 if __name__ == "__main__":
-    model = OtterForConditionalGeneration.from_pretrained(
-        "luodian/otter-9b-dc-hf",
-    )
-    model.text_tokenizer.padding_side = "left"
-    tokenizer = model.text_tokenizer
-    image_processor = transformers.CLIPImageProcessor()
-    model.eval()
-
-    while True:
-        video_url = "dc_demo.mp4"  # Replace with the path to your video file
-
-        frames_list = get_image(video_url)
-
-        prompts_input = input("Enter prompts (comma-separated): ")
-        prompts = [prompt.strip() for prompt in prompts_input.split(",")]
-
-        for prompt in prompts:
-            print(f"\nPrompt: {prompt}")
-            response = get_response(frames_list, prompt, model, image_processor)
-            print(f"Response: {response}")
-
-        if prompts_input.lower() == "quit":
-            break
+  # ------------------- Main Function -------------------
+  load_bit = "fp16"
+  if load_bit == "fp16":
+      precision = {"torch_dtype": torch.float16}
+  elif load_bit == "bf16":
+      precision = {"torch_dtype": torch.bfloat16}
+  elif load_bit == "fp32":
+      precision = {"torch_dtype": torch.float32}
+  
+  # This model version is trained on MIMIC-IT DC dataset.
+  model = OtterForConditionalGeneration.from_pretrained("luodian/otter-9b-dc-hf", device_map="auto", **precision)
+  model.text_tokenizer.padding_side = "left"
+  tokenizer = model.text_tokenizer
+  image_processor = transformers.CLIPImageProcessor()
+  model.eval()
+  
+  while True:
+      video_url = "demo.mp4"  # Replace with the path to your video file
+  
+      frames_list = get_image(video_url)
+  
+      prompts_input = input("Enter prompts (comma-separated): ")
+      prompts = [prompt.strip() for prompt in prompts_input.split(",")]
+  
+      for prompt in prompts:
+          print(f"\nPrompt: {prompt}")
+          response = get_response(frames_list, prompt, model, image_processor)
+          print(f"Response: {response}")
+  
+      if prompts_input.lower() == "quit":
+          break
 ```