fixed issues with converting and saving models. Cleaned keys. Improved testing for cycle load saving.

toolkit/config_modules.py

@@ -67,6 +67,8 @@ class TrainConfig:
         self.noise_scheduler = kwargs.get('noise_scheduler', 'ddpm')
         self.steps: int = kwargs.get('steps', 1000)
         self.lr = kwargs.get('lr', 1e-6)
+        self.unet_lr = kwargs.get('unet_lr', self.lr)
+        self.text_encoder_lr = kwargs.get('text_encoder_lr', self.lr)
         self.optimizer = kwargs.get('optimizer', 'adamw')
         self.lr_scheduler = kwargs.get('lr_scheduler', 'constant')
         self.max_denoising_steps: int = kwargs.get('max_denoising_steps', 50)