DreamBooth / LoRA 微调与出图服务化

# -*- coding: utf-8 -*-
"""最小 LoRA 推理：给一个通用模型挂上训好的挂件，出图。

环境：
    pip install diffusers transformers accelerate peft torch safetensors
要点：
    load_lora_weights 只是在运行时把 BA 叠加到对应层上，
    磁盘上的 base model 权重一个字节都不会被改写。
"""
import torch
from diffusers import StableDiffusionPipeline

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

# ① 先装底模（base model / checkpoint），也就是那位通才画师
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=DTYPE).to(DEVICE)

# ② 挂上 LoRA 权重。几十 MB 的 .safetensors，秒级加载。
pipe.load_lora_weights(
    "./lora_out",                     # 目录，或直接给 .safetensors 文件路径
    weight_name="glasses_lora.safetensors",
    adapter_name="glasses",
)

# ③ 挂件强度：1.0 是训练时的原始强度，画风过冲就往下调到 0.6~0.8
pipe.set_adapters(["glasses"], adapter_weights=[0.8])

generator = torch.Generator(device="cpu").manual_seed(1024)

# ④ prompt 里必须喊出训练时用的专属 token，否则挂件挂了也等于没挂
image = pipe(
    prompt="a portrait of sks_glasses man, front view, soft studio light",
    negative_prompt="lowres, blurry, distorted face",
    num_inference_steps=25,
    guidance_scale=7.5,
    generator=generator,
).images[0]

image.save("lora_out.png")

# ⑤ 卸下挂件，pipe 立刻恢复成原来的通才画师
pipe.unload_lora_weights()
print("saved lora_out.png")

# -*- coding: utf-8 -*-
"""DreamBooth 训练启动器：把参数集中在一处，拼出 accelerate 命令并执行。

依赖 diffusers 官方示例脚本 train_dreambooth.py：
    git clone https://github.com/huggingface/diffusers
    pip install -r diffusers/examples/dreambooth/requirements.txt
    accelerate config default

显存：fp16 + 8bit Adam 下约 16~24 GB；24 GB 以上才建议同时训 text_encoder。
"""
import os
import subprocess
import sys

# ======================================================================
# 一、路径
# ======================================================================
# 底模。可以是 HuggingFace 仓库名，也可以是本地 diffusers 格式目录。
# 如果手里只有一个 .ckpt/.safetensors 单文件，先用
# convert_original_stable_diffusion_to_diffusers.py 转成 diffusers 目录。
PRETRAINED = os.environ.get("SD_BASE", "runwayml/stable-diffusion-v1-5")

# 你的主体图片目录。103 张格子衫照片就放这里，统一裁成 512x512。
INSTANCE_DIR = "./train_images"

# 类别先验图片目录。留空目录即可，脚本会自己用底模生成填满。
CLASS_DIR = "./class_images"

# 训练产物目录
OUTPUT_DIR = "./dreambooth_out"

# ======================================================================
# 二、两个 prompt：DreamBooth 的灵魂
# ======================================================================
# 专属 token + 类别词。Plaidshirtprogrammer 是刻意挑的「不存在的词」，
# 它在 tokenizer 里会被切成几个没什么语义的碎片，不会和已有概念打架。
# 反例：写 "a photo of a dog" 会把模型里所有的狗都拽向你这几张图。
INSTANCE_PROMPT = "a photo of Plaidshirtprogrammer shirt"

# 类别先验 prompt：不含专属 token，只有类别词。
# 它的作用是让模型一边学新主体、一边被「普通衬衫长什么样」拽住，
# 避免语言漂移（learned prior preservation）。
CLASS_PROMPT = "a photo of shirt"

# 训练中途的预览图 prompt，用来肉眼判断学到第几步最像
VALIDATION_PROMPT = "a photo of Plaidshirtprogrammer shirt, on a wooden hanger"

# ======================================================================
# 三、超参数
# ======================================================================
RESOLUTION = 512            # 与数据集裁切尺寸一致；改大显存暴涨
TRAIN_BATCH_SIZE = 1        # 单卡基本只能 1，用梯度累积换等效 batch
GRAD_ACCUM = 1
LEARNING_RATE = 2e-6        # DreamBooth 改的是全量 U-Net，lr 必须极小
LR_SCHEDULER = "constant"
LR_WARMUP_STEPS = 0
MAX_TRAIN_STEPS = 1200      # 经验值：实例图片数 x 10~15
NUM_CLASS_IMAGES = 200      # 类别先验图数量，一般 100~200
PRIOR_LOSS_WEIGHT = 1.0     # 先验损失权重，调大越保守、越不容易过拟合
CHECKPOINT_STEPS = 400      # 每隔多少步存一次中间权重，便于回退挑点
SEED = 1337


def build_cmd():
    """拼 accelerate launch 命令。改参数只需要动上面的常量。"""
    cmd = [
        "accelerate", "launch",
        "diffusers/examples/dreambooth/train_dreambooth.py",
        "--pretrained_model_name_or_path=%s" % PRETRAINED,
        "--instance_data_dir=%s" % INSTANCE_DIR,
        "--output_dir=%s" % OUTPUT_DIR,
        '--instance_prompt=%s' % INSTANCE_PROMPT,
        # 开启类别先验保留，下面三个参数必须成组出现
        "--with_prior_preservation",
        "--prior_loss_weight=%s" % PRIOR_LOSS_WEIGHT,
        "--class_data_dir=%s" % CLASS_DIR,
        '--class_prompt=%s' % CLASS_PROMPT,
        "--num_class_images=%d" % NUM_CLASS_IMAGES,
        # 训练规模
        "--resolution=%d" % RESOLUTION,
        "--train_batch_size=%d" % TRAIN_BATCH_SIZE,
        "--gradient_accumulation_steps=%d" % GRAD_ACCUM,
        "--learning_rate=%s" % LEARNING_RATE,
        "--lr_scheduler=%s" % LR_SCHEDULER,
        "--lr_warmup_steps=%d" % LR_WARMUP_STEPS,
        "--max_train_steps=%d" % MAX_TRAIN_STEPS,
        "--checkpointing_steps=%d" % CHECKPOINT_STEPS,
        "--seed=%d" % SEED,
        # 省显存三件套
        "--mixed_precision=fp16",
        "--gradient_checkpointing",
        "--use_8bit_adam",
        # 中途预览
        '--validation_prompt=%s' % VALIDATION_PROMPT,
        "--num_validation_images=2",
        "--validation_steps=200",
    ]
    # 显存 24 GB 以上再打开：同时微调文本编码器，专属 token 学得更牢，
    # 但也更容易把整个词表带偏。
    if os.environ.get("TRAIN_TEXT_ENCODER") == "1":
        cmd.append("--train_text_encoder")
    return cmd


def main():
    os.makedirs(INSTANCE_DIR, exist_ok=True)
    os.makedirs(CLASS_DIR, exist_ok=True)
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    imgs = [f for f in os.listdir(INSTANCE_DIR)
            if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))]
    if not imgs:
        print("实例图片目录是空的：%s" % INSTANCE_DIR)
        return 1
    print("实例图片 %d 张，建议步数 %d~%d，当前设置 %d"
          % (len(imgs), len(imgs) * 10, len(imgs) * 15, MAX_TRAIN_STEPS))

    cmd = build_cmd()
    print("\n".join(cmd))
    return subprocess.call(cmd)


if __name__ == "__main__":
    sys.exit(main())

# -*- coding: utf-8 -*-
"""LoRA 训练启动器：读一份 JSON 配置，拼出训练命令。

依赖 diffusers 官方 LoRA 示例：
    diffusers/examples/dreambooth/train_dreambooth_lora.py
显存：fp16 + 8bit Adam 下 8 GB 即可跑 512 分辨率。
为什么只要 8 GB：底模全程冻结、不需要为它保存优化器状态和梯度，
显存里只多出 rank 那一小撮可训练参数。
"""
import json
import os
import subprocess
import sys

CONFIG_PATH = os.environ.get("LORA_CONFIG", "lora_config.json")

DEFAULTS = {
    "pretrained_model": "runwayml/stable-diffusion-v1-5",
    "instance_data_dir": "./train",
    "output_dir": "./lora_out",
    "instance_prompt": "a photo of sks_glasses man",
    "resolution": 512,
    "train_batch_size": 2,
    "gradient_accumulation_steps": 1,
    "max_train_steps": 1600,
    "checkpointing_steps": 200,
    "rank": 32,
    "learning_rate": 1e-4,
    "text_encoder_lr": 5e-5,
    "lr_scheduler": "cosine_with_restarts",
    "lr_warmup_steps": 100,
    "lr_num_cycles": 3,
    "seed": 1337,
    "train_text_encoder": True,
    "mixed_precision": "fp16",
    "use_8bit_adam": True,
}


def load_config(path):
    """文件不存在就用默认值，并顺手写一份出来供修改。"""
    if not os.path.exists(path):
        with open(path, "w", encoding="utf-8") as f:
            json.dump(DEFAULTS, f, ensure_ascii=False, indent=2)
        print("已生成默认配置：%s" % path)
        return dict(DEFAULTS)
    with open(path, encoding="utf-8") as f:
        cfg = dict(DEFAULTS)
        cfg.update(json.load(f))
        return cfg


def check(cfg):
    """训练前先把最容易翻车的三件事拦下来。"""
    problems = []
    d = cfg["instance_data_dir"]
    if not os.path.isdir(d):
        problems.append("训练图目录不存在：%s" % d)
    else:
        imgs = [f for f in os.listdir(d)
                if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))]
        txts = [f for f in os.listdir(d) if f.lower().endswith(".txt")]
        if not imgs:
            problems.append("训练图目录里没有图片")
        # kohya 风格的数据集要求每张图配一个同名 txt 描述文件
        missing = [f for f in imgs
                   if os.path.splitext(f)[0] + ".txt" not in txts]
        if txts and missing:
            problems.append("缺少同名 txt 描述：%s" % ", ".join(missing[:5]))
    if cfg["rank"] > 128:
        problems.append("rank=%d 偏大，几十张图的小数据集基本必过拟合"
                        % cfg["rank"])
    if cfg["resolution"] % 8 != 0:
        problems.append("resolution 必须是 8 的倍数")
    return problems


def build_cmd(cfg):
    cmd = [
        "accelerate", "launch",
        "diffusers/examples/dreambooth/train_dreambooth_lora.py",
        "--pretrained_model_name_or_path=%s" % cfg["pretrained_model"],
        "--instance_data_dir=%s" % cfg["instance_data_dir"],
        "--output_dir=%s" % cfg["output_dir"],
        '--instance_prompt=%s' % cfg["instance_prompt"],
        "--resolution=%d" % cfg["resolution"],
        "--train_batch_size=%d" % cfg["train_batch_size"],
        "--gradient_accumulation_steps=%d" % cfg["gradient_accumulation_steps"],
        "--max_train_steps=%d" % cfg["max_train_steps"],
        "--checkpointing_steps=%d" % cfg["checkpointing_steps"],
        # 低秩分解的秩：w = w0 + BA 里 B 是 d x r、A 是 r x d
        "--rank=%d" % cfg["rank"],
        "--learning_rate=%s" % cfg["learning_rate"],
        "--text_encoder_lr=%s" % cfg["text_encoder_lr"],
        # 余弦退火重启：学习率周期性跳回高位，帮模型跳出局部最优
        "--lr_scheduler=%s" % cfg["lr_scheduler"],
        "--lr_warmup_steps=%d" % cfg["lr_warmup_steps"],
        "--lr_num_cycles=%d" % cfg["lr_num_cycles"],
        "--seed=%d" % cfg["seed"],
        "--mixed_precision=%s" % cfg["mixed_precision"],
        "--gradient_checkpointing",
    ]
    if cfg["use_8bit_adam"]:
        # AdamW8bit：把优化器状态从 fp32 压到 8bit，省下约 2/3 的优化器显存
        cmd.append("--use_8bit_adam")
    if cfg["train_text_encoder"]:
        cmd.append("--train_text_encoder")
    return cmd


def main():
    cfg = load_config(CONFIG_PATH)
    problems = check(cfg)
    if problems:
        for p in problems:
            print("[检查未通过] %s" % p)
        return 1
    cmd = build_cmd(cfg)
    print("\n".join(cmd))
    return subprocess.call(cmd)


if __name__ == "__main__":
    sys.exit(main())

# -*- coding: utf-8 -*-
"""把训练产物转成 WebUI 能加载的格式，并放进正确的目录。

两条路完全不同，别混：
    DreamBooth 产物  是一个 diffusers 目录（unet/ vae/ text_encoder/ ...）
                     → 合并成单文件 .safetensors → models/Stable-diffusion/
    LoRA 产物        本来就是单个 .safetensors
                     → 直接复制 → models/Lora/

转换脚本来自 diffusers 仓库：
    diffusers/scripts/convert_diffusers_to_original_stable_diffusion.py
"""
import argparse
import os
import shutil
import subprocess
import sys

CONVERT_SCRIPT = ("diffusers/scripts/"
                  "convert_diffusers_to_original_stable_diffusion.py")


def convert_dreambooth(model_dir, webui_root, name, half=True):
    """diffusers 目录 → 单文件 .safetensors → models/Stable-diffusion/"""
    if not os.path.isdir(os.path.join(model_dir, "unet")):
        raise SystemExit("不像 diffusers 目录，缺 unet/ 子目录：%s" % model_dir)

    dst_dir = os.path.join(webui_root, "models", "Stable-diffusion")
    os.makedirs(dst_dir, exist_ok=True)
    dst = os.path.join(dst_dir, name + ".safetensors")

    cmd = [sys.executable, CONVERT_SCRIPT,
           "--model_path", model_dir,
           "--checkpoint_path", dst,
           "--use_safetensors"]
    if half:
        # fp16 存盘，2 GB 出头；不加这个参数是 4 GB 起步
        cmd.append("--half")
    print(" ".join(cmd))
    rc = subprocess.call(cmd)
    if rc != 0:
        raise SystemExit("转换失败，退出码 %d" % rc)
    print("已写入 %s（%.1f MB）" % (dst, os.path.getsize(dst) / 1e6))
    return dst


def install_lora(lora_file, webui_root, name=None):
    """LoRA 不需要转换，复制到 models/Lora/ 即可。"""
    if not os.path.isfile(lora_file):
        raise SystemExit("找不到 LoRA 权重：%s" % lora_file)
    dst_dir = os.path.join(webui_root, "models", "Lora")
    os.makedirs(dst_dir, exist_ok=True)
    fname = (name + ".safetensors") if name else os.path.basename(lora_file)
    dst = os.path.join(dst_dir, fname)
    shutil.copyfile(lora_file, dst)
    print("已写入 %s（%.1f MB）" % (dst, os.path.getsize(dst) / 1e6))
    return dst


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--webui", required=True,
                    help="stable-diffusion-webui 根目录")
    ap.add_argument("--dreambooth", help="DreamBooth 输出的 diffusers 目录")
    ap.add_argument("--lora", help="LoRA 输出的 .safetensors 文件")
    ap.add_argument("--name", default="my_model", help="落盘后的文件名（不含后缀）")
    ap.add_argument("--fp32", action="store_true", help="不压 fp16")
    args = ap.parse_args()

    if not args.dreambooth and not args.lora:
        raise SystemExit("--dreambooth 与 --lora 至少给一个")

    if args.dreambooth:
        convert_dreambooth(args.dreambooth, args.webui, args.name,
                           half=not args.fp32)
    if args.lora:
        install_lora(args.lora, args.webui, args.name)

    print("\n放好之后回到 WebUI：左上角模型下拉框点刷新，"
          "LoRA 在生成页的 Lora 标签页里；\n"
          "出图时 prompt 里必须带上训练时的专属 token。")


if __name__ == "__main__":
    main()

# -*- coding: utf-8 -*-
"""异步出图服务：提交任务拿号 → 后台画 → 凭号查结果。

这是一套通用的异步长任务模式，和「出图」本身无关：
凡是单次耗时超过十几秒的活儿（转码、导出、批量分析），
都应该走 提交/查询 两个接口，而不是让 HTTP 连接干等。

只用标准库跑起来，不引第三方 Web 框架，便于直接阅读：
    python3 async_draw_service.py
    curl -X POST localhost:8000/frontend/create -d '{"prompt":"一只柯基"}'
    curl -X POST localhost:8000/frontend/query  -d '{"job_id":"..."}'

生产环境把 QUEUE/RESULT 两个字典换成 Redis 或对象存储即可横向扩机器。
"""
import json
import os
import random
import string
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

LISTEN = ("0.0.0.0", int(os.environ.get("PORT", "8000")))

# 任务队列与结果表。单机内存版，重启即丢；生产环境换 Redis/COS。
QUEUE = []                 # 待处理任务，元素是 dict
RESULT = {}                # job_id -> 结果 dict
LOCK = threading.Lock()

# 一次出图的经验耗时，用来给前端估算剩余时间
AVG_COST_SEC = 40

ALPHABET = string.ascii_letters + string.digits


def gen_job_id(n=32):
    """任务号要足够长、足够随机：它同时充当「凭号取件」的凭证。
    用自增 ID 的话，别人改个数字就能看到你的图。"""
    return "".join(random.choice(ALPHABET) for _ in range(n))


# ====================================================================
# 出图后端：这里替换成你自己的实现
# ====================================================================
def render_image(prompt, negative_prompt, seed):
    """真正干活的地方。三选一：

    1) 本地 diffusers：pipe(prompt=...).images[0].save(path)
    2) 本机 WebUI 的 /sdapi/v1/txt2img 接口
    3) 云厂商文生图 API（返回 base64 或临时 URL）

    这里用 sleep 占位，保证这份代码不依赖 GPU 也能跑通整条链路。
    """
    time.sleep(AVG_COST_SEC * 0.1)
    return {"prompt": prompt, "negative_prompt": negative_prompt, "seed": seed}


def upload_to_object_storage(job_id, payload):
    """把成品图传到对象存储，返回带签名、会过期的临时 URL。

    两个要点：
    ① 不要把图片 base64 塞进查询接口的响应里——几 MB 的 JSON 会拖垮客户端；
    ② 一律用预签名 URL 并设过期时间，桶本身保持私有读。
    """
    path = "aidraw/%s.png" % job_id
    expire = int(time.time()) + 3600
    sign = os.environ.get("COS_SECRET_ID", "")[:6] or "unsigned"
    return "https://example-bucket.cos.example.com/%s?e=%d&s=%s" % (
        path, expire, sign)


# ====================================================================
# 后台 worker：单线程逐个消费，天然就是限流
# ====================================================================
def worker_loop():
    while True:
        job = None
        with LOCK:
            if QUEUE:
                job = QUEUE.pop(0)
        if job is None:
            time.sleep(1)
            continue

        job_id = job["job_id"]
        begin = time.time()
        with LOCK:
            RESULT[job_id] = {"job_id": job_id, "job_status": "RUNNING"}
        try:
            payload = render_image(job.get("prompt", ""),
                                   job.get("negative_prompt", ""),
                                   job.get("seed"))
            url = upload_to_object_storage(job_id, payload)
            done = {
                "job_id": job_id,
                "session_id": job.get("session_id", ""),
                "job_status": "FINISHED",
                "result_url": url,
                "cost_time": int((time.time() - begin) * 1000),
            }
        except Exception as exc:                       # noqa: BLE001
            # 失败也必须写进结果表，否则客户端会永远轮询下去
            done = {
                "job_id": job_id,
                "job_status": "FAILED",
                "error": "%s: %s" % (type(exc).__name__, exc),
                "cost_time": int((time.time() - begin) * 1000),
            }
        with LOCK:
            RESULT[job_id] = done


# ====================================================================
# HTTP 接口
# ====================================================================
class Handler(BaseHTTPRequestHandler):
    def _read_json(self):
        length = int(self.headers.get("Content-Length") or 0)
        raw = self.rfile.read(length) if length else b"{}"
        try:
            return json.loads(raw.decode("utf-8") or "{}")
        except json.JSONDecodeError:
            return {}

    def _send(self, obj, code=200):
        body = json.dumps(obj, ensure_ascii=False).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_POST(self):                                  # noqa: N802
        req = self._read_json()

        # 接口一：提交任务，立刻返回任务号。这一步绝不能等图画完。
        if self.path == "/frontend/create":
            prompt = (req.get("prompt") or "").strip()
            if not prompt:
                return self._send({"error": "prompt 不能为空"}, 400)
            if len(prompt) > 256:
                # 云厂商文生图接口普遍限制 256 字符，前置校验省一次失败调用
                return self._send({"error": "prompt 超过 256 字符"}, 400)

            job = {
                "job_id": gen_job_id(),
                "session_id": req.get("session_id", ""),
                "prompt": prompt,
                "negative_prompt": req.get("negative_prompt", ""),
                "seed": req.get("seed"),
            }
            with LOCK:
                QUEUE.append(job)
                queue_len = len(QUEUE)
            return self._send({
                "job_id": job["job_id"],
                "session_id": job["session_id"],
                "job_status": "PENDING",
                "queue_len": queue_len,
                # 把估算的等待时间一起返回，前端才能画进度条
                "remain_time": queue_len * AVG_COST_SEC,
            })

        # 接口二：凭号查结果。查不到就说明还在队列里。
        if self.path == "/frontend/query":
            job_id = req.get("job_id") or ""
            with LOCK:
                ret = RESULT.get(job_id)
                pending = any(j["job_id"] == job_id for j in QUEUE)
            if ret:
                return self._send(ret)
            if pending:
                return self._send({"job_id": job_id, "job_status": "PENDING"})
            return self._send({"job_id": job_id, "job_status": "NOT_FOUND"}, 404)

        return self._send({"error": "no such route"}, 404)

    def log_message(self, fmt, *args):
        print("[%s] %s" % (time.strftime("%H:%M:%S"), fmt % args))


def main():
    threading.Thread(target=worker_loop, daemon=True).start()
    print("listening on %s:%d" % LISTEN)
    ThreadingHTTPServer(LISTEN, Handler).serve_forever()


if __name__ == "__main__":
    main()

# -*- coding: utf-8 -*-
"""异步出图服务的客户端调用模板：提交 → 轮询 → 取图。

小程序端 wx.request、Web 端 fetch、Go 端 http.Post，逻辑完全一样，
只是语法不同。这份 Python 版可以直接拿来做联调和压测。
"""
import json
import os
import time
import urllib.error
import urllib.request

# ======================================================================
# 参数区
# ======================================================================
# TODO 1：服务端地址
BASE_URL = os.environ.get("DRAW_API", "http://127.0.0.1:8000")

# TODO 2：鉴权。凭证一律走环境变量，绝不写进源码或前端包里。
TOKEN = os.environ.get("DRAW_API_TOKEN", "")

# TODO 3：出图参数
PROMPT = "一只戴墨镜的柯基，工作室灯光，高细节"
NEGATIVE_PROMPT = "低分辨率, 模糊, 水印"

# TODO 4：轮询策略。
# 首次等待要给足（出图本来就要几十秒），间隔用退避而不是死等 1 秒，
# 否则 100 个用户在线就是每秒 100 次无效查询。
FIRST_WAIT = 8          # 提交后先睡多久再开始问
POLL_INTERVAL = 3       # 初始轮询间隔（秒）
POLL_BACKOFF = 1.4      # 每次乘这个系数，上限 POLL_MAX
POLL_MAX = 15
TIMEOUT_SEC = 180       # 总超时，超了就报错退出，不能无限等
# ======================================================================


def post(path, payload):
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(
        BASE_URL + path, data=data,
        headers={"Content-Type": "application/json"})
    if TOKEN:
        req.add_header("Authorization", "Bearer " + TOKEN)
    try:
        with urllib.request.urlopen(req, timeout=20) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as err:
        body = err.read().decode("utf-8", "replace")
        return {"job_status": "HTTP_%d" % err.code, "error": body}


def submit():
    """第一步：提交任务，只拿号，不等图。"""
    ret = post("/frontend/create", {
        "prompt": PROMPT,
        "negative_prompt": NEGATIVE_PROMPT,
        "session_id": "demo-session",
    })
    if not ret.get("job_id"):
        raise SystemExit("提交失败：%s" % ret)
    print("已取号 job_id=%s 队列长度=%s 预计 %ss"
          % (ret["job_id"], ret.get("queue_len"), ret.get("remain_time")))
    return ret["job_id"]


def poll(job_id):
    """第二步：凭号轮询，直到 FINISHED / FAILED / 超时。"""
    deadline = time.time() + TIMEOUT_SEC
    interval = POLL_INTERVAL
    time.sleep(FIRST_WAIT)

    while time.time() < deadline:
        ret = post("/frontend/query", {"job_id": job_id})
        status = ret.get("job_status")
        print("  status=%s" % status)

        if status == "FINISHED":
            return ret
        if status in ("FAILED", "NOT_FOUND"):
            raise SystemExit("任务失败：%s" % ret)

        time.sleep(interval)
        interval = min(interval * POLL_BACKOFF, POLL_MAX)

    raise SystemExit("超过 %ds 仍未完成，job_id=%s" % (TIMEOUT_SEC, job_id))


def download(url, out_path):
    """第三步：结果是一个会过期的临时 URL，要落盘就趁早。"""
    with urllib.request.urlopen(url, timeout=60) as resp:
        body = resp.read()
    with open(out_path, "wb") as f:
        f.write(body)
    print("已保存 %s（%d 字节）" % (out_path, len(body)))


def main():
    job_id = submit()
    ret = poll(job_id)
    print("耗时 %s ms" % ret.get("cost_time"))
    print("图片地址 %s" % ret.get("result_url"))
    # TODO 5：真实环境打开这一行把图拉下来
    # download(ret["result_url"], "%s.png" % job_id)


if __name__ == "__main__":
    main()

{
  "_说明": "LoRA 训练参数模板。复制成 lora_config.json，改完 TODO 再跑 train_lora.py。",

  "pretrained_model": "TODO 底模路径。想让 LoRA 和某个画风底模配套，就填那个底模",
  "instance_data_dir": "TODO 训练图目录。每张图配一个同名 .txt 描述文件",
  "output_dir": "./lora_out",
  "instance_prompt": "TODO 专属 token + 类别词，例如 a photo of sks_glasses man",

  "_分辨率": "训练分辨率要和数据集裁切尺寸一致。拉到 768 显存翻倍，SD v1 底模收益有限",
  "resolution": 512,

  "_批次": "8G 显存填 1~2；等效批次靠 gradient_accumulation_steps 放大，显存不涨",
  "train_batch_size": 2,
  "gradient_accumulation_steps": 1,

  "_步数": "经验值 = 图片数 x repeat x epoch。20 张图跑 1600 步是常见起点；过拟合就往下砍",
  "max_train_steps": 1600,

  "_存盘频率": "每 200 步存一版，训完横向挑最好的那一版，而不是只信最后一版",
  "checkpointing_steps": 200,

  "_rank": "低秩分解的秩 r。8~16 学画风，32~64 学具体人物/物体，>128 小数据集必过拟合",
  "rank": 32,

  "_学习率": "LoRA 的 lr 比全量微调大两个量级。unet 用 1e-4，文本编码器要更小，否则词表被带偏",
  "learning_rate": 1e-4,
  "text_encoder_lr": 5e-5,

  "_调度器": "余弦退火重启：lr 周期性跳回高位，帮模型跳出局部最优。cycles 一般 2~4",
  "lr_scheduler": "cosine_with_restarts",
  "lr_warmup_steps": 100,
  "lr_num_cycles": 3,

  "_优化器": "AdamW8bit 把优化器状态压到 8bit，是 8G 显存能跑起来的关键一环",
  "use_8bit_adam": true,
  "mixed_precision": "fp16",

  "_文本编码器": "true 时专属 token 学得更牢，但画风容易外溢到别的词；只学画风可以关掉",
  "train_text_encoder": true,

  "_种子": "固定住才能做参数对比实验：只改一个参数、seed 不变，差异才归因得清",
  "seed": 1337
}

# -*- coding: utf-8 -*-
"""异步出图服务的客户端调用模板：提交 → 轮询 → 取图。

小程序端 wx.request、Web 端 fetch、Go 端 http.Post，逻辑完全一样，
只是语法不同。这份 Python 版可以直接拿来做联调和压测。
"""
import json
import os
import time
import urllib.error
import urllib.request

# ======================================================================
# 参数区
# ======================================================================
# TODO 1：服务端地址
BASE_URL = os.environ.get("DRAW_API", "http://127.0.0.1:8000")

# TODO 2：鉴权。凭证一律走环境变量，绝不写进源码或前端包里。
TOKEN = os.environ.get("DRAW_API_TOKEN", "")

# TODO 3：出图参数
PROMPT = "一只戴墨镜的柯基，工作室灯光，高细节"
NEGATIVE_PROMPT = "低分辨率, 模糊, 水印"

# TODO 4：轮询策略。
# 首次等待要给足（出图本来就要几十秒），间隔用退避而不是死等 1 秒，
# 否则 100 个用户在线就是每秒 100 次无效查询。
FIRST_WAIT = 8          # 提交后先睡多久再开始问
POLL_INTERVAL = 3       # 初始轮询间隔（秒）
POLL_BACKOFF = 1.4      # 每次乘这个系数，上限 POLL_MAX
POLL_MAX = 15
TIMEOUT_SEC = 180       # 总超时，超了就报错退出，不能无限等
# ======================================================================


def post(path, payload):
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(
        BASE_URL + path, data=data,
        headers={"Content-Type": "application/json"})
    if TOKEN:
        req.add_header("Authorization", "Bearer " + TOKEN)
    try:
        with urllib.request.urlopen(req, timeout=20) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as err:
        body = err.read().decode("utf-8", "replace")
        return {"job_status": "HTTP_%d" % err.code, "error": body}


def submit():
    """第一步：提交任务，只拿号，不等图。"""
    ret = post("/frontend/create", {
        "prompt": PROMPT,
        "negative_prompt": NEGATIVE_PROMPT,
        "session_id": "demo-session",
    })
    if not ret.get("job_id"):
        raise SystemExit("提交失败：%s" % ret)
    print("已取号 job_id=%s 队列长度=%s 预计 %ss"
          % (ret["job_id"], ret.get("queue_len"), ret.get("remain_time")))
    return ret["job_id"]


def poll(job_id):
    """第二步：凭号轮询，直到 FINISHED / FAILED / 超时。"""
    deadline = time.time() + TIMEOUT_SEC
    interval = POLL_INTERVAL
    time.sleep(FIRST_WAIT)

    while time.time() < deadline:
        ret = post("/frontend/query", {"job_id": job_id})
        status = ret.get("job_status")
        print("  status=%s" % status)

        if status == "FINISHED":
            return ret
        if status in ("FAILED", "NOT_FOUND"):
            raise SystemExit("任务失败：%s" % ret)

        time.sleep(interval)
        interval = min(interval * POLL_BACKOFF, POLL_MAX)

    raise SystemExit("超过 %ds 仍未完成，job_id=%s" % (TIMEOUT_SEC, job_id))


def download(url, out_path):
    """第三步：结果是一个会过期的临时 URL，要落盘就趁早。"""
    with urllib.request.urlopen(url, timeout=60) as resp:
        body = resp.read()
    with open(out_path, "wb") as f:
        f.write(body)
    print("已保存 %s（%d 字节）" % (out_path, len(body)))


def main():
    job_id = submit()
    ret = poll(job_id)
    print("耗时 %s ms" % ret.get("cost_time"))
    print("图片地址 %s" % ret.get("result_url"))
    # TODO 5：真实环境打开这一行把图拉下来
    # download(ret["result_url"], "%s.png" % job_id)


if __name__ == "__main__":
    main()