improved correction of pred norm by targeting the prior

toolkit/config_modules.py

@@ -306,6 +306,7 @@ class ModelConfig:
         self.refiner_name_or_path = kwargs.get('refiner_name_or_path', None)
         self._original_refiner_name_or_path = self.refiner_name_or_path
         self.refiner_start_at = kwargs.get('refiner_start_at', 0.5)
+        self.lora_path = kwargs.get('lora_path', None)
 
         # only for SDXL models for now
         self.use_text_encoder_1: bool = kwargs.get('use_text_encoder_1', True)

toolkit/stable_diffusion_model.py

@@ -281,6 +281,12 @@ class StableDiffusion:
         self.unet.requires_grad_(False)
         self.unet.eval()
 
+        # load any loras we have
+        if self.model_config.lora_path is not None:
+            pipe.load_lora_weights(self.model_config.lora_path, adapter_name="lora1")
+            pipe.fuse_lora()
+            self.unet.fuse_lora()
+
         self.tokenizer = tokenizer
         self.text_encoder = text_encoder
         self.pipeline = pipe