This PR upgrades the space to Llama 3.2V 11B Cot

It uses a more recent model from Meta.

Click on _Merge_ to add this feature.

README.md

@@ -1,14 +1,14 @@
 ---
-title: Llava Llama-3 8B
+title: Llama 3.2V 11B Cot
 emoji: 🔥
 colorFrom: yellow
 colorTo: green
 sdk: gradio
-sdk_version: 4.42.0
+sdk_version: 5.0.1
 app_file: app.py
 pinned: true
 license: apache-2.0
-short_description: Meta Llama3 8b with Llava Multimodal capabilities
+short_description: Meta Llama3 3.2V 11B Cot Multimodal
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).

app.py

@@ -1,105 +1,97 @@
-import time
+from transformers import MllamaForConditionalGeneration, AutoProcessor, TextIteratorStreamer
+from PIL import Image
+import requests
+import torch
 from threading import Thread
-
 import gradio as gr
-import torch
-from PIL import Image
-from transformers import AutoProcessor, LlavaForConditionalGeneration
-from transformers import TextIteratorStreamer
-
+from gradio import FileData
+import time
 import spaces
-
-
-PLACEHOLDER = """
-<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
-   <img src="https://cdn-uploads.huggingface.co/production/uploads/64ccdc322e592905f922a06e/DDIW0kbWmdOQWwy4XMhwX.png" style="width: 80%; max-width: 550px; height: auto; opacity: 0.55;  "> 
-   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">LLaVA-Llama-3-8B</h1>
-   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Llava-Llama-3-8b is a LLaVA model fine-tuned from Meta-Llama-3-8B-Instruct and CLIP-ViT-Large-patch14-336 with ShareGPT4V-PT and InternVL-SFT by XTuner</p>
-</div>
-"""
-
-
-model_id = "xtuner/llava-llama-3-8b-v1_1-transformers"
-
-processor = AutoProcessor.from_pretrained(model_id)
-
-model = LlavaForConditionalGeneration.from_pretrained(
-    model_id,
-    torch_dtype=torch.float16,
-    low_cpu_mem_usage=True,
-)
-
-model.to("cuda:0")
-model.generation_config.eos_token_id = 128009
+import re
+ckpt = "Xkev/Llama-3.2V-11B-cot"
+model = MllamaForConditionalGeneration.from_pretrained(ckpt,
+    torch_dtype=torch.bfloat16).to("cuda")
+processor = AutoProcessor.from_pretrained(ckpt)
 
 
 @spaces.GPU
-def bot_streaming(message, history):
-    print(message)
-    if message["files"]:
-        # message["files"][-1] is a Dict or just a string
-        if type(message["files"][-1]) == dict:
-            image = message["files"][-1]["path"]
-        else:
-            image = message["files"][-1]
+def bot_streaming(message, history, max_new_tokens=250):
+    
+    txt = message["text"]
+    ext_buffer = f"{txt}"
+    
+    messages= [] 
+    images = []
+    
+
+    for i, msg in enumerate(history): 
+        if isinstance(msg[0], tuple):
+            messages.append({"role": "user", "content": [{"type": "text", "text": history[i+1][0]}, {"type": "image"}]})
+            messages.append({"role": "assistant", "content": [{"type": "text", "text": history[i+1][1]}]})
+            images.append(Image.open(msg[0][0]).convert("RGB"))
+        elif isinstance(history[i-1], tuple) and isinstance(msg[0], str):
+            # messages are already handled
+            pass
+        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str): # text only turn
+            messages.append({"role": "user", "content": [{"type": "text", "text": msg[0]}]})
+            messages.append({"role": "assistant", "content": [{"type": "text", "text": msg[1]}]})
+
+    # add current message
+    if len(message["files"]) == 1:
+        
+        if isinstance(message["files"][0], str): # examples
+            image = Image.open(message["files"][0]).convert("RGB")
+        else: # regular input
+            image = Image.open(message["files"][0]["path"]).convert("RGB")
+        images.append(image)
+        messages.append({"role": "user", "content": [{"type": "text", "text": txt}, {"type": "image"}]})
     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]
-    try:
-        if image is None:
-            # Handle the case where image is None
-            gr.Error("You need to upload an image for LLaVA to work.")
-    except NameError:
-        # Handle the case where 'image' is not defined at all
-        gr.Error("You need to upload an image for LLaVA to work.")
+        messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})
 
-    prompt = f"<|start_header_id|>user<|end_header_id|>\n\n<image>\n{message['text']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-    # print(f"prompt: {prompt}")
-    image = Image.open(image)
-    inputs = processor(prompt, image, return_tensors='pt').to(0, torch.float16)
 
-    streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": False, "skip_prompt": True})
-    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024, do_sample=False)
+    texts = processor.apply_chat_template(messages, add_generation_prompt=True)
 
+    if images == []:
+        inputs = processor(text=texts, return_tensors="pt").to("cuda")
+    else:
+        inputs = processor(text=texts, images=images, return_tensors="pt").to("cuda")
+        
+    streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)
+
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
+    generated_text = ""
+    
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
-
-    text_prompt = f"<|start_header_id|>user<|end_header_id|>\n\n{message['text']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
-    # print(f"text_prompt: {text_prompt}")
-
     buffer = ""
-    time.sleep(0.5)
+    
     for new_text in streamer:
-        # find <|eot_id|> and remove it from the new_text
-        if "<|eot_id|>" in new_text:
-            new_text = new_text.split("<|eot_id|>")[0]
         buffer += new_text
-
-        # generated_text_without_prompt = buffer[len(text_prompt):]
         generated_text_without_prompt = buffer
-        # print(generated_text_without_prompt)
-        time.sleep(0.06)
-        # print(f"new_text: {generated_text_without_prompt}")
-        yield generated_text_without_prompt
-
-
-chatbot=gr.Chatbot(placeholder=PLACEHOLDER,scale=1)
-chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)
-with gr.Blocks(fill_height=True, ) as demo:
-    gr.ChatInterface(
-    fn=bot_streaming,
-    title="LLaVA Llama-3-8B",
-    examples=[{"text": "What is on the flower?", "files": ["./bee.jpg"]},
-              {"text": "How to make this pastry?", "files": ["./baklava.png"]}],
-    description="Try [LLaVA Llama-3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). Upload an image and start chatting about it, or simply try one of the examples below. If you don't upload an image, you will receive an error.",
-    stop_btn="Stop Generation",
-    multimodal=True,
-    textbox=chat_input,
-    chatbot=chatbot,
-    )
-
-demo.queue(api_open=False)
-demo.launch(show_api=False, share=False)
+        time.sleep(0.01)
+    
+        buffer = re.sub(r"<(\w+)>", r"(Here begins the \1 stage)", buffer)  
+        buffer = re.sub(r"</(\w+)>", r"(Here ends the \1 stage)", buffer)  
+    
+        yield buffer
+
+
+demo = gr.ChatInterface(fn=bot_streaming, title="LLaVA-CoT",
+      textbox=gr.MultimodalTextbox(), 
+      additional_inputs = [gr.Slider(
+              minimum=512,
+              maximum=1024,
+              value=512,
+              step=1,
+              label="Maximum number of new tokens to generate",
+          )
+        ],
+      examples=[[{"text": "What is on the flower?", "files": ["./bee.jpg"]},512],
+                [{"text": "How to make this pastry?", "files": ["./baklava.png"]},512]],
+      cache_examples=False,
+      description="Upload an image, and start chatting about it. To learn more about LLaVA-CoT, visit [our GitHub page](https://github.com/PKU-YuanGroup/LLaVA-CoT). Note: Since Gradio currently does not support displaying the special markings in the output, we have replaced it with the expression (Here begins the X phase).",
+      stop_btn="Stop Generation", 
+      fill_height=True,
+    multimodal=True)
+    
+demo.launch(debug=True)

requirements.txt

@@ -1,5 +1,3 @@
 torch
-git+https://github.com/huggingface/transformers.git
 spaces
-pillow
-accelerate
+git+https://github.com/huggingface/transformers.git