Added training for an experimental decoratgor embedding. Allow for turning off guidance embedding on flux (for unreleased model). Various bug fixes and modifications

toolkit/config_modules.py

@@ -227,6 +227,11 @@ class EmbeddingConfig:
         self.trigger_class_name = kwargs.get('trigger_class_name', None)  # used for inverted masked prior
 
 
+class DecoratorConfig:
+    def __init__(self, **kwargs):
+        self.num_tokens: str = kwargs.get('num_tokens', 4)
+
+
 ContentOrStyleType = Literal['balanced', 'style', 'content']
 LossTarget = Literal['noise', 'source', 'unaugmented', 'differential_noise']
 
@@ -393,6 +398,8 @@ class TrainConfig:
         self.do_paramiter_swapping = kwargs.get('do_paramiter_swapping', False)
         # 0.1 is 10% of the parameters active at a time lower is less vram, higher is more
         self.paramiter_swapping_factor = kwargs.get('paramiter_swapping_factor', 0.1)
+        # bypass the guidance embedding for training. For open flux with guidance embedding
+        self.bypass_guidance_embedding = kwargs.get('bypass_guidance_embedding', False)
 
 
 class ModelConfig:
@@ -458,6 +465,7 @@ class ModelConfig:
         # for targeting a specific layers
         self.ignore_if_contains: Optional[List[str]] = kwargs.get("ignore_if_contains", None)
         self.only_if_contains: Optional[List[str]] = kwargs.get("only_if_contains", None)
+        self.quantize_kwargs = kwargs.get("quantize_kwargs", {})
         
         if self.ignore_if_contains is not None or self.only_if_contains is not None:
             if not self.is_flux:
@@ -914,4 +922,6 @@ def validate_configs(
         if save_config.save_format != 'diffusers':
             # make it diffusers
             save_config.save_format = 'diffusers'
-        
+        if model_config.use_flux_cfg:
+            # bypass the embedding
+            train_config.bypass_guidance_embedding = True

toolkit/custom_adapter.py

@@ -205,7 +205,8 @@ class CustomAdapter(torch.nn.Module):
         elif self.adapter_type == 'single_value':
             self.single_value_adapter = SingleValueAdapter(self, self.sd_ref(), num_values=self.config.num_tokens)
         elif self.adapter_type == 'redux':
-            self.redux_adapter = ReduxImageEncoder(1152, 4096, self.device, torch_dtype)
+            vision_hidden_size = self.vision_encoder.config.hidden_size
+            self.redux_adapter = ReduxImageEncoder(vision_hidden_size, 4096, self.device, torch_dtype)
         else:
             raise ValueError(f"unknown adapter type: {self.adapter_type}")
 

toolkit/models/decorator.py

@@ -0,0 +1,33 @@
+import torch
+
+
+class Decorator(torch.nn.Module):
+    def __init__(
+        self,
+        num_tokens: int = 4,
+        token_size: int = 4096,
+    ) -> None:
+        super().__init__()
+
+        self.weight: torch.nn.Parameter = torch.nn.Parameter(
+            torch.randn(num_tokens, token_size)
+        )
+        # ensure it is float32
+        self.weight.data = self.weight.data.float()
+
+    def forward(self, text_embeds: torch.Tensor, is_unconditional=False) -> torch.Tensor:
+        # make sure the param is float32
+        if self.weight.dtype != text_embeds.dtype:
+            self.weight.data = self.weight.data.float()
+        # expand batch to match text_embeds
+        batch_size = text_embeds.shape[0]
+        decorator_embeds = self.weight.unsqueeze(0).expand(batch_size, -1, -1)
+        if is_unconditional:
+            # zero pad the decorator embeds
+            decorator_embeds = torch.zeros_like(decorator_embeds)
+
+        if decorator_embeds.dtype != text_embeds.dtype:
+            decorator_embeds = decorator_embeds.to(text_embeds.dtype)
+        text_embeds = torch.cat((text_embeds, decorator_embeds), dim=-2)
+
+        return text_embeds

toolkit/models/flux.py

@@ -0,0 +1,35 @@
+
+# forward that bypasses the guidance embedding so it can be avoided during training.
+from functools import partial
+
+
+def guidance_embed_bypass_forward(self, timestep, guidance, pooled_projection):
+    timesteps_proj = self.time_proj(timestep)
+    timesteps_emb = self.timestep_embedder(
+        timesteps_proj.to(dtype=pooled_projection.dtype))  # (N, D)
+    pooled_projections = self.text_embedder(pooled_projection)
+    conditioning = timesteps_emb + pooled_projections
+    return conditioning
+
+# bypass the forward function
+
+
+def bypass_flux_guidance(transformer):
+    if hasattr(transformer.time_text_embed, '_bfg_orig_forward'):
+        return
+    # dont bypass if it doesnt have the guidance embedding
+    if not hasattr(transformer.time_text_embed, 'guidance_embedder'):
+        return
+    transformer.time_text_embed._bfg_orig_forward = transformer.time_text_embed.forward
+    transformer.time_text_embed.forward = partial(
+        guidance_embed_bypass_forward, transformer.time_text_embed
+    )
+
+# restore the forward function
+
+
+def restore_flux_guidance(transformer):
+    if not hasattr(transformer.time_text_embed, '_bfg_orig_forward'):
+        return
+    transformer.time_text_embed.forward = transformer.time_text_embed._bfg_orig_forward
+    del transformer.time_text_embed._bfg_orig_forward

toolkit/optimizers/adafactor.py

@@ -264,6 +264,10 @@ class Adafactor(torch.optim.Optimizer):
                 if grad.is_sparse:
                     raise RuntimeError(
                         "Adafactor does not support sparse gradients.")
+                
+                # if p has atts _scale then it is quantized. We need to divide the grad by the scale
+                # if hasattr(p, "_scale"):
+                #     grad = grad / p._scale
 
                 state = self.state[p]
                 grad_shape = grad.shape

toolkit/optimizers/automagic.py

@@ -101,7 +101,7 @@ class Automagic(torch.optim.Optimizer):
         if 'avg_lr' in param_state:
             lr = param_state["avg_lr"]
         else:
-            lr = param_state["lr"]
+            lr = 0.0
         return lr
 
     def _get_group_lr(self, group):
@@ -332,4 +332,4 @@ class Automagic(torch.optim.Optimizer):
                 state['lr_mask'] = Auto8bitTensor(sd_mask)
                 del state_dict['state'][idx]['lr_mask']
                 idx += 1
-        super().load_state_dict(state_dict, strict)
+        super().load_state_dict(state_dict)

toolkit/pipelines.py

@@ -14,6 +14,7 @@ from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import
 from diffusers.utils import is_torch_xla_available
 from k_diffusion.external import CompVisVDenoiser, CompVisDenoiser
 from k_diffusion.sampling import get_sigmas_karras, BrownianTreeNoiseSampler
+from toolkit.models.flux import bypass_flux_guidance, restore_flux_guidance
 
 
 if is_torch_xla_available():
@@ -1235,6 +1236,8 @@ class FluxWithCFGPipeline(FluxPipeline):
         callback_on_step_end_tensor_inputs: List[str] = ["latents"],
         max_sequence_length: int = 512,
     ):
+        # bypass the guidance embedding if there is one
+        bypass_flux_guidance(self.transformer)
 
         height = height or self.default_sample_size * self.vae_scale_factor
         width = width or self.default_sample_size * self.vae_scale_factor
@@ -1410,6 +1413,7 @@ class FluxWithCFGPipeline(FluxPipeline):
 
         # Offload all models
         self.maybe_free_model_hooks()
+        restore_flux_guidance(self.transformer)
 
         if not return_dict:
             return (image,)

toolkit/sd_device_states_presets.py

@@ -39,6 +39,7 @@ def get_train_sd_device_state_preset(
         train_lora: bool = False,
         train_adapter: bool = False,
         train_embedding: bool = False,
+        train_decorator: bool = False,
         train_refiner: bool = False,
         unload_text_encoder: bool = False,
         require_grads: bool = True,
@@ -89,6 +90,14 @@ def get_train_sd_device_state_preset(
         preset['unet']['requires_grad'] = False
         preset['unet']['device'] = device
         preset['text_encoder']['device'] = device
+    
+    if train_decorator:
+        preset['text_encoder']['training'] = False
+        preset['text_encoder']['requires_grad'] = False
+        preset['text_encoder']['device'] = device
+        preset['unet']['training'] = True
+        preset['unet']['requires_grad'] = False
+        preset['unet']['device'] = device
 
     if unload_text_encoder:
         preset['text_encoder']['training'] = False

toolkit/stable_diffusion_model.py

@@ -31,6 +31,7 @@ from library.model_util import convert_unet_state_dict_to_sd, convert_text_encod
 from toolkit import train_tools
 from toolkit.config_modules import ModelConfig, GenerateImageConfig
 from toolkit.metadata import get_meta_for_safetensors
+from toolkit.models.decorator import Decorator
 from toolkit.paths import REPOS_ROOT, KEYMAPS_ROOT
 from toolkit.prompt_utils import inject_trigger_into_prompt, PromptEmbeds, concat_prompt_embeds
 from toolkit.reference_adapter import ReferenceAdapter
@@ -59,6 +60,7 @@ from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextModelWithProjecti
 
 from toolkit.paths import ORIG_CONFIGS_ROOT, DIFFUSERS_CONFIGS_ROOT
 from huggingface_hub import hf_hub_download
+from toolkit.models.flux import bypass_flux_guidance, restore_flux_guidance
 
 from optimum.quanto import freeze, qfloat8, quantize, QTensor, qint4
 from typing import TYPE_CHECKING
@@ -165,6 +167,7 @@ class StableDiffusion:
         # to hold network if there is one
         self.network = None
         self.adapter: Union['ControlNetModel', 'T2IAdapter', 'IPAdapter', 'ReferenceAdapter', None] = None
+        self.decorator: Union[Decorator, None] = None
         self.is_xl = model_config.is_xl
         self.is_v2 = model_config.is_v2
         self.is_ssd = model_config.is_ssd
@@ -668,7 +671,7 @@ class StableDiffusion:
                 patch_dequantization_on_save(transformer)
                 quantization_type = qfloat8
                 print("Quantizing transformer")
-                quantize(transformer, weights=quantization_type)
+                quantize(transformer, weights=quantization_type, **self.model_config.quantize_kwargs)
                 freeze(transformer)
                 transformer.to(self.device_torch)
             else:
@@ -1209,6 +1212,11 @@ class StableDiffusion:
                         conditional_embeds,
                         unconditional_embeds,
                     )
+                    
+                    if self.decorator is not None:
+                        # apply the decorator to the embeddings
+                        conditional_embeds.text_embeds = self.decorator(conditional_embeds.text_embeds)
+                        unconditional_embeds.text_embeds = self.decorator(unconditional_embeds.text_embeds, is_unconditional=True)
 
                     if self.adapter is not None and isinstance(self.adapter, IPAdapter) \
                             and gen_config.adapter_image_path is not None:
@@ -1566,6 +1574,7 @@ class StableDiffusion:
             rescale_cfg=None,
             return_conditional_pred=False,
             guidance_embedding_scale=1.0,
+            bypass_guidance_embedding=False,
             **kwargs,
     ):
         conditional_pred = None
@@ -1842,13 +1851,16 @@ class StableDiffusion:
 
                         # # handle guidance
                         if self.unet.config.guidance_embeds:
-                            if isinstance(guidance_scale, list):
-                                guidance = torch.tensor(guidance_scale, device=self.device_torch)
+                            if isinstance(guidance_embedding_scale, list):
+                                guidance = torch.tensor(guidance_embedding_scale, device=self.device_torch)
                             else:
-                                guidance = torch.tensor([guidance_scale], device=self.device_torch)
+                                guidance = torch.tensor([guidance_embedding_scale], device=self.device_torch)
                                 guidance = guidance.expand(latents.shape[0])
                         else:
                             guidance = None
+                    
+                    if bypass_guidance_embedding:
+                        bypass_flux_guidance(self.unet)
 
                     cast_dtype = self.unet.dtype
                     # with torch.amp.autocast(device_type='cuda', dtype=cast_dtype):
@@ -1879,6 +1891,9 @@ class StableDiffusion:
                         pw=2,
                         c=latent_model_input.shape[1],
                     )
+                    
+                    if bypass_guidance_embedding:
+                        restore_flux_guidance(self.unet)
                 elif self.is_v3:
                     noise_pred = self.unet(
                         hidden_states=latent_model_input.to(self.device_torch, self.torch_dtype),