Телеграм-бот для генерации изображений на Stable Diffusion в Docker

Телеграм-бот генерации изображений — это связка из трех частей: сам бот принимает текст от пользователя, отдельный сервис генерации запускает модель, а модель Stable Diffusion превращает промпт в картинку. В этой статье я разберу именно проектную связку: как разложить ее на сервисы, зачем тут Docker, как поднять Stable Diffusion через библиотеку diffusers и как пережить долгую генерацию, не заблокировав бота.

Это не туториал по базовому боту (обработка команд, токен от BotFather) — предполагаю, что с эхо-ботом вы уже знакомы. Фокус тут на архитектуре и на интеграции с тяжелой генеративной моделью.

Архитектура: три роли, два контейнера

Соблазн написать все в одном скрипте велик, но генерация изображения на GPU занимает секунды и десятки секунд, а Telegram ждет быстрых ответов. Поэтому роли разводят.

Бот и сервис генерации — это два разных процесса и два контейнера Docker. Бот легкий и всегда отзывчив, сервис генерации держит в памяти тяжелую модель. Общаются они по HTTP внутри Docker-сети: бот шлет POST с промптом, сервис возвращает PNG.

Такое разделение дает три вещи: бота можно перезапускать, не выгружая модель из видеопамяти; сервис генерации можно вынести на отдельную машину с GPU; нагрузку на генерацию видно и ей можно управлять отдельно от логики диалога.

Почему именно Docker

Stable Diffusion тянет за собой хрупкое окружение: конкретные версии PyTorch, CUDA-драйверов, diffusers, transformers. На соседней машине эта сборка легко ломается. Docker фиксирует ее целиком.

• Изоляция окружения. Версии CUDA и Python внутри образа не конфликтуют с системными.
• GPU-зависимости. Через NVIDIA Container Toolkit контейнер получает доступ к видеокарте (--gpus all или блок devices в compose), а образ строится от CUDA-базы.
• Воспроизводимость. Тот же образ поднимется на сервере так же, как у вас локально; веса модели выносим в volume, чтобы не тянуть их при каждой пересборке.

Как запустить Stable Diffusion: diffusers или AUTOMATIC1111

Есть два рабочих пути дать боту доступ к модели. Выбор зависит от того, нужен ли вам полный контроль из кода или готовый веб-интерфейс.

Ниже я иду по пути diffusers: это библиотека Hugging Face, вокруг которой удобно построить свой HTTP-сервис и не зависеть от чужого интерфейса. Исходная версия этой статьи обращалась к стороннему локальному API по адресу вида http://localhost:9000/ping с ручным опросом статуса — я заменил это на прямой сервис, где очередь и таймауты под контролем.

Минимальный сервис генерации

Начнем с законченного сервиса, который принимает промпт и возвращает картинку. Он серизализует работу на GPU: одна видеокарта считает одну генерацию за раз, поэтому запросы выстраиваем в очередь через asyncio.Lock.

import io
import asyncio
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel
import torch
from diffusers import StableDiffusionXLPipeline

app = FastAPI()
gpu_lock = asyncio.Lock()  # один GPU - генерации идут по очереди

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
).to("cuda")


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


def render(prompt: str, steps: int) -> bytes:
    image = pipe(prompt=prompt, num_inference_steps=steps).images[0]
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    return buf.getvalue()


@app.post("/generate")
async def generate(job: Job):
    if not 1 <= job.steps <= 50:
        raise HTTPException(status_code=422, detail="steps 1..50")
    async with gpu_lock:
        png = await asyncio.to_thread(render, job.prompt, job.steps)
    return Response(content=png, media_type="image/png")

Запустив сервис (uvicorn service:app --host 0.0.0.0 --port 9000) и отправив POST на /generate с телом {"prompt": "a red fox in the snow", "steps": 30}, вы получите обратно байты PNG. Разберем ключевые места.

• from_pretrained(...).to("cuda") один раз загружает веса SDXL в видеопамять при старте, а не на каждый запрос — загрузка занимает секунды и делать ее многократно нельзя.
• torch_dtype=torch.float16 держит модель в половинной точности: вдвое меньше видеопамяти при почти том же качестве.
• gpu_lock не дает двум запросам одновременно занять GPU, иначе получите ошибку нехватки видеопамяти.

Долгая генерация: не блокируем цикл событий

Тут прячется типовая ошибка. Кажется логичным вызвать модель прямо в асинхронном эндпоинте:

@app.post("/generate")
async def generate(job: Job):
    png = render(job.prompt, job.steps)   # блокирует весь процесс
    return Response(content=png, media_type="image/png")

Симптом: пока одна генерация считается 20 секунд, сервис не отвечает вообще ни на что — render синхронный и держит цикл событий asyncio. Проверки здоровья отваливаются по таймауту, а не только соседние запросы.

Исправление — увести тяжелый вызов в отдельный поток, вернув управление циклу событий:

async with gpu_lock:
    png = await asyncio.to_thread(render, job.prompt, job.steps)

Теперь цикл событий свободен: он обслуживает проверки здоровья и ставит следующие запросы в очередь на gpu_lock, пока поток считает картинку. Это и есть простая очередь: один воркер, запросы ждут блокировку по очереди. Для десятков одновременных пользователей одного лока мало — тогда переходят к схеме с идентификатором задачи: сервис сразу возвращает task_id, кладет работу в очередь (например, отдельный воркер или брокер вроде Redis), а бот периодически спрашивает статус. Для учебного проекта на одного-двух пользователей блокировки достаточно.

Бот: принять промпт, дождаться, отдать картинку

Бот держим тонким. Он не знает про модель — только шлет HTTP-запрос сервису и работает с таймаутом, ведь генерация долгая.

import os
import httpx
from telegram import Update
from telegram.ext import (
    Application, CommandHandler, MessageHandler, filters, ContextTypes,
)

TOKEN = os.environ["TELEGRAM_TOKEN"]          # токен только из окружения
SERVICE_URL = os.environ.get("SERVICE_URL", "http://generator:9000")


async def start(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text("Пришли текст - верну картинку.")


async def on_prompt(update: Update, ctx: ContextTypes.DEFAULT_TYPE):
    prompt = update.message.text.strip()
    if len(prompt) > 500:
        await update.message.reply_text("Слишком длинный промпт, до 500 символов.")
        return
    note = await update.message.reply_text("Генерирую, это занимает до минуты...")
    try:
        async with httpx.AsyncClient(timeout=180) as client:
            resp = await client.post(
                f"{SERVICE_URL}/generate",
                json={"prompt": prompt, "steps": 30},
            )
        resp.raise_for_status()
        await update.message.reply_photo(photo=resp.content)
    except httpx.TimeoutException:
        await update.message.reply_text("Не уложился в таймаут, попробуй позже.")
    finally:
        await note.delete()


def main():
    app = Application.builder().token(TOKEN).build()
    app.add_handler(CommandHandler("start", start))
    app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, on_prompt))
    app.run_polling()


if __name__ == "__main__":
    main()

Что тут важно. Таймаут httpx (180 секунд) сознательно больше, чем на генерацию картинки, но он есть: без него зависший сервис навсегда подвесит обработчик. Сообщение «Генерирую…» удаляется в finally при любом исходе. Ограничение длины промпта — это тоже защита: чужой ввод нельзя считать безопасным.

Токен и ключи — только через окружение

Токен бота — это доступ к нему целиком. Захардкоженный в коде токен утечет в историю git при первом же коммите. Правильный вариант — читать из переменной окружения (os.environ["TELEGRAM_TOKEN"]), а значение хранить вне репозитория.

# .env - добавлен в .gitignore, в репозиторий не попадает
TELEGRAM_TOKEN=123456:AA...ваш_токен

Тот же принцип для любых ключей (например, токена Hugging Face для приватных весов): переменные окружения и файл .env в .gitignore, а не строки в коде. Если токен все же засветился — отзовите его у BotFather и выпустите новый.

Собираем в Docker Compose

Compose поднимает оба сервиса одной командой и прокидывает GPU в контейнер генерации.

services:
  generator:
    build: ./generator
    environment:
      - HF_HOME=/models
    volumes:
      - ./models:/models          # веса кешируются между пересборками
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]

  bot:
    build: ./bot
    environment:
      - TELEGRAM_TOKEN=${TELEGRAM_TOKEN}
      - SERVICE_URL=http://generator:9000
    depends_on:
      - generator

Бот обращается к сервису по имени generator — это DNS-имя внутри Docker-сети, порт 9000 наружу можно не публиковать. Токен подставляется из .env через ${TELEGRAM_TOKEN}, в compose-файл его не пишем. Запуск: docker compose up --build.

Выводы

• Разведите роли: тонкий бот и отдельный сервис генерации — два контейнера, общаются по HTTP внутри Docker-сети.
• Docker решает конкретную боль Stable Diffusion: фиксирует связку CUDA + PyTorch + diffusers и делает сборку воспроизводимой; веса выносите в volume.
• diffusers дает полный контроль из кода, готовый API AUTOMATIC1111 — быстрый прототип; выбор по потребности в настройке.
• Генерация долгая: уводите вызов модели в поток (asyncio.to_thread), серизализуйте GPU через lock, на боте держите таймаут и статус-сообщение.
• Токен и ключи — только из переменных окружения; .env в .gitignore, длину и содержимое пользовательского промпта проверяйте.

Где применяется / связь с практикой

Такой проект — хороший вход в прикладной ML: вы касаетесь диффузионных моделей, инференса на GPU, упаковки модели в сервис и очередей. Это ровно те навыки, которые нужны, когда генеративную модель надо не просто запустить в блокноте, а довести до сервиса, который выдерживает нагрузку.

FAQ

Обязательна ли видеокарта NVIDIA?
Для приемлемой скорости — да, SDXL на CPU считается минуты и требует много оперативной памяти. Как учебный вариант CPU-инференс работает, но для бота с живыми пользователями нужен GPU.

Можно ли взять модель посвежее SDXL?
Да, diffusers поддерживает и линейку SD 3.x через свой пайплайн — меняется класс пайплайна и идентификатор весов, а архитектура сервиса и бота остается прежней. Точные имена весов и требования к видеопамяти сверяйте на странице модели.

Как ограничить нагрузку, чтобы бота не завалили запросами?
Введите ограничение частоты на пользователя (например, один активный запрос на chat_id) и очередь с максимальной длиной: при переполнении сразу отвечайте отказом, а не копите задачи, которые все равно не успеете посчитать.