20231218_japanese_stable_vlm.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/aicuai/GenAI-Steam/blob/main/20231218_japanese_stable_vlm.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K92bGOlkLNC5"
      },
      "outputs": [],
      "source": [
        "#@title Setup\n",
        "!nvidia-smi\n",
        "#!pip install ftfy regex tqdm gradio transformers sentencepiece 'accelerate>=0.12.0' 'bitsandbytes>=0.31.5'\n",
        "!pip install ftfy regex tqdm gradio sentencepiece transformers accelerate 'bitsandbytes>=0.31.5'"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Reboot kernel for accelarate package\n",
        "import os\n",
        "os.kill(os.getpid(), 9)"
      ],
      "metadata": {
        "id": "aZGcejXM2j3X"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uNNRyH3WLw84"
      },
      "outputs": [],
      "source": [
        "# @title Login HuggingFace\n",
        "# @markdown [Put HuggingFrace User Access Tokens here](https://huggingface.co/settings/tokens)\n",
        "# @markdown  StabilityAI 社の HuggingFace で [JVLM のライセンス](https://huggingface.co/stabilityai/japanese-stable-vlm) に合意するのを忘れずに。\n",
        "HUGGINGFACE_TOKEN = ''  # @param {type: \"string\"}\n",
        "\n",
        "# !huggingface-cli login\n",
        "!huggingface-cli login --token $HUGGINGFACE_TOKEN"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mKC-h58pLyOu",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Load Japanese Stable VLM\n",
        "import torch\n",
        "from transformers import AutoTokenizer, AutoModelForVision2Seq, AutoImageProcessor\n",
        "\n",
        "from PIL import Image\n",
        "import requests\n",
        "\n",
        "# helper function to format input prompts\n",
        "TASK2INSTRUCTION = {\n",
        "    \"caption\": \"画像を詳細に述べてください。\",\n",
        "    \"tag\": \"与えられた単語を使って、画像を詳細に述べてください。\",\n",
        "    \"vqa\": \"与えられた画像を下に、質問に答えてください。\",\n",
        "}\n",
        "\n",
        "\n",
        "def build_prompt(task=\"caption\", input=None, sep=\"\\n\\n### \"):\n",
        "    assert (\n",
        "        task in TASK2INSTRUCTION\n",
        "    ), f\"Please choose from {list(TASK2INSTRUCTION.keys())}\"\n",
        "    if task in [\"tag\", \"vqa\"]:\n",
        "        assert input is not None, \"Please fill in `input`!\"\n",
        "        if task == \"tag\" and isinstance(input, list):\n",
        "            input = \"、\".join(input)\n",
        "    else:\n",
        "        assert input is None, f\"`{task}` mode doesn't support to input questions\"\n",
        "    sys_msg = \"以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。\"\n",
        "    p = sys_msg\n",
        "    roles = [\"指示\", \"応答\"]\n",
        "    instruction = TASK2INSTRUCTION[task]\n",
        "    msgs = [\": \\n\" + instruction, \": \\n\"]\n",
        "    if input:\n",
        "        roles.insert(1, \"入力\")\n",
        "        msgs.insert(1, \": \\n\" + input)\n",
        "    for role, msg in zip(roles, msgs):\n",
        "        p += sep + role + msg\n",
        "    return p\n",
        "\n",
        "\n",
        "# load model\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "load_in = \"int8\" # @param [\"fp32\", \"fp16\", \"int8\"]\n",
        "# @markdown If you use Colab free plan, please set `load_in` to `int8`. But, please remember that `int8` degrades the performance. In general, `fp32` is better than `fp16` and `fp16` is better than `int8`.\n",
        "# @markdown  Google Colab Proを購入している方は V100環境でfp16ならロードできると思いますが、安定した実行環境のためには int8 が無難かもしれません。\n",
        "\n",
        "model_kwargs = {\"trust_remote_code\": True, \"low_cpu_mem_usage\": True}\n",
        "if load_in == \"fp16\":\n",
        "  model_kwargs[\"variant\"] = \"fp16\"\n",
        "  model_kwargs[\"torch_dtype\"] = torch.float16\n",
        "elif load_in == \"int8\":\n",
        "  model_kwargs[\"variant\"] = \"fp16\"\n",
        "  model_kwargs[\"load_in_8bit\"] = True\n",
        "  model_kwargs[\"max_memory\"] = f'{int(torch.cuda.mem_get_info()[0]/1024**3)-2}GB'\n",
        "\n",
        "model = AutoModelForVision2Seq.from_pretrained(\"stabilityai/japanese-stable-vlm\", **model_kwargs)\n",
        "processor = AutoImageProcessor.from_pretrained(\"stabilityai/japanese-stable-vlm\")\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"stabilityai/japanese-stable-vlm\")\n",
        "if load_in != \"int8\":\n",
        "  model.to(device)\n",
        "model = model.eval()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u2biq4NwMPMR",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Prepare for the demo\n",
        "\n",
        "\n",
        "@torch.inference_mode()\n",
        "def inference_fn(\n",
        "    image,\n",
        "    task,\n",
        "    prompt,\n",
        "    min_len,\n",
        "    max_len,\n",
        "    beam_size,\n",
        "    len_penalty,\n",
        "    repetition_penalty,\n",
        "    top_p,\n",
        "    decoding_method,\n",
        "    num_return_sequences=3,\n",
        "):\n",
        "    prompt = build_prompt(task=task, input=prompt if not task == \"caption\" else None)\n",
        "    print(f\"instruction: {prompt}\")\n",
        "    inputs = processor(images=image, return_tensors=\"pt\")\n",
        "    text_encoding = tokenizer(prompt, add_special_tokens=False, return_tensors=\"pt\")\n",
        "    inputs.update(text_encoding)\n",
        "    generation_kwargs = {\n",
        "        \"do_sample\": decoding_method == \"Nucleus sampling\",\n",
        "        \"length_penalty\": float(len_penalty),\n",
        "        \"repetition_penalty\": float(repetition_penalty),\n",
        "        \"num_beams\": beam_size,\n",
        "        \"max_new_tokens\": max_len,\n",
        "        \"min_length\": min_len,\n",
        "        \"top_p\": top_p,\n",
        "        \"num_return_sequences\": 3,\n",
        "    }\n",
        "    outputs = model.generate(\n",
        "        **inputs.to(device, dtype=model.dtype), **generation_kwargs\n",
        "    )\n",
        "    generated = [\n",
        "        txt.strip() for txt in tokenizer.batch_decode(outputs, skip_special_tokens=True)\n",
        "    ]\n",
        "    if num_return_sequences > 1:\n",
        "        generated = \"\\n\".join([f\"{i}: {g}\" for i, g in enumerate(generated)])\n",
        "    else:\n",
        "        generated = generated[0]\n",
        "    del inputs\n",
        "    del outputs\n",
        "    torch.cuda.empty_cache()\n",
        "    return generated\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MPX7dCw1MAof",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title Launch the demo\n",
        "import gradio as gr\n",
        "\n",
        "with gr.Blocks() as demo:\n",
        "    gr.Markdown(f\"# Japanese VLM Demo\")\n",
        "    gr.Markdown(\n",
        "        \"\"\"[Japanese Stable VLM](https://huggingface.co/stabilityai/japanese-stable-vlm) is a Japanese vision-language model by [Stability AI](https://ja.stability.ai/).\n",
        "                - Blog: https://ja.stability.ai/blog/japanese-stable-vlm\n",
        "                - Twitter: https://twitter.com/StabilityAI_JP\n",
        "                - Discord: https://discord.com/invite/StableJP\"\"\"\n",
        "    )\n",
        "    with gr.Row():\n",
        "        with gr.Column():\n",
        "            input_image = gr.Image(type=\"pil\", label=\"image\")\n",
        "            task = gr.Radio(\n",
        "                choices=list(TASK2INSTRUCTION.keys()), value=0, label=\"task\"\n",
        "            )\n",
        "            prompt = gr.Textbox(label=\"input\", value=\"\")\n",
        "            with gr.Accordion(label=\"Configs\", open=False):\n",
        "                min_len = gr.Slider(\n",
        "                    minimum=1,\n",
        "                    maximum=50,\n",
        "                    value=1,\n",
        "                    step=1,\n",
        "                    interactive=True,\n",
        "                    label=\"Min Length\",\n",
        "                )\n",
        "\n",
        "                max_len = gr.Slider(\n",
        "                    minimum=10,\n",
        "                    maximum=100,\n",
        "                    value=65,\n",
        "                    step=5,\n",
        "                    interactive=True,\n",
        "                    label=\"Max New Tokens\",\n",
        "                )\n",
        "\n",
        "                sampling = gr.Radio(\n",
        "                    choices=[\"Beam search\", \"Nucleus sampling\"],\n",
        "                    value=\"Beam search\",\n",
        "                    label=\"Text Decoding Method\",\n",
        "                    interactive=True,\n",
        "                )\n",
        "\n",
        "                top_p = gr.Slider(\n",
        "                    minimum=0.5,\n",
        "                    maximum=1.0,\n",
        "                    value=0.9,\n",
        "                    step=0.1,\n",
        "                    interactive=True,\n",
        "                    label=\"Top p\",\n",
        "                )\n",
        "\n",
        "                beam_size = gr.Slider(\n",
        "                    minimum=1,\n",
        "                    maximum=10,\n",
        "                    value=5,\n",
        "                    step=1,\n",
        "                    interactive=True,\n",
        "                    label=\"Beam Size\",\n",
        "                )\n",
        "\n",
        "                len_penalty = gr.Slider(\n",
        "                    minimum=-1,\n",
        "                    maximum=2,\n",
        "                    value=1,\n",
        "                    step=0.2,\n",
        "                    interactive=True,\n",
        "                    label=\"Length Penalty\",\n",
        "                )\n",
        "\n",
        "                repetition_penalty = gr.Slider(\n",
        "                    minimum=-1,\n",
        "                    maximum=3,\n",
        "                    value=1.5,\n",
        "                    step=0.2,\n",
        "                    interactive=True,\n",
        "                    label=\"Repetition Penalty\",\n",
        "                )\n",
        "                num_return_sequences = gr.Number(\n",
        "                    value=3, label=\"Number of Outputs\", precision=0\n",
        "                )\n",
        "            # button\n",
        "            input_button = gr.Button(value=\"Submit\")\n",
        "        with gr.Column():\n",
        "            output = gr.Textbox(label=\"Output\")\n",
        "\n",
        "    inputs = [\n",
        "        input_image,\n",
        "        task,\n",
        "        prompt,\n",
        "        min_len,\n",
        "        max_len,\n",
        "        beam_size,\n",
        "        len_penalty,\n",
        "        repetition_penalty,\n",
        "        top_p,\n",
        "        sampling,\n",
        "        num_return_sequences,\n",
        "    ]\n",
        "    input_button.click(inference_fn, inputs=inputs, outputs=[output])\n",
        "    prompt.submit(inference_fn, inputs=inputs, outputs=[output])\n",
        "    img2txt_examples = gr.Examples(\n",
        "        examples=[\n",
        "            [\n",
        "                \"https://images.unsplash.com/photo-1582538885592-e70a5d7ab3d3?ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D&auto=format&fit=crop&w=1770&q=80\",\n",
        "                \"caption\",\n",
        "                \"\",\n",
        "                1,\n",
        "                32,\n",
        "                5,\n",
        "                1.0,\n",
        "                1.5,\n",
        "                0.9,\n",
        "                \"Beam search\",\n",
        "                1,\n",
        "            ],\n",
        "            [\n",
        "                \"https://images.unsplash.com/photo-1589467397966-5e600cb93d3d?auto=format&fit=crop&q=60&w=900&ixlib=rb-4.0.3&ixid=M3wxMjA3fDB8MHxzZWFyY2h8MjB8fEphcGFuJTIwc3RyZWV0fGVufDB8fDB8fHww\",\n",
        "                \"vqa\",\n",
        "                \"道路に書かれた速度制限は？\",\n",
        "                1,\n",
        "                32,\n",
        "                5,\n",
        "                1.0,\n",
        "                1.5,\n",
        "                0.9,\n",
        "                \"Beam search\",\n",
        "                1,\n",
        "            ],\n",
        "        ],\n",
        "        inputs=inputs,\n",
        "    )\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    demo.launch(share=True, debug=True, show_error=True)\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "V100",
      "private_outputs": true,
      "provenance": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}