Fixed bug where changing image size in config wouldnt change anything…

…. Also Unet sometimes not loading

led/pipelines/led_pipeline.py

@@ -33,6 +33,7 @@ class LEDPipeline(DiffusionPipeline):
     def __init__(self, unet=None, scheduler=None, base_pipeline='ddim', backend=None, num_cond_steps=800, image_size=512):
         super().__init__()
         # make sure scheduler can always be converted to DDIM (basically)
+        self.unet = unet
         if scheduler is None:
             scheduler = self._set_default_scheduler()
         if unet is None:

led/trainers/led_trainer.py

@@ -269,6 +269,7 @@ def train(self):
                     pipeline = LEDPipeline(
                         unet=unwrap_model,
                         scheduler=self.noise_scheduler,
+                        image_size=self.config.data.image_size,
                     )
                     generator = torch.Generator(device=pipeline.device).manual_seed(0)
                     # run pipeline in inference (sample random noise and denoise)