【大模型系列 02】Diffusion For PyTorch 昇腾迁移

源码链接

https://gitee.com/ascend/ModelZoo-PyTorch/raw/master/PyTorch/built-in/diffusion/diffusers

Diffusers for PyTorch

概述

简述

扩散模型 (Diffusion Models) 是一种生成模型，可生成各种各样的高分辨率图像。Diffusers 是Huggingface发布的模型套件，包含基于扩散模型的多种下游任务训练与推理，可用于生成图像、音频，甚至分子的 3D 结构。

• 参考实现：

  url=https://github.com/huggingface/diffusers
  commit_id=6bbde99c4ac3cfaae3973ef6413e0cf9a8189ddb
  

• 适配昇腾 AI 处理器的实现：

  url=https://gitee.com/ascend/ModelZoo-PyTorch.git
  code_path=PyTorch/built-in/diffusion/
  

准备训练环境

准备环境

• 当前模型支持的 PyTorch 版本和已知三方库依赖如下表所示。

  表 1 版本支持表

  Torch_Version	三方库依赖版本
  PyTorch 1.8	diffusers== 0.17.0、accelerate==0.19.0

• 环境准备指导。

  请参考《Pytorch框架训练环境准备》搭建torch环境。

• 安装依赖。

  在模型根目录下执行命令，安装模型对应PyTorch版本需要的依赖。

  pip install -e .                    # 安装diffusers
  cd examples/text_to_image/           # 根据下游任务安装对应依赖
  pip install -r requirements.txt  
  

• 修改三方库代码：

  # 1. 修改${python路径}/python3.7/site-packages/accelerate/accelerator.py的432行为：
  self.scaler = torch.npu.amp.GradScaler(**kwargs)
  
  # 2. 修改${python路径}/python3.7/site-packages/accelerate/utils/dataclasses.py，给类GradScalerKwargs添加属性：
  dynamic: bool = True
  

• 卸载safetensors（如有）：

  pip uninstall safetensors
  

• 安装Megatron-LM，参考链接。

准备数据集

1. 获取数据集。

   联网情况下，数据集会自动下载。

   无网络情况下，用户需自行获取arrow格式的pokemon数据集，并在shell启动脚本时传入--local_data_dir参数，参数值为本地数据集路径，填写一级目录，数据结构如下：

   $dataset
   ├── dataset_dict.json
   └── train
       ├── data-00000-of-00001.arrow
       ├── dataset_info.json
       └── state.json
   

   说明：
   该数据集的训练过程脚本只作为一种参考示例。

获取预训练模型

联网情况下，预训练模型会自动下载。无网络时，用户可访问huggingface官网自行下载，文件namespace如下：

CompVis/stable-diffusion-v1-4
runwayml/stable-diffusion-v1-5
stabilityai/stable-diffusion-2
stabilityai/stable-diffusion-2-1

获取对应的预训练模型后，在shell启动脚本时传入--model_name参数，参数值为本地预训练模型路径，填写一级目录。

开始训练

训练模型

本节以文生图下游任务为例，展示模型训练方法，其余下游任务controlnet、dreambooth、textual inversion等需要替换启动脚本。

1. 进入解压后的源码包根目录。

   cd /${模型文件夹名称} 
   

2. 运行训练脚本。

   该模型支持单机单卡训练和单机8卡训练。

   • 单机单卡训练

     bash test/train_full_1p_text_to_image.sh  # 单卡精度
     

   • 单机8卡训练

     bash test/train_full_8p_text_to_image.sh # 8卡精度
     
     bash test/train_full_8p_text_to_image_sd1-5_fp16.sh  # 8卡精度，SD1.5，fp16
     bash test/train_performance_8p_text_to_image_sd1-5_fp16.sh # 8卡性能，SD1.5，fp16
     

   模型训练脚本参数说明如下。

   train_text_to_image.py：
   --local_data_dir                    //本地数据集路径
   --pretrained_model_name_or_path     //预训练模型名称或者地址
   --dataset_name                      //数据集名称
   --dataset_config_name               //数据集配置     
   --train_data_dir                    //符合huggingface结构的训练数据集
   --image_column                      //图片所在列
   --caption_column                    //图片caption所在列
   --max_train_samples                 //最大训练样本数
   --validation_prompts                //验证提示词
   --output_dir                        //输出路径
   --resolution                        //分辨率
   --num_train_epochs                  //训练epoch数
   --gradient_accumulation_steps	    //梯度累计步数
   --mixed_precision				   //精度模式
   

   训练完成后，权重文件保存在test/output路径下，并输出模型训练精度和性能信息。

训练结果展示

表 2 训练结果展示表(910B，SD2.1）

NAME	clip_score(use_ema)	FPS	batch_size	AMP_Type	Torch_Version
1p-竞品A	\	1.313	1	fp32	1.13
8p-竞品A	0.321	13.278	4	fp32	1.13
1p-NPU-910B	\	1.312	1	fp32	1.8
8p-NPU-910B	0.319	13.389	4	fp32	1.8
8p-竞品A	\	54.000	3	fp16	1.13
8p-NPU-910	\	24.000	3	fp16	1.8

表3 训练支持场景

SD版本/AMP_Type	fp16	fp32
SD1.5	支持，需设置–mixd_precision=“fp16”	支持，需设置–mixd_precision=“no”
SD2.1	不支持	支持，同上

说明：

910A仅支持fp16训练，在训练时必须指定–mixd_precision=“fp16”；910B同时支持fp16与fp32训练。

推理

文生图

参考实现：(https://huggingface.co/docs/diffusers/using-diffusers/conditional_image_generation)

预训练模型准备

联网情况下，预训练模型会自动下载。无网络时，用户可访问huggingface官网自行下载，文件namespace如下：

CompVis/ldm-text2im-large-256

获得对应的预训练模型后，修改以下代码中的地址为本地地址即可

generator = DiffusionPipeline.from_pretrained("CompVis/ldm-text2im-large-256",torch_dtype=torch.float16)

运行推理

python test_infer/text-to-image.py

修改prompt等操作需要对代码进行修改

文本指导图生图

参考实现：(https://huggingface.co/docs/diffusers/using-diffusers/img2img)

预训练模型准备

联网情况下，预训练模型会自动下载。无网络时，用户可访问huggingface官网自行下载，文件namespace如下：

nitrosocke/Ghibli-Diffusion

获得对应的预训练模型后，修改以下代码中的地址为本地地址即可

pipe = StableDiffusionImg2ImgPipeline.from_pretrained("nitrosocke/Ghibli-Diffusion").to(device)

运行推理

修改test_infer/text-guide-img-to-img.py中url为本地图片地址

python test_infer/text-guide-img-to-img.py

修改prompt等操作需要对代码进行修改

文本指导图像修复

参考实现：(https://huggingface.co/docs/diffusers/using-diffusers/inpaint)

预训练模型准备

联网情况下，预训练模型会自动下载。无网络时，用户可访问huggingface官网自行下载，文件namespace如下：

runwayml/stable-diffusion-inpainting

获得对应的预训练模型后，修改以下代码中的地址为本地地址即可

pipeline = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting")

运行推理

python test_infer/text-guide-image-inpainting.py

修改prompt等操作需要对代码进行修改

纹理反转

参考实现：(https://huggingface.co/docs/diffusers/using-diffusers/textual_inversion_inference)

预训练模型准备

联网情况下，预训练模型会自动下载。无网络时，用户可访问huggingface官网自行下载，文件namespace如下：

runwayml/stable-diffusion-v1-5

获得对应的预训练模型后，修改以下代码中的地址为本地地址即可

pretrained_model_name = "runwayml/stable-diffusion-v1-5"

数据集准备

联网情况下，数据会自动下载。无网络时，用户可访问huggingface官网自行下载，文件namespace如下。用户也可参考该数据集自行准备数据集：

sd-concepts-library/cat-toy

获得对应的预训练模型后，修改以下代码中的地址为本地地址即可

repo_id = "sd-concepts-library/cat-toy"

运行推理

python  test_infer/textual-inversion.py

修改prompt等操作需要对代码进行修改

文本指导图像深度生成

参考实现：(https://huggingface.co/docs/diffusers/using-diffusers/depth2img)

预训练模型准备

联网情况下，预训练模型会自动下载。无网络时，用户可访问huggingface官网自行下载，文件namespace如下：

stabilityai/stable-diffusion-2-depth

获得对应的预训练模型后，修改以下代码中的地址为本地地址即可

pipe = StableDiffusionDepth2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-2-depth").to("npu")

运行推理

python test_infer/text-guide-depth-to-image.py

修改prompt等操作需要对代码进行修改

无条件图像生成

参考实现：(https://huggingface.co/docs/diffusers/using-diffusers/unconditional_image_generation)

预训练模型准备

联网情况下，预训练模型会自动下载。无网络时，用户可访问huggingface官网自行下载，文件namespace如下：

anton-l/ddpm-butterflies-128

获得对应的预训练模型后，修改以下代码中的地址为本地地址即可

generator = DiffusionPipeline.from_pretrained("anton-l/ddpm-butterflies-128")

运行推理

python test_infer/unconditional-image-generation.py

修改prompt等操作需要对代码进行修改

版本说明

变更

2023.06.20：首次发布。

FAQ

1. 使用训练后的权重推理，如果出现NSFW检测，需要在推理前，关闭模型中的NSFW检测，具体做法：

   1) 找到模型文件model_index.json，将其中的requires_safety_checker参数设置为false
   2) 删除safetychecker参数及其对应的参数值