Logo
Project Icon

ml-mdm

开源框架实现高分辨率文本到图像生成模型

Matryoshka Diffusion Models文本到图像生成高分辨率图像合成神经网络模型深度学习Github开源项目

ml_mdm是一个Python开源项目,实现了Matryoshka扩散模型技术用于文本到图像生成。该框架支持训练单个像素空间模型生成高达1024x1024分辨率的图像,开源了U-Net和嵌套U-Net的实现。项目提供预训练模型、Web演示和CC12M数据集上的训练教程,为高分辨率图像和视频合成提供完整解决方案。

GithubGithub论文论文
介绍相关项目

ml_mdm - 套娃扩散模型

ml_mdm是一个Python包,用于高效训练高质量的文本到图像扩散模型 — 由Luke Carlson、Jiatao Gu、Shuangfei Zhai和Navdeep Jaitly推出。

本软件项目伴随研究论文Matryoshka Diffusion Models。

Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Josh Susskind, Navdeep Jaitly

[论文] [引用]

mdm文本到图像输出

目录

章节描述
简介套娃扩散模型的简要概述
安装开始使用ml_mdm训练模型和生成样本
预训练模型下载我们预训练模型(64, 356, 1024)的链接
网页演示使用我们的网页界面生成图像
代码库结构Python模块概述
概念核心概念和设计原则
教程在CC12m上逐步训练MDM模型

简介

扩散模型是生成高质量图像和视频的事实标准方法,但由于计算和优化挑战,学习高维模型仍然是一项艰巨的任务。

ml_mdm是一个端到端的高分辨率图像和视频合成框架 — 它以我们的技术命名:套娃扩散模型。

值得注意的是,我们可以在高达1024x1024像素的分辨率下训练单个像素空间模型,使用仅包含1200万张图像的CC12M数据集展示了强大的零样本泛化能力。

mdm多尺度流程

安装

默认安装依赖项在pyproject.toml中定义,选择这些依赖项是为了让您即使在只有CPU的机器上也能安装此库。

用户已经在Python 3.9、3.10和cuda_12、cuda-11.8上运行过这个代码库

> pip install -e .

开发人员还应该使用pre-commit install设置pre-commit。

运行测试用例

> pytest   # 将运行所有测试用例 - 包括需要GPU的测试
> pytest  -m "not gpu"  # 运行只需要CPU的测试用例

预训练模型

我们已经上传了模型检查点到:

注意:我们发布的模型是在从Flickr收集的5000万对文本-图像对上训练的。在这个仓库中,我们提供了下载CC12M的脚本,以及在CC12M数据上训练等效模型的配置。

您可以随意下载模型或继续往下阅读以训练自己的模型。一旦预训练模型在本地下载完成,您可以在我们的网页演示中使用它,将其作为参数传递给训练、采样等。

export ASSET_PATH=https://docs-assets.developer.apple.com/ml-research/models/mdm

curl $ASSET_PATH/flickr64/vis_model.pth --output vis_model_64x64.pth
curl $ASSET_PATH/flickr256/vis_model.pth --output vis_model_256x256.pth
curl $ASSET_PATH/flickr1024/vis_model.pth --output vis_model_1024x1024.pth

网页演示

下载检查点后,您可以使用以下命令运行自己的网页演示实例:

torchrun --standalone --nproc_per_node=1  ml_mdm/clis/generate_sample.py --port $YOUR_PORT

image

代码库

模块描述
ml_mdm.models核心模型实现
ml_mdm.diffusion模型管道,例如DDPM
ml_mdm.config使用simple parsing将配置数据类与相关模型、管道和命令行接口连接
ml_mdm.clis项目中的所有命令行工具,最相关的是train_parallel.py
tests/单元测试和样本训练文件

概念

ml_mdm.models

在ml_mdm.models子模块中,我们开源了以下实现:

  • U-Nets
  • 嵌套U-Nets

ml_mdm.config

ml_mdm.config包含核心配置和命令行接口逻辑。这个代码库中的许多模型、命令行接口和函数都是通过传入一个dataclass对象来配置的。我们使用SimpleParsing动态创建命令行接口,并允许使用--config_path传入yaml config表示。

本质上,simple_parsing会将所有传入的命令行参数和yaml文件转换为干净的配置类,如ml_mdm.reader.ReaderConfig、ml_mdm.diffusion.DiffusionConfig。

教程

使用预训练检查点生成自己的图像

安装ml_mdm后,将这些检查点下载到仓库目录中。

curl https://docs-assets.developer.apple.com/ml-research/models/mdm/flickr64/vis_model.pth --output vis_model_64x64.pth
curl https://docs-assets.developer.apple.com/ml-research/models/mdm/flickr256/vis_model.pth --output vis_model_256x256.pth

网页演示将使用相应的配置加载每个模型:

  • vis_model_64x64.pth将使用configs/models/cc12m_64x64.yaml中的设置加载
  • vis_model_256x256.pth将使用configs/models/cc12m_256x256.yaml中的设置加载
  • vis_model_1024x1024.pth将使用configs/models/cc12m_1024x1024.yaml中的设置加载

在演示中,您可以更改各种设置并查看模型的内部结构。通过替换$YOUR_PORT设置您想使用的端口,然后运行:

torchrun --standalone --nproc_per_node=1  ml_mdm/clis/generate_sample.py --port $YOUR_PORT

在虚拟数据上训练

如果您只想在不下载大型数据集的情况下逐步完成训练模型和运行管道的过程,我们为您准备了一个最小示例。它使用tests/test_files/中的虚拟数据

随意尝试直接在命令行或通过编辑配置yaml文件来更改各种--args

torchrun --standalone --nproc_per_node=1 ml_mdm/clis/train_parallel.py \
 --file-list=tests/test_files/sample_training_0.tsv \
 --multinode=0 \
  --output-dir=outputs    --config_path configs/models/cc12m_64x64.yaml \
  -num_diffusion_steps=10 \
	--num-training-steps=10

您应该看到一个outputs/vis_model_000100.pth文件。现在让我们做一些更有意义的事:

让我们在CC12m上训练一个MDM模型

1. 数据准备:

(可选)使用此示例参数下载CC12m的前1000个文件

该脚本基于img2dataset的CC12M脚本。

curl https://storage.googleapis.com/conceptual_12m/cc12m.tsv | head -n 1000 > cc12m_index.tsv

# 向文件添加标题
sed -i '1s/^/url\tcaption\n/'  cc12m_index.tsv

注意:如果您想要所有的cc12m,请从调用中删除| head -n 1000

然后准备并分割为训练/验证集

此脚本需要img2dataset,运行pip install '.[data_prep]'或只运行pip install img2dataset

python3 -m ml_mdm.clis.scrape_cc12m \
  --cc12m_index cc12m_index.tsv \
  --cc12m_local_dir cc12m_download

运行此命令后,您将看到以下文件:

training.0.tsv # 训练索引文件
validation.tsv # 验证索引文件
cc12m_download/
   00000.parquet  00000.tar  00000.tsv  00000_stats.json  validation.tsv
   00001.parquet ....

2. 训练

现在我们有了训练文件,我们可以选择一个模型配置并传递任何额外的训练参数:

# 修改torchrun参数以适应您的GPU设置
torchrun --standalone --nproc_per_node=8 ml_mdm/clis/train_parallel.py \
  --file-list=training_0.tsv \
  --multinode=0 --output-dir=/mnt/data/outputs \
  --config_path configs/models/cc12m_64x64.yaml \
  --num-training-steps=100   --warmup-steps 10

注意:configs/models/cc12m_64x64.yaml包含更多参数,查看它以了解更多详情。 如果你已下载预训练模型，可以将 --pretrained-vision-file 参数设置为指向磁盘上的模型位置

训练完成后，你会在 --output-dir 参数定义的文件夹中找到模型：

2024-07-22:17:58:46,649 INFO     [model_ema.py:33] 保存 EMA 模型文件：/mnt/data/outputs/vis_model_000100.pth
2024-07-22:17:58:47,448 INFO     [unet.py:794] 保存模型文件：/mnt/data/outputs/vis_model_noema_000100.pth

3. 从模型中采样

现在我们有了训练好的模型，我们可以从扩散模型中生成样本：

torchrun --standalone --nproc_per_node=1 ml_mdm/clis/generate_batch.py \
  --config_path configs/models/cc12m_64x64.yaml \
  --min-examples 3 --test-file-list validation.tsv \
  --sample-image-size 64 --model-file /mnt/data/outputs/vis_model_000100.pth

如果你想跳过训练步骤，可以将 --model-file 更新为指向我们的预训练模型之一

数据集存储

对于长期存储，你可以选择将数据上传到 s3://{your_bucket}/datasets/{datasetname}/*.[tar,tsv]。

然后更新 configs/datasets/cc12m.yaml 以指向你的 s3 路径。

# configs/datasets/cc12m.yaml
train:
  files:
    - s3://mlx/datasets/cc12m-64x64/images_00.*.tsv
eval:
  files:
    - s3://mlx/datasets/cc12m-64x64/validation.tsv

# configs/datasets/reader_config.yaml
reader_config:
  append_eos: true
  bucket: ${your_bucket} # 添加你的 s3 存储桶
  endpoint_url: None # boto 将自动推断端点

然后你可以使用我们的数据集下载助手：

python -m ml_mdm.clis.download_tar_from_index \
  --dataset-config-file configs/datasets/cc12m.yaml \
  --subset train --download_tar

python -m ml_mdm.clis.download_tar_from_index \
  --dataset-config-file configs/datasets/cc12m.yaml \
  --subset eval --download_tar

S3 数据集选择

查看 configs/datasets/cc12m.yaml。

代码允许提供多个正则表达式。请记住，这些正则表达式不是通配符 -- 它们是来自 Python re 库的正则表达式。所以如果你想在训练中只使用 WebImage 中 1000 个 tar 文件中的 100 个，你可以这样做：

train:
  files:
    - s3://mlx/datasets/example-dataset-100M_64px/example-dataset-100M-00[0-1]..-[0-9]*-of-01000.tsv
eval:
  files:
    - s3://mlx/datasets/example-dataset-100M_64px/validation.tsv

你也可以混合搭配文件。所以如果你想合并 CC12m 和 ImageNet，你可以创建一个新的 yaml 文件，内容如下：

train:
  files:
    - s3://mlx/datasets/imagenet-64px/imagenet-train-000??-of-00100.tsv
    - s3://mlx/datasets/cc12m-64x64/images_00.*.tsv
eval:
  files:
    - s3://mlx/datasets/cc12m-64x64/validation.tsv

数据集结构

S3 存储桶包含一系列以这种格式的文件，查看 ml_mdm/clis/scrape_cc12m.py 以生成你自己的。

2023-04-01 01:31:30   36147200 images_00000.tar
2023-05-10 11:34:49    1108424 images_00000.tsv
2023-04-01 01:31:26   36454400 images_00001.tar
2023-05-10 11:34:49    1109588 images_00001.tsv
2023-04-01 01:31:53   36116480 images_00002.tar
...

这些文件的最小表示可以在 tests/test_files/ 中找到。

引用

如果你发现我们的工作有用，请考虑引用我们：

@misc{gu2023matryoshkadiffusionmodels,
      title={Matryoshka Diffusion Models},
      author={Jiatao Gu and Shuangfei Zhai and Yizhe Zhang and Josh Susskind and Navdeep Jaitly},
      year={2023},
      eprint={2310.15111},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2310.15111},
}
相关项目
项目侧边栏1项目侧边栏2
推荐项目
Project Cover

豆包MarsCode

豆包 MarsCode 是一款革命性的编程助手，通过AI技术提供代码补全、单测生成、代码解释和智能问答等功能，支持100+编程语言，与主流编辑器无缝集成，显著提升开发效率和代码质量。

Project Cover

AI写歌

Suno AI是一个革命性的AI音乐创作平台，能在短短30秒内帮助用户创作出一首完整的歌曲。无论是寻找创作灵感还是需要快速制作音乐，Suno AI都是音乐爱好者和专业人士的理想选择。

Project Cover

有言AI

有言平台提供一站式AIGC视频创作解决方案，通过智能技术简化视频制作流程。无论是企业宣传还是个人分享，有言都能帮助用户快速、轻松地制作出专业级别的视频内容。

Project Cover

Kimi

Kimi AI助手提供多语言对话支持，能够阅读和理解用户上传的文件内容，解析网页信息，并结合搜索结果为用户提供详尽的答案。无论是日常咨询还是专业问题，Kimi都能以友好、专业的方式提供帮助。

Project Cover

阿里绘蛙

绘蛙是阿里巴巴集团推出的革命性AI电商营销平台。利用尖端人工智能技术，为商家提供一键生成商品图和营销文案的服务，显著提升内容创作效率和营销效果。适用于淘宝、天猫等电商平台，让商品第一时间被种草。

Project Cover

吐司

探索Tensor.Art平台的独特AI模型，免费访问各种图像生成与AI训练工具，从Stable Diffusion等基础模型开始，轻松实现创新图像生成。体验前沿的AI技术，推动个人和企业的创新发展。

Project Cover

SubCat字幕猫

SubCat字幕猫APP是一款创新的视频播放器，它将改变您观看视频的方式！SubCat结合了先进的人工智能技术，为您提供即时视频字幕翻译，无论是本地视频还是网络流媒体，让您轻松享受各种语言的内容。

Project Cover

美间AI

美间AI创意设计平台，利用前沿AI技术，为设计师和营销人员提供一站式设计解决方案。从智能海报到3D效果图，再到文案生成，美间让创意设计更简单、更高效。

Project Cover

AIWritePaper论文写作

AIWritePaper论文写作是一站式AI论文写作辅助工具，简化了选题、文献检索至论文撰写的整个过程。通过简单设定，平台可快速生成高质量论文大纲和全文，配合图表、参考文献等一应俱全，同时提供开题报告和答辩PPT等增值服务，保障数据安全，有效提升写作效率和论文质量。

使用协议隐私政策广告服务
投诉举报邮箱: service@vectorlightyear.com
@2024 懂AI·鲁ICP备2024100362号-6·鲁公网安备37021002001498号