运行SD模型 Running Diffusion Models

在此示例中，我们演示了如何编写自己的Stable Diffusion工具，从终端运行它，并将其部署到 fal 作为可扩展的生产级 HTTP API，而无需更改任何代码。

从小处着手

我们可以从导入 fal 开始，并定义一个名为 MODEL_NAME 的全局变量来表示我们要使用哪个模型（它可以是任何 SD/SDXL 模型或任何经过微调的模型）来自 HF 模型中心：

import fal

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

然后继续使用缓存function[link]，当模型尚不存在时，它会将模型加载到 GPU 中。当我们连续多次调用我们的工具（或在短时间内命中大量 API 请求）时，这应该可以帮助我们节省大量时间

@fal.cached
def load_model(model_name: str = MODEL_NAME) -> object:
    """Load the stable diffusion model into the GPU and return the
    diffusers pipeline."""

    from diffusers import DiffusionPipeline

    pipe = DiffusionPipeline.from_pretrained(model_name)
    return pipe.to("cuda")

实施说明

应用程序可以通过将其作为输入参数并传递给 load_model 来延迟加载多个不同的 SD 模型。@fal.cached可以感知输入，因此如果您传递不同的输入（例如 SD1.5 与 SDXL），它实际上会重新加载代码并为您提供新模型（旧模型仍将保留在缓存中，以防您需要它）。

获取输入并返回输出

为了通过 serve=True 启用自动 Web 端点 fal 优惠，我们必须通过 Pydantic 以结构化的方式定义我们的输入和输出。虽然这看起来像是 Web 的东西，但实际上没有什么可以阻止您为 CLI 使用相同的 I/O，这正是我们要做的。

import fal
from fal.toolkit import Image
from pydantic import BaseModel

# [...]

class DiffusionOptions(BaseModel):
    prompt: str
    steps: int = 30

class Result(BaseModel):
    image: Image

稳定扩散应用程序

我们可以使用必要的软件包（本例中只有 diffusers 和 transformers）注释我们的推理函数，并通过设置 serve=True 将其标记为服务函数。此工作流程可以在不同的 GPU 上运行，请查看 https://fal.ai/pricing 以获取选项列表。

import fal

@fal.function(
    "virtualenv",
    requirements=[
        "diffusers[torch]",
        "transformers",
    ],
    serve=True,
    machine_type="GPU",
    keep_alive=60,
)
def run_stable_diffusion(options: DiffusionOptions) -> Result:
    # 加载Diffusers管道
    pipe = load_model()

    # 执行推理过程
    result = pipe(options.prompt, num_inference_steps=options.steps)

    # 上传图像并返回
    image = Image.from_pil(result.images[0])
    return Result(image=image)

推理逻辑本身应该很容易理解，但如果我们需要总结三个步骤，则每次调用时此函数都会执行：

获取包含实际模型的扩散器管道。虽然第一次调用会有点昂贵（约 15 秒），但所有后续缓存调用都将免费。
使用给定的选项运行管道以执行推理并生成图像。
将图像上传到 fal 的存储服务器，并返回结果对象。

在 CLI 中使用该应用程序

要尝试使用新的稳定扩散应用程序，您可以为其编写一个非常小的界面并开始在本地运行它。

from argparse import ArgumentParser

[...]

def main(argv: list[str] | None = None) -> None:
    parser = ArgumentParser()
    parser.add_argument("prompt", type=str)
    parser.add_argument("--steps", type=int, default=40)
    args = parser.parse_args(argv)

    local_diffusion = run_stable_diffusion.on(serve=False)
    result = local_diffusion(
        DiffusionOptions(
            prompt=args.prompt,
            steps=args.steps,
        )
    )
    print(
        f"Image generation is complete. Access your "
        f"image through the URL: {result.image.url}"
    )

if __name__ == "__main__":
    main()

您可能已经注意到，我们正在创建一个名为local_diffusion的新函数，方法是在通过 Python 执行调用时将serve属性设置为false。这样做是为了确保我们的应用程序在通过run_stable_diffusion()运行（或部署）时既可以作为 Web 应用程序运行，也可以通过 Python 进行交互。

$ python app.py "a cat on the moon" --steps=50
[...]
Image generation is complete. Access your image through the URL: $URL

生产化您的 API

要以 HTTP API 的形式与他人共享此应用程序，您所要做的就是调用 fal 的 serve 命令并让它将函数部署到无服务器运行时。每个 HTTP 请求都会自动唤醒服务器（如果还没有），处理请求，挂起一段时间以防有其他后续请求（在定义的 keep_alive 内），最后关闭自身以防止在空闲时产生费用。

bash

$ fal deploy t.py::run_stable_diffusion --app-name stable-diffusion
Registered a new revision for function 'stable-diffusion' (revision='[...]').
URL: https://fal.run/$USER/stable-diffusion

只要您设置了凭据[link]，您就可以从任何地方调用此 API（从您的终端或您自己的前端）：

bash

$ curl $APP_URL \
  -H 'Authorization: Key $FAL_KEY' \
  -H 'Content-Type: application/json' \
  -H 'Accept: application/json, */*;q=0.5' \
  -d '{"prompt":"a cat on the moon"}'

json

{
  "image": {
    "url": "...",
    "content_type": "image/png",
    "file_name": "...",
    "file_size": "...",
    "width": 512,
    "height": 512
  }
}

上次更新于 2024 年 6 月 21 日