tutorial.md

# SD2.1 MIGraphX Pipeline

使用 MIGraphX 后端，执行 SD2.1 文生图任务。

## 安装与环境配置

首先安装第三方依赖：
```bash
pip install -r requirements.txt
```

然后执行以下命令:
```bash
source /opt/dtk/env.sh
export PYTHONPATH=/opt/dtk/lib:$PYTHONPATH
python -c "import migraphx"
```

若回显 `ModuleNotFoundError: No module named 'migraphx'`，则说明没有安装 MIGraphX，请移步 [DAS资源下载目录](https://das.sourcefind.cn:55011/portal/#/installation?id=f0c8268b-756c-11ef-95cc-005056904552&type=frame) 下载安装包并安装 MIGraphX。


接下来，安装migraphx_diffusers：
```bash
pip install dist/migraphx_diffusers-${version}-py3-none-any.whl
```

最后，设置环境变量:
```bash
source set_env.sh /path/to/your/rocblas_lib
```

## 使用说明


### 模型准备

下载官方模型：
```bash
huggingface-cli download stabilityai/stable-diffusion-2-1-base --local-dir ./stable-diffusion-2-1-base --local-dir-use-symlinks False
```

导出 ONNX 模型：
```bash
# pip install optimum[exporters]
# optimum-cli export onnx --model ./stable-diffusion-2-1-base ./stable-diffusion-2-1-base-onnx --task text-to-image

python tools/export_onnxs.py --pipeline-dir stable-diffusion-2-1-base
```

模型优化：

```bash
bash tools/onnx_optimize.sh ./stable-diffusion-2-1-base ./stable-diffusion-2-1-base-opt
```


### 快速开始

```python
from diffusers import DiffusionPipeline
import migraphx_diffusers

pipe = DiffusionPipeline.from_pretrained(
    "./stable-diffusion-2-1-base-opt", 
    migraphx_config=migraphx_diffusers.DEFAULT_ARGS['sd2.1']
)
pipe.to("cuda")

image = pipe("sunflower").images[0]
image.save("sd2.1_output_512.png")
```

> :memo: **注意：**  在执行 `from_pretrained` 函数前必须执行 `import migraphx_diffusers`，使 MIGraphX 模型的加载逻辑生效。

生成结果：

<img src="assets/sd2.1_output_512.png" width="512" />

<br>

也可直接使用 `tools/run_pipe.py` 脚本进行快速的推理：

```bash
python tools/run_pipe.py -m ./stable-diffusion-2-1-base-opt -p "sunflower"
```

脚本参数说明：

| 参数 | 说明 | 类型 | 默认值 |
| --- | --- | --- | --- |
| `-m` / `--model-dir` | **必选**，pipeline 模型路径 | str | None |
| `--force-compile` | 可选，是否强制重新编译模型 | bool | False |
| `--num-images-per-prompt` | 可选，一条提示词一次生成图片的数量 | int | 1 |
| `--img-size` | 可选，生成图像尺寸，如果不设置，则跟随各 pipeline 默认的图像尺寸参数 | int | None | 
| `-p` / `--prompt` | **必选**，提示词，描述图片内容、风格、生成要求等 | str | None |
| `-n` / `--negative-prompt` | 可选，反向提示词，例如 "ugly" | str | None |
| `-t` / `--num-inference-steps` | 可选，生成图片时迭代多少步 | int | 50 |
| `--seed` | 可选，随机数种子 | int | 42 |
| `--save-prefix` | 可选，保存图片的前缀 | str | None |


### 参数配置

此项目使用 MIGraphX 静态 shape 推理，在加载各个子模型时，需要 `batch`、`img_size` 两个参数来固定模型的 shape；此外还需要 `model_dtype` 与 `force_compile` 两个参数，前者用来固定模型推理时采用的数据精度，后者用来控制当 .mxr 模型已经存在的情况下，是否强制重新编译。

+ `batch`: 模型输入的 batch 大小，默认为 1
+ `img_size`: 模型输入的图片尺寸，sdxl 默认为 1024，sd2.1 默认为 512
+ `model_dtype`: 模型推理时使用的数据精度，默认为 `fp16` 
+ `force_compile`: 是否强制重新编译模型，默认为 `false`

加载 pipeline 时，这些参数通过 `migraphx_config` 传递给 `DiffusionPipeline.from_pretrained` 函数，在此函数内部，程序会解析出以上四个参数，并在调用各个子模型的 `from_pretrained` 函数时传入这四个参数。

各 pipeline 默认的参数配置存放于文件 [migraphx_diffusers/default_args.json](migraphx_diffusers/default_args.json) 中，可通过 `migraphx_diffusers.DEFAULT_ARGS['sd2.1']` 获取 SD2.1 的默认参数配置，其内容如下：

```json
{
    "sd2.1":{
        "use_migraphx_models": [
            "text_encoder", 
            "unet", 
            "vae_decoder"
        ],
        "common_args": {
            "batch": 1,
            "img_size": 512,
            "model_dtype": "fp16",
            "force_compile": false
        },
        "model_args": {
            "text_encoder": {
                "batch": 1
            }
        }
    }
}
```

其中：  

+ `use_migraphx_models`: 指定在推理时，哪些子模型使用 MIGraphX 后端进行推理，此列表以外的子模型，使用 PyTorch 进行推理。
+ `common_args`: MIGraphX 加载各子模型的通用参数；
+ `model_args`: 若某个子模型不想按照通用参数加载模型，可在此字段中单独进行配置，例如：
    ```json
    "model_args": {
        "vae_decoder":{
            "model_dtype": "fp32",
            "force_compile": true
        } 
    }
    ```

> :memo: **注意：**  pipeline 的 batch 含义为  **1** 条 prompt 生成 batch 张图片。 prompt 数量决定文本编码器的 batch，也就是说，不管 pipeline 是单 batch 推理还是多 batch 推理，文本编码器的 batch 均为 1，所以在默认配置文件中，将 `text_encoder` 的 batch 固定为 1。


### 模型加载与缓存

各子模型加载流程如下：

<img src="assets/model_load.png">

加载与缓存逻辑：  
+ 若模型名在 `use_migraphx_models` 列表中，则使用 MIGraphX 后端：
  + 若存在 .mxr 缓存且 `force_compile` 为 false，则直接从 .mxr 缓存文件中加载模型；
  + 否则判断是否存在 .onnx 文件，若存在则解析 .onnx 文件并编译，且对编译后的模型保存为 .mxr 文件，不存在则抛出异常。
+ 否则使用 PyTorch 后端：
  + 若存在 .safetensors 文件，则从中加载权重，否则抛出异常。

### 目录结构与命名规则

整体的目录结构遵循 [huggingface](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/tree/main) 的格式，以 sd2.1 为例，其目录结构如下：

```
├── ./stable-diffusion-2-1-base-opt/
    ├── model_index.json
    ├── feature_extractor
    │   └── preprocessor_config.json
    ├── scheduler/
    │   └── scheduler_config.json
    ├── text_encoder/
    │   ├── config.json
    │   └── model.onnx
    ├── tokenizer/
    │   ├── merges.txt
    │   ├── special_tokens_map.json
    │   ├── tokenizer_config.json
    │   └── vocab.json
    ├── unet/
    │   ├── config.json
    │   ├── model.onnx
    │   └── model.onnx.data
    └── vae_decoder/
        ├── config.json
        └── model.onnx
```

每个子模型的 ONNX 文件均需命名为 model.onnx，若模型较大，可以只将模型结构存放于 model.onnx 内，而将权重数据外存为 model.onnx.data 文件。

若想使用编译好的 mxr 文件，则需要将 mxr 文件重命名为 `model_{model_dtype}_b{batch}_{img_size}_gpu.mxr`，然后将其放置于对应的子模型目录下。最后加载模型时，还应且将 `force_compile` 设置为 False。

如果某个子模型想使用 PyTorch 推理，则需要在 `use_migraphx_models` 列表中移除该子模型的名称，并在对应的子目录下放置 safetensors 文件，名称与 [huggingface](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/tree/main) 保持一致即可。

### 多 batch 推理

当前仅支持一个提示词同时生成一张图片，即 batch=1 推理，多 batch 推理功能适配中。

### 不同的图像尺寸

在默认配置中，sdxl生成图像尺寸为 1024x1024, 如果想要生成其他尺寸的图像，只需修改 `img_size` 字段即可。

示例：

```python
from diffusers import DiffusionPipeline
import migraphx_diffusers

migraphx_config = migraphx_diffusers.DEFAULT_ARGS['sd2.1']
migraphx_config['common_args']['img_size'] = 1024  # 修改图像尺寸

pipe = DiffusionPipeline.from_pretrained(
    "./stable-diffusion-2-1-base-opt/", 
    migraphx_config=migraphx_config
)
pipe.to("cuda")

image = pipe("sunflower").images[0]
image.save("sd2.1_output_1024.png")
```

生成结果：

<img src="assets/sd2.1_output_1024.png" width="512" />



## Tools

1. [以通用方式加载并运行pipeline](tools/README.md#以通用方式加载并运行pipeline)
2. [批量生成图片](tools/README.md#批量生成图片)
3. [统计各模块耗时](tools/README.md#统计各模块耗时)
4. [端到端性能测试](tools/README.md#SD2.1端到端性能测试)
5. [模型精度评估](tools/README.md#模型精度评估)


## 样例展示

SD2.1 生成样例：

<table>

<tr align="center" valign="middle">
<td>自然风景</td>
<td>人物肖像</td>
</tr>
<tr align="center" valign="middle">
<td><img src="examples/sd2.1-images-512/0-NatureScenery/theme_0_example_2_image_0.png" width="300" /></td>
<td><img src="examples/sd2.1-images-512/1-HumanPortrait/theme_1_example_3_image_0.png" width="300" /></td>
</tr>

<tr align="center" valign="middle">
<td>科幻与奇幻</td>
<td>生物</td>
</tr>
<tr align="center" valign="middle">
<td><img src="examples/sd2.1-images-512/2-ScienceFictionAndFantasy/theme_2_example_3_image_0.png" width="300" /></td>
<td><img src="examples/sd2.1-images-512/3-Creature/theme_3_example_2_image_0.png" width="300" /></td>
</tr>

<tr align="center" valign="middle">
<td>建筑与空间</td>
<td>抽象与艺术</td>
</tr>
<tr align="center" valign="middle">
<td><img src="examples/sd2.1-images-512/4-ArchitectureAndSpace/theme_4_example_1_image_0.png" width="300" /></td>
<td><img src="examples/sd2.1-images-512/5-AbstractionAndArt/theme_5_example_3_image_0.png" width="300" /></td>
</tr>

<tr align="center" valign="middle">
<td>日常生活</td>
<td>历史与复古</td>
</tr>
<tr align="center" valign="middle">
<td><img src="examples/sd2.1-images-512/6-DailyLife/theme_6_example_0_image_0.png" width="300" /></td>
<td><img src="examples/sd2.1-images-512/7-HistoryAndRetro/theme_7_example_2_image_0.png" width="300" /></td>
</tr>

<tr align="center" valign="middle">
<td>暗黑与怪诞</td>
<td>科技与数码</td>
</tr>
<tr align="center" valign="middle">
<td><img src="examples/sd2.1-images-512/8-DarkAndGrotesque/theme_8_example_0_image_0.png" width="300" /></td>
<td><img src="examples/sd2.1-images-512/9-TechnologyAndDigital/theme_9_example_1_image_0.png" width="300" /></td>
</tr>

</table>

更多生成示例: [examples](examples)


## 参考
1. [https://github.com/ROCm/AMDMIGraphX/tree/develop/examples/diffusion/python_stable_diffusion_xl](https://github.com/ROCm/AMDMIGraphX/tree/develop/examples/diffusion/python_stable_diffusion_xl)
2. [https://huggingface.co/docs/diffusers/using-diffusers/custom_pipeline_overview](https://huggingface.co/docs/diffusers/using-diffusers/custom_pipeline_overview)
3. [https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)