Minor changes; add train_timestep_fraction

config/base.py

@@ -32,14 +32,15 @@ def get_config():
     train.cfg = True
     train.adv_clip_max = 10
     train.clip_range = 1e-4
+    train.timestep_fraction = 1.0
 
     # sampling
     config.sample = sample = ml_collections.ConfigDict()
-    sample.num_steps = 5
+    sample.num_steps = 10
     sample.eta = 1.0
     sample.guidance_scale = 5.0
     sample.batch_size = 1
-    sample.num_batches_per_epoch = 1
+    sample.num_batches_per_epoch = 2
 
     # prompting
     config.prompt_fn = "imagenet_animals"
@@ -49,7 +50,7 @@ def get_config():
     config.reward_fn = "jpeg_compressibility"
 
     config.per_prompt_stat_tracking = ml_collections.ConfigDict()
-    config.per_prompt_stat_tracking.buffer_size = 64
+    config.per_prompt_stat_tracking.buffer_size = 16
     config.per_prompt_stat_tracking.min_count = 16
 
     return config

config/dgx.py

@@ -8,20 +8,25 @@ base = imp.load_source("base", os.path.join(os.path.dirname(__file__), "base.py"
 def get_config():
     config = base.get_config()
 
-    config.mixed_precision = "no"
+    config.pretrained.model = "runwayml/stable-diffusion-v1-5"
+
+    config.mixed_precision = "fp16"
     config.allow_tf32 = True
     config.use_lora = False
 
     config.train.batch_size = 4
-    config.train.gradient_accumulation_steps = 8
-    config.train.learning_rate = 1e-5
-    config.train.clip_range = 1.0
+    config.train.gradient_accumulation_steps = 2
+    config.train.learning_rate = 3e-5
+    config.train.clip_range = 1e-4
 
     # sampling
     config.sample.num_steps = 50
-    config.sample.batch_size = 16
-    config.sample.num_batches_per_epoch = 2
+    config.sample.batch_size = 8
+    config.sample.num_batches_per_epoch = 4
 
-    config.per_prompt_stat_tracking = None
+    config.per_prompt_stat_tracking = {
+        "buffer_size": 16,
+        "min_count": 16,
+    }
 
     return config