fixes

app.py

@@ -30,66 +30,91 @@ def sample_frames(video_file, num_frames):
     video.release()
     return frames
 
-@spaces.GPU
 def bot_streaming(message, history):
 
-  txt = message.text
-  ext_buffer = f"user\n{txt} assistant"
+  txt = message["text"]
+  ext_buffer = f"USER: {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:
-    # 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]
+      
+    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
+
+    for i, item in enumerate(history):
+        if all(isinstance(sub_item, str) for sub_item in item):
+            latest_text_only_index = i
 
-  if message.files is None:
+    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:
       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 = sample_frames(image[0], 32)
-        image = None
-        prompt = f"<|im_start|>user <video>\n{message.text}<|im_end|><|im_start|>assistant"
+        video_list = sample_frames(image[0], 12)
+        
+        prompt = f"USER: <video> {message.text} ASSISTANT:"
     elif image[0].endswith(image_extensions):
-        image = Image.open(image[0]).convert("RGB")
-        video = None
-        prompt = f"<|im_start|>user <image>\n{message.text}<|im_end|><|im_start|>assistant"
+        image_list.append(Image.open(image[0]).convert("RGB"))
+        msg = message["text"]
+        prompt =  f"USER: <image> {message.text} ASSISTANT:"
 
   elif len(image) > 1:
-    image_list = []
-    user_prompt = message.text
+    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):        
-        frames = sample_frames(img, 6)
-        for frame in frames:
-          image_list.append(frame)
-      
-    toks = "<image>" * len(image_list)
-    prompt = "<|im_start|>user"+ toks + f"\n{user_prompt}<|im_end|><|im_start|>assistant"
+      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 = ""
 
-    image = image_list
-    video = None
+    if image_list != []:
+      image_tokens = "<image>" * len(image_list)
+    if video_list != []:   
+      
+      toks = len(video_list) 
+      video_tokens = "<video>" * toks
+      
+    
 
+    prompt = f"USER: {image_tokens}{video_tokens} {user_prompt} ASSISTANT:"
 
-  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)
+  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)
   generated_text = ""
 
   thread = Thread(target=model.generate, kwargs=generation_kwargs)
@@ -101,10 +126,10 @@ def bot_streaming(message, history):
   for new_text in streamer:
     
     buffer += new_text
-    
-    generated_text_without_prompt = buffer[len(ext_buffer):]
+    print("new_text", new_text)
+    #generated_text_without_prompt = buffer[len(ext_buffer):][:-1]
     time.sleep(0.01)
-    yield generated_text_without_prompt
+    yield buffer #generated_text_without_prompt
 
 
 demo = gr.ChatInterface(fn=bot_streaming, title="LLaVA Onevision", examples=[