Minor changes; add train_timestep_fraction

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