Fix other things...

rct_diffusion_pipeline.py

@@ -29,6 +29,7 @@ class RCTDiffusionPipeline(DiffusionPipeline):
 
         # channels for 1 image
         self.num_channels = int(self.unet.config.in_channels / 4)
+        self.load_dictionaries_from_dataset()
     
     def load_dictionaries_from_dataset(self):
         dataset = load_dataset('frutiemax/rct_dataset')

test_pipeline.py

@@ -1,17 +1,17 @@
 from rct_diffusion_pipeline import RCTDiffusionPipeline
 from diffusers import UNet2DConditionModel, DDPMScheduler, AutoencoderKL
-
+import torch
 
 torch_device = "cuda"
 
-unet = UNet2DConditionModel(sample_size=64, in_channels=16, out_channels=16, \
+unet = UNet2DConditionModel(sample_size=32, in_channels=16, out_channels=16, \
                         down_block_types=('CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'DownBlock2D'),\
                               up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D'), cross_attention_dim=160,
                             block_out_channels=(64, 128, 256), norm_num_groups=32)
+unet = unet.to('cuda', dtype=torch.float16)
 scheduler = DDPMScheduler(num_train_timesteps=20)
-vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae", use_safetensors=True)
-vae.tile_sample_min_size = 256
-
+vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", use_safetensors=True)
+vae = vae.to('cuda', dtype=torch.float16)
 
 pipeline = RCTDiffusionPipeline(unet, scheduler, vae)
 output = pipeline([[('aleppo pine tree', 1.0)]], [[('dark green', 1.0)]])

train_model.py

@@ -119,7 +119,6 @@ def train_model(batch_size=4, epochs=100, scheduler_num_timesteps=20, save_model
         num_training_steps=num_images * epochs
     )
     model = RCTDiffusionPipeline(unet, scheduler, vae)
-    model.load_dictionaries_from_dataset()
     labels = convert_labels(dataset, model, num_images)
     del model