#
SANA-Sprint

SANA-Sprint: One-Step Diffusion with Continuous-Time Consistency Distillation from NVIDIA, MIT HAN Lab, and Hugging Face by Junsong Chen, Shuchen Xue, Yuyang Zhao, Jincheng Yu, Sayak Paul, Junyu Chen, Han Cai, Enze Xie, Song Han

The abstract from the paper is:

This paper presents SANA-Sprint, an efficient diffusion model for ultra-fast text-to-image (T2I) generation. SANA-Sprint is built on a pre-trained foundation model and augmented with hybrid distillation, dramatically reducing inference steps from 20 to 1-4. We introduce three key innovations: (1) We propose a training-free approach that transforms a pre-trained flow-matching model for continuous-time consistency distillation (sCM), eliminating costly training from scratch and achieving high training efficiency. Our hybrid distillation strategy combines sCM with latent adversarial distillation (LADD): sCM ensures alignment with the teacher model, while LADD enhances single-step generation fidelity. (2) SANA-Sprint is a unified step-adaptive model that achieves high-quality generation in 1-4 steps, eliminating step-specific training and improving efficiency. (3) We integrate ControlNet with SANA-Sprint for real-time interactive image generation, enabling instant visual feedback for user interaction. SANA-Sprint establishes a new Pareto frontier in speed-quality tradeoffs, achieving state-of-the-art performance with 7.59 FID and 0.74 GenEval in only 1 step — outperforming FLUX-schnell (7.94 FID / 0.71 GenEval) while being 10× faster (0.1s vs 1.1s on H100). It also achieves 0.1s (T2I) and 0.25s (ControlNet) latency for 1024×1024 images on H100, and 0.31s (T2I) on an RTX 4090, showcasing its exceptional efficiency and potential for AI-powered consumer applications (AIPC). Code and pre-trained models will be open-sourced.

Make sure to check out the Schedulers guide to learn how to explore the tradeoff between scheduler speed and quality, and see the reuse components across pipelines section to learn how to efficiently load the same components into multiple pipelines.

This pipeline was contributed by lawrence-cj, shuchen Xue and Enze Xie. The original codebase can be found here. The original weights can be found under hf.co/Efficient-Large-Model.

Available models:

Model | Recommended dtype
--- | ---
`Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers` | `torch.bfloat16`
`Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers` | `torch.bfloat16`

Refer to this collection for more information.

Note: The recommended dtype mentioned is for the transformer weights. The text encoder must stay in `torch.bfloat16` and VAE weights must stay in `torch.bfloat16` or `torch.float32` for the model to work correctly. Please refer to the inference example below to see how to load the model with the recommended dtype.

##
Quantization

Quantization helps reduce the memory requirements of very large models by storing model weights in a lower precision data type. However, quantization may have varying impact on video quality depending on the video model.

Refer to the Quantization overview to learn more about supported quantization backends and selecting a quantization backend that supports your use case. The example below demonstrates how to load a quantized SanaSprintPipeline for inference with bitsandbytes.

    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,
        torch_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,
        torch_dtype=torch.bfloat16,
    )

    pipeline = SanaSprintPipeline.from_pretrained(
        "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers",
        text_encoder=text_encoder_8bit,
        transformer=transformer_8bit,
        torch_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")

##
Setting max_timesteps

Users can tweak the `max_timesteps` value for experimenting with the visual quality of the generated outputs. The default `max_timesteps` value was obtained with an inference-time search process. For more details about it, check out the paper.

##
SanaSprintPipeline

###
class diffusers.SanaSprintPipeline

< source >

( tokenizer: typing.Union[transformers.models.gemma.tokenization_gemma.GemmaTokenizer, transformers.models.gemma.tokenization_gemma_fast.GemmaTokenizerFast] text_encoder: Gemma2PreTrainedModel vae: AutoencoderDC transformer: SanaTransformer2DModel scheduler: DPMSolverMultistepScheduler )

Pipeline for text-to-image generation using SANA-Sprint.

####
__call__

< source >

( prompt: typing.Union[str, typing.List[str]] = None num_inference_steps: int = 2 timesteps: typing.List[int] = None max_timesteps: float = 1.5708 intermediate_timesteps: float = 1.3 guidance_scale: float = 4.5 num_images_per_prompt: typing.Optional[int] = 1 height: int = 1024 width: int = 1024 eta: float = 0.0 generator: typing.Union[torch._C.Generator, typing.List[torch._C.Generator], NoneType] = None latents: typing.Optional[torch.Tensor] = None prompt_embeds: typing.Optional[torch.Tensor] = None prompt_attention_mask: typing.Optional[torch.Tensor] = None output_type: typing.Optional[str] = 'pil' return_dict: bool = True clean_caption: bool = False use_resolution_binning: bool = True attention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None callback_on_step_end: typing.Optional[typing.Callable[[int, int, typing.Dict], NoneType]] = None callback_on_step_end_tensor_inputs: typing.List[str] = ['latents'] max_sequence_length: int = 300 complex_human_instruction: typing.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`

Parameters

- prompt (`str` or `List[str]`, optional) — The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead.
- num_inference_steps (`int`, optional, defaults to 20) — The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.
- max_timesteps (`float`, optional, defaults to 1.57080) — The maximum timestep value used in the SCM scheduler.
- intermediate_timesteps (`float`, optional, defaults to 1.3) — The intermediate timestep value used in SCM scheduler (only used when num_inference_steps=2).
- timesteps (`List[int]`, optional) — Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order.
- guidance_scale (`float`, optional, defaults to 4.5) — Guidance scale as defined in Classifier-Free Diffusion Guidance. `guidance_scale` is defined as `w` of equation 2. of Imagen Paper. Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality.
- num_images_per_prompt (`int`, optional, defaults to 1) — The number of images to generate per prompt.
- height (`int`, optional, defaults to self.unet.config.sample_size) — The height in pixels of the generated image.
- width (`int`, optional, defaults to self.unet.config.sample_size) — The width in pixels of the generated image.
- eta (`float`, optional, defaults to 0.0) — Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to schedulers.DDIMScheduler, will be ignored for others.
- generator (`torch.Generator` or `List[torch.Generator]`, optional) — One or a list of torch generator(s) to make generation deterministic.
- latents (`torch.Tensor`, optional) — Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image generation. Can be used to tweak the same generation with different prompts. If not provided, a latents tensor will ge generated by sampling using the supplied random `generator`.
- prompt_embeds (`torch.Tensor`, optional) — Pre-generated text embeddings. Can be used to easily tweak text inputs, e.g. prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument.
- prompt_attention_mask (`torch.Tensor`, optional) — Pre-generated attention mask for text embeddings.
- output_type (`str`, optional, defaults to `"pil"`) — The output format of the generate image. Choose between PIL: `PIL.Image.Image` or `np.array`.
- return_dict (`bool`, optional, defaults to `True`) — Whether or not to return a `~pipelines.stable_diffusion.IFPipelineOutput` instead of a plain tuple.
- attention_kwargs — A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in diffusers.models.attention_processor.
- clean_caption (`bool`, optional, defaults to `True`) — Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to be installed. If the dependencies are not installed, the embeddings will be created from the raw prompt.
- use_resolution_binning (`bool` defaults to `True`) — If set to `True`, the requested height and width are first mapped to the closest resolutions using `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to the requested resolution. Useful for generating non-square images.
- callback_on_step_end (`Callable`, optional) — A function that calls at the end of each denoising steps during the inference. The function is called with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, optional) — The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class.
- max_sequence_length (`int` defaults to `300`) — Maximum sequence length to use with the `prompt`.
- complex_human_instruction (`List[str]`, optional) — Instructions for complex human attention: https://github.com/NVlabs/Sana/blob/main/configs/sana_app_config/Sana_1600M_app.yaml#L55.

Returns

SanaPipelineOutput or `tuple`

If `return_dict` is `True`, SanaPipelineOutput is returned, otherwise a `tuple` is returned where the first element is a list with the generated images

Function invoked when calling the pipeline for generation.

Examples:

    >>> 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")

####
disable_vae_slicing

< source >

( )

Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to computing decoding in one step.

####
disable_vae_tiling

< source >

( )

Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to computing decoding in one step.

####
enable_vae_slicing

< source >

( )

Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.

####
enable_vae_tiling

< source >

( )

Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow processing larger images.

####
encode_prompt

< source >

( prompt: typing.Union[str, typing.List[str]] num_images_per_prompt: int = 1 device: typing.Optional[torch.device] = None prompt_embeds: typing.Optional[torch.Tensor] = None prompt_attention_mask: typing.Optional[torch.Tensor] = None clean_caption: bool = False max_sequence_length: int = 300 complex_human_instruction: typing.Optional[typing.List[str]] = None lora_scale: typing.Optional[float] = None )

Parameters

- prompt (`str` or `List[str]`, optional) — prompt to be encoded
- num_images_per_prompt (`int`, optional, defaults to 1) — number of images that should be generated per prompt
- device — (`torch.device`, optional): torch device to place the resulting embeddings on
- prompt_embeds (`torch.Tensor`, optional) — Pre-generated text embeddings. Can be used to easily tweak text inputs, e.g. prompt weighting. If not provided, text embeddings will be generated from `prompt` input argument.
- clean_caption (`bool`, defaults to `False`) — If `True`, the function will preprocess and clean the provided caption before encoding.
- max_sequence_length (`int`, defaults to 300) — Maximum sequence length to use for the prompt.
- complex_human_instruction (`list[str]`, defaults to `complex_human_instruction`) — If `complex_human_instruction` is not empty, the function will use the complex Human instruction for the prompt.

Encodes the prompt into text encoder hidden states.

##
SanaPipelineOutput

###
class diffusers.pipelines.sana.pipeline_output.SanaPipelineOutput

< source >

( images: typing.Union[typing.List[PIL.Image.Image], numpy.ndarray] )

Parameters

- images (`List[PIL.Image.Image]` or `np.ndarray`) — List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width, num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.

Output class for Sana pipelines.

< > Update on GitHub