huggingface_finetune_clip_runner.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "88820440",
   "metadata": {},
   "source": [
    "# CLIP training using HuggingFace libs\n",
    "\n",
    "This notebook demonstrates how to finetune the CLIP model that was used for Stable Diffusion v1.4-1.5 (`clip-vit-large-patch14-336` aka `ViT-L/14@336px`) using a local dataset stored as `.jpg/.jpeg` and `.txt` file pairs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84dac1a7",
   "metadata": {},
   "source": [
    "## Setup\n",
    "Run these first. Assumes that PyTorch is already installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f1a7e4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q datasets pillow\n",
    "# we need v4.26 of transformers - as of writing pip only provides up to v4.25\n",
    "!pip install -q git+https://github.com/huggingface/transformers\n",
    "print(\"--\\nDONE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b6b1671",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import os\n",
    "import pathlib\n",
    "from typing import Generator\n",
    "\n",
    "def collect_captioned_images(root_folder: str) -> Generator[tuple[str,str], None, None]:\n",
    "    image_paths = []\n",
    "    captions = []\n",
    "    \n",
    "    for directory, _, filenames in os.walk(root_folder):\n",
    "        image_extensions = ['.jpg', '.jpeg']\n",
    "        image_filenames = [f for f in filenames if os.path.splitext(f)[1] in image_extensions]\n",
    "        for image_filename in image_filenames:\n",
    "            caption_filename = os.path.splitext(image_filename)[0] + '.txt'\n",
    "            caption_path = os.path.join(directory, caption_filename)\n",
    "            if not os.path.exists(caption_path):\n",
    "                continue\n",
    "\n",
    "            with open(caption_path, 'r') as f:\n",
    "                caption = f.read().replace('\\n', ' ')\n",
    "\n",
    "                image_path = os.path.join(directory, image_filename)\n",
    "                yield image_path, caption\n",
    "\n",
    "                \n",
    "def convert_text_image_pairs_to_huggingface_json(root_folder, out_json):\n",
    "    out_folder = os.path.dirname(root_folder)\n",
    "    pathlib.Path(out_folder).mkdir(parents=True, exist_ok=True)\n",
    "    with open(out_json, \"w\") as f:\n",
    "        written_count = 0\n",
    "        for image_path, caption in collect_captioned_images(root_folder):\n",
    "            line_dict = {\"image\":image_path, \"caption\":caption}\n",
    "            json_line = json.dumps(line_dict, indent=None, separators=(\",\",\":\"))\n",
    "            #print(json_line)\n",
    "            f.write(json_line + \"\\n\")\n",
    "            written_count += 1\n",
    "        print(f\"wrote {written_count} lines to {out_json}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5e24a69",
   "metadata": {},
   "source": [
    "## Convert the data folder of text/image pairs to a huggingface dataset-compatible json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd589bc7",
   "metadata": {},
   "source": [
    "Replace `root_folder` in the next cell with the top-level folder containing your images, and `out_json` with a path to where the json file representing the image/caption pairs in that folder should be saved.\n",
    "\n",
    "Note this only works with pairs of the form `filename.jpg`/`filename.txt` or `filename.jpeg`/`filename.txt`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ab7d69f",
   "metadata": {},
   "outputs": [],
   "source": [
    "root_folder = \"/Users/damian/2.current/stablediffusion/buzzybee/fullsize\"\n",
    "out_json = \"/Users/damian/2.current/stablediffusion/buzzybee.json\"\n",
    "convert_text_image_pairs_to_huggingface_json(root_folder, out_json)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec0e2ac6",
   "metadata": {},
   "source": [
    "Test that it worked by running the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3e42ee0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test loading it back in\n",
    "from datasets import load_dataset\n",
    "dataset = load_dataset(\"json\", data_files=out_json)\n",
    "print(f\"first image: {dataset['train'][0]['image']}, caption: '{dataset['train'][0]['caption']}'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19184e7c",
   "metadata": {},
   "source": [
    "## Run the finetuning\n",
    "\n",
    "### Configuration\n",
    "\n",
    "`repo_id` - The starting point for finetuning. By default this uses the `openai/clip-vit-large-patch14-336` pre-trained CLIP weights. This is what Stable Diffusion versions up to 1.5 used. Another option you might want to consider is `laion/CLIP-ViT-H-14-laion2B-s32B-b79K`, which was used for Stable Diffusion 2.0 onwards.\n",
    "\n",
    "`output_folder` - Where to store the output. The saving process writes multiple files to this folder, so it should be empty.\n",
    "\n",
    "`batch_size` - Training batch size. Don't go lower than 8 - try 32 or 64 (unless you only have a few images).\n",
    "\n",
    "`num_train_epochs` - How many epochs to train. With <500 images each epoch on a 3090 takes a few minutes - do a small number, say `3` to start with, and check the loss when it's done before increasing the number of epochs. With 3 epochs my loss went down to around 2. After 10 epochs it was down to 0.63. Be careful not to over-fit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbb7b053",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "repo_id =  \"openai/clip-vit-large-patch14-336\"\n",
    "output_folder = \"/workspace/output/clip-finetuned\"\n",
    "batch_size = 8\n",
    "num_train_epochs = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f598c759",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(f\"Finetuning {repo_id} for {num_train_epochs} epochs with batch size {batch_size}, and then saving output to {output_folder}.\")\n",
    "!python huggingface_finetune_clip.py \\\n",
    "    --output_dir {output_folder} \\\n",
    "    --model_name_or_path {repo_id} \\\n",
    "    --train_file {out_json} \\\n",
    "    --image_column image \\\n",
    "    --overwrite_output_dir=True \\\n",
    "    --max_seq_length=77 \\ \n",
    "    --num_train_epochs={num_train_epochs} \\\n",
    "    --caption_column caption \\\n",
    "    --remove_unused_columns=False \\\n",
    "    --do_train \\\n",
    "    --per_device_train_batch_size={batch_size} \\\n",
    "    --learning_rate=\"5e-5\" --warmup_steps=\"0\" --weight_decay 0.1 \n",
    "print(\"--\\nDONE\")\n",
    "print(f\"If it worked, trained data should be in {output_folder}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcb56b3c",
   "metadata": {},
   "source": [
    "If it all worked, your finetuned CLIP model is in the `output_folder` defined above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "378ca0fa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}