Hugging Face's logo

Diffusers documentation

SANA-Sprint

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
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 = 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 = ['latents']max_sequence_length: int = 300complex_human_instruction: list = ["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 = 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 = ['latents']max_sequence_length: int = 300complex_human_instruction: list = ["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→