微调稳定扩散模型并使用Intel Gaudi生成图像

在本教程中，我们将介绍如何微调稳定扩散模型并使用Intel Gaudi (HPU) 生成图像。此外，我们还将展示如何将现有的HPU示例改编为使用Ray。一旦您学习了如何进行适配，您可以通过“Ray-ify”来自optimum-habana/examples 和 Model References 的示例，轻松访问更多为HPU开发的模型和优化！

现在，让我们看看如何“Ray-ify”这个稳定扩散示例。

配置

运行此示例需要安装Gaudi/Gaudi2的节点。Gaudi和Gaudi2都有8个HPU。

我们建议使用预构建的容器来运行这些示例。要运行容器，您需要Docker。有关安装说明，请参阅 安装Docker引擎。

接下来，请按照 使用容器运行 的说明安装Gaudi驱动程序和容器运行时。

然后，启动Gaudi容器：

docker pull vault.habana.ai/gaudi-docker/1.15.1/ubuntu22.04/habanalabs/pytorch-installer-2.2.0:latest
docker run -it --runtime=habana -e HABANA_VISIBLE_DEVICES=all -e OMPI_MCA_btl_vader_single_copy_mechanism=none --cap-add=sys_nice --net=host --ipc=host vault.habana.ai/gaudi-docker/1.15.1/ubuntu22.04/habanalabs/pytorch-installer-2.2.0:latest

在容器内，克隆 Optimum-Habana 并安装依赖项：

git clone https://github.com/huggingface/optimum-habana.git
pip install ray[train,serve] optimum-habana
cd optimum-habana/
pip install -r examples/stable-diffusion/requirements.txt
pip install -r examples/stable-diffusion/training/requirements.txt

使用文本反演进行微调

首先，让我们开始微调。检查 examples/stable-diffusion/training/textual_inversion.py，该文件在HPU上微调Stable Diffusion模型。您可以参考这份文档，并尝试一次而不使用Ray。

为了在Ray上运行此脚本，我们需要进行一些更改。但别担心，这其实非常简单。基本上，我们只需要识别主要训练循环，并在TorchTrainer中运行它。

首先，检查文件末尾的这一块：

if __name__ == "__main__":
    main()

最初，如果使用多个工作节点，该脚本将由MPI启动。但是在Ray中，我们应该设置TorchTrainer并提供一个主函数，即本示例中的main()。

因此，进行这些更改变得非常简单：

if __name__ == "__main__":
    import ray
    from ray import train
    from ray.train import ScalingConfig, Checkpoint, CheckpointConfig, RunConfig
    from ray.train.torch import TorchTrainer, TorchConfig

    ray.init(address="auto")

    # 配置计算资源
    # 在ScalingConfig中，为每个工作节点要求一个HPU
    scaling_config = ScalingConfig(num_workers=1, resources_per_worker={"CPU": 1, "HPU": 1})
    # 在TorchConfig中将后端设置为hccl
    torch_config = TorchConfig(backend = "hccl")
    # 初始化一个Ray TorchTrainer
    trainer = TorchTrainer(
        train_loop_per_worker=main,
        torch_config=torch_config,
        scaling_config=scaling_config,
    )

    result = trainer.fit()

在我们尝试运行之前，我们需要检查main函数以确保它可以这样工作。浏览该函数后，明显它不接受任何输入参数，但它调用parse_args来获取所有配置。最初，这些配置是在命令行通过MPI设置的。但由于我们切换到Ray来启动工作节点，命令行参数不再可用。因此，parse_args应该在主程序中调用，并传递给main函数。

除此之外，没有其他必要的更改。通过插入以下代码，您现在可以在Ray上运行该脚本。

# 替换以下行：
# def main():
# args = parse_args()
# 用这些台词：
def main(config):
    args = config["args"]

# 替换以下行：
# if __name__ == "__main__":
#     main()
# 用这些台词：
if __name__ == "__main__":
    import ray
    from ray import train
    from ray.train import ScalingConfig, Checkpoint, CheckpointConfig, RunConfig
    from ray.train.torch import TorchTrainer, TorchConfig

    ray.init(address="auto")

    # 配置计算资源
    # 在ScalingConfig中，要求每个worker配备一个HPU。
    scaling_config = ScalingConfig(num_workers=1, resources_per_worker={"CPU": 1, "HPU": 1})
    # 将TorchConfig中的后端设置为hccl
    torch_config = TorchConfig(backend = "hccl")
    # 初始化一个 Ray TorchTrainer
    trainer = TorchTrainer(
        train_loop_per_worker=main,
		train_loop_config={"args": parse_args()},
        torch_config=torch_config,
        scaling_config=scaling_config,
    )

    result = trainer.fit()

最后一件事：记得在命令行参数中使用绝对路径。原因与我们将 parse_args 移出的原因类似，Ray 的工作节点不共享当前工作目录。现在，您可以在 Ray 上运行 Stable Diffusion 的微调了！ 示例命令：

python /root/optimum-habana/examples/stable-diffusion/training/textual_inversion.py \
  --pretrained_model_name_or_path runwayml/stable-diffusion-v1-5 \
  --train_data_dir "/root/cat" \
  --learnable_property object \
  --placeholder_token "<cat-toy>" \
  --initializer_token toy \
  --resolution 512 \
  --train_batch_size 4 \
  --max_train_steps 3000 \
  --learning_rate 5.0e-04 \
  --scale_lr \
  --lr_scheduler constant \
  --lr_warmup_steps 0 \
  --output_dir /tmp/textual_inversion_cat \
  --save_as_full_pipeline \
  --gaudi_config_name Habana/stable-diffusion \
  --throughput_warmup_steps 3

示例命令的输出：

训练已根据配置开始：
╭───────────────────────────────────────────────╮
│ 训练配置                                     │
├───────────────────────────────────────────────┤
│ train_loop_config/args   ...t_warmup_steps=3) │
╰───────────────────────────────────────────────╯
(RayTrainWorker pid=15683) 正在为以下设置进程组： env:// [rank=0, world_size=1]
(TorchTrainer pid=15530) 启动分布式工作进程：
(TorchTrainer pid=15530) - (ip=172.17.0.2, pid=15683) world_rank=0, local_rank=0, node_rank=0
(RayTrainWorker pid=15683) [警告|utils.py:185] 2024-05-09 05:21:13,961 >> optimum-habana v1.10.4 已为 SynapseAI v1.14.0 验证，但发现 habana-frameworks v1.15.1.15，这可能会导致不确定行为！
(RayTrainWorker pid=15683) [警告|utils.py:198] 2024-05-09 05:21:15,401 >> optimum-habana v1.10.4 已为 SynapseAI v1.14.0 验证，但驱动版本为 v1.15.0，这可能会导致不确定行为！
(RayTrainWorker pid=15683) /usr/local/lib/python3.10/dist-packages/diffusers/utils/outputs.py:63: 用户警告：torch.utils._pytree._register_pytree_node 已弃用。请使用 torch.utils._pytree.register_pytree_node 替代。
(RayTrainWorker pid=15683)   torch.utils._pytree._register_pytree_node(
(RayTrainWorker pid=15683) /usr/local/lib/python3.10/dist-packages/huggingface_hub/file_download.py:1132: 未来警告：`resume_download` 已弃用，并将在版本 1.0.0 中移除。下载通常会恢复。如果要强制新下载，请使用 `force_download=True`。
(RayTrainWorker pid=15683)   warnings.warn(
(RayTrainWorker pid=15683) 05/09/2024 05:21:15 - 信息 - __main__ - 分布式环境：MULTI_HPU  后端：hccl
(RayTrainWorker pid=15683) 进程数量：1
(RayTrainWorker pid=15683) 进程索引：0
(RayTrainWorker pid=15683) 本地进程索引：0
(RayTrainWorker pid=15683) 设备：hpu
(RayTrainWorker pid=15683) 
(RayTrainWorker pid=15683) 混合精度类型：bf16
(RayTrainWorker pid=15683) 
(RayTrainWorker pid=15683) {'timestep_spacing', 'rescale_betas_zero_snr', 'dynamic_thresholding_ratio', 'prediction_type', 'thresholding', 'sample_max_value', 'clip_sample_range'} 未在配置中找到。值将初始化为默认值。
(RayTrainWorker pid=15683) ============================= HABANA PT 桥接配置 =========================== 
(RayTrainWorker pid=15683)  PT_HPU_LAZY_MODE = 1
(RayTrainWorker pid=15683)  PT_RECIPE_CACHE_PATH = 
(RayTrainWorker pid=15683)  PT_CACHE_FOLDER_DELETE = 0
(RayTrainWorker pid=15683)  PT_HPU_RECIPE_CACHE_CONFIG = 
(RayTrainWorker pid=15683)  PT_HPU_MAX_COMPOUND_OP_SIZE = 9223372036854775807
(RayTrainWorker pid=15683)  PT_HPU_LAZY_ACC_PAR_MODE = 0
(RayTrainWorker pid=15683)  PT_HPU_ENABLE_REFINE_DYNAMIC_SHAPES = 0
(RayTrainWorker pid=15683) ---------------------------: 系统配置 :---------------------------
(RayTrainWorker pid=15683) CPU核心数 : 152
(RayTrainWorker pid=15683) CPU内存       : 1056440348 KB
(RayTrainWorker pid=15683) ------------------------------------------------------------------------------
(RayTrainWorker pid=15683) {'scaling_factor', 'force_upcast'} 未在配置中找到。值将初始化为默认值。
(RayTrainWorker pid=15683) {'addition_embed_type', 'resnet_out_scale_factor', 'transformer_layers_per_block', 'attention_type', 'conv_out_kernel', 'time_embedding_type', 'addition_embed_type_num_heads', 'dropout', 'only_cross_attention', 'projection_class_embeddings_input_dim', 'num_attention_heads', 'upcast_attention', 'reverse_transformer_layers_per_block', 'resnet_time_scale_shift', 'resnet_skip_time_act', 'time_embedding_act_fn', 'class_embeddings_concat', 'time_cond_proj_dim', 'num_class_embeds', 'mid_block_type', 'cross_attention_norm', 'addition_time_embed_dim', 'dual_cross_attention', 'mid_block_only_cross_attention', 'encoder_hid_dim_type', 'time_embedding_dim', 'class_embed_type', 'conv_in_kernel', 'timestep_post_act', 'encoder_hid_dim', 'use_linear_projection'} 未在配置中找到。值将初始化为默认值。
(RayTrainWorker pid=15683) 05/09/2024 05:21:20 - 信息 - __main__ - ***** 开始训练 *****
(RayTrainWorker pid=15683) 05/09/2024 05:21:20 - 信息 - __main__ -   示例数量 = 600
(RayTrainWorker pid=15683) 05/09/2024 05:21:20 - 信息 - __main__ -   训练轮数 = 20
(RayTrainWorker pid=15683) 05/09/2024 05:21:20 - 信息 - __main__ -   每个设备的瞬时批处理大小 = 4
(RayTrainWorker pid=15683) 05/09/2024 05:21:20 - 信息 - __main__ -   总训练批处理大小（包括并行、分布式和累积） = 4
(RayTrainWorker pid=15683) 05/09/2024 05:21:20 - 信息 - __main__ -   梯度累积步数 = 1
(RayTrainWorker pid=15683) 05/09/2024 05:21:20 - 信息 - __main__ -   总优化步骤 = 3000
步骤：   0%|          | 0/3000 [00:00<?, ?it/s]/usr/local/lib/python3.10/dist-packages/habana_frameworks/torch/hpu/__init__.py:158: 用户警告：torch.hpu.setDeterministic 已弃用，并将在下一个版本中移除。请使用 torch.use_deterministic_algorithms 替代。
(RayTrainWorker pid=15683)   warnings.warn(
(RayTrainWorker pid=15683) [2024-05-09 05:21:20,951] [信息] [real_accelerator.py:178:get_accelerator] 正在将 ds_accelerator 设置为 hpu（自动检测）
步骤：   0%|          | 1/3000 [00:51<42:39:29, 51.21s/it, loss=0.138, lr=0.002]
...
...
...
(RayTrainWorker pid=15683) 05/09/2024 05:40:25 - 信息 - __main__ - 状态已保存至 /tmp/textual_inversion_cat/checkpoint-3000
步骤： 100%|██████████| 3000/3000 [19:04<00:00,  3.20it/s, loss=0.0261, lr=0.002]05/09/2024 05:40:25 - 信息 - __main__ - 吞吐量 = 11.083350060523712 样本/秒
(RayTrainWorker pid=15683) 05/09/2024 05:40:25 - 信息 - __main__ - 训练运行时间 = 1081.622427743976 秒
(RayTrainWorker pid=15683) /usr/local/lib/python3.10/dist-packages/huggingface_hub/file_download.py:1132: 未来警告：`resume_download` 已弃用，并将在版本 1.0.0 中移除。下载通常会恢复。如果要强制新下载，请使用 `force_download=True`。
(RayTrainWorker pid=15683)   warnings.warn(
(RayTrainWorker pid=15683) {'use_habana', 'requires_safety_checker', 'image_encoder', 'use_hpu_graphs', 'gaudi_config', 'bf16_full_eval'} 未在配置中找到。值将初始化为默认值。
(RayTrainWorker pid=15683) 
加载管道组件...:   0%|          | 0/7 [00:00<?, ?it/s]从 runwayml/stable-diffusion-v1-5 的 `feature_extractor` 子文件夹加载 feature_extractor 作为 CLIPImageProcessor。
(RayTrainWorker pid=15683) 从 runwayml/stable-diffusion-v1-5 的 `safety_checker` 子文件夹加载 safety_checker 作为 StableDiffusionSafetyChecker。
(RayTrainWorker pid=15683) 
加载管道组件...: 100%|██████████| 7/7 [00:00<00:00, 26.90it/s]
(RayTrainWorker pid=15683) [信息|pipeline_utils.py:193] 2024-05-09 05:40:25,683 >> 正在使用 CPU 运行。
(RayTrainWorker pid=15683) 配置已保存至 /tmp/textual_inversion_cat/vae/config.json
(RayTrainWorker pid=15683) 模型权重已保存至 /tmp/textual_inversion_cat/vae/diffusion_pytorch_model.safetensors
(RayTrainWorker pid=15683) 配置已保存至 /tmp/textual_inversion_cat/unet/config.json
(RayTrainWorker pid=15683) 模型权重已保存至 /tmp/textual_inversion_cat/unet/diffusion_pytorch_model.safetensors
(RayTrainWorker pid=15683) 配置已保存至 /tmp/textual_inversion_cat/scheduler/scheduler_config.json
(RayTrainWorker pid=15683) 配置已保存至 /tmp/textual_inversion_cat/model_index.json
(RayTrainWorker pid=15683) 05/09/2024 05:40:31 - 信息 - __main__ - 正在保存嵌入
步骤：100%|██████████| 3000/3000 [19:10<00:00,  2.61it/s, loss=0.0261, lr=0.002]

训练在 2024-05-09 05:40:33 完成，共经过 0 次迭代。总运行时间：19分钟30秒
2024-05-09 05:40:33,116	信息 tune.py:1007 -- 将所有结果文件和实验状态的最新版本写入 '/root/ray_results/TorchTrainer_2024-05-09_05-21-02'，耗时 0.0022秒。

Ray的一个优点是它易于扩展。在这个例子中，我们可以通过更改ScalingConfig中的num_workers轻松地将训练扩展到多个工作节点。Torch的分布式环境将会在Ray中自动初始化。

在Ray上服务微调的模型

现在我们已经对稳定扩散模型进行了微调，我们可以将其用于图像生成。下面的代码加载微调后的模型并生成一张图像。

import torch
from optimum.habana.diffusers import GaudiStableDiffusionPipeline
model_id = "/tmp/textual_inversion_cat/"
pipe = GaudiStableDiffusionPipeline.from_pretrained(
  model_id,
  torch_dtype=torch.bfloat16,
  use_habana=True,
  use_hpu_graphs=True,
  gaudi_config="Habana/stable-diffusion",
)
prompt = "a <cat-toy> is dancing on the grass."
image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
image.save("cat-backpack.png")

我们可以很容易地使用Ray Serve将其部署为HTTP服务。以下代码是根据这个示例修改的。将其保存为gaudi_sd_deploy.py，并使用serve run gaudi_sd_deploy:entrypoint来启动Serve应用程序。

import torch
from optimum.habana.diffusers import GaudiStableDiffusionPipeline
from io import BytesIO
from fastapi import FastAPI
from fastapi.responses import Response

from ray import serve
from ray.serve.handle import DeploymentHandle


app = FastAPI()


@serve.deployment(num_replicas=1)
@serve.ingress(app)
class APIIngress:
    def __init__(self, diffusion_model_handle: DeploymentHandle) -> None:
        self.handle = diffusion_model_handle

    @app.get(
        "/imagine",
        responses={200: {"content": {"image/png": {}}}},
        response_class=Response,
    )
    async def generate(self, prompt: str, img_size: int = 512):
        assert len(prompt), "prompt parameter cannot be empty"

        image = await self.handle.generate.remote(prompt, img_size=img_size)
        file_stream = BytesIO()
        image.save(file_stream, "PNG")
        return Response(content=file_stream.getvalue(), media_type="image/png")


@serve.deployment(
    ray_actor_options={"resources": {"HPU": 1}}
)
class GaudiStableDiffusion:
    def __init__(self, model_id):
        self.pipe = GaudiStableDiffusionPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
            use_habana=True,
            use_hpu_graphs=True,
            gaudi_config="Habana/stable-diffusion",
        )

    def generate(self, prompt: str, img_size: int = 512):
        assert len(prompt), "prompt parameter cannot be empty"

        image = self.pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
		return image


entrypoint = APIIngress.bind(GaudiStableDiffusion.bind("/tmp/textual_inversion_cat/"))

成功部署此Serve应用程序后，运行以下代码以生成图像。

import requests

prompt = "a <cat-toy> is dancing on the grass."
input = "%20".join(prompt.split(" "))
resp = requests.get(f"http://127.0.0.1:8000/imagine?prompt={input}")
with open("output.png", 'wb') as f:
    f.write(resp.content)

这里是一个示例图像：

from IPython.display import Image
Image(filename='ouput.png')