WIP creating textual inversion training script

toolkit/config_modules.py

@@ -1,6 +1,6 @@
 import os
 import time
-from typing import List, Optional
+from typing import List, Optional, Literal
 import random
 
 
@@ -50,6 +50,13 @@ class NetworkConfig:
         self.conv_alpha: float = kwargs.get('conv_alpha', self.conv)
 
 
+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.save_format = kwargs.get('save_format', 'safetensors')
+
 class TrainConfig:
     def __init__(self, **kwargs):
         self.noise_scheduler = kwargs.get('noise_scheduler', 'ddpm')
@@ -146,6 +153,20 @@ class SliderConfig:
                 self.targets.append(target)
 
 
+class DatasetConfig:
+    caption_type: Literal["txt", "caption"] = 'txt'
+
+    def __init__(self, **kwargs):
+        self.type = kwargs.get('type', 'image')  # sd, slider, reference
+        self.folder_path: str = kwargs.get('folder_path', None)
+        self.default_caption: str = kwargs.get('default_caption', None)
+        self.caption_type: str = kwargs.get('caption_type', None)
+        self.random_scale: bool = kwargs.get('random_scale', False)
+        self.random_crop: bool = kwargs.get('random_crop', False)
+        self.resolution: int = kwargs.get('resolution', 512)
+        self.scale: float = kwargs.get('scale', 1.0)
+
+
 class GenerateImageConfig:
     def __init__(
             self,

toolkit/data_loader.py

@@ -1,23 +1,33 @@
 import os
 import random
+from typing import List
 
 import cv2
 import numpy as np
 from PIL import Image
 from PIL.ImageOps import exif_transpose
 from torchvision import transforms
-from torch.utils.data import Dataset
+from torch.utils.data import Dataset, DataLoader, ConcatDataset
 from tqdm import tqdm
 import albumentations as A
 
+from toolkit.config_modules import DatasetConfig
+from toolkit.dataloader_mixins import CaptionMixin
 
-class ImageDataset(Dataset):
+
+class ImageDataset(Dataset, CaptionMixin):
     def __init__(self, config):
         self.config = config
         self.name = self.get_config('name', 'dataset')
         self.path = self.get_config('path', required=True)
         self.scale = self.get_config('scale', 1)
         self.random_scale = self.get_config('random_scale', False)
+        self.include_prompt = self.get_config('include_prompt', False)
+        self.default_prompt = self.get_config('default_prompt', '')
+        if self.include_prompt:
+            self.caption_type = self.get_config('caption_type', 'txt')
+        else:
+            self.caption_type = None
         # we always random crop if random scale is enabled
         self.random_crop = self.random_scale if self.random_scale else self.get_config('random_crop', False)
 
@@ -81,7 +91,11 @@ class ImageDataset(Dataset):
 
         img = self.transform(img)
 
-        return img
+        if self.include_prompt:
+            prompt = self.get_caption_item(index)
+            return img, prompt
+        else:
+            return img
 
 
 class Augments:
@@ -268,3 +282,101 @@ class PairedImageDataset(Dataset):
         img = self.transform(img)
 
         return img, prompt, (self.neg_weight, self.pos_weight)
+
+
+class AiToolkitDataset(Dataset, CaptionMixin):
+    def __init__(self, dataset_config: 'DatasetConfig'):
+        self.dataset_config = dataset_config
+        self.folder_path = dataset_config.folder_path
+        self.caption_type = dataset_config.caption_type
+        self.default_caption = dataset_config.default_caption
+        self.random_scale = dataset_config.random_scale
+        self.scale = dataset_config.scale
+        # we always random crop if random scale is enabled
+        self.random_crop = self.random_scale if self.random_scale else dataset_config.random_crop
+        self.resolution = dataset_config.resolution
+
+        # get the file list
+        self.file_list = [
+            os.path.join(self.folder_path, file) for file in os.listdir(self.folder_path) if
+            file.lower().endswith(('.jpg', '.jpeg', '.png', '.webp'))
+        ]
+
+        # this might take a while
+        print(f"  -  Preprocessing image dimensions")
+        new_file_list = []
+        bad_count = 0
+        for file in tqdm(self.file_list):
+            img = Image.open(file)
+            if int(min(img.size) * self.scale) >= self.resolution:
+                new_file_list.append(file)
+            else:
+                bad_count += 1
+
+        print(f"  -  Found {len(self.file_list)} images")
+        print(f"  -  Found {bad_count} images that are too small")
+        assert len(self.file_list) > 0, f"no images found in {self.folder_path}"
+
+        self.transform = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5]),  # normalize to [-1, 1]
+        ])
+
+    def __len__(self):
+        return len(self.file_list)
+
+    def __getitem__(self, index):
+        img_path = self.file_list[index]
+        img = exif_transpose(Image.open(img_path)).convert('RGB')
+
+        # Downscale the source image first
+        img = img.resize((int(img.size[0] * self.scale), int(img.size[1] * self.scale)), Image.BICUBIC)
+        min_img_size = min(img.size)
+
+        if self.random_crop:
+            if self.random_scale and min_img_size > self.resolution:
+                if min_img_size < self.resolution:
+                    print(
+                        f"Unexpected values: min_img_size={min_img_size}, self.resolution={self.resolution}, image file={img_path}")
+                    scale_size = self.resolution
+                else:
+                    scale_size = random.randint(self.resolution, int(min_img_size))
+                img = img.resize((scale_size, scale_size), Image.BICUBIC)
+            img = transforms.RandomCrop(self.resolution)(img)
+        else:
+            img = transforms.CenterCrop(min_img_size)(img)
+            img = img.resize((self.resolution, self.resolution), Image.BICUBIC)
+
+        img = self.transform(img)
+
+        if self.caption_type is not None:
+            prompt = self.get_caption_item(index)
+            return img, prompt
+        else:
+            return img
+
+
+def get_dataloader_from_datasets(dataset_options, batch_size=1):
+    if dataset_options is None or len(dataset_options) == 0:
+        return None
+
+    datasets = []
+    for dataset_option in dataset_options:
+        if isinstance(dataset_option, DatasetConfig):
+            config = dataset_option
+        else:
+            config = DatasetConfig(**dataset_option)
+        if config.type == 'image':
+            dataset = AiToolkitDataset(config)
+            datasets.append(dataset)
+        else:
+            raise ValueError(f"invalid dataset type: {config.type}")
+
+    concatenated_dataset = ConcatDataset(datasets)
+    data_loader = DataLoader(
+        concatenated_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        num_workers=2
+    )
+    return data_loader

toolkit/dataloader_mixins.py

@@ -0,0 +1,43 @@
+import os
+
+
+class CaptionMixin:
+    def get_caption_item(self, index):
+        if not hasattr(self, 'caption_type'):
+            raise Exception('caption_type not found on class instance')
+        if not hasattr(self, 'file_list'):
+            raise Exception('file_list not found on class instance')
+        img_path_or_tuple = self.file_list[index]
+        if isinstance(img_path_or_tuple, tuple):
+            # check if either has a prompt file
+            path_no_ext = os.path.splitext(img_path_or_tuple[0])[0]
+            prompt_path = path_no_ext + '.txt'
+            if not os.path.exists(prompt_path):
+                path_no_ext = os.path.splitext(img_path_or_tuple[1])[0]
+                prompt_path = path_no_ext + '.txt'
+        else:
+            img_path = img_path_or_tuple
+            # see if prompt file exists
+            path_no_ext = os.path.splitext(img_path)[0]
+            prompt_path = path_no_ext + '.txt'
+
+        if os.path.exists(prompt_path):
+            with open(prompt_path, 'r', encoding='utf-8') as f:
+                prompt = f.read()
+                # remove any newlines
+                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
+            if hasattr(self, 'default_prompt'):
+                prompt = self.default_prompt
+            if hasattr(self, 'default_caption'):
+                prompt = self.default_caption
+        return prompt

toolkit/embedding.py

@@ -0,0 +1,185 @@
+import json
+import os
+from collections import OrderedDict
+
+import safetensors
+import torch
+from typing import TYPE_CHECKING
+
+from safetensors.torch import save_file
+
+from toolkit.metadata import get_meta_for_safetensors
+
+if TYPE_CHECKING:
+    from toolkit.stable_diffusion_model import StableDiffusion
+    from toolkit.config_modules import EmbeddingConfig
+
+
+# this is a frankenstein mix of automatic1111 and my own code
+
+class Embedding:
+    def __init__(
+            self,
+            sd: 'StableDiffusion',
+            embed_config: 'EmbeddingConfig'
+    ):
+        self.name = embed_config.trigger
+        self.sd = sd
+        self.embed_config = embed_config
+        # setup our embedding
+        # Add the placeholder token in tokenizer
+        placeholder_tokens = [self.embed_config.trigger]
+
+        # add dummy tokens for multi-vector
+        additional_tokens = []
+        for i in range(1, self.embed_config.tokens):
+            additional_tokens.append(f"{self.embed_config.trigger}_{i}")
+        placeholder_tokens += additional_tokens
+
+        num_added_tokens = self.sd.tokenizer.add_tokens(placeholder_tokens)
+        if num_added_tokens != self.embed_config.tokens:
+            raise ValueError(
+                f"The tokenizer already contains the token {self.embed_config.trigger}. Please pass a different"
+                " `placeholder_token` that is not already in the tokenizer."
+            )
+
+        # Convert the initializer_token, placeholder_token to ids
+        init_token_ids = self.sd.tokenizer.encode(self.embed_config.init_words, add_special_tokens=False)
+        # if length of token ids is more than number of orm embedding tokens fill with *
+        if len(init_token_ids) > self.embed_config.tokens:
+            init_token_ids = init_token_ids[:self.embed_config.tokens]
+        elif len(init_token_ids) < self.embed_config.tokens:
+            pad_token_id = self.sd.tokenizer.encode(["*"], add_special_tokens=False)
+            init_token_ids += pad_token_id * (self.embed_config.tokens - len(init_token_ids))
+
+        self.placeholder_token_ids = self.sd.tokenizer.convert_tokens_to_ids(placeholder_tokens)
+
+        # Resize the token embeddings as we are adding new special tokens to the tokenizer
+        # todo SDXL has 2 text encoders, need to do both for all of this
+        self.sd.text_encoder.resize_token_embeddings(len(self.sd.tokenizer))
+
+        # Initialise the newly added placeholder token with the embeddings of the initializer token
+        token_embeds = self.sd.text_encoder.get_input_embeddings().weight.data
+        with torch.no_grad():
+            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
+        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
+    def get_embedding_string(self):
+        return self.embedding_tokens
+
+    def get_trainable_params(self):
+        # todo only get this one as we could have more than one
+        return self.sd.text_encoder.get_input_embeddings().parameters()
+
+    # make setter and getter for vec
+    @property
+    def vec(self):
+        # should we get params instead
+        # create vector from token embeds
+        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],
+            dim=0
+        )
+        return new_vector
+
+    @vec.setter
+    def vec(self, new_vector):
+        # shape is (1, 768) for SD 1.5 for 1 token
+        token_embeds = self.sd.text_encoder.get_input_embeddings().weight.data
+        for i in range(new_vector.shape[0]):
+            # apply the weights to the placeholder tokens while preserving gradient
+            token_embeds[self.placeholder_token_ids[i]] = new_vector[i].clone()
+            x = 1
+
+    def save(self, filename):
+        # todo check to see how to get the vector out of the embedding
+
+        embedding_data = {
+            "string_to_token": {"*": 265},
+            "string_to_param": {"*": self.vec},
+            "name": self.name,
+            "step": 0,
+            # todo get these
+            "sd_checkpoint": None,
+            "sd_checkpoint_name": None,
+            "notes": None,
+        }
+        if filename.endswith('.pt'):
+            torch.save(embedding_data, filename)
+        elif filename.endswith('.bin'):
+            torch.save(embedding_data, filename)
+        elif filename.endswith('.safetensors'):
+            # save the embedding as a safetensors file
+            state_dict = {"emb_params": self.vec}
+            # add all embedding data (except string_to_param), to metadata
+            metadata = OrderedDict({k: json.dumps(v) for k, v in embedding_data.items() if k != "string_to_param"})
+            metadata["string_to_param"] = {"*": "emb_params"}
+            save_meta = get_meta_for_safetensors(metadata, name=self.name)
+            save_file(state_dict, filename, metadata=save_meta)
+
+    def load_embedding_from_file(self, file_path, device):
+        # full path
+        path = os.path.realpath(file_path)
+        filename = os.path.basename(path)
+        name, ext = os.path.splitext(filename)
+        ext = ext.upper()
+        if ext in ['.PNG', '.WEBP', '.JXL', '.AVIF']:
+            _, second_ext = os.path.splitext(name)
+            if second_ext.upper() == '.PREVIEW':
+                return
+
+        if ext in ['.BIN', '.PT']:
+            data = torch.load(path, map_location="cpu")
+        elif ext in ['.SAFETENSORS']:
+            # rebuild the embedding from the safetensors file if it has it
+            tensors = {}
+            with safetensors.torch.safe_open(path, framework="pt", device="cpu") as f:
+                metadata = f.metadata()
+                for k in f.keys():
+                    tensors[k] = f.get_tensor(k)
+            # data = safetensors.torch.load_file(path, device="cpu")
+            if metadata and 'string_to_param' in metadata and 'emb_params' in tensors:
+                # our format
+                def try_json(v):
+                    try:
+                        return json.loads(v)
+                    except:
+                        return v
+
+                data = {k: try_json(v) for k, v in metadata.items()}
+                data['string_to_param'] = {'*': tensors['emb_params']}
+            else:
+                # old format
+                data = tensors
+        else:
+            return
+
+        # textual inversion embeddings
+        if 'string_to_param' in data:
+            param_dict = data['string_to_param']
+            if hasattr(param_dict, '_parameters'):
+                param_dict = getattr(param_dict,
+                                     '_parameters')  # fix for torch 1.12.1 loading saved file from torch 1.11
+            assert len(param_dict) == 1, 'embedding file has multiple terms in it'
+            emb = next(iter(param_dict.items()))[1]
+        # diffuser concepts
+        elif type(data) == dict and type(next(iter(data.values()))) == torch.Tensor:
+            assert len(data.keys()) == 1, 'embedding file has multiple terms in it'
+
+            emb = next(iter(data.values()))
+            if len(emb.shape) == 1:
+                emb = emb.unsqueeze(0)
+        else:
+            raise Exception(
+                f"Couldn't identify {filename} as neither textual inversion embedding nor diffuser concept.")
+
+        self.vec = emb.detach().to(device, dtype=torch.float32)

toolkit/stable_diffusion_model.py

@@ -520,6 +520,11 @@ class StableDiffusion:
                         noise_pred_text - noise_pred_uncond
                 )
 
+                # https://github.com/huggingface/diffusers/blob/7a91ea6c2b53f94da930a61ed571364022b21044/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py#L775
+                if guidance_rescale > 0.0:
+                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+
         return noise_pred
 
     # ref: https://github.com/huggingface/diffusers/blob/0bab447670f47c28df60fbd2f6a0f833f75a16f5/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L746