Hugging Face's logo

Diffusers documentation

SANA-Sprint

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

SANA-Sprint

LoRA

来自 NVIDIA、MIT HAN Lab 和 Hugging Face 的 SANA-Sprint: One-Step Diffusion with Continuous-Time Consistency Distillation，作者为 Junsong Chen、Shuchen Xue、Yuyang Zhao、Jincheng Yu、Sayak Paul、Junyu Chen、Han Cai、Enze Xie、Song Han

该论文的摘要如下：

本文提出了 SANA-Sprint，一种用于超快速文本到图像（T2I）生成的高效扩散模型。SANA-Sprint 基于预训练的基础模型构建，并辅以混合蒸馏技术，将推理步数从 20 步大幅减少至 1-4 步。我们引入了三项关键创新：(1) 我们提出了一种无需训练的方法，将预训练的流匹配模型转化为连续时间一致性蒸馏（sCM）模型，从而消除了从头开始的高昂训练成本，并实现了极高的训练效率。我们的混合蒸馏策略将 sCM 与潜在对抗蒸馏（LADD）相结合：sCM 确保与教师模型的对齐，而 LADD 则增强了单步生成的保真度。(2) SANA-Sprint 是一个统一的步数自适应模型，能够在 1-4 步内实现高质量生成，消除了针对特定步数的训练需求，从而提升了效率。(3) 我们将 ControlNet 与 SANA-Sprint 集成，用于实时交互式图像生成，使用户交互能够即时获得视觉反馈。SANA-Sprint 在速度-质量权衡中确立了新的帕累托前沿，仅用 1 步便实现了最先进的性能，FID 为 7.59，GenEval 为 0.74——在 H100 上比 FLUX-schnell（FID 7.94 / GenEval 0.71）快 10 倍（0.1 秒 vs 1.1 秒）。此外，它在 H100 上生成 1024×1024 图像的延迟为 0.1 秒（T2I）和 0.25 秒（ControlNet），在 RTX 4090 上为 0.31 秒（T2I），展示了其卓越的效率以及在 AI 驱动的消费者应用（AIPC）中的潜力。代码和预训练模型将开源。

该流水线由 lawrence-cj、shuchen Xue 和 Enze Xie 贡献。原始代码库可在 此处 找到。原始权重可在 hf.co/Efficient-Large-Model 下找到。

可用模型：

模型推荐的 dtype
Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffuserstorch.bfloat16
Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffuserstorch.bfloat16

有关更多信息，请参阅此集合。

注意：上述推荐的 dtype 是针对 transformer 权重的。为了让模型正确运行，文本编码器必须保持为 torch.bfloat16，VAE 权重必须保持为 torch.bfloat16 或 torch.float32。请参阅下方的推理示例，了解如何以推荐的 dtype 加载模型。

量化

量化通过以较低精度的数据类型存储模型权重，有助于降低超大模型的内存需求。不过，量化对视频质量的影响会因视频模型的不同而有所差异。

请参阅 量化 概述，以了解支持的量化后端以及如何选择支持您使用场景的量化后端。以下示例演示了如何使用 bitsandbytes 加载用于推理的量化 SanaSprintPipeline。

import torch
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig, SanaTransformer2DModel, SanaSprintPipeline
from transformers import BitsAndBytesConfig as BitsAndBytesConfig, AutoModel

quant_config = BitsAndBytesConfig(load_in_8bit=True)
text_encoder_8bit = AutoModel.from_pretrained(
    "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers",
    subfolder="text_encoder",
    quantization_config=quant_config,
    dtype=torch.bfloat16,
)

quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
transformer_8bit = SanaTransformer2DModel.from_pretrained(
    "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.bfloat16,
)

pipeline = SanaSprintPipeline.from_pretrained(
    "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers",
    text_encoder=text_encoder_8bit,
    transformer=transformer_8bit,
    dtype=torch.bfloat16,
    device_map="balanced",
)

prompt = "a tiny astronaut hatching from an egg on the moon"
image = pipeline(prompt).images[0]
image.save("sana.png")

设置 max_timesteps

用户可以调整 max_timesteps 的值，以实验生成输出的视觉质量。默认的 max_timesteps 值是通过推理时搜索过程获得的。有关更多详细信息，请参阅该论文。

图像到图像

SanaSprintImg2ImgPipeline 是一个用于图像到图像生成的流水线。它接收一张输入图像和一段提示词，并根据输入图像和提示词生成一张新图像。

import torch
from diffusers import SanaSprintImg2ImgPipeline
from diffusers.utils.loading_utils import load_image

image = load_image(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png"
)

pipe = SanaSprintImg2ImgPipeline.from_pretrained(
    "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers", 
    dtype=torch.bfloat16)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

image = pipe(
    prompt="a cute pink bear", 
    image=image, 
    strength=0.5, 
    height=832, 
    width=480
).images[0]
image.save("output.png")

SanaSprintPipeline

类 diffusers.SanaSprintPipeline

< >

( tokenizer: GemmaTokenizertext_encoder: Gemma2PreTrainedModelvae: AutoencoderDCtransformer: SanaTransformer2DModelscheduler: DPMSolverMultistepScheduler )

使用 SANA-Sprint 进行文本到图像生成的 Pipeline。

__call__

< >

( prompt: str | list[str] = Nonenum_inference_steps: int = 2timesteps: list[int] = Nonemax_timesteps: float = 1.5708intermediate_timesteps: float = 1.3guidance_scale: float = 4.5num_images_per_prompt: int | None = 1height: int = 1024width: int = 1024eta: float = 0.0generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonelatents: typing.Optional[torch.Tensor] = Noneprompt_embeds: typing.Optional[torch.Tensor] = Noneprompt_attention_mask: typing.Optional[torch.Tensor] = Noneoutput_type: str | None = 'pil'return_dict: bool = Trueclean_caption: bool = Falseuse_resolution_binning: bool = Trueattention_kwargs: dict[str, typing.Any] | None = Nonecallback_on_step_end: typing.Optional[typing.Callable[[int, int], NoneType]] = Nonecallback_on_step_end_tensor_inputs: list[str] = ['latents']max_sequence_length: int = 300complex_human_instruction: list[str] = ["Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:", '- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.', '- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.', 'Here are examples of how to transform or refine prompts:', '- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.', '- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.', 'Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:', 'User Prompt: '] ) → SanaPipelineOutput or tuple

参数

  • prompt (str 或 list[str]，可选) — 用于引导图像生成的提示词或提示词列表。如果未定义，则必须改为传入 prompt_embeds。
  • num_inference_steps (int, 可选，默认为 20) — 去噪步数。更多的去噪步数通常能生成更高质量的图像，但代价是推理速度变慢。
  • max_timesteps (float，可选，默认值为 1.57080) — SCM 调度器中使用的最大时间步值。
  • intermediate_timesteps（float，可选，默认值为 1.3）— SCM 调度器中使用的中间时间步数值（仅在 num_inference_steps=2 时使用）。
  • timesteps (list[int], 可选) — 用于支持 timesteps 参数的调度器去噪过程的自定义时间步，该参数位于其 set_timesteps 方法中。如果未定义，则使用传入 num_inference_steps 时的默认行为。必须按降序排列。
  • guidance_scale (float, 可选, ，默认值为 4.5）—— 通过设置 guidance_scale > 1. 更高的 guidance_scale 值，可以促使模型生成与 prompt 更加一致的图像，但代价是图像质量会降低。

    Guidance-distilled 模型近似了 guidance_scale > 1 时的真实 classifier-free guidance。有关详细信息，请参阅 论文。

  • num_images_per_prompt (int，可选，默认为 1) — 每个提示词生成的图像数量。
  • height (int，可选，默认为 self.unet.config.sample_size) — 生成图像的高度（以像素为单位）。
  • width (int，可选，默认为 self.unet.config.sample_size) — 生成图像的宽度（以像素为单位）。
  • eta（float，可选，默认值为 0.0）— 对应 DDIM 论文中的参数 eta (η)：https://huggingface.co/papers/2010.02502。仅适用于 schedulers.DDIMScheduler，对其他调度器将被忽略。
  • generator (torch.Generator 或 list[torch.Generator]，可选) — 一个或一组 torch 生成器，用于使生成过程具有确定性。
  • latents (torch.Tensor, 可选) — 预先生成的噪声潜变量，从高斯分布中采样得到，用作图像生成的输入。可用于使用不同的提示词对同一生成过程进行微调。如果未提供，将通过使用提供的随机generator进行采样来生成一个潜变量张量。
  • prompt_embeds (torch.Tensor, 可选) — 预生成的文本嵌入。可用于轻松调整文本输入，例如 提示词加权。如果未提供，文本嵌入将从 prompt 输入参数生成。
  • prompt_attention_mask (torch.Tensor，可选) — 为文本嵌入预生成的注意力掩码。
  • output_type (str，可选，默认为 "pil") — 生成图像的输出格式。可选值为 PIL：PIL.Image.Image 或 np.array。
  • return_dict (bool，可选，默认为 True) — 是否返回 ~pipelines.stable_diffusion.IFPipelineOutput 而不是普通元组。
  • attention_kwargs — 一个 kwargs 字典，如果指定，将传递给 AttentionProcessor，该处理器在 diffusers.models.attention_processor 中定义，并位于 self.processor 下。
  • clean_caption (bool，可选，默认为 True) — 是否在创建嵌入之前清理字幕。需要安装 beautifulsoup4 和 ftfy。如果未安装这些依赖项，嵌入将直接从原始提示创建。
  • use_resolution_binning (bool，默认为 True) — 如果设置为 True，请求的高度和宽度会先通过 ASPECT_RATIO_1024_BIN 映射到最接近的分辨率。生成的潜变量解码为图像后，会被重新调整回请求的分辨率。适用于生成非正方形图像。
  • callback_on_step_end (Callable, 可选) — 一个在推理过程中每个去噪步骤结束时调用的函数。该函数以以下参数被调用：callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)。callback_kwargs 将包含由 callback_on_step_end_tensor_inputs 指定的所有张量列表。
  • callback_on_step_end_tensor_inputs (list，可选) — callback_on_step_end 函数的张量输入列表。列表中指定的张量将作为 callback_kwargs 参数传递。您只能包含流水线类 ._callback_tensor_inputs 属性中列出的变量。
  • max_sequence_length（int，默认为 300）— 与 prompt 配合使用的最大序列长度。
  • complex_human_instruction (list[str], 可选) — 用于复杂人类注意力的指令： https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55。

返回值

SanaPipelineOutput 或 tuple

如果 return_dict 为 True，则返回 SanaPipelineOutput，否则返回一个 tuple，其中第一个元素是包含生成图像的列表

调用生成 pipeline 时调用的函数。

示例：

>>> import torch
>>> from diffusers import SanaSprintPipeline

>>> pipe = SanaSprintPipeline.from_pretrained(
...     "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers", torch_dtype=torch.bfloat16
... )
>>> pipe.to("cuda")

>>> image = pipe(prompt="a tiny astronaut hatching from an egg on the moon")[0]
>>> image[0].save("output.png")

encode_prompt

< >

( prompt: str | list[str]num_images_per_prompt: int = 1device: typing.Optional[torch.device] = Noneprompt_embeds: typing.Optional[torch.Tensor] = Noneprompt_attention_mask: typing.Optional[torch.Tensor] = Noneclean_caption: bool = Falsemax_sequence_length: int = 300complex_human_instruction: list[str] | None = Nonelora_scale: float | None = None )

参数

  • prompt (str 或 list[str]，可选) — 要编码的提示
  • num_images_per_prompt (int，可选，默认为 1) — 每个提示词应生成的图像数量
  • device — (torch.device，可选)： 用于放置生成的嵌入向量的 torch 设备
  • prompt_embeds (torch.Tensor, 可选) — 预生成的文本嵌入。可用于轻松调整文本输入，例如 提示词加权。如果未提供，文本嵌入将从 prompt 输入参数生成。
  • clean_caption (bool，默认为 False) — 如果为 True，该函数将在编码之前对提供的 caption 进行预处理和清理。
  • max_sequence_length (int，默认为 300) — 用于提示词的最大序列长度。
  • complex_human_instruction (list[str]，默认为 complex_human_instruction) — 如果 complex_human_instruction 不为空，该函数将使用复杂的 Human 指令作为提示词。

将提示词编码为文本编码器的隐藏状态。

SanaSprintImg2ImgPipeline

类 diffusers.SanaSprintImg2ImgPipeline

< >

( tokenizer: GemmaTokenizertext_encoder: Gemma2PreTrainedModelvae: AutoencoderDCtransformer: SanaTransformer2DModelscheduler: DPMSolverMultistepScheduler )

使用 SANA-Sprint 进行文本到图像生成的 Pipeline。

__call__

< >

( prompt: str | list[str] = Nonenum_inference_steps: int = 2timesteps: list[int] = Nonemax_timesteps: float = 1.5708intermediate_timesteps: float = 1.3guidance_scale: float = 4.5image: typing.Union[PIL.Image.Image, numpy.ndarray, torch.Tensor, list[PIL.Image.Image], list[numpy.ndarray], list[torch.Tensor]] = Nonestrength: float = 0.6num_images_per_prompt: int | None = 1height: int = 1024width: int = 1024eta: float = 0.0generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonelatents: typing.Optional[torch.Tensor] = Noneprompt_embeds: typing.Optional[torch.Tensor] = Noneprompt_attention_mask: typing.Optional[torch.Tensor] = Noneoutput_type: str | None = 'pil'return_dict: bool = Trueclean_caption: bool = Falseuse_resolution_binning: bool = Trueattention_kwargs: dict[str, typing.Any] | None = Nonecallback_on_step_end: typing.Optional[typing.Callable[[int, int], NoneType]] = Nonecallback_on_step_end_tensor_inputs: list[str] = ['latents']max_sequence_length: int = 300complex_human_instruction: list[str] = ["Given a user prompt, generate an 'Enhanced prompt' that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:", '- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.', '- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.', 'Here are examples of how to transform or refine prompts:', '- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.', '- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.', 'Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:', 'User Prompt: '] ) → SanaPipelineOutput or tuple

参数

  • prompt (str 或 list[str]，可选) — 用于引导图像生成的提示词或提示词列表。如果未定义，则必须改为传入 prompt_embeds。
  • num_inference_steps (int, 可选，默认为 20) — 去噪步数。更多的去噪步数通常能生成更高质量的图像，但代价是推理速度变慢。
  • max_timesteps (float，可选，默认值为 1.57080) — SCM 调度器中使用的最大时间步值。
  • intermediate_timesteps（float，可选，默认值为 1.3）— SCM 调度器中使用的中间时间步数值（仅在 num_inference_steps=2 时使用）。
  • timesteps (list[int], 可选) — 用于支持 timesteps 参数的调度器去噪过程的自定义时间步，该参数位于其 set_timesteps 方法中。如果未定义，则使用传入 num_inference_steps 时的默认行为。必须按降序排列。
  • guidance_scale (float，可选，默认值为 4.5) — 如 Classifier-Free Diffusion Guidance 中定义的引导尺度。 guidance_scale 定义为 Imagen Paper 中公式 2 的 w。通过设置 guidance_scale > 1 来启用引导尺度。较高的引导尺度会促使生成与文本 prompt 紧密关联的图像， 通常以牺牲图像质量为代价。
  • image (torch.Tensor、PIL.Image.Image、np.ndarray、list[torch.Tensor]、list[PIL.Image.Image] 或 list[np.ndarray]) — Image、numpy 数组或张量，表示用作起点的图像批次。也可以将图像潜变量作为 image 传入，但如果直接传入潜变量，则不会再次进行编码。
  • strength (float，可选，默认值为 0.6) — 指示对参考 image 进行变换的程度。取值必须在 0 到 1 之间。image 用作起始点，strength 越高，添加的噪声越多。去噪步数取决于初始添加的噪声量。当 strength 为 1 时，添加的噪声达到最大值，去噪过程将运行 num_inference_steps 中指定的完整迭代次数。值为 1 时，基本上会忽略 image。
  • num_images_per_prompt (int，可选，默认为 1) — 每个提示词生成的图像数量。
  • height (int，可选，默认为 self.unet.config.sample_size) — 生成图像的高度（以像素为单位）。
  • width (int，可选，默认为 self.unet.config.sample_size) — 生成图像的宽度（以像素为单位）。
  • eta (float, 可选，默认为 0.0) — 对应 DDIM 论文中的参数 eta (η)：https://arxiv.org/abs/2010.02502。仅适用于 schedulers.DDIMScheduler，对其他调度器将被忽略。
  • generator (torch.Generator 或 list[torch.Generator]，可选) — 一个或一组 torch 生成器，用于使生成过程具有确定性。
  • latents (torch.Tensor, 可选) — 预先生成的噪声潜变量，从高斯分布中采样得到，用作图像生成的输入。可用于使用不同的提示词对同一生成过程进行微调。如果未提供，将通过使用提供的随机generator进行采样来生成一个潜变量张量。
  • prompt_embeds (torch.Tensor, 可选) — 预生成的文本嵌入。可用于轻松调整文本输入，例如 提示词加权。如果未提供，文本嵌入将从 prompt 输入参数生成。
  • prompt_attention_mask (torch.Tensor，可选) — 为文本嵌入预生成的注意力掩码。
  • output_type (str，可选，默认为 "pil") — 生成图像的输出格式。可选值为 PIL：PIL.Image.Image 或 np.array。
  • return_dict (bool，可选，默认为 True) — 是否返回 ~pipelines.stable_diffusion.IFPipelineOutput 而不是普通元组。
  • attention_kwargs — 一个 kwargs 字典，如果指定，将传递给 AttentionProcessor，该处理器在 diffusers.models.attention_processor 中定义，并位于 self.processor 下。
  • clean_caption (bool，可选，默认为 True) — 是否在创建嵌入之前清理字幕。需要安装 beautifulsoup4 和 ftfy。如果未安装这些依赖项，嵌入将直接从原始提示创建。
  • use_resolution_binning (bool，默认为 True) — 如果设置为 True，请求的高度和宽度会先通过 ASPECT_RATIO_1024_BIN 映射到最接近的分辨率。生成的潜变量解码为图像后，会被重新调整回请求的分辨率。适用于生成非正方形图像。
  • callback_on_step_end (Callable, 可选) — 一个在推理过程中每个去噪步骤结束时调用的函数。该函数以以下参数被调用：callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)。callback_kwargs 将包含由 callback_on_step_end_tensor_inputs 指定的所有张量列表。
  • callback_on_step_end_tensor_inputs (list，可选) — callback_on_step_end 函数的张量输入列表。列表中指定的张量将作为 callback_kwargs 参数传递。您只能包含流水线类 ._callback_tensor_inputs 属性中列出的变量。
  • max_sequence_length（int，默认为 300）— 与 prompt 配合使用的最大序列长度。
  • complex_human_instruction (list[str], 可选) — 用于复杂人类注意力的指令： https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55。

返回值

SanaPipelineOutput 或 tuple

如果 return_dict 为 True，则返回 SanaPipelineOutput，否则返回一个 tuple，其中第一个元素是包含生成图像的列表

调用生成 pipeline 时调用的函数。

示例：

>>> import torch
>>> from diffusers import SanaSprintImg2ImgPipeline
>>> from diffusers.utils.loading_utils import load_image

>>> pipe = SanaSprintImg2ImgPipeline.from_pretrained(
...     "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers", torch_dtype=torch.bfloat16
... )
>>> pipe.to("cuda")

>>> image = load_image(
...     "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png"
... )


>>> image = pipe(prompt="a cute pink bear", image=image, strength=0.5, height=832, width=480).images[0]
>>> image[0].save("output.png")

encode_prompt

< >

( prompt: str | list[str]num_images_per_prompt: int = 1device: typing.Optional[torch.device] = Noneprompt_embeds: typing.Optional[torch.Tensor] = Noneprompt_attention_mask: typing.Optional[torch.Tensor] = Noneclean_caption: bool = Falsemax_sequence_length: int = 300complex_human_instruction: list[str] | None = Nonelora_scale: float | None = None )

参数

  • prompt (str 或 list[str]，可选) — 要编码的提示
  • num_images_per_prompt (int，可选，默认为 1) — 每个提示词应生成的图像数量
  • device — (torch.device，可选)： 用于放置生成的嵌入向量的 torch 设备
  • prompt_embeds (torch.Tensor, 可选) — 预生成的文本嵌入。可用于轻松调整文本输入，例如 提示词加权。如果未提供，文本嵌入将从 prompt 输入参数生成。
  • clean_caption (bool，默认为 False) — 如果为 True，该函数将在编码之前对提供的 caption 进行预处理和清理。
  • max_sequence_length (int，默认为 300) — 用于提示词的最大序列长度。
  • complex_human_instruction (list[str]，默认为 complex_human_instruction) — 如果 complex_human_instruction 不为空，该函数将使用复杂的 Human 指令作为提示词。

将提示词编码为文本编码器的隐藏状态。

SanaPipelineOutput

class diffusers.pipelines.sana.pipeline_output.SanaPipelineOutput

< >

( images: list[PIL.Image.Image] | numpy.ndarray )

参数

  • images (list[PIL.Image.Image] 或 np.ndarray) — 长度为 batch_size 的去噪 PIL 图像列表，或形状为 (batch_size, height, width, num_channels) 的 numpy 数组。PIL 图像或 numpy 数组表示扩散流水线生成的去噪图像。

Sana 管道的输出类。

Update on GitHub

←Sana Sana Video→