Adding svdquant required model type conversion script and fix bugs (#222)

* fix little typo in sana-sprint app.

Signed-off-by: lawrence-cj <[email protected]>

* update README.md

Signed-off-by: lawrence-cj <[email protected]>

* add convert sana to svdquant quantization required model type scripts;

Signed-off-by: lawrence-cj <[email protected]>

* fix model path bugs;

Signed-off-by: lawrence-cj <[email protected]>

* remove unused args;

Signed-off-by: lawrence-cj <[email protected]>

* FSDP does not support resume from ddp checkpoint;

Signed-off-by: lawrence-cj <[email protected]>

* pre-commit;

Signed-off-by: lawrence-cj <[email protected]>

* pre-commit;

Signed-off-by: lawrence-cj <[email protected]>

* update workflow;

Signed-off-by: lawrence-cj <[email protected]>

* fix timestep dtype;

Signed-off-by: lawrence-cj <[email protected]>

---------

Signed-off-by: lawrence-cj <[email protected]>

Author: lawrence-cj

.github/workflows/bot-autolint.yaml

@@ -11,7 +11,7 @@ on:
       - unlabeled
 # run only one unit test for a branch / tag.
 concurrency:
-  group: ci-lint-${{ github.ref }}
+  group: ci-lint-${{ github.head_ref || github.ref }}
   cancel-in-progress: true
 jobs:
   lint-by-label:
@@ -21,7 +21,7 @@ jobs:
       - name: Check out Git repository
         uses: actions/checkout@v4
         with:
-          token: ${{ secrets.PAT }}
+          token: ${{ secrets.GITHUB_TOKEN }}
           ref: ${{ github.event.pull_request.head.ref }}
       - name: Set up Python
         uses: actions/setup-python@v5

README.md

@@ -10,7 +10,8 @@
   <a href="https://nvlabs.github.io/Sana/"><img src="https://img.shields.io/static/v1?label=Project&message=Github&color=blue&logo=github-pages"></a> &ensp;
   <a href="https://hanlab.mit.edu/projects/sana/"><img src="https://img.shields.io/static/v1?label=Page&message=MIT&color=darkred&logo=github-pages"></a> &ensp;
   <a href="https://arxiv.org/abs/2410.10629"><img src="https://img.shields.io/static/v1?label=Arxiv&message=Sana&color=red&logo=arxiv"></a> &ensp;
-  <a href="https://nv-sana.mit.edu/"><img src="https://img.shields.io/static/v1?label=Demo:6x3090&message=MIT&color=yellow"></a> &ensp;
+  <a href="https://nv-sana.mit.edu/"><img src="https://img.shields.io/static/v1?label=Demo:5x3090&message=SANA&color=yellow"></a> &ensp;
+  <a href="https://nv-sana.mit.edu/sprint/"><img src="https://img.shields.io/static/v1?label=Demo:1x3090&message=SANA-Sprint&color=yellow"></a> &ensp;
   <a href="https://nv-sana.mit.edu/4bit/"><img src="https://img.shields.io/static/v1?label=Demo:1x3090&message=4bit&color=yellow"></a> &ensp;
   <a href="https://nv-sana.mit.edu/ctrlnet/"><img src="https://img.shields.io/static/v1?label=Demo:1x3090&message=ControlNet&color=yellow"></a> &ensp;
   <a href="https://replicate.com/chenxwh/sana"><img src="https://img.shields.io/static/v1?label=API:H100&message=Replicate&color=pink"></a> &ensp;
@@ -25,7 +26,7 @@
 
 ### 🚶 Basic:
 
-**Demo**: [SANA-1.5](https://nv-sana.mit.edu/) | [SANA-ControlNet](https://nv-sana.mit.edu/ctrlnet/) | [SANA-4bit](https://nv-sana.mit.edu/4bit/) | [SANA-Sprint (Coming)](<>) <br>
+**Demo**: [SANA-1.5](https://nv-sana.mit.edu/) | [SANA-ControlNet](https://nv-sana.mit.edu/ctrlnet/) | [SANA-4bit](https://nv-sana.mit.edu/4bit/) | [SANA-Sprint](https://nv-sana.mit.edu/sprint/) <br>
 **ComfyUI**: [ComfyUI Guidance](asset/docs/ComfyUI/comfyui.md) <br>
 **Model Zoo:** [Model Card Collects All Models](asset/docs/model_zoo.md) <br>
 **Env Preparation:** [One-Click Env Install](#-1-dependencies-and-installation) <br>

app/app_sana_sprint.py

@@ -159,17 +159,10 @@ def get_args():
     parser.add_argument(
         "--model_path",
         nargs="?",
-        default="hf://Efficient-Large-Model/SANA_Sprint_1.6B_1024px/checkpoints/SANA_Sprint_1.6B_1024px.pth",
+        default="hf://Efficient-Large-Model/Sana_Sprint_1.6B_1024px/checkpoints/Sana_Sprint_1.6B_1024px.pth",
         type=str,
         help="Path to the model file (positional)",
     )
-    parser.add_argument("--output", default="./", type=str)
-    parser.add_argument("--bs", default=1, type=int)
-    parser.add_argument("--image_size", default=1024, type=int)
-    parser.add_argument("--cfg_scale", default=3.0, type=float)
-    parser.add_argument("--seed", default=42, type=int)
-    parser.add_argument("--step", default=-1, type=int)
-    parser.add_argument("--custom_image_size", default=None, type=int)
     parser.add_argument("--share", action="store_true")
     parser.add_argument(
         "--shield_model_path",
@@ -184,7 +177,6 @@ def get_args():
 args = get_args()
 
 if torch.cuda.is_available():
-    weight_dtype = torch.float16
     model_path = args.model_path
     pipe = SanaSprintPipeline(args.config)
     pipe.from_pretrained(model_path)
@@ -300,7 +292,7 @@ def generate(
     )
 
 
-model_size = "1.6" if "1600M" in args.model_path else "0.6"
+model_size = "1.6" if "1.6B" in args.model_path else "0.6"
 title = f"""
     <div style='display: flex; align-items: center; justify-content: center; text-align: center;'>
         <img src="https://nvlabs.github.io/Sana/Sprint/asset/SANA-Sprint.png" width="50%" alt="logo"/>
@@ -309,7 +301,7 @@ def generate(
 DESCRIPTION = f"""
         <p><span style="font-size: 36px; font-weight: bold;">SANA-Sprint-{model_size}B</span><span style="font-size: 20px; font-weight: bold;">{args.image_size}px</span></p>
         <p style="font-size: 16px; font-weight: bold;">SANA-Sprint: One-Step Diffusion with Continuous-Time Consistency Distillation</p>
-        <p><span style="font-size: 16px;"><a href="https://arxiv.org/abs/2503.09641">[Paper]</a></span> <span style="font-size: 16px;"><a href="https://github.com/NVlabs/Sana">[Github(coming soon)]</a></span> <span style="font-size: 16px;"><a href="https://nvlabs.github.io/Sana">[Project]</a></span</p>
+        <p><span style="font-size: 16px;"><a href="https://arxiv.org/abs/2503.09641">[Paper]</a></span> <span style="font-size: 16px;"><a href="https://github.com/NVlabs/Sana">[Github]</a></span> <span style="font-size: 16px;"><a href="https://nvlabs.github.io/Sana">[Project]</a></span</p>
         <p style="font-size: 16px; font-weight: bold;">Powered by <a href="https://hanlab.mit.edu/projects/dc-ae">DC-AE</a> with 32x latent space, </p>running on node {socket.gethostname()}.
         <p style="font-size: 16px; font-weight: bold;">Unsafe word will give you a 'Red Heart' in the image instead.</p>
         """

asset/docs/4bit_sana.md

@@ -19,6 +19,22 @@
 
 Follow the official [SVDQuant-Nunchaku](https://github.com/mit-han-lab/nunchaku) repository to set up the environment. The guidance can be found [here](https://github.com/mit-han-lab/nunchaku?tab=readme-ov-file#installation).
 
+### 1-1. Quantize Sana with SVDQuant-4bit (Optional)
+
+1. Convert pth to SVDQuant required safetensor
+
+```
+python tools/convert_sana_to_svdquant.py \
+      --orig_ckpt_path Efficient-Large-Model/SANA1.5_1.6B_1024px/checkpoints/SANA1.5_1.6B_1024px.pth \
+      --model_type SanaMS1.5_1600M_P1_D20 \
+      --dtype bf16 \
+      --dump_path output/SANA1.5_1.6B_1024px_svdquant_diffusers \
+      --save_full_pipeline
+```
+
+2. follow the guidance to compress model
+   [Quantization guidance](https://github.com/mit-han-lab/deepcompressor/tree/main/examples/diffusion)
+
 ### 2. Code snap for inference
 
 Here we show the code snippet for SanaPipeline. For SanaPAGPipeline, please refer to the [SanaPAGPipeline](https://github.com/mit-han-lab/nunchaku/blob/main/examples/sana_1600m_pag.py) section.

configs/sana1-5_config/1024ms/Sana_1600M_1024px_AdamW_fsdp.yaml

@@ -19,7 +19,8 @@ model:
   image_size: 1024
   mixed_precision: bf16
   fp32_attention: true
-  load_from: hf://Efficient-Large-Model/SANA1.5_1.6B_1024px/checkpoints/SANA1.5_1.6B_1024px.pth
+  #  load_from: hf://Efficient-Large-Model/SANA1.5_1.6B_1024px/checkpoints/SANA1.5_1.6B_1024px.pth
+  load_from:
   aspect_ratio_type: ASPECT_RATIO_1024
   multi_scale: true
   attn_type: linear

configs/sana_sprint_config/1024ms/SanaSprint_1600M_1024px_allqknorm_bf16_scm_ladd.yaml

@@ -20,7 +20,7 @@ model:
   mixed_precision: bf16
   fp32_attention: true
   teacher_model: hf://Efficient-Large-Model/Sana_Sprint_1.6B_1024px_teacher/checkpoints/Sana_Sprint_1.6B_1024px_teacher.pth
-  load_from: hf://Efficient-Large-Model/SANA_Sprint_1.6B_1024px/checkpoints/Sana_Sprint_1.6B_1024px.pth
+  load_from: hf://Efficient-Large-Model/Sana_Sprint_1.6B_1024px/checkpoints/Sana_Sprint_1.6B_1024px.pth
   resume_from:
   aspect_ratio_type: ASPECT_RATIO_1024
   multi_scale: true

diffusion/model/nets/sana_multi_scale.py

@@ -22,7 +22,7 @@
 from timm.models.layers import DropPath
 
 from diffusion.model.builder import MODELS
-from diffusion.model.nets.basic_modules import DWMlp, GLUMBConv, MBConvPreGLU, Mlp
+from diffusion.model.nets.basic_modules import DWMlp, GLUMBConv, Mlp
 from diffusion.model.nets.sana import Sana, get_2d_sincos_pos_embed
 from diffusion.model.nets.sana_blocks import (
     Attention,
@@ -279,9 +279,9 @@ def forward(self, x, timestep, y, mask=None, data_info=None, return_logvar=False
         bs = x.shape[0]
         x = x.to(self.dtype)
         if self.timestep_norm_scale_factor != 1.0:
-            timestep = (timestep.float() / self.timestep_norm_scale_factor).to(self.dtype)
+            timestep = (timestep.float() / self.timestep_norm_scale_factor).to(torch.float32)
         else:
-            timestep = timestep.long().to(self.dtype)
+            timestep = timestep.long().to(torch.float32)
         y = y.to(self.dtype)
         self.h, self.w = x.shape[-2] // self.patch_size, x.shape[-1] // self.patch_size
         x = self.x_embedder(x)
@@ -322,6 +322,7 @@ def forward(self, x, timestep, y, mask=None, data_info=None, return_logvar=False
             y = self.attention_y_norm(y)
 
         if mask is not None:
+            mask = mask.to(torch.int16)
             mask = mask.repeat(y.shape[0] // mask.shape[0], 1) if mask.shape[0] != y.shape[0] else mask
             mask = mask.squeeze(1).squeeze(1)
             if _xformers_available:
@@ -478,34 +479,34 @@ def SanaMS_1600M_P2_D20(**kwargs):
     return SanaMS(depth=20, hidden_size=2240, patch_size=2, num_heads=20, **kwargs)
 
 
-# TrigFlow/sCM model
 @MODELS.register_module()
-def SanaMSCM_600M_P1_D28(**kwargs):
-    return SanaMSCM(depth=28, hidden_size=1152, patch_size=1, num_heads=16, **kwargs)
+def SanaMS_2400M_P1_D30(**kwargs):
+    return SanaMS(depth=30, hidden_size=2240, patch_size=1, num_heads=20, **kwargs)
 
 
 @MODELS.register_module()
-def SanaMSCM_1600M_P1_D20(**kwargs):
-    return SanaMSCM(depth=20, hidden_size=2240, patch_size=1, num_heads=20, **kwargs)
+def SanaMS_3200M_P1_D40(**kwargs):
+    return SanaMS(depth=40, hidden_size=2240, patch_size=1, num_heads=20, **kwargs)
 
 
 @MODELS.register_module()
-def SanaMSCM_2400M_P1_D30(**kwargs):
-    # 30 layers, 2400M
-    return SanaMSCM(depth=30, hidden_size=2240, patch_size=1, num_heads=20, **kwargs)
+def SanaMS_4800M_P1_D60(**kwargs):
+    # 60 layers, 4800M
+    return SanaMS(depth=60, hidden_size=2240, patch_size=1, num_heads=20, **kwargs)
 
 
+# TrigFlow/sCM model
 @MODELS.register_module()
-def SanaMS_2400M_P1_D30(**kwargs):
-    return SanaMS(depth=30, hidden_size=2240, patch_size=1, num_heads=20, **kwargs)
+def SanaMSCM_600M_P1_D28(**kwargs):
+    return SanaMSCM(depth=28, hidden_size=1152, patch_size=1, num_heads=16, **kwargs)
 
 
 @MODELS.register_module()
-def SanaMS_3200M_P1_D40(**kwargs):
-    return SanaMS(depth=40, hidden_size=2240, patch_size=1, num_heads=20, **kwargs)
+def SanaMSCM_1600M_P1_D20(**kwargs):
+    return SanaMSCM(depth=20, hidden_size=2240, patch_size=1, num_heads=20, **kwargs)
 
 
 @MODELS.register_module()
-def SanaMS_4800M_P1_D60(**kwargs):
-    # 60 layers, 4800M
-    return SanaMS(depth=60, hidden_size=2240, patch_size=1, num_heads=20, **kwargs)
+def SanaMSCM_2400M_P1_D30(**kwargs):
+    # 30 layers, 2400M
+    return SanaMSCM(depth=30, hidden_size=2240, patch_size=1, num_heads=20, **kwargs)

scripts/inference_sana_sprint.py

@@ -212,7 +212,7 @@ class SanaInference(SanaConfig):
     ] = "configs/sana_sprint_config/1024ms/SanaSprint_1600M_1024px_allqknorm_bf16_scm_ladd.yaml"  # config
     model_path: Optional[
         str
-    ] = "hf://Efficient-Large-Model/SANA_Sprint_1.6B_1024px/checkpoints/SANA_Sprint_1.6B_1024px.pth"
+    ] = "hf://Efficient-Large-Model/Sana_Sprint_1.6B_1024px/checkpoints/Sana_Sprint_1.6B_1024px.pth"
     work_dir: Optional[str] = None
     txt_file: str = "asset/samples/samples_mini.txt"
     json_file: Optional[str] = None