New image generation img2img. various tweaks and fixes

toolkit/config_modules.py

@@ -266,6 +266,8 @@ class TrainConfig:
         self.reg_weight = kwargs.get('reg_weight', 1.0)
         self.num_train_timesteps = kwargs.get('num_train_timesteps', 1000)
         self.random_noise_shift = kwargs.get('random_noise_shift', 0.0)
+        # automatically adapte the vae scaling based on the image norm
+        self.adaptive_scaling_factor = kwargs.get('adaptive_scaling_factor', False)
 
         # dropout that happens before encoding. It functions independently per text encoder
         self.prompt_dropout_prob = kwargs.get('prompt_dropout_prob', 0.0)

toolkit/dataloader_mixins.py

@@ -387,7 +387,7 @@ class CaptionProcessingDTOMixin:
 
         # join back together
         caption = ', '.join(token_list)
-        caption = inject_trigger_into_prompt(caption, trigger, to_replace_list, add_if_not_present)
+        # caption = inject_trigger_into_prompt(caption, trigger, to_replace_list, add_if_not_present)
 
         if self.dataset_config.random_triggers:
             num_triggers = self.dataset_config.random_triggers_max

toolkit/guidance.py

@@ -407,6 +407,7 @@ def get_guided_loss_polarity(
         batch: 'DataLoaderBatchDTO',
         noise: torch.Tensor,
         sd: 'StableDiffusion',
+        scaler=None,
         **kwargs
 ):
     dtype = get_torch_dtype(sd.torch_dtype)
@@ -473,7 +474,10 @@ def get_guided_loss_polarity(
 
     loss = loss.mean([1, 2, 3])
     loss = loss.mean()
-    loss.backward()
+    if scaler is not None:
+        scaler.scale(loss).backward()
+    else:
+        loss.backward()
 
     # detach it so parent class can run backward on no grads without throwing error
     loss = loss.detach()
@@ -590,6 +594,7 @@ def get_guidance_loss(
         unconditional_embeds: Optional[PromptEmbeds] = None,
         mask_multiplier=None,
         prior_pred=None,
+        scaler=None,
         **kwargs
 ):
     # TODO add others and process individual batch items separately
@@ -621,6 +626,7 @@ def get_guidance_loss(
             batch,
             noise,
             sd,
+            scaler=scaler,
             **kwargs
         )
     elif guidance_type == "tnt":

toolkit/sampler.py

@@ -41,9 +41,12 @@ sd_config = {
     "prediction_type": "epsilon",
     "sample_max_value": 1.0,
     "set_alpha_to_one": False,
-    "skip_prk_steps": False,
-    "steps_offset": 1,
-    "timestep_spacing": "trailing",
+    # "skip_prk_steps": False,  # for training
+    "skip_prk_steps": True,
+    # "steps_offset": 1,
+    "steps_offset": 0,
+    # "timestep_spacing": "trailing", # for training
+    "timestep_spacing": "leading",
     "trained_betas": None
 }