🎉Glad to announce SANA-Sprint is available!

[lawrence-cj]

Our SANA-Sprint models focus on timestep distillation, achieving high-quality generation with 1-4 inference steps. Refer to SANA-Sprint Manual for more details.

After the PR merged: huggingface/diffusers#11074

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

# test sana sprint
from diffusers import SanaSprintPipeline
import torch

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

prompt = "a tiny astronaut hatching from an egg on the moon"

image = pipeline(prompt=prompt, num_inference_steps=2).images[0]
image.save("test_out.png")

[sr228822]

Excited to try this out but hitting an error

(sana) ~ $ python example.py 
/home/ec2-user/miniconda3/envs/sana/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:211: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.
  @torch.library.impl_abstract("xformers_flash::flash_fwd")
/home/ec2-user/miniconda3/envs/sana/lib/python3.10/site-packages/xformers/ops/fmha/flash.py:344: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.
  @torch.library.impl_abstract("xformers_flash::flash_bwd")
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00,  7.60it/s]
Loading pipeline components...: 100%|████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:13<00:00,  2.72s/it]
Set timesteps: tensor([1.5708, 1.3000, 0.0000], device='cuda:0')
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:01<00:00,  1.20it/s]
/home/ec2-user/miniconda3/envs/sana/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana_sprint.py:872: UserWarning: CUDA out of memory. Tried to allocate 4.50 GiB. GPU 0 has a total capacity of 14.56 GiB of which 2.69 GiB is free. Including non-PyTorch memory, this process has 11.87 GiB memory in use. Of the allocated memory 9.42 GiB is allocated by PyTorch, and 2.33 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables). 
Try to use VAE tiling for large images. For example: 
pipe.vae.enable_tiling(tile_sample_min_width=512, tile_sample_min_height=512)
  warnings.warn(
Traceback (most recent call last):
  File "/home/ec2-user/example.py", line 13, in <module>
    image = pipeline(prompt=prompt, num_inference_steps=2).images[0]
  File "/home/ec2-user/miniconda3/envs/sana/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/home/ec2-user/miniconda3/envs/sana/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana_sprint.py", line 878, in __call__
    image = self.image_processor.resize_and_crop_tensor(image, orig_width, orig_height)
UnboundLocalError: local variable 'image' referenced before assignment

Using the env from https://github.com/NVlabs/Sana/blob/main/environment_setup.sh

(sana) ~ $ pip freeze | grep -E 'diffusers|torch|xformers|transformers'
came-pytorch==0.1.3
diffusers @ git+https://github.com/huggingface/diffusers@a7d53a59394d5d8367826663601b69828e9f74fc
pytorch-fid==0.3.0
torch==2.4.0+cu121
torchaudio==2.4.0
torchvision==0.19.0
transformers==4.50.0
xformers==0.0.27.post2

[thewhatifproject]

Is it planned a support for controlnet?

[lawrence-cj]

@sr228822
Seems you got an error of "/home/ec2-user/miniconda3/envs/sana/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana_sprint.py:872: UserWarning: CUDA out of memory. Tried to allocate 4.50 GiB. GPU 0 has a total capacity of 14.56 GiB of which 2.69 GiB is free. Including non-PyTorch memory, this process has 11.87 GiB memory in use. Of the allocated memory 9.42 GiB is allocated by PyTorch, and 2.33 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables). "

And you also get a hint of "Try to use VAE tiling for large images. For example:
pipe.vae.enable_tiling(tile_sample_min_width=512, tile_sample_min_height=512)"

Please try to use vae tiling before sampling, eg:

# test sana sprint
from diffusers import SanaSprintPipeline
import torch

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

pipe.vae.enable_tiling(tile_sample_min_width=512, tile_sample_min_height=512)

prompt = "a tiny astronaut hatching from an egg on the moon"

image = pipeline(prompt=prompt, num_inference_steps=2).images[0]
image.save("test_out.png")

[lawrence-cj]

Is it planned a support for controlnet?

We will make it public when it's ready ASAP. No specific date for it. @thewhatifproject

[sr228822]

@sr228822 Seems you got an error of "/home/ec2-user/miniconda3/envs/sana/lib/python3.10/site-packages/diffusers/pipelines/sana/pipeline_sana_sprint.py:872: UserWarning: CUDA out of memory.

Ahh jeesh 🤦 I don't know how i missed that Cuda memory line .. the image variable stack trace further down got me looking in the wrong place. I just tried with a larger 48GB VRAM instance and its working fine now! I'll make note of the VAE enable_tiling tip in case we want to try and fit onto slimmer GPUs.

Thanks @lawrence-cj !

[atonderski]

Thanks for the release, very impressive speed! I notice in your graphs that in the 1-step setting the VAE consumes almost all the time(!), with 0.12s for vae and only 0.03s for the sana transformer.

What are your thoughts on this, would sana-sprint make sense with a lighter VAE, maybe even at lower compression rate? Or do you see some other way to overcome the VAE bottleneck?

[lawrence-cj]

Lower compression rate won't speed up the VAE part a lot. IMO, a lighter VAE would definitely help. @atonderski

[EladWarshawsky]

Not to be a bother, but any controlnet updates for sana-sprint?

[bertaveira]

Thank you for releasing the model. I am wondering if there are any plans to distill and release a Sana-sprint for 2k or 4k in the near future?