28、创建基于Stable Diffusion模型的分布式文本到图像AI系统（上）

创建基于Stable Diffusion模型的分布式文本到图像AI系统（上）

在当今的AI领域，图像生成模型如DALL - E、Midjourney等吸引了全球的目光。这些模型能够根据简单的文本提示生成全新的图像，而Stable Diffusion作为2022年发布的流行图像生成模型，为我们构建文本到图像的AI系统提供了强大的支持。下面将详细介绍如何使用Stable Diffusion模型创建一个分布式的文本到图像AI系统。

技术要求

• Python虚拟环境 ：需要一个Python虚拟环境，可参考常见的Python开发环境搭建方法来创建。
• 硬件要求 ：为了正确运行Stable Diffusion模型，建议使用配备至少16GB RAM的现代计算机，理想情况下有一个8GB VRAM的专用GPU。对于Mac用户，配备M1 Pro或M2 Pro芯片的机型也很合适。如果没有这样的机器也没关系，只是图像生成会变慢且效果不佳。
• Redis服务器 ：运行工作进程需要在本地计算机上启动一个Redis服务器。最简单的方法是将其作为Docker容器运行。如果从未使用过Docker，可阅读官方文档中的入门教程（https://docs.docker.com/get-started/ ）。启动Redis服务器的命令如下：

$ docker run -d --name worker-redis -p 6379:6379 redis

• 代码示例 ：所有代码示例可在GitHub仓库（https://github.com/PacktPublishing/Building-Data-Science-Applications-with-FastAPI-Second-Edition/tree/main/chapter14 ）中找到。

使用Stable Diffusion从文本提示生成图像

首先，安装所需工具：

(venv) $ pip install accelerate diffusers

安装完成后，就可以使用Hugging Face的扩散模型了。

在Python脚本中实现模型

以下是一个能够实例化模型并运行图像生成的类的实现，这里采用了延迟加载模式，将模型加载和图像生成方法分开。

加载模型（load_model方法） ：

text_to_image.py
class TextToImage:
    pipe: StableDiffusionPipeline | None = None
    def load_model(self) -> None:
        # Enable CUDA GPU
        if torch.cuda.is_available():
            device = "cuda"
        # Enable Apple Silicon (M1) GPU
        elif torch.backends.mps.is_available():
            device = "mps"
        # Fallback to CPU
        else:
            device = "cpu"
        pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
        pipe.to(device)
        self.pipe = pipe

此方法首先会根据计算机的硬件情况，选择最有效的方式来运行模型。优先检查是否有CUDA（NVIDIA GPU）或MPS（Apple Silicon）设备可用，如果没有则回退到CPU。然后创建一个Hugging Face提供的StableDiffusionPipeline管道，并指定要从Hugging Face hub下载的模型，这里选择了“runwayml/stable - diffusion - v1 - 5”，其详细信息可在https://huggingface.co/runwayml/stable - diffusion - v1 - 5 查看。

生成图像（generate方法） ：

    def generate(
        self,
        prompt: str,
        *,
        negative_prompt: str | None = None,
        num_steps: int = 50,
        callback: Callable[[int, int, torch.FloatTensor], None] | None = None,
    ) -> Image.Image:
        if not self.pipe:
            raise RuntimeError("Pipeline is not loaded")
        return self.pipe(
            prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_steps,
            guidance_scale=9.0,
            callback=callback,
        ).images[0]

这个方法接受四个参数：
- prompt ：描述要生成图像的文本提示。
- negative_prompt ：可选提示，用于告诉模型绝对不想要的内容。
- num_steps ：模型应运行的推理步骤数。步骤越多，图像质量越好，但每次迭代都会延迟推理时间。默认值50在速度和质量之间提供了较好的平衡。
- callback ：可选函数，在每次迭代步骤时调用，有助于了解生成进度，并可能执行更多逻辑，如将进度保存到数据库。

方法签名中的星号（*）表示其后的参数只能作为关键字参数传递，例如： .generate("PROMPT", negative_prompt="NEGATIVE", num_steps = 10) 。这有助于使函数更加清晰和自解释。

将这些参数传递给 pipe 对象，它能够为每个提示生成多个图像，操作结果是一个Pillow图像列表，这里默认只生成一张图像，所以直接返回第一个。

执行Python脚本

在示例的底部添加了一个小的主脚本，方便你亲自尝试：

if __name__ == "__main__":
    text_to_image = TextToImage()
    text_to_image.load_model()
    def callback(step: int, _timestep, _tensor):
        print(f"🚀 Step {step}")
    image = text_to_image.generate(
        "A Renaissance castle in the Loire Valley",
        negative_prompt="low quality, ugly",
        callback=callback,
    )
    image.save("output.png")

第一次运行此脚本时，Hugging Face会将几GB的文件下载到计算机上，这就是Stable Diffusion模型。推理开始后，会看到一个进度条显示还剩多少推理步骤，同时会输出回调函数的打印语句。

不同计算机生成单张图像所需的时间不同：
| 硬件情况 | 生成时间 | 内存使用 |
| ---- | ---- | ---- |
| 现代NVIDIA GPU（8GB RAM）或配备M1 Pro芯片的Mac | 约1分钟（50推理步骤） | 合理 |
| CPU | 约5 - 10分钟 | 最多16GB |

如果推理速度过慢，可以尝试减少 num_steps 参数。

创建Dramatiq工作进程并定义图像生成任务

由于图像生成操作可能需要几分钟且消耗大量内存，不能直接在REST API服务器上运行图像生成模型。因此，我们需要定义一个独立于服务器进程的工作进程来处理图像生成任务。

在Web开发中，工作进程是一个在后台持续运行、等待传入任务的进程，任务通常由Web服务器根据用户操作发送。为了实现Web服务器和工作进程之间的通信，需要一个队列，它接收并堆叠来自Web服务器的消息，供工作进程读取，这就是Web - 队列 - 工作进程架构，其架构图如下：

graph LR
    A[Web服务器] --> B[队列]
    B --> C[工作进程]
    C --> D[生成图像]

在Python中，我们使用Dramatiq这个轻量级但强大的现代后台任务处理库，并使用Redis作为消息代理。

实现工作进程

首先安装所需依赖：

(venv) $ pip install "dramatiq[redis]"

设置Dramatiq工作进程涉及两件事：
1. 设置代理类型和URL。
2. 使用 @dramatiq.actor 装饰器包装函数来定义任务。

由于需要加载庞大的Stable Diffusion模型，我们实现了一个中间件，以便在工作进程启动时加载模型：

worker.py
class TextToImageMiddleware(Middleware):
    def __init__(self) -> None:
        super().__init__()
        self.text_to_image = TextToImage()
    def after_process_boot(self, broker):
        self.text_to_image.load_model()
        return super().after_process_boot(broker)

text_to_image_middleware = TextToImageMiddleware()
redis_broker = RedisBroker(host="localhost")
redis_broker.add_middleware(text_to_image_middleware)
dramatiq.set_broker(redis_broker)

定义了一个 TextToImageMiddleware 类，继承自Dramatiq的 Middleware 类，关键的 after_process_boot 方法会在工作进程启动后加载Stable Diffusion模型。

接下来定义图像生成任务：

@dramatiq.actor()
def text_to_image_task(
    prompt: str, *, negative_prompt: str | None = None, num_steps: int = 50
):
    image = text_to_image_middleware.text_to_image.generate(
        prompt, negative_prompt=negative_prompt, num_steps=num_steps
    )
    image.save(f"{uuid.uuid4()}.png")

Dramatiq任务实际上是用 @dramatiq.actor 装饰的普通函数。需要注意的是，任务参数必须是可序列化的数据，因为在从服务器调度任务时，参数会被存储在队列存储中，Dramatiq会将参数内部序列化为JSON。函数体调用 TextToImage 实例生成图像，并将其保存到磁盘，使用UUID作为文件名以避免文件覆盖。

启动工作进程

在启动工作进程之前，确保已经按照技术要求部分启动了Redis服务器。然后使用以下命令启动Dramatiq工作进程：

(venv) $ dramatiq -p 1 -t 1 chapter14.basic.worker

Dramatiq提供了命令行工具来启动工作进程，主位置参数是工作模块的点路径，类似于Uvicorn的使用方式。还设置了两个可选参数 -p 和 -t ，分别控制Dramatiq启动的进程数和线程数。默认情况下，它会启动10个进程，每个进程有8个线程，共80个工作进程可拉取和执行任务，但对于Stable Diffusion模型不适用，原因如下：
- 线程安全问题 ：进程中的每个线程共享相同的内存空间，如果两个或更多线程尝试生成图像，会在内存中对相同对象进行读写操作，导致并发问题，因此每个进程应只启动一个线程（ -t 1 ）。
- 内存占用问题 ：每个进程都需要在内存中加载模型，如果启动8个进程，会加载模型8次，这会占用大量内存，可能导致计算机内存不足，所以使用 -p 1 只启动一个进程。如果想尝试并行化，可以尝试 -p 2 启动两个进程，但要确保计算机能够处理。

运行上述命令后，会看到Stable Diffusion管道检查模型文件是否下载的输出，这表明模型已正确加载。

创建基于Stable Diffusion模型的分布式文本到图像AI系统（下）

存储和服务文件在对象存储中

在生成图像之后，我们需要考虑如何存储和服务这些文件。对象存储是一种可扩展且可靠的解决方案，适合存储大量的图像文件。

选择对象存储服务

常见的对象存储服务有Amazon S3、Google Cloud Storage、Microsoft Azure Blob Storage等。这里以Amazon S3为例进行说明。

配置AWS凭证

首先，需要在本地配置AWS凭证，以便能够与S3服务进行交互。可以通过AWS CLI进行配置：

aws configure

按照提示输入AWS Access Key ID、AWS Secret Access Key、默认区域和输出格式。

上传图像到S3

在Python中，可以使用 boto3 库来与S3进行交互。以下是一个上传图像到S3的示例代码：

import boto3

def upload_image_to_s3(image_path, bucket_name, object_name):
    s3 = boto3.client('s3')
    try:
        s3.upload_file(image_path, bucket_name, object_name)
        print(f"Image uploaded to S3: {object_name}")
    except Exception as e:
        print(f"Error uploading image to S3: {e}")

# 示例调用
image_path = "output.png"
bucket_name = "your-bucket-name"
object_name = "generated_images/output.png"
upload_image_to_s3(image_path, bucket_name, object_name)

服务图像

上传到S3的图像可以通过生成预签名URL来提供服务。预签名URL允许在一定时间内访问私有对象。以下是生成预签名URL的示例代码：

import boto3

def generate_presigned_url(bucket_name, object_name, expiration=3600):
    s3 = boto3.client('s3')
    try:
        response = s3.generate_presigned_url('get_object',
                                             Params={'Bucket': bucket_name,
                                                     'Key': object_name},
                                             ExpiresIn=expiration)
        return response
    except Exception as e:
        print(f"Error generating presigned URL: {e}")
        return None

# 示例调用
bucket_name = "your-bucket-name"
object_name = "generated_images/output.png"
presigned_url = generate_presigned_url(bucket_name, object_name)
print(f"Presigned URL: {presigned_url}")

构建API让用户生成自己的图像

现在我们已经有了图像生成的核心功能和工作进程，接下来需要构建一个API，让用户可以通过发送文本提示来生成自己的图像。

使用FastAPI构建API

FastAPI是一个快速（高性能）的Python Web框架，适合构建API。以下是一个简单的FastAPI应用示例：

from fastapi import FastAPI
import dramatiq
from worker import text_to_image_task

app = FastAPI()

@app.post("/generate-image/")
async def generate_image(prompt: str, negative_prompt: str = None, num_steps: int = 50):
    task = text_to_image_task.send(prompt, negative_prompt=negative_prompt, num_steps=num_steps)
    return {"task_id": task.message_id}

运行FastAPI应用

使用Uvicorn来运行FastAPI应用：

uvicorn main:app --reload

调用API

可以使用 curl 或其他工具来调用API：

curl -X POST "http://127.0.0.1:8000/generate-image/" -H "Content-Type: application/json" -d '{"prompt": "A beautiful sunset", "negative_prompt": "blurry", "num_steps": 60}'

API会返回一个任务ID，通过这个任务ID可以跟踪任务的状态。

监控和扩展系统

为了确保系统的可靠性和可扩展性，我们需要对系统进行监控，并在必要时进行扩展。

监控任务状态

Dramatiq提供了一些工具来监控任务的状态。可以使用 dramatiq-dashboard 来可视化任务的状态和统计信息。安装 dramatiq-dashboard ：

pip install dramatiq-dashboard

启动 dramatiq-dashboard ：

dramatiq-dashboard -b redis://localhost:6379

在浏览器中访问 http://localhost:5000 ，可以看到任务的状态和统计信息。

扩展工作进程

如果系统的负载增加，可以通过增加工作进程的数量来扩展系统。可以使用 -p 参数来增加进程数：

(venv) $ dramatiq -p 2 -t 1 chapter14.basic.worker

同时，要确保计算机有足够的资源来处理增加的负载。

总结

通过以上步骤，我们成功创建了一个分布式的文本到图像AI系统。从使用Stable Diffusion模型从文本提示生成图像，到创建Dramatiq工作进程处理图像生成任务，再到存储和服务图像在对象存储中，最后构建API让用户可以生成自己的图像。整个系统具有可靠性和可扩展性，能够满足不同规模的需求。

以下是整个系统的流程图：

graph LR
    A[用户] --> B[API服务器]
    B --> C[队列]
    C --> D[工作进程]
    D --> E[Stable Diffusion模型]
    E --> F[生成图像]
    F --> G[对象存储（如S3）]
    G --> H[返回图像URL给用户]

在实际应用中，可以根据具体需求对系统进行进一步的优化和扩展，例如添加更多的模型参数调整、优化任务调度算法等。希望这个教程能帮助你构建出自己的文本到图像AI系统。