Bug fixes and little improvements here and there.

toolkit/config_modules.py

@@ -129,13 +129,13 @@ class NetworkConfig:
                 self.conv = 4
 
 
-AdapterTypes = Literal['t2i', 'ip', 'ip+', 'clip', 'ilora', 'photo_maker']
+AdapterTypes = Literal['t2i', 'ip', 'ip+', 'clip', 'ilora', 'photo_maker', 'control_net']
 
 CLIPLayer = Literal['penultimate_hidden_states', 'image_embeds', 'last_hidden_state']
 
 class AdapterConfig:
     def __init__(self, **kwargs):
-        self.type: AdapterTypes = kwargs.get('type', 't2i')  # t2i, ip, clip
+        self.type: AdapterTypes = kwargs.get('type', 't2i')  # t2i, ip, clip, control_net
         self.in_channels: int = kwargs.get('in_channels', 3)
         self.channels: List[int] = kwargs.get('channels', [320, 640, 1280, 1280])
         self.num_res_blocks: int = kwargs.get('num_res_blocks', 2)
@@ -530,6 +530,8 @@ class DatasetConfig:
         self.prefetch_factor: int = kwargs.get('prefetch_factor', 2)
         self.extra_values: List[float] = kwargs.get('extra_values', [])
         self.square_crop: bool = kwargs.get('square_crop', False)
+        # apply same augmentations to control images. Usually want this true unless special case
+        self.replay_transforms: bool = kwargs.get('replay_transforms', True)
 
 
 def preprocess_dataset_raw_config(raw_config: List[dict]) -> List[dict]:

toolkit/dataloader_mixins.py

@@ -860,7 +860,8 @@ class AugmentationFileItemDTOMixin:
         # only store the spatial transforms
         augmented_params['transforms'] = [t for t in augmented_params['transforms'] if t['__class_fullname__'].split('.')[-1] in spatial_transforms]
 
-        self.aug_replay_spatial_transforms = augmented_params
+        if self.dataset_config.replay_transforms:
+            self.aug_replay_spatial_transforms = augmented_params
 
         # convert back to RGB tensor
         augmented = cv2.cvtColor(augmented, cv2.COLOR_BGR2RGB)

toolkit/guidance.py

@@ -240,7 +240,7 @@ def get_direct_guidance_loss(
 
     noise_pred_uncond, noise_pred_cond = torch.chunk(prediction, 2, dim=0)
 
-    guidance_scale = 1.0
+    guidance_scale = 1.25
     guidance_pred = noise_pred_uncond + guidance_scale * (
             noise_pred_cond - noise_pred_uncond
     )
@@ -586,6 +586,8 @@ def get_guided_tnt(
 
     loss = prior_loss + this_loss - that_loss
 
+    loss = loss.mean()
+
     loss.backward()
 
     # detach it so parent class can run backward on no grads without throwing error

toolkit/optimizer.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import Adafactor
+from transformers import Adafactor, AdamW
 
 
 def get_optimizer(
@@ -69,7 +69,7 @@ def get_optimizer(
         if 'relative_step' not in optimizer_params:
             optimizer_params['relative_step'] = False
         if 'scale_parameter' not in optimizer_params:
-            optimizer_params['scale_parameter'] = True
+            optimizer_params['scale_parameter'] = False
         if 'warmup_init' not in optimizer_params:
             optimizer_params['warmup_init'] = False
         optimizer = Adafactor(params, lr=float(learning_rate), eps=1e-6, **optimizer_params)

toolkit/stable_diffusion_model.py

@@ -39,7 +39,8 @@ from toolkit.pipelines import CustomStableDiffusionXLPipeline, CustomStableDiffu
     StableDiffusionKDiffusionXLPipeline, StableDiffusionXLRefinerPipeline
 from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, T2IAdapter, DDPMScheduler, \
     StableDiffusionXLAdapterPipeline, StableDiffusionAdapterPipeline, DiffusionPipeline, \
-    StableDiffusionXLImg2ImgPipeline, LCMScheduler, Transformer2DModel, AutoencoderTiny, ControlNetModel
+    StableDiffusionXLImg2ImgPipeline, LCMScheduler, Transformer2DModel, AutoencoderTiny, ControlNetModel, \
+    StableDiffusionXLControlNetPipeline, StableDiffusionControlNetPipeline
 import diffusers
 from diffusers import \
     AutoencoderKL, \
@@ -497,6 +498,12 @@ class StableDiffusion:
                     else:
                         Pipe = StableDiffusionAdapterPipeline
                     extra_args['adapter'] = self.adapter
+                elif isinstance(self.adapter, ControlNetModel):
+                    if self.is_xl:
+                        Pipe = StableDiffusionXLControlNetPipeline
+                    else:
+                        Pipe = StableDiffusionControlNetPipeline
+                    extra_args['controlnet'] = self.adapter
                 elif isinstance(self.adapter, ReferenceAdapter):
                     # pass the noise scheduler to the adapter
                     self.adapter.noise_scheduler = noise_scheduler
@@ -588,6 +595,10 @@ class StableDiffusion:
                             validation_image = validation_image.resize((gen_config.width * 2, gen_config.height * 2))
                             extra['image'] = validation_image
                             extra['adapter_conditioning_scale'] = gen_config.adapter_conditioning_scale
+                        if isinstance(self.adapter, ControlNetModel):
+                            validation_image = validation_image.resize((gen_config.width, gen_config.height))
+                            extra['image'] = validation_image
+                            extra['controlnet_conditioning_scale'] = gen_config.adapter_conditioning_scale
                         if isinstance(self.adapter, IPAdapter) or isinstance(self.adapter, ClipVisionAdapter):
                             transform = transforms.Compose([
                                 transforms.ToTensor(),
@@ -967,6 +978,16 @@ class StableDiffusion:
                 if do_classifier_free_guidance and item.shape[0] != text_embeddings.text_embeds.shape[0]:
                     kwargs['down_intrablock_additional_residuals'][idx] = torch.cat([item] * 2, dim=0)
 
+        # handle controlnet
+        if 'down_block_additional_residuals' in kwargs and 'mid_block_additional_residual' in kwargs:
+            # go through each item and concat if doing cfg and it doesnt have the same shape
+            for idx, item in enumerate(kwargs['down_block_additional_residuals']):
+                if do_classifier_free_guidance and item.shape[0] != text_embeddings.text_embeds.shape[0]:
+                    kwargs['down_block_additional_residuals'][idx] = torch.cat([item] * 2, dim=0)
+            for idx, item in enumerate(kwargs['mid_block_additional_residual']):
+                if do_classifier_free_guidance and item.shape[0] != text_embeddings.text_embeds.shape[0]:
+                    kwargs['mid_block_additional_residual'][idx] = torch.cat([item] * 2, dim=0)
+
         def scale_model_input(model_input, timestep_tensor):
             if is_input_scaled:
                 return model_input
@@ -1383,11 +1404,13 @@ class StableDiffusion:
         # move to device and dtype
         image_list = [image.to(self.device, dtype=self.torch_dtype) for image in image_list]
 
+        VAE_SCALE_FACTOR = 2 ** (len(self.vae.config['block_out_channels']) - 1)
+
         # resize images if not divisible by 8
         for i in range(len(image_list)):
             image = image_list[i]
-            if image.shape[1] % 8 != 0 or image.shape[2] % 8 != 0:
-                image_list[i] = Resize((image.shape[1] // 8 * 8, image.shape[2] // 8 * 8))(image)
+            if image.shape[1] % VAE_SCALE_FACTOR != 0 or image.shape[2] % VAE_SCALE_FACTOR != 0:
+                image_list[i] = Resize((image.shape[1] // VAE_SCALE_FACTOR * VAE_SCALE_FACTOR, image.shape[2] // VAE_SCALE_FACTOR * VAE_SCALE_FACTOR))(image)
 
         images = torch.stack(image_list)
         if isinstance(self.vae, AutoencoderTiny):
@@ -1756,6 +1779,9 @@ class StableDiffusion:
             elif isinstance(self.adapter, T2IAdapter):
                 requires_grad = self.adapter.adapter.conv_in.weight.requires_grad
                 adapter_device = self.adapter.device
+            elif isinstance(self.adapter, ControlNetModel):
+                requires_grad = self.adapter.conv_in.training
+                adapter_device = self.adapter.device
             elif isinstance(self.adapter, ClipVisionAdapter):
                 requires_grad = self.adapter.embedder.training
                 adapter_device = self.adapter.device

toolkit/style.py

@@ -158,7 +158,7 @@ def get_style_model_and_losses(
 ):
     # content_layers = ['conv_4']
     # style_layers = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
-    content_layers = ['conv2_2', 'conv3_2', 'conv4_2', 'conv5_2']
+    content_layers = ['conv2_2', 'conv3_2', 'conv4_2']
     style_layers = ['conv2_1', 'conv3_1', 'conv4_1']
     cnn = models.vgg19(pretrained=True).features.to(device, dtype=dtype).eval()
     # set all weights in the model to our dtype

toolkit/util/adafactor_stochastic_rounding.py

@@ -81,7 +81,8 @@ def step_adafactor(self, closure=None):
             lr = self._get_lr(group, state)
 
             beta2t = 1.0 - math.pow(state["step"], group["decay_rate"])
-            update = (grad ** 2) + group["eps"][0]
+            eps = group["eps"][0] if isinstance(group["eps"], list) else group["eps"]
+            update = (grad ** 2) + eps
             if factored:
                 exp_avg_sq_row = state["exp_avg_sq_row"]
                 exp_avg_sq_col = state["exp_avg_sq_col"]