wkea/blockdiffusion-api

BlockDiffusion

目标：训练一个专用于生成 16×16《我的世界》纹理 的小型 Diffusion（UNet）模型（PyTorch），支持 提示词条件生成 与 RGBA 透明通道。

Hugging Face Spaces（在线API + 伪私有）

• —Space入口文件：app.py
• —伪私有方式：在 Space 的 Settings -> Secrets 配置 SPACE_API_KEY（若网页端不允许下划线则用 SPACEAPIKEY），所有推理请求必须携带该 Key 才会执行。
• —checkpoint来源（二选一）：
• —配置 BD_CKPT_PATH（Space仓库内的文件路径）
• —或配置 BD_CKPT_REPO + BD_CKPT_FILENAME（从 Hugging Face Hub 下载；私有仓库可再配 HF_TOKEN）

当前能力（已实现）

• —16×16 专用 UNet 扩散：避免过深/过度下采样（16→8→4）
• —噪声调度：linear / cosine
• —loss：基础为 MSE，支持可选 SNR 加权（提升稳定性/质量）
• —训练/采样解耦：Diffusion 公式与采样逻辑独立于训练器
• —文本条件（提示词）：
• —支持 captions.jsonl/csv（图文配对）
• —支持“从文件名自动提取提示词”的弱标注训练（见 train_oneclick.py）
• —支持 CFG（训练丢条件 + 采样 scale）
• —RGB/RGBA：channels=3/4，RGBA 输出 PNG 透明通道

主要入口

• —一键训练（推荐）：train_oneclick.py
• —无需命令行参数，只修改脚本顶部常量
• —默认从 DATA_DIR 扫描图片，并按文件名生成提示词训练
• —默认支持 自动续训：扫描 outputs/<RUN_NAME>/checkpoints/step_*.pt 选择最大 step 继续
• —可参数化训练：train.py（保留）
• —采样：sample.py（支持 --prompt + CFG）
• —WebUI 采样预览：webui.py（浏览器里快速对比 checkpoint）

数据与提示词规则

方式 A：文件名弱标注（train_oneclick.py）

适配命名：HASH-tag_tag_tag.png、HASH-side.png、HASH-top.png 等。

提取规则（实现于 blockdiffusion/data/filename_captions.py）：

• —去扩展名
• —取第一个 - 后的部分
• —_/- 视为分隔符（转空格），转小写

方式 B：显式 captions（train.py --captions）

• —captions.jsonl：每行 {"file":"相对data_dir的路径或绝对路径","text":"提示词"}
• —captions.csv：两列 file,text（含表头）

输出与预览

• —checkpoints：outputs/<RUN_NAME>/checkpoints/step_XXXX.pt
• —训练中采样图：outputs/<RUN_NAME>/samples/step_XXXX.png
• —采样脚本输出：默认 samples/samples.png

核心模块导航（给接手者）

• —模型文档：见 blockdiffusion/models/README.md（只讲模型：结构/输入输出/条件/ckpt 格式）
• —blockdiffusion/models/unet16.py：16×16 UNet（支持条件向量注入）
• —blockdiffusion/models/text_unet16.py：TextCondUNet16（tokens/mask→cond→UNet）
• —blockdiffusion/text/simple_tokenizer.py：轻量 tokenizer/vocab（可随 checkpoint 保存）
• —blockdiffusion/text/text_encoder.py：轻量 TextEncoder（masked mean pooling + MLP）
• —blockdiffusion/diffusion/schedules.py：beta schedule（linear/cosine）
• —blockdiffusion/diffusion/gaussian_diffusion.py：训练损失 + DDPM/DDIM 采样 + CFG 训练丢条件
• —blockdiffusion/train/trainer.py：AMP/EMA/checkpoint/采样、梯度累积、warmup+余弦、channels_last
• —blockdiffusion/utils/ema.py：EMA（续训时自动对齐 device）
• —blockdiffusion/utils/image_io.py：保存 RGB/RGBA 网格 PNG

常用提速/提质开关（修改 train_oneclick.py 顶部常量）

• —提速：NUM_WORKERS、PREFETCH_FACTOR、PERSISTENT_WORKERS、CHANNELS_LAST
• —显存/吞吐：BATCH_SIZE、GRAD_ACCUM_STEPS
• —质量/稳定：SNR_GAMMA、LR_WARMUP_STEPS、LR_MIN_RATIO、COND_DROP_PROB、SAMPLE_CFG_SCALE