{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a205ddd9",
   "metadata": {},
   "source": [
    "# 魔搭社区 AIGC 系列课程 - 可控生成技术\n",
    "\n",
    "本实验以 **Diffusion-Templates** 为框架，系统介绍图像生成模型的多种可控生成技术，并演示如何自行训练一个可控生成模块。\n",
    "\n",
    "相关资料：\n",
    "\n",
    "* 开源代码：[DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio)\n",
    "* 技术报告：[arXiv](https://arxiv.org/abs/2604.24351)\n",
    "* 项目主页：[GitHub](https://modelscope.github.io/diffusion-templates-web/)\n",
    "* 文档参考：[English Version](https://diffsynth-studio-doc.readthedocs.io/en/latest/Diffusion_Templates/Introducing_Diffusion_Templates.html)、[中文版](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/Diffusion_Templates/Introducing_Diffusion_Templates.html)\n",
    "* 在线体验：[魔搭社区创空间](https://modelscope.cn/studios/DiffSynth-Studio/Diffusion-Templates)\n",
    "* 模型集：[ModelScope](https://modelscope.cn/collections/DiffSynth-Studio/KleinBase4B-Templates)、[ModelScope 国际站](https://modelscope.ai/collections/DiffSynth-Studio/KleinBase4B-Templates)、[HuggingFace](https://huggingface.co/collections/DiffSynth-Studio/kleinbase4b-templates)\n",
    "* 数据集：[ModelScope](https://modelscope.cn/collections/DiffSynth-Studio/ImagePulseV2)、[ModelScope 国际站](https://modelscope.cn/collections/DiffSynth-Studio/ImagePulseV2)、[HuggingFace](https://huggingface.co/collections/DiffSynth-Studio/imagepulsev2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c556f6de",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install diffsynth==2.0.15 transformers==5.8.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acbd35c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffsynth.diffusion.template import TemplatePipeline\n",
    "from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig\n",
    "import torch\n",
    "from modelscope import dataset_snapshot_download, snapshot_download\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "vram_config = {\n",
    "    \"offload_dtype\": \"disk\",\n",
    "    \"offload_device\": \"disk\",\n",
    "    \"onload_dtype\": torch.float8_e4m3fn,\n",
    "    \"onload_device\": \"cpu\",\n",
    "    \"preparing_dtype\": torch.float8_e4m3fn,\n",
    "    \"preparing_device\": \"cuda\",\n",
    "    \"computation_dtype\": torch.bfloat16,\n",
    "    \"computation_device\": \"cuda\",\n",
    "}\n",
    "\n",
    "def show_images(images, resolution):\n",
    "    images = [i.resize((resolution, resolution)).convert(\"RGB\") for i in images]\n",
    "    images = [np.array(i) for i in images]\n",
    "    images = np.concat(images, axis=1)\n",
    "    images = Image.fromarray(images)\n",
    "    return images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d58a54f2",
   "metadata": {},
   "source": [
    "首先，加载基础模型 [black-forest-labs/FLUX.2-klein-base-4B](https://modelscope.cn/models/black-forest-labs/FLUX.2-klein-base-4B)。这是一个参数量为 4B 的图像生成模型，本实验后续所有可控生成模块都会挂载到这个基础模型之上。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bb3f260",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = Flux2ImagePipeline.from_pretrained(\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    device=\"cuda\",\n",
    "    model_configs=[\n",
    "        ModelConfig(model_id=\"black-forest-labs/FLUX.2-klein-base-4B\", origin_file_pattern=\"transformer/*.safetensors\", **vram_config),\n",
    "        ModelConfig(model_id=\"black-forest-labs/FLUX.2-klein-4B\", origin_file_pattern=\"text_encoder/*.safetensors\", **vram_config),\n",
    "        ModelConfig(model_id=\"black-forest-labs/FLUX.2-klein-4B\", origin_file_pattern=\"vae/diffusion_pytorch_model.safetensors\"),\n",
    "    ],\n",
    "    tokenizer_config=ModelConfig(model_id=\"black-forest-labs/FLUX.2-klein-4B\", origin_file_pattern=\"tokenizer/\"),\n",
    "    vram_limit=torch.cuda.mem_get_info(\"cuda\")[1] / (1024 ** 3) - 0.5,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b0ed288",
   "metadata": {},
   "source": [
    "## 图像结构控制\n",
    "\n",
    "[ControlNet](https://arxiv.org/abs/2302.05543) 是最早的一批 Diffusion 可控生成技术，可用**深度图、边缘图、姿态图**等结构性条件对生成画面进行**逐像素级**的控制。\n",
    "\n",
    "以 Template 格式加载 [DiffSynth-Studio/Template-KleinBase4B-ControlNet](https://modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-ControlNet)，即可在保留输入结构的前提下，用不同的提示词生成不同风格的画面。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96c1225f",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = TemplatePipeline.from_pretrained(\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    device=\"cuda\",\n",
    "    model_configs=[ModelConfig(model_id=\"DiffSynth-Studio/Template-KleinBase4B-ControlNet\")],\n",
    "    lazy_loading=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b906f029",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_snapshot_download(\n",
    "    \"DiffSynth-Studio/examples_in_diffsynth\",\n",
    "    allow_file_pattern=[\"templates/*\"],\n",
    "    local_dir=\"data/examples\",\n",
    ")\n",
    "image = template(\n",
    "    pipe,\n",
    "    prompt=\"A cat is sitting on a stone, bathed in bright sunshine.\",\n",
    "    seed=0, cfg_scale=4, num_inference_steps=50,\n",
    "    template_inputs = [{\n",
    "        \"image\": Image.open(\"data/examples/templates/image_depth.jpg\"),\n",
    "        \"prompt\": \"A cat is sitting on a stone, bathed in bright sunshine.\",\n",
    "    }],\n",
    "    negative_template_inputs = [{\n",
    "        \"image\": Image.open(\"data/examples/templates/image_depth.jpg\"),\n",
    "        \"prompt\": \"\",\n",
    "    }],\n",
    ")\n",
    "image.save(\"image_ControlNet_sunshine.jpg\")\n",
    "image = template(\n",
    "    pipe,\n",
    "    prompt=\"A cat is sitting on a stone, surrounded by colorful magical particles.\",\n",
    "    seed=0, cfg_scale=4, num_inference_steps=50,\n",
    "    template_inputs = [{\n",
    "        \"image\": Image.open(\"data/examples/templates/image_depth.jpg\"),\n",
    "        \"prompt\": \"A cat is sitting on a stone, surrounded by colorful magical particles.\",\n",
    "    }],\n",
    "    negative_template_inputs = [{\n",
    "        \"image\": Image.open(\"data/examples/templates/image_depth.jpg\"),\n",
    "        \"prompt\": \"\",\n",
    "    }],\n",
    ")\n",
    "image.save(\"image_ControlNet_magic.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af45ce68",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_images([\n",
    "    Image.open(\"data/examples/templates/image_depth.jpg\"),\n",
    "    Image.open(\"image_ControlNet_sunshine.jpg\"),\n",
    "    Image.open(\"image_ControlNet_magic.jpg\"),\n",
    "], resolution=256)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bb8d720",
   "metadata": {},
   "source": [
    "## 数值属性控制\n",
    "\n",
    "[AttriCtrl](https://arxiv.org/abs/2508.02151) 是一类**数值型**可控生成模型，能够将连续的数值属性作为控制条件注入生成过程。\n",
    "\n",
    "运行以下代码，加载 [DiffSynth-Studio/Template-KleinBase4B-SoftRGB](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-SoftRGB)，通过输入 R/G/B 数值精确控制画面的整体色调。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3352ae2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = TemplatePipeline.from_pretrained(\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    device=\"cuda\",\n",
    "    model_configs=[ModelConfig(model_id=\"DiffSynth-Studio/Template-KleinBase4B-SoftRGB\")],\n",
    "    lazy_loading=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b6a871c",
   "metadata": {},
   "outputs": [],
   "source": [
    "image = template(\n",
    "    pipe,\n",
    "    prompt=\"A cat is sitting on a stone.\",\n",
    "    seed=0, cfg_scale=4, num_inference_steps=50,\n",
    "    template_inputs = [{\"R\": 128/255, \"G\": 128/255, \"B\": 128/255}],\n",
    ")\n",
    "image.save(\"image_rgb_normal.jpg\")\n",
    "image = template(\n",
    "    pipe,\n",
    "    prompt=\"A cat is sitting on a stone.\",\n",
    "    seed=0, cfg_scale=4, num_inference_steps=50,\n",
    "    template_inputs = [{\"R\": 208/255, \"G\": 185/255, \"B\": 138/255}],\n",
    ")\n",
    "image.save(\"image_rgb_warm.jpg\")\n",
    "image = template(\n",
    "    pipe,\n",
    "    prompt=\"A cat is sitting on a stone.\",\n",
    "    seed=0, cfg_scale=4, num_inference_steps=50,\n",
    "    template_inputs = [{\"R\": 94/255, \"G\": 163/255, \"B\": 174/255}],\n",
    ")\n",
    "image.save(\"image_rgb_cold.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00f4174f",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_images([\n",
    "    Image.open(\"image_rgb_normal.jpg\"),\n",
    "    Image.open(\"image_rgb_warm.jpg\"),\n",
    "    Image.open(\"image_rgb_cold.jpg\"),\n",
    "], resolution=256)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebf205dd",
   "metadata": {},
   "source": [
    "## 图像编辑\n",
    "\n",
    "图像编辑模型是一类**通用性较强**的可控生成模型：给定一张原图和一段编辑指令，即可对原图进行局部或整体修改。\n",
    "\n",
    "运行以下代码，加载 [DiffSynth-Studio/Template-KleinBase4B-Edit](https://www.modelscope.cn/models/DiffSynth-Studio/Template-KleinBase4B-Edit)。该模型通过 **KV-Cache** 复用输入图像的注意力键值，从而快速完成编辑，推理速度较快。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "730bb8bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = TemplatePipeline.from_pretrained(\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    device=\"cuda\",\n",
    "    model_configs=[ModelConfig(model_id=\"DiffSynth-Studio/Template-KleinBase4B-Edit\")],\n",
    "    lazy_loading=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f8da7ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_snapshot_download(\n",
    "    \"DiffSynth-Studio/examples_in_diffsynth\",\n",
    "    allow_file_pattern=[\"templates/*\"],\n",
    "    local_dir=\"data/examples\",\n",
    ")\n",
    "image = template(\n",
    "    pipe,\n",
    "    prompt=\"Put a hat on this cat.\",\n",
    "    seed=0, cfg_scale=4, num_inference_steps=50,\n",
    "    template_inputs = [{\n",
    "        \"image\": Image.open(\"data/examples/templates/image_reference.jpg\"),\n",
    "        \"prompt\": \"Put a hat on this cat.\",\n",
    "    }],\n",
    "    negative_template_inputs = [{\n",
    "        \"image\": Image.open(\"data/examples/templates/image_reference.jpg\"),\n",
    "        \"prompt\": \"\",\n",
    "    }],\n",
    ")\n",
    "image.save(\"image_Edit_hat.jpg\")\n",
    "image = template(\n",
    "    pipe,\n",
    "    prompt=\"Make the cat turn its head to look to the right.\",\n",
    "    seed=0, cfg_scale=4, num_inference_steps=50,\n",
    "    template_inputs = [{\n",
    "        \"image\": Image.open(\"data/examples/templates/image_reference.jpg\"),\n",
    "        \"prompt\": \"Make the cat turn its head to look to the right.\",\n",
    "    }],\n",
    "    negative_template_inputs = [{\n",
    "        \"image\": Image.open(\"data/examples/templates/image_reference.jpg\"),\n",
    "        \"prompt\": \"\",\n",
    "    }],\n",
    ")\n",
    "image.save(\"image_Edit_head.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16fa68bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_images([\n",
    "    Image.open(\"data/examples/templates/image_reference.jpg\"),\n",
    "    Image.open(\"image_Edit_hat.jpg\"),\n",
    "    Image.open(\"image_Edit_head.jpg\"),\n",
    "], resolution=256)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cd453b9",
   "metadata": {},
   "source": [
    "## 风格控制\n",
    "\n",
    "实现图像风格控制的最直接方式，是训练一个风格 [LoRA](https://arxiv.org/abs/2106.09685)——但每种风格都需要单独训练，成本较高。为此我们训练了一个特殊的 [Image-to-LoRA](https://arxiv.org/abs/2606.13809) 模型，它可以**根据输入的参考图像即时生成一份 LoRA 权重**，免去了传统的风格训练过程。\n",
    "\n",
    "运行以下代码，加载 [DiffSynth-Studio/KleinBase4B-i2L-v2](https://www.modelscope.cn/models/DiffSynth-Studio/KleinBase4B-i2L-v2)，用参考图像动态生成 LoRA，从而控制画面风格。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51ae1734",
   "metadata": {},
   "outputs": [],
   "source": [
    "from modelscope import snapshot_download\n",
    "\n",
    "template = TemplatePipeline.from_pretrained(\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    device=\"cuda\",\n",
    "    model_configs=[ModelConfig(model_id=\"DiffSynth-Studio/KleinBase4B-i2L-v2\")],\n",
    "    lazy_loading=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fcc038d",
   "metadata": {},
   "outputs": [],
   "source": [
    "snapshot_download(\"DiffSynth-Studio/KleinBase4B-i2L-v2\", allow_file_pattern=\"assets/*\", local_dir=\"data\")\n",
    "images = [Image.open(f\"data/assets/image_1_{i}.jpg\") for i in range(4)]\n",
    "image = template(\n",
    "    pipe,\n",
    "    prompt=\"A cat is sitting on a stone\",\n",
    "    seed=42, cfg_scale=4, num_inference_steps=50,\n",
    "    template_inputs = [{\"image\": images}],\n",
    "    negative_template_inputs = [{\"image\": [Image.fromarray(np.zeros_like(np.array(i)) + 128) for i in images]}],\n",
    ")\n",
    "image.save(\"image_KleinBase4B-i2L-v2_1.jpg\")\n",
    "images = [Image.open(f\"data/assets/image_3_{i}.jpg\") for i in range(4)]\n",
    "image = template(\n",
    "    pipe,\n",
    "    prompt=\"A cat is sitting on a stone\",\n",
    "    seed=42, cfg_scale=4, num_inference_steps=50,\n",
    "    template_inputs = [{\"image\": images}],\n",
    "    negative_template_inputs = [{\"image\": [Image.fromarray(np.zeros_like(np.array(i)) + 128) for i in images]}],\n",
    ")\n",
    "image.save(\"image_KleinBase4B-i2L-v2_2.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4a6609b",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_images([\n",
    "    Image.open(\"data/assets/image_1_2.jpg\"),\n",
    "    Image.open(\"image_KleinBase4B-i2L-v2_1.jpg\"),\n",
    "    Image.open(\"data/assets/image_3_0.jpg\"),\n",
    "    Image.open(\"image_KleinBase4B-i2L-v2_2.jpg\"),\n",
    "], resolution=256)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7918117f",
   "metadata": {},
   "source": [
    "## 训练可控生成模型\n",
    "\n",
    "**Diffusion-Templates 框架允许开发者训练任意结构的可控生成模型**——只要给定模型定义、数据处理逻辑和数据集，即可接入统一的训练流程。下面我们从零训练一个**亮度控制模型**，让画面按指定的亮度数值生成。\n",
    "\n",
    "第一步，编写模型结构代码（包含数值编码器、KV-Cache 生成主干和数据标注器）："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c25c94f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "code = \"\"\"\n",
    "import torch, math, os\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "class SingleValueEncoder(torch.nn.Module):\n",
    "    def __init__(self, dim_in=256, dim_out=4096, length=32):\n",
    "        super().__init__()\n",
    "        self.length = length\n",
    "        self.prefer_value_embedder = torch.nn.Sequential(torch.nn.Linear(dim_in, dim_out), torch.nn.SiLU(), torch.nn.Linear(dim_out, dim_out))\n",
    "        self.positional_embedding = torch.nn.Parameter(torch.randn(self.length, dim_out))\n",
    "\n",
    "    def get_timestep_embedding(self, timesteps, embedding_dim, max_period=10000):\n",
    "        half_dim = embedding_dim // 2\n",
    "        exponent = -math.log(max_period) * torch.arange(0, half_dim, dtype=torch.float32, device=timesteps.device) / half_dim\n",
    "        emb = timesteps[:, None].float() * torch.exp(exponent)[None, :]\n",
    "        emb = torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1)\n",
    "        return emb\n",
    "\n",
    "    def forward(self, value, dtype):\n",
    "        emb = self.get_timestep_embedding(value * 1000, 256).to(dtype)\n",
    "        emb = self.prefer_value_embedder(emb).squeeze(0)\n",
    "        base_embeddings = emb.expand(self.length, -1)\n",
    "        positional_embedding = self.positional_embedding.to(dtype=base_embeddings.dtype, device=base_embeddings.device)\n",
    "        learned_embeddings = base_embeddings + positional_embedding\n",
    "        return learned_embeddings\n",
    "\n",
    "\n",
    "# 主干模型结构（将输入的数值转换为 KV-Cache 向量）\n",
    "class ValueFormatModel(torch.nn.Module):\n",
    "    def __init__(self, num_double_blocks=5, num_single_blocks=20, dim=3072, num_heads=24, length=512):\n",
    "        super().__init__()\n",
    "        self.block_names = [f\"double_{i}\" for i in range(num_double_blocks)] + [f\"single_{i}\" for i in range(num_single_blocks)]\n",
    "        self.proj_k = torch.nn.ModuleDict({block_name: SingleValueEncoder(dim_out=dim, length=length) for block_name in self.block_names})\n",
    "        self.proj_v = torch.nn.ModuleDict({block_name: SingleValueEncoder(dim_out=dim, length=length) for block_name in self.block_names})\n",
    "        self.num_heads = num_heads\n",
    "        self.length = length\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def process_inputs(self, pipe, scale, **kwargs):\n",
    "        return {\"value\": torch.Tensor([scale]).to(dtype=pipe.torch_dtype, device=pipe.device)}\n",
    "\n",
    "    def forward(self, value, **kwargs):\n",
    "        kv_cache = {}\n",
    "        for block_name in self.block_names:\n",
    "            k = self.proj_k[block_name](value, value.dtype)\n",
    "            k = k.view(1, self.length, self.num_heads, -1)\n",
    "            v = self.proj_v[block_name](value, value.dtype)\n",
    "            v = v.view(1, self.length, self.num_heads, -1)\n",
    "            kv_cache[block_name] = (k, v)\n",
    "        return {\"kv_cache\": kv_cache}\n",
    "\n",
    "\n",
    "# 将图像数据转换为模型输入（根据图像中的 RGB 数值计算亮度）\n",
    "class DataAnnotator(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def __call__(self, image, **kwargs):\n",
    "        image = Image.open(image)\n",
    "        image = np.array(image)\n",
    "        return {\"scale\": image.astype(np.float32).mean() / 255}\n",
    "\n",
    "\n",
    "TEMPLATE_MODEL = ValueFormatModel\n",
    "TEMPLATE_MODEL_PATH = \"model.safetensors\" if \"model.safetensors\" in os.listdir(os.path.dirname(__file__)) else None\n",
    "TEMPLATE_DATA_PROCESSOR = DataAnnotator\n",
    "\"\"\"\n",
    "\n",
    "import os\n",
    "\n",
    "os.makedirs(\"models/template_brightness\", exist_ok=True)\n",
    "with open(\"models/template_brightness/model.py\", \"w\", encoding=\"utf-8\") as f:\n",
    "    f.write(code.strip())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9105dffc",
   "metadata": {},
   "source": [
    "第二步，下载并预处理数据集，同时生成训练所需的 metadata："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94557ebb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json, os\n",
    "from modelscope import dataset_snapshot_download\n",
    "\n",
    "# 下载数据集\n",
    "dataset_snapshot_download(\n",
    "    \"DiffSynth-Studio/ImagePulseV2-TextImage\",\n",
    "    local_dir=\"data/ImagePulseV2-TextImage\",\n",
    "    allow_file_pattern=\"data/1770381050168240056.tar.gz\"\n",
    ")\n",
    "\n",
    "# 解压数据集\n",
    "os.makedirs(\"data/dataset\", exist_ok=True)\n",
    "os.system(\"tar zxvf data/ImagePulseV2-TextImage/data/1770381050168240056.tar.gz -C data/dataset\")\n",
    "\n",
    "# 生成数据集 metadata\n",
    "dataset_path = \"data/dataset/1770381050168240056\"\n",
    "metadata = []\n",
    "for file_name in os.listdir(dataset_path):\n",
    "    if file_name.endswith(\".json\"):\n",
    "        with open(os.path.join(dataset_path, file_name), \"r\") as f:\n",
    "            data = json.load(f)\n",
    "        data[\"template_inputs\"] = {\"image\": os.path.join(dataset_path, data[\"image\"])}\n",
    "        metadata.append(data)\n",
    "with open(\"data/dataset/metadata.json\", \"w\") as f:\n",
    "    json.dump(metadata, f, indent=4, ensure_ascii=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d48dd079",
   "metadata": {},
   "source": [
    "第三步，启动训练："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49c3d4fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# 训练脚本\n",
    "code = \"\"\"\n",
    "import torch, os, argparse, accelerate\n",
    "from diffsynth.core import UnifiedDataset\n",
    "from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig\n",
    "from diffsynth.diffusion import *\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "\n",
    "class Flux2ImageTrainingModule(DiffusionTrainingModule):\n",
    "    def __init__(\n",
    "        self,\n",
    "        model_paths=None, model_id_with_origin_paths=None,\n",
    "        tokenizer_path=None,\n",
    "        trainable_models=None,\n",
    "        lora_base_model=None, lora_target_modules=\"\", lora_rank=32, lora_checkpoint=None,\n",
    "        preset_lora_path=None, preset_lora_model=None,\n",
    "        use_gradient_checkpointing=True,\n",
    "        use_gradient_checkpointing_offload=False,\n",
    "        extra_inputs=None,\n",
    "        fp8_models=None,\n",
    "        offload_models=None,\n",
    "        template_model_id_or_path=None,\n",
    "        resume_from_checkpoint=None, remove_prefix_in_ckpt=None,\n",
    "        enable_lora_hot_loading=False,\n",
    "        device=\"cpu\",\n",
    "        task=\"sft\",\n",
    "    ):\n",
    "        super().__init__()\n",
    "        # Load models\n",
    "        model_configs = self.parse_model_configs(model_paths, model_id_with_origin_paths, fp8_models=fp8_models, offload_models=offload_models, device=device)\n",
    "        tokenizer_config = self.parse_path_or_model_id(tokenizer_path, default_value=ModelConfig(model_id=\"black-forest-labs/FLUX.2-dev\", origin_file_pattern=\"tokenizer/\"))\n",
    "        self.pipe = Flux2ImagePipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device, model_configs=model_configs, tokenizer_config=tokenizer_config)\n",
    "        self.pipe = self.load_training_template_model(self.pipe, template_model_id_or_path, use_gradient_checkpointing, use_gradient_checkpointing_offload)\n",
    "        self.pipe = self.split_pipeline_units(task, self.pipe, trainable_models, lora_base_model, remove_unnecessary_params=True)\n",
    "        self.resume_from_checkpoint(resume_from_checkpoint, remove_prefix_in_ckpt)\n",
    "        if enable_lora_hot_loading: self.pipe.dit = self.pipe.enable_lora_hot_loading(self.pipe.dit)\n",
    "\n",
    "        # Training mode\n",
    "        self.switch_pipe_to_training_mode(\n",
    "            self.pipe, trainable_models,\n",
    "            lora_base_model, lora_target_modules, lora_rank, lora_checkpoint,\n",
    "            preset_lora_path, preset_lora_model,\n",
    "            task=task,\n",
    "        )\n",
    "\n",
    "        # Other configs\n",
    "        self.use_gradient_checkpointing = use_gradient_checkpointing\n",
    "        self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload\n",
    "        self.extra_inputs = extra_inputs.split(\",\") if extra_inputs is not None else []\n",
    "        self.fp8_models = fp8_models\n",
    "        self.task = task\n",
    "        self.task_to_loss = {\n",
    "            \"sft:data_process\": lambda pipe, *args: args,\n",
    "            \"direct_distill:data_process\": lambda pipe, *args: args,\n",
    "            \"sft\": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi),\n",
    "            \"sft:train\": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi),\n",
    "            \"direct_distill\": lambda pipe, inputs_shared, inputs_posi, inputs_nega: DirectDistillLoss(pipe, **inputs_shared, **inputs_posi),\n",
    "            \"direct_distill:train\": lambda pipe, inputs_shared, inputs_posi, inputs_nega: DirectDistillLoss(pipe, **inputs_shared, **inputs_posi),\n",
    "        }\n",
    "\n",
    "    def get_pipeline_inputs(self, data):\n",
    "        inputs_posi = {\"prompt\": data[\"prompt\"]}\n",
    "        inputs_nega = {\"negative_prompt\": \"\"}\n",
    "        inputs_shared = {\n",
    "            # Assume you are using this pipeline for inference,\n",
    "            # please fill in the input parameters.\n",
    "            \"input_image\": data[\"image\"],\n",
    "            \"height\": data[\"image\"].size[1],\n",
    "            \"width\": data[\"image\"].size[0],\n",
    "            # Please do not modify the following parameters\n",
    "            # unless you clearly know what this will cause.\n",
    "            \"embedded_guidance\": 1.0,\n",
    "            \"cfg_scale\": 1,\n",
    "            \"rand_device\": self.pipe.device,\n",
    "            \"use_gradient_checkpointing\": self.use_gradient_checkpointing,\n",
    "            \"use_gradient_checkpointing_offload\": self.use_gradient_checkpointing_offload,\n",
    "        }\n",
    "        inputs_shared = self.parse_extra_inputs(data, self.extra_inputs, inputs_shared)\n",
    "        return inputs_shared, inputs_posi, inputs_nega\n",
    "\n",
    "    def forward(self, data, inputs=None):\n",
    "        if inputs is None: inputs = self.get_pipeline_inputs(data)\n",
    "        inputs = self.transfer_data_to_device(inputs, self.pipe.device, self.pipe.torch_dtype)\n",
    "        for unit in self.pipe.units:\n",
    "            inputs = self.pipe.unit_runner(unit, self.pipe, *inputs)\n",
    "        loss = self.task_to_loss[self.task](self.pipe, *inputs)\n",
    "        return loss\n",
    "\n",
    "\n",
    "def flux2_parser():\n",
    "    parser = argparse.ArgumentParser(description=\"Simple example of a training script.\")\n",
    "    parser = add_general_config(parser)\n",
    "    parser = add_image_size_config(parser)\n",
    "    parser.add_argument(\"--tokenizer_path\", type=str, default=None, help=\"Path to tokenizer.\")\n",
    "    parser.add_argument(\"--initialize_model_on_cpu\", default=False, action=\"store_true\", help=\"Whether to initialize models on CPU.\")\n",
    "    return parser\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    parser = flux2_parser()\n",
    "    args = parser.parse_args()\n",
    "\n",
    "    accelerator = accelerate.Accelerator(\n",
    "        gradient_accumulation_steps=args.gradient_accumulation_steps,\n",
    "        kwargs_handlers=[accelerate.DistributedDataParallelKwargs(find_unused_parameters=args.find_unused_parameters)],\n",
    "    )\n",
    "    dataset = UnifiedDataset(\n",
    "        base_path=args.dataset_base_path,\n",
    "        metadata_path=args.dataset_metadata_path,\n",
    "        repeat=args.dataset_repeat,\n",
    "        data_file_keys=args.data_file_keys.split(\",\"),\n",
    "        main_data_operator=UnifiedDataset.default_image_operator(\n",
    "            base_path=args.dataset_base_path,\n",
    "            max_pixels=args.max_pixels,\n",
    "            height=args.height,\n",
    "            width=args.width,\n",
    "            height_division_factor=16,\n",
    "            width_division_factor=16,\n",
    "        )\n",
    "    )\n",
    "    model = Flux2ImageTrainingModule(\n",
    "        model_paths=args.model_paths,\n",
    "        model_id_with_origin_paths=args.model_id_with_origin_paths,\n",
    "        tokenizer_path=args.tokenizer_path,\n",
    "        trainable_models=args.trainable_models,\n",
    "        lora_base_model=args.lora_base_model,\n",
    "        lora_target_modules=args.lora_target_modules,\n",
    "        lora_rank=args.lora_rank,\n",
    "        lora_checkpoint=args.lora_checkpoint,\n",
    "        preset_lora_path=args.preset_lora_path,\n",
    "        preset_lora_model=args.preset_lora_model,\n",
    "        use_gradient_checkpointing=args.use_gradient_checkpointing,\n",
    "        use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload,\n",
    "        extra_inputs=args.extra_inputs,\n",
    "        fp8_models=args.fp8_models,\n",
    "        offload_models=args.offload_models,\n",
    "        template_model_id_or_path=args.template_model_id_or_path,\n",
    "        resume_from_checkpoint=args.resume_from_checkpoint,\n",
    "        remove_prefix_in_ckpt=args.remove_prefix_in_ckpt,\n",
    "        enable_lora_hot_loading=args.enable_lora_hot_loading,\n",
    "        task=args.task,\n",
    "        device=\"cpu\" if (args.initialize_model_on_cpu or args.enable_model_cpu_offload) else accelerator.device,\n",
    "    )\n",
    "    model_logger = ModelLogger(\n",
    "        args.output_path,\n",
    "        remove_prefix_in_ckpt=args.remove_prefix_in_ckpt,\n",
    "        enable_tensorboard_log=args.enable_tensorboard_log,\n",
    "        enable_swanlab_log=args.enable_swanlab_log,\n",
    "        swanlab_project=args.swanlab_project,\n",
    "        enable_wandb_log=args.enable_wandb_log,\n",
    "        wandb_project=args.wandb_project,\n",
    "    )\n",
    "    launcher_map = {\n",
    "        \"sft:data_process\": launch_data_process_task,\n",
    "        \"direct_distill:data_process\": launch_data_process_task,\n",
    "        \"sft\": launch_training_task,\n",
    "        \"sft:train\": launch_training_task,\n",
    "        \"direct_distill\": launch_training_task,\n",
    "        \"direct_distill:train\": launch_training_task,\n",
    "    }\n",
    "    launcher_map[args.task](accelerator, dataset, model, model_logger, args=args)\n",
    "\"\"\".strip()\n",
    "with open(\"train.py\", \"w\", encoding=\"utf-8\") as f:\n",
    "    f.write(code)\n",
    "\n",
    "# 启动训练任务\n",
    "cmd = \"\"\"\n",
    "accelerate launch train.py \\\n",
    "  --dataset_base_path data/dataset/1770381050168240056 \\\n",
    "  --dataset_metadata_path data/dataset/metadata.json \\\n",
    "  --extra_inputs \"template_inputs\" \\\n",
    "  --max_pixels 1048576 \\\n",
    "  --dataset_repeat 1 \\\n",
    "  --model_id_with_origin_paths \"black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors\" \\\n",
    "  --template_model_id_or_path \"DiffSynth-Studio/Template-KleinBase4B-Brightness:\" \\\n",
    "  --tokenizer_path \"black-forest-labs/FLUX.2-klein-4B:tokenizer/\" \\\n",
    "  --learning_rate 1e-4 \\\n",
    "  --num_epochs 1 \\\n",
    "  --remove_prefix_in_ckpt \"pipe.template_model.\" \\\n",
    "  --output_path \"models/template_brightness_training\" \\\n",
    "  --trainable_models \"template_model\" \\\n",
    "  --use_gradient_checkpointing \\\n",
    "  --find_unused_parameters \\\n",
    "  --fp8_models \"black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors\"\n",
    "\"\"\"\n",
    "os.system(cmd)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e91a608",
   "metadata": {},
   "source": [
    "训练完成后，将得到的权重与前面写好的模型定义一起打包到 `models/template_brightness` 目录，形成一个完整的 Template 模型："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3dd83dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "\n",
    "shutil.copy(\n",
    "    \"models/template_brightness_training/epoch-0.safetensors\",\n",
    "    \"models/template_brightness/model.safetensors\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1619f72c",
   "metadata": {},
   "source": [
    "加载训练好的模型，通过传入不同的 `scale` 数值生成明暗不同的图像："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5acc60f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = TemplatePipeline.from_pretrained(\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    device=\"cuda\",\n",
    "    model_configs=[ModelConfig(\"models/template_brightness\")],\n",
    "    lazy_loading=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "427b288d",
   "metadata": {},
   "outputs": [],
   "source": [
    "image = template(\n",
    "    pipe,\n",
    "    prompt=\"A cat is sitting on a stone.\",\n",
    "    seed=0, cfg_scale=4, num_inference_steps=50,\n",
    "    template_inputs = [{\"scale\": 0.7}],\n",
    "    negative_template_inputs = [{\"scale\": 0.5}]\n",
    ")\n",
    "image.save(\"image_Brightness_light.jpg\")\n",
    "image = template(\n",
    "    pipe,\n",
    "    prompt=\"A cat is sitting on a stone.\",\n",
    "    seed=0, cfg_scale=4, num_inference_steps=50,\n",
    "    template_inputs = [{\"scale\": 0.3}],\n",
    "    negative_template_inputs = [{\"scale\": 0.5}]\n",
    ")\n",
    "image.save(\"image_Brightness_dark.jpg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "603a288e",
   "metadata": {},
   "outputs": [],
   "source": [
    "show_images([\n",
    "    Image.open(\"image_Brightness_light.jpg\"),\n",
    "    Image.open(\"image_Brightness_dark.jpg\"),\n",
    "], resolution=256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08d3ff52",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "class",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}