
在 Segmind SSD-1B 试用该模型,体验⚡最快推理速度。您也可以在 🤗 Spaces 上试用。
Segmind Stable Diffusion 模型(SSD-1B)是 Stable Diffusion XL(SDXL)的蒸馏版,体积缩小50%,在保持高质量文本到图像生成能力的同时,提供60%的速度提升。它在多样化的数据集上进行了训练,包括 Grit 和 Midjourney 抓取数据,以增强其根据文本提示创建各种视觉内容的能力。
该模型采用知识蒸馏策略,依次借鉴了多个专家模型的经验,包括 SDXL、ZavyChromaXL 和 JuggernautXL,以融合它们的优势并生成令人印象深刻的视觉输出。
特别感谢 HF 团队🤗,尤其是 Sayak、Patrick 和 Poli 在这项工作中的合作与指导。

import os
import argparse
import torch
from PIL import Image
from diffusers import DiffusionPipeline
from openmind import is_torch_npu_available
from openmind_hub import snapshot_download
def parse_args():
parser = argparse.ArgumentParser(description="Eval the model")
parser.add_argument(
"--model_name_or_path",
type=str,
help="Path to the model",
default="jeffding/SSD-1B-openmind",
)
args = parser.parse_args()
return args
def main():
args = parse_args()
model_path = args.model_name_or_path
if is_torch_npu_available():
device = "npu:0"
else:
device = "cpu"
pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
pipe.to(device)
prompt = "An astronaut riding a green horse."
image = pipe(prompt=prompt).images[0]
image.save("astronaut_rides_horse.png")
if __name__ == '__main__':
main()
此模型可通过 🧨 Diffusers 库使用。
请确保通过运行以下命令从源代码安装 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")
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]文本到图像生成: 该模型擅长根据文本提示生成图像,可支持广泛的创意应用。
为速度而蒸馏: 专为效率设计,此模型可提供 60% 的速度提升,使其成为实时应用以及需要快速图像生成场景的实用选择。
多样化训练数据: 经过多样化数据集训练,该模型能够处理各种文本提示并有效生成相应图像。
知识蒸馏: 通过从多个专家模型中蒸馏知识,Segmind 稳定扩散模型融合了它们的优势并最大限度地减少了其局限性,从而实现了性能提升。
SSD-1B 模型是一个拥有 13 亿参数的模型,它从基础 SDXL 模型中移除了多个层。

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

SSD-1B 可支持以下输出分辨率:
我们观察到,SSD-1B 比基础 SDXL 模型快高达 60%。以下是在 A100 80GB 上的对比结果。

以下是在 RTX 4090 GPU 上的加速指标。

出于研发目的,SSD-1B 模型可通过 Segmind AI 平台访问。如需更多信息和访问详情,请访问 Segmind。
Segmind Stable Diffusion 模型适用于多个领域的研究和实际应用,包括:
艺术与设计:可用于生成艺术品、设计方案及其他创意内容,为创作过程提供灵感并增强创意。
教育:该模型可应用于教育工具,为教学和学习创建视觉内容。
研究:研究人员可利用该模型探索生成式模型,评估其性能,并推动文本到图像生成技术的发展边界。
安全内容生成:提供安全可控的内容生成方式,降低生成有害或不当内容的风险。
偏差与局限性分析:研究人员和开发人员可使用该模型探究其局限性和偏差,有助于更好地理解生成式模型的行为。
Segmind Stable Diffusion 模型还可直接与 🧨 Diffusers 库的训练脚本配合使用,进行进一步训练,包括:
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_hubexport 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_hubexport 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_hubSSD-1B 模型不适用于创建人物、事件或现实世界信息的真实或准确呈现。它不用于需要高精度和准确性的任务。
局限性与偏见 SSD-1B 模型在实现绝对照片真实感方面存在一些挑战,尤其是在人物描绘上。由于其自编码方法,该模型在整合清晰文本和保持复杂构图的保真度方面存在困难,但这些障碍为未来的改进铺平了道路。重要的是,尽管模型接触到多样化的数据集并非解决根深蒂固的社会和数字偏见的万能药,但这代表了迈向更公平技术的基础一步。鼓励用户在了解其当前局限性的情况下使用这一开创性工具,营造有意识参与的环境,并期待其持续发展。
@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}
}