Demos/14b. 尝试使用 LoRA 微调 Stable Diffusion 模型 - 精简版.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7e3e66d1-137d-4acd-8b3c-1c8047c5fd9a",
   "metadata": {},
   "source": [
    "# b. 尝试使用 LoRA 微调 Stable Diffusion 模型（文生图）- 精简版\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",
    "\n",
    "注意，[版本 a](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Demos/14a.%20尝试使用%20LoRA%20微调%20Stable%20Diffusion%20模型.ipynb) 对于保存和加载使用的方法与版本 b 不同，精简版使用 PEFT 库直接完成，与之前的知识[《14. PEFT：在大模型中快速应用 LoRA》](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Guide/14.%20PEFT：在大模型中快速应用%20LoRA.md)对齐。\n",
    "\n",
    "另外，我添加了一些 Emoji 在输出中，为之后的 Toy 脚本做准备 :)\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": "33364619-4768-4837-ad1c-2b72edc0f957",
   "metadata": {},
   "source": [
    "## 安装必要的库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b4e7d63-f337-4b40-840e-77678f29ef1c",
   "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": "09608f50-24e2-4860-b68a-320b25207f98",
   "metadata": {},
   "source": [
    "## 导入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63e43e2c-c6f9-4981-8ac8-a8f89de845e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ========== 标准库模块 ==========\n",
    "import os\n",
    "import math\n",
    "import glob\n",
    "import shutil\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",
    "\n",
    "# Transformers (Hugging Face)\n",
    "from transformers import CLIPTextModel, CLIPTokenizer, CLIPModel, CLIPProcessor\n",
    "\n",
    "# Diffusers (Hugging Face)\n",
    "from diffusers import (\n",
    "    AutoencoderKL,\n",
    "    DDPMScheduler,\n",
    "    UNet2DConditionModel,\n",
    "    DiffusionPipeline\n",
    ")\n",
    "from diffusers.optimization import get_scheduler\n",
    "from diffusers.training_utils import compute_snr\n",
    "\n",
    "# ========== LoRA 模型库 ==========\n",
    "from peft import LoraConfig, get_peft_model, PeftModel\n",
    "\n",
    "# ========== 面部检测库 ==========\n",
    "from deepface import DeepFace\n",
    "\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbbf9948-94b5-463b-a8e7-f48d999762e3",
   "metadata": {},
   "source": [
    "## 准备项目\n",
    "\n",
    "### 设置路径\n",
    "这里的参数不需要修改，如果你想自定义文件夹，后续我会出一个更通用的脚本供你学习。\n",
    "\n",
    "当你看见✅时，代表项目已经准备好。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ae237bf-6d30-4efa-b6a9-7a6d22827238",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 项目名称和数据集名称\n",
    "project_name = \"Brad\"\n",
    "dataset_name = \"Brad\"\n",
    "\n",
    "# 根目录和主要目录\n",
    "root_dir = \"./\"  # 当前目录\n",
    "main_dir = os.path.join(root_dir, \"SD\")  # 主目录\n",
    "\n",
    "# 项目目录\n",
    "project_dir = os.path.join(main_dir, project_name)  # 项目目录\n",
    "\n",
    "# 数据集和模型路径\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",
    "output_folder = os.path.join(project_dir, \"logs\")  # 存放 model checkpoints 和 validation 的文件夹\n",
    "\n",
    "# prompt 文件路径\n",
    "validation_prompt_name = \"validation_prompt.txt\"\n",
    "validation_prompt_path = os.path.join(prompts_folder, validation_prompt_name)\n",
    "\n",
    "# 模型检查点路径\n",
    "model_path = os.path.join(project_dir, \"logs\", \"checkpoint-last\")\n",
    "\n",
    "# 其他路径设置\n",
    "zip_file = os.path.join(\"./\", \"data/14/Datasets.zip\")\n",
    "inference_path = os.path.join(project_dir, \"inference\")  # 保存推理结果的文件夹\n",
    "\n",
    "os.makedirs(images_folder, exist_ok=True)\n",
    "os.makedirs(prompts_folder, exist_ok=True)\n",
    "os.makedirs(output_folder, exist_ok=True)\n",
    "os.makedirs(inference_path, exist_ok=True)\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",
    "    print(f\"✅ 项目 {project_name} 已准备好！\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cdaa70a-5e3a-4b67-8743-3780d360d1a4",
   "metadata": {},
   "source": [
    "## 定义一些有用的函数和类"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0040ab7-0da7-4458-b627-af55dd701e67",
   "metadata": {},
   "source": [
    "### 数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36c7d433-bb97-4cc1-a3b6-7c4c2ea22832",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 图片后缀，不用关注\n",
    "IMAGE_EXTENSIONS = [\".png\", \".jpg\", \".jpeg\", \".webp\", \".bmp\", \".PNG\", \".JPG\", \".JPEG\", \".WEBP\", \".BMP\"]\n",
    "\n",
    "class Text2ImageDataset(torch.utils.data.Dataset):\n",
    "    \"\"\"\n",
    "    (1) 目标:\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(os.path.join(images_folder, f\"*{ext}\")))\n",
    "        self.image_paths = sorted(self.image_paths)\n",
    "\n",
    "        # 加载对应的文本标注，依次读取每个文本文件中的内容\n",
    "        caption_paths = sorted(glob.glob(os.path.join(captions_folder, \"*.txt\")))\n",
    "        captions = []\n",
    "        for p in caption_paths:\n",
    "            with open(p, \"r\", encoding=\"utf-8\") as f:\n",
    "                captions.append(f.readline().strip())\n",
    "\n",
    "        # 确保图像和文本标注数量一致\n",
    "        if len(captions) != len(self.image_paths):\n",
    "            raise ValueError(\"图像数量与文本标注数量不一致，请检查数据集。\")\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",
    "            # 返回一个全零的张量和空的输入 ID 以避免崩溃\n",
    "            tensor = torch.zeros((3, resolution, resolution))\n",
    "            input_id = torch.zeros_like(input_id)\n",
    "        \n",
    "        return tensor, input_id  # 返回处理后的图像和相应的文本标注\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.image_paths)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c18b966-f7e0-415d-b49a-70773a33fb7b",
   "metadata": {},
   "source": [
    "### 加载 LoRA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e413064d-c177-4857-944f-0030ff51089c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_lora_model(lora_config, pretrained_model_name_or_path, model_path=None, resume=False, merge_lora=False):\n",
    "    \"\"\"\n",
    "    (1) 目标:\n",
    "        - 加载完整的 Stable Diffusion 模型，包括 LoRA 层，并根据需要合并 LoRA 权重。这包括 Tokenizer、噪声调度器、UNet、VAE 和文本编码器。\n",
    "\n",
    "    (2) 参数:\n",
    "        - lora_config: LoraConfig, LoRA 的配置对象\n",
    "        - pretrained_model_name_or_path: str, Hugging Face 上的模型名称或路径\n",
    "        - model_path: str, 预训练模型的路径\n",
    "        - resume: bool, 是否从上一次训练中恢复\n",
    "        - merge_lora: bool, 是否在推理时合并 LoRA 权重\n",
    "\n",
    "    (3) 返回:\n",
    "        - tokenizer: CLIPTokenizer\n",
    "        - noise_scheduler: DDPMScheduler\n",
    "        - unet: UNet2DConditionModel\n",
    "        - vae: AutoencoderKL\n",
    "        - text_encoder: CLIPTextModel\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",
    "    # 如果设置为继续训练，则加载上一次的模型权重\n",
    "    if resume:\n",
    "        if model_path is None or not os.path.exists(model_path):\n",
    "            raise ValueError(\"当 resume 设置为 True 时，必须提供有效的 model_path\")\n",
    "        # 使用 PEFT 的 from_pretrained 方法加载 LoRA 模型\n",
    "        text_encoder = PeftModel.from_pretrained(text_encoder, os.path.join(model_path, \"text_encoder\"))\n",
    "        unet = PeftModel.from_pretrained(unet, os.path.join(model_path, \"unet\"))\n",
    "\n",
    "        # 确保 UNet 的可训练参数的 requires_grad 为 True\n",
    "        for param in unet.parameters():\n",
    "            if param.requires_grad is False:\n",
    "                param.requires_grad = True\n",
    "        \n",
    "        # 确保文本编码器的可训练参数的 requires_grad 为 True\n",
    "        for param in text_encoder.parameters():\n",
    "            if param.requires_grad is False:\n",
    "                param.requires_grad = True\n",
    "                \n",
    "        print(f\"✅ 已从 {model_path} 恢复模型权重\")\n",
    "\n",
    "    else:\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",
    "        print(\"📊 Text Encoder 可训练参数:\")\n",
    "        text_encoder.print_trainable_parameters()\n",
    "        print(\"📊 UNet 可训练参数:\")\n",
    "        unet.print_trainable_parameters()\n",
    "    \n",
    "    if merge_lora:\n",
    "        # 合并 LoRA 权重到基础模型，仅在推理时调用\n",
    "        text_encoder = text_encoder.merge_and_unload()\n",
    "        unet = unet.merge_and_unload()\n",
    "\n",
    "        # 切换为评估模式\n",
    "        text_encoder.eval()\n",
    "        unet.eval()\n",
    "\n",
    "    # 冻结 VAE 参数\n",
    "    vae.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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19f62fe6-8b07-4210-84ed-7c4f4368c179",
   "metadata": {},
   "source": [
    "### 准备优化器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "491cd8f9-678f-464b-a386-3dd3eca3b70e",
   "metadata": {},
   "outputs": [],
   "source": [
    "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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "574c7c21-5dc3-484a-8fa1-5689be53684f",
   "metadata": {},
   "source": [
    "### 定义 collate_fn 函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25bec1c8-4535-4cd4-9c97-e571da1dc861",
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_fn(examples):\n",
    "    pixel_values = []\n",
    "    input_ids = []\n",
    "    \n",
    "    for tensor, input_id in examples:\n",
    "        pixel_values.append(tensor)\n",
    "        input_ids.append(input_id)\n",
    "    \n",
    "    pixel_values = torch.stack(pixel_values, dim=0).float()\n",
    "    input_ids = torch.stack(input_ids, dim=0)\n",
    "    \n",
    "    # 如果你喜欢列表推导式的话，使用下面的方法\n",
    "    #pixel_values = torch.stack([example[0] for example in examples], dim=0).float()\n",
    "    #input_ids = torch.stack([example[1] for example in examples], dim=0)\n",
    "    \n",
    "    return {\"pixel_values\": pixel_values, \"input_ids\": input_ids}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "838f6274-fb4b-4b7c-9936-06252f45f42e",
   "metadata": {},
   "source": [
    "## 参数设置\n",
    "\n",
    "### 1. 设备配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfe24c04-4638-466a-a909-8e44e7b43560",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设备配置\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# For Mac M1, M2...\n",
    "# DEVICE = torch.device(\"mps\" if torch.backends.mps.is_available() else (\"cuda\" if torch.cuda.is_available() else \"cpu\"))\n",
    "\n",
    "print(f\"🖥 当前使用的设备: {DEVICE}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "561a41d7-f426-406c-9fbd-9ba3afe5ed7d",
   "metadata": {},
   "source": [
    "### 2. 图像预处理与数据增强"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3967ac77-7b2d-410d-975e-af9433145e6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 训练图像的分辨率\n",
    "resolution = 512\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",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3c9aec3-6c15-4e98-bf83-3b423c5dab40",
   "metadata": {},
   "source": [
    "### 3. 模型与训练参数配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a050558-444c-4d7c-bccd-4b001829fad7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 训练相关参数\n",
    "train_batch_size = 2  # 训练批次大小，即每次训练中处理的样本数量\n",
    "weight_dtype = torch.bfloat16  # 权重数据类型，使用 bfloat16 以节省内存并加快计算速度\n",
    "snr_gamma = 5  # SNR 参数，用于信噪比加权损失的调节系数\n",
    "\n",
    "# 设置随机数种子以确保可重复性\n",
    "seed = 1126  # 随机数种子\n",
    "torch.manual_seed(seed)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    \n",
    "# Stable Diffusion LoRA 的微调参数\n",
    "\n",
    "# 优化器参数\n",
    "unet_learning_rate = 1e-4  # UNet 的学习率，控制 UNet 参数更新的步长\n",
    "text_encoder_learning_rate = 1e-4  # 文本编码器的学习率，控制文本嵌入层的参数更新步长\n",
    "\n",
    "# 学习率调度器参数\n",
    "lr_scheduler_name = \"cosine_with_restarts\"  # 设置学习率调度器为 Cosine annealing with restarts，逐渐减少学习率并定期重启\n",
    "lr_warmup_steps = 100  # 学习率预热步数，在最初的 100 步中逐渐增加学习率到最大值\n",
    "max_train_steps = 200  # 总训练步数，决定了整个训练过程的迭代次数\n",
    "num_cycles = 3  # Cosine 调度器的周期数量，在训练期间会重复 3 次学习率周期性递减并重启\n",
    "\n",
    "# 预训练的 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": "1fdeedb5-b8ce-4590-9dcc-559729edebd2",
   "metadata": {},
   "source": [
    "## 微调前的准备"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e46cd03-d4bd-4094-88f3-036fcf16b872",
   "metadata": {},
   "source": [
    "### 1. 数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d335cfc-007d-4d47-878b-432cf8d1f4ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 初始化 tokenizer，用于加载数据集\n",
    "tokenizer = CLIPTokenizer.from_pretrained(\n",
    "    pretrained_model_name_or_path,\n",
    "    subfolder=\"tokenizer\"\n",
    ")\n",
    "\n",
    "# 准备数据集\n",
    "dataset = Text2ImageDataset(\n",
    "    images_folder=images_folder,\n",
    "    captions_folder=captions_folder,\n",
    "    transform=train_transform,\n",
    "    tokenizer=tokenizer,\n",
    ")\n",
    "\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,\n",
    ")\n",
    "\n",
    "print(\"✅ 数据集准备完成！\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db606893-440b-431b-bdcc-f81695d3d679",
   "metadata": {},
   "source": [
    "### 2. 模型和优化器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d00a8bf-4103-4e22-93e5-14e749a84376",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 准备模型\n",
    "tokenizer, noise_scheduler, unet, vae, text_encoder = prepare_lora_model(\n",
    "    lora_config,\n",
    "    pretrained_model_name_or_path,\n",
    "    model_path,\n",
    "    resume=False,  # 根据需要设置为 True 以从 checkpoint 恢复\n",
    "    merge_lora=False  # 是否合并 LoRA 权重\n",
    ")\n",
    "\n",
    "# 准备优化器\n",
    "optimizer = prepare_optimizer(\n",
    "    unet, \n",
    "    text_encoder, \n",
    "    unet_learning_rate=unet_learning_rate, \n",
    "    text_encoder_learning_rate=text_encoder_learning_rate\n",
    ")\n",
    "\n",
    "# 设置学习率调度器\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=num_cycles\n",
    ")\n",
    "\n",
    "print(\"✅ 模型和优化器准备完成！可以开始训练。\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8f3f649-bc32-441a-93b0-4fc784e0fa3b",
   "metadata": {},
   "source": [
    "## 开始微调"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4117f19a-fd07-4828-9c0a-1b91ae658986",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 禁用并行化，避免警告\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "# 初始化\n",
    "global_step = 0\n",
    "best_face_score = float(\"inf\")  # 初始化为正无穷大，存储最佳面部相似度分数\n",
    "\n",
    "# 进度条显示训练进度\n",
    "progress_bar = tqdm(\n",
    "    range(max_train_steps),  # 根据 num_training_steps 设置\n",
    "    desc=\"训练步骤\",\n",
    ")\n",
    "\n",
    "# 训练循环\n",
    "for epoch in range(math.ceil(max_train_steps / len(train_dataloader))):\n",
    "    # 如果你想在训练中增加评估，那在循环中增加 train() 是有必要的\n",
    "    unet.train()\n",
    "    text_encoder.train()\n",
    "    \n",
    "    for step, batch in enumerate(train_dataloader):\n",
    "        if global_step >= max_train_steps:\n",
    "            break\n",
    "        \n",
    "        # 编码图像为潜在表示（latent）\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",
    "        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (latents.shape[0],), device=DEVICE).long()\n",
    "        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)\n",
    "\n",
    "        # 获取文本的嵌入表示\n",
    "        encoder_hidden_states = text_encoder(batch[\"input_ids\"].to(DEVICE))[0]\n",
    "\n",
    "        # 计算目标值\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)[0]\n",
    "\n",
    "        # 计算损失\n",
    "        if not snr_gamma:\n",
    "            loss = F.mse_loss(model_pred.float(), target.float(), reduction=\"mean\")\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",
    "        if global_step % 100 == 0 or global_step == max_train_steps:\n",
    "            print(f\"🔥 步骤 {global_step}, 损失: {loss.item()}\")\n",
    "\n",
    "        # 保存中间检查点，当前简单设置为每 500 步保存一次\n",
    "        if global_step % 500 == 0:\n",
    "            save_path = os.path.join(output_folder, f\"checkpoint-{global_step}\")\n",
    "            os.makedirs(save_path, exist_ok=True)\n",
    "\n",
    "            # 使用 save_pretrained 保存 PeftModel\n",
    "            unet.save_pretrained(os.path.join(save_path, \"unet\"))\n",
    "            text_encoder.save_pretrained(os.path.join(save_path, \"text_encoder\"))\n",
    "            print(f\"💾 已保存中间模型到 {save_path}\")\n",
    "\n",
    "# 保存最终模型到 checkpoint-last\n",
    "save_path = os.path.join(output_folder, \"checkpoint-last\")\n",
    "os.makedirs(save_path, exist_ok=True)\n",
    "unet.save_pretrained(os.path.join(save_path, \"unet\"))\n",
    "text_encoder.save_pretrained(os.path.join(save_path, \"text_encoder\"))\n",
    "print(f\"💾 已保存最终模型到 {save_path}\")\n",
    "\n",
    "print(\"🎉 微调完成！\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d83c69b-bb4a-430b-8580-f897874a69eb",
   "metadata": {},
   "source": [
    "## 生成图像和测试\n",
    "\n",
    "### 加载用于验证的 prompts\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53cce6c3-d106-4317-8f28-4dc7be4c7b9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_validation_prompts(validation_prompt_path):\n",
    "    \"\"\"\n",
    "    (1) 目标:\n",
    "        - 加载验证提示文本。\n",
    "\n",
    "    (2) 参数:\n",
    "        - validation_prompt_path: str, 验证提示文件的路径\n",
    "\n",
    "    (3) 返回:\n",
    "        - validation_prompt: list, 验证提示的字符串列表，每一行就是一个prompt\n",
    "    \"\"\"\n",
    "    with open(validation_prompt_path, \"r\", encoding=\"utf-8\") as f:\n",
    "        validation_prompt = [line.strip() for line in f.readlines()]\n",
    "    return validation_prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0908312c-d673-4f23-9039-e42b01f05f45",
   "metadata": {},
   "source": [
    "### 定义生成图像的函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b41e53b-7222-4cff-97de-5d8a7bef482e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_images(pipeline, prompts, num_inference_steps=50, guidance_scale=7.5, output_folder=\"inference\", generator=None):\n",
    "    \"\"\"\n",
    "    (1) 目标:\n",
    "        - 使用 DiffusionPipeline 生成图像，保存到指定文件夹并返回生成的图像列表。\n",
    "\n",
    "    (2) 参数:\n",
    "        - pipeline: DiffusionPipeline, 已加载并配置好的 Pipeline\n",
    "        - prompts: list, 文本提示列表\n",
    "        - num_inference_steps: int, 推理步骤数，越高图像质量越好，但推理时间也会增加\n",
    "        - guidance_scale: float, 决定文本提示对生成图像的影响程度\n",
    "        - output_folder: str, 保存生成图像的文件夹路径\n",
    "        - generator: torch.Generator, 控制生成随机数的种子，确保图像生成的一致性。如果不提供，生成的图像每次可能不同\n",
    "\n",
    "    (3) 返回:\n",
    "        - 生成的图像列表，同时图像也会保存到指定文件夹。\n",
    "    \"\"\"\n",
    "    print(\"🎨 正在生成图像...\")\n",
    "    os.makedirs(output_folder, exist_ok=True)\n",
    "    generated_images = []\n",
    "    \n",
    "    for i, prompt in enumerate(tqdm(prompts, desc=\"生成图像中\")):\n",
    "        # 使用 pipeline 生成图像\n",
    "        image = pipeline(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator).images[0]\n",
    "        \n",
    "        # 保存图像到指定文件夹\n",
    "        save_file = os.path.join(output_folder, f\"generated_{i+1}.png\")\n",
    "        image.save(save_file)\n",
    "        \n",
    "        # 将图像保存到列表中，稍后返回\n",
    "        generated_images.append(image)\n",
    "    \n",
    "    print(f\"✅ 已生成并保存 {len(prompts)} 张图像到 {output_folder}\")\n",
    "    \n",
    "    return generated_images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f71ff0b1-f183-4f45-b333-405577ce5958",
   "metadata": {},
   "source": [
    "### 定义评估函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcc5cfa8-7186-431f-923c-ca40d7a120db",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(lora_config):\n",
    "    \"\"\"\n",
    "    加载模型、生成图像并评估。\n",
    "    \n",
    "    主要步骤：\n",
    "    1. 加载验证文本提示（prompts）用于生成图像。\n",
    "    2. 加载和准备 LoRA 微调后的模型。\n",
    "    3. 使用 DiffusionPipeline 生成图像。\n",
    "    4. 评估生成图像的人脸相似度、CLIP 评分和无面部图像数量。\n",
    "    5. 打印评估结果。\n",
    "    \"\"\"\n",
    "    print(\"📂 加载验证提示...\")\n",
    "    validation_prompts = load_validation_prompts(validation_prompt_path)\n",
    "\n",
    "    print(\"🔧 准备 LoRA 模型...\")\n",
    "    # 准备 LoRA 模型（用于推理，合并权重）\n",
    "    tokenizer, noise_scheduler, unet, vae, text_encoder = prepare_lora_model(\n",
    "        lora_config,\n",
    "        pretrained_model_name_or_path,\n",
    "        model_path=model_path,\n",
    "        resume=True,  # 从检查点恢复\n",
    "        merge_lora=True  # 合并 LoRA 权重\n",
    "    )\n",
    "\n",
    "    # 创建 DiffusionPipeline 并更新其组件\n",
    "    print(\"🔄 创建 DiffusionPipeline...\")\n",
    "    pipeline = DiffusionPipeline.from_pretrained(\n",
    "        pretrained_model_name_or_path,\n",
    "        unet=unet,  # 传递基础模型\n",
    "        text_encoder=text_encoder,  # 传递基础模型\n",
    "        torch_dtype=weight_dtype,\n",
    "        safety_checker=None,\n",
    "    )\n",
    "    pipeline = pipeline.to(DEVICE)\n",
    "\n",
    "    # 加载 CLIP 模型和处理器\n",
    "    print(\"🎯 加载 CLIP 模型...\")\n",
    "    clip_model_name = \"openai/clip-vit-base-patch32\"\n",
    "    clip_model = CLIPModel.from_pretrained(clip_model_name).to(DEVICE)\n",
    "    clip_processor = CLIPProcessor.from_pretrained(clip_model_name)\n",
    "\n",
    "    # CLIP 模型设置为评估模式\n",
    "    clip_model.eval()\n",
    "\n",
    "    # 设置随机数种子\n",
    "    generator = torch.Generator(device=DEVICE)\n",
    "    generator.manual_seed(seed)\n",
    "\n",
    "    # 加载训练图像的面部嵌入\n",
    "    print(\"📂 加载训练图像的面部嵌入...\")\n",
    "    train_image_paths = sorted([\n",
    "        p for p in glob.glob(os.path.join(images_folder, \"*\")) \n",
    "        if any(p.endswith(ext) for ext in IMAGE_EXTENSIONS)\n",
    "    ])\n",
    "    train_emb_list = []\n",
    "    for img_path in tqdm(train_image_paths, desc=\"提取训练图像面部嵌入\"):\n",
    "        face_representation = DeepFace.represent(\n",
    "            img_path, \n",
    "            detector_backend=\"ssd\",\n",
    "            model_name=\"GhostFaceNet\",\n",
    "            enforce_detection=False\n",
    "        )\n",
    "        if face_representation:\n",
    "            embedding = face_representation[0]['embedding']\n",
    "            train_emb_list.append(embedding)\n",
    "\n",
    "    if len(train_emb_list) == 0:\n",
    "        print(\"⚠️ 未能提取到任何训练图像的面部嵌入。\")\n",
    "        train_emb = torch.tensor([]).to(DEVICE)\n",
    "    else:\n",
    "        train_emb = torch.tensor(train_emb_list).to(DEVICE)\n",
    "\n",
    "    # 生成图像\n",
    "    generated_images = generate_images(\n",
    "        pipeline=pipeline,\n",
    "        prompts=validation_prompts,\n",
    "        num_inference_steps=30,\n",
    "        guidance_scale=7.5,\n",
    "        output_folder=inference_path,\n",
    "        generator=generator\n",
    "    )\n",
    "\n",
    "    # 评估生成的图像，mis记录无法检测到面部的图像数量\n",
    "    face_score, clip_score, mis = 0, 0, 0  # 初始化评估分数和计数\n",
    "    valid_emb = []\n",
    "    print(\"📊 正在计算评估分数...\")\n",
    "\n",
    "    for i, image in enumerate(tqdm(generated_images, desc=\"评估图像中\")):\n",
    "        # 使用 DeepFace 检测面部特征\n",
    "        opencvImage = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)\n",
    "        emb = DeepFace.represent(\n",
    "            opencvImage,\n",
    "            detector_backend=\"ssd\",\n",
    "            model_name=\"GhostFaceNet\",\n",
    "            enforce_detection=False,\n",
    "        )\n",
    "        if not emb or emb[0].get('face_confidence', 0) == 0:\n",
    "            mis += 1  # 无法检测到面部的图像数量\n",
    "            continue\n",
    "\n",
    "        # 计算 CLIP 分数\n",
    "        current_prompt = validation_prompts[i]\n",
    "        inputs = clip_processor(text=current_prompt, 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",
    "\n",
    "        # 收集有效的面部嵌入\n",
    "        valid_emb.append(emb[0]['embedding'])\n",
    "\n",
    "    # 如果没有有效的面部嵌入，则返回默认分数\n",
    "    if len(valid_emb) == 0:\n",
    "        print(\"⚠️ 无法检测到面部嵌入！\")\n",
    "        return 0, 0, mis\n",
    "\n",
    "    # 计算面部相似度分数（使用欧氏距离）\n",
    "    valid_emb = torch.tensor(valid_emb).to(DEVICE)\n",
    "    valid_emb = valid_emb / valid_emb.norm(p=2, dim=-1, keepdim=True)\n",
    "    train_emb = train_emb / train_emb.norm(p=2, dim=-1, keepdim=True)\n",
    "    face_distance = torch.cdist(valid_emb, train_emb, p=2).mean().item()\n",
    "    face_score = face_distance  # 平均欧氏距离作为面部相似性分数\n",
    "    clip_score /= (len(validation_prompts) - mis) if (len(validation_prompts) - mis) > 0 else 1\n",
    "    print(\"📈 评估完成！\")\n",
    "\n",
    "    # 打印评估结果\n",
    "    print(f\"✅ 面部相似度评分 (平均欧氏距离): {face_score:.4f} (越低越好，表示生成图像与训练图像更相似)\")\n",
    "    print(f\"✅ CLIP 评分 (平均相似度): {clip_score:.4f} (越高越好，表示生成图像与文本提示的相关性更强)\")\n",
    "    print(f\"✅ 无面部图像数量: {mis} (无法检测到面部的生成图像数量)\")\n",
    "\n",
    "# 调用函数执行\n",
    "evaluate(lora_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21673bd9-ac20-47c6-9195-456fd05c7df5",
   "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
}