Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,125 Bytes
d710af7 2d637d7 d710af7 2d637d7 d710af7 2d637d7 d710af7 2d637d7 d710af7 2d637d7 d710af7 2d637d7 d710af7 55dbdd5 2d637d7 d710af7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
from transformers import MllamaForConditionalGeneration, AutoProcessor, TextIteratorStreamer
from PIL import Image
import requests
import torch
from threading import Thread
import gradio as gr
from gradio import FileData
import time
import spaces
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, 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:
messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})
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")
generation_kwargs = dict(inputs, max_new_tokens=max_new_tokens)
with torch.no_grad():
output = model.generate(**inputs, **generation_kwargs)
output_text = processor.decode(output[0][inputs['input_ids'].shape[1]:]).replace('<|eot_id|>', '')
yield output_text
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",
)
],
cache_examples=False,
description="Upload an image, and start chatting about it. To learn more about LLaVA-CoT, visit [oir GitHub page](https://github.com/PKU-YuanGroup/LLaVA-CoT). ",
stop_btn="Stop Generation",
fill_height=True,
multimodal=True)
demo.launch(debug=True) |