Fixes and longer prompts

extensions_built_in/sd_trainer/SDTrainer.py

@@ -4,6 +4,7 @@ from diffusers import T2IAdapter
 from toolkit.basic import value_map
 from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO
 from toolkit.ip_adapter import IPAdapter
+from toolkit.prompt_utils import PromptEmbeds
 from toolkit.stable_diffusion_model import StableDiffusion, BlankNetwork
 from toolkit.train_tools import get_torch_dtype, apply_snr_weight
 import gc
@@ -27,6 +28,10 @@ class SDTrainer(BaseSDTrainProcess):
     def __init__(self, process_id: int, job, config: OrderedDict, **kwargs):
         super().__init__(process_id, job, config, **kwargs)
         self.assistant_adapter: Union['T2IAdapter', None]
+        self.do_prior_prediction = False
+        if self.train_config.inverted_mask_prior:
+            self.do_prior_prediction = True
+
 
     def before_model_load(self):
         pass
@@ -135,6 +140,40 @@ class SDTrainer(BaseSDTrainProcess):
     def preprocess_batch(self, batch: 'DataLoaderBatchDTO'):
         return batch
 
+    def get_prior_prediction(
+            self,
+            noisy_latents: torch.Tensor,
+            conditional_embeds: PromptEmbeds,
+            match_adapter_assist: bool,
+            network_weight_list: list,
+            timesteps: torch.Tensor,
+            pred_kwargs: dict,
+            batch: 'DataLoaderBatchDTO',
+            noise: torch.Tensor,
+            **kwargs
+    ):
+        # do a prediction here so we can match its output with network multiplier set to 0.0
+        with torch.no_grad():
+            dtype = get_torch_dtype(self.train_config.dtype)
+            # dont use network on this
+            self.network.multiplier = 0.0
+            self.sd.unet.eval()
+            prior_pred = self.sd.predict_noise(
+                latents=noisy_latents.to(self.device_torch, dtype=dtype).detach(),
+                conditional_embeddings=conditional_embeds.to(self.device_torch, dtype=dtype).detach(),
+                timestep=timesteps,
+                guidance_scale=1.0,
+                **pred_kwargs  # adapter residuals in here
+            )
+            self.sd.unet.train()
+            prior_pred = prior_pred.detach()
+            # remove the residuals as we wont use them on prediction when matching control
+            if match_adapter_assist and 'down_block_additional_residuals' in pred_kwargs:
+                del pred_kwargs['down_block_additional_residuals']
+            # restore network
+            self.network.multiplier = network_weight_list
+        return prior_pred
+
     def hook_train_loop(self, batch: 'DataLoaderBatchDTO'):
 
         self.timer.start('preprocess_batch')
@@ -287,28 +326,18 @@ class SDTrainer(BaseSDTrainProcess):
                         pred_kwargs['down_block_additional_residuals'] = down_block_additional_residuals
 
             prior_pred = None
-            if (has_adapter_img and self.assistant_adapter and match_adapter_assist) or self.train_config.inverted_mask_prior:
+            if (has_adapter_img and self.assistant_adapter and match_adapter_assist) or self.do_prior_prediction:
                 with self.timer('prior predict'):
-                    # do a prediction here so we can match its output with network multiplier set to 0.0
+                    prior_pred = self.get_prior_prediction(
-                    with torch.no_grad():
+                        noisy_latents=noisy_latents,
-                        # dont use network on this
+                        conditional_embeds=conditional_embeds,
-                        network.multiplier = 0.0
+                        match_adapter_assist=match_adapter_assist,
-                        self.sd.unet.eval()
+                        network_weight_list=network_weight_list,
-                        prior_pred = self.sd.predict_noise(
+                        timesteps=timesteps,
-                            latents=noisy_latents.to(self.device_torch, dtype=dtype).detach(),
+                        pred_kwargs=pred_kwargs,
-                            conditional_embeddings=conditional_embeds.to(self.device_torch, dtype=dtype).detach(),
+                        noise=noise,
-                            timestep=timesteps,
+                        batch=batch,
-                            guidance_scale=1.0,
+                    )
-                            **pred_kwargs  # adapter residuals in here
-                        )
-                        self.sd.unet.train()
-                        prior_pred = prior_pred.detach()
-                        # remove the residuals as we wont use them on prediction when matching control
-                        if match_adapter_assist and 'down_block_additional_residuals' in pred_kwargs:
-                            del pred_kwargs['down_block_additional_residuals']
-                        # restore network
-                        network.multiplier = network_weight_list
-
 
             if has_adapter_img and self.adapter and isinstance(self.adapter, IPAdapter):
                 with self.timer('encode_adapter'):

toolkit/dataloader_mixins.py

@@ -55,6 +55,18 @@ transforms_dict = {
 
 caption_ext_list = ['txt', 'json', 'caption']
 
+def clean_caption(caption):
+    # remove any newlines
+    caption = caption.replace('\n', ', ')
+    # remove new lines for all operating systems
+    caption = caption.replace('\r', ', ')
+    caption_split = caption.split(',')
+    # remove empty strings
+    caption_split = [p.strip() for p in caption_split if p.strip()]
+    # join back together
+    caption = ', '.join(caption_split)
+    return caption
+
 
 class CaptionMixin:
     def get_caption_item(self: 'AiToolkitDataset', index):
@@ -91,15 +103,7 @@ class CaptionMixin:
                     if 'caption' in prompt:
                         prompt = prompt['caption']
 
-                # remove any newlines
+                prompt = clean_caption(prompt)
-                prompt = prompt.replace('\n', ', ')
-                # remove new lines for all operating systems
-                prompt = prompt.replace('\r', ', ')
-                prompt_split = prompt.split(',')
-                # remove empty strings
-                prompt_split = [p.strip() for p in prompt_split if p.strip()]
-                # join back together
-                prompt = ', '.join(prompt_split)
         else:
             prompt = ''
             # get default_prompt if it exists on the class instance
@@ -135,6 +139,10 @@ class BucketsMixin:
                 batch = bucket.file_list_idx[start_idx:end_idx]
                 self.batch_indices.append(batch)
 
+    def shuffle_buckets(self: 'AiToolkitDataset'):
+        for key, bucket in self.buckets.items():
+            random.shuffle(bucket.file_list_idx)
+
     def setup_buckets(self: 'AiToolkitDataset', quiet=False):
         if not hasattr(self, 'file_list'):
             raise Exception(f'file_list not found on class instance {self.__class__.__name__}')
@@ -206,6 +214,7 @@ class BucketsMixin:
             self.buckets[bucket_key].file_list_idx.append(idx)
 
         # print the buckets
+        self.shuffle_buckets()
         self.build_batch_indices()
         if not quiet:
             print(f'Bucket sizes for {self.dataset_path}:')

toolkit/stable_diffusion_model.py

@@ -327,7 +327,6 @@ class StableDiffusion:
                 scheduler=noise_scheduler,
                 **extra_args
             ).to(self.device_torch)
-            # force turn that (ruin your images with obvious green and red dots) the #$@@ off!!!
             pipeline.watermark = None
         else:
             pipeline = Pipe(
@@ -372,7 +371,8 @@ class StableDiffusion:
                             extra['adapter_conditioning_scale'] = gen_config.adapter_conditioning_scale
                         if isinstance(self.adapter, IPAdapter):
                             transform = transforms.Compose([
-                                transforms.Resize(gen_config.width, interpolation=transforms.InterpolationMode.BILINEAR),
+                                transforms.Resize(gen_config.width,
+                                                  interpolation=transforms.InterpolationMode.BILINEAR),
                                 transforms.PILToTensor(),
                             ])
                             validation_image = transform(validation_image)
@@ -395,14 +395,15 @@ class StableDiffusion:
                         unconditional_embeds,
                     )
 
-                    if self.adapter is not None and isinstance(self.adapter, IPAdapter) and gen_config.adapter_image_path is not None:
+                    if self.adapter is not None and isinstance(self.adapter,
+                                                               IPAdapter) and gen_config.adapter_image_path is not None:
                         # apply the image projection
                         conditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors(validation_image)
-                        unconditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors(validation_image, True)
+                        unconditional_clip_embeds = self.adapter.get_clip_image_embeds_from_tensors(validation_image,
+                                                                                                    True)
                         conditional_embeds = self.adapter(conditional_embeds, conditional_clip_embeds)
                         unconditional_embeds = self.adapter(unconditional_embeds, unconditional_clip_embeds)
 
-
                     # 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
@@ -668,7 +669,15 @@ class StableDiffusion:
         # return latents_steps
         return latents
 
-    def encode_prompt(self, prompt, prompt2=None, num_images_per_prompt=1, force_all=False) -> PromptEmbeds:
+    def encode_prompt(
+            self,
+            prompt,
+            prompt2=None,
+            num_images_per_prompt=1,
+            force_all=False,
+            long_prompts=False,
+            max_length=None
+    ) -> PromptEmbeds:
         # sd1.5 embeddings are (bs, 77, 768)
         prompt = prompt
         # if it is not a list, make it one
@@ -695,12 +704,14 @@ class StableDiffusion:
                     num_images_per_prompt=num_images_per_prompt,
                     use_text_encoder_1=use_encoder_1,
                     use_text_encoder_2=use_encoder_2,
+                    truncate=not long_prompts,
+                    max_length=max_length,
                 )
             )
         else:
             return PromptEmbeds(
                 train_tools.encode_prompts(
-                    self.tokenizer, self.text_encoder, prompt
+                    self.tokenizer, self.text_encoder, prompt, truncate=not long_prompts, max_length=max_length
                 )
             )
 

toolkit/train_tools.py

@@ -447,29 +447,78 @@ if TYPE_CHECKING:
 
 
 def text_tokenize(
-        tokenizer: 'CLIPTokenizer',  # 普通ならひとつ、XLならふたつ！
+        tokenizer: 'CLIPTokenizer',
         prompts: list[str],
+        truncate: bool = True,
+        max_length: int = None,
+        max_length_multiplier: int = 4,
 ):
-    return tokenizer(
+    # allow fo up to 4x the max length for long prompts
+    if max_length is None:
+        if truncate:
+            max_length = tokenizer.model_max_length
+        else:
+            # allow up to 4x the max length for long prompts
+            max_length = tokenizer.model_max_length * max_length_multiplier
+
+    input_ids = tokenizer(
         prompts,
-        padding="max_length",
+        padding='max_length',
-        max_length=tokenizer.model_max_length,
+        max_length=max_length,
         truncation=True,
         return_tensors="pt",
     ).input_ids
 
+    if truncate or max_length == tokenizer.model_max_length:
+        return input_ids
+    else:
+        # remove additional padding
+        num_chunks = input_ids.shape[1] // tokenizer.model_max_length
+        chunks = torch.chunk(input_ids, chunks=num_chunks, dim=1)
+
+        # New list to store non-redundant chunks
+        non_redundant_chunks = []
+
+        for chunk in chunks:
+            if not chunk.eq(chunk[0, 0]).all():  # Check if all elements in the chunk are the same as the first element
+                non_redundant_chunks.append(chunk)
+
+        input_ids = torch.cat(non_redundant_chunks, dim=1)
+        return input_ids
+
 
 # https://github.com/huggingface/diffusers/blob/78922ed7c7e66c20aa95159c7b7a6057ba7d590d/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py#L334-L348
 def text_encode_xl(
         text_encoder: Union['CLIPTextModel', 'CLIPTextModelWithProjection'],
         tokens: torch.FloatTensor,
         num_images_per_prompt: int = 1,
+        max_length: int = 77,  # not sure what default to put here, always pass one?
+        truncate: bool = True,
 ):
-    prompt_embeds = text_encoder(
+    if truncate:
-        tokens.to(text_encoder.device), output_hidden_states=True
+        # normal short prompt 77 tokens max
-    )
+        prompt_embeds = text_encoder(
-    pooled_prompt_embeds = prompt_embeds[0]
+            tokens.to(text_encoder.device), output_hidden_states=True
-    prompt_embeds = prompt_embeds.hidden_states[-2]  # always penultimate layer
+        )
+        pooled_prompt_embeds = prompt_embeds[0]
+        prompt_embeds = prompt_embeds.hidden_states[-2]  # always penultimate layer
+    else:
+        # handle long prompts
+        prompt_embeds_list = []
+        tokens = tokens.to(text_encoder.device)
+        pooled_prompt_embeds = None
+        for i in range(0, tokens.shape[-1], max_length):
+            # todo run it through the in a single batch
+            section_tokens = tokens[:, i: i + max_length]
+            embeds = text_encoder(section_tokens, output_hidden_states=True)
+            pooled_prompt_embed = embeds[0]
+            if pooled_prompt_embeds is None:
+                # we only want the first ( I think??)
+                pooled_prompt_embeds = pooled_prompt_embed
+            prompt_embed = embeds.hidden_states[-2]  # always penultimate layer
+            prompt_embeds_list.append(prompt_embed)
+
+        prompt_embeds = torch.cat(prompt_embeds_list, dim=1)
 
     bs_embed, seq_len, _ = prompt_embeds.shape
     prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
@@ -485,7 +534,9 @@ def encode_prompts_xl(
         prompts2: Union[list[str], None],
         num_images_per_prompt: int = 1,
         use_text_encoder_1: bool = True,  # sdxl
-        use_text_encoder_2: bool = True  # sdxl
+        use_text_encoder_2: bool = True,  # sdxl
+        truncate: bool = True,
+        max_length=None,
 ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
     # text_encoder and text_encoder_2's penuultimate layer's output
     text_embeds_list = []
@@ -502,9 +553,14 @@ def encode_prompts_xl(
         if idx == 1 and not use_text_encoder_2:
             prompt_list_to_use = ["" for _ in prompts]
 
-        text_tokens_input_ids = text_tokenize(tokenizer, prompt_list_to_use)
+        text_tokens_input_ids = text_tokenize(tokenizer, prompt_list_to_use, truncate=truncate, max_length=max_length)
+        # set the max length for the next one
+        if idx == 0:
+            max_length = text_tokens_input_ids.shape[-1]
+
         text_embeds, pooled_text_embeds = text_encode_xl(
-            text_encoder, text_tokens_input_ids, num_images_per_prompt
+            text_encoder, text_tokens_input_ids, num_images_per_prompt, max_length=tokenizer.model_max_length,
+            truncate=truncate
         )
 
         text_embeds_list.append(text_embeds)
@@ -517,18 +573,36 @@ def encode_prompts_xl(
     return torch.concat(text_embeds_list, dim=-1), pooled_text_embeds
 
 
-def text_encode(text_encoder: 'CLIPTextModel', tokens):
+# ref for long prompts https://github.com/huggingface/diffusers/issues/2136
-    return text_encoder(tokens.to(text_encoder.device))[0]
+def text_encode(text_encoder: 'CLIPTextModel', tokens, truncate: bool = True, max_length=None):
+    if max_length is None and not truncate:
+        raise ValueError("max_length must be set if truncate is True")
 
+    tokens = tokens.to(text_encoder.device)
+
+    if truncate:
+        return text_encoder(tokens)[0]
+    else:
+        # handle long prompts
+        prompt_embeds_list = []
+        for i in range(0, tokens.shape[-1], max_length):
+            prompt_embeds = text_encoder(tokens[:, i: i + max_length])[0]
+            prompt_embeds_list.append(prompt_embeds)
+
+        return torch.cat(prompt_embeds_list, dim=1)
 
 
 def encode_prompts(
         tokenizer: 'CLIPTokenizer',
-        text_encoder: 'CLIPTokenizer',
+        text_encoder: 'CLIPTextModel',
         prompts: list[str],
+        truncate: bool = True,
+        max_length=None,
 ):
-    text_tokens = text_tokenize(tokenizer, prompts)
+    if max_length is None:
-    text_embeddings = text_encode(text_encoder, text_tokens)
+        max_length = tokenizer.model_max_length
+    text_tokens = text_tokenize(tokenizer, prompts, truncate=truncate, max_length=max_length)
+    text_embeddings = text_encode(text_encoder, text_tokens, truncate=truncate, max_length=max_length)
 
     return text_embeddings