One-DM：用于手写文本生成的单样本扩散模仿器

🌟 简介

• 我们提出了一种用于风格化手写文本生成的单样本扩散模仿器（One-DM），该方法仅需一个参考样本作为风格输入，即可模仿其书写风格生成任意内容的手写文本。
• 以往的最先进方法由于学习风格的能力有限，难以从单个样本中准确提取用户的书写风格。为解决这一问题，我们引入了参考样本的高频成分来增强书写风格的提取。所提出的风格增强模块能够有效捕捉书写风格模式，并抑制背景噪声的干扰。
• 在英语、中文和日语手写数据集上的大量实验表明，我们的方法仅使用一个风格参考样本，其性能甚至超过了使用15倍以上参考样本的先前方法。

所提出的One-DM概述

🌠 最新消息

• [2025/06/26] 🔥🔥🔥 DiffBrush，一种用于整行文本生成的新型最先进方法，已被ICCV 2025接收。
• [2024/10/24] 我们已在Google Drive和百度网盘上提供了训练好的One-DM检查点 :)
• [2024/09/16] 本工作被Synced（机器之心）报道。
• [2024/09/07]🔥🔥🔥 我们开源了首个能够生成手写字词的One-DM版本。（后续支持中文和日语的版本将很快发布。)

🔨 环境要求

conda create -n One-DM python=3.8 -y
conda activate One-DM
# 安装所有依赖
conda env create -f environment.yml

☀️ 数据集

我们提供了英语数据集，可通过Google Drive | 百度网盘 | ShiZhi AI获取。请下载这些数据集，解压后将文件移至/data。

🐳 模型库

模型	Google Drive	百度网盘	ShiZhi AI
预训练One-DM	Google Drive	百度网盘	ShiZhi AI
预训练OCR模型	Google Drive	百度网盘	ShiZhi AI
预训练Resnet18	Google Drive	百度网盘	ShiZhi AI

注意： 请下载这些权重，并将其移动到/model_zoo。（如果您无法访问Hugging Face上提供的预训练VAE模型，请参阅置顶问题以获取指导。）

🏋️ 训练与测试

• 在英语数据集上训练

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 train.py \
    --feat_model model_zoo/RN18_class_10400.pth \
    --log English

• 在英语数据集上微调

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 train_finetune.py \
    --one_dm ./Saved/IAM64_scratch/English-timestamp/model/epoch-ckpt.pt \
    --ocr_model ./model_zoo/vae_HTR138.pth --log English

注意： 请根据您自己的路径修改timestamp和epoch。

• 在英语数据集上测试

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 test.py \
   --one_dm ./Saved/IAM64_finetune/English-timestamp/model/epoch-ckpt.pt \
   --generate_type oov_u --dir ./Generated/English

注意： 请根据您自己的路径修改timestamp和epoch。

📺 展示

• 与工业级图像生成方法在手写文本生成上的对比

• 与工业级图像生成方法在中文手写生成上的对比

• 英语手写文本生成

• 中文和日语手写生成

❤️ 引用

如果您觉得我们的工作富有启发性，或在您的研究中使用了我们的代码库，请引用我们的工作：

@inproceedings{one-dm2024,
  title={用于手写文本生成的单样本扩散模仿器},
  author={戴刚、张一凡、柯曲辉、郭强亚、黄双平},
  booktitle={欧洲计算机视觉会议},
  year={2024}
}

⭐ 星图

One-DM 快速上手指南

One-DM (One-Shot Diffusion Mimicker) 是一个用于手写文本生成的开源项目。它仅需单张参考样本即可模仿其书写风格，生成任意内容的个性化手写文本。该项目在英文、中文和日文数据集上均表现出色。

环境准备

• 操作系统: Linux (推荐 Ubuntu)
• Python 版本: 3.8
• 深度学习框架: PyTorch 1.13+
• 硬件要求: 支持 CUDA 的 NVIDIA GPU (训练建议多卡，推理单卡即可)

安装步骤

1. 创建并激活虚拟环境

使用 Conda 创建名为 One-DM 的独立环境：

conda create -n One-DM python=3.8 -y
conda activate One-DM

2. 安装依赖

在项目根目录下执行以下命令安装所有必要依赖：

conda env create -f environment.yml

3. 下载数据集与预训练模型

为了获得最佳体验，建议优先使用国内镜像源下载资源。

数据集 (英文) 下载后解压，并将文件夹移动至项目根目录下的 /data 文件夹中。

• 百度网盘: 点击下载 (提取码: 4vsv)
• 始智 AI (ShiZhi): 点击下载

预训练模型 (Model Zoo) 下载以下三个文件，并统一移动至项目根目录下的 /model_zoo 文件夹中：

1. One-DM 主模型: One-DM-ckpt.pt
   • 始智 AI 下载
2. OCR 模型: vae_HTR138.pth
   • 始智 AI 下载
3. ResNet18 特征模型: RN18_class_10400.pth
   • 始智 AI 下载

注意: 请确保目录结构如下：

One-DM/
├── data/          # 存放解压后的数据集
├── model_zoo/     # 存放下载的 .pth 和 .pt 权重文件
├── train.py
├── test.py
└── ...

基本使用

以下命令演示如何加载预训练模型进行推理测试（生成手写文本）。

假设你已经完成了上述“下载预训练模型”步骤，并且模型位于 ./model_zoo 目录下。由于官方提供的测试脚本默认指向训练保存的路径（./Saved/...），你需要先确认是否有直接对应的测试权重，或者将下载的 One-DM-ckpt.pt 重命名并放置到脚本预期的路径，亦或直接修改命令中的 --one_dm 参数路径。

最简单的推理命令示例：

请将 ./model_zoo/One-DM-ckpt.pt 替换为你实际下载的模型文件路径。如果官方测试脚本 test.py 强依赖特定的目录结构（如包含 epoch 信息的文件夹），建议先将下载的权重文件移动到对应结构或修改代码中的加载逻辑。

若直接使用官方提供的测试命令模板（需调整路径以匹配你下载的模型）：

CUDA_VISIBLE_DEVICES=0 torchrun --nproc_per_node=1 test.py \
    --one_dm ./model_zoo/One-DM-ckpt.pt \
    --generate_type oov_u \
    --dir ./Generated/English

注：如果运行报错提示找不到文件路径格式不匹配（因为官方示例中路径包含时间戳和 epoch），请检查 test.py 源码，看是否可以直接传入 .pt 文件路径，或者按照官方说明将模型放入 ./Saved/IAM64_finetune/English-[timestamp]/model/ 目录下并命名为 epoch-ckpt.pt。

生成结果查看： 运行成功后，生成的手写文本图片将保存在 ./Generated/English 目录中。

---

更多高级用法（如微调训练、多卡训练）请参考项目原始 README 中的 Training & Test 章节。