Demos/14a. 尝试使用 LoRA 微调 Stable Diffusion 模型.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a13831c9-2a09-40e7-98ae-3ee6be8a52cd",
   "metadata": {},
   "source": [
    "# a. 尝试使用 LoRA 微调 Stable Diffusion 模型（文生图）\n",
    "\n",
    "> [HW10: Stable Diffusion Fine-tuning](https://colab.research.google.com/drive/1dI_-HVggxyIwDVoreymviwg6ZOvEHiLS?usp=sharing#scrollTo=CnJtiRaRuTFX) 中文镜像版\n",
    ">\n",
    "> 这次不是1:1的镜像，我重构了主要的代码逻辑，以便于阅读。对于可以访问 Colab 的同学来说，一样建议使用下面这份代码进行学习。\n",
    ">\n",
    "> 修改：\n",
    "> 1. Colab 中的代码每次训练时先导入了 checkpoint-last 文件夹下的 unet 和 text_encoder，这就意味着每次重新运行是强制 resume 的。\n",
    "> 我对这个行为进行了修改，你现在可以通过指定 prepare_lora_model 中的 resume 参数来选择是否接着之前的进行训练。\n",
    "> 建议设置为 False，每次重新运行才能对不同参数有直观的感受。\n",
    ">\n",
    "> 2. 原代码中导入了 PEFT 并设置了 LoRA 参数却没有用到，出于学习而非提交作业的目的，我会重新写这一模块。\n",
    "> 重写之后的可训练参数层与数量已经经过验证，与 [yahcreeper/GenAI-HW10-Model](https://huggingface.co/yahcreeper/GenAI-HW10-Model/tree/main/checkpoint-last) 完全一致，可放心使用当前版本（关于可训练数量，可以参考[《14. PEFT：在大模型中快速应用 LoRA.md》](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Guide/14.%20PEFT：在大模型中快速应用%20LoRA.md)使用 `print_trainable_parameters()` 进行打印）。\n",
    ">\n",
    "> 3. 会增加注释以便理解。\n",
    ">\n",
    "> 指导文章：[16. 用 LoRA 微调 Stable Diffusion：拆开炼丹炉，动手实现你的第一次 AI 绘画](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Guide/16.%20用%20LoRA%20微调%20Stable%20Diffusion：拆开炼丹炉，动手实现你的第一次%20AI%20绘画.md)\n",
    "> \n",
    "\n",
    "你将微调你自己的 Stable Diffusion 模型，从给定的文本描述生成自定义图像，在初始化项目之后，你将可以在 `SD/Datasets/Brad` 下看到图像和对应的文本描述。\n",
    "\n",
    "当前代码使用到的数据集 `dataset.zip` 已经放在 Demos/data/14 下面。\n",
    "\n",
    "你需要注意的是，这份代码倾向于如何微调而非加载数据集或者评估，对于初学者，可以仅关注训练参数和LoRA部分，其余部分的代码细节不用关注。\n",
    "\n",
    "这次的学习可以在两种版本中二选一，学习效果一致。\n",
    "\n",
    "[版本 b - 精简版](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Demos/14b.%20尝试使用%20LoRA%20微调%20Stable%20Diffusion%20模型%20-%20精简版.ipynb)\n",
    "\n",
    "这里还有一个简单的 [🎡 SD LoRA 脚本](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/CodePlayground/sd_lora.py)供你尝试，详见：[CodePlayground](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/CodePlayground/README.md#当前的玩具)，点击 `►` 或对应的文本展开。\n",
    "\n",
    "在线链接（精简版）：[Kaggle](https://www.kaggle.com/code/aidemos/14b-lora-stable-diffusion) | [Colab](https://colab.research.google.com/drive/1idmnaQZwRhjUPw7ToEXlVo82Mihfl_aA?usp=sharing)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8a8deb6-6844-400c-bb10-388d89d3a020",
   "metadata": {},
   "source": [
    "## 安装必要的库\n",
    "\n",
    "本单元中将安装一些库用于微调。\n",
    "\n",
    "安装大约需要 5 分钟。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3629cf04-6988-483d-bfc9-509f8b59a9f5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install timm\n",
    "!pip install fairscale\n",
    "!pip install transformers\n",
    "!pip install requests\n",
    "!pip install accelerate\n",
    "!pip install diffusers\n",
    "!pip install einop\n",
    "!pip install safetensors\n",
    "!pip install voluptuous\n",
    "!pip install jax\n",
    "!pip install jaxlib\n",
    "!pip install peft\n",
    "!pip install deepface==0.0.90\n",
    "!pip install tensorflow==2.9.0  # 为了避免最后评估阶段使用deepface时的错误，这里选择降级版本\n",
    "!pip install keras\n",
    "!pip install opencv-python\n",
    "#!pip install gdown"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5a92e17-0911-4f29-9bc1-6f81e1dc287d",
   "metadata": {},
   "source": [
    "## 导入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bd2c745-a37a-42d8-940c-2cdcbc784b0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ========== 标准库模块 ==========\n",
    "import os\n",
    "import math\n",
    "import random\n",
    "import glob\n",
    "import shutil\n",
    "import argparse\n",
    "from pathlib import Path\n",
    "import subprocess\n",
    "\n",
    "# ========== 第三方库 ==========\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from PIL import Image\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "# ========== 深度学习相关库 ==========\n",
    "from torchvision import transforms\n",
    "from torchvision.utils import save_image\n",
    "\n",
    "# Transformers (Hugging Face)\n",
    "from transformers import AutoProcessor, AutoModel, CLIPTextModel, CLIPTokenizer\n",
    "\n",
    "# Diffusers (Hugging Face)\n",
    "from diffusers import (\n",
    "    AutoencoderKL,\n",
    "    DDPMScheduler,\n",
    "    DiffusionPipeline,\n",
    "    StableDiffusionPipeline,\n",
    "    UNet2DConditionModel\n",
    ")\n",
    "from diffusers.optimization import get_scheduler\n",
    "from diffusers.utils import convert_state_dict_to_diffusers\n",
    "from diffusers.training_utils import compute_snr\n",
    "from diffusers.utils.torch_utils import is_compiled_module\n",
    "\n",
    "# ========== LoRA 模型库 ==========\n",
    "from peft import LoraConfig, get_peft_model\n",
    "\n",
    "# ========== 面部检测库 ==========\n",
    "from deepface import DeepFace\n",
    "\n",
    "# ========== IPython 和 Widgets ==========\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import display, Markdown, clear_output\n",
    "\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62438d83-b6c5-43a3-afe9-30ceebc50d46",
   "metadata": {},
   "source": [
    "## 准备项目\n",
    "\n",
    "直接运行代码，这个模块不用修改参数，不用关心这里的代码细节，除非你对交互感兴趣。\n",
    "\n",
    "注意：记得点击**初始化项目**\n",
    "\n",
    "当你看见✅时，代表项目已经准备好。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fbde642-16f6-4728-8b7b-05d220644de0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建项目名称输入框\n",
    "project_name_widget = widgets.Text(\n",
    "    value=\"Brad\",\n",
    "    description=\"项目名称:\",\n",
    "    style={'description_width': 'initial'}\n",
    ")\n",
    "\n",
    "# 创建数据集名称输入框\n",
    "dataset_name_widget = widgets.Text(\n",
    "    value=\"Brad\",  # \"Brad-512\", \"Anne-512\"\n",
    "    description=\"数据集:\",\n",
    "    style={'description_width': 'initial'},\n",
    "    disabled=True\n",
    ")\n",
    "\n",
    "# 创建提交按钮\n",
    "submit_button = widgets.Button(description=\"初始化项目\")\n",
    "\n",
    "# 输出区域\n",
    "output_area = widgets.Output()\n",
    "\n",
    "# 定义项目初始化逻辑\n",
    "def initialize_project(b):\n",
    "    global project_name, dataset_name, root_dir, main_dir, project_dir, model_path, images_folder, prompts_folder, captions_folder\n",
    "    with output_area:\n",
    "        clear_output()\n",
    "        project_name = project_name_widget.value.strip()\n",
    "        dataset_name = dataset_name_widget.value.strip()\n",
    "        \n",
    "        if not project_name or any(c in project_name for c in \" .()\\\"'\\\\\") or project_name.count(\"/\") > 1:\n",
    "            print(\"请输入有效的项目名称。\")\n",
    "        else:\n",
    "            # 创建项目目录\n",
    "            project_base = project_name if \"/\" not in project_name else project_name[:project_name.rfind(\"/\")]\n",
    "            project_subfolder = project_name if \"/\" not in project_name else project_name[project_name.rfind(\"/\")+1:]\n",
    "            root_dir = \"./\"  # 当前目录\n",
    "            main_dir = os.path.join(root_dir, \"SD\")  # 主目录\n",
    "            project_dir = os.path.join(main_dir, project_name)  # 项目目录\n",
    "            \n",
    "            # 确保目录存在\n",
    "            os.makedirs(main_dir, exist_ok=True)\n",
    "            os.makedirs(project_dir, exist_ok=True)\n",
    "            \n",
    "            # 定义数据集和模型路径\n",
    "            zip_file = os.path.join(\"./\", \"data/14/Datasets.zip\")\n",
    "            model_path = os.path.join(project_dir, \"logs\", \"checkpoint-last\")\n",
    "            images_folder = os.path.join(main_dir, \"Datasets\", dataset_name)\n",
    "            prompts_folder = os.path.join(main_dir, \"Datasets\", \"prompts\")\n",
    "            captions_folder = images_folder\n",
    "\n",
    "            # 检查并解压数据集\n",
    "            print(\"📂 正在检查并解压样例数据集...\")\n",
    "            \n",
    "            if not os.path.exists(zip_file):\n",
    "                print(\"❌ 未找到数据集压缩文件 Datasets.zip！\")\n",
    "                print(\"请下载数据集:\\nhttps://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Demos/data/14/Datasets.zip\\n并放在 ./data/14 文件夹下\")\n",
    "            else:\n",
    "                subprocess.run(f\"unzip -q -o {zip_file} -d {main_dir}\", shell=True)\n",
    "                \n",
    "                # # 克隆模型文件\n",
    "                print(\"📂 默认不克隆LoRA模型文件而是自己训练，因为下载需要占用2G且等价于你先进行1000个step训练。\\n\")\n",
    "    \n",
    "                # 创建必要的文件夹\n",
    "                os.makedirs(images_folder, exist_ok=True)\n",
    "    \n",
    "                print(f\"✅ 项目 {project_name} 已准备好！\")\n",
    "\n",
    "# 将提交按钮与项目初始化逻辑绑定\n",
    "submit_button.on_click(initialize_project)\n",
    "\n",
    "# 显示控件\n",
    "display(project_name_widget, dataset_name_widget, submit_button, output_area)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b44210b9-245d-4640-ae1c-d4a55420c984",
   "metadata": {},
   "source": [
    "## 设置参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f253cc2c-42eb-476d-aa9b-a740c432caf0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 请勿随意更改以下参数，除非你自己知道改的参数对应的什么，否则可能因 GPU 内存不足导致进程崩溃。\n",
    "output_folder = os.path.join(project_dir, \"logs\")  # 存放 model checkpoints 和 validation 的文件夹\n",
    "seed = 1126  # 随机数种子\n",
    "train_batch_size = 2  # 训练批次大小，即每次训练中处理的样本数量\n",
    "resolution = 512  # 训练图像的分辨率\n",
    "weight_dtype = torch.bfloat16  # 权重数据类型，使用 bfloat16 以节省内存并加快计算速度\n",
    "snr_gamma = 5  # SNR 参数，用于信噪比加权损失的调节系数"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "255ee76f-e838-4296-bf6d-fe175d9e08b8",
   "metadata": {},
   "source": [
    "### Stable Diffusion LoRA 的微调参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a07a3378-e074-4793-ac5e-f8ba7bab2056",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 预训练的 Stable Diffusion 模型路径，用于加载模型进行微调\n",
    "pretrained_model_name_or_path = \"stablediffusionapi/cyberrealistic-41\"\n",
    "\n",
    "# LoRA 配置\n",
    "lora_config = LoraConfig(\n",
    "    r=32,  # LoRA 的秩，即低秩矩阵的维度，决定了参数调整的自由度\n",
    "    lora_alpha=16,  # 缩放系数，控制 LoRA 权重对模型的影响\n",
    "    target_modules=[\n",
    "        \"q_proj\", \"v_proj\", \"k_proj\", \"out_proj\",  # 指定 Text encoder 的 LoRA 应用对象（用于调整注意力机制中的投影矩阵）\n",
    "        \"to_k\", \"to_q\", \"to_v\", \"to_out.0\"  # 指定 UNet 的 LoRA 应用对象（用于调整 UNet 中的注意力机制）\n",
    "    ],\n",
    "    lora_dropout=0  # LoRA dropout 概率，0 表示不使用 dropout\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "806d27ec-5425-454a-bb6c-506db30fe69c",
   "metadata": {},
   "source": [
    "### 其他设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1b00988-3201-4c16-b363-17528f5cadb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 学习率调度器部分\n",
    "lr_scheduler_name = \"cosine_with_restarts\" # 设置学习率的调度器\n",
    "lr_warmup_steps = 100 # 设置学习率预热步数\n",
    "\n",
    "# prompt 处理\n",
    "validation_prompt_name = \"validation_prompt.txt\"\n",
    "validation_prompt_path = os.path.join(prompts_folder, validation_prompt_name)\n",
    "with open(validation_prompt_path, \"r\") as f:\n",
    "    validation_prompt = [line.strip() for line in f.readlines()]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a74e1b82-00a7-4e4e-92cb-8aaf50aa2a60",
   "metadata": {},
   "source": [
    "直接运行代码，交互式修改参数，不用关心这里的代码细节，除非你对交互感兴趣。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a70fdd8b-3e5b-4657-bacd-206b087375b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 学习率部分\n",
    "learning_rate_markdown = Markdown(\"▶️ **Learning Rate**<br>学习率是影响结果最重要的参数，当前默认三者一致。<br>如果你想慢速训练且有大量图像，或者如果 LoRA rank（有些文生图UI可能用的是dim） 和 alpha 值较高，建议将 unet 的学习率调至 5e-5 或更低。<br>需要注意的是，实际上learning_rate在这里不起作用，我们只使用对应模块的学习率。<br>这里进行展示是为了让你之后在炼丹炉上能够更好的理解设置。\")\n",
    "learning_rate_widget = widgets.FloatText(\n",
    "    value=1e-4,\n",
    "    description='learning_rate:',\n",
    "    style={'description_width': '200px'},\n",
    "    layout=widgets.Layout(width='300px'),\n",
    "    disabled=True  # 禁用手动修改\n",
    ")\n",
    "\n",
    "unet_learning_rate_widget = widgets.FloatText(\n",
    "    value=1e-4,\n",
    "    description='unet_learning_rate:',\n",
    "    style={'description_width': '200px'},\n",
    "    layout=widgets.Layout(width='300px')\n",
    ")\n",
    "\n",
    "text_encoder_learning_rate_widget = widgets.FloatText(\n",
    "    value=1e-4,\n",
    "    description='text_encoder_learning_rate:',\n",
    "    style={'description_width': '200px'},\n",
    "    layout=widgets.Layout(width='300px')\n",
    ")\n",
    "\n",
    "# 训练步骤部分\n",
    "max_train_steps_markdown = Markdown(\"▶️ **Steps**<br>选择训练步骤和每次验证生成的图像数量\")\n",
    "max_train_steps_widget = widgets.IntSlider(\n",
    "    value=200,\n",
    "    min=200,\n",
    "    max=2000,\n",
    "    step=100,\n",
    "    description='max_train_steps:',\n",
    "    continuous_update=False,\n",
    "    style={'description_width': '200px'},\n",
    "    layout=widgets.Layout(width='400px')\n",
    ")\n",
    "\n",
    "\n",
    "validation_prompt_num_widget = widgets.IntSlider(\n",
    "    value=3,\n",
    "    min=1,\n",
    "    max=5,\n",
    "    step=1,\n",
    "    description='validation_prompt_num:',\n",
    "    style={'description_width': '200px'},\n",
    "    layout=widgets.Layout(width='400px')\n",
    ")\n",
    "\n",
    "\n",
    "validation_step_ratio_widget = widgets.FloatSlider(\n",
    "    value=1,\n",
    "    min=0,\n",
    "    max=1,\n",
    "    step=0.1,\n",
    "    description='validation_step_ratio:',\n",
    "    style={'description_width': '200px'},\n",
    "    layout=widgets.Layout(width='400px')\n",
    ")\n",
    "\n",
    "\n",
    "validation_prompt_widget = widgets.Text(\n",
    "    value = validation_prompt_path,\n",
    "    description='validation_prompt_path:',\n",
    "    style={'description_width': '200px'},\n",
    "    layout=widgets.Layout(width='550px'),\n",
    "    disabled=True  # 禁用手动修改\n",
    ")\n",
    "\n",
    "# 创建输出区域来显示打印内容\n",
    "output_area = widgets.Output()\n",
    "\n",
    "# 获取用户输入值并打印配置\n",
    "def on_button_click(b):\n",
    "    global unet_learning_rate, text_encoder_learning_rate, max_train_steps, validation_prompt, validation_prompt_num, validation_step_ratio\n",
    "    with output_area:  # 使用 Output 小部件捕获输出\n",
    "        output_area.clear_output()  # 清除之前的输出\n",
    "        unet_learning_rate = unet_learning_rate_widget.value\n",
    "        text_encoder_learning_rate = text_encoder_learning_rate_widget.value\n",
    "        max_train_steps = max_train_steps_widget.value\n",
    "        validation_prompt_num = validation_prompt_num_widget.value\n",
    "        validation_step_ratio = validation_step_ratio_widget.value\n",
    "\n",
    "        # 打印配置\n",
    "        print(f\"UNet 学习率: {unet_learning_rate}\")\n",
    "        print(f\"文本编码器学习率: {text_encoder_learning_rate}\")\n",
    "        print(f\"训练步骤: {max_train_steps}\")\n",
    "        print(f\"验证prompt数量: {validation_prompt_num}\")\n",
    "        print(f\"验证步骤比例: {validation_step_ratio}\")\n",
    "        print(f\"用于验证的prompt文件位置: {validation_prompt_path}\")\n",
    "\n",
    "# 创建提交按钮\n",
    "submit_button = widgets.Button(description=\"提交配置\")\n",
    "submit_button.on_click(on_button_click)\n",
    "\n",
    "# 显示带有说明的所有小部件\n",
    "display(learning_rate_markdown, learning_rate_widget, unet_learning_rate_widget, text_encoder_learning_rate_widget)\n",
    "display(max_train_steps_markdown, max_train_steps_widget, validation_prompt_num_widget, validation_step_ratio_widget, validation_prompt_widget)\n",
    "\n",
    "# 显示提交按钮\n",
    "display(submit_button)\n",
    "\n",
    "# 显示输出区域\n",
    "display(output_area)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d98bedfc-4764-4a33-b883-cc1e8f3a89dc",
   "metadata": {},
   "source": [
    "## 定义一些有用的函数和类\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "953edb1b-a08f-4cd0-876d-a0b352ba7fb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "IMAGE_EXTENSIONS = [\".png\", \".jpg\", \".jpeg\", \".webp\", \".bmp\", \".PNG\", \".JPG\", \".JPEG\", \".WEBP\", \".BMP\"]\n",
    "\n",
    "# 数据增强\n",
    "train_transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR),  # 调整图像大小\n",
    "        transforms.CenterCrop(resolution),  # 中心裁剪图像\n",
    "        transforms.RandomHorizontalFlip(),  # 随机水平翻转\n",
    "        transforms.ToTensor(),  # 将图像转换为张量\n",
    "    ]\n",
    ")\n",
    "class Text2ImageDataset(torch.utils.data.Dataset):\n",
    "    \"\"\"\n",
    "    (1) 目标:\n",
    "        - 用于构建文本到图像模型的微调数据集\n",
    "\n",
    "    \"\"\"\n",
    "    def __init__(self, images_folder, captions_folder, transform, tokenizer):\n",
    "        \"\"\"\n",
    "        (2) 参数:\n",
    "            - images_folder: str, 图像文件夹路径\n",
    "            - captions_folder: str, 标注文件夹路径\n",
    "            - transform: function, 将原始图像转换为 torch.tensor\n",
    "            - tokenizer: CLIPTokenizer, 将文本标注转为 word ids\n",
    "        \"\"\"\n",
    "        # 初始化图像路径列表，并根据指定的扩展名找到所有图像文件\n",
    "        self.image_paths = []\n",
    "        for ext in IMAGE_EXTENSIONS:\n",
    "            self.image_paths.extend(glob.glob(f\"{images_folder}/*{ext}\"))\n",
    "        self.image_paths = sorted(self.image_paths)\n",
    "\n",
    "        # 通过 DeepFace 库提取每张图像的面部嵌入（特征向量）\n",
    "        self.train_emb = torch.tensor([DeepFace.represent(img_path, detector_backend=\"ssd\", model_name=\"GhostFaceNet\", enforce_detection=False)[0]['embedding'] for img_path in self.image_paths])\n",
    "\n",
    "        # 加载对应的文本标注，依次读取每个文本文件中的内容\n",
    "        caption_paths = sorted(glob.glob(f\"{captions_folder}/*txt\"))\n",
    "        captions = []\n",
    "        for p in caption_paths:\n",
    "            with open(p, \"r\") as f:\n",
    "                captions.append(f.readline())\n",
    "\n",
    "        # 使用 tokenizer 将文本标注转换为 word ids\n",
    "        inputs = tokenizer(\n",
    "            captions, max_length=tokenizer.model_max_length, padding=\"max_length\", truncation=True, return_tensors=\"pt\"\n",
    "        )\n",
    "        self.input_ids = inputs.input_ids\n",
    "        self.transform = transform\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_path = self.image_paths[idx]\n",
    "        input_id = self.input_ids[idx]\n",
    "        try:\n",
    "            # 加载图像并将其转换为 RGB 模式，然后应用数据增强\n",
    "            image = Image.open(img_path).convert(\"RGB\")\n",
    "            tensor = self.transform(image)\n",
    "        except Exception as e:\n",
    "            print(f\"无法加载图像路径: {img_path}, 错误: {e}\")\n",
    "            return None\n",
    "\n",
    "        return tensor, input_id  # 返回处理后的图像和相应的文本标注\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.image_paths)\n",
    "    \n",
    "def prepare_lora_model(lora_config, pretrained_model_name_or_path=\"runwayml/stable-diffusion-v1-5\", model_path=None, resume=False):\n",
    "    \"\"\"\n",
    "    (1) 目标:\n",
    "        - 用于加载完整的 Stable Diffusion 模型，包括 Lora 层，并冻结非 Lora 参数。这包括 Tokenizer、噪声调度器、UNet、VAE 和 文本编码器。\n",
    "\n",
    "    (2) 参数:\n",
    "        - pretrained_model_name_or_path: str, 从 Hugging Face 获取的模型名称\n",
    "        - model_path: str, 预训练模型的路径\n",
    "        - resume: bool, 是否从上一次训练中恢复\n",
    "\n",
    "    (3) 返回:\n",
    "        - 输出: Tokenizer, 噪声调度器, UNet, VAE, 文本编码器\n",
    "\n",
    "    \"\"\"\n",
    "    # 加载噪声调度器，用于控制扩散模型的噪声添加和移除过程\n",
    "    noise_scheduler = DDPMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder=\"scheduler\")\n",
    "\n",
    "    # 加载 Tokenizer，用于将文本标注转换为 tokens\n",
    "    tokenizer = CLIPTokenizer.from_pretrained(\n",
    "        pretrained_model_name_or_path,\n",
    "        subfolder=\"tokenizer\"\n",
    "    )\n",
    "\n",
    "    # 加载 CLIP 文本编码器，用于将文本标注转换为特征向量\n",
    "    text_encoder = CLIPTextModel.from_pretrained(\n",
    "        pretrained_model_name_or_path,\n",
    "        torch_dtype=weight_dtype,\n",
    "        subfolder=\"text_encoder\"\n",
    "    )\n",
    "\n",
    "    # 加载 VAE 模型，用于在扩散模型中处理图像的潜在表示\n",
    "    vae = AutoencoderKL.from_pretrained(\n",
    "        pretrained_model_name_or_path,\n",
    "        subfolder=\"vae\"\n",
    "    )\n",
    "\n",
    "    # 加载 UNet 模型，负责处理扩散模型中的图像生成和推理过程\n",
    "    unet = UNet2DConditionModel.from_pretrained(\n",
    "        pretrained_model_name_or_path,\n",
    "        torch_dtype=weight_dtype,\n",
    "        subfolder=\"unet\"\n",
    "    )\n",
    "\n",
    "    # 将 LoRA 配置应用到 text_encoder 和 unet\n",
    "    text_encoder = get_peft_model(text_encoder, lora_config)\n",
    "    unet = get_peft_model(unet, lora_config)\n",
    "    \n",
    "    # 打印可训练参数数量\n",
    "    text_encoder.print_trainable_parameters()\n",
    "    unet.print_trainable_parameters()\n",
    "    \n",
    "    # 如果设置为继续训练，则加载上一次的模型权重，当然，你可以修改 model_path 来指定其他的路径\n",
    "    if resume:\n",
    "        if model_path is None or not os.path.exists(model_path):\n",
    "            raise ValueError(\"当resume设置为True的时候你必须提供有效的model_path\")\n",
    "            \n",
    "        # 加载上次训练的模型权重，注意这里只加载权重，而不是覆盖整个模型，覆盖：model = torch.load(...)\n",
    "        text_encoder = torch.load(os.path.join(model_path, \"text_encoder.pt\"))\n",
    "        unet = torch.load(os.path.join(model_path, \"unet.pt\"))\n",
    "    \n",
    "    # 冻结 VAE 参数\n",
    "    vae.requires_grad_(False)\n",
    "\n",
    "    # 使用 get_peft_model() 后会自动冻结其中的非 LoRA 层，所以可以注释掉\n",
    "    # # 冻结 UNet 和文本编码器中的非 Lora 参数\n",
    "    # for name, param in unet.named_parameters():\n",
    "    #     if \"lora\" in name:\n",
    "    #         param.requires_grad_(True)\n",
    "    #     else:\n",
    "    #         param.requires_grad_(False)\n",
    "    # for name, param in text_encoder.named_parameters():\n",
    "    #     if \"lora\" in name:\n",
    "    #         param.requires_grad_(True)\n",
    "    #     else:\n",
    "    #         param.requires_grad_(False)\n",
    "\n",
    "    # 将模型移动到 GPU 上并设置权重的数据类型\n",
    "    unet.to(DEVICE, dtype=weight_dtype)\n",
    "    vae.to(DEVICE, dtype=weight_dtype)\n",
    "    text_encoder.to(DEVICE, dtype=weight_dtype)\n",
    "    \n",
    "    return tokenizer, noise_scheduler, unet, vae, text_encoder\n",
    "\n",
    "\n",
    "def prepare_optimizer(unet, text_encoder, unet_learning_rate=5e-4, text_encoder_learning_rate=1e-4):\n",
    "    \"\"\"\n",
    "    (1) 目标:\n",
    "        - 为 UNet 和文本编码器的可训练参数分别设置优化器，并指定不同的学习率。\n",
    "\n",
    "    (2) 参数:\n",
    "        - unet: UNet2DConditionModel, Hugging Face 的 UNet 模型\n",
    "        - text_encoder: CLIPTextModel, Hugging Face 的文本编码器\n",
    "        - unet_learning_rate: float, UNet 的学习率\n",
    "        - text_encoder_learning_rate: float, 文本编码器的学习率\n",
    "\n",
    "    (3) 返回:\n",
    "        - 输出: 优化器 Optimizer\n",
    "    \"\"\"\n",
    "    # 筛选出 UNet 中需要训练的 Lora 层参数\n",
    "    unet_lora_layers = [p for p in unet.parameters() if p.requires_grad]\n",
    "    \n",
    "    # 筛选出文本编码器中需要训练的 Lora 层参数\n",
    "    text_encoder_lora_layers = [p for p in text_encoder.parameters() if p.requires_grad]\n",
    "    \n",
    "    # 将需要训练的参数分组并设置不同的学习率\n",
    "    trainable_params = [\n",
    "        {\"params\": unet_lora_layers, \"lr\": unet_learning_rate},\n",
    "        {\"params\": text_encoder_lora_layers, \"lr\": text_encoder_learning_rate}\n",
    "    ]\n",
    "    \n",
    "    # 使用 AdamW 优化器\n",
    "    optimizer = torch.optim.AdamW(trainable_params)\n",
    "    \n",
    "    return optimizer\n",
    "\n",
    "\n",
    "def evaluate(pretrained_model_name_or_path, weight_dtype, seed, unet_path, text_encoder_path, validation_prompt, output_folder, train_emb):\n",
    "    \"\"\"\n",
    "    (1) 目标:\n",
    "        - 用于加载给定路径中的 UNet 和文本编码器模型，生成图像并计算面部相似性、CLIP 分数和无面部图像的数量。\n",
    "\n",
    "    (2) 参数:\n",
    "        - pretrained_model_name_or_path: str, Hugging Face 的模型名称\n",
    "        - weight_dtype: torch.dtype, 模型的权重数据类型\n",
    "        - seed: int, 随机数种子\n",
    "        - unet_path: str, UNet 模型检查点路径\n",
    "        - text_encoder_path: str, 文本编码器模型检查点路径\n",
    "        - validation_prompt: list, 存储验证文本的字符串列表\n",
    "        - output_folder: str, 保存生成图像的文件夹\n",
    "        - train_emb: tensor, 训练图像的面部特征\n",
    "\n",
    "    (3) 返回:\n",
    "        - face_score: float, 面部相似性评分\n",
    "        - clip_score: float, CLIP 分数\n",
    "        - mis: int, 无法检测到面部的图像数量\n",
    "\n",
    "    \"\"\"\n",
    "    # 加载 DiffusionPipeline 并设置 UNet 和文本编码器\n",
    "    pipeline = DiffusionPipeline.from_pretrained(\n",
    "        pretrained_model_name_or_path,\n",
    "        torch_dtype=weight_dtype,\n",
    "        safety_checker=None,\n",
    "    )\n",
    "    pipeline.unet = torch.load(unet_path)  # 加载 UNet 模型\n",
    "    pipeline.text_encoder = torch.load(text_encoder_path)  # 加载文本编码器模型\n",
    "    pipeline = pipeline.to(DEVICE)\n",
    "\n",
    "    # 加载 CLIP 模型用于计算图像与文本的相似性分数\n",
    "    clip_model_name = \"openai/clip-vit-base-patch32\"\n",
    "    clip_model = AutoModel.from_pretrained(clip_model_name).to(DEVICE)\n",
    "    clip_processor = AutoProcessor.from_pretrained(clip_model_name)\n",
    "    \n",
    "    # 推理生成图像\n",
    "    with torch.no_grad():\n",
    "        generator = torch.Generator(device=DEVICE)\n",
    "        generator.manual_seed(seed)\n",
    "        face_score = 0\n",
    "        clip_score = 0\n",
    "        mis = 0  # 记录无法检测到面部的图像数量\n",
    "        print(\"正在生成验证图像......\")\n",
    "        images = []\n",
    "        for i in range(0, len(validation_prompt), 4):\n",
    "            images.extend(pipeline(validation_prompt[i:min(i + 4, len(validation_prompt))], num_inference_steps=30, generator=generator).images)\n",
    "        \n",
    "        # 计算面部相似性和 CLIP 分数\n",
    "        print(\"正在计算验证分数......\")\n",
    "        valid_emb = []\n",
    "        for i, image in enumerate(tqdm(images)):\n",
    "            save_file = f\"{output_folder}/valid_image_{i}.png\"\n",
    "            image.save(save_file)  # 将生成的图像保存到指定目录\n",
    "            opencvImage = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)\n",
    "\n",
    "            # 使用 DeepFace 检测面部特征\n",
    "            emb = DeepFace.represent(\n",
    "                opencvImage,\n",
    "                detector_backend=\"ssd\",\n",
    "                model_name=\"GhostFaceNet\",\n",
    "                enforce_detection=False,\n",
    "            )\n",
    "            # 如果无法检测到面部，计入 mis\n",
    "            if emb == [] or emb[0]['face_confidence'] == 0:\n",
    "                mis += 1\n",
    "                continue\n",
    "\n",
    "            # 计算 CLIP 分数，衡量图像与文本的相似度\n",
    "            emb = emb[0]\n",
    "            inputs = clip_processor(text=validation_prompt[i], images=image, return_tensors=\"pt\").to(DEVICE)\n",
    "            with torch.no_grad():\n",
    "                outputs = clip_model(**inputs)\n",
    "            sim = outputs.logits_per_image  # 获取图像与文本的相似性得分\n",
    "            clip_score += sim.item()\n",
    "            valid_emb.append(emb['embedding'])\n",
    "\n",
    "        # 如果生成的图片都没有人脸，直接 0 分\n",
    "        if len(valid_emb) == 0:\n",
    "            return 0, 0, mis\n",
    "        \n",
    "        # 计算面部相似性分数\n",
    "        valid_emb = torch.tensor(valid_emb)\n",
    "        valid_emb = (valid_emb / torch.norm(valid_emb, p=2, dim=-1)[:, None]).cuda()\n",
    "        train_emb = (train_emb / torch.norm(train_emb, p=2, dim=-1)[:, None]).cuda()\n",
    "        face_score = torch.cdist(valid_emb, train_emb, p=2).mean().item()\n",
    "        clip_score /= len(validation_prompt) - mis\n",
    "    \n",
    "    return face_score, clip_score, mis   # 返回面部相似性、CLIP 分数和无面部图像的数量"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71b834a7-3f78-4d6f-9432-17444e7c29bd",
   "metadata": {},
   "source": [
    "## 准备微调所需的模型和优化器\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f73b8290-6b15-4171-bc1b-403de3abe81b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 决定是否继续之前的训练\n",
    "resume = False\n",
    "\n",
    "# 准备微调所需的 tokenizer, 噪声调度器, UNet, VAE 和文本编码器\n",
    "tokenizer, noise_scheduler, unet, vae, text_encoder = prepare_lora_model(lora_config, pretrained_model_name_or_path, model_path, resume)\n",
    "\n",
    "# 准备优化器，用于更新 UNet 和 文本编码器中经过 LoRA 层调整的可训练参数\n",
    "optimizer = prepare_optimizer(unet, text_encoder, unet_learning_rate, text_encoder_learning_rate)\n",
    "\n",
    "# 设置学习率调度器，控制学习率的变化\n",
    "# get_scheduler 根据指定的调度器类型（例如线性、余弦等），为优化器提供一个动态调整学习率的策略\n",
    "# num_warmup_steps: 在初始阶段进行学习率预热的步数\n",
    "# num_training_steps: 总的训练步数\n",
    "# num_cycles: 用于余弦调度器，指定余弦衰减周期的数量\n",
    "lr_scheduler = get_scheduler(\n",
    "    lr_scheduler_name,\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=lr_warmup_steps,\n",
    "    num_training_steps=max_train_steps,\n",
    "    num_cycles=3\n",
    ")\n",
    "\n",
    "# 准备数据集\n",
    "# Text2ImageDataset 会根据指定的图像文件夹和文本标注文件夹，加载并处理数据\n",
    "# 具体：使用指定的 transform 来预处理图像数据，并使用 tokenizer 来将文本转换为 tokens\n",
    "dataset = Text2ImageDataset(\n",
    "    images_folder=images_folder,\n",
    "    captions_folder=captions_folder,\n",
    "    transform=train_transform,\n",
    "    tokenizer=tokenizer,\n",
    ")\n",
    "\n",
    "# 设置数据加载器，负责将数据集分批次加载到模型中进行训练\n",
    "# collate_fn: 自定义的批处理函数，将数据集中每个样本的图像和文本标注合并成批次\n",
    "def collate_fn(examples):\n",
    "    pixel_values = []\n",
    "    input_ids = []\n",
    "    \n",
    "    # 将每个样本中的图像 tensor 和文本标注 input_ids 分别收集到列表中\n",
    "    for tensor, input_id in examples:\n",
    "        pixel_values.append(tensor)  # 将图像张量收集到 pixel_values 列表中\n",
    "        input_ids.append(input_id)  # 将文本标注收集到 input_ids 列表中\n",
    "    \n",
    "    # 将列表中的所有图像张量堆叠成一个大的 batch tensor\n",
    "    pixel_values = torch.stack(pixel_values, dim=0).float()  # 形状为 (batch_size, channels, height, width)\n",
    "    input_ids = torch.stack(input_ids, dim=0)  # 形状为 (batch_size, max_seq_length)\n",
    "    \n",
    "    # 返回批处理后的图像和文本标注\n",
    "    return {\"pixel_values\": pixel_values, \"input_ids\": input_ids}\n",
    "\n",
    "# 使用 PyTorch DataLoader 加载数据集，设置每个批次的大小和线程数\n",
    "train_dataloader = torch.utils.data.DataLoader(\n",
    "    dataset,\n",
    "    shuffle=True,  # 打乱数据\n",
    "    collate_fn=collate_fn,  # 自定义批处理函数\n",
    "    batch_size=train_batch_size,  # 每个批次的样本数量\n",
    "    num_workers=8,  # 使用 8 个线程并行加载数据，加快数据加载速度\n",
    ")\n",
    "\n",
    "print(\"准备完成！\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61b27875-b3b2-4b03-a554-4675667e4415",
   "metadata": {},
   "source": [
    "## 开始微调\n",
    "\n",
    "最后是开始进行微调的部分：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8836a890-8e06-4aed-bd11-71df8fd5bf14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 开始微调 Stable Diffusion 模型\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "torch.backends.cuda.enable_mem_efficient_sdp(False)\n",
    "torch.backends.cuda.enable_flash_sdp(False)\n",
    "\n",
    "# 进度条显示训练进度\n",
    "progress_bar = tqdm(\n",
    "    range(0, max_train_steps),\n",
    "    initial=0,\n",
    "    desc=\"步骤\",\n",
    ")\n",
    "\n",
    "# 初始化全局步数和其他控制变量\n",
    "global_step = 0\n",
    "num_epochs = math.ceil(max_train_steps / len(train_dataloader))  # 根据最大训练步数计算总共的 epoch 数量\n",
    "validation_step = int(max_train_steps * validation_step_ratio)  # 根据设置的验证步数比例计算每隔多少步进行验证\n",
    "best_face_score = float(\"inf\")  # 初始化为正无穷大，存储最佳面部相似度分数\n",
    "\n",
    "# 训练循环，遍历所有的 epoch\n",
    "for epoch in range(num_epochs):\n",
    "    # 将模型切换到训练模式\n",
    "    unet.train()\n",
    "    text_encoder.train()\n",
    "    \n",
    "    # 遍历数据加载器，进行逐步训练\n",
    "    for step, batch in enumerate(train_dataloader):\n",
    "        if global_step >= max_train_steps:\n",
    "            break  # 当达到最大训练步数时，终止训练\n",
    "        \n",
    "        # 编码图像为潜在表示（latent），通过 VAE 对图像进行编码\n",
    "        latents = vae.encode(batch[\"pixel_values\"].to(DEVICE, dtype=weight_dtype)).latent_dist.sample()\n",
    "        latents = latents * vae.config.scaling_factor  # 根据 VAE 的缩放因子调整潜在空间\n",
    "\n",
    "        # 为潜在表示添加噪声，生成带噪声的图像\n",
    "        noise = torch.randn_like(latents)  # 生成与潜在表示相同形状的随机噪声\n",
    "        bsz = latents.shape[0]  # 批次大小\n",
    "        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device).long()  # 随机选择时间步\n",
    "        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)  # 添加噪声到潜在表示\n",
    "\n",
    "        # 获取文本的嵌入表示\n",
    "        encoder_hidden_states = text_encoder(batch[\"input_ids\"].to(latents.device), return_dict=False)[0]\n",
    "        \n",
    "        # 计算目标值，基于预测类型（epsilon 或 v_prediction）\n",
    "        if noise_scheduler.config.prediction_type == \"epsilon\":\n",
    "            target = noise  # 预测噪声\n",
    "        elif noise_scheduler.config.prediction_type == \"v_prediction\":\n",
    "            target = noise_scheduler.get_velocity(latents, noise, timesteps)  # 预测速度向量\n",
    "\n",
    "        # UNet 模型预测，输入噪声潜在空间、时间步和文本嵌入\n",
    "        model_pred = unet(noisy_latents, timesteps, encoder_hidden_states, return_dict=False)[0]\n",
    "        \n",
    "        # 计算损失，基于是否设置了 snr_gamma，这个参数在最开始的部分进行了设置\n",
    "        if not snr_gamma:\n",
    "            loss = F.mse_loss(model_pred.float(), target.float(), reduction=\"mean\")  # 使用均方误差 (MSE) 作为损失\n",
    "        else:\n",
    "            # 计算信噪比 (SNR) 并根据 SNR 加权 MSE 损失\n",
    "            snr = compute_snr(noise_scheduler, timesteps)\n",
    "            mse_loss_weights = torch.stack([snr, snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0]\n",
    "            if noise_scheduler.config.prediction_type == \"epsilon\":\n",
    "                mse_loss_weights = mse_loss_weights / snr\n",
    "            elif noise_scheduler.config.prediction_type == \"v_prediction\":\n",
    "                mse_loss_weights = mse_loss_weights / (snr + 1)\n",
    "            \n",
    "            # 计算加权的 MSE 损失\n",
    "            loss = F.mse_loss(model_pred.float(), target.float(), reduction=\"none\")\n",
    "            loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights\n",
    "            loss = loss.mean()\n",
    "\n",
    "        # 反向传播\n",
    "        loss.backward()  # 计算梯度\n",
    "        optimizer.step()  # 更新模型参数\n",
    "        lr_scheduler.step()  # 更新学习率\n",
    "        optimizer.zero_grad()  # 清零梯度\n",
    "        progress_bar.update(1)  # 更新进度条\n",
    "        global_step += 1  # 更新全局步数\n",
    "\n",
    "\n",
    "        # 验证模型性能\n",
    "        if global_step % validation_step == 0 or global_step == max_train_steps:\n",
    "            # 保存当前检查点（包含 UNet 和 文本编码器的模型参数）\n",
    "            save_path = os.path.join(output_folder, f\"checkpoint-last\")\n",
    "            unet_path = os.path.join(save_path, \"unet.pt\")\n",
    "            text_encoder_path = os.path.join(save_path, \"text_encoder.pt\")\n",
    "            print(f\"正在保存检查点到 {save_path} ......\")\n",
    "            os.makedirs(save_path, exist_ok=True)\n",
    "            torch.save(unet, unet_path)\n",
    "            torch.save(text_encoder, text_encoder_path)\n",
    "            \n",
    "            # 进行验证\n",
    "            save_path = os.path.join(output_folder, f\"checkpoint-{global_step}\")\n",
    "            os.makedirs(save_path, exist_ok=True)\n",
    "            face_score, clip_score, mis = evaluate(\n",
    "                pretrained_model_name_or_path=pretrained_model_name_or_path,\n",
    "                weight_dtype=weight_dtype,\n",
    "                seed=seed,\n",
    "                unet_path=unet_path,\n",
    "                text_encoder_path=text_encoder_path,\n",
    "                validation_prompt=validation_prompt[:validation_prompt_num],  # 验证提示文本\n",
    "                output_folder=save_path,  # 保存验证结果的文件夹\n",
    "                train_emb=dataset.train_emb  # 训练数据中的面部特征嵌入\n",
    "            )\n",
    "            print(\"步骤:\", global_step, \"面部相似度评分:\", face_score, \"CLIP评分:\", clip_score, \"无面部图像数量:\", mis)\n",
    "            \n",
    "            # 如果当前面部相似度评分优于之前的最佳记录，保存最佳模型\n",
    "            if face_score < best_face_score:\n",
    "                best_face_score = face_score\n",
    "                save_path = os.path.join(output_folder, f\"checkpoint-best\")\n",
    "                unet_path = os.path.join(save_path, \"unet.pt\")\n",
    "                text_encoder_path = os.path.join(save_path, \"text_encoder.pt\")\n",
    "                os.makedirs(save_path, exist_ok=True)\n",
    "                torch.save(unet, unet_path)  # 保存最佳的 UNet 模型\n",
    "                torch.save(text_encoder, text_encoder_path)  # 保存最佳的文本编码器模型\n",
    "\n",
    "print(\"微调完成！\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdf90ba1-3511-425c-b278-a3c6da887ed5",
   "metadata": {},
   "source": [
    "## 生成图像和测试\n",
    "\n",
    "微调过程已经完成。接下来将测试模型。\n",
    "\n",
    "首先加载之前保存的微调模型的检查点，并计算面部相似度、CLIP 评分以及没有面部的图像数量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb883783-f35b-46b7-9279-2cf96fdd5d4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置用于推理的 checkpoint 路径\n",
    "checkpoint_path = os.path.join(output_folder, f\"checkpoint-last\")  # 设置推理时使用最后保存的 checkpoint\n",
    "\n",
    "# 设置路径\n",
    "unet_path = os.path.join(checkpoint_path, \"unet.pt\")\n",
    "text_encoder_path = os.path.join(checkpoint_path, \"text_encoder.pt\")\n",
    "inference_path = os.path.join(project_dir, \"inference\")\n",
    "os.makedirs(inference_path, exist_ok=True)\n",
    "\n",
    "# 获取训练图像的路径\n",
    "train_image_paths = []  # 初始化存储所有训练图像路径的列表\n",
    "for ext in IMAGE_EXTENSIONS:\n",
    "    # 通过扩展名遍历图像文件夹中的所有图像\n",
    "    train_image_paths.extend(glob.glob(f\"{images_folder}/*{ext}\"))\n",
    "train_image_paths = sorted(train_image_paths)  # 对图像路径进行排序\n",
    "\n",
    "# 初始化列表存储面部特征嵌入\n",
    "train_emb_list = []\n",
    "\n",
    "# 遍历训练图像路径并提取面部特征嵌入\n",
    "for img_path in train_image_paths:\n",
    "    # 使用 DeepFace 从每张图像中提取面部特征\n",
    "    face_representation = DeepFace.represent(\n",
    "        img_path, \n",
    "        detector_backend=\"ssd\",  # 使用 ssd 检测器\n",
    "        model_name=\"GhostFaceNet\",  # 使用 GhostFaceNet 模型\n",
    "        enforce_detection=False  # 关闭强制检测\n",
    "    )\n",
    "    \n",
    "    # 如果提取到的特征非空，则获取嵌入向量\n",
    "    if face_representation:\n",
    "        embedding = face_representation[0]['embedding']\n",
    "        train_emb_list.append(embedding)\n",
    "\n",
    "# 将所有面部嵌入转换为 tensor\n",
    "train_emb = torch.tensor(train_emb_list)\n",
    "\n",
    "# 调用 evaluate 函数评估模型的性能\n",
    "face_score, clip_score, mis = evaluate(\n",
    "    pretrained_model_name_or_path=pretrained_model_name_or_path,  # 预训练模型的名称或路径\n",
    "    weight_dtype=weight_dtype,  # 模型的权重数据类型\n",
    "    seed=seed,  # 随机数种子，确保可重复的推理结果\n",
    "    unet_path=unet_path,  # 训练保存的 UNet 模型的权重文件路径\n",
    "    text_encoder_path=text_encoder_path,  # 训练保存的文本编码器模型的权重文件路径\n",
    "    validation_prompt=validation_prompt,  # 验证时使用的文本提示，输入模型生成图像\n",
    "    output_folder=inference_path,  # 推理时生成的图像保存到指定的文件夹\n",
    "    train_emb=train_emb  # 使用的训练图像的面部特征嵌入，用于评估面部相似度\n",
    ")\n",
    "\n",
    "# 打印评估结果\n",
    "print(f\"面部相似度评分 (平均欧氏距离): {face_score:.4f} (越低越好，表示生成图像与训练图像更相似)\")\n",
    "print(f\"CLIP 评分 (平均相似度): {clip_score:.4f} (越高越好，表示生成图像与文本提示的相关性更强)\")\n",
    "print(f\"无面部图像数量: {mis} (无法检测到面部的生成图像数量)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b2cea9f-314a-49ed-9b7c-20a590b298bb",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}