Add support for fine-tuning Hidream O1 (#831)

* Initial support for hidream. Lora keys likely need work

* Fix saving for hidream-o1

* Remove dependence on flash attention for hidream o1

* Fix gradient checkpointing for hidream o1

* A lot of fixes for hidream. Handle loading and saving as comfy model.

* Omit layers not used in comfy. Fix issue with lora loading keys in comfy

* Version bumpo

toolkit/config_modules.py

@@ -577,6 +577,11 @@ class TrainConfig:
         self.do_blank_stabilization = kwargs.get('do_blank_stabilization', False)
         
         self.audio_loss_multiplier = kwargs.get("audio_loss_multiplier", 1.0)
+        
+        # will throw detailed error when it goes over
+        self.max_loss_debug: bool = kwargs.get("max_loss_debug", False)
+        # will clip the loss to this amount to prevent wild outliers
+        self.max_loss: Optional[float] = kwargs.get("max_loss", None)
 
 
 ModelArch = Literal['sd1', 'sd2', 'sd3', 'sdxl', 'pixart', 'pixart_sigma', 'auraflow', 'flux', 'flex1', 'flex2', 'lumina2', 'vega', 'ssd', 'wan21']