Added training for an experimental decoratgor embedding. Allow for turning off guidance embedding on flux (for unreleased model). Various bug fixes and modifications

toolkit/optimizers/adafactor.py

@@ -264,6 +264,10 @@ class Adafactor(torch.optim.Optimizer):
                 if grad.is_sparse:
                     raise RuntimeError(
                         "Adafactor does not support sparse gradients.")
+                
+                # if p has atts _scale then it is quantized. We need to divide the grad by the scale
+                # if hasattr(p, "_scale"):
+                #     grad = grad / p._scale
 
                 state = self.state[p]
                 grad_shape = grad.shape

toolkit/optimizers/automagic.py

@@ -101,7 +101,7 @@ class Automagic(torch.optim.Optimizer):
         if 'avg_lr' in param_state:
             lr = param_state["avg_lr"]
         else:
-            lr = param_state["lr"]
+            lr = 0.0
         return lr
 
     def _get_group_lr(self, group):
@@ -332,4 +332,4 @@ class Automagic(torch.optim.Optimizer):
                 state['lr_mask'] = Auto8bitTensor(sd_mask)
                 del state_dict['state'][idx]['lr_mask']
                 idx += 1
-        super().load_state_dict(state_dict, strict)
+        super().load_state_dict(state_dict)