Add model checkpoint and config files

checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6a5262e8184d15d3876a7376e559e6a5643cad669864495687214397133875d3
+size 10801103890

checkpoint_.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6a5262e8184d15d3876a7376e559e6a5643cad669864495687214397133875d3
+size 10801103890

config.yaml

@@ -0,0 +1,106 @@
+run_dir: .exp/diffusion/imagenet_512/dc_ae_f32c32_in_1.0/dit_xl_1/bs_1024_lr_2e-4_fp16
+seed: 0
+allow_tf32: true
+resolution: 512
+amp: fp16
+cfg_scale: 1.0
+evaluate_split: test
+evaluate_dir_name: null
+num_save_images: 64
+save_all_images: false
+save_images_at_all_procs: false
+evaluate_dataset: sample_class
+sample_class:
+  name: sample_class
+  batch_size: 32
+  n_worker: 8
+  train_drop_last: true
+  seed: 0
+  num_classes: 1000
+  num_samples: 50000
+autoencoder: dc-ae-f32c32-in-1.0
+autoencoder_dtype: fp32
+scaling_factor: 0.3189
+model: dit
+dit:
+  name: DiT
+  input_size: 16
+  patch_size: 1
+  in_channels: 32
+  hidden_size: 1152
+  depth: 28
+  num_heads: 16
+  mlp_ratio: 4.0
+  post_norm: false
+  class_dropout_prob: 0.1
+  num_classes: 1000
+  learn_sigma: true
+  unconditional: false
+  use_checkpoint: true
+  pretrained_path: null
+  pretrained_source: dc-ae
+  eval_scheduler: GaussianDiffusion
+  num_inference_steps: 250
+  train_scheduler: GaussianDiffusion
+uvit:
+  name: UViT
+  input_size: 32
+  patch_size: 2
+  in_channels: 4
+  hidden_size: 1152
+  depth: 28
+  num_heads: 16
+  mlp_ratio: 4.0
+  mlp_time_embed: false
+  qkv_bias: false
+  act_layer: gelu
+  use_checkpoint: true
+  class_dropout_prob: 0.1
+  num_classes: 1000
+  pretrained_path: null
+  pretrained_source: dc-ae
+  eval_scheduler: DPM_Solver
+  num_inference_steps: 30
+  train_scheduler: DPM_Solver
+  attn_mode: null
+compute_fid: true
+fid:
+  save_path: null
+  ref_path: assets/data/fid/imagenet_512_train.npz
+compute_inception_score: true
+inception_score: {}
+train_dataset: latent_imagenet
+latent_imagenet:
+  name: latent_imagenet
+  batch_size: 32
+  n_worker: 8
+  train_drop_last: true
+  seed: 0
+  data_dir: assets/data/latent/dc_ae_f32c32_in_1.0/imagenet_512
+resume: true
+resume_path: null
+resume_schedule: true
+num_epochs: null
+max_steps: 3000000
+clip_grad: null
+num_store_images: 64
+save_checkpoint_steps: 1000
+evaluate_steps: 20000
+optimizer:
+  name: adamw
+  lr: 0.0002
+  warmup_lr: 0.0
+  weight_decay: 0.0
+  no_wd_keys: []
+  betas:
+  - 0.9
+  - 0.999
+lr_scheduler:
+  name: constant
+  warmup_steps: 1000
+log: true
+wandb_entity: null
+wandb_project: null
+ema_decay: 0.9999
+ema_warmup_steps: 2000
+evaluate_ema: true

model.txt

@@ -0,0 +1,47 @@
+DiT(
+  (x_embedder): PatchEmbed(
+    (proj): Conv2d(32, 1152, kernel_size=(1, 1), stride=(1, 1))
+    (norm): Identity()
+  )
+  (t_embedder): TimestepEmbedder(
+    (mlp): Sequential(
+      (0): Linear(in_features=256, out_features=1152, bias=True)
+      (1): SiLU()
+      (2): Linear(in_features=1152, out_features=1152, bias=True)
+    )
+  )
+  (y_embedder): LabelEmbedder(
+    (embedding_table): Embedding(1001, 1152)
+  )
+  (blocks): ModuleList(
+    (0-27): 28 x DiTBlock(
+      (norm1): LayerNorm((1152,), eps=1e-06, elementwise_affine=False)
+      (attn): Attention(
+        (qkv): Linear(in_features=1152, out_features=3456, bias=True)
+        (attn_drop): Dropout(p=0.0, inplace=False)
+        (proj): Linear(in_features=1152, out_features=1152, bias=True)
+        (proj_drop): Dropout(p=0.0, inplace=False)
+      )
+      (norm2): LayerNorm((1152,), eps=1e-06, elementwise_affine=False)
+      (mlp): Mlp(
+        (fc1): Linear(in_features=1152, out_features=4608, bias=True)
+        (act): GELU(approximate='tanh')
+        (drop1): Dropout(p=0, inplace=False)
+        (fc2): Linear(in_features=4608, out_features=1152, bias=True)
+        (drop2): Dropout(p=0, inplace=False)
+      )
+      (adaLN_modulation): Sequential(
+        (0): SiLU()
+        (1): Linear(in_features=1152, out_features=6912, bias=True)
+      )
+    )
+  )
+  (final_layer): FinalLayer(
+    (norm_final): LayerNorm((1152,), eps=1e-06, elementwise_affine=False)
+    (linear): Linear(in_features=1152, out_features=64, bias=True)
+    (adaLN_modulation): Sequential(
+      (0): SiLU()
+      (1): Linear(in_features=1152, out_features=2304, bias=True)
+    )
+  )
+)