Moved SD batch processing to a shared method and added it for use in slider training. Still testing if it affects quality over sampling

toolkit/config_modules.py

@@ -77,6 +77,7 @@ class TrainConfig:
         self.skip_first_sample = kwargs.get('skip_first_sample', False)
         self.gradient_checkpointing = kwargs.get('gradient_checkpointing', True)
         self.weight_jitter = kwargs.get('weight_jitter', 0.0)
+        self.merge_network_on_save = kwargs.get('merge_network_on_save', False)
 
 
 class ModelConfig:

toolkit/llvae.py

@@ -101,14 +101,16 @@ if __name__ == '__main__':
     from PIL import Image
     import torchvision.transforms as transforms
     user_path = os.path.expanduser('~')
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    dtype = torch.float32
 
     input_path = os.path.join(user_path, "Pictures/sample_2_512.png")
     output_path = os.path.join(user_path, "Pictures/sample_2_512_llvae.png")
     img = Image.open(input_path)
     img_tensor = transforms.ToTensor()(img)
-    img_tensor = img_tensor.unsqueeze(0)
+    img_tensor = img_tensor.unsqueeze(0).to(device=device, dtype=dtype)
     print("input_shape: ", list(img_tensor.shape))
-    vae = LosslessLatentVAE(in_channels=3, latent_depth=8)
+    vae = LosslessLatentVAE(in_channels=3, latent_depth=8, dtype=dtype).to(device=device, dtype=dtype)
     latent = vae.encode(img_tensor)
     print("latent_shape: ", list(latent.shape))
     out_tensor = vae.decode(latent)

toolkit/stable_diffusion_model.py

@@ -152,6 +152,11 @@ class StableDiffusion:
             steps_offset=1
         )
 
+        # move the betas alphas and  alphas_cumprod to device. Sometimed they get stuck on cpu, not sure why
+        scheduler.betas = scheduler.betas.to(self.device_torch)
+        scheduler.alphas = scheduler.alphas.to(self.device_torch)
+        scheduler.alphas_cumprod = scheduler.alphas_cumprod.to(self.device_torch)
+
         model_path = self.model_config.name_or_path
         if 'civitai.com' in self.model_config.name_or_path:
             # load is a civit ai model, use the loader.
@@ -323,6 +328,13 @@ class StableDiffusion:
 
                     # todo do we disable text encoder here as well if disabled for model, or only do that for training?
                     if self.is_xl:
+                        # fix guidance rescale for sdxl
+                        # was trained on 0.7 (I believe)
+
+                        grs = gen_config.guidance_rescale
+                        if grs is None or grs < 0.00001:
+                            grs = 0.7
+
                         img = pipeline(
                             prompt=gen_config.prompt,
                             prompt_2=gen_config.prompt_2,
@@ -332,7 +344,7 @@ class StableDiffusion:
                             width=gen_config.width,
                             num_inference_steps=gen_config.num_inference_steps,
                             guidance_scale=gen_config.guidance_scale,
-                            guidance_rescale=gen_config.guidance_rescale,
+                            guidance_rescale=grs,
                         ).images[0]
                     else:
                         img = pipeline(
@@ -619,6 +631,27 @@ class StableDiffusion:
 
         return latents
 
+    def decode_latents(
+            self,
+            latents: torch.Tensor,
+            device=None,
+            dtype=None
+    ):
+        if device is None:
+            device = self.device
+        if dtype is None:
+            dtype = self.torch_dtype
+
+        # Move to vae to device if on cpu
+        if self.vae.device == 'cpu':
+            self.vae.to(self.device)
+        latents = latents.to(device, dtype=dtype)
+        latents = latents / 0.18215
+        images = self.vae.decode(latents).sample
+        images = images.to(device, dtype=dtype)
+
+        return images
+
     def encode_image_prompt_pairs(
             self,
             prompt_list: List[str],