{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8d170e51-0ecd-484b-b7ad-2907fc6bd0ff",
   "metadata": {},
   "source": [
    "# Hunyuan-DiT 训练示例教程"
   ]
  },
  {
   "attachments": {
    "e79c217f-2bd2-4be2-b856-a6679f8b8b37.png": {
     "image/png": "[encoded data, 316956 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "17e41dcd-ff7f-434f-9c65-dbe786cfffe9",
   "metadata": {},
   "source": [
    "\n",
    "![image.png](attachment:e79c217f-2bd2-4be2-b856-a6679f8b8b37.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65088730-4653-4f7e-8b86-9f626baa3db3",
   "metadata": {},
   "source": [
    "## 1. 进入环境 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "518f5e7a-a109-4c90-b872-a55d90967e2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# （务必）开局必运行，否则找不到代码环境\n",
    "import os\n",
    "\n",
    "current_directory = os.path.basename(os.getcwd())\n",
    "if current_directory != \"HunyuanDiT\":\n",
    "    os.chdir('./HunyuanDiT')\n",
    "    print(\"已进入\")\n",
    "else:\n",
    "    print(\"已经在对应位置\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "899f753e-f651-4bdd-a982-91d070d09944",
   "metadata": {},
   "source": [
    "使用下面的命令下载从huggingface下载模型权重 （包含文本编码器clip、t5，图像编解码器vae权重等）\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1c6463d-562e-42d5-8bab-2e35f6e0baf8",
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir ckpts\n",
    "! huggingface-cli download Tencent-Hunyuan/HunyuanDiT-v1.2 --local-dir ./ckpts"
   ]
  },
  {
   "attachments": {
    "04f2a54f-c673-4446-a76c-0059d9d6b58f.png": {
     "image/png": "[encoded data, 353800 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "c5910361-97bc-4793-9c7e-0824ed9d29df",
   "metadata": {},
   "source": [
    "从这里手动下载蒸馏版模型并放到本地 (路径 \"/HunyuanDiT/ckpts/t2i/model/pytorch_model_distill.pt\")  \n",
    "https://huggingface.co/Tencent-Hunyuan/Distillation-v1.2/tree/main\n",
    "![image.png](attachment:04f2a54f-c673-4446-a76c-0059d9d6b58f.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fea6921c-7a8a-4b77-99fb-cc69f3ed9259",
   "metadata": {},
   "source": [
    "## 2. 数据集准备"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "773a6f8b-bd68-46a7-8b1c-cb0949068bb0",
   "metadata": {},
   "source": [
    "安装数据加载器依赖包 （用于训练时读取图片和文本描述数据）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9d51d06-142c-4585-9da1-e71a2aa4d2f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -e ./IndexKits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3af6344f-6e07-4fc3-9fd8-e787d9e0903d",
   "metadata": {},
   "source": [
    "样例图片数据下载"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1aab015-12a1-486a-a7c5-8771f61ace3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "! wget -O ./dataset/data_demo.zip https://dit.hunyuan.tencent.com/download/HunyuanDiT/data_demo.zip\n",
    "! unzip -o ./dataset/data_demo.zip -d ./dataset\n",
    "! mkdir ./dataset/porcelain/arrows ./dataset/porcelain/jsons\n",
    "! echo \"样例数据下载完成！\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "476d39a0-91a2-4442-99e2-2edaa400d518",
   "metadata": {},
   "source": [
    "图片描述器 （数据打标）"
   ]
  },
  {
   "attachments": {
    "1711751f-3e50-432e-9c9d-73593899aece.png": {
     "image/png": "[encoded data, 2328180 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "2d318aa4-d32b-4e5c-81c9-10c9fd3b6381",
   "metadata": {},
   "source": [
    "![image.png](attachment:1711751f-3e50-432e-9c9d-73593899aece.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "208a9295-fb33-40f3-b3e4-aff7e4f9ee4a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# (可不运行) 随便推一张数据集中的图，看看ai给出的描述情况\n",
    "! python mllm/caption_demo.py --mode \"caption_zh\" --image_file \"mllm/images/demo1.png\" --model_path \"./ckpts/captioner\"\n",
    "# 其中--image_file \"mllm/images/demo1.png\"替换为你要测试的图片样例路径"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d97be2ff-98f7-4769-acb4-a8e29df1c523",
   "metadata": {},
   "source": [
    "### 以青花瓷数据集为例\n",
    "PS：如果想用自己的数据训练，可以直接简单替换./dataset/porcelain/images下的所有图片为自己上传的图片后，重新运行以下代码 （小白推荐）  \n",
    "或者自己仿照./dataset/porcelain的文件夹，建立自己的数据集文件夹，以深入了解其思想"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d7ac423-1b5b-4683-894c-6d2de99ce0c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 谨慎运行以下代码！会删除所有原先图片数据的指令, 以方便自行上传图片数据训练\n",
    "! rm -rf ./dataset/porcelain/images/*  && echo \"清空原有数据完成\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d083d05a-e790-468e-a0f1-7d4f0a057c42",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 这里仅仅是例子，把wukong的demo数据集用于学习\n",
    "! cp ../demo_images_wukong/*.png ./dataset/porcelain/images/ && echo \"复制数据完成\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36cdaf2c-8c36-4abf-aefc-b3a9a5c3dc23",
   "metadata": {},
   "outputs": [],
   "source": [
    "### 把图片路径记录在csv文件\n",
    "! python mllm/make_csv.py --img_dir \"./dataset/porcelain/images\" --input_file \"./dataset/porcelain/csvfile/images.csv\"\n",
    "\n",
    "### 批量推理csv中每张图片的文本描述\n",
    "! python mllm/caption_demo.py --mode \"caption_zh\" --input_file \"./dataset/porcelain/csvfile/images.csv\" --output_file \"./dataset/porcelain/csvfile/image_text.csv\" --model_path \"./ckpts/captioner\"\n",
    "\n",
    "# 把csv表转成.arrow格式文件 （indexkits包使用的数据读取格式） \n",
    "! python ./hydit/data_loader/csv2arrow.py \\\n",
    "        ./dataset/porcelain/csvfile/image_text.csv \\\n",
    "        ./dataset/porcelain/arrows 1\n",
    "\n",
    "! echo \"数据打标 & 转换记录完成!\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb236c62-89cf-46b3-bea7-9c21c998e4e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 单分辨率数据处理\n",
    "! idk base -c dataset/yamls/porcelain.yaml -t dataset/porcelain/jsons/porcelain.json\n",
    "\n",
    "# 多种分辨率数据处理     \n",
    "! idk multireso -c dataset/yamls/porcelain_mt.yaml -t dataset/porcelain/jsons/porcelain_mt.json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af6ce54f-1f0f-4f05-b0a1-7499b0450590",
   "metadata": {},
   "source": [
    "### 3. 训练模型"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4cb0238-42a7-4471-981d-f7633540ec71",
   "metadata": {},
   "source": [
    "### 3.1 lora模型训练\n",
    "具体训练配置可以自行修改lora/train_lora.sh脚本文件  \n",
    "当前rank默认设置为128，你可以根据在自己的数据集上训练的结果测试情况调大、或者调小  \n",
    "如果发现lora训练效果不够好，建议优先清洗 & 处理数据集中不理想的样本  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c711e0d2-1988-4086-b403-91dd1df6e2c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "! PYTHONPATH=./ sh lora/train_lora.sh --index-file dataset/porcelain/jsons/porcelain.json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e03119f-0773-4806-9f12-ead9d883ae07",
   "metadata": {},
   "source": [
    "### 3.2 直接通过代码测试\n",
    "可自行修改想要使用的训练步数检查点，当前取step1500（默认训练2000步）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63d392a1-173d-441a-9375-82eb4cb81014",
   "metadata": {},
   "outputs": [],
   "source": [
    "! python sample_t2i.py \\\n",
    "    --prompt \"青花瓷风格，一只猫在追蝴蝶\" \\\n",
    "    --no-enhance \\\n",
    "    --load-key distill"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b493a6c-6f52-4175-9442-29f55447c54d",
   "metadata": {},
   "outputs": [],
   "source": [
    "! python sample_t2i.py \\\n",
    "    --prompt \"青花瓷风格，一只猫在追蝴蝶\" \\\n",
    "    --no-enhance \\\n",
    "    --load-key distill \\\n",
    "    --lora-ckpt ./log_EXP/001-lora_porcelain_ema_rank128/checkpoints/0001500.pt/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c90265dc-09aa-4d73-b895-3290c5d43268",
   "metadata": {},
   "source": [
    "### 3.3 移动到ComfyUI测试\n",
    "把训练好的lora模型权重文件移动到ComfyUI的对应加载模型文件夹\n",
    "可以自行修改复制后的权重文件名称xxx_step1500.safetensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "149b7552-e19d-4383-a537-5c5e3cb566bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "! cp ./log_EXP/001-lora_porcelain_ema_rank128/checkpoints/0001500.pt/adapter_model.safetensors /path/to/ComfyUI/models/loras/xxx_step1500.safetensors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c02d6b79-f7ad-4c8c-9f24-86577275dca5",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 享受混元DiT文生图！"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (base)",
   "language": "python",
   "name": "base"
  },
  "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}