{
 "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",
    "这里还有一个简单的 [🎡 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/1E7kF00jcjrJMax5iP2DBD86PGysRT8Ux?usp=sharing)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33364619-4768-4837-ad1c-2b72edc0f957",
   "metadata": {},
   "source": [
    "## 安装必要的库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3b4e7d63-f337-4b40-840e-77678f29ef1c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: http://mirrors.aliyun.com/pypi/simple\n",
      "Requirement already satisfied: transformers==4.56.2 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (4.56.2)\n",
      "Requirement already satisfied: diffusers==0.37.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (0.37.1)\n",
      "Requirement already satisfied: peft==0.17.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (0.17.1)\n",
      "Requirement already satisfied: accelerate==1.13.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (1.13.0)\n",
      "Requirement already satisfied: safetensors==0.7.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (0.7.0)\n",
      "Requirement already satisfied: opencv-python in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (4.13.0.92)\n",
      "Requirement already satisfied: requests in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (2.33.1)\n",
      "Requirement already satisfied: deepface==0.0.100 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (0.0.100)\n",
      "Requirement already satisfied: tf-keras in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (2.21.0)\n",
      "Requirement already satisfied: filelock in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from transformers==4.56.2) (3.29.0)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.34.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from transformers==4.56.2) (0.36.2)\n",
      "Requirement already satisfied: numpy>=1.17 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from transformers==4.56.2) (2.4.4)\n",
      "Requirement already satisfied: packaging>=20.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from transformers==4.56.2) (26.1)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from transformers==4.56.2) (6.0.3)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from transformers==4.56.2) (2026.4.4)\n",
      "Requirement already satisfied: tokenizers<=0.23.0,>=0.22.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from transformers==4.56.2) (0.22.2)\n",
      "Requirement already satisfied: tqdm>=4.27 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from transformers==4.56.2) (4.67.3)\n",
      "Requirement already satisfied: importlib-metadata in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from diffusers==0.37.1) (9.0.0)\n",
      "Requirement already satisfied: httpx<1.0.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from diffusers==0.37.1) (0.28.1)\n",
      "Requirement already satisfied: Pillow in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from diffusers==0.37.1) (10.2.0)\n",
      "Requirement already satisfied: psutil in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from peft==0.17.1) (7.2.2)\n",
      "Requirement already satisfied: torch>=1.13.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from peft==0.17.1) (2.7.1)\n",
      "Requirement already satisfied: pandas>=0.23.4 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from deepface==0.0.100) (2.3.3)\n",
      "Requirement already satisfied: gdown>=3.10.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from deepface==0.0.100) (6.0.0)\n",
      "Requirement already satisfied: tensorflow>=1.9.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from deepface==0.0.100) (2.21.0)\n",
      "Requirement already satisfied: keras>=2.2.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from deepface==0.0.100) (3.14.1)\n",
      "Requirement already satisfied: Flask>=1.1.2 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from deepface==0.0.100) (3.1.3)\n",
      "Requirement already satisfied: flask-cors>=4.0.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from deepface==0.0.100) (6.0.2)\n",
      "Requirement already satisfied: mtcnn>=0.1.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from deepface==0.0.100) (1.0.0)\n",
      "Requirement already satisfied: retina-face>=0.0.14 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from deepface==0.0.100) (0.0.17)\n",
      "Requirement already satisfied: fire>=0.4.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from deepface==0.0.100) (0.7.1)\n",
      "Requirement already satisfied: gunicorn>=20.1.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from deepface==0.0.100) (26.0.0)\n",
      "Requirement already satisfied: lightphe>=0.0.15 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from deepface==0.0.100) (0.0.25)\n",
      "Requirement already satisfied: lightdsa>=0.0.3 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from deepface==0.0.100) (0.0.3)\n",
      "Requirement already satisfied: python-dotenv>=1.0.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from deepface==0.0.100) (1.2.2)\n",
      "Requirement already satisfied: anyio in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from httpx<1.0.0->diffusers==0.37.1) (4.13.0)\n",
      "Requirement already satisfied: certifi in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from httpx<1.0.0->diffusers==0.37.1) (2026.2.25)\n",
      "Requirement already satisfied: httpcore==1.* in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from httpx<1.0.0->diffusers==0.37.1) (1.0.9)\n",
      "Requirement already satisfied: idna in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from httpx<1.0.0->diffusers==0.37.1) (3.11)\n",
      "Requirement already satisfied: h11>=0.16 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from httpcore==1.*->httpx<1.0.0->diffusers==0.37.1) (0.16.0)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from huggingface-hub<1.0,>=0.34.0->transformers==4.56.2) (2025.3.0)\n",
      "Requirement already satisfied: hf-xet<2.0.0,>=1.1.3 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from huggingface-hub<1.0,>=0.34.0->transformers==4.56.2) (1.4.3)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from huggingface-hub<1.0,>=0.34.0->transformers==4.56.2) (4.15.0)\n",
      "Requirement already satisfied: charset_normalizer<4,>=2 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from requests) (3.4.7)\n",
      "Requirement already satisfied: urllib3<3,>=1.26 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from requests) (2.6.3)\n",
      "Requirement already satisfied: absl-py>=1.0.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (2.4.0)\n",
      "Requirement already satisfied: astunparse>=1.6.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (1.6.3)\n",
      "Requirement already satisfied: flatbuffers>=25.9.23 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (25.12.19)\n",
      "Requirement already satisfied: gast!=0.5.0,!=0.5.1,!=0.5.2,>=0.2.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (0.7.0)\n",
      "Requirement already satisfied: google_pasta>=0.1.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (0.2.0)\n",
      "Requirement already satisfied: libclang>=13.0.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (18.1.1)\n",
      "Requirement already satisfied: opt_einsum>=2.3.2 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (3.4.0)\n",
      "Requirement already satisfied: protobuf<8.0.0,>=6.31.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (7.34.1)\n",
      "Requirement already satisfied: setuptools in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (82.0.1)\n",
      "Requirement already satisfied: six>=1.12.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (1.17.0)\n",
      "Requirement already satisfied: termcolor>=1.1.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (3.3.0)\n",
      "Requirement already satisfied: wrapt>=1.11.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (2.1.2)\n",
      "Requirement already satisfied: grpcio<2.0,>=1.24.3 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (1.80.0)\n",
      "Requirement already satisfied: h5py<3.15.0,>=3.11.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (3.14.0)\n",
      "Requirement already satisfied: ml_dtypes<1.0.0,>=0.5.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from tensorflow>=1.9.0->deepface==0.0.100) (0.5.4)\n",
      "Requirement already satisfied: wheel<1.0,>=0.23.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from astunparse>=1.6.0->tensorflow>=1.9.0->deepface==0.0.100) (0.47.0)\n",
      "Requirement already satisfied: blinker>=1.9.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from Flask>=1.1.2->deepface==0.0.100) (1.9.0)\n",
      "Requirement already satisfied: click>=8.1.3 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from Flask>=1.1.2->deepface==0.0.100) (8.3.2)\n",
      "Requirement already satisfied: itsdangerous>=2.2.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from Flask>=1.1.2->deepface==0.0.100) (2.2.0)\n",
      "Requirement already satisfied: jinja2>=3.1.2 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from Flask>=1.1.2->deepface==0.0.100) (3.1.6)\n",
      "Requirement already satisfied: markupsafe>=2.1.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from Flask>=1.1.2->deepface==0.0.100) (3.0.3)\n",
      "Requirement already satisfied: werkzeug>=3.1.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from Flask>=1.1.2->deepface==0.0.100) (3.1.8)\n",
      "Requirement already satisfied: beautifulsoup4 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from gdown>=3.10.1->deepface==0.0.100) (4.14.3)\n",
      "Requirement already satisfied: rich in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from keras>=2.2.0->deepface==0.0.100) (15.0.0)\n",
      "Requirement already satisfied: namex in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from keras>=2.2.0->deepface==0.0.100) (0.1.0)\n",
      "Requirement already satisfied: optree in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from keras>=2.2.0->deepface==0.0.100) (0.19.1)\n",
      "Requirement already satisfied: lightecc in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from lightdsa>=0.0.3->deepface==0.0.100) (0.0.6)\n",
      "Requirement already satisfied: sympy in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from lightdsa>=0.0.3->deepface==0.0.100) (1.14.0)\n",
      "Requirement already satisfied: joblib>=1.4.2 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from mtcnn>=0.1.0->deepface==0.0.100) (1.5.3)\n",
      "Requirement already satisfied: lz4>=4.3.3 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from mtcnn>=0.1.0->deepface==0.0.100) (4.4.5)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from pandas>=0.23.4->deepface==0.0.100) (2.9.0.post0)\n",
      "Requirement already satisfied: pytz>=2020.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from pandas>=0.23.4->deepface==0.0.100) (2026.1.post1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from pandas>=0.23.4->deepface==0.0.100) (2026.1)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from sympy->lightdsa>=0.0.3->deepface==0.0.100) (1.3.0)\n",
      "Requirement already satisfied: networkx in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (3.6.1)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.6.77 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (12.6.77)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.6.77 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (12.6.77)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.6.80 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (12.6.80)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==9.5.1.17 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (9.5.1.17)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.6.4.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (12.6.4.1)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.3.0.4 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (11.3.0.4)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.7.77 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (10.3.7.77)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.7.1.2 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (11.7.1.2)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.5.4.2 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (12.5.4.2)\n",
      "Requirement already satisfied: nvidia-cusparselt-cu12==0.6.3 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (0.6.3)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.26.2 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (2.26.2)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.6.77 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (12.6.77)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12==12.6.85 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (12.6.85)\n",
      "Requirement already satisfied: nvidia-cufile-cu12==1.11.1.6 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (1.11.1.6)\n",
      "Requirement already satisfied: triton==3.3.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from torch>=1.13.0->peft==0.17.1) (3.3.1)\n",
      "Requirement already satisfied: soupsieve>=1.6.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from beautifulsoup4->gdown>=3.10.1->deepface==0.0.100) (2.8.3)\n",
      "Requirement already satisfied: zipp>=3.20 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from importlib-metadata->diffusers==0.37.1) (3.23.1)\n",
      "Requirement already satisfied: PySocks!=1.5.7,>=1.5.6 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from requests[socks]->gdown>=3.10.1->deepface==0.0.100) (1.7.1)\n",
      "Requirement already satisfied: markdown-it-py>=2.2.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from rich->keras>=2.2.0->deepface==0.0.100) (4.0.0)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from rich->keras>=2.2.0->deepface==0.0.100) (2.20.0)\n",
      "Requirement already satisfied: mdurl~=0.1 in /root/autodl-tmp/AI-Guide-and-Demos-zh_CN/.venv/lib/python3.12/site-packages (from markdown-it-py>=2.2.0->rich->keras>=2.2.0->deepface==0.0.100) (0.1.2)\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install \\\n",
    "    \"transformers==4.56.2\" \\\n",
    "    \"diffusers==0.37.1\" \\\n",
    "    \"peft==0.17.1\" \\\n",
    "    \"accelerate==1.13.0\" \\\n",
    "    \"safetensors==0.7.0\" \\\n",
    "    \"opencv-python\" \\\n",
    "    \"requests\" \\\n",
    "    \"deepface==0.0.100\" \\\n",
    "    \"tf-keras\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09608f50-24e2-4860-b68a-320b25207f98",
   "metadata": {},
   "source": [
    "## 导入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "63e43e2c-c6f9-4981-8ac8-a8f89de845e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "I0000 00:00:1780455999.460469     715 port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "I0000 00:00:1780456001.016503     715 cpu_feature_guard.cc:227] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "I0000 00:00:1780456004.455499     715 port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n"
     ]
    }
   ],
   "source": [
    "# ========== 标准库模块 ==========\n",
    "import os\n",
    "# 设置模型下载镜像\n",
    "os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'\n",
    "os.environ['TF_USE_LEGACY_KERAS'] = '1'  # 让 tf 2.16+ 走 Keras 2 兼容层，deepface 需要\n",
    "\n",
    "try:\n",
    "    import tensorflow as _tf\n",
    "    _tf.config.set_visible_devices([], \"GPU\")\n",
    "except Exception:\n",
    "    pass\n",
    "\n",
    "\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": 3,
   "id": "4ae237bf-6d30-4efa-b6a9-7a6d22827238",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "📂 正在检查并解压样例数据集...\n",
      "✅ 项目 Brad 已准备好！\n"
     ]
    }
   ],
   "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",
    "# 推理用的 checkpoint: \"best\"（按 face_score 选出的最佳）/ \"last\"（训练末尾）\n",
    "checkpoint_choice = \"best\"\n",
    "model_path = os.path.join(project_dir, \"logs\", f\"checkpoint-{checkpoint_choice}\")\n",
    "if not os.path.exists(model_path):\n",
    "    fallback = os.path.join(project_dir, \"logs\", \"checkpoint-last\")\n",
    "    if os.path.exists(fallback):\n",
    "        print(f\"⚠ {model_path} 不存在，回退到 {fallback}\")\n",
    "        model_path = fallback\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",
    "    os.makedirs(os.path.dirname(zip_file), exist_ok=True)\n",
    "    print(\"📥 数据集 Datasets.zip 未找到，开始下载…\")\n",
    "    subprocess.run([\n",
    "        \"wget\", \"-nc\", \"-O\", zip_file,\n",
    "        \"https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/raw/master/Demos/data/14/Datasets.zip\"\n",
    "    ], check=True)\n",
    "\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": 4,
   "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",
    "    用于构建文本到图像模型的微调数据集。\n",
    "    \"\"\"\n",
    "    def __init__(self, images_folder, captions_folder, transform, tokenizer):\n",
    "        \"\"\"\n",
    "        初始化数据集。\n",
    "\n",
    "        参数:\n",
    "            images_folder (str): 图像文件夹路径\n",
    "            captions_folder (str): 标注文件夹路径\n",
    "            transform (Callable): 将原始图像转换为 torch.Tensor 的变换函数\n",
    "            tokenizer (CLIPTokenizer): 将文本标注转为 token 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": 5,
   "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",
    "    加载完整的 Stable Diffusion 模型，包括 LoRA 层，并根据需要合并 LoRA 权重。\n",
    "\n",
    "    包括 Tokenizer、噪声调度器、UNet、VAE 和文本编码器。\n",
    "\n",
    "    参数:\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",
    "    返回:\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",
    "        variant=\"fp16\"\n",
    "    )\n",
    "\n",
    "    # 加载 VAE 模型，用于在扩散模型中处理图像的潜在表示\n",
    "    vae = AutoencoderKL.from_pretrained(\n",
    "        pretrained_model_name_or_path,\n",
    "        subfolder=\"vae\",\n",
    "        variant=\"fp16\"\n",
    "    )\n",
    "\n",
    "    # 加载 UNet 模型，负责处理扩散模型中的图像生成和推理过程\n",
    "    unet = UNet2DConditionModel.from_pretrained(\n",
    "        pretrained_model_name_or_path,\n",
    "        torch_dtype=weight_dtype,\n",
    "        subfolder=\"unet\",\n",
    "        variant=\"fp16\"\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": 6,
   "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",
    "    为 UNet 和文本编码器的可训练参数分别设置优化器，并指定不同的学习率。\n",
    "\n",
    "    参数:\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",
    "    返回:\n",
    "        torch.optim.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": 7,
   "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": 8,
   "id": "bfe24c04-4638-466a-a909-8e44e7b43560",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🖥 当前使用的设备: cuda\n"
     ]
    }
   ],
   "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": 9,
   "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": 10,
   "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 = 2000  # 总训练步数，决定了整个训练过程的迭代次数\n",
    "num_cycles = 3  # Cosine 调度器的周期数量，在训练期间会重复 3 次学习率周期性递减并重启\n",
    "\n",
    "# 预训练的 Stable Diffusion 模型路径，用于加载模型进行微调\n",
    "pretrained_model_name_or_path = \"stablediffusionapi/cyberrealistic-v42\"  \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": 11,
   "id": "7d335cfc-007d-4d47-878b-432cf8d1f4ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ 数据集准备完成！\n"
     ]
    }
   ],
   "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": 12,
   "id": "0d00a8bf-4103-4e22-93e5-14e749a84376",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`torch_dtype` is deprecated! Use `dtype` instead!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"margin:0; white-space:pre; font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;\"></pre>"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARN  Python GIL is enabled: Multi-gpu quant acceleration for MoE models is sub-optimal and multi-core accelerated cpu packing is also disabled. We recommend Python >= 3.13.3t with Pytorch > 2.8 for mult-gpu quantization and multi-cpu packing with env `PYTHON_GIL=0`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"margin:0; white-space:pre; font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;\"></pre>"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"margin:0; white-space:pre; font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;\"></pre>"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARN  Feature `utils/Perplexity` requires python GIL or Python >= 3.13.3T (T for Threading-Free edition of Python) plus Torch 2.8. Feature is currently skipped/disabled.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"margin:0; white-space:pre; font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;\"></pre>"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"margin:0; white-space:pre; font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;\"></pre>"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO  ENV: Auto setting PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' for memory saving.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"margin:0; white-space:pre; font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;\"></pre>"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"margin:0; white-space:pre; font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;\"></pre>"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO  ENV: Auto setting CUDA_DEVICE_ORDER=PCI_BUS_ID for correctness.          \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"margin:0; white-space:pre; font-family:ui-monospace,SFMono-Regular,Menlo,Consolas,monospace;\"></pre>"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "📊 Text Encoder 可训练参数:\n",
      "trainable params: 2,359,296 || all params: 125,419,776 || trainable%: 1.8811\n",
      "📊 UNet 可训练参数:\n",
      "trainable params: 6,377,472 || all params: 865,898,436 || trainable%: 0.7365\n",
      "✅ 模型和优化器准备完成！可以开始训练。\n"
     ]
    }
   ],
   "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": "markdown",
   "id": "ecd8f8f8",
   "metadata": {},
   "source": [
    "### 验证辅助函数（用 face_score 选出最佳检查点）\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d73e295b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a3b0d9b1b1a40adaff15a9f72e78ab6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "提取训练图像面部嵌入:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "训练集面部嵌入提取完成: 100 / 100 张\n"
     ]
    }
   ],
   "source": [
    "# 训练过程中用 face_score 评估，保存表现最好的一版为 checkpoint-best\n",
    "def compute_face_score(unet, text_encoder, train_emb, validation_prompt_path,\n",
    "                       output_folder=None, num_prompts=None):\n",
    "    \"\"\"用当前训练中的 unet/text_encoder 生成验证图，返回 (face_score, clip_score, mis)。\"\"\"\n",
    "    if train_emb is None or len(train_emb) == 0:\n",
    "        return float(\"inf\"), 0.0, 0\n",
    "    unet.eval(); text_encoder.eval()\n",
    "    pipeline = DiffusionPipeline.from_pretrained(\n",
    "        pretrained_model_name_or_path,\n",
    "        unet=unet, text_encoder=text_encoder,\n",
    "        torch_dtype=weight_dtype, safety_checker=None, variant=\"fp16\",\n",
    "    ).to(DEVICE)\n",
    "\n",
    "    with open(validation_prompt_path, \"r\", encoding=\"utf-8\") as f:\n",
    "        prompts = [line.strip() for line in f if line.strip()]\n",
    "    if num_prompts is not None:\n",
    "        prompts = prompts[:num_prompts]\n",
    "\n",
    "    # 加载 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",
    "    clip_model.eval()\n",
    "\n",
    "    gen = torch.Generator(device=DEVICE).manual_seed(seed)\n",
    "    images = []\n",
    "    with torch.no_grad():\n",
    "        for p in prompts:\n",
    "            images.append(pipeline(p, num_inference_steps=30, generator=gen).images[0])\n",
    "\n",
    "    unet.train(); text_encoder.train()\n",
    "\n",
    "    face_score, clip_score, mis = 0, 0, 0\n",
    "    valid = []\n",
    "    for i, image in enumerate(tqdm(images, desc=\"验证图像\")):\n",
    "        if output_folder is not None:\n",
    "            os.makedirs(output_folder, exist_ok=True)\n",
    "            image.save(os.path.join(output_folder, f\"valid_image_{i}.png\"))\n",
    "        opencvImage = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)\n",
    "        emb = DeepFace.represent(opencvImage, detector_backend=\"ssd\",\n",
    "                                 model_name=\"GhostFaceNet\", enforce_detection=False)\n",
    "        if not emb or emb[0].get(\"face_confidence\", 0) == 0:\n",
    "            mis += 1\n",
    "            continue\n",
    "        inputs = clip_processor(text=prompts[i], images=image, return_tensors=\"pt\").to(DEVICE)\n",
    "        with torch.no_grad():\n",
    "            outputs = clip_model(**inputs)\n",
    "        clip_score += outputs.logits_per_image.item()\n",
    "        valid.append(emb[0][\"embedding\"])\n",
    "\n",
    "    del pipeline, clip_model, clip_processor\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "    if not valid:\n",
    "        return float(\"inf\"), 0.0, mis\n",
    "    v = torch.tensor(valid).to(DEVICE)\n",
    "    v = v / v.norm(p=2, dim=-1, keepdim=True)\n",
    "    t = train_emb / train_emb.norm(p=2, dim=-1, keepdim=True)\n",
    "    face_score = torch.cdist(v, t, p=2).mean().item()\n",
    "    clip_score /= max(1, len(prompts) - mis)\n",
    "    return face_score, clip_score, mis\n",
    "\n",
    "# 提取训练集图像的面部嵌入，作为验证时比对的基准\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",
    "emb_list = []\n",
    "for img_path in tqdm(train_image_paths, desc=\"提取训练图像面部嵌入\"):\n",
    "    face_representation = DeepFace.represent(img_path, detector_backend=\"ssd\",\n",
    "                                             model_name=\"GhostFaceNet\", enforce_detection=False)\n",
    "    if face_representation:\n",
    "        emb_list.append(face_representation[0][\"embedding\"])\n",
    "train_face_emb = torch.tensor(emb_list).to(DEVICE) if emb_list else None\n",
    "print(f\"训练集面部嵌入提取完成: {len(emb_list)} / {len(train_image_paths)} 张\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4117f19a-fd07-4828-9c0a-1b91ae658986",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7fc0a1e53c0f42c89441fb1fecba2e3d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "训练步骤:   0%|          | 0/2000 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🔥 步骤 100, 损失: 0.018014751374721527\n",
      "🔥 步骤 200, 损失: 0.05831506848335266\n",
      "🔥 步骤 300, 损失: 0.14568862318992615\n",
      "🔥 步骤 400, 损失: 0.13069385290145874\n",
      "🔥 步骤 500, 损失: 0.05475645884871483\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b6558a11467149f2bf48a4fc0577044e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/6 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Expected types for unet: (<class 'diffusers.models.unets.unet_2d_condition.UNet2DConditionModel'>,), got <class 'peft.peft_model.PeftModel'>.\n",
      "Expected types for text_encoder: (<class 'transformers.models.clip.modeling_clip.CLIPTextModel'>,), got <class 'peft.peft_model.PeftModel'>.\n",
      "You have disabled the safety checker for <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> by passing `safety_checker=None`. Ensure that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling it only for use-cases that involve analyzing network behavior or auditing its results. For more information, please have a look at https://github.com/huggingface/diffusers/pull/254 .\n",
      "Using a slow image processor as `use_fast` is unset and a slow processor was saved with this model. `use_fast=True` will be the default behavior in v4.52, even if the model was saved with a slow processor. This will result in minor differences in outputs. You'll still be able to use a slow processor with `use_fast=False`.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c0d80ac52da84d3a9ba856e1c864d282",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "27c479983ef14aa88dc4982129ae0518",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e111101abb5a49beacc3721fee19aad3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9a7201ac0edf488ca08830a9a27631d3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "验证图像:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "📊 步骤 500: face_score = 1.2877, CLIP = 31.0312, mis = 0 (当前最佳 face_score = inf)\n",
      "💾 face_score 改善，已保存 checkpoint-best\n",
      "💾 已保存中间模型到 ./SD/Brad/logs/checkpoint-500\n",
      "🔥 步骤 600, 损失: 0.06504807621240616\n",
      "🔥 步骤 700, 损失: 0.03095773234963417\n",
      "🔥 步骤 800, 损失: 0.006262064911425114\n",
      "🔥 步骤 900, 损失: 0.01825116015970707\n",
      "🔥 步骤 1000, 损失: 0.04013074189424515\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8d3344b958854e80bd2609243bbb0511",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/6 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Expected types for unet: (<class 'diffusers.models.unets.unet_2d_condition.UNet2DConditionModel'>,), got <class 'peft.peft_model.PeftModel'>.\n",
      "Expected types for text_encoder: (<class 'transformers.models.clip.modeling_clip.CLIPTextModel'>,), got <class 'peft.peft_model.PeftModel'>.\n",
      "You have disabled the safety checker for <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> by passing `safety_checker=None`. Ensure that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling it only for use-cases that involve analyzing network behavior or auditing its results. For more information, please have a look at https://github.com/huggingface/diffusers/pull/254 .\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ea0b2d1108b420a8b432672b47a225b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7f7da20f2cf2447bba632dafdb2e5f06",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2df88e15c92e4bd78cb7b2e7e1c0d73e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "784fd1ec7cb54c12bcd650a063ccf3f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "验证图像:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "📊 步骤 1000: face_score = 1.2223, CLIP = 30.0137, mis = 0 (当前最佳 face_score = 1.2877)\n",
      "💾 face_score 改善，已保存 checkpoint-best\n",
      "💾 已保存中间模型到 ./SD/Brad/logs/checkpoint-1000\n",
      "🔥 步骤 1100, 损失: 0.11263954639434814\n",
      "🔥 步骤 1200, 损失: 0.027997460216283798\n",
      "🔥 步骤 1300, 损失: 0.06858616322278976\n",
      "🔥 步骤 1400, 损失: 0.022438760846853256\n",
      "🔥 步骤 1500, 损失: 0.041379399597644806\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "54820b4d9b1c4e21823a23a35fa6a1a1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/6 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Expected types for unet: (<class 'diffusers.models.unets.unet_2d_condition.UNet2DConditionModel'>,), got <class 'peft.peft_model.PeftModel'>.\n",
      "Expected types for text_encoder: (<class 'transformers.models.clip.modeling_clip.CLIPTextModel'>,), got <class 'peft.peft_model.PeftModel'>.\n",
      "You have disabled the safety checker for <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> by passing `safety_checker=None`. Ensure that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling it only for use-cases that involve analyzing network behavior or auditing its results. For more information, please have a look at https://github.com/huggingface/diffusers/pull/254 .\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ddb45c4d7c04d1baf0e1f96a9d1d200",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ed7e73bd56942a49fd639cec0552822",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ad4ccb90cfc04a73bf5af33823eb9da2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1520c8a67e194456860dfafb238c3f13",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "验证图像:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "📊 步骤 1500: face_score = 1.2325, CLIP = 29.0263, mis = 0 (当前最佳 face_score = 1.2223)\n",
      "💾 已保存中间模型到 ./SD/Brad/logs/checkpoint-1500\n",
      "🔥 步骤 1600, 损失: 0.020244833081960678\n",
      "🔥 步骤 1700, 损失: 0.020814694464206696\n",
      "🔥 步骤 1800, 损失: 0.07016117870807648\n",
      "🔥 步骤 1900, 损失: 0.08010156452655792\n",
      "🔥 步骤 2000, 损失: 0.1658126711845398\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "10be169fad2a414f9e3d24e022941290",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/6 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Expected types for unet: (<class 'diffusers.models.unets.unet_2d_condition.UNet2DConditionModel'>,), got <class 'peft.peft_model.PeftModel'>.\n",
      "Expected types for text_encoder: (<class 'transformers.models.clip.modeling_clip.CLIPTextModel'>,), got <class 'peft.peft_model.PeftModel'>.\n",
      "You have disabled the safety checker for <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> by passing `safety_checker=None`. Ensure that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling it only for use-cases that involve analyzing network behavior or auditing its results. For more information, please have a look at https://github.com/huggingface/diffusers/pull/254 .\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4951cacff76c4f00bc2faea2ce9b1b50",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d1f0d1cc95e4b2c814c7d05a2d38bf7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0cabcdb0d9324105ba0c16d4149fb402",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3164c277910641ac87188d34a44f97d6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "验证图像:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "📊 步骤 2000: face_score = 1.2220, CLIP = 29.4632, mis = 0 (当前最佳 face_score = 1.2223)\n",
      "💾 face_score 改善，已保存 checkpoint-best\n",
      "💾 已保存中间模型到 ./SD/Brad/logs/checkpoint-2000\n",
      "💾 已保存最终模型到 ./SD/Brad/logs/checkpoint-last\n",
      "🎉 微调完成！\n"
     ]
    }
   ],
   "source": [
    "# 禁用并行化，避免警告\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "# 初始化\n",
    "global_step = 0\n",
    "best_face_score = float(\"inf\")  # 初始化为正无穷大，存储最佳面部相似度分数\n",
    "validation_step_ratio = 0.25  # 每 25% 训练进度做一次 face_score 验证\n",
    "validation_prompt_num = 3  # 验证时使用的 prompt 数\n",
    "validation_step = max(1, int(max_train_steps * validation_step_ratio))\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",
    "        # 周期性验证：face_score 改善则保存 checkpoint-best\n",
    "        if global_step % validation_step == 0 or global_step == max_train_steps:\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",
    "\n",
    "            # 验证图保存到 checkpoint-{global_step}/\n",
    "            val_output_folder = os.path.join(output_folder, f\"checkpoint-{global_step}\")\n",
    "            face_score, clip_score, mis = compute_face_score(\n",
    "                unet, text_encoder, train_face_emb, validation_prompt_path,\n",
    "                output_folder=val_output_folder, num_prompts=validation_prompt_num,\n",
    "            )\n",
    "            print(f\"📊 步骤 {global_step}: face_score = {face_score:.4f}, CLIP = {clip_score:.4f}, mis = {mis} (当前最佳 face_score = {best_face_score:.4f})\")\n",
    "            if face_score < best_face_score:\n",
    "                best_face_score = face_score\n",
    "                save_path = os.path.join(output_folder, \"checkpoint-best\")\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\"💾 face_score 改善，已保存 checkpoint-best\")\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": 15,
   "id": "53cce6c3-d106-4317-8f28-4dc7be4c7b9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_validation_prompts(validation_prompt_path):\n",
    "    \"\"\"\n",
    "    加载验证提示文本。\n",
    "\n",
    "    参数:\n",
    "        validation_prompt_path (str): 验证提示文件的路径\n",
    "\n",
    "    返回:\n",
    "        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": 16,
   "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",
    "    使用 DiffusionPipeline 生成图像，保存到指定文件夹并返回生成的图像列表。\n",
    "\n",
    "    参数:\n",
    "        pipeline (DiffusionPipeline): 已加载并配置好的 Pipeline\n",
    "        prompts (list): 文本提示列表\n",
    "        num_inference_steps (int): 推理步骤数，越高图像质量越好，但推理时间也会增加\n",
    "        guidance_scale (float): 决定文本提示对生成图像的影响程度\n",
    "        output_folder (str): 保存生成图像的文件夹路径\n",
    "        generator (torch.Generator | None): 控制生成随机数的种子\n",
    "\n",
    "    返回:\n",
    "        list: 生成的图像列表，同时图像也会保存到指定文件夹\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": 17,
   "id": "fcc5cfa8-7186-431f-923c-ca40d7a120db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "📂 加载验证提示...\n",
      "🔧 准备 LoRA 模型...\n",
      "✅ 已从 ./SD/Brad/logs/checkpoint-best 恢复模型权重\n",
      "🔄 创建 DiffusionPipeline...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c97f2889a3a342028ca2950fffe6c55c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/6 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You have disabled the safety checker for <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> by passing `safety_checker=None`. Ensure that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling it only for use-cases that involve analyzing network behavior or auditing its results. For more information, please have a look at https://github.com/huggingface/diffusers/pull/254 .\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🎯 加载 CLIP 模型...\n",
      "📂 加载训练图像的面部嵌入...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "755c19042c90452dba739245e6a8f96b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "提取训练图像面部嵌入:   0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🎨 正在生成图像...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "466ae3646ed045e7afe28dc4699df4c1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "生成图像中:   0%|          | 0/25 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f95a1de933f747d4bb9a5f4afde3ccb3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "10dd35b2d156428fb521448e7b43c4f2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dde14ed4cf71445e860be4fc5d3d31d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "07c4b8cfe72249ce818090f0e3b26092",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "79e02491649b4dfd9380d7c44e0ac19f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e69083ed2d5643bf997b359f882e94f0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f693b3126bd244c4b0e9be9e083c13c9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "55e064b95df245cbbcf143783fcb0595",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4ad5c9361b6240a0b4a0c1c90cc9182a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "71eb3317c1ac4282a4108a251fc9ab67",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "066f5c812b7e4ac5980997d92740aeee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "78f2f73b4bae45e89457f2359521d6cd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "76ece1a20e8c4030b6be0e97b94b1b6a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "218e229b031c4c81aeb1e59752abd45a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3c1444b4ba424876b4d0f325a2944c12",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a2674b8591aa473fa5d31e0cadd4bb43",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1defe52dd1de4987b41ddd211435a795",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7ed7f336a6a04491937d4f7b3dace7f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d6e5af9534742b2b7fe97e1f9f83b57",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d73af30ac004892b23bc2d181ff687a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7b6b19e516844270af6bbac3a08cdf45",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c2c0057cf6454e5c816156207d8e4786",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f282c4e24e84a4da7b9cc00c3a6a650",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5e7cded1cd1d4cfb8b316f36a9ee99ba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d7c6a267ff234ab29decba144cfa0690",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/30 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ 已生成并保存 25 张图像到 ./SD/Brad/inference\n",
      "📊 正在计算评估分数...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "59f94da5cb6f4c3989201614a80921de",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "评估图像中:   0%|          | 0/25 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "📈 评估完成！\n",
      "✅ 面部相似度评分 (平均欧氏距离): 1.1969 (越低越好，表示生成图像与训练图像更相似)\n",
      "✅ CLIP 评分 (平均相似度): 28.3417 (越高越好，表示生成图像与文本提示的相关性更强)\n",
      "✅ 无面部图像数量: 1 (无法检测到面部的生成图像数量)\n"
     ]
    }
   ],
   "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",
    "    参数:\n",
    "        lora_config (LoraConfig): LoRA 的配置对象\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",
    "        variant=\"fp16\",\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": "ai",
   "language": "python",
   "name": "ai"
  },
  "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.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}