LLaVA-1.5をDiscord botに組み込んでみた

はじめに

おはこんにちばんは、えるです❢　

前回のnoteで遊んだLLaVA-1.5ですが、 Discord botに組み込んでみると便利かも？と思ったのでやってみました

LLaVA-1.5については前回noteを、Discord botについては過去のnoteなども参考にしています

画像の解釈が可能なLLM 『LLaVA-1.5』で遊ぶ｜える (note.com)

Discord bot with chatGPT 作成日記｜える (note.com)

実装方法

いつものようにGoogleColabで作ります

１）LLaVAのコピー

!git clone https://github.com/haotian-liu/LLaVA.git
%cd /content/LLaVA

２）Discord botのトークン設定

import os
token = "ここにトークンを入力"#@param {type:"string"}
os.environ['TOKEN'] = token

３）インストール

!pip install --upgrade pip
!pip install -e .
!pip install discord.py

４）実行

!python3 main.py

main.pyの中身は、LLaVA-1.5のCLIを参考に以下の感じにしました
画像はひとまず公式サンプルを指定しています

import discord
import os
import torch
from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
from llava.conversation import conv_templates, SeparatorStyle
from llava.model.builder import load_pretrained_model
from llava.utils import disable_torch_init
from llava.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
from PIL import Image
import requests
from io import BytesIO
from transformers import TextStreamer

def load_image(image_file):
    if image_file.startswith('http') or image_file.startswith('https'):
        response = requests.get(image_file)
        image = Image.open(BytesIO(response.content)).convert('RGB')
    else:
        image = Image.open(image_file).convert('RGB')
    return image

def llava_exe(text, tokenizer, model, context_len, image_tensor, conv, first_flag):

  if first_flag is not None:
      # first message
      if model.config.mm_use_im_start_end:
          text = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + text

      else:
          text = DEFAULT_IMAGE_TOKEN + '\n' + text
      conv.append_message(conv.roles[0], text)
      first_flag = None
  else:
      # later messages
      conv.append_message(conv.roles[0], text)
  conv.append_message(conv.roles[1], None)
  prompt = conv.get_prompt()

  input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
  stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
  keywords = [stop_str]
  stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
  streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

  with torch.inference_mode():
      output_ids = model.generate(
          input_ids,
          images=image_tensor,
          do_sample=True,
          temperature=0.2,
          max_new_tokens=1024,
          streamer=streamer,
          use_cache=True,
          stopping_criteria=[stopping_criteria])

  outputs = tokenizer.decode(output_ids[0, input_ids.shape[1]:]).strip()
  conv.messages[-1][-1] = outputs

  #print("\n", {"prompt": prompt, "outputs": outputs}, "\n")
  return outputs, conv, first_flag


disable_torch_init()
model_path = "liuhaotian/llava-v1.5-13b"
model_name = get_model_name_from_path(model_path)
#7b, 4bit = GPU 7GB
#13b = GPU 30GB
tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name, False, True) 
#tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit)

tokenizer_ini, model_ini, image_processor_ini,context_len_ini = tokenizer, model, image_processor, context_len

conv_mode = "llava_v1"
conv = conv_templates[conv_mode].copy()
roles = conv.roles

image = load_image("https://llava-vl.github.io/static/images/view.jpg") #画像のリンクパスを指定
image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().cuda()
first_flag = True

 #Discord initial
intents = discord.Intents.default()
intents.message_content = True
client = discord.Client(intents=intents)

#token and API seting
token = os.environ['TOKEN']


 #Discord event
@client.event
async def on_ready():
  print(f'We have logged in as {client.user}')

@client.event
async def on_message(message):
  global first_flag, tokenizer, model, image_processor, context_len,conv,image_tensor
  if message.author.bot == True:
    return

  elif message.attachments:
      for attachment in message.attachments:
          img_url = attachment.url
          tokenizer, model, image_processor, context_len = tokenizer_ini, model_ini, image_processor_ini,context_len_ini
          conv = conv_templates[conv_mode].copy()
          roles = conv.roles          
          first_flag=True

          image = load_image(img_url)
          image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().cuda()
          await message.channel.send("---image input & conversation initialize---")

  else:
    text = message.content
    output, conv, first_flag = llava_exe(text, tokenizer, model, context_len, image_tensor, conv, first_flag)
    await message.channel.send(output)

client.run(token)

注意点1：
画像の差し替えがなかなか上手くいかず、、
現状は画像を差し替え時に初期化処理を入れています
公式のGradioデモのコードを参考に実装するといいのかな？　とか思いつつまだ出来てません

注意点2：
処理の関係で、Discord上の添付は画像と見なして動作するので、画像以外を添付すると現状おそらく誤動作します

注意点３：
Globalいっぱい使ってるんで怒られそう（？）

結果

例えば以下の感じで動きました❢　
公式のGradioデモがまだColabで動かせていないこともあり、スマホ民の私としてはしばらくこれで使おうかなあとか思っています

画像の添付と添付後の画像参照