Tạo một máy chủ

Các pipeline của Diffusers có thể được sử dụng như một engine suy luận cho server. Nó hỗ trợ các yêu cầu đồng thời và đa luồng để tạo ảnh, có thể được yêu cầu bởi nhiều người dùng cùng một lúc.

Hướng dẫn này sẽ cho bạn thấy cách sử dụng StableDiffusion3Pipeline trên một server, nhưng bạn có thể tự do sử dụng bất kỳ pipeline nào bạn muốn.

Bắt đầu bằng cách điều hướng đến thư mục examples/server và cài đặt tất cả các phụ thuộc.

pip install .
pip install -f requirements.txt

Chạy máy chủ bằng lệnh sau.

python server.py

Máy chủ được truy cập tại http://localhost:8000. Bạn có thể gọi mô hình này bằng lệnh sau.

curl -X POST -H "Content-Type: application/json" --data '{"model": "something", "prompt": "a kitten in front of a fireplace"}' http://localhost:8000/v1/images/generations

Nếu bạn cần nâng cấp một số phụ thuộc, bạn có thể sử dụng pip-tools hoặc uv. Ví dụ, nâng cấp các phụ thuộc bằng uv bằng lệnh sau.

uv pip compile requirements.in -o requirements.txt

Máy chủ được xây dựng bằng FastAPI. Điểm cuối cho v1/images/generations được hiển thị bên dưới.

@app.post("/v1/images/generations")
async def generate_image(image_input: TextToImageInput):
    try:
        loop = asyncio.get_event_loop()
        scheduler = shared_pipeline.pipeline.scheduler.from_config(shared_pipeline.pipeline.scheduler.config)
        pipeline = StableDiffusion3Pipeline.from_pipe(shared_pipeline.pipeline, scheduler=scheduler)
        generator = torch.Generator(device="cuda")  # or "mps", "xpu", "cpu"
        generator.manual_seed(random.randint(0, 10000000))
        output = await loop.run_in_executor(None, lambda: pipeline(image_input.prompt, generator = generator))
        logger.info(f"output: {output}")
        image_url = save_image(output.images[0])
        return {"data": [{"url": image_url}]}
    except Exception as e:
        if isinstance(e, HTTPException):
            raise e
        elif hasattr(e, 'message'):
            raise HTTPException(status_code=500, detail=e.message + traceback.format_exc())
        raise HTTPException(status_code=500, detail=str(e) + traceback.format_exc())

Hàm generate_image được định nghĩa là bất đồng bộ với từ khóa async để FastAPI biết rằng những gì xảy ra trong hàm này không nhất thiết sẽ trả về kết quả ngay lập tức. Khi gặp một điểm trong hàm mà cần chờ một Task khác, luồng chính sẽ quay lại để xử lý các yêu cầu HTTP khác. Điều này được thể hiện trong đoạn mã bên dưới với từ khóa await.

output = await loop.run_in_executor(None, lambda: pipeline(image_input.prompt, generator = generator))

Tại thời điểm này, việc thực thi hàm pipeline được đặt lên một thread mới, và thread chính sẽ thực hiện các tác vụ khác cho đến khi nhận được kết quả từ pipeline.

Một khía cạnh quan trọng khác của việc triển khai này là tạo một pipeline từ shared_pipeline. Mục tiêu đằng sau điều này là tránh việc nạp mô hình cơ sở lên GPU nhiều hơn một lần, đồng thời vẫn cho phép mỗi yêu cầu mới đang chạy trên một thread riêng có bộ sinh (generator) và bộ lập lịch (scheduler) của riêng mình. Cụ thể, bộ lập lịch không an toàn cho thread, và sẽ gây ra các lỗi như: IndexError: index 21 is out of bounds for dimension 0 with size 21 nếu bạn cố gắng sử dụng cùng một bộ lập lịch trên nhiều thread.