AI 绘画提示词网站：从技术选型到生产部署

2. 异步推理服务实现

Flask 接口层使用 Celery 实现任务队列，避免阻塞主线程。模型加载部分做了缓存处理，防止重复加载浪费资源：

from flask import Flask, request
from celery import Celery
import torch
from diffusers import StableDiffusionPipeline

app = Flask(__name__)
celery = Celery('tasks', broker='redis://localhost:6379/0')

# 带缓存的模型加载
model_cache = {}

def load_model(model_name: str, device: str):
    if model_name not in model_cache:
        try:
            pipe = StableDiffusionPipeline.from_pretrained(
                model_name, torch_dtype=torch.float16
            ).to(device)
            model_cache[model_name] = pipe
        except Exception as e:
            raise RuntimeError(f"Model loading failed: {str(e)}")
    return model_cache[model_name]

@celery.task
def generate_image_task(prompt: str):
    try:
        pipe = load_model("runwayml/stable-diffusion-v1-5", "cuda")
        return pipe(prompt).images[0]
    except torch.cuda.OutOfMemoryError:
        # 显存不足时自动降级到 CPU
        pipe = load_model("runwayml/stable-diffusion-v1-5", "cpu")
        return pipe(prompt).images[0]

@app.route('/generate', methods=['POST'])
def generate():
    prompt = request.json.get('prompt')
    task = generate_image_task.delay(prompt)
    return {'task_id': task.id}

性能优化：从单机到分布式

GPU 资源调度策略

通过测试发现 batch_size=4 时达到最佳性价比，再往上显存压力过大：

Nginx 负载均衡配置

为了支撑高并发，我们在前端加了 Nginx 做反向代理，配置了健康检查和长连接：

upstream sd_backend {
    least_conn;
    server 192.168.1.10:5000 max_fails=3 fail_timeout=30s;
    server 192.168.1.11:5000 max_fails=3 fail_timeout=30s;
    keepalive 32;
}

server {
    location /generate {
        proxy_pass http://sd_backend;
        proxy_read_timeout 300s;
        proxy_buffering off;
    }
}

避坑指南：血泪经验总结

安全防护：必须过滤危险提示词，防止注入攻击或生成违规内容：

import re

def sanitize_prompt(prompt: str) -> str:
    # 过滤注入攻击和 NSFW 内容
    pattern = r"(?:\.\./|\\|\0|eval\(|system\(|裸体 | 暴力)"
    return re.sub(pattern, "[REDACTED]", prompt)

中文分词陷阱：CLIP 模型的 tokenizer 对中文按字切分，直接分词会拆解中文字符，需要特殊处理：

from transformers import CLIPTokenizer

tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")

# 错误做法：直接分词会拆解中文字符
tokens = tokenizer.tokenize("水墨画")  # ['水', '墨', '画</w>']

# 正确做法：添加特殊标记
tokens = tokenizer.tokenize("水墨画", add_special_tokens=True)

延伸思考：垂直领域优化方案

通过 LoRA 微调可以显著提升特定领域的提示词生成质量。比如收集 500-1000 组领域相关提示词（如'国画风格'），通用模型生成的'龙'偏向西方 dragon 形象，而 LoRA 微调后能自动生成符合中国龙特征的图像。

训练配置大致如下：

train:
  base_model: "runwayml/stable-diffusion-v1-5"
  rank: 64
  epochs: 10
  learning_rate: 1e-4

最后留给大家一个开放性问题：当用户量从 1000 激增到 10 万时，如何在实时性（<500ms）和模型精度（高画质）之间找到平衡点？或许分布式推理方案能给你启发。我在实际部署中发现，异步任务调度机制对高并发场景特别友好，值得借鉴。