许可证: mit
流水线标签: 无条件图像生成
标签:

• 模型中心混合
• pytorch模型中心混合
  数据集:
• Artificio/WikiArt

该模型已通过PytorchModelHubMixin集成推送至Hub：

• 库地址: https://hf-mirror.com/kaupane/DiT-Wikiart-Small

模型描述

此模型是基于Wikiart数据集（https://huggingface.co/datasets/Artificio/WikiArt）从头训练的DiT（扩散变换器），旨在根据艺术流派和艺术风格生成艺术图像。

模型架构

该模型大体遵循论文《Scalable Diffusion Models with Transformers》中描述的经典DiT架构，并做了以下微调：

1. 将ImageNet类别嵌入替换为Wikiart流派和风格嵌入；
2. 使用后归一化而非前归一化；
3. 移除了最终的线性层；
4. 将sin-cos-2d位置嵌入替换为可学习的位置嵌入；
5. 模型仅预测噪声，不学习sigma；
6. 所有模型变体的patch_size均设为2；
7. 模型设有不同尺寸配置。
   如需了解更多细节，可查看本仓库中的modeling_dit_wikiart.py文件。

模型包含三个变体：

• S（小）：num_blocks=8，hidden_size=384，num_heads=6，总参数量20M；
• B（基础）：num_blocks=12，hidden_size=640，num_heads=10，总参数量90M；
• L（大）：num_blocks=16，hidden_size=896，num_heads=14，总参数量234M。

训练过程

• 数据集：所有模型变体均在103K的Wikiart数据集上训练，并采用水平翻转进行数据增强。
• 优化器：使用默认设置的AdamW。
• 学习率：前1%的步骤采用线性预热，学习率最高达3e-4，随后通过余弦衰减降至零。
• 训练周期与批次大小：
  • S：96个周期，批次大小176；
  • B：120个周期，批次大小192；
  • L：144个周期，批次大小192。
• 设备：
  • S：单块RTX 4060ti 16G，耗时24小时；
  • B：单块RTX 4060ti 16G，耗时90小时；
  • L：单块RTX 4090D 24G运行48小时，随后单块RTX 4060ti 16G运行100小时。
• 损失曲线：所有变体在首个周期内损失从1.0000以上骤降至约0.2000，随后缓慢下降，至第20个周期时达到loss=0.1600。最终，DiT-S的损失为0.1590，DiT-B为0.1525，DiT-L为0.1510。

性能与局限

• 模型展现出理解流派与风格并生成视觉吸引力画作的基本能力（初看之下）。
• 局限包括：
  • 难以理解复杂结构（如人脸、建筑等）；
  • 在生成数据集中罕见的流派或风格时可能出现模态崩溃（例如极简主义风格或浮世绘流派）；
  • 分辨率限制为256x256；
  • 基于Wikiart数据集训练，无法生成超出该范围的内容。

使用方法

使用前需安装"huggingface_hub"库，并从"Files and versions"下载modeling_dit_wikiart.py以获取模型定义。随后可通过以下代码调用模型：

from modeling_dit_wikiart import DiTWikiartModel  

model = DiTWikiartModel.from_pretrained("kaupane/DiT-Wikiart-Small")  
num_samples = 8  
noisy_latents = torch.randn(num_samples,4,32,32)  
predicted_noise = model(noisy_latents)  
print(predicted_noise)  

该模型需搭配stabilityai/sd-vae-ft-ema使用。