flux-finetune.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3ee6dc33-60a0-4b86-b641-88d54a913cc4",
   "metadata": {},
   "source": [
    "## Setup the Workspace\n",
    "To get started, we need to install the required packages for the workspace. Run the code in the following cell to install everything, and then the second code cell to login with your HuggingFace credentials. We need to do this since HuggingFace hosts all the used models. \n",
    "\n",
    "### Copy and paste the next code cell into a terminal window - does not work with line magic!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b28616e6-63b6-4826-9f63-999bf55dcb57",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-08-21T21:18:00.107478Z",
     "iopub.status.busy": "2024-08-21T21:18:00.107226Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/notebooks/ai-toolkit\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.11/dist-packages/IPython/core/magics/osm.py:417: UserWarning: using dhist requires you to install the `pickleshare` library.\n",
      "  self.shell.db['dhist'] = compress_dhist(dhist)[-100:]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting git+https://github.com/huggingface/diffusers.git (from -r requirements.txt (line 4))\n",
      "  Cloning https://github.com/huggingface/diffusers.git to /tmp/pip-req-build-fy4h4ra9\n",
      "  Running command git clone --filter=blob:none --quiet https://github.com/huggingface/diffusers.git /tmp/pip-req-build-fy4h4ra9\n",
      "  Resolved https://github.com/huggingface/diffusers.git to commit c2916175186e2b6d9c2d09b13a753cc47f5d9e19\n",
      "  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h  Preparing metadata (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: torch in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 1)) (2.1.1+cu121)\n",
      "Requirement already satisfied: torchvision in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 2)) (0.16.1+cu121)\n",
      "Requirement already satisfied: safetensors in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 3)) (0.4.0)\n",
      "Requirement already satisfied: transformers in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 5)) (4.35.2)\n",
      "Collecting lycoris-lora==1.8.3 (from -r requirements.txt (line 6))\n",
      "  Downloading lycoris_lora-1.8.3.tar.gz (96 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m96.5/96.5 kB\u001b[0m \u001b[31m3.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hCollecting flatten_json (from -r requirements.txt (line 7))\n",
      "  Downloading flatten_json-0.1.14-py3-none-any.whl.metadata (4.2 kB)\n",
      "Requirement already satisfied: pyyaml in /usr/lib/python3/dist-packages (from -r requirements.txt (line 8)) (5.4.1)\n",
      "Collecting oyaml (from -r requirements.txt (line 9))\n",
      "  Downloading oyaml-1.0-py2.py3-none-any.whl.metadata (1.2 kB)\n",
      "Requirement already satisfied: tensorboard in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 10)) (2.15.1)\n",
      "Collecting kornia (from -r requirements.txt (line 11))\n",
      "  Downloading kornia-0.7.3-py2.py3-none-any.whl.metadata (7.7 kB)\n",
      "Collecting invisible-watermark (from -r requirements.txt (line 12))\n",
      "  Downloading invisible_watermark-0.2.0-py3-none-any.whl.metadata (8.2 kB)\n",
      "Collecting einops (from -r requirements.txt (line 13))\n",
      "  Downloading einops-0.8.0-py3-none-any.whl.metadata (12 kB)\n",
      "Requirement already satisfied: accelerate in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 14)) (0.24.1)\n",
      "Collecting toml (from -r requirements.txt (line 15))\n",
      "  Downloading toml-0.10.2-py2.py3-none-any.whl.metadata (7.1 kB)\n",
      "Collecting albumentations (from -r requirements.txt (line 16))\n",
      "  Downloading albumentations-1.4.14-py3-none-any.whl.metadata (38 kB)\n",
      "Requirement already satisfied: pydantic in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 17)) (1.10.14)\n",
      "Requirement already satisfied: omegaconf in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 18)) (2.3.0)\n",
      "Collecting k-diffusion (from -r requirements.txt (line 19))\n",
      "  Downloading k_diffusion-0.1.1.post1-py3-none-any.whl.metadata (3.9 kB)\n",
      "Collecting open_clip_torch (from -r requirements.txt (line 20))\n",
      "  Downloading open_clip_torch-2.26.1-py3-none-any.whl.metadata (31 kB)\n",
      "Requirement already satisfied: timm in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 21)) (0.9.7)\n",
      "Collecting prodigyopt (from -r requirements.txt (line 22))\n",
      "  Downloading prodigyopt-1.0-py3-none-any.whl.metadata (1.2 kB)\n",
      "Collecting controlnet_aux==0.0.7 (from -r requirements.txt (line 23))\n",
      "  Downloading controlnet_aux-0.0.7.tar.gz (202 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m202.4/202.4 kB\u001b[0m \u001b[31m9.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hCollecting python-dotenv (from -r requirements.txt (line 24))\n",
      "  Downloading python_dotenv-1.0.1-py3-none-any.whl.metadata (23 kB)\n",
      "Requirement already satisfied: bitsandbytes in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 25)) (0.41.2)\n",
      "Collecting hf_transfer (from -r requirements.txt (line 26))\n",
      "  Downloading hf_transfer-0.1.8-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (1.7 kB)\n",
      "Collecting lpips (from -r requirements.txt (line 27))\n",
      "  Downloading lpips-0.1.4-py3-none-any.whl.metadata (10 kB)\n",
      "Collecting pytorch_fid (from -r requirements.txt (line 28))\n",
      "  Downloading pytorch_fid-0.3.0-py3-none-any.whl.metadata (5.3 kB)\n",
      "Collecting optimum-quanto (from -r requirements.txt (line 29))\n",
      "  Downloading optimum_quanto-0.2.4-py3-none-any.whl.metadata (13 kB)\n",
      "Requirement already satisfied: sentencepiece in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 30)) (0.1.99)\n",
      "Requirement already satisfied: huggingface_hub in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 31)) (0.20.3)\n",
      "Requirement already satisfied: importlib_metadata in /usr/lib/python3/dist-packages (from controlnet_aux==0.0.7->-r requirements.txt (line 23)) (4.6.4)\n",
      "Requirement already satisfied: scipy in /usr/local/lib/python3.11/dist-packages (from controlnet_aux==0.0.7->-r requirements.txt (line 23)) (1.11.2)\n",
      "Requirement already satisfied: opencv-python in /usr/local/lib/python3.11/dist-packages (from controlnet_aux==0.0.7->-r requirements.txt (line 23)) (4.8.0.76)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from controlnet_aux==0.0.7->-r requirements.txt (line 23)) (3.13.1)\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.11/dist-packages (from controlnet_aux==0.0.7->-r requirements.txt (line 23)) (1.26.3)\n",
      "Requirement already satisfied: Pillow in /usr/local/lib/python3.11/dist-packages (from controlnet_aux==0.0.7->-r requirements.txt (line 23)) (9.5.0)\n",
      "Requirement already satisfied: scikit-image in /usr/local/lib/python3.11/dist-packages (from controlnet_aux==0.0.7->-r requirements.txt (line 23)) (0.21.0)\n",
      "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.11/dist-packages (from torch->-r requirements.txt (line 1)) (4.9.0)\n",
      "Requirement already satisfied: sympy in /usr/local/lib/python3.11/dist-packages (from torch->-r requirements.txt (line 1)) (1.12)\n",
      "Requirement already satisfied: networkx in /usr/local/lib/python3.11/dist-packages (from torch->-r requirements.txt (line 1)) (3.2.1)\n",
      "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from torch->-r requirements.txt (line 1)) (3.1.3)\n",
      "Requirement already satisfied: fsspec in /usr/local/lib/python3.11/dist-packages (from torch->-r requirements.txt (line 1)) (2023.6.0)\n",
      "Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.11/dist-packages (from torch->-r requirements.txt (line 1)) (2.1.0)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.11/dist-packages (from torchvision->-r requirements.txt (line 2)) (2.31.0)\n",
      "Collecting huggingface_hub (from -r requirements.txt (line 31))\n",
      "  Downloading huggingface_hub-0.24.6-py3-none-any.whl.metadata (13 kB)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.11/dist-packages (from diffusers==0.31.0.dev0->-r requirements.txt (line 4)) (2023.12.25)\n",
      "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.11/dist-packages (from transformers->-r requirements.txt (line 5)) (23.2)\n",
      "Requirement already satisfied: tokenizers<0.19,>=0.14 in /usr/local/lib/python3.11/dist-packages (from transformers->-r requirements.txt (line 5)) (0.15.1)\n",
      "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.11/dist-packages (from transformers->-r requirements.txt (line 5)) (4.66.1)\n",
      "Requirement already satisfied: six in /usr/lib/python3/dist-packages (from flatten_json->-r requirements.txt (line 7)) (1.16.0)\n",
      "Requirement already satisfied: absl-py>=0.4 in /usr/local/lib/python3.11/dist-packages (from tensorboard->-r requirements.txt (line 10)) (2.1.0)\n",
      "Requirement already satisfied: grpcio>=1.48.2 in /usr/local/lib/python3.11/dist-packages (from tensorboard->-r requirements.txt (line 10)) (1.60.0)\n",
      "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.11/dist-packages (from tensorboard->-r requirements.txt (line 10)) (2.26.2)\n",
      "Requirement already satisfied: google-auth-oauthlib<2,>=0.5 in /usr/local/lib/python3.11/dist-packages (from tensorboard->-r requirements.txt (line 10)) (1.2.0)\n",
      "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.11/dist-packages (from tensorboard->-r requirements.txt (line 10)) (3.5.2)\n",
      "Requirement already satisfied: protobuf<4.24,>=3.19.6 in /usr/local/lib/python3.11/dist-packages (from tensorboard->-r requirements.txt (line 10)) (4.23.4)\n",
      "Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.11/dist-packages (from tensorboard->-r requirements.txt (line 10)) (69.0.3)\n",
      "Requirement already satisfied: tensorboard-data-server<0.8.0,>=0.7.0 in /usr/local/lib/python3.11/dist-packages (from tensorboard->-r requirements.txt (line 10)) (0.7.2)\n",
      "Requirement already satisfied: werkzeug>=1.0.1 in /usr/local/lib/python3.11/dist-packages (from tensorboard->-r requirements.txt (line 10)) (3.0.1)\n",
      "Collecting kornia-rs>=0.1.0 (from kornia->-r requirements.txt (line 11))\n",
      "  Downloading kornia_rs-0.1.5-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (8.7 kB)\n",
      "Requirement already satisfied: PyWavelets>=1.1.1 in /usr/local/lib/python3.11/dist-packages (from invisible-watermark->-r requirements.txt (line 12)) (1.5.0)\n",
      "Requirement already satisfied: psutil in /usr/local/lib/python3.11/dist-packages (from accelerate->-r requirements.txt (line 14)) (5.9.8)\n",
      "Collecting pydantic (from -r requirements.txt (line 17))\n",
      "  Downloading pydantic-2.8.2-py3-none-any.whl.metadata (125 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m125.2/125.2 kB\u001b[0m \u001b[31m9.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting albucore>=0.0.13 (from albumentations->-r requirements.txt (line 16))\n",
      "  Downloading albucore-0.0.13-py3-none-any.whl.metadata (3.1 kB)\n",
      "Collecting eval-type-backport (from albumentations->-r requirements.txt (line 16))\n",
      "  Downloading eval_type_backport-0.2.0-py3-none-any.whl.metadata (2.2 kB)\n",
      "Collecting opencv-python-headless>=4.9.0.80 (from albumentations->-r requirements.txt (line 16))\n",
      "  Downloading opencv_python_headless-4.10.0.84-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (20 kB)\n",
      "Collecting annotated-types>=0.4.0 (from pydantic->-r requirements.txt (line 17))\n",
      "  Downloading annotated_types-0.7.0-py3-none-any.whl.metadata (15 kB)\n",
      "Collecting pydantic-core==2.20.1 (from pydantic->-r requirements.txt (line 17))\n",
      "  Downloading pydantic_core-2.20.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (6.6 kB)\n",
      "Requirement already satisfied: antlr4-python3-runtime==4.9.* in /usr/local/lib/python3.11/dist-packages (from omegaconf->-r requirements.txt (line 18)) (4.9.3)\n",
      "Collecting clean-fid (from k-diffusion->-r requirements.txt (line 19))\n",
      "  Downloading clean_fid-0.1.35-py3-none-any.whl.metadata (36 kB)\n",
      "Collecting clip-anytorch (from k-diffusion->-r requirements.txt (line 19))\n",
      "  Downloading clip_anytorch-2.6.0-py3-none-any.whl.metadata (8.4 kB)\n",
      "Collecting dctorch (from k-diffusion->-r requirements.txt (line 19))\n",
      "  Downloading dctorch-0.1.2-py3-none-any.whl.metadata (607 bytes)\n",
      "Collecting jsonmerge (from k-diffusion->-r requirements.txt (line 19))\n",
      "  Downloading jsonmerge-1.9.2-py3-none-any.whl.metadata (21 kB)\n",
      "Collecting torchdiffeq (from k-diffusion->-r requirements.txt (line 19))\n",
      "  Downloading torchdiffeq-0.2.4-py3-none-any.whl.metadata (440 bytes)\n",
      "Collecting torchsde (from k-diffusion->-r requirements.txt (line 19))\n",
      "  Downloading torchsde-0.2.6-py3-none-any.whl.metadata (5.3 kB)\n",
      "Requirement already satisfied: wandb in /usr/local/lib/python3.11/dist-packages (from k-diffusion->-r requirements.txt (line 19)) (0.15.10)\n",
      "Collecting ftfy (from open_clip_torch->-r requirements.txt (line 20))\n",
      "  Downloading ftfy-6.2.3-py3-none-any.whl.metadata (7.8 kB)\n",
      "INFO: pip is looking at multiple versions of optimum-quanto to determine which version is compatible with other requirements. This could take a while.\n",
      "Collecting optimum-quanto (from -r requirements.txt (line 29))\n",
      "  Downloading optimum_quanto-0.2.3-py3-none-any.whl.metadata (13 kB)\n",
      "  Downloading optimum_quanto-0.2.2-py3-none-any.whl.metadata (10 kB)\n",
      "  Downloading optimum_quanto-0.2.1-py3-none-any.whl.metadata (10 kB)\n",
      "Collecting pytorch_fid (from -r requirements.txt (line 28))\n",
      "  Downloading pytorch-fid-0.2.1.tar.gz (14 kB)\n",
      "  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hINFO: pip is still looking at multiple versions of optimum-quanto to determine which version is compatible with other requirements. This could take a while.\n",
      "  Downloading pytorch-fid-0.2.0.tar.gz (11 kB)\n",
      "  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Downloading pytorch-fid-0.1.1.tar.gz (9.3 kB)\n",
      "  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hINFO: This is taking longer than usual. You might need to provide the dependency resolver with stricter constraints to reduce runtime. See https://pip.pypa.io/warnings/backtracking for guidance. If you want to abort this run, press Ctrl + C.\n",
      "Collecting lpips (from -r requirements.txt (line 27))\n",
      "  Downloading lpips-0.1.3-py3-none-any.whl.metadata (10 kB)\n",
      "Collecting timm (from -r requirements.txt (line 21))\n",
      "  Downloading timm-1.0.8-py3-none-any.whl.metadata (53 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m53.8/53.8 kB\u001b[0m \u001b[31m8.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-1.0.7-py3-none-any.whl.metadata (47 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m47.5/47.5 kB\u001b[0m \u001b[31m6.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-1.0.3-py3-none-any.whl.metadata (43 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m43.6/43.6 kB\u001b[0m \u001b[31m9.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-0.9.16-py3-none-any.whl.metadata (38 kB)\n",
      "  Downloading timm-0.9.12-py3-none-any.whl.metadata (60 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m60.6/60.6 kB\u001b[0m \u001b[31m10.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-0.9.11-py3-none-any.whl.metadata (60 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m60.5/60.5 kB\u001b[0m \u001b[31m11.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-0.9.10-py3-none-any.whl.metadata (59 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m59.8/59.8 kB\u001b[0m \u001b[31m10.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-0.9.9-py3-none-any.whl.metadata (59 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m59.8/59.8 kB\u001b[0m \u001b[31m9.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-0.9.8-py3-none-any.whl.metadata (59 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m59.3/59.3 kB\u001b[0m \u001b[31m11.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-0.9.7-py3-none-any.whl.metadata (58 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.8/58.8 kB\u001b[0m \u001b[31m7.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-0.9.6-py3-none-any.whl.metadata (58 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.6/58.6 kB\u001b[0m \u001b[31m10.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-0.9.5-py3-none-any.whl.metadata (69 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m69.4/69.4 kB\u001b[0m \u001b[31m11.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-0.9.2-py3-none-any.whl.metadata (68 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m68.5/68.5 kB\u001b[0m \u001b[31m10.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-0.9.1-py3-none-any.whl.metadata (68 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m68.5/68.5 kB\u001b[0m \u001b[31m11.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-0.9.0-py3-none-any.whl.metadata (68 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m68.5/68.5 kB\u001b[0m \u001b[31m10.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-0.6.13-py3-none-any.whl.metadata (38 kB)\n",
      "  Downloading timm-0.6.12-py3-none-any.whl.metadata (37 kB)\n",
      "  Downloading timm-0.6.11-py3-none-any.whl.metadata (36 kB)\n",
      "  Downloading timm-0.6.7-py3-none-any.whl.metadata (33 kB)\n",
      "  Downloading timm-0.6.5-py3-none-any.whl.metadata (45 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m46.0/46.0 kB\u001b[0m \u001b[31m9.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading timm-0.5.4-py3-none-any.whl.metadata (36 kB)\n",
      "  Downloading timm-0.4.12-py3-none-any.whl.metadata (30 kB)\n",
      "  Downloading timm-0.4.9-py3-none-any.whl.metadata (27 kB)\n",
      "  Downloading timm-0.4.5-py3-none-any.whl.metadata (24 kB)\n",
      "  Downloading timm-0.3.4-py3-none-any.whl.metadata (20 kB)\n",
      "  Downloading timm-0.3.3-py3-none-any.whl.metadata (20 kB)\n",
      "  Downloading timm-0.3.2-py3-none-any.whl.metadata (19 kB)\n",
      "  Downloading timm-0.3.1-py3-none-any.whl.metadata (19 kB)\n",
      "  Downloading timm-0.3.0-py3-none-any.whl.metadata (19 kB)\n",
      "  Downloading timm-0.2.1-py3-none-any.whl.metadata (14 kB)\n",
      "  Downloading timm-0.1.30-py3-none-any.whl.metadata (37 kB)\n",
      "  Downloading timm-0.1.28-py3-none-any.whl.metadata (37 kB)\n",
      "  Downloading timm-0.1.26-py3-none-any.whl.metadata (39 kB)\n",
      "  Downloading timm-0.1.24-py3-none-any.whl.metadata (38 kB)\n",
      "  Downloading timm-0.1.22-py3-none-any.whl.metadata (38 kB)\n",
      "  Downloading timm-0.1.20-py3-none-any.whl.metadata (38 kB)\n",
      "  Downloading timm-0.1.18-py3-none-any.whl.metadata (34 kB)\n",
      "  Downloading timm-0.1.16-py3-none-any.whl.metadata (30 kB)\n",
      "  Downloading timm-0.1.14-py3-none-any.whl.metadata (20 kB)\n",
      "  Downloading timm-0.1.12-py3-none-any.whl.metadata (18 kB)\n",
      "  Downloading timm-0.1.10-py3-none-any.whl.metadata (18 kB)\n",
      "  Downloading timm-0.1.8-py3-none-any.whl.metadata (17 kB)\n",
      "  Downloading timm-0.1.6-py3-none-any.whl.metadata (16 kB)\n",
      "  Downloading timm-0.1.4-py3-none-any.whl.metadata (16 kB)\n",
      "  Downloading timm-0.1.2-py3-none-any.whl.metadata (16 kB)\n",
      "  Downloading timm-0.1.1-py3-none-any.whl.metadata (16 kB)\n",
      "Collecting open_clip_torch (from -r requirements.txt (line 20))\n",
      "  Downloading open_clip_torch-2.24.0-py3-none-any.whl.metadata (30 kB)\n",
      "  Downloading open_clip_torch-2.23.0-py3-none-any.whl.metadata (30 kB)\n",
      "  Downloading open_clip_torch-2.22.0-py3-none-any.whl.metadata (29 kB)\n",
      "Collecting protobuf<4.24,>=3.19.6 (from tensorboard->-r requirements.txt (line 10))\n",
      "  Downloading protobuf-3.20.3-py2.py3-none-any.whl.metadata (720 bytes)\n",
      "Collecting open_clip_torch (from -r requirements.txt (line 20))\n",
      "  Downloading open_clip_torch-2.20.0-py3-none-any.whl.metadata (46 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m46.8/46.8 kB\u001b[0m \u001b[31m9.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading open_clip_torch-2.19.0-py3-none-any.whl.metadata (46 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m46.8/46.8 kB\u001b[0m \u001b[31m10.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading open_clip_torch-2.18.0-py3-none-any.whl.metadata (44 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.1/44.1 kB\u001b[0m \u001b[31m10.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading open_clip_torch-2.17.2-py3-none-any.whl.metadata (43 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m43.9/43.9 kB\u001b[0m \u001b[31m11.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading open_clip_torch-2.17.1-py3-none-any.whl.metadata (43 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m43.9/43.9 kB\u001b[0m \u001b[31m9.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading open_clip_torch-2.16.2-py3-none-any.whl.metadata (42 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.8/42.8 kB\u001b[0m \u001b[31m9.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading open_clip_torch-2.16.1-py3-none-any.whl.metadata (42 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.8/42.8 kB\u001b[0m \u001b[31m8.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading open_clip_torch-2.16.0-py3-none-any.whl.metadata (42 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.8/42.8 kB\u001b[0m \u001b[31m9.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading open_clip_torch-2.15.0-py3-none-any.whl.metadata (42 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.7/42.7 kB\u001b[0m \u001b[31m7.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading open_clip_torch-2.14.0-py3-none-any.whl.metadata (42 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.6/42.6 kB\u001b[0m \u001b[31m8.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading open_clip_torch-2.13.0-py3-none-any.whl.metadata (42 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.2/42.2 kB\u001b[0m \u001b[31m8.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Downloading open_clip_torch-2.12.0-py3-none-any.whl.metadata (41 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m41.2/41.2 kB\u001b[0m \u001b[31m6.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h"
     ]
    }
   ],
   "source": [
    "git clone https://github.com/ostris/ai-toolkit.git\n",
    "cd ai-toolkit\n",
    "git submodule update --init --recursive\n",
    "python3 -m venv venv\n",
    "source venv/bin/activate\n",
    "pip3 install -r requirements.txt\n",
    "pip install peft\n",
    "pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118\n",
    "huggingface-cli login"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6614b901",
   "metadata": {},
   "source": [
    "## Label your images\n",
    "\n",
    "Now that the environment is setup, we need to caption our images. We have provided a quick loop to caption each image you upload.\n",
    "\n",
    "To ensure this works properly, upload your images into a single directory. Paste the name in for the variable `your dir` in the second code cell below, run the cells, and the script will automatically generate the labeled text files. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1612c3f9-902b-43fd-bcc2-942a20777369",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-08-21T21:40:28.311150Z",
     "iopub.status.busy": "2024-08-21T21:40:28.310644Z",
     "iopub.status.idle": "2024-08-21T21:40:39.592701Z",
     "shell.execute_reply": "2024-08-21T21:40:39.591997Z",
     "shell.execute_reply.started": "2024-08-21T21:40:28.311130Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: einops in /usr/local/lib/python3.11/dist-packages (0.8.0)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mRequirement already satisfied: transformers in /usr/local/lib/python3.11/dist-packages (4.35.2)\n",
      "Collecting transformers\n",
      "  Downloading transformers-4.44.1-py3-none-any.whl.metadata (43 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m43.7/43.7 kB\u001b[0m \u001b[31m2.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from transformers) (3.13.1)\n",
      "Collecting huggingface-hub<1.0,>=0.23.2 (from transformers)\n",
      "  Downloading huggingface_hub-0.24.6-py3-none-any.whl.metadata (13 kB)\n",
      "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.11/dist-packages (from transformers) (1.26.3)\n",
      "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.11/dist-packages (from transformers) (23.2)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /usr/lib/python3/dist-packages (from transformers) (5.4.1)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.11/dist-packages (from transformers) (2023.12.25)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.11/dist-packages (from transformers) (2.31.0)\n",
      "Collecting safetensors>=0.4.1 (from transformers)\n",
      "  Downloading safetensors-0.4.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (3.8 kB)\n",
      "Collecting tokenizers<0.20,>=0.19 (from transformers)\n",
      "  Downloading tokenizers-0.19.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (6.7 kB)\n",
      "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.11/dist-packages (from transformers) (4.66.1)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.11/dist-packages (from huggingface-hub<1.0,>=0.23.2->transformers) (2023.6.0)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.11/dist-packages (from huggingface-hub<1.0,>=0.23.2->transformers) (4.9.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests->transformers) (3.3)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (2.0.7)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests->transformers) (2020.6.20)\n",
      "Downloading transformers-4.44.1-py3-none-any.whl (9.5 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m9.5/9.5 MB\u001b[0m \u001b[31m58.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hDownloading huggingface_hub-0.24.6-py3-none-any.whl (417 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m417.5/417.5 kB\u001b[0m \u001b[31m80.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading safetensors-0.4.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (435 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m435.4/435.4 kB\u001b[0m \u001b[31m41.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading tokenizers-0.19.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.6 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.6/3.6 MB\u001b[0m \u001b[31m119.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: safetensors, huggingface-hub, tokenizers, transformers\n",
      "  Attempting uninstall: safetensors\n",
      "    Found existing installation: safetensors 0.4.0\n",
      "    Uninstalling safetensors-0.4.0:\n",
      "      Successfully uninstalled safetensors-0.4.0\n",
      "  Attempting uninstall: huggingface-hub\n",
      "    Found existing installation: huggingface-hub 0.20.3\n",
      "    Uninstalling huggingface-hub-0.20.3:\n",
      "      Successfully uninstalled huggingface-hub-0.20.3\n",
      "  Attempting uninstall: tokenizers\n",
      "    Found existing installation: tokenizers 0.15.1\n",
      "    Uninstalling tokenizers-0.15.1:\n",
      "      Successfully uninstalled tokenizers-0.15.1\n",
      "  Attempting uninstall: transformers\n",
      "    Found existing installation: transformers 4.35.2\n",
      "    Uninstalling transformers-4.35.2:\n",
      "      Successfully uninstalled transformers-4.35.2\n",
      "Successfully installed huggingface-hub-0.24.6 safetensors-0.4.4 tokenizers-0.19.1 transformers-4.44.1\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install einops\n",
    "!pip install -U transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0c72e61a-c28b-4557-b1a7-25da06618fa1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-08-21T21:40:52.626680Z",
     "iopub.status.busy": "2024-08-21T21:40:52.626119Z",
     "iopub.status.idle": "2024-08-21T21:40:54.079886Z",
     "shell.execute_reply": "2024-08-21T21:40:54.079257Z",
     "shell.execute_reply.started": "2024-08-21T21:40:52.626658Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from unittest.mock import patch\n",
    "from transformers.dynamic_module_utils import get_imports\n",
    "import os\n",
    "def fixed_get_imports(filename: str | os.PathLike) -> list[str]:\n",
    "    if not str(filename).endswith(\"modeling_florence2.py\"):\n",
    "        return get_imports(filename)\n",
    "    imports = get_imports(filename)\n",
    "    imports.remove(\"flash_attn\")\n",
    "    return imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9b668d04-cfd6-4d02-aad3-ca7135aaa7d8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-08-21T21:40:54.610485Z",
     "iopub.status.busy": "2024-08-21T21:40:54.609940Z",
     "iopub.status.idle": "2024-08-21T21:41:33.509539Z",
     "shell.execute_reply": "2024-08-21T21:41:33.508768Z",
     "shell.execute_reply.started": "2024-08-21T21:40:54.610464Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-08-21 21:40:55.361167: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2024-08-21 21:40:55.361229: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2024-08-21 21:40:55.362432: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2024-08-21 21:40:55.369643: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-08-21 21:40:56.256670: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "362f21f6b08d428f9fafc215e0911975",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/1.54G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.11/dist-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return self.fget.__get__(instance, owner)()\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a14ef7ba282f4811bb804f5c6655aaf4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "generation_config.json:   0%|          | 0.00/51.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "28baff1d72fc47e997cf6aac6f417c0a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "preprocessor_config.json:   0%|          | 0.00/806 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4a34e0727c514a0f988aac4c6ee5a142",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "processing_florence2.py:   0%|          | 0.00/46.4k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "A new version of the following files was downloaded from https://huggingface.co/microsoft/Florence-2-large:\n",
      "- processing_florence2.py\n",
      ". Make sure to double-check they do not contain any added malicious code. To avoid downloading new versions of the code file, you can pin a revision.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5c4521c7da814434a9321aea220146b2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/34.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac5720a2a18248f39d704a024b9ded87",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "vocab.json:   0%|          | 0.00/1.10M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "78752a3941f04e919d9f398f6b30a248",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.11/dist-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "ename": "UnidentifiedImageError",
     "evalue": "cannot identify image file 'fold/.aitk_size.json'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mUnidentifiedImageError\u001b[0m                    Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[2], line 25\u001b[0m\n\u001b[1;32m     23\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m i\u001b[38;5;241m.\u001b[39msplit(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m'\u001b[39m)[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]\u001b[38;5;241m==\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtxt\u001b[39m\u001b[38;5;124m'\u001b[39m:\n\u001b[1;32m     24\u001b[0m     \u001b[38;5;28;01mcontinue\u001b[39;00m\n\u001b[0;32m---> 25\u001b[0m image \u001b[38;5;241m=\u001b[39m \u001b[43mImage\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mopen\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43mf\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;132;43;01m{\u001b[39;49;00m\u001b[43myour_dir\u001b[49m\u001b[38;5;132;43;01m}\u001b[39;49;00m\u001b[38;5;124;43m/\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[43mi\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     27\u001b[0m inputs \u001b[38;5;241m=\u001b[39m processor(text\u001b[38;5;241m=\u001b[39mprompt, images\u001b[38;5;241m=\u001b[39mimage, return_tensors\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpt\u001b[39m\u001b[38;5;124m\"\u001b[39m)\u001b[38;5;241m.\u001b[39mto(device, torch_dtype)\n\u001b[1;32m     29\u001b[0m generated_ids \u001b[38;5;241m=\u001b[39m model\u001b[38;5;241m.\u001b[39mgenerate(\n\u001b[1;32m     30\u001b[0m   input_ids\u001b[38;5;241m=\u001b[39minputs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124minput_ids\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[1;32m     31\u001b[0m   pixel_values\u001b[38;5;241m=\u001b[39minputs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpixel_values\u001b[39m\u001b[38;5;124m\"\u001b[39m],\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     34\u001b[0m   do_sample\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m     35\u001b[0m )\n",
      "File \u001b[0;32m/usr/local/lib/python3.11/dist-packages/PIL/Image.py:3298\u001b[0m, in \u001b[0;36mopen\u001b[0;34m(fp, mode, formats)\u001b[0m\n\u001b[1;32m   3296\u001b[0m     warnings\u001b[38;5;241m.\u001b[39mwarn(message)\n\u001b[1;32m   3297\u001b[0m msg \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcannot identify image file \u001b[39m\u001b[38;5;132;01m%r\u001b[39;00m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m%\u001b[39m (filename \u001b[38;5;28;01mif\u001b[39;00m filename \u001b[38;5;28;01melse\u001b[39;00m fp)\n\u001b[0;32m-> 3298\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m UnidentifiedImageError(msg)\n",
      "\u001b[0;31mUnidentifiedImageError\u001b[0m: cannot identify image file 'fold/.aitk_size.json'"
     ]
    }
   ],
   "source": [
    "!pip install -U oyaml transformers einops albumentations python-dotenv\n",
    "\n",
    "import requests\n",
    "import torch\n",
    "from PIL import Image\n",
    "from transformers import AutoProcessor, AutoModelForCausalLM \n",
    "import os\n",
    "\n",
    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32\n",
    "\n",
    "\n",
    "\n",
    "model_id = 'microsoft/Florence-2-large'\n",
    "with patch(\"transformers.dynamic_module_utils.get_imports\", fixed_get_imports): #workaround for unnecessary flash_attn requirement\n",
    "            model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation=\"sdpa\", torch_dtype='auto',trust_remote_code=True)\n",
    "processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)\n",
    "model.to(device)\n",
    "\n",
    "your_dir = 'fold'\n",
    "\n",
    "\n",
    "prompt = \"<MORE_DETAILED_CAPTION>\"\n",
    "for i in os.listdir(f'{your_dir}'):\n",
    "    if i.split('.')[-1]=='txt':\n",
    "        continue\n",
    "    image = Image.open(f'{your_dir}/'+i)\n",
    "\n",
    "    inputs = processor(text=prompt, images=image, return_tensors=\"pt\").to(device, torch_dtype)\n",
    "\n",
    "    generated_ids = model.generate(\n",
    "      input_ids=inputs[\"input_ids\"],\n",
    "      pixel_values=inputs[\"pixel_values\"],\n",
    "      max_new_tokens=1024,\n",
    "      num_beams=3,\n",
    "      do_sample=False\n",
    "    )\n",
    "    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]\n",
    "\n",
    "    parsed_answer = processor.post_process_generation(generated_text, task=\"<MORE_DETAILED_CAPTION>\", image_size=(image.width, image.height))\n",
    "    # print(parsed_answer)\n",
    "    with open(f'{your_dir}/'+f\"{i.split('.')[0]}.txt\", \"w\") as f:\n",
    "        f.write(parsed_answer[\"<MORE_DETAILED_CAPTION>\"])\n",
    "        f.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af8d3540",
   "metadata": {},
   "source": [
    "## The Image folder directory \n",
    "\n",
    "The format of your image folder should look like the following now:\n",
    "\n",
    "---|\n",
    "\n",
    "  Your Image Directory\n",
    "   \n",
    "  |\n",
    "\n",
    "------- img1.png\n",
    "\n",
    "------- img1.txt\n",
    "\n",
    "------- img2.png\n",
    "\n",
    "------- img2.txt\n",
    "\n",
    "...\n",
    "\n",
    "Remember The images and text files must follow the same naming convention \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87f6d6a9",
   "metadata": {},
   "source": [
    "## Configure the config yaml file\n",
    "\n",
    "We are going to use the `config/examples/train_lora_flux_24gb.yaml` file to order this model training. \n",
    "\n",
    "The most important lines we are going to edit are going to be found on lines 5 -where we change the name, 30 - where we add the path to our image directory, and 69 and 70 - where we can edit the height and width to reflect our training images. Edit these lines to correspondingly attune the trainer to run on your images. \n",
    "\n",
    "Additionally, we may want to edit the prompts. Several of the prompts refer to animals or scenes, so if we are trying to capture a specific person, we may want to edit these to better inform the model. We can also further control these generated samples using the guidance scale and sample steps values on lines 87-88.\n",
    "\n",
    "We can further optimize training the model by editing the batch size, on line 37, and the gradient accumulation steps, line 39, if we want to more quickly train the FLUX.1 model. If we are training on a multi-GPU or H100, we can raise these values up slightly, but we otherwise recommend they be left the same. Be wary raising them may cause an Out Of Memory error. \n",
    "\n",
    "On line 38, we can change the number of training steps. They recommend between 500 and 4000, so we are going in the middle with 2500. We got good results with this value. It will checkpoint every 250 steps, but we can also change this value on line 22 if needed. \n",
    "\n",
    "Finally, we can change the model from dev to schnell by pasting the HuggingFace id for schnell in on line 62 ('black-forest-labs/FLUX.1-schnell'). Now that everything has been set up, we can run the training!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "37744dd5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-08-21T21:41:39.617226Z",
     "iopub.status.busy": "2024-08-21T21:41:39.616631Z",
     "iopub.status.idle": "2024-08-21T21:41:39.629149Z",
     "shell.execute_reply": "2024-08-21T21:41:39.628596Z",
     "shell.execute_reply.started": "2024-08-21T21:41:39.617205Z"
    }
   },
   "outputs": [],
   "source": [
    "your_dir = '/notebooks/fold'\n",
    "lora_name = 'my_new_lora'\n",
    "batch_size = 1\n",
    "training_steps = 500\n",
    "name_or_path = \"black-forest-labs/FLUX.1-dev\"\n",
    "width = 1024\n",
    "height = 1024\n",
    "sample_steps = 20\n",
    "\n",
    "yaml = f\"\"\"---\n",
    "job: extension\n",
    "config:\n",
    "  # this name will be the folder and filename name\n",
    "  name: {lora_name}\n",
    "  process:\n",
    "    - type: 'sd_trainer'\n",
    "      # root folder to save training sessions/samples/weights\n",
    "      training_folder: \"output\"\n",
    "      # uncomment to see performance stats in the terminal every N steps\n",
    "#      performance_log_every: 1000\n",
    "      device: cuda:0\n",
    "      # if a trigger word is specified, it will be added to captions of training data if it does not already exist\n",
    "      # alternatively, in your captions you can add [trigger] and it will be replaced with the trigger word\n",
    "#      trigger_word: \"p3r5on\"\n",
    "      network:\n",
    "        type: \"lora\"\n",
    "        linear: 16\n",
    "        linear_alpha: 16\n",
    "      save:\n",
    "        dtype: float16 # precision to save\n",
    "        save_every: 250 # save every this many steps\n",
    "        max_step_saves_to_keep: 4 # how many intermittent saves to keep\n",
    "      datasets:\n",
    "        # datasets are a folder of images. captions need to be txt files with the same name as the image\n",
    "        # for instance image2.jpg and image2.txt. Only jpg, jpeg, and png are supported currently\n",
    "        # images will automatically be resized and bucketed into the resolution specified\n",
    "        # on windows, escape back slashes with another backslash so\n",
    "        # \"C:\\\\path\\\\to\\\\images\\\\folder\"\n",
    "        - folder_path: {your_dir}\n",
    "          caption_ext: \"txt\"\n",
    "          caption_dropout_rate: 0.05  # will drop out the caption 5% of time\n",
    "          shuffle_tokens: false  # shuffle caption order, split by commas\n",
    "          cache_latents_to_disk: true  # leave this true unless you know what you're doing\n",
    "          resolution: [1024]  # flux enjoys multiple resolutions\n",
    "      train:\n",
    "        batch_size: {batch_size}\n",
    "        steps: {training_steps}  # total number of steps to train 500 - 4000 is a good range\n",
    "        gradient_accumulation_steps: 1\n",
    "        train_unet: true\n",
    "        train_text_encoder: false  # probably won't work with flux\n",
    "        gradient_checkpointing: true  # need the on unless you have a ton of vram\n",
    "        noise_scheduler: \"flowmatch\" # for training only\n",
    "        optimizer: \"adamw8bit\"\n",
    "        lr: 1e-4\n",
    "        # uncomment this to skip the pre training sample\n",
    "#        skip_first_sample: true\n",
    "        # uncomment to completely disable sampling\n",
    "#        disable_sampling: true\n",
    "        # uncomment to use new vell curved weighting. Experimental but may produce better results\n",
    "        linear_timesteps: true\n",
    "\n",
    "        # ema will smooth out learning, but could slow it down. Recommended to leave on.\n",
    "        ema_config:\n",
    "          use_ema: true\n",
    "          ema_decay: 0.99\n",
    "\n",
    "        # will probably need this if gpu supports it for flux, other dtypes may not work correctly\n",
    "        dtype: bf16\n",
    "      model:\n",
    "        # huggingface model name or path\n",
    "        name_or_path: {name_or_path}\n",
    "        is_flux: true\n",
    "        quantize: true  # run 8bit mixed precision\n",
    "#        low_vram: true  # uncomment this if the GPU is connected to your monitors. It will use less vram to quantize, but is slower.\n",
    "      sample:\n",
    "        sampler: \"flowmatch\" # must match train.noise_scheduler\n",
    "        sample_every: 250 # sample every this many steps\n",
    "        width: {width}\n",
    "        height: {height}\n",
    "        prompts:\n",
    "          # you can add [trigger] to the prompts here and it will be replaced with the trigger word\n",
    "#          - \"[trigger] holding a sign that says 'I LOVE PROMPTS!'\"\\\n",
    "          - \"woman with red hair, playing chess at the park, bomb going off in the background\"\n",
    "          - \"a woman holding a coffee cup, in a beanie, sitting at a cafe\"\n",
    "          - \"a horse is a DJ at a night club, fish eye lens, smoke machine, lazer lights, holding a martini\"\n",
    "          - \"a man showing off his cool new t shirt at the beach, a shark is jumping out of the water in the background\"\n",
    "          - \"a bear building a log cabin in the snow covered mountains\"\n",
    "          - \"woman playing the guitar, on stage, singing a song, laser lights, punk rocker\"\n",
    "          - \"hipster man with a beard, building a chair, in a wood shop\"\n",
    "          - \"photo of a man, white background, medium shot, modeling clothing, studio lighting, white backdrop\"\n",
    "          - \"a man holding a sign that says, 'this is a sign'\"\n",
    "          - \"a bulldog, in a post apocalyptic world, with a shotgun, in a leather jacket, in a desert, with a motorcycle\"\n",
    "        neg: \"\"  # not used on flux\n",
    "        seed: 42\n",
    "        walk_seed: true\n",
    "        guidance_scale: 4\n",
    "        sample_steps: {sample_steps}\n",
    "# you can add any additional meta info here. [name] is replaced with config name at top\n",
    "meta:\n",
    "  name: \"[name]\"\n",
    "  version: '1.0'\n",
    "\"\"\"\n",
    "\n",
    "with open('/notebooks/ai-toolkit/config/examples/mylora.yaml', 'w') as file:\n",
    "    file.write(yaml)\n",
    "    file.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "869aa836-1053-4983-90e2-d4dee0f43f69",
   "metadata": {},
   "source": [
    "## Train the model\n",
    "\n",
    "Now that everything is done, we are ready to train our LoRA! This needs to be pasted into the terminal and uses the virtual env, will not work with line magic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "de891fc3-979d-4ca7-bb42-952210930342",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-08-21T21:42:32.466074Z",
     "iopub.status.busy": "2024-08-21T21:42:32.465480Z",
     "iopub.status.idle": "2024-08-21T21:42:32.468786Z",
     "shell.execute_reply": "2024-08-21T21:42:32.468248Z",
     "shell.execute_reply.started": "2024-08-21T21:42:32.466050Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    " python3 /notebooks/ai-toolkit/run.py /notebooks/ai-toolkit/config/examples/mylora.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2677f93-0f7b-4431-b07a-1d91b45be398",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Run inference with the new FLUX.1 LoRA\n",
    "\n",
    "If you left your LoRA with the same name as default, then we can run the following cell as is to generate an image with the subject of the training!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4293b7e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U diffusers accelerate transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d7f0bbde-6764-42d7-b616-df90f2842e3b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/paperspace/Downloads/ai-toolkit/venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Loading pipeline components...:   0%|          | 0/7 [00:00<?, ?it/s]You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers\n",
      "Loading checkpoint shards: 100%|██████████| 2/2 [00:02<00:00,  1.24s/it]it/s]\n",
      "Loading pipeline components...: 100%|██████████| 7/7 [00:08<00:00,  1.27s/it]\n",
      "100%|██████████| 50/50 [01:13<00:00,  1.47s/it]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from diffusers import DiffusionPipeline\n",
    "\n",
    "model_id = 'black-forest-labs/FLUX.1-dev'\n",
    "adapter_id = f'output/{lora_name}/{lora_name}.safetensors'\n",
    "pipeline = DiffusionPipeline.from_pretrained(model_id)\n",
    "pipeline.load_lora_weights(adapter_id)\n",
    "\n",
    "prompt = \"ethnographic photography of man at a picnic\"\n",
    "negative_prompt = \"blurry, cropped, ugly\"\n",
    "\n",
    "pipeline.to('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')\n",
    "image = pipeline(\n",
    "    prompt=prompt,\n",
    "    num_inference_steps=50,\n",
    "    generator=torch.Generator(device='cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu').manual_seed(1641421826),\n",
    "    width=1152,\n",
    "    height=768,\n",
    ").images[0]\n",
    "# image.save(\"output.png\", format=\"PNG\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bcf99335-95d8-4456-a2b0-c7dbe223604d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-08-21T21:28:12.539986Z",
     "iopub.status.busy": "2024-08-21T21:28:12.539616Z",
     "iopub.status.idle": "2024-08-21T21:28:12.542376Z",
     "shell.execute_reply": "2024-08-21T21:28:12.541939Z",
     "shell.execute_reply.started": "2024-08-21T21:28:12.539964Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "display(image)"
   ]
  }
 ],
 "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}