diff --git a/configs/repa_flatten_condit22_fixt_xl.yaml b/configs/repa_improved_ddt_xlen22de6_256.yaml
similarity index 96%
rename from configs/repa_flatten_condit22_fixt_xl.yaml
rename to configs/repa_improved_ddt_xlen22de6_256.yaml
index 6f3a6ce..3e702e2 100644
--- a/configs/repa_flatten_condit22_fixt_xl.yaml
+++ b/configs/repa_improved_ddt_xlen22de6_256.yaml
@@ -42,14 +42,14 @@ model:
       precompute: true
       weight_path: /mnt/bn/wangshuai6/models/sd-vae-ft-ema/
   denoiser:
-    class_path: src.models.denoiser.flatten_condit_dit_fixt.FlattenConDiT
+    class_path: src.models.denoiser.decoupled_improved_dit.DDT
     init_args:
       in_channels: 4
       patch_size: 2
       num_groups: 16
       hidden_size: &hidden_dim 1152
       num_blocks: 28
-      num_cond_blocks: 22
+      num_encoder_blocks: 22
       num_classes: 1000
   conditioner:
     class_path: src.models.conditioner.LabelConditioner
diff --git a/configs/repa_flatten_condit22_fixt_xl512.yaml b/configs/repa_improved_ddt_xlen22de6_512.yaml
similarity index 96%
rename from configs/repa_flatten_condit22_fixt_xl512.yaml
rename to configs/repa_improved_ddt_xlen22de6_512.yaml
index ce59338..5c6b665 100644
--- a/configs/repa_flatten_condit22_fixt_xl512.yaml
+++ b/configs/repa_improved_ddt_xlen22de6_512.yaml
@@ -42,14 +42,14 @@ model:
       precompute: true
       weight_path: /mnt/bn/wangshuai6/models/sd-vae-ft-ema/
   denoiser:
-    class_path: src.models.denoiser.flatten_condit_dit_fixt.FlattenConDiT
+    class_path: src.models.denoiser.decoupled_improved_dit.DDT
     init_args:
       in_channels: 4
       patch_size: 2
       num_groups: 16
       hidden_size: &hidden_dim 1152
       num_blocks: 28
-      num_cond_blocks: 22
+      num_encoder_blocks: 22
       num_classes: 1000
   conditioner:
     class_path: src.models.conditioner.LabelConditioner
diff --git a/configs/repa_flatten_dit_fixt_large.yaml b/configs/repa_improved_dit_large.yaml
similarity index 96%
rename from configs/repa_flatten_dit_fixt_large.yaml
rename to configs/repa_improved_dit_large.yaml
index 7dc9611..a6a51b9 100644
--- a/configs/repa_flatten_dit_fixt_large.yaml
+++ b/configs/repa_improved_dit_large.yaml
@@ -1,7 +1,7 @@
 # lightning.pytorch==2.4.0
 seed_everything: true
 tags:
-  exp: &exp repa_flatten_dit_fixt_large
+  exp: &exp repa_improved_dit_large
 torch_hub_dir: /mnt/bn/wangshuai6/torch_hub
 huggingface_cache_dir: null
 trainer:
@@ -42,7 +42,7 @@ model:
       precompute: true
       weight_path: /mnt/bn/wangshuai6/models/sd-vae-ft-ema/
   denoiser:
-    class_path: src.models.denoiser.flatten_dit_fixt.FlattenDiT
+    class_path: src.models.denoiser.improved_dit.DiT
     init_args:
       in_channels: 4
       patch_size: 2
diff --git a/configs/repa_flatten_dit_fixt_xl.yaml b/configs/repa_improved_dit_xl.yaml
similarity index 96%
rename from configs/repa_flatten_dit_fixt_xl.yaml
rename to configs/repa_improved_dit_xl.yaml
index 2bc8606..593c048 100644
--- a/configs/repa_flatten_dit_fixt_xl.yaml
+++ b/configs/repa_improved_dit_xl.yaml
@@ -1,7 +1,7 @@
 # lightning.pytorch==2.4.0
 seed_everything: true
 tags:
-  exp: &exp repa_flatten_dit_fixt_xl
+  exp: &exp repa_improved_dit_xlen22de6_512
 torch_hub_dir: /mnt/bn/wangshuai6/torch_hub
 huggingface_cache_dir: null
 trainer:
@@ -42,7 +42,7 @@ model:
       precompute: true
       weight_path: /mnt/bn/wangshuai6/models/sd-vae-ft-ema/
   denoiser:
-    class_path: src.models.denoiser.flatten_dit_fixt.FlattenDiT
+    class_path: src.models.denoiser.improved_dit.DiT
     init_args:
       in_channels: 4
       patch_size: 2