Fix some small things

configs/models/cc12m_1024x1024.yaml

@@ -102,6 +102,21 @@ unet_config:
   temporal_mode: false
   temporal_positional_encoding: false
   temporal_spatial_ds: false
+diffusion_config:
+  sampler_config:
+    num_diffusion_steps: 1000
+    reproject_signal: False
+    schedule_type: DEEPFLOYD
+    prediction_type: V_PREDICTION
+    loss_target_type: DDPM
+    beta_start: 0.0001
+    beta_end: 0.02
+    threshold_function: CLIP
+    rescale_schedule: 1.0
+    schedule_shifted: True
+  model_output_scale: 0.0
+  use_vdm_loss_weights: False
+  no_use_residual: true
 
 # import defaults
 # reader-config-file: configs/datasets/reader_config.yaml

configs/models/cc12m_256x256.yaml

@@ -36,7 +36,7 @@ unet_config:
   conditioning_feature_dim: -1
   conditioning_feature_proj_dim: -1
   freeze_inner_unet: false
-  initialize_inner_with_pretrained: None
+  initialize_inner_with_pretrained: null
   inner_config:
     attention_levels: [1, 2]
     conditioning_feature_dim: -1
@@ -76,6 +76,21 @@ unet_config:
   temporal_mode: false
   temporal_positional_encoding: false
   temporal_spatial_ds: false
+diffusion_config:
+  sampler_config:
+    num_diffusion_steps: 1000
+    reproject_signal: False
+    schedule_type: DEEPFLOYD
+    prediction_type: V_PREDICTION
+    loss_target_type: DDPM
+    beta_start: 0.0001
+    beta_end: 0.02
+    threshold_function: CLIP
+    rescale_schedule: 1.0
+    schedule_shifted: True
+  model_output_scale: 0.0
+  use_vdm_loss_weights: False
+  no_use_residual: true
 
 reader_config:
   image_size: 256

ml_mdm/clis/generate_sample.py

@@ -138,7 +138,7 @@ def generate(
         prompt = input_template.format(prompt=prompt)
     if len(negative_template) > 0:
         negative_prompt = negative_prompt + negative_template
-    print(f"Postive: {prompt} / Negative: {negative_prompt}")
+    print(f"Positive: {prompt} / Negative: {negative_prompt}")
 
     if not os.path.exists(ckpt_name):
         logging.info(f"Did not generate because {ckpt_name} does not exist")
@@ -478,7 +478,7 @@ def main(args):
 
             with gr.Column(scale=2):
                 with gr.Accordion(
-                    "Addditional outputs", open=False, elem_id="output-accordion"
+                    "Additional outputs", open=False, elem_id="output-accordion"
                 ):
                     with gr.Row(equal_height=True):
                         output_text = gr.Textbox(value=None, label="System output")

ml_mdm/models/nested_unet.py

@@ -76,6 +76,9 @@ class Nested4UNetConfig(Nested3UNetConfig):
 
 
 def download(vision_model_path):
+    from ml_mdm.utils import fix_old_checkpoints
+    fix_old_checkpoints.mimic_old_modules()
+
     import os
 
     from distributed import get_local_rank

ml_mdm/samplers.py

@@ -597,7 +597,7 @@ def get_xt(self, x0, eps, g, scales):
             x_t += [
                 super().get_xt(
                     self.get_image_rescaled(x, s)
-                    if not self._config.schedule_shifted
+                    if self._config.schedule_shifted
                     else x,
                     e,
                     gi,
@@ -611,7 +611,7 @@ def get_prediction_targets(self, x0, eps, g, g_last, scales, prediction_type=Non
             tgt += [
                 super().get_prediction_targets(
                     self.get_image_rescaled(x, s)
-                    if not self._config.schedule_shifted
+                    if self._config.schedule_shifted
                     else x,
                     e,
                     gi,
@@ -668,7 +668,7 @@ def get_xt_minus_1(
                         need_noise=time_step != 1,
                         ddim_eta=ddim_eta,
                         clip_fn=self.clip_sample,
-                        image_scale=s if not self._config.schedule_shifted else 1,
+                        image_scale=s if self._config.schedule_shifted else 1,
                     )
                     for x, p, g, g_last, s in zip(x_t, p_t, g_t, g_s, scales)
                 ]
@@ -693,7 +693,7 @@ def _postprocess(
     ):
         scales = [
             x_t[i].size(-1) / x_t[-1].size(-1)
-            if not self._config.schedule_shifted
+            if self._config.schedule_shifted
             else 1
             for i in range(len(x_t))
         ]