diff --git a/app/sana_pipeline.py b/app/sana_pipeline.py
index 051b2f6..c710c9a 100644
--- a/app/sana_pipeline.py
+++ b/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
diff --git a/asset/docs/model_zoo.md b/asset/docs/model_zoo.md
index 1405919..55d76f2 100644
--- a/asset/docs/model_zoo.md
+++ b/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")
-```
\ No newline at end of file
+```
diff --git a/scripts/inference.py b/scripts/inference.py
index a00167f..fb55953 100755
--- a/scripts/inference.py
+++ b/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
diff --git a/scripts/inference_dpg.py b/scripts/inference_dpg.py
index 8e12c77..cae8b72 100644
--- a/scripts/inference_dpg.py
+++ b/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
diff --git a/scripts/inference_geneval.py b/scripts/inference_geneval.py
index 1eee141..9c8b9a4 100644
--- a/scripts/inference_geneval.py
+++ b/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
diff --git a/scripts/inference_image_reward.py b/scripts/inference_image_reward.py
index 40b2fb7..a1b3d1f 100644
--- a/scripts/inference_image_reward.py
+++ b/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
diff --git a/scripts/interface.py b/scripts/interface.py
index 35ef47c..a8d9215 100755
--- a/scripts/interface.py
+++ b/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