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

flink9streamer

于 2025-09-28 14:56:40 发布

阅读量34

点赞数

CC 4.0 BY-SA版权

分类专栏： FastAPI构建高效数据科学应用文章标签： Stable Diffusion 分布式系统文本到图像

本文链接：https://blog.youkuaiyun.com/flink9streamer/article/details/152360960

FastAPI构建高效数据科学应用专栏收录该内容

32 篇文章 ¥499.90

订阅专栏¥69.90

会员秒杀 ¥9.9 重磅福利

超级会员免费看

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

1. 工作进程中的任务调度

我们可以在工作进程中尝试调度任务。具体操作步骤如下：
1. 打开一个新的命令行，启动Python交互式shell（确保已启用Python虚拟环境）。
2. 运行以下命令：

(venv) $ python
>>> from chapter14.basic.worker import text_to_image_task
>>> text_to_image_task.send("A Renaissance castle in the Loire Valley")

执行上述代码后，会输出如下信息：

Message(queue_name='default', actor_name='text_to_image_task', 
args=('A Renaissance castle in the Loire Valley',), kwargs={}, 
options={'redis_message_id': '663df44a-cfc1-4f13-8457-05d8181290c1'}, 
message_id='bf57d112-6c20-49bc-a926-682ca43ea7ea', message_
timestamp=1675324585644)

这里使用了任务函数的 send 方法而不是直接调用它，这是告诉Dramatiq将任务发送到队列的方式。回到工作进程终端，会看到Stable Diffusion开始生成图像，片刻后图像将保存到磁盘。若在短时间内连续发送两个任务，Dramatiq会依次处理它们。

2. 实现REST API

为了让用户能与工作进程交互来调度任务，我们需要一个安全的接口，REST API是不错的选择，它能轻松集成到网站或移动应用等软件中。以下是实现图像生成任务发送到队列的简单API端点代码：

# api.py
from fastapi import FastAPI
from pydantic import BaseModel, Field
from uuid import UUID4
from dramatiq import Message
from chapter14.basic.worker import text_to_image_task
from fastapi import status

class ImageGenerationInput(BaseModel):
    prompt: str
    negative_prompt: str | None
    num_steps: int = Field(50, gt=0, le=50)

class ImageGenerationOutput(BaseModel):
    task_id: UUID4

app = FastAPI()

@app.post(
    "/image-generation",
    response_model=ImageGenerationOutput,
    status_code=status.HTTP_202_ACCEPTED,
)
async def post_image_generation(input: ImageGenerationInput) -> ImageGenerationOutput:
    task: Message = text_to_image_task.send(
        input.prompt, negative_prompt=input.negative_prompt, num_steps=input.num_steps
    )
    return ImageGenerationOutput(task_id=task.message_id)

此实现中，定义了合适的Pydantic模型来结构化和验证端点有效负载，数据会直接用于向Dramatiq发送任务。输出仅包含消息ID，HTTP状态码设置为202，表示请求已被接受，但处理尚未完成或开始。启动工作进程和此API后，就可以通过HTTP调用触发图像生成。

3. 数据库和对象存储中存储结果

虽然我们实现了后台工作进程进行大量计算和API来调度任务，但用户无法了解任务进度和获取最终结果。下面来解决这些问题。

3.1 工作进程和API之间的数据共享

工作进程在后台执行API请求的计算，但它无法与API服务器通信。因为可能存在多个服务器进程，甚至它们可能运行在不同物理服务器上，进程不能直接通信，需要一个中央数据源供进程读写数据。

一种解决方法是使用调度任务的同一代理，工作进程将结果写入代理，API从中读取。但这种方法有局限性，如Redis等代理不适合长期可靠存储数据，为限制内存使用，需要删除旧数据。

因此，我们选择使用数据库来存储图像生成请求和结果，实现工作进程和API之间的信息共享。

3.2 定义SQLAlchemy模型

首先定义一个SQLAlchemy模型来存储单个图像生成任务：

# models.py
from sqlalchemy.orm import Mapped, mapped_column
from sqlalchemy import Integer, DateTime, Text, String
from datetime import datetime
from sqlalchemy.orm import declarative_base

Base = declarative_base()

class GeneratedImage(Base):
    __tablename__ = "generated_images"
    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime, nullable=False, default=datetime.now
    )
    progress: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    prompt: Mapped[str] = mapped_column(Text, nullable=False)
    negative_prompt: Mapped[str | None] = mapped_column(Text, nullable=True)
    num_steps: Mapped[int] = mapped_column(Integer, nullable=False)
    file_name: Mapped[str | None] = mapped_column(String(255), nullable=True)

这里定义了自增ID作为主键，还添加了 prompt 、 negative_prompt 和 num_steps 列，对应工作进程任务的参数。 progress 列用于存储生成任务的当前进度， file_name 列用于存储系统中实际的文件名。

3.3 调整API以将图像生成任务保存到数据库

有了这个模型，API中调度图像生成的方法有所改变。不再直接将任务发送到工作进程，而是先在数据库中创建一行，使用该对象的ID作为工作进程任务的输入。端点实现如下：

# api.py
from fastapi import Depends, status
from sqlalchemy.ext.asyncio import AsyncSession
from chapter14.complete.models import GeneratedImage
from chapter14.complete import schemas
from chapter14.basic.worker import text_to_image_task
from sqlalchemy.future import select

@app.post(
    "/generated-images",
    response_model=schemas.GeneratedImageRead,
    status_code=status.HTTP_201_CREATED,
)
async def create_generated_image(
    generated_image_create: schemas.GeneratedImageCreate,
    session: AsyncSession = Depends(get_async_session),
):
    image = GeneratedImage(**generated_image_create.dict())
    session.add(image)
    await session.commit()
    text_to_image_task.send(image.id)
    return image

此代码中，将新创建对象的ID作为 text_to_image_task 的参数。用户请求的响应是 GeneratedImage 模型的表示，包含请求的提示信息和ID，用户可通过ID查询特定请求的数据和状态。

3.4 调整工作进程以从数据库读取和更新图像生成任务

需要更改任务的实现，使其从数据库中检索对象而不是直接读取参数。具体步骤如下：
1. 使用任务参数中的ID从数据库中检索 GeneratedImage 对象：

# worker.py
import dramatiq
from chapter14.complete.models import GeneratedImage
from sqlalchemy.ext.asyncio import async_session_maker
from sqlalchemy.future import select
import asyncio

def get_image(id: int) -> GeneratedImage:
    async def _get_image(id: int) -> GeneratedImage:
        async with async_session_maker() as session:
            select_query = select(GeneratedImage).where(GeneratedImage.id == id)
            result = await session.execute(select_query)
            image = result.scalar_one_or_none()
            if image is None:
                raise Exception("Image does not exist")
            return image
    return asyncio.run(_get_image(id))

@dramatiq.actor()
def text_to_image_task(image_id: int):
    image = get_image(image_id)

这里使用了 get_image 辅助函数，由于Dramatiq不能原生运行异步函数，所以使用 asyncio.run 手动调度异步函数的执行。

定义回调函数更新生成任务的进度：

# worker.py
def update_progress(image: GeneratedImage, step: int):
    async def _update_progress(image: GeneratedImage, step: int):
        async with async_session_maker() as session:
            image.progress = int((step / image.num_steps) * 100)
            session.add(image)
            await session.commit()
    asyncio.run(_update_progress(image, step))

@dramatiq.actor()
def text_to_image_task(image_id: int):
    image = get_image(image_id)
    def callback(step: int, _timestep, _tensor):
        update_progress(image, step)

调用 TextToImage 模型生成图像，生成随机文件名并上传到对象存储：

# worker.py
import uuid
from chapter14.complete.storage import Storage
from chapter14.complete import settings

@dramatiq.actor()
def text_to_image_task(image_id: int):
    image = get_image(image_id)
    def callback(step: int, _timestep, _tensor):
        update_progress(image, step)
    image_output = text_to_image_middleware.text_to_image.generate(
        image.prompt,
        negative_prompt=image.negative_prompt,
        num_steps=image.num_steps,
        callback=callback,
    )
    file_name = f"{uuid.uuid4()}.png"
    storage = Storage()
    storage.upload_image(image_output, file_name, settings.storage_bucket)
    update_file_name(image, file_name)

通过以上步骤，工作进程能从API接收图像生成任务，定期更新任务进度并设置结果文件名，用户可通过API的GET请求查看任务状态。

4. 对象存储中存储和提供文件

最后要解决生成图像的存储问题，需要可靠存储图像并让用户能从互联网轻松检索。传统的文件存储方式不适用于动态文件和复杂架构，因此引入对象存储。

对象存储将每个文件作为一个对象存储，包含实际数据和元数据，如名称、大小、类型和唯一ID。其优点是便于在多个物理机器上分布文件，用户只需请求特定文件，存储系统会负责从物理磁盘加载文件。

在云时代，对象存储很受欢迎，如Amazon S3的API已成为行业标准，大多数云对象存储都与之兼容，还有开源实现如MinIO。使用通用的S3 API，项目中可以使用相同的代码和库与任何对象存储提供商交互。

4.1 实现对象存储辅助类

使用Python的MinIO客户端库与任何S3兼容的存储交互，首先安装库：

(venv) $ pip install minio

然后实现一个类来处理对象存储操作：

# storage.py
import io
from minio import Minio
from datetime import timedelta
from PIL import Image

class Storage:
    def __init__(self) -> None:
        self.client = Minio(
            settings.storage_endpoint,
            access_key=settings.storage_access_key,
            secret_key=settings.storage_secret_key,
        )

    def ensure_bucket(self, bucket_name: str):
        bucket_exists = self.client.bucket_exists(bucket_name)
        if not bucket_exists:
            self.client.make_bucket(bucket_name)

    def upload_image(self, image: Image, object_name: str, bucket_name: str):
        self.ensure_bucket(bucket_name)
        image_data = io.BytesIO()
        image.save(image_data, format="PNG")
        image_data.seek(0)
        image_data_length = len(image_data.getvalue())
        self.client.put_object(
            bucket_name,
            object_name,
            image_data,
            length=image_data_length,
            content_type="image/png",
        )

    def get_presigned_url(
        self,
        object_name: str,
        bucket_name: str,
        *,
        expires: timedelta = timedelta(days=7)
    ) -> str:
        return self.client.presigned_get_object(
            bucket_name, object_name, expires=expires
        )

这个类包含以下方法：
- ensure_bucket ：确保对象存储中创建了正确的存储桶。
- upload_image ：将图像上传到存储桶，先将Pillow Image对象转换为字节流，再上传到指定存储桶。
- get_presigned_url ：生成具有临时访问权限的URL，用于安全地向用户提供文件访问权限。

4.2 在工作进程中使用对象存储辅助类

在任务实现中使用 Storage 类上传生成的图像：

# worker.py
from chapter14.complete.storage import Storage
from chapter14.complete import settings

@dramatiq.actor()
def text_to_image_task(image_id: int):
    # ...
    file_name = f"{uuid.uuid4()}.png"
    storage = Storage()
    storage.upload_image(image_output, file_name, settings.storage_bucket)
    # ...

4.3 在服务器上生成预签名URL

在API端实现一个新的端点，用于为给定的 GeneratedImage 生成预签名URL：

# server.py
from fastapi import Depends, HTTPException, status
from chapter14.complete.models import GeneratedImage
from chapter14.complete.storage import Storage
from chapter14.complete import schemas

@app.get("/generated-images/{id}/url")
async def get_generated_image_url(
    image: GeneratedImage = Depends(get_generated_image_or_404),
    storage: Storage = Depends(get_storage),
) -> schemas.GeneratedImageURL:
    if image.file_name is None:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Image is not available yet. Please try again later.",
        )
    url = storage.get_presigned_url(image.file_name, settings.storage_bucket)
    return schemas.GeneratedImageURL(url=url)

在生成URL前，会检查 GeneratedImage 对象的 file_name 属性是否设置，若未设置表示任务未完成。使用依赖注入获取 Storage 实例，确保代码的可维护性。

5. 运行图像生成系统

首先，需要设置项目的环境变量，包括数据库URL和S3凭证。为简单起见，使用SQLite数据库和MinIO playground进行S3存储：

# .env
DATABASE_URL=sqlite+aiosqlite:///chapter14.db
STORAGE_ENDPOINT=play.min.io
STORAGE_ACCESS_KEY=Q3AM3UQ867SPQQA43P2F
STORAGE_SECRET_KEY=zuf+tfteSlswRu7BJ86wekitnifILbZam1KYY3TG
STORAGE_BUCKET=fastapi-book-text-to-image

确保Redis服务器正在运行，然后依次运行FastAPI服务器和工作进程：

(venv) $ uvicorn chapter14.complete.api:app
(venv) $ dramatiq -p 1 -t 1 chapter14.complete.worker

系统准备好后，使用HTTPie发起请求启动新任务：

$ http POST http://localhost:8000/generated-images prompt="a sunset over a beach"

会得到如下响应：

{
    "created_at": "2023-02-13T08:24:45.954240",
    "file_name": null,
    "id": 1,
    "negative_prompt": null,
    "num_steps": 50,
    "progress": 0,
    "prompt": "a sunset over a beach"
}

表示新的图像生成任务已创建，进度为0%。可以通过以下请求查询任务状态：

http GET http://localhost:8000/generated-images/1

随着任务进行，进度会更新，任务完成后会显示文件名。当进度为100%时，可以请求生成预签名URL：

$ http GET http://localhost:8000/generated-images/1/url

会得到包含预签名URL的响应，用户可通过该URL访问生成的图像。

总结

通过以上步骤，我们创建了一个基于Stable Diffusion模型的分布式文本到图像AI系统，包括任务调度、REST API实现、数据库和对象存储的使用，用户可以通过API请求图像生成并查看任务状态和结果。整个系统的架构可以用以下mermaid流程图表示：

graph LR
    A[用户请求] --> B[FastAPI API]
    B --> C[Dramatiq任务队列]
    C --> D[工作进程]
    D --> E[Stable Diffusion生成图像]
    E --> F[数据库存储任务信息]
    E --> G[对象存储存储图像]
    B --> H[返回任务状态和结果]
    H --> A[用户获取结果]

这个系统的主要步骤总结如下表：
| 步骤 | 操作 | 代码文件 |
| ---- | ---- | ---- |
| 任务调度 | 在工作进程中调度任务 | worker.py |
| 实现REST API | 创建API端点处理图像生成请求 | api.py |
| 数据库存储 | 定义SQLAlchemy模型，调整API和工作进程与数据库交互 | models.py, api.py, worker.py |
| 对象存储 | 实现对象存储辅助类，上传图像和生成预签名URL | storage.py, worker.py, server.py |
| 运行系统 | 设置环境变量，启动服务器和工作进程，发起请求 | .env, uvicorn, dramatiq, httpie |

通过这些操作，我们构建了一个完整的图像生成系统，可在不同环境中运行并扩展。

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

6. 系统运行的详细流程和示例

在前面的内容中，我们已经介绍了系统各个部分的实现和配置。下面，让我们详细梳理一下系统运行的流程，并给出更多的示例。

6.1 系统运行流程

整个图像生成系统的运行流程可以概括为以下几个主要步骤：
1. 环境准备 ：设置环境变量，包括数据库URL和S3凭证，确保Redis服务器运行。
2. 启动服务 ：启动FastAPI服务器和工作进程。
3. 发起请求 ：用户通过HTTP请求向API发起图像生成任务。
4. 任务处理 ：API将任务信息存储到数据库，并将任务ID发送给工作进程。
5. 图像生成 ：工作进程从数据库获取任务信息，使用Stable Diffusion生成图像。
6. 结果存储 ：将生成的图像上传到对象存储，并更新数据库中的任务信息。
7. 结果获取 ：用户通过API查询任务状态和结果，获取预签名URL访问图像。

这个流程可以用以下mermaid流程图更直观地表示：

graph LR
    classDef startend fill:#F5EBFF,stroke:#BE8FED,stroke-width:2px
    classDef process fill:#E5F6FF,stroke:#73A6FF,stroke-width:2px
    classDef decision fill:#FFF6CC,stroke:#FFBC52,stroke-width:2px

    A([环境准备]):::startend --> B([启动服务]):::startend
    B --> C(用户发起请求):::process
    C --> D(API存储任务信息):::process
    D --> E(API发送任务ID给工作进程):::process
    E --> F(工作进程获取任务信息):::process
    F --> G(工作进程生成图像):::process
    G --> H(上传图像到对象存储):::process
    H --> I(更新数据库任务信息):::process
    I --> J{任务是否完成?}:::decision
    J -->|是| K(用户查询结果):::process
    J -->|否| F
    K --> L(获取预签名URL):::process
    L --> M([用户访问图像]):::startend

6.2 详细示例

以下是一个更详细的示例，展示如何一步步运行系统并获取图像。

步骤1：环境准备
创建 .env 文件，内容如下：

DATABASE_URL=sqlite+aiosqlite:///chapter14.db
STORAGE_ENDPOINT=play.min.io
STORAGE_ACCESS_KEY=Q3AM3UQ867SPQQA43P2F
STORAGE_SECRET_KEY=zuf+tfteSlswRu7BJ86wekitnifILbZam1KYY3TG
STORAGE_BUCKET=fastapi-book-text-to-image

步骤2：启动服务
启动FastAPI服务器：

(venv) $ uvicorn chapter14.complete.api:app

启动工作进程：

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

步骤3：发起请求
使用HTTPie发起图像生成请求：

$ http POST http://localhost:8000/generated-images prompt="A beautiful forest at sunset"

响应如下：

{
    "created_at": "2024-01-01T12:00:00.000000",
    "file_name": null,
    "id": 2,
    "negative_prompt": null,
    "num_steps": 50,
    "progress": 0,
    "prompt": "A beautiful forest at sunset"
}

步骤4：查询任务状态
可以定期查询任务状态：

http GET http://localhost:8000/generated-images/2

随着任务进行，进度会逐渐更新：

{
    "created_at": "2024-01-01T12:00:00.000000",
    "file_name": null,
    "id": 2,
    "negative_prompt": null,
    "num_steps": 50,
    "progress": 30,
    "prompt": "A beautiful forest at sunset"
}

步骤5：获取预签名URL
当任务完成，进度为100%时：

{
    "created_at": "2024-01-01T12:00:00.000000",
    "file_name": "12345678-1234-1234-1234-1234567890ab.png",
    "id": 2,
    "negative_prompt": null,
    "num_steps": 50,
    "progress": 100,
    "prompt": "A beautiful forest at sunset"
}

此时，发起请求获取预签名URL：

$ http GET http://localhost:8000/generated-images/2/url

响应如下：

{
    "url": "https://play.min.io/fastapi-book-text-to-image/12345678-1234-1234-1234-1234567890ab.png?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=Q3AM3UQ867SPQQA43P2F%2F20240101%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240101T123000Z&X-Amz-Expires=604800&X-Amz-SignedHeaders=host&X-Amz-Signature=abcdef1234567890abcdef1234567890abcdef12"
}

用户可以通过这个URL访问生成的图像。

7. 系统的优势和可扩展性

这个基于Stable Diffusion模型的分布式文本到图像AI系统具有以下几个显著的优势，并且具备良好的可扩展性。

7.1 系统优势

分布式架构 ：采用分布式架构，将任务调度、图像生成和数据存储分离，提高了系统的性能和可靠性。工作进程可以在不同的物理机器上运行，充分利用资源。
REST API接口 ：提供REST API接口，方便与其他软件集成，如网站、移动应用等。用户可以通过简单的HTTP请求发起图像生成任务。
数据存储分离 ：使用数据库存储任务信息，对象存储存储图像，避免了传统文件存储的局限性，便于数据的管理和维护。
安全访问 ：通过生成预签名URL的方式，确保用户只能在有限的时间内访问图像，提高了数据的安全性。

7.2 可扩展性

增加工作进程 ：可以根据需求增加工作进程的数量，提高系统的并发处理能力，处理更多的图像生成任务。
更换存储提供商 ：由于使用了通用的S3 API，很容易更换对象存储提供商，如从MinIO切换到Amazon S3或Google Cloud Storage。
扩展功能 ：可以在现有基础上扩展功能，如增加图像编辑、图像分类等功能，满足不同用户的需求。

8. 注意事项和常见问题解答

在使用这个系统的过程中，有一些注意事项和常见问题需要关注。

8.1 注意事项

环境变量配置 ：确保环境变量配置正确，特别是数据库URL和S3凭证。错误的配置可能导致系统无法正常运行。
Redis服务器 ：确保Redis服务器正常运行，它是任务队列的基础。
对象存储权限 ：确保对象存储的访问权限设置正确，避免出现文件无法上传或访问的问题。

8.2 常见问题解答

问题1：任务一直处于未开始状态
- 解答：检查Redis服务器是否正常运行，工作进程是否启动，以及数据库连接是否正常。
问题2：无法上传图像到对象存储
- 解答：检查对象存储的访问凭证是否正确，存储桶是否存在，以及网络连接是否正常。
问题3：预签名URL无法访问
- 解答：检查URL是否过期，对象存储的配置是否正确，以及文件是否存在于存储桶中。

总结回顾

通过本文，我们详细介绍了如何创建一个基于Stable Diffusion模型的分布式文本到图像AI系统。从任务调度、REST API实现，到数据库和对象存储的使用，再到系统的运行和扩展，我们全面地阐述了系统的各个方面。

整个系统的实现步骤可以总结如下：
1. 在工作进程中进行任务调度，使用Dramatiq将任务放入队列。
2. 实现REST API，处理用户的图像生成请求，并将任务信息存储到数据库。
3. 定义SQLAlchemy模型，用于数据库存储任务信息，调整API和工作进程与数据库交互。
4. 实现对象存储辅助类，将生成的图像上传到对象存储，并生成预签名URL供用户访问。
5. 配置环境变量，启动FastAPI服务器和工作进程，通过HTTP请求发起图像生成任务。

这个系统具有分布式架构、REST API接口、数据存储分离和安全访问等优势，并且具备良好的可扩展性。在使用过程中，需要注意环境变量配置、Redis服务器和对象存储权限等问题。

希望本文能帮助你构建和运行自己的图像生成系统，如果你有任何问题或建议，欢迎在评论区留言。

最后，再次用mermaid流程图回顾一下系统的整体架构：

graph LR
    A[用户请求] --> B[FastAPI API]
    B --> C[Dramatiq任务队列]
    C --> D[工作进程]
    D --> E[Stable Diffusion生成图像]
    E --> F[数据库存储任务信息]
    E --> G[对象存储存储图像]
    B --> H[返回任务状态和结果]
    H --> A[用户获取结果]

系统主要步骤总结如下表：
| 步骤 | 操作 | 代码文件 |
| ---- | ---- | ---- |
| 任务调度 | 在工作进程中调度任务 | worker.py |
| 实现REST API | 创建API端点处理图像生成请求 | api.py |
| 数据库存储 | 定义SQLAlchemy模型，调整API和工作进程与数据库交互 | models.py, api.py, worker.py |
| 对象存储 | 实现对象存储辅助类，上传图像和生成预签名URL | storage.py, worker.py, server.py |
| 运行系统 | 设置环境变量，启动服务器和工作进程，发起请求 | .env, uvicorn, dramatiq, httpie |

通过这些步骤和操作，你可以构建一个完整、可靠且可扩展的图像生成系统。