{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55f9ccf6-d430-4c0a-aa60-12d8d4425c60",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8370ded1-4832-46b6-944f-235c032673b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25f44308-a145-4af2-8fd2-81dc50eaf472",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the dataset\n",
    "from datasets import load_dataset\n",
    "ds = load_dataset(\"hdparmar/irish-traditional-tunes\", split=\"train\")\n",
    "sample = ds[0]\n",
    "display(sample[\"image\"])\n",
    "print(sample[\"text\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eac20890-d221-449f-9cd6-a71c9099bff5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import hf_hub_download\n",
    "ckpt_path = hf_hub_download(repo_id=\"riffusion/riffusion-model-v1\", filename=\"riffusion-model-v1.ckpt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "985a83ee-43d4-4d95-8904-4a601d3a8bca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2xA6000:\n",
    "BATCH_SIZE = 4\n",
    "N_GPUS = 2\n",
    "ACCUMULATE_BATCHES = 2\n",
    "GPU_LIST = \",\".join((str(x) for x in range(N_GPUS))) + \",\"\n",
    "print(f\"Using GPUs: {GPU_LIST}\")\n",
    "\n",
    "import os\n",
    "os.environ[\"BATCH_SIZE\"] = str(BATCH_SIZE) \n",
    "os.environ[\"N_GPUS\"] = str(N_GPUS) \n",
    "os.environ[\"ACCUMULATE_BATCHES\"] = str(ACCUMULATE_BATCHES)\n",
    "os.environ[\"GPU_LIST\"] = GPU_LIST\n",
    "os.environ[\"CKPT_PATH\"] = ckpt_path\n",
    "\n",
    "!echo \"$BATCH_SIZE\"\n",
    "!echo \"$N_GPUS\"\n",
    "!echo \"$ACCUMULATE_BATCHES\"\n",
    "!echo \"$GPU_LIST\"\n",
    "!echo \"$CKPT_PATH\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64d80dca-e7c6-438f-965c-3322d0cb33ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run training\n",
    "!(python main.py \\\n",
    "    -t \\\n",
    "    --base configs/stable-diffusion/pokemon.yaml \\\n",
    "    --gpus \"$GPU_LIST\" \\\n",
    "    --scale_lr False \\\n",
    "    --num_nodes 1 \\\n",
    "    --check_val_every_n_epoch 6 \\\n",
    "    --finetune_from \"$CKPT_PATH\" \\\n",
    "    data.params.batch_size=\"$BATCH_SIZE\" \\\n",
    "    lightning.trainer.accumulate_grad_batches=\"$ACCUMULATE_BATCHES\" \\\n",
    "    data.params.validation.params.n_gpus=\"$N_GPUS\" \\\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "438dd0c2-674b-408e-b8d6-8b789dfcd2ee",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}