SSD 1B

🚀 Segmind Stable Diffusion 1B (SSD-1B) 模型卡片

Segmind Stable Diffusion 1B (SSD-1B) 是一个基于知识蒸馏的文本到图像生成模型，相比 Stable Diffusion XL (SDXL) 体积缩小 50%，速度提升 60%，同时保持了高质量的图像生成能力。

🚀 快速开始

兼容性说明

• AUTOMATIC1111 兼容性已添加，支持文件 点击此处 获取。

在线体验

• 你可以在 Segmind SSD-1B 体验该模型，以获得 ⚡ 最快推理速度。也可以在 🤗 Spaces 上尝试。

安装依赖

确保通过以下命令从源代码安装 diffusers：

pip install git+https://github.com/huggingface/diffusers

此外，请安装 transformers、safetensors 和 accelerate：

pip install transformers accelerate safetensors

使用示例

from diffusers import StableDiffusionXLPipeline
import torch

# 从预训练模型加载管道
pipe = StableDiffusionXLPipeline.from_pretrained("segmind/SSD-1B", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
# 将模型移动到 GPU 上
pipe.to("cuda")

# 如果使用的 PyTorch 版本低于 2.0，可以启用 xformers 以提高内存效率
# pipe.enable_xformers_memory_efficient_attention()

# 定义文本提示
prompt = "An astronaut riding a green horse"  # 你的文本提示
neg_prompt = "ugly, blurry, poor quality"  # 负向提示

# 生成图像
image = pipe(prompt=prompt, negative_prompt=neg_prompt).images[0]

更新说明

• 现在该模型可以在 ComfyUI 中使用。
• 建议使用负向提示，并将 CFG 设置为 9.0 左右以获得最佳质量！

✨ 主要特性

• 文本到图像生成：该模型擅长根据文本提示生成图像，可应用于各种创意领域。
• 蒸馏加速：经过知识蒸馏，模型体积缩小 50%，速度提升 60%，适用于实时应用场景。
• 多样化训练数据：在多种数据集上进行训练，能够处理各种文本提示并生成相应的高质量图像。
• 知识蒸馏策略：借鉴多个专家模型的优势，结合它们的长处，减少局限性，提高性能。

📦 安装指南

环境依赖

• 确保已安装 diffusers、transformers、safetensors 和 accelerate 库。安装命令如下：

pip install git+https://github.com/huggingface/diffusers
pip install transformers accelerate safetensors

💻 使用示例

基础用法

from diffusers import StableDiffusionXLPipeline
import torch

pipe = StableDiffusionXLPipeline.from_pretrained("segmind/SSD-1B", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
pipe.to("cuda")
# if using torch < 2.0
# pipe.enable_xformers_memory_efficient_attention()
prompt = "An astronaut riding a green horse" # Your prompt here
neg_prompt = "ugly, blurry, poor quality" # Negative prompt here
image = pipe(prompt=prompt, negative_prompt=neg_prompt).images[0]

高级用法

LoRA 微调

export MODEL_NAME="segmind/SSD-1B"
export VAE_NAME="madebyollin/sdxl-vae-fp16-fix"
export DATASET_NAME="lambdalabs/pokemon-blip-captions"

accelerate launch train_text_to_image_lora_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --pretrained_vae_model_name_or_path=$VAE_NAME \
  --dataset_name=$DATASET_NAME --caption_column="text" \
  --resolution=1024 --random_flip \
  --train_batch_size=1 \
  --num_train_epochs=2 --checkpointing_steps=500 \
  --learning_rate=1e-04 --lr_scheduler="constant" --lr_warmup_steps=0 \
  --mixed_precision="fp16" \
  --seed=42 \
  --output_dir="sd-pokemon-model-lora-ssd" \
  --validation_prompt="cute dragon creature" --report_to="wandb" \
  --push_to_hub

全量微调

export MODEL_NAME="segmind/SSD-1B"
export VAE_NAME="madebyollin/sdxl-vae-fp16-fix"
export DATASET_NAME="lambdalabs/pokemon-blip-captions"

accelerate launch train_text_to_image_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --pretrained_vae_model_name_or_path=$VAE_NAME \
  --dataset_name=$DATASET_NAME \
  --enable_xformers_memory_efficient_attention \
  --resolution=512 --center_crop --random_flip \
  --proportion_empty_prompts=0.2 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 --gradient_checkpointing \
  --max_train_steps=10000 \
  --use_8bit_adam \
  --learning_rate=1e-06 --lr_scheduler="constant" --lr_warmup_steps=0 \
  --mixed_precision="fp16" \
  --report_to="wandb" \
  --validation_prompt="a cute Sundar Pichai creature" --validation_epochs 5 \
  --checkpointing_steps=5000 \
  --output_dir="ssd-pokemon-model" \
  --push_to_hub

Dreambooth LoRA

export MODEL_NAME="segmind/SSD-1B"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="lora-trained-xl"
export VAE_PATH="madebyollin/sdxl-vae-fp16-fix"

accelerate launch train_dreambooth_lora_sdxl.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --pretrained_vae_model_name_or_path=$VAE_PATH \
  --output_dir=$OUTPUT_DIR \
  --mixed_precision="fp16" \
  --instance_prompt="a photo of sks dog" \
  --resolution=1024 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --learning_rate=1e-5 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=25 \
  --seed="0" \
  --push_to_hub

📚 详细文档

模型描述

模型架构

SSD-1B 模型是一个 13 亿参数的模型，从基础 SDXL 模型中移除了若干层。

训练信息

以下是训练过程中使用的关键超参数：

• 步数：251000
• 学习率：1e-5
• 批次大小：32
• 梯度累积步数：4
• 图像分辨率：1024
• 混合精度：fp16

多分辨率支持

SSD-1B 支持以下输出分辨率：

• 1024 x 1024 (1:1 方形)
• 1152 x 896 (9:7)
• 896 x 1152 (7:9)
• 1216 x 832 (19:13)
• 832 x 1216 (13:19)
• 1344 x 768 (7:4 横向)
• 768 x 1344 (4:7 纵向)
• 1536 x 640 (12:5 横向)
• 640 x 1536 (5:12 纵向)

速度比较

• 与基础 SDXL 模型相比，SSD-1B 速度最多可提升 60%。以下是在 A100 80GB 上的比较结果：
• 以下是在 RTX 4090 GPU 上的加速指标：

模型来源

如需进行研究和开发，可通过 Segmind AI 平台访问 SSD-1B 模型。更多信息和访问详情，请访问 Segmind。

用途

直接使用

Segmind Stable Diffusion 模型适用于多个领域的研究和实际应用，包括：

• 艺术与设计：可用于生成艺术作品、设计和其他创意内容，为创作过程提供灵感。
• 教育：可应用于教育工具，为教学和学习创建视觉内容。
• 研究：研究人员可使用该模型探索生成模型，评估其性能，推动文本到图像生成技术的发展。
• 安全内容生成：提供了一种安全可控的内容生成方式，降低有害或不当输出的风险。
• 偏差和局限性分析：研究人员和开发人员可使用该模型探索其局限性和偏差，有助于更好地理解生成模型的行为。

下游使用

Segmind Stable Diffusion 模型还可以直接与 🧨 Diffusers 库的训练脚本一起使用，进行进一步训练，包括：

• LoRA
• 全量微调
• Dreambooth LoRA

不适用场景

SSD-1B 模型不适合创建人物、事件或现实世界信息的事实或准确表示，不用于需要高精度和准确性的任务。

🔧 技术细节

蒸馏方法

📣 阅读我们的 技术报告 以获取有关蒸馏方法的更多详细信息。

局限性和偏差

SSD-1B 模型在实现绝对的照片级真实感方面存在一些挑战，特别是在人物描绘方面。由于其自动编码方法，在融入清晰文本和保持复杂构图的保真度方面存在困难，这些问题为未来的改进提供了方向。重要的是，尽管模型接触了多样化的数据集，但这并不能完全消除根深蒂固的社会和数字偏差，不过这是迈向更公平技术的基础一步。建议用户在使用此开创性工具时，了解其当前的局限性，促进有意识的使用，并期待其持续发展。

📄 许可证

本模型使用 Apache 2.0 许可证。

引用

@misc{gupta2024progressive,
      title={Progressive Knowledge Distillation Of Stable Diffusion XL Using Layer Level Loss}, 
      author={Yatharth Gupta and Vishnu V. Jaddipal and Harish Prabhala and Sayak Paul and Patrick Von Platen},
      year={2024},
      eprint={2401.02677},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}