Tied in ant tested TI script

toolkit/config_modules.py

@@ -54,9 +54,10 @@ class EmbeddingConfig:
     def __init__(self, **kwargs):
         self.trigger = kwargs.get('trigger', 'custom_embedding')
         self.tokens = kwargs.get('tokens', 4)
-        self.init_words = kwargs.get('init_phrase', '*')
+        self.init_words = kwargs.get('init_words', '*')
         self.save_format = kwargs.get('save_format', 'safetensors')
 
+
 class TrainConfig:
     def __init__(self, **kwargs):
         self.noise_scheduler = kwargs.get('noise_scheduler', 'ddpm')
@@ -75,6 +76,7 @@ class TrainConfig:
         self.optimizer_params = kwargs.get('optimizer_params', {})
         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)
 
 
 class ModelConfig:
@@ -165,6 +167,7 @@ class DatasetConfig:
         self.random_crop: bool = kwargs.get('random_crop', False)
         self.resolution: int = kwargs.get('resolution', 512)
         self.scale: float = kwargs.get('scale', 1.0)
+        self.buckets: bool = kwargs.get('buckets', False)
 
 
 class GenerateImageConfig:

toolkit/data_loader.py

@@ -14,6 +14,13 @@ import albumentations as A
 from toolkit.config_modules import DatasetConfig
 from toolkit.dataloader_mixins import CaptionMixin
 
+BUCKET_STEPS = 64
+
+def get_bucket_sizes_for_resolution(resolution: int) -> List[int]:
+    # make sure resolution is divisible by 8
+    if resolution % 8 != 0:
+        resolution = resolution - (resolution % 8)
+
 
 class ImageDataset(Dataset, CaptionMixin):
     def __init__(self, config):
@@ -357,6 +364,7 @@ class AiToolkitDataset(Dataset, CaptionMixin):
 
 
 def get_dataloader_from_datasets(dataset_options, batch_size=1):
+    # TODO do bucketing
     if dataset_options is None or len(dataset_options) == 0:
         return None
 

toolkit/embedding.py

@@ -25,7 +25,9 @@ class Embedding:
     ):
         self.name = embed_config.trigger
         self.sd = sd
+        self.trigger = embed_config.trigger
         self.embed_config = embed_config
+        self.step = 0
         # setup our embedding
         # Add the placeholder token in tokenizer
         placeholder_tokens = [self.embed_config.trigger]
@@ -64,10 +66,7 @@ class Embedding:
             for initializer_token_id, token_id in zip(init_token_ids, self.placeholder_token_ids):
                 token_embeds[token_id] = token_embeds[initializer_token_id].clone()
 
-        # this doesnt seem to be used again
-        self.token_embeds = token_embeds
-
-        # replace "[name] with this. This triggers it in the text encoder
+        # replace "[name] with this. on training. This is automatically generated in pipeline on inference
         self.embedding_tokens = " ".join(self.sd.tokenizer.convert_ids_to_tokens(self.placeholder_token_ids))
 
     # returns the string to have in the prompt to trigger the embedding
@@ -86,7 +85,7 @@ class Embedding:
         token_embeds = self.sd.text_encoder.get_input_embeddings().weight.data
         # stack the tokens along batch axis adding that axis
         new_vector = torch.stack(
-            [token_embeds[token_id].unsqueeze(0) for token_id in self.placeholder_token_ids],
+            [token_embeds[token_id] for token_id in self.placeholder_token_ids],
             dim=0
         )
         return new_vector
@@ -100,6 +99,39 @@ class Embedding:
             token_embeds[self.placeholder_token_ids[i]] = new_vector[i].clone()
             x = 1
 
+    # diffusers automatically expands the token meaning test123 becomes test123 test123_1 test123_2 etc
+    # however, on training we don't use that pipeline, so we have to do it ourselves
+    def inject_embedding_to_prompt(self, prompt, expand_token=False, to_replace_list=None, add_if_not_present=True):
+        output_prompt = prompt
+        default_replacements = [self.name, self.trigger, "[name]", "[trigger]", self.embedding_tokens]
+
+        replace_with = self.embedding_tokens if expand_token else self.trigger
+        if to_replace_list is None:
+            to_replace_list = default_replacements
+        else:
+            to_replace_list += default_replacements
+
+        # remove duplicates
+        to_replace_list = list(set(to_replace_list))
+
+        # replace them all
+        for to_replace in to_replace_list:
+            # replace it
+            output_prompt = output_prompt.replace(to_replace, replace_with)
+
+        # see how many times replace_with is in the prompt
+        num_instances = prompt.count(replace_with)
+
+        if num_instances == 0 and add_if_not_present:
+            # add it to the beginning of the prompt
+            output_prompt = replace_with + " " + output_prompt
+
+        if num_instances > 1:
+            print(
+                f"Warning: {self.name} token appears {num_instances} times in prompt {output_prompt}. This may cause issues.")
+
+        return output_prompt
+
     def save(self, filename):
         # todo check to see how to get the vector out of the embedding
 
@@ -107,7 +139,7 @@ class Embedding:
             "string_to_token": {"*": 265},
             "string_to_param": {"*": self.vec},
             "name": self.name,
-            "step": 0,
+            "step": self.step,
             # todo get these
             "sd_checkpoint": None,
             "sd_checkpoint_name": None,
@@ -182,4 +214,7 @@ class Embedding:
             raise Exception(
                 f"Couldn't identify {filename} as neither textual inversion embedding nor diffuser concept.")
 
+        if 'step' in data:
+            self.step = int(data['step'])
+
         self.vec = emb.detach().to(device, dtype=torch.float32)

toolkit/stable_diffusion_model.py

@@ -435,7 +435,7 @@ class StableDiffusion:
             text_embeddings = train_tools.concat_prompt_embeddings(
                 unconditional_embeddings,  # negative embedding
                 conditional_embeddings,  # positive embedding
-                latents.shape[0],  # batch size
+                1,  # batch size
             )
         elif text_embeddings is None and conditional_embeddings is not None:
             # not doing cfg
@@ -506,6 +506,17 @@ class StableDiffusion:
 
             latent_model_input = self.noise_scheduler.scale_model_input(latent_model_input, timestep)
 
+            # check if we need to concat timesteps
+            if isinstance(timestep, torch.Tensor):
+                ts_bs = timestep.shape[0]
+                if ts_bs != latent_model_input.shape[0]:
+                    if ts_bs == 1:
+                        timestep = torch.cat([timestep] * latent_model_input.shape[0])
+                    elif ts_bs * 2 == latent_model_input.shape[0]:
+                        timestep = torch.cat([timestep] * 2)
+                    else:
+                        raise ValueError(f"Batch size of latents {latent_model_input.shape[0]} must be the same or half the batch size of timesteps {timestep.shape[0]}")
+
             # predict the noise residual
             noise_pred = self.unet(
                 latent_model_input,