feat(ml): diffusion in latent space

Author: beniz

examples/example_ddpm_noglasses2glasses.json

@@ -32,7 +32,7 @@
         "dropout": false,
         "nblocks": 2,
         "netE": "resnet_256",
-        "netG": "unet_mha",
+        "netG": "hdit",
         "ngf": 64,
         "norm": "instance",
         "padding_type": "reflect",
@@ -132,11 +132,11 @@
             "mask_square_B": false,
             "rand_mask_A": true
         },
-        "crop_size": 128,
+        "crop_size": 256,
         "dataset_mode": "self_supervised_labeled_mask",
         "direction": "BtoA",
         "inverted_mask": false,
-        "load_size": 128,
+        "load_size": 256,
         "max_dataset_size": 1000000000,
         "num_threads": 4,
         "online_context_pixels": 0,
@@ -167,8 +167,8 @@
             "aim_port": 53800,
             "aim_server": "http://localhost",
             "diff_fake_real": false,
-            "env": "noglasses2glasses",
-            "freq": 10000,
+            "env": "noglasses2glasses_hdit_latent",
+            "freq": 1000,
             "id": 1,
             "ncols": 0,
             "networks": false,
@@ -235,7 +235,7 @@
         "epoch_count": 1,
         "export_jit": false,
         "gan_mode": "lsgan",
-        "iter_size": 16,
+        "iter_size": 8,
         "load_iter": 0,
         "lr_decay_iters": 50,
         "lr_policy": "linear",
@@ -282,12 +282,12 @@
     "ddp_port": "12355",
     "gpu_ids": "0",
     "model_type": "palette",
-    "name": "noglasses2glasses",
+    "name": "noglasses2glasses_hdit_latent",
     "phase": "train",
     "suffix": "",
     "test_batch_size": 1,
     "warning_mode": false,
-    "with_amp": false,
-    "with_tf32": false,
+    "with_amp": true,
+    "with_tf32": true,
     "with_torch_compile": false
 }

models/base_diffusion_model.py

@@ -219,6 +219,19 @@ def modify_commandline_options_train(parser):
             help="the range of probabilities for dropping the canny for each frame",
         )
 
+        parser.add_argument(
+            "--alg_diffusion_latent_dc_ae_path",
+            type=str,
+            default="",
+            help="Path to the pretrained DC-AE model for latent space encoding. If empty, this feature is disabled.",
+        )
+        parser.add_argument(
+            "--alg_diffusion_latent_dc_ae_torch_dtype",
+            type=str,
+            default="float32",
+            help="Torch dtype for the DC-AE model.",
+        )
+
         return parser
 
     def __init__(self, opt, rank):