Update app.py

app.py

@@ -29,93 +29,67 @@ def sample_frames(video_file, num_frames):
             frames.append(pil_img)
     video.release()
     return frames
-    
+
 @spaces.GPU
 def bot_streaming(message, history):
 
-  txt = message["text"]
-  ext_buffer = f"USER: {txt} ASSISTANT: "
+  txt = message.text
+  ext_buffer = f"user\n{txt} assistant"
 
-  if message["files"]:
-    if len(message["files"]) == 1:
+  if message.files:
+    if len(message.files) == 1:
       image = [message.files[0].path]
     # interleaved images or video
-    elif len(message["files"]) > 1:
-      image = [msg["path"] for msg in message["files"]]
+    elif len(message.files) > 1:
+      image = [msg.path for msg in message.files]
   else:
-      
-    def has_file_data(lst):
-      return any(isinstance(item, FileData) for sublist in lst if isinstance(sublist, tuple) for item in sublist)
-
-    def extract_paths(lst):
-        return [item["path"] for sublist in lst if isinstance(sublist, tuple) for item in sublist if isinstance(item, FileData)]
-
-    latest_text_only_index = -1
+    # if there's no image uploaded for this turn, look for images in the past turns
+    # kept inside tuples, take the last one
+    for hist in history:
+      if type(hist[0])==tuple:
+        image = hist[0][0]
 
-    for i, item in enumerate(history):
-        if all(isinstance(sub_item, str) for sub_item in item):
-            latest_text_only_index = i
-
-    image = [path for i, item in enumerate(history) if i < latest_text_only_index and has_file_data(item) for path in extract_paths(item)]
-
-  if message["files"] is None:
+  if message.files is None:
       gr.Error("You need to upload an image or video for LLaVA to work.")
       
   video_extensions = ("avi", "mp4", "mov", "mkv", "flv", "wmv", "mjpeg")
   image_extensions = Image.registered_extensions()
   image_extensions = tuple([ex for ex, f in image_extensions.items()])
-  image_list = []
-  video_list = []
-
-  print("media", image)
   if len(image) == 1:
     if image[0].endswith(video_extensions):
         
-        video_list = sample_frames(image[0], 12)
-        
-        prompt = f"USER: <video> {message.text} ASSISTANT:"
+        video = sample_frames(image[0], 32)
+        image = None
+        prompt = f"<|im_start|>user <video>\n{message.text}<|im_end|><|im_start|>assistant"
     elif image[0].endswith(image_extensions):
-        image_list.append(Image.open(image[0]).convert("RGB"))
-        msg = message["text"]
-        prompt =  f"USER: <image> {message.text} ASSISTANT:"
+        image = Image.open(image[0]).convert("RGB")
+        video = None
+        prompt = f"<|im_start|>user <image>\n{message.text}<|im_end|><|im_start|>assistant"
 
   elif len(image) > 1:
-    user_prompt = message["text"]
+    image_list = []
+    user_prompt = message.text
 
     for img in image:
       if img.endswith(image_extensions):
         img = Image.open(img).convert("RGB")
         image_list.append(img)
 
-      elif img.endswith(video_extensions):
-        video_list.append(sample_frames(img, 7))
-        #for frame in sample_frames(img, 6):
-          #video_list.append(frame)
-        
-    image_tokens = ""
-    video_tokens = ""
-
-    if image_list != []:
-      image_tokens = "<image>" * len(image_list)
-    if video_list != []:   
+      elif img.endswith(video_extensions):        
+        frames = sample_frames(img, 6)
+        for frame in frames:
+          image_list.append(frame)
       
-      toks = len(video_list) 
-      video_tokens = "<video>" * toks
-      
-    
+    toks = "<image>" * len(image_list)
+    prompt = "<|im_start|>user"+ toks + f"\n{user_prompt}<|im_end|><|im_start|>assistant"
 
-    prompt = f"USER: {image_tokens}{video_tokens} {user_prompt} ASSISTANT:"
+    image = image_list
+    video = None
 
-  if image_list != [] and video_list != []:
-    inputs = processor(text=prompt, images=image_list, videos=video_list, padding=True, return_tensors="pt").to("cuda",torch.float16)
-  elif image_list != [] and video_list == []:
-    inputs = processor(text=prompt, images=image_list, padding=True, return_tensors="pt").to("cuda", torch.float16)
-  elif image_list == [] and video_list != []:
-    inputs = processor(text=prompt, videos=video_list, padding=True, return_tensors="pt").to("cuda", torch.float16)
-  
-  
-  streamer = TextIteratorStreamer(processor, **{"max_new_tokens": 200, "skip_special_tokens": True, "clean_up_tokenization_spaces":True})
-  generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=200)
+
+  inputs = processor(text=prompt, images=image, videos=video, return_tensors="pt").to("cuda", torch.float16)
+  streamer = TextIteratorStreamer(processor, **{"max_new_tokens": 200, "skip_special_tokens": True})
+  generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=100)
   generated_text = ""
 
   thread = Thread(target=model.generate, kwargs=generation_kwargs)
@@ -127,10 +101,10 @@ def bot_streaming(message, history):
   for new_text in streamer:
     
     buffer += new_text
-    print("new_text", new_text)
-    #generated_text_without_prompt = buffer[len(ext_buffer):][:-1]
+    
+    generated_text_without_prompt = buffer[len(ext_buffer):]
     time.sleep(0.01)
-    yield buffer #generated_text_without_prompt
+    yield generated_text_without_prompt
 
 
 demo = gr.ChatInterface(fn=bot_streaming, title="LLaVA Onevision", examples=[