mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
925 lines
34 KiB
Python
925 lines
34 KiB
Python
# This code was copied from
|
|
# https://github.com/huggingface/diffusers/blob/main/examples/textual_inversion/textual_inversion.py
|
|
# on January 2, 2023
|
|
# and modified slightly by Lincoln Stein (@lstein) to work with InvokeAI
|
|
|
|
"""
|
|
This is the backend to "textual_inversion.py"
|
|
"""
|
|
|
|
import logging
|
|
import math
|
|
import os
|
|
import random
|
|
from argparse import Namespace
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import datasets
|
|
import diffusers
|
|
import numpy as np
|
|
import PIL
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torch.utils.checkpoint
|
|
import transformers
|
|
from accelerate import Accelerator
|
|
from accelerate.logging import get_logger
|
|
from accelerate.utils import ProjectConfiguration, set_seed
|
|
from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
|
|
from diffusers.optimization import get_scheduler
|
|
from diffusers.utils import check_min_version
|
|
from diffusers.utils.import_utils import is_xformers_available
|
|
from huggingface_hub import HfFolder, Repository, whoami
|
|
from packaging import version
|
|
from PIL import Image
|
|
from torch.utils.data import Dataset
|
|
from torchvision import transforms
|
|
from tqdm.auto import tqdm
|
|
from transformers import CLIPTextModel, CLIPTokenizer
|
|
|
|
# invokeai stuff
|
|
from invokeai.app.services.config import InvokeAIAppConfig, PagingArgumentParser
|
|
from invokeai.app.services.config.config_default import get_config
|
|
from invokeai.backend.install.install_helper import initialize_record_store
|
|
from invokeai.backend.model_manager import BaseModelType, ModelType
|
|
|
|
if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"):
|
|
PIL_INTERPOLATION = {
|
|
"linear": PIL.Image.Resampling.BILINEAR,
|
|
"bilinear": PIL.Image.Resampling.BILINEAR,
|
|
"bicubic": PIL.Image.Resampling.BICUBIC,
|
|
"lanczos": PIL.Image.Resampling.LANCZOS,
|
|
"nearest": PIL.Image.Resampling.NEAREST,
|
|
}
|
|
else:
|
|
PIL_INTERPOLATION = {
|
|
"linear": PIL.Image.LINEAR,
|
|
"bilinear": PIL.Image.BILINEAR,
|
|
"bicubic": PIL.Image.BICUBIC,
|
|
"lanczos": PIL.Image.LANCZOS,
|
|
"nearest": PIL.Image.NEAREST,
|
|
}
|
|
# ------------------------------------------------------------------------------
|
|
|
|
|
|
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
|
check_min_version("0.10.0.dev0")
|
|
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
def save_progress(text_encoder, placeholder_token_id, accelerator, placeholder_token, save_path):
|
|
logger.info("Saving embeddings")
|
|
learned_embeds = accelerator.unwrap_model(text_encoder).get_input_embeddings().weight[placeholder_token_id]
|
|
learned_embeds_dict = {placeholder_token: learned_embeds.detach().cpu()}
|
|
torch.save(learned_embeds_dict, save_path)
|
|
|
|
|
|
def parse_args() -> Namespace:
|
|
config = get_config()
|
|
parser = PagingArgumentParser(description="Textual inversion training")
|
|
general_group = parser.add_argument_group("General")
|
|
model_group = parser.add_argument_group("Models and Paths")
|
|
image_group = parser.add_argument_group("Training Image Location and Options")
|
|
trigger_group = parser.add_argument_group("Trigger Token")
|
|
training_group = parser.add_argument_group("Training Parameters")
|
|
checkpointing_group = parser.add_argument_group("Checkpointing and Resume")
|
|
integration_group = parser.add_argument_group("Integration")
|
|
general_group.add_argument(
|
|
"--front_end",
|
|
"--gui",
|
|
dest="front_end",
|
|
action="store_true",
|
|
default=False,
|
|
help="Activate the text-based graphical front end for collecting parameters. Aside from --root_dir, other parameters will be ignored.",
|
|
)
|
|
general_group.add_argument(
|
|
"--root_dir",
|
|
"--root",
|
|
type=Path,
|
|
default=config.root_path,
|
|
help="Path to the invokeai runtime directory",
|
|
)
|
|
general_group.add_argument(
|
|
"--logging_dir",
|
|
type=Path,
|
|
default="logs",
|
|
help=(
|
|
"[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
|
|
" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
|
|
),
|
|
)
|
|
general_group.add_argument(
|
|
"--output_dir",
|
|
type=Path,
|
|
default=f"{config.root_path}/text-inversion-model",
|
|
help="The output directory where the model predictions and checkpoints will be written.",
|
|
)
|
|
model_group.add_argument(
|
|
"--model",
|
|
type=str,
|
|
default="sd-1/main/stable-diffusion-v1-5",
|
|
help="Name of the diffusers model to train against.",
|
|
)
|
|
model_group.add_argument(
|
|
"--revision",
|
|
type=str,
|
|
default=None,
|
|
required=False,
|
|
help="Revision of pretrained model identifier from huggingface.co/models.",
|
|
)
|
|
|
|
model_group.add_argument(
|
|
"--tokenizer_name",
|
|
type=str,
|
|
default=None,
|
|
help="Pretrained tokenizer name or path if not the same as model_name",
|
|
)
|
|
image_group.add_argument(
|
|
"--train_data_dir",
|
|
type=Path,
|
|
default=None,
|
|
help="A folder containing the training data.",
|
|
)
|
|
image_group.add_argument(
|
|
"--resolution",
|
|
type=int,
|
|
default=512,
|
|
help=(
|
|
"The resolution for input images, all the images in the train/validation dataset will be resized to this"
|
|
" resolution"
|
|
),
|
|
)
|
|
image_group.add_argument(
|
|
"--center_crop",
|
|
action="store_true",
|
|
help="Whether to center crop images before resizing to resolution",
|
|
)
|
|
trigger_group.add_argument(
|
|
"--placeholder_token",
|
|
"--trigger_term",
|
|
dest="placeholder_token",
|
|
type=str,
|
|
default=None,
|
|
help='A token to use as a placeholder for the concept. This token will trigger the concept when included in the prompt as "<trigger>".',
|
|
)
|
|
trigger_group.add_argument(
|
|
"--learnable_property",
|
|
type=str,
|
|
choices=["object", "style"],
|
|
default="object",
|
|
help="Choose between 'object' and 'style'",
|
|
)
|
|
trigger_group.add_argument(
|
|
"--initializer_token",
|
|
type=str,
|
|
default="*",
|
|
help="A symbol to use as the initializer word.",
|
|
)
|
|
checkpointing_group.add_argument(
|
|
"--checkpointing_steps",
|
|
type=int,
|
|
default=500,
|
|
help=(
|
|
"Save a checkpoint of the training state every X updates. These checkpoints are only suitable for resuming"
|
|
" training using `--resume_from_checkpoint`."
|
|
),
|
|
)
|
|
checkpointing_group.add_argument(
|
|
"--resume_from_checkpoint",
|
|
type=Path,
|
|
default=None,
|
|
help=(
|
|
"Whether training should be resumed from a previous checkpoint. Use a path saved by"
|
|
' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
|
|
),
|
|
)
|
|
checkpointing_group.add_argument(
|
|
"--save_steps",
|
|
type=int,
|
|
default=500,
|
|
help="Save learned_embeds.bin every X updates steps.",
|
|
)
|
|
training_group.add_argument(
|
|
"--repeats",
|
|
type=int,
|
|
default=100,
|
|
help="How many times to repeat the training data.",
|
|
)
|
|
training_group.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
|
|
training_group.add_argument(
|
|
"--train_batch_size",
|
|
type=int,
|
|
default=16,
|
|
help="Batch size (per device) for the training dataloader.",
|
|
)
|
|
training_group.add_argument("--num_train_epochs", type=int, default=100)
|
|
training_group.add_argument(
|
|
"--max_train_steps",
|
|
type=int,
|
|
default=5000,
|
|
help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
|
|
)
|
|
training_group.add_argument(
|
|
"--gradient_accumulation_steps",
|
|
type=int,
|
|
default=1,
|
|
help="Number of updates steps to accumulate before performing a backward/update pass.",
|
|
)
|
|
training_group.add_argument(
|
|
"--gradient_checkpointing",
|
|
action="store_true",
|
|
help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
|
|
)
|
|
training_group.add_argument(
|
|
"--learning_rate",
|
|
type=float,
|
|
default=1e-4,
|
|
help="Initial learning rate (after the potential warmup period) to use.",
|
|
)
|
|
training_group.add_argument(
|
|
"--scale_lr",
|
|
action="store_true",
|
|
default=True,
|
|
help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
|
|
)
|
|
training_group.add_argument(
|
|
"--lr_scheduler",
|
|
type=str,
|
|
default="constant",
|
|
help=(
|
|
'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
|
|
' "constant", "constant_with_warmup"]'
|
|
),
|
|
)
|
|
training_group.add_argument(
|
|
"--lr_warmup_steps",
|
|
type=int,
|
|
default=500,
|
|
help="Number of steps for the warmup in the lr scheduler.",
|
|
)
|
|
training_group.add_argument(
|
|
"--adam_beta1",
|
|
type=float,
|
|
default=0.9,
|
|
help="The beta1 parameter for the Adam optimizer.",
|
|
)
|
|
training_group.add_argument(
|
|
"--adam_beta2",
|
|
type=float,
|
|
default=0.999,
|
|
help="The beta2 parameter for the Adam optimizer.",
|
|
)
|
|
training_group.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
|
|
training_group.add_argument(
|
|
"--adam_epsilon",
|
|
type=float,
|
|
default=1e-08,
|
|
help="Epsilon value for the Adam optimizer",
|
|
)
|
|
training_group.add_argument(
|
|
"--mixed_precision",
|
|
type=str,
|
|
default="no",
|
|
choices=["no", "fp16", "bf16"],
|
|
help=(
|
|
"Whether to use mixed precision. Choose"
|
|
"between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
|
|
"and an Nvidia Ampere GPU."
|
|
),
|
|
)
|
|
training_group.add_argument(
|
|
"--allow_tf32",
|
|
action="store_true",
|
|
help=(
|
|
"Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
|
|
" https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
|
|
),
|
|
)
|
|
training_group.add_argument(
|
|
"--local_rank",
|
|
type=int,
|
|
default=-1,
|
|
help="For distributed training: local_rank",
|
|
)
|
|
parser.add_argument(
|
|
"--enable_xformers_memory_efficient_attention",
|
|
action="store_true",
|
|
help="Whether or not to use xformers.",
|
|
)
|
|
|
|
integration_group.add_argument(
|
|
"--only_save_embeds",
|
|
action="store_true",
|
|
default=False,
|
|
help="Save only the embeddings for the new concept.",
|
|
)
|
|
integration_group.add_argument(
|
|
"--hub_model_id",
|
|
type=str,
|
|
default=None,
|
|
help="The name of the repository to keep in sync with the local `output_dir`.",
|
|
)
|
|
integration_group.add_argument(
|
|
"--report_to",
|
|
type=str,
|
|
default="tensorboard",
|
|
help=(
|
|
'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
|
|
' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
|
|
),
|
|
)
|
|
integration_group.add_argument(
|
|
"--push_to_hub",
|
|
action="store_true",
|
|
help="Whether or not to push the model to the Hub.",
|
|
)
|
|
integration_group.add_argument(
|
|
"--hub_token",
|
|
type=str,
|
|
default=None,
|
|
help="The token to use to push to the Model Hub.",
|
|
)
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
imagenet_templates_small = [
|
|
"a photo of a {}",
|
|
"a rendering of a {}",
|
|
"a cropped photo of the {}",
|
|
"the photo of a {}",
|
|
"a photo of a clean {}",
|
|
"a photo of a dirty {}",
|
|
"a dark photo of the {}",
|
|
"a photo of my {}",
|
|
"a photo of the cool {}",
|
|
"a close-up photo of a {}",
|
|
"a bright photo of the {}",
|
|
"a cropped photo of a {}",
|
|
"a photo of the {}",
|
|
"a good photo of the {}",
|
|
"a photo of one {}",
|
|
"a close-up photo of the {}",
|
|
"a rendition of the {}",
|
|
"a photo of the clean {}",
|
|
"a rendition of a {}",
|
|
"a photo of a nice {}",
|
|
"a good photo of a {}",
|
|
"a photo of the nice {}",
|
|
"a photo of the small {}",
|
|
"a photo of the weird {}",
|
|
"a photo of the large {}",
|
|
"a photo of a cool {}",
|
|
"a photo of a small {}",
|
|
]
|
|
|
|
imagenet_style_templates_small = [
|
|
"a painting in the style of {}",
|
|
"a rendering in the style of {}",
|
|
"a cropped painting in the style of {}",
|
|
"the painting in the style of {}",
|
|
"a clean painting in the style of {}",
|
|
"a dirty painting in the style of {}",
|
|
"a dark painting in the style of {}",
|
|
"a picture in the style of {}",
|
|
"a cool painting in the style of {}",
|
|
"a close-up painting in the style of {}",
|
|
"a bright painting in the style of {}",
|
|
"a cropped painting in the style of {}",
|
|
"a good painting in the style of {}",
|
|
"a close-up painting in the style of {}",
|
|
"a rendition in the style of {}",
|
|
"a nice painting in the style of {}",
|
|
"a small painting in the style of {}",
|
|
"a weird painting in the style of {}",
|
|
"a large painting in the style of {}",
|
|
]
|
|
|
|
|
|
class TextualInversionDataset(Dataset):
|
|
def __init__(
|
|
self,
|
|
data_root,
|
|
tokenizer,
|
|
learnable_property="object", # [object, style]
|
|
size=512,
|
|
repeats=100,
|
|
interpolation="bicubic",
|
|
flip_p=0.5,
|
|
set="train",
|
|
placeholder_token="*",
|
|
center_crop=False,
|
|
):
|
|
self.data_root = Path(data_root)
|
|
self.tokenizer = tokenizer
|
|
self.learnable_property = learnable_property
|
|
self.size = size
|
|
self.placeholder_token = placeholder_token
|
|
self.center_crop = center_crop
|
|
self.flip_p = flip_p
|
|
|
|
self.image_paths = [
|
|
self.data_root / file_path
|
|
for file_path in self.data_root.iterdir()
|
|
if file_path.is_file()
|
|
and file_path.name.endswith((".png", ".PNG", ".jpg", ".JPG", ".jpeg", ".JPEG", ".gif", ".GIF"))
|
|
]
|
|
|
|
self.num_images = len(self.image_paths)
|
|
self._length = self.num_images
|
|
|
|
if set == "train":
|
|
self._length = self.num_images * repeats
|
|
|
|
self.interpolation = {
|
|
"linear": PIL_INTERPOLATION["linear"],
|
|
"bilinear": PIL_INTERPOLATION["bilinear"],
|
|
"bicubic": PIL_INTERPOLATION["bicubic"],
|
|
"lanczos": PIL_INTERPOLATION["lanczos"],
|
|
}[interpolation]
|
|
|
|
self.templates = imagenet_style_templates_small if learnable_property == "style" else imagenet_templates_small
|
|
self.flip_transform = transforms.RandomHorizontalFlip(p=self.flip_p)
|
|
|
|
def __len__(self) -> int:
|
|
return self._length
|
|
|
|
def __getitem__(self, i):
|
|
example = {}
|
|
image = Image.open(self.image_paths[i % self.num_images])
|
|
|
|
if not image.mode == "RGB":
|
|
image = image.convert("RGB")
|
|
|
|
placeholder_string = self.placeholder_token
|
|
text = random.choice(self.templates).format(placeholder_string)
|
|
|
|
example["input_ids"] = self.tokenizer(
|
|
text,
|
|
padding="max_length",
|
|
truncation=True,
|
|
max_length=self.tokenizer.model_max_length,
|
|
return_tensors="pt",
|
|
).input_ids[0]
|
|
|
|
# default to score-sde preprocessing
|
|
img = np.array(image).astype(np.uint8)
|
|
|
|
if self.center_crop:
|
|
crop = min(img.shape[0], img.shape[1])
|
|
(
|
|
h,
|
|
w,
|
|
) = (
|
|
img.shape[0],
|
|
img.shape[1],
|
|
)
|
|
img = img[(h - crop) // 2 : (h + crop) // 2, (w - crop) // 2 : (w + crop) // 2]
|
|
|
|
image = Image.fromarray(img)
|
|
image = image.resize((self.size, self.size), resample=self.interpolation)
|
|
|
|
image = self.flip_transform(image)
|
|
image = np.array(image).astype(np.uint8)
|
|
image = (image / 127.5 - 1.0).astype(np.float32)
|
|
|
|
example["pixel_values"] = torch.from_numpy(image).permute(2, 0, 1)
|
|
return example
|
|
|
|
|
|
def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: Optional[str] = None):
|
|
if token is None:
|
|
token = HfFolder.get_token()
|
|
if organization is None:
|
|
username = whoami(token)["name"]
|
|
return f"{username}/{model_id}"
|
|
else:
|
|
return f"{organization}/{model_id}"
|
|
|
|
|
|
def do_textual_inversion_training(
|
|
config: InvokeAIAppConfig,
|
|
model: str,
|
|
train_data_dir: Path,
|
|
output_dir: Path,
|
|
placeholder_token: str,
|
|
initializer_token: str,
|
|
save_steps: int = 500,
|
|
only_save_embeds: bool = False,
|
|
tokenizer_name: Optional[str] = None,
|
|
learnable_property: str = "object",
|
|
repeats: int = 100,
|
|
seed: Optional[int] = None,
|
|
resolution: int = 512,
|
|
center_crop: bool = False,
|
|
train_batch_size: int = 16,
|
|
num_train_epochs: int = 100,
|
|
max_train_steps: int = 5000,
|
|
gradient_accumulation_steps: int = 1,
|
|
gradient_checkpointing: bool = False,
|
|
learning_rate: float = 1e-4,
|
|
scale_lr: bool = True,
|
|
lr_scheduler: str = "constant",
|
|
lr_warmup_steps: int = 500,
|
|
adam_beta1: float = 0.9,
|
|
adam_beta2: float = 0.999,
|
|
adam_weight_decay: float = 1e-02,
|
|
adam_epsilon: float = 1e-08,
|
|
push_to_hub: bool = False,
|
|
hub_token: Optional[str] = None,
|
|
logging_dir: Path = Path("logs"),
|
|
mixed_precision: str = "fp16",
|
|
allow_tf32: bool = False,
|
|
report_to: str = "tensorboard",
|
|
local_rank: int = -1,
|
|
checkpointing_steps: int = 500,
|
|
resume_from_checkpoint: Optional[Path] = None,
|
|
enable_xformers_memory_efficient_attention: bool = False,
|
|
hub_model_id: Optional[str] = None,
|
|
**kwargs,
|
|
) -> None:
|
|
assert model, "Please specify a base model with --model"
|
|
assert train_data_dir, "Please specify a directory containing the training images using --train_data_dir"
|
|
assert placeholder_token, "Please specify a trigger term using --placeholder_token"
|
|
env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
|
|
if env_local_rank != -1 and env_local_rank != local_rank:
|
|
local_rank = env_local_rank
|
|
|
|
# setting up things the way invokeai expects them
|
|
if not os.path.isabs(output_dir):
|
|
output_dir = config.root_path / output_dir
|
|
|
|
logging_dir = output_dir / logging_dir
|
|
|
|
accelerator_config = ProjectConfiguration()
|
|
accelerator_config.logging_dir = logging_dir
|
|
accelerator = Accelerator(
|
|
gradient_accumulation_steps=gradient_accumulation_steps,
|
|
mixed_precision=mixed_precision,
|
|
log_with=report_to,
|
|
project_config=accelerator_config,
|
|
)
|
|
|
|
# Make one log on every process with the configuration for debugging.
|
|
logging.basicConfig(
|
|
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
|
datefmt="%m/%d/%Y %H:%M:%S",
|
|
level=logging.INFO,
|
|
)
|
|
logger.info(accelerator.state, main_process_only=False)
|
|
if accelerator.is_local_main_process:
|
|
datasets.utils.logging.set_verbosity_warning()
|
|
transformers.utils.logging.set_verbosity_warning()
|
|
diffusers.utils.logging.set_verbosity_info()
|
|
else:
|
|
datasets.utils.logging.set_verbosity_error()
|
|
transformers.utils.logging.set_verbosity_error()
|
|
diffusers.utils.logging.set_verbosity_error()
|
|
|
|
# If passed along, set the training seed now.
|
|
if seed is not None:
|
|
set_seed(seed)
|
|
|
|
# Handle the repository creation
|
|
if accelerator.is_main_process:
|
|
if push_to_hub:
|
|
if hub_model_id is None:
|
|
repo_name = get_full_repo_name(Path(output_dir).name, token=hub_token)
|
|
else:
|
|
repo_name = hub_model_id
|
|
repo = Repository(output_dir, clone_from=repo_name)
|
|
|
|
with open(os.path.join(output_dir, ".gitignore"), "w+") as gitignore:
|
|
if "step_*" not in gitignore:
|
|
gitignore.write("step_*\n")
|
|
if "epoch_*" not in gitignore:
|
|
gitignore.write("epoch_*\n")
|
|
elif output_dir is not None:
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
model_records = initialize_record_store(config)
|
|
base, type, name = model.split("/") # note frontend still returns old-style keys
|
|
try:
|
|
model_config = model_records.search_by_attr(
|
|
model_name=name, model_type=ModelType(type), base_model=BaseModelType(base)
|
|
)[0]
|
|
except IndexError:
|
|
raise Exception(f"Unknown model {model}")
|
|
model_path = config.models_path / model_config.path
|
|
|
|
pipeline_args = {"local_files_only": True}
|
|
if tokenizer_name:
|
|
tokenizer = CLIPTokenizer.from_pretrained(tokenizer_name, **pipeline_args)
|
|
else:
|
|
tokenizer = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer", **pipeline_args)
|
|
|
|
# Load scheduler and models
|
|
noise_scheduler = DDPMScheduler.from_pretrained(model_path, subfolder="scheduler", **pipeline_args)
|
|
text_encoder = CLIPTextModel.from_pretrained(
|
|
model_path,
|
|
subfolder="text_encoder",
|
|
**pipeline_args,
|
|
)
|
|
vae = AutoencoderKL.from_pretrained(
|
|
model_path,
|
|
subfolder="vae",
|
|
**pipeline_args,
|
|
)
|
|
unet = UNet2DConditionModel.from_pretrained(
|
|
model_path,
|
|
subfolder="unet",
|
|
**pipeline_args,
|
|
)
|
|
|
|
# Add the placeholder token in tokenizer
|
|
num_added_tokens = tokenizer.add_tokens(placeholder_token)
|
|
if num_added_tokens == 0:
|
|
raise ValueError(
|
|
f"The tokenizer already contains the token {placeholder_token}. Please pass a different"
|
|
" `placeholder_token` that is not already in the tokenizer."
|
|
)
|
|
|
|
# Convert the initializer_token, placeholder_token to ids
|
|
token_ids = tokenizer.encode(initializer_token, add_special_tokens=False)
|
|
# Check if initializer_token is a single token or a sequence of tokens
|
|
if len(token_ids) > 1:
|
|
raise ValueError(
|
|
f"The initializer token must be a single token. Provided initializer={initializer_token}. Token ids={token_ids}"
|
|
)
|
|
|
|
initializer_token_id = token_ids[0]
|
|
placeholder_token_id = tokenizer.convert_tokens_to_ids(placeholder_token)
|
|
|
|
# Resize the token embeddings as we are adding new special tokens to the tokenizer
|
|
text_encoder.resize_token_embeddings(len(tokenizer))
|
|
|
|
# Initialise the newly added placeholder token with the embeddings of the initializer token
|
|
token_embeds = text_encoder.get_input_embeddings().weight.data
|
|
token_embeds[placeholder_token_id] = token_embeds[initializer_token_id]
|
|
|
|
# Freeze vae and unet
|
|
vae.requires_grad_(False)
|
|
unet.requires_grad_(False)
|
|
# Freeze all parameters except for the token embeddings in text encoder
|
|
text_encoder.text_model.encoder.requires_grad_(False)
|
|
text_encoder.text_model.final_layer_norm.requires_grad_(False)
|
|
text_encoder.text_model.embeddings.position_embedding.requires_grad_(False)
|
|
|
|
if gradient_checkpointing:
|
|
# Keep unet in train mode if we are using gradient checkpointing to save memory.
|
|
# The dropout cannot be != 0 so it doesn't matter if we are in eval or train mode.
|
|
unet.train()
|
|
text_encoder.gradient_checkpointing_enable()
|
|
unet.enable_gradient_checkpointing()
|
|
|
|
if enable_xformers_memory_efficient_attention:
|
|
if is_xformers_available():
|
|
unet.enable_xformers_memory_efficient_attention()
|
|
else:
|
|
raise ValueError("xformers is not available. Make sure it is installed correctly")
|
|
|
|
# Enable TF32 for faster training on Ampere GPUs,
|
|
# cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
|
|
if allow_tf32:
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
|
|
if scale_lr:
|
|
learning_rate = learning_rate * gradient_accumulation_steps * train_batch_size * accelerator.num_processes
|
|
|
|
# Initialize the optimizer
|
|
optimizer = torch.optim.AdamW(
|
|
text_encoder.get_input_embeddings().parameters(), # only optimize the embeddings
|
|
lr=learning_rate,
|
|
betas=(adam_beta1, adam_beta2),
|
|
weight_decay=adam_weight_decay,
|
|
eps=adam_epsilon,
|
|
)
|
|
|
|
# Dataset and DataLoaders creation:
|
|
train_dataset = TextualInversionDataset(
|
|
data_root=train_data_dir,
|
|
tokenizer=tokenizer,
|
|
size=resolution,
|
|
placeholder_token=placeholder_token,
|
|
repeats=repeats,
|
|
learnable_property=learnable_property,
|
|
center_crop=center_crop,
|
|
set="train",
|
|
)
|
|
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)
|
|
|
|
# Scheduler and math around the number of training steps.
|
|
overrode_max_train_steps = False
|
|
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
|
|
if max_train_steps is None:
|
|
max_train_steps = num_train_epochs * num_update_steps_per_epoch
|
|
overrode_max_train_steps = True
|
|
|
|
scheduler = get_scheduler(
|
|
lr_scheduler,
|
|
optimizer=optimizer,
|
|
num_warmup_steps=lr_warmup_steps * gradient_accumulation_steps,
|
|
num_training_steps=max_train_steps * gradient_accumulation_steps,
|
|
)
|
|
|
|
# Prepare everything with our `accelerator`.
|
|
text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
|
|
text_encoder, optimizer, train_dataloader, scheduler
|
|
)
|
|
|
|
# For mixed precision training we cast the unet and vae weights to half-precision
|
|
# as these models are only used for inference, keeping weights in full precision is not required.
|
|
weight_dtype = torch.float32
|
|
if accelerator.mixed_precision == "fp16":
|
|
weight_dtype = torch.float16
|
|
elif accelerator.mixed_precision == "bf16":
|
|
weight_dtype = torch.bfloat16
|
|
|
|
# Move vae and unet to device and cast to weight_dtype
|
|
unet.to(accelerator.device, dtype=weight_dtype)
|
|
vae.to(accelerator.device, dtype=weight_dtype)
|
|
|
|
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
|
|
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / gradient_accumulation_steps)
|
|
if overrode_max_train_steps:
|
|
max_train_steps = num_train_epochs * num_update_steps_per_epoch
|
|
# Afterwards we recalculate our number of training epochs
|
|
num_train_epochs = math.ceil(max_train_steps / num_update_steps_per_epoch)
|
|
|
|
# We need to initialize the trackers we use, and also store our configuration.
|
|
# The trackers initializes automatically on the main process.
|
|
if accelerator.is_main_process:
|
|
params = locals()
|
|
for k in params: # init_trackers() doesn't like objects
|
|
params[k] = str(params[k]) if isinstance(params[k], object) else params[k]
|
|
accelerator.init_trackers("textual_inversion", config=params)
|
|
|
|
# Train!
|
|
total_batch_size = train_batch_size * accelerator.num_processes * gradient_accumulation_steps
|
|
|
|
logger.info("***** Running training *****")
|
|
logger.info(f" Num examples = {len(train_dataset)}")
|
|
logger.info(f" Num Epochs = {num_train_epochs}")
|
|
logger.info(f" Instantaneous batch size per device = {train_batch_size}")
|
|
logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
|
|
logger.info(f" Gradient Accumulation steps = {gradient_accumulation_steps}")
|
|
logger.info(f" Total optimization steps = {max_train_steps}")
|
|
global_step = 0
|
|
first_epoch = 0
|
|
resume_step = None
|
|
|
|
# Potentially load in the weights and states from a previous save
|
|
if resume_from_checkpoint:
|
|
if resume_from_checkpoint != "latest":
|
|
path = os.path.basename(resume_from_checkpoint)
|
|
else:
|
|
# Get the most recent checkpoint
|
|
dirs = os.listdir(output_dir)
|
|
dirs = [d for d in dirs if d.startswith("checkpoint")]
|
|
dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
|
|
path = dirs[-1] if len(dirs) > 0 else None
|
|
|
|
if path is None:
|
|
accelerator.print(f"Checkpoint '{resume_from_checkpoint}' does not exist. Starting a new training run.")
|
|
resume_from_checkpoint = None
|
|
else:
|
|
accelerator.print(f"Resuming from checkpoint {path}")
|
|
accelerator.load_state(os.path.join(output_dir, path))
|
|
global_step = int(path.split("-")[1])
|
|
|
|
resume_global_step = global_step * gradient_accumulation_steps
|
|
first_epoch = global_step // num_update_steps_per_epoch
|
|
resume_step = resume_global_step % (num_update_steps_per_epoch * gradient_accumulation_steps)
|
|
|
|
# Only show the progress bar once on each machine.
|
|
progress_bar = tqdm(
|
|
range(global_step, max_train_steps),
|
|
disable=not accelerator.is_local_main_process,
|
|
)
|
|
progress_bar.set_description("Steps")
|
|
|
|
# keep original embeddings as reference
|
|
orig_embeds_params = accelerator.unwrap_model(text_encoder).get_input_embeddings().weight.data.clone()
|
|
|
|
for epoch in range(first_epoch, num_train_epochs):
|
|
text_encoder.train()
|
|
for step, batch in enumerate(train_dataloader):
|
|
# Skip steps until we reach the resumed step
|
|
if resume_step and resume_from_checkpoint and epoch == first_epoch and step < resume_step:
|
|
if step % gradient_accumulation_steps == 0:
|
|
progress_bar.update(1)
|
|
continue
|
|
|
|
with accelerator.accumulate(text_encoder):
|
|
# Convert images to latent space
|
|
latents = vae.encode(batch["pixel_values"].to(dtype=weight_dtype)).latent_dist.sample().detach()
|
|
latents = latents * 0.18215
|
|
|
|
# Sample noise that we'll add to the latents
|
|
noise = torch.randn_like(latents)
|
|
bsz = latents.shape[0]
|
|
# Sample a random timestep for each image
|
|
timesteps = torch.randint(
|
|
0,
|
|
noise_scheduler.config.num_train_timesteps,
|
|
(bsz,),
|
|
device=latents.device,
|
|
)
|
|
timesteps = timesteps.long()
|
|
|
|
# Add noise to the latents according to the noise magnitude at each timestep
|
|
# (this is the forward diffusion process)
|
|
noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
|
|
|
|
# Get the text embedding for conditioning
|
|
encoder_hidden_states = text_encoder(batch["input_ids"])[0].to(dtype=weight_dtype)
|
|
|
|
# Predict the noise residual
|
|
model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
|
|
|
|
# Get the target for loss depending on the prediction type
|
|
if noise_scheduler.config.prediction_type == "epsilon":
|
|
target = noise
|
|
elif noise_scheduler.config.prediction_type == "v_prediction":
|
|
target = noise_scheduler.get_velocity(latents, noise, timesteps)
|
|
else:
|
|
raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
|
|
|
|
loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
|
|
|
|
accelerator.backward(loss)
|
|
|
|
optimizer.step()
|
|
scheduler.step()
|
|
optimizer.zero_grad()
|
|
|
|
# Let's make sure we don't update any embedding weights besides the newly added token
|
|
index_no_updates = torch.arange(len(tokenizer)) != placeholder_token_id
|
|
with torch.no_grad():
|
|
accelerator.unwrap_model(text_encoder).get_input_embeddings().weight[index_no_updates] = (
|
|
orig_embeds_params[index_no_updates]
|
|
)
|
|
|
|
# Checks if the accelerator has performed an optimization step behind the scenes
|
|
if accelerator.sync_gradients:
|
|
progress_bar.update(1)
|
|
global_step += 1
|
|
if global_step % save_steps == 0:
|
|
save_path = os.path.join(output_dir, f"learned_embeds-steps-{global_step}.bin")
|
|
save_progress(
|
|
text_encoder,
|
|
placeholder_token_id,
|
|
accelerator,
|
|
placeholder_token,
|
|
save_path,
|
|
)
|
|
|
|
if global_step % checkpointing_steps == 0:
|
|
if accelerator.is_main_process:
|
|
save_path = os.path.join(output_dir, f"checkpoint-{global_step}")
|
|
accelerator.save_state(save_path)
|
|
logger.info(f"Saved state to {save_path}")
|
|
|
|
logs = {"loss": loss.detach().item(), "lr": scheduler.get_last_lr()[0]}
|
|
progress_bar.set_postfix(**logs)
|
|
accelerator.log(logs, step=global_step)
|
|
|
|
if global_step >= max_train_steps:
|
|
break
|
|
|
|
# Create the pipeline using using the trained modules and save it.
|
|
accelerator.wait_for_everyone()
|
|
if accelerator.is_main_process:
|
|
if push_to_hub and only_save_embeds:
|
|
logger.warn("Enabling full model saving because --push_to_hub=True was specified.")
|
|
save_full_model = True
|
|
else:
|
|
save_full_model = not only_save_embeds
|
|
if save_full_model:
|
|
pipeline = StableDiffusionPipeline.from_pretrained(
|
|
model_path,
|
|
text_encoder=accelerator.unwrap_model(text_encoder),
|
|
vae=vae,
|
|
unet=unet,
|
|
tokenizer=tokenizer,
|
|
**pipeline_args,
|
|
)
|
|
pipeline.save_pretrained(output_dir)
|
|
# Save the newly trained embeddings
|
|
save_path = os.path.join(output_dir, "learned_embeds.bin")
|
|
save_progress(
|
|
text_encoder,
|
|
placeholder_token_id,
|
|
accelerator,
|
|
placeholder_token,
|
|
save_path,
|
|
)
|
|
|
|
if push_to_hub:
|
|
repo.push_to_hub(commit_message="End of training", blocking=False, auto_lfs_prune=True)
|
|
|
|
accelerator.end_training()
|