1. add 4096 aspect ratio;

2. pre-commit

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

app/sana_pipeline.py

@@ -26,7 +26,12 @@
 
 
 from diffusion import DPMS, FlowEuler
-from diffusion.data.datasets.utils import ASPECT_RATIO_512_TEST, ASPECT_RATIO_1024_TEST, ASPECT_RATIO_2048_TEST
+from diffusion.data.datasets.utils import (
+    ASPECT_RATIO_512_TEST,
+    ASPECT_RATIO_1024_TEST,
+    ASPECT_RATIO_2048_TEST,
+    ASPECT_RATIO_4096_TEST,
+)
 from diffusion.model.builder import build_model, get_tokenizer_and_text_encoder, get_vae, vae_decode
 from diffusion.model.utils import get_weight_dtype, prepare_prompt_ar, resize_and_crop_tensor
 from diffusion.utils.config import SanaConfig, model_init_config

asset/docs/model_zoo.md

@@ -78,6 +78,7 @@ image[0].save('sana.png')
 ```
 
 #### 2). For 4K models
+
 4K models need [patch_conv](https://github.com/mit-han-lab/patch_conv) to avoid OOM issue.(80GB GPU is recommended)
 
 run `pip install patch_conv` first, then
@@ -113,4 +114,4 @@ image = pipe(
 )[0]
 
 image[0].save("sana_4K.png")
-```
+```

scripts/inference.py

@@ -35,7 +35,12 @@
 warnings.filterwarnings("ignore")  # ignore warning
 
 from diffusion import DPMS, FlowEuler, SASolverSampler
-from diffusion.data.datasets.utils import ASPECT_RATIO_512_TEST, ASPECT_RATIO_1024_TEST, ASPECT_RATIO_2048_TEST
+from diffusion.data.datasets.utils import (
+    ASPECT_RATIO_512_TEST,
+    ASPECT_RATIO_1024_TEST,
+    ASPECT_RATIO_2048_TEST,
+    ASPECT_RATIO_4096_TEST,
+)
 from diffusion.model.builder import build_model, get_tokenizer_and_text_encoder, get_vae, vae_decode
 from diffusion.model.utils import get_weight_dtype, prepare_prompt_ar
 from diffusion.utils.config import SanaConfig, model_init_config

scripts/inference_dpg.py

@@ -38,6 +38,7 @@
     ASPECT_RATIO_512_TEST,
     ASPECT_RATIO_1024_TEST,
     ASPECT_RATIO_2048_TEST,
+    ASPECT_RATIO_4096_TEST,
     get_chunks,
 )
 from diffusion.model.builder import build_model, get_tokenizer_and_text_encoder, get_vae, vae_decode

scripts/inference_geneval.py

@@ -36,7 +36,12 @@
 warnings.filterwarnings("ignore")  # ignore warning
 
 from diffusion import DPMS, FlowEuler, SASolverSampler
-from diffusion.data.datasets.utils import ASPECT_RATIO_512_TEST, ASPECT_RATIO_1024_TEST, ASPECT_RATIO_2048_TEST
+from diffusion.data.datasets.utils import (
+    ASPECT_RATIO_512_TEST,
+    ASPECT_RATIO_1024_TEST,
+    ASPECT_RATIO_2048_TEST,
+    ASPECT_RATIO_4096_TEST,
+)
 from diffusion.model.builder import build_model, get_tokenizer_and_text_encoder, get_vae, vae_decode
 from diffusion.model.utils import get_weight_dtype, prepare_prompt_ar
 from diffusion.utils.config import SanaConfig, model_init_config

scripts/inference_image_reward.py

@@ -33,7 +33,12 @@
 from tqdm import tqdm
 
 from diffusion import DPMS, FlowEuler, SASolverSampler
-from diffusion.data.datasets.utils import ASPECT_RATIO_512_TEST, ASPECT_RATIO_1024_TEST, ASPECT_RATIO_2048_TEST
+from diffusion.data.datasets.utils import (
+    ASPECT_RATIO_512_TEST,
+    ASPECT_RATIO_1024_TEST,
+    ASPECT_RATIO_2048_TEST,
+    ASPECT_RATIO_4096_TEST,
+)
 from diffusion.model.builder import build_model, get_tokenizer_and_text_encoder, get_vae, vae_decode
 from diffusion.model.utils import get_weight_dtype, prepare_prompt_ar
 from diffusion.utils.config import SanaConfig, model_init_config

scripts/interface.py

@@ -34,7 +34,12 @@
 
 from asset.examples import examples
 from diffusion import DPMS, FlowEuler, SASolverSampler
-from diffusion.data.datasets.utils import ASPECT_RATIO_512_TEST, ASPECT_RATIO_1024_TEST, ASPECT_RATIO_2048_TEST
+from diffusion.data.datasets.utils import (
+    ASPECT_RATIO_512_TEST,
+    ASPECT_RATIO_1024_TEST,
+    ASPECT_RATIO_2048_TEST,
+    ASPECT_RATIO_4096_TEST,
+)
 from diffusion.model.builder import build_model, get_tokenizer_and_text_encoder, get_vae, vae_decode
 from diffusion.model.utils import get_weight_dtype, prepare_prompt_ar, resize_and_crop_tensor
 from diffusion.utils.config import SanaConfig, model_init_config