mirror of
https://github.com/Sygil-Dev/sygil-webui.git
synced 2024-12-15 22:42:14 +03:00
840 lines
38 KiB
Python
840 lines
38 KiB
Python
# This file is part of stable-diffusion-webui (https://github.com/sd-webui/stable-diffusion-webui/).
|
|
|
|
# Copyright 2022 sd-webui team.
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published by
|
|
# the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Affero General Public License for more details.
|
|
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
|
# base webui import and utils.
|
|
from webui_streamlit import st
|
|
from sd_utils import *
|
|
|
|
# streamlit imports
|
|
from streamlit import StopException
|
|
from streamlit_tensorboard import st_tensorboard
|
|
|
|
#other imports
|
|
from transformers import CLIPTextModel, CLIPTokenizer
|
|
|
|
# Temp imports
|
|
|
|
import argparse
|
|
import itertools
|
|
import math
|
|
import os
|
|
import random
|
|
#import datetime
|
|
#from pathlib import Path
|
|
#from typing import Optional
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torch.utils.checkpoint
|
|
from torch.utils.data import Dataset
|
|
|
|
import PIL
|
|
from accelerate import Accelerator
|
|
from accelerate.logging import get_logger
|
|
from accelerate.utils import set_seed
|
|
from diffusers import AutoencoderKL, DDPMScheduler, LMSDiscreteScheduler, StableDiffusionPipeline, UNet2DConditionModel#, PNDMScheduler
|
|
from diffusers.optimization import get_scheduler
|
|
#from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
|
|
from pipelines.stable_diffusion.no_check import NoCheck
|
|
from huggingface_hub import HfFolder, whoami#, Repository
|
|
from PIL import Image
|
|
from torchvision import transforms
|
|
from tqdm.auto import tqdm
|
|
from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer
|
|
from slugify import slugify
|
|
import json
|
|
import os, subprocess
|
|
from io import StringIO
|
|
import sys
|
|
|
|
|
|
# end of imports
|
|
#---------------------------------------------------------------------------------------------------------------
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
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",
|
|
set="train",
|
|
placeholder_token="*",
|
|
center_crop=False,
|
|
templates=None
|
|
):
|
|
|
|
self.data_root = data_root
|
|
self.tokenizer = tokenizer
|
|
self.learnable_property = learnable_property
|
|
self.size = size
|
|
self.placeholder_token = placeholder_token
|
|
self.center_crop = center_crop
|
|
|
|
self.image_paths = [os.path.join(self.data_root, file_path) for file_path in os.listdir(self.data_root) if file_path.lower().endswith(('.png', '.jpg', '.jpeg'))]
|
|
|
|
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.Image.LINEAR,
|
|
"bilinear": PIL.Image.Resampling.BILINEAR,
|
|
"bicubic": PIL.Image.Resampling.BICUBIC,
|
|
"lanczos": PIL.Image.Resampling.LANCZOS,
|
|
}[interpolation]
|
|
|
|
self.templates = templates
|
|
self.cache = {}
|
|
self.tokenized_templates = [self.tokenizer(
|
|
text.format(self.placeholder_token),
|
|
padding="max_length",
|
|
truncation=True,
|
|
max_length=self.tokenizer.model_max_length,
|
|
return_tensors="pt",
|
|
).input_ids[0] for text in self.templates]
|
|
|
|
def __len__(self):
|
|
return self._length
|
|
|
|
def get_example(self, image_path, flipped):
|
|
if image_path in self.cache:
|
|
return self.cache[image_path]
|
|
|
|
example = {}
|
|
image = Image.open(image_path)
|
|
|
|
if not image.mode == "RGB":
|
|
image = image.convert("RGB")
|
|
|
|
# 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 = transforms.RandomHorizontalFlip(p=1 if flipped else 0)(image)
|
|
image = np.array(image).astype(np.uint8)
|
|
image = (image / 127.5 - 1.0).astype(np.float32)
|
|
example["key"] = "-".join([image_path, "-", str(flipped)])
|
|
example["pixel_values"] = torch.from_numpy(image).permute(2, 0, 1)
|
|
|
|
self.cache[image_path] = example
|
|
return example
|
|
|
|
def __getitem__(self, i):
|
|
flipped = random.choice([False, True])
|
|
example = self.get_example(self.image_paths[i % self.num_images], flipped)
|
|
example["input_ids"] = random.choice(self.tokenized_templates)
|
|
return example
|
|
|
|
|
|
def freeze_params(params):
|
|
for param in params:
|
|
param.requires_grad = False
|
|
|
|
|
|
def save_resume_file(basepath, args, extra = {}):
|
|
info = {"args": st.session_state['textual_inversion']["args"]}
|
|
info["args"].update(extra)
|
|
with open(f"{basepath}/resume.json", "w") as f:
|
|
json.dump(info, f, indent=4)
|
|
|
|
class Checkpointer:
|
|
def __init__(
|
|
self,
|
|
accelerator,
|
|
vae,
|
|
unet,
|
|
tokenizer,
|
|
placeholder_token,
|
|
placeholder_token_id,
|
|
templates,
|
|
output_dir,
|
|
random_sample_batches,
|
|
sample_batch_size,
|
|
stable_sample_batches,
|
|
seed
|
|
):
|
|
self.accelerator = accelerator
|
|
self.vae = vae
|
|
self.unet = unet
|
|
self.tokenizer = tokenizer
|
|
self.placeholder_token = placeholder_token
|
|
self.placeholder_token_id = placeholder_token_id
|
|
self.templates = templates
|
|
self.output_dir = output_dir
|
|
self.seed = seed
|
|
self.random_sample_batches = random_sample_batches
|
|
self.sample_batch_size = sample_batch_size
|
|
self.stable_sample_batches = stable_sample_batches
|
|
|
|
@torch.no_grad()
|
|
def checkpoint(self, step, text_encoder, save_samples=True, path=None):
|
|
print("Saving checkpoint for step %d..." % step)
|
|
with torch.autocast("cuda"):
|
|
if path is None:
|
|
checkpoints_path = f"{self.output_dir}/checkpoints"
|
|
os.makedirs(checkpoints_path, exist_ok=True)
|
|
|
|
unwrapped = self.accelerator.unwrap_model(text_encoder)
|
|
|
|
# Save a checkpoint
|
|
learned_embeds = unwrapped.get_input_embeddings().weight[self.placeholder_token_id]
|
|
learned_embeds_dict = {self.placeholder_token: learned_embeds.detach().cpu()}
|
|
|
|
filename = f"%s_%d.bin" % (slugify(self.placeholder_token), step)
|
|
if path is not None:
|
|
torch.save(learned_embeds_dict, path)
|
|
else:
|
|
torch.save(learned_embeds_dict, f"{checkpoints_path}/{filename}")
|
|
torch.save(learned_embeds_dict, f"{checkpoints_path}/last.bin")
|
|
del unwrapped
|
|
del learned_embeds
|
|
|
|
|
|
@torch.no_grad()
|
|
def save_samples(self, step, text_encoder, height, width, guidance_scale, eta, num_inference_steps):
|
|
samples_path = f"{self.output_dir}/samples"
|
|
os.makedirs(samples_path, exist_ok=True)
|
|
|
|
#if "checker" not in server_state['textual_inversion']:
|
|
#with server_state_lock['textual_inversion']["checker"]:
|
|
server_state['textual_inversion']["checker"] = NoCheck()
|
|
|
|
#if "unwrapped" not in server_state['textual_inversion']:
|
|
# with server_state_lock['textual_inversion']["unwrapped"]:
|
|
server_state['textual_inversion']["unwrapped"] = self.accelerator.unwrap_model(text_encoder)
|
|
|
|
#if "pipeline" not in server_state['textual_inversion']:
|
|
# with server_state_lock['textual_inversion']["pipeline"]:
|
|
# Save a sample image
|
|
server_state['textual_inversion']["pipeline"] = StableDiffusionPipeline(
|
|
text_encoder=server_state['textual_inversion']["unwrapped"],
|
|
vae=self.vae,
|
|
unet=self.unet,
|
|
tokenizer=self.tokenizer,
|
|
scheduler=LMSDiscreteScheduler(
|
|
beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
|
|
),
|
|
safety_checker=NoCheck(),
|
|
feature_extractor=CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32"),
|
|
).to("cuda")
|
|
|
|
server_state['textual_inversion']["pipeline"].enable_attention_slicing()
|
|
|
|
if self.stable_sample_batches > 0:
|
|
stable_latents = torch.randn(
|
|
(self.sample_batch_size, server_state['textual_inversion']["pipeline"].unet.in_channels, height // 8, width // 8),
|
|
device=server_state['textual_inversion']["pipeline"].device,
|
|
generator=torch.Generator(device=server_state['textual_inversion']["pipeline"].device).manual_seed(self.seed),
|
|
)
|
|
|
|
stable_prompts = [choice.format(self.placeholder_token) for choice in (self.templates * self.sample_batch_size)[:self.sample_batch_size]]
|
|
|
|
# Generate and save stable samples
|
|
for i in range(0, self.stable_sample_batches):
|
|
samples = server_state['textual_inversion']["pipeline"](
|
|
prompt=stable_prompts,
|
|
height=384,
|
|
latents=stable_latents,
|
|
width=384,
|
|
guidance_scale=guidance_scale,
|
|
eta=eta,
|
|
num_inference_steps=num_inference_steps,
|
|
output_type='pil'
|
|
)["sample"]
|
|
for idx, im in enumerate(samples):
|
|
filename = f"stable_sample_%d_%d_step_%d.png" % (i+1, idx+1, step)
|
|
im.save(f"{samples_path}/{filename}")
|
|
del samples
|
|
del stable_latents
|
|
|
|
prompts = [choice.format(self.placeholder_token) for choice in random.choices(self.templates, k=self.sample_batch_size)]
|
|
# Generate and save random samples
|
|
for i in range(0, self.random_sample_batches):
|
|
samples = server_state['textual_inversion']["pipeline"](
|
|
prompt=prompts,
|
|
height=384,
|
|
width=384,
|
|
guidance_scale=guidance_scale,
|
|
eta=eta,
|
|
num_inference_steps=num_inference_steps,
|
|
output_type='pil'
|
|
)["sample"]
|
|
for idx, im in enumerate(samples):
|
|
filename = f"step_%d_sample_%d_%d.png" % (step, i+1, idx+1)
|
|
im.save(f"{samples_path}/{filename}")
|
|
del samples
|
|
|
|
del server_state['textual_inversion']["checker"]
|
|
del server_state['textual_inversion']["unwrapped"]
|
|
del server_state['textual_inversion']["pipeline"]
|
|
torch.cuda.empty_cache()
|
|
|
|
@retry(RuntimeError, tries=5)
|
|
def textual_inversion():
|
|
print ("Running textual inversion.")
|
|
if "pipeline" in server_state["textual_inversion"]:
|
|
del server_state['textual_inversion']["checker"]
|
|
del server_state['textual_inversion']["unwrapped"]
|
|
del server_state['textual_inversion']["pipeline"]
|
|
|
|
global_step_offset = 0
|
|
if st.session_state['textual_inversion']['args']["resume_from"]:
|
|
basepath = f"{st.session_state['textual_inversion']['args']['resume_from']}"
|
|
print("Resuming state from %s" % st.session_state['textual_inversion']['args']['resume_from'])
|
|
with open(f"{basepath}/resume.json", 'r') as f:
|
|
state = json.load(f)
|
|
global_step_offset = state["args"].get("global_step", 0)
|
|
|
|
print("We've trained %d steps so far" % global_step_offset)
|
|
else:
|
|
basepath = f"{st.session_state['textual_inversion']['args']['output_dir']}/{slugify(st.session_state['textual_inversion']['args']['placeholder_token'])}"
|
|
os.makedirs(basepath, exist_ok=True)
|
|
|
|
|
|
accelerator = Accelerator(
|
|
gradient_accumulation_steps=st.session_state['textual_inversion']['args']['gradient_accumulation_steps'],
|
|
mixed_precision=st.session_state['textual_inversion']['args']['mixed_precision']
|
|
)
|
|
|
|
# If passed along, set the training seed.
|
|
if st.session_state['textual_inversion']['args']['seed']:
|
|
set_seed(st.session_state['textual_inversion']['args']['seed'])
|
|
|
|
#if "tokenizer" not in server_state["textual_inversion"]:
|
|
# Load the tokenizer and add the placeholder token as a additional special token
|
|
#with server_state_lock['textual_inversion']["tokenizer"]:
|
|
if st.session_state['textual_inversion']['args']['tokenizer_name']:
|
|
server_state['textual_inversion']["tokenizer"] = CLIPTokenizer.from_pretrained(st.session_state['textual_inversion']['args']['tokenizer_name'])
|
|
elif st.session_state['textual_inversion']['args']['pretrained_model_name_or_path']:
|
|
server_state['textual_inversion']["tokenizer"] = CLIPTokenizer.from_pretrained(
|
|
st.session_state['textual_inversion']['args']['pretrained_model_name_or_path'] + '/tokenizer'
|
|
)
|
|
|
|
# Add the placeholder token in tokenizer
|
|
num_added_tokens = server_state['textual_inversion']["tokenizer"].add_tokens(st.session_state['textual_inversion']['args']['placeholder_token'])
|
|
if num_added_tokens == 0:
|
|
st.error(
|
|
f"The tokenizer already contains the token {st.session_state['textual_inversion']['args']['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 = server_state['textual_inversion']["tokenizer"].encode(st.session_state['textual_inversion']['args']['initializer_token'], add_special_tokens=False)
|
|
# Check if initializer_token is a single token or a sequence of tokens
|
|
if len(token_ids) > 1:
|
|
st.error("The initializer token must be a single token.")
|
|
|
|
initializer_token_id = token_ids[0]
|
|
placeholder_token_id = server_state['textual_inversion']["tokenizer"].convert_tokens_to_ids(st.session_state['textual_inversion']['args']['placeholder_token'])
|
|
|
|
# Load models and create wrapper for stable diffusion
|
|
text_encoder = CLIPTextModel.from_pretrained(
|
|
st.session_state['textual_inversion']['args']['pretrained_model_name_or_path'] + '/text_encoder',
|
|
)
|
|
vae = AutoencoderKL.from_pretrained(
|
|
st.session_state['textual_inversion']['args']['pretrained_model_name_or_path'] + '/vae',
|
|
)
|
|
unet = UNet2DConditionModel.from_pretrained(
|
|
st.session_state['textual_inversion']['args']['pretrained_model_name_or_path'] + '/unet',
|
|
)
|
|
|
|
base_templates = imagenet_style_templates_small if st.session_state['textual_inversion']['args']['learnable_property'] == "style" else imagenet_templates_small
|
|
if st.session_state['textual_inversion']['args']['custom_templates']:
|
|
templates = st.session_state['textual_inversion']['args']['custom_templates'].split(";")
|
|
else:
|
|
templates = base_templates
|
|
|
|
slice_size = unet.config.attention_head_dim // 2
|
|
unet.set_attention_slice(slice_size)
|
|
|
|
# Resize the token embeddings as we are adding new special tokens to the tokenizer
|
|
text_encoder.resize_token_embeddings(len(server_state['textual_inversion']["tokenizer"]))
|
|
|
|
# Initialise the newly added placeholder token with the embeddings of the initializer token
|
|
token_embeds = text_encoder.get_input_embeddings().weight.data
|
|
|
|
if "resume_checkpoint" in st.session_state['textual_inversion']['args']:
|
|
if st.session_state['textual_inversion']['args']['resume_checkpoint'] is not None:
|
|
token_embeds[placeholder_token_id] = torch.load(st.session_state['textual_inversion']['args']['resume_checkpoint'])[st.session_state['textual_inversion']['args']['placeholder_token']]
|
|
else:
|
|
token_embeds[placeholder_token_id] = token_embeds[initializer_token_id]
|
|
|
|
# Freeze vae and unet
|
|
freeze_params(vae.parameters())
|
|
freeze_params(unet.parameters())
|
|
# Freeze all parameters except for the token embeddings in text encoder
|
|
params_to_freeze = itertools.chain(
|
|
text_encoder.text_model.encoder.parameters(),
|
|
text_encoder.text_model.final_layer_norm.parameters(),
|
|
text_encoder.text_model.embeddings.position_embedding.parameters(),
|
|
)
|
|
freeze_params(params_to_freeze)
|
|
|
|
checkpointer = Checkpointer(
|
|
accelerator=accelerator,
|
|
vae=vae,
|
|
unet=unet,
|
|
tokenizer=server_state['textual_inversion']["tokenizer"],
|
|
placeholder_token=st.session_state['textual_inversion']['args']['placeholder_token'],
|
|
placeholder_token_id=placeholder_token_id,
|
|
templates=templates,
|
|
output_dir=basepath,
|
|
sample_batch_size=st.session_state['textual_inversion']['args']['sample_batch_size'],
|
|
random_sample_batches=st.session_state['textual_inversion']['args']['random_sample_batches'],
|
|
stable_sample_batches=st.session_state['textual_inversion']['args']['stable_sample_batches'],
|
|
seed=st.session_state['textual_inversion']['args']['seed']
|
|
)
|
|
|
|
if st.session_state['textual_inversion']['args']['scale_lr']:
|
|
st.session_state['textual_inversion']['args']['learning_rate'] = (
|
|
st.session_state['textual_inversion']['args']['learning_rate'] * st.session_state['textual_inversion'][
|
|
'args']['gradient_accumulation_steps'] * st.session_state['textual_inversion']['args']['train_batch_size'] * accelerator.num_processes
|
|
)
|
|
|
|
# Initialize the optimizer
|
|
optimizer = torch.optim.AdamW(
|
|
text_encoder.get_input_embeddings().parameters(), # only optimize the embeddings
|
|
lr=st.session_state['textual_inversion']['args']['learning_rate'],
|
|
betas=(st.session_state['textual_inversion']['args']['adam_beta1'], st.session_state['textual_inversion']['args']['adam_beta2']),
|
|
weight_decay=st.session_state['textual_inversion']['args']['adam_weight_decay'],
|
|
eps=st.session_state['textual_inversion']['args']['adam_epsilon'],
|
|
)
|
|
|
|
# TODO (patil-suraj): laod scheduler using args
|
|
noise_scheduler = DDPMScheduler(
|
|
beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000, tensor_format="pt"
|
|
)
|
|
|
|
train_dataset = TextualInversionDataset(
|
|
data_root=st.session_state['textual_inversion']['args']['train_data_dir'],
|
|
tokenizer=server_state['textual_inversion']["tokenizer"],
|
|
size=st.session_state['textual_inversion']['args']['resolution'],
|
|
placeholder_token=st.session_state['textual_inversion']['args']['placeholder_token'],
|
|
repeats=st.session_state['textual_inversion']['args']['repeats'],
|
|
learnable_property=st.session_state['textual_inversion']['args']['learnable_property'],
|
|
center_crop=st.session_state['textual_inversion']['args']['center_crop'],
|
|
set="train",
|
|
templates=templates
|
|
)
|
|
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=st.session_state['textual_inversion']['args']['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) / st.session_state['textual_inversion']['args']['gradient_accumulation_steps'])
|
|
if st.session_state['textual_inversion']['args']['max_train_steps'] is None:
|
|
st.session_state['textual_inversion']['args']['max_train_steps'] = st.session_state['textual_inversion']['args']['num_train_epochs'] * num_update_steps_per_epoch
|
|
overrode_max_train_steps = True
|
|
|
|
lr_scheduler = get_scheduler(
|
|
st.session_state['textual_inversion']['args']['lr_scheduler'],
|
|
optimizer=optimizer,
|
|
num_warmup_steps=st.session_state['textual_inversion']['args']['lr_warmup_steps'] * st.session_state['textual_inversion']['args']['gradient_accumulation_steps'],
|
|
num_training_steps=st.session_state['textual_inversion']['args']['max_train_steps'] * st.session_state['textual_inversion']['args']['gradient_accumulation_steps'],
|
|
)
|
|
|
|
text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
|
|
text_encoder, optimizer, train_dataloader, lr_scheduler
|
|
)
|
|
|
|
# Move vae and unet to device
|
|
vae.to(accelerator.device)
|
|
unet.to(accelerator.device)
|
|
|
|
# Keep vae and unet in eval mode as we don't train these
|
|
vae.eval()
|
|
unet.eval()
|
|
|
|
# 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) / st.session_state['textual_inversion']['args']['gradient_accumulation_steps'])
|
|
if overrode_max_train_steps:
|
|
st.session_state['textual_inversion']['args']['max_train_steps'] = st.session_state['textual_inversion']['args']['num_train_epochs'] * num_update_steps_per_epoch
|
|
# Afterwards we recalculate our number of training epochs
|
|
st.session_state['textual_inversion']['args']['num_train_epochs'] = math.ceil(st.session_state['textual_inversion']['args']['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:
|
|
accelerator.init_trackers("textual_inversion", config=st.session_state['textual_inversion']['args'])
|
|
|
|
# Train!
|
|
total_batch_size = st.session_state['textual_inversion']['args']['train_batch_size'] * accelerator.num_processes * st.session_state[
|
|
'textual_inversion']['args']['gradient_accumulation_steps']
|
|
|
|
logger.info("***** Running training *****")
|
|
logger.info(f" Num examples = {len(train_dataset)}")
|
|
logger.info(f" Num Epochs = {st.session_state['textual_inversion']['args']['num_train_epochs']}")
|
|
logger.info(f" Instantaneous batch size per device = {st.session_state['textual_inversion']['args']['train_batch_size']}")
|
|
logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
|
|
logger.info(f" Gradient Accumulation steps = {st.session_state['textual_inversion']['args']['gradient_accumulation_steps']}")
|
|
logger.info(f" Total optimization steps = {st.session_state['textual_inversion']['args']['max_train_steps']}")
|
|
# Only show the progress bar once on each machine.
|
|
progress_bar = tqdm(range(st.session_state['textual_inversion']['args']['max_train_steps']), disable=not accelerator.is_local_main_process)
|
|
progress_bar.set_description("Steps")
|
|
global_step = 0
|
|
encoded_pixel_values_cache = {}
|
|
|
|
try:
|
|
for epoch in range(st.session_state['textual_inversion']['args']['num_train_epochs']):
|
|
text_encoder.train()
|
|
for step, batch in enumerate(train_dataloader):
|
|
with accelerator.accumulate(text_encoder):
|
|
# Convert images to latent space
|
|
key = "|".join(batch["key"])
|
|
if encoded_pixel_values_cache.get(key, None) is None:
|
|
encoded_pixel_values_cache[key] = vae.encode(batch["pixel_values"]).latent_dist
|
|
latents = encoded_pixel_values_cache[key].sample().detach().half() * 0.18215
|
|
|
|
# Sample noise that we'll add to the latents
|
|
noise = torch.randn(latents.shape).to(latents.device)
|
|
bsz = latents.shape[0]
|
|
# Sample a random timestep for each image
|
|
timesteps = torch.randint(0, noise_scheduler.num_train_timesteps, (bsz,), device=latents.device).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]
|
|
|
|
# Predict the noise residual
|
|
noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
|
|
|
|
loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
|
|
accelerator.backward(loss)
|
|
|
|
# Zero out the gradients for all token embeddings except the newly added
|
|
# embeddings for the concept, as we only want to optimize the concept embeddings
|
|
if accelerator.num_processes > 1:
|
|
grads = text_encoder.module.get_input_embeddings().weight.grad
|
|
else:
|
|
grads = text_encoder.get_input_embeddings().weight.grad
|
|
# Get the index for tokens that we want to zero the grads for
|
|
index_grads_to_zero = torch.arange(len(server_state['textual_inversion']["tokenizer"])) != placeholder_token_id
|
|
grads.data[index_grads_to_zero, :] = grads.data[index_grads_to_zero, :].fill_(0)
|
|
|
|
optimizer.step()
|
|
lr_scheduler.step()
|
|
optimizer.zero_grad()
|
|
|
|
# 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 % st.session_state['textual_inversion']['args']['checkpoint_frequency'] == 0 and global_step > 0 and accelerator.is_main_process:
|
|
checkpointer.checkpoint(global_step + global_step_offset, text_encoder)
|
|
save_resume_file(basepath, st.session_state['textual_inversion']['args'], {
|
|
"global_step": global_step + global_step_offset,
|
|
"resume_checkpoint": f"{basepath}/checkpoints/last.bin"
|
|
})
|
|
checkpointer.save_samples(
|
|
global_step + global_step_offset,
|
|
text_encoder,
|
|
st.session_state['textual_inversion']['args']['resolution'], st.session_state['textual_inversion']['args'][
|
|
'resolution'], 7.5, 0.0, st.session_state['textual_inversion']['args']['sample_steps'])
|
|
|
|
logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
|
|
progress_bar.set_postfix(**logs)
|
|
accelerator.log(logs, step=global_step)
|
|
|
|
if global_step >= st.session_state['textual_inversion']['args']['max_train_steps']:
|
|
break
|
|
|
|
accelerator.wait_for_everyone()
|
|
|
|
# Create the pipeline using using the trained modules and save it.
|
|
if accelerator.is_main_process:
|
|
print("Finished! Saving final checkpoint and resume state.")
|
|
checkpointer.checkpoint(
|
|
global_step + global_step_offset,
|
|
text_encoder,
|
|
path=f"{basepath}/learned_embeds.bin"
|
|
)
|
|
|
|
save_resume_file(basepath, st.session_state['textual_inversion']['args'], {
|
|
"global_step": global_step + global_step_offset,
|
|
"resume_checkpoint": f"{basepath}/checkpoints/last.bin"
|
|
})
|
|
|
|
accelerator.end_training()
|
|
|
|
except KeyboardInterrupt:
|
|
if accelerator.is_main_process:
|
|
print("Interrupted, saving checkpoint and resume state...")
|
|
checkpointer.checkpoint(global_step + global_step_offset, text_encoder)
|
|
save_resume_file(basepath, st.session_state['textual_inversion']['args'], {
|
|
"global_step": global_step + global_step_offset,
|
|
"resume_checkpoint": f"{basepath}/checkpoints/last.bin"
|
|
})
|
|
quit()
|
|
|
|
|
|
def layout():
|
|
|
|
with st.form("textual-inversion"):
|
|
#st.info("Under Construction. :construction_worker:")
|
|
#parser = argparse.ArgumentParser(description="Simple example of a training script.")
|
|
|
|
set_page_title("Textual Inversion - Stable Diffusion Playground")
|
|
|
|
config_tab, output_tab = st.tabs(["Textual Inversion Config", "Ouput"])
|
|
|
|
with config_tab:
|
|
col1, col2, col3, col4, col5 = st.columns(5, gap='large')
|
|
|
|
if "textual_inversion" not in st.session_state:
|
|
st.session_state['textual_inversion'] = {}
|
|
|
|
if "textual_inversion" not in server_state:
|
|
server_state["textual_inversion"] = {}
|
|
|
|
if "args" not in st.session_state['textual_inversion']:
|
|
st.session_state['textual_inversion']["args"] = {}
|
|
|
|
|
|
with col1:
|
|
st.session_state['textual_inversion']["args"]["pretrained_model_name_or_path"] = st.text_input("Pretrained Model Path",
|
|
value=st.session_state["defaults"].textual_inversion.pretrained_model_name_or_path,
|
|
help="Path to pretrained model or model identifier from huggingface.co/models.")
|
|
|
|
st.session_state['textual_inversion']["args"]["tokenizer_name"] = st.text_input("Tokenizer Name",
|
|
value=st.session_state["defaults"].textual_inversion.tokenizer_name,
|
|
help="Pretrained tokenizer name or path if not the same as model_name")
|
|
|
|
st.session_state['textual_inversion']["args"]["train_data_dir"] = st.text_input("train_data_dir", value="", help="A folder containing the training data.")
|
|
|
|
st.session_state['textual_inversion']["args"]["placeholder_token"] = st.text_input("Placeholder Token", value="", help="A token to use as a placeholder for the concept.")
|
|
|
|
st.session_state['textual_inversion']["args"]["initializer_token"] = st.text_input("Initializer Token", value="", help="A token to use as initializer word.")
|
|
|
|
st.session_state['textual_inversion']["args"]["learnable_property"] = st.selectbox("Learnable Property", ["object", "style"], index=0, help="Choose between 'object' and 'style'")
|
|
|
|
st.session_state['textual_inversion']["args"]["repeats"] = int(st.text_input("Number of times to Repeat", value=100, help="How many times to repeat the training data."))
|
|
|
|
with col2:
|
|
st.session_state['textual_inversion']["args"]["output_dir"] = st.text_input("Output Directory",
|
|
value=str(os.path.join("outputs", "textual_inversion")),
|
|
help="The output directory where the model predictions and checkpoints will be written.")
|
|
|
|
st.session_state['textual_inversion']["args"]["seed"] = seed_to_int(st.text_input("Seed", help="A seed for reproducible training."))
|
|
|
|
st.session_state['textual_inversion']["args"]["resolution"] = int(st.text_input("Resolution", value=512,
|
|
help="The resolution for input images, all the images in the train/validation dataset will be resized to this resolution"))
|
|
|
|
st.session_state['textual_inversion']["args"]["center_crop"] = st.checkbox("Center Image", value=True, help="Whether to center crop images before resizing to resolution")
|
|
|
|
st.session_state['textual_inversion']["args"]["train_batch_size"] = int(st.text_input("Train Batch Size", value=1, help="Batch size (per device) for the training dataloader."))
|
|
|
|
st.session_state['textual_inversion']["args"]["num_train_epochs"] = int(st.text_input("Number of Steps to Train", value=100, help="Number of steps to train."))
|
|
|
|
st.session_state['textual_inversion']["args"]["max_train_steps"] = int(st.text_input("Max Number of Steps to Train", value=5000,
|
|
help="Total number of training steps to perform. If provided, overrides 'Number of Steps to Train'."))
|
|
|
|
with col3:
|
|
st.session_state['textual_inversion']["args"]["gradient_accumulation_steps"] = int(st.text_input("Gradient Accumulation Steps", value=1,
|
|
help="Number of updates steps to accumulate before performing a backward/update pass."))
|
|
|
|
st.session_state['textual_inversion']["args"]["learning_rate"] = float(st.text_input("Learning Rate", value=5.0e-04,
|
|
help="Initial learning rate (after the potential warmup period) to use."))
|
|
|
|
st.session_state['textual_inversion']["args"]["scale_lr"] = st.checkbox("Scale Learning Rate", value=True,
|
|
help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.")
|
|
|
|
st.session_state['textual_inversion']["args"]["lr_scheduler"] = st.text_input("Learning Rate Scheduler", value="constant",
|
|
help=("The scheduler type to use. Choose between ['linear', 'cosine', 'cosine_with_restarts', 'polynomial',"
|
|
" 'constant', 'constant_with_warmup']" ))
|
|
|
|
st.session_state['textual_inversion']["args"]["lr_warmup_steps"] = int(st.text_input("Learning Rate Warmup Steps", value=500, help="Number of steps for the warmup in the lr scheduler."))
|
|
|
|
st.session_state['textual_inversion']["args"]["adam_beta1"] = float(st.text_input("Adam Beta 1", value=0.9, help="The beta1 parameter for the Adam optimizer."))
|
|
|
|
st.session_state['textual_inversion']["args"]["adam_beta2"] = float(st.text_input("Adam Beta 2", value=0.999, help="The beta2 parameter for the Adam optimizer."))
|
|
|
|
st.session_state['textual_inversion']["args"]["adam_weight_decay"] = float(st.text_input("Adam Weight Decay", value=1e-2, help="Weight decay to use."))
|
|
|
|
st.session_state['textual_inversion']["args"]["adam_epsilon"] = float(st.text_input("Adam Epsilon", value=1e-08, help="Epsilon value for the Adam optimizer"))
|
|
|
|
with col4:
|
|
st.session_state['textual_inversion']["args"]["mixed_precision"] = st.selectbox("Mixed Precision", ["no", "fp16", "bf16"], index=1,
|
|
help="Whether to use mixed precision. Choose" "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
|
|
"and an Nvidia Ampere GPU.")
|
|
|
|
st.session_state['textual_inversion']["args"]["local_rank"] = int(st.text_input("Local Rank", value=1, help="For distributed training: local_rank"))
|
|
|
|
st.session_state['textual_inversion']["args"]["checkpoint_frequency"] = int(st.text_input("Checkpoint Frequency", value=500, help="How often to save a checkpoint and sample image"))
|
|
|
|
# stable_sample_batches is crashing when saving the samples so for now I will disable it util its fixed.
|
|
#st.session_state['textual_inversion']["args"]["stable_sample_batches"] = int(st.text_input("Stable Sample Batches", value=0,
|
|
#help="Number of fixed seed sample batches to generate per checkpoint"))
|
|
|
|
st.session_state['textual_inversion']["args"]["stable_sample_batches"] = 0
|
|
|
|
st.session_state['textual_inversion']["args"]["random_sample_batches"] = int(st.text_input("Random Sample Batches", value=2,
|
|
help="Number of random seed sample batches to generate per checkpoint"))
|
|
|
|
st.session_state['textual_inversion']["args"]["sample_batch_size"] = int(st.text_input("Sample Batch Size", value=1, help="Number of samples to generate per batch"))
|
|
|
|
st.session_state['textual_inversion']["args"]["sample_steps"] = int(st.text_input("Sample Steps", value=100,
|
|
help="Number of steps for sample generation. Higher values will result in more detailed samples, but longer runtimes."))
|
|
|
|
st.session_state['textual_inversion']["args"]["custom_templates"] = st.text_input("Custom Templates", value="",
|
|
help="A semicolon-delimited list of custom template to use for samples, using {} as a placeholder for the concept.")
|
|
with col5:
|
|
st.session_state['textual_inversion']["args"]["resume_from"] = st.text_input(label="Resume From", help="Path to a directory to resume training from (ie, logs/token_name)")
|
|
|
|
#st.session_state['textual_inversion']["args"]["resume_checkpoint"] = st.file_uploader("Resume Checkpoint", type=["bin"],
|
|
#help="Path to a specific checkpoint to resume training from (ie, logs/token_name/checkpoints/something.bin).")
|
|
|
|
#st.session_state['textual_inversion']["args"]["config"] = st.file_uploader("Config File", type=["json"],
|
|
#help="Path to a JSON configuration file containing arguments for invoking this script."
|
|
#"If resume_from is given, its resume.json takes priority over this.")
|
|
#
|
|
if "resume_from" in st.session_state['textual_inversion']["args"]:
|
|
if st.session_state['textual_inversion']["args"]["resume_from"]:
|
|
if os.path.exists(os.path.join(st.session_state['textual_inversion']['args']['resume_from'], "resume.json")):
|
|
with open(os.path.join(st.session_state['textual_inversion']['args']['resume_from'], "resume.json"), 'rt') as f:
|
|
st.session_state['textual_inversion']["args"] = json.load(f)["args"]
|
|
#print(st.session_state['textual_inversion']["args"])
|
|
|
|
#elif st.session_state['textual_inversion']["args"]["config"] is not None:
|
|
#with open(st.session_state['textual_inversion']["args"]["config"], 'rt') as f:
|
|
#args = parser.parse_args(namespace=argparse.Namespace(**json.load(f)["args"]))
|
|
|
|
env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
|
|
if env_local_rank != -1 and env_local_rank != st.session_state['textual_inversion']["args"]["local_rank"]:
|
|
st.session_state['textual_inversion']["args"]["local_rank"] = env_local_rank
|
|
|
|
if st.session_state['textual_inversion']["args"]["train_data_dir"] is None:
|
|
st.error("You must specify --train_data_dir")
|
|
|
|
if st.session_state['textual_inversion']["args"]["pretrained_model_name_or_path"] is None:
|
|
st.error("You must specify --pretrained_model_name_or_path")
|
|
|
|
if st.session_state['textual_inversion']["args"]["placeholder_token"] is None:
|
|
st.error("You must specify --placeholder_token")
|
|
|
|
if st.session_state['textual_inversion']["args"]["initializer_token"] is None:
|
|
st.error("You must specify --initializer_token")
|
|
|
|
if st.session_state['textual_inversion']["args"]["output_dir"] is None:
|
|
st.error("You must specify --output_dir")
|
|
|
|
# add a spacer and the submit button for the form.
|
|
|
|
st.session_state["textual_inversion"]["message"] = st.empty()
|
|
st.session_state["textual_inversion"]["progress_bar"] = st.empty()
|
|
|
|
st.write("---")
|
|
|
|
submit = st.form_submit_button(help="")
|
|
if submit:
|
|
if "pipe" in st.session_state:
|
|
del st.session_state["pipe"]
|
|
if "model" in st.session_state:
|
|
del st.session_state["model"]
|
|
|
|
st.session_state["textual_inversion"]["message"].info("Textual Inversion Running. For more info check the progress on your console or the Ouput Tab.")
|
|
|
|
try:
|
|
try:
|
|
textual_inversion()
|
|
except RuntimeError:
|
|
if "pipeline" in server_state["textual_inversion"]:
|
|
del server_state['textual_inversion']["checker"]
|
|
del server_state['textual_inversion']["unwrapped"]
|
|
del server_state['textual_inversion']["pipeline"]
|
|
|
|
textual_inversion()
|
|
|
|
except StopException:
|
|
print(f"Received Streamlit StopException")
|
|
|
|
st.session_state["textual_inversion"]["message"].empty()
|
|
|
|
with output_tab:
|
|
#st.info("Under Construction. :construction_worker:")
|
|
|
|
# Start TensorBoard
|
|
st_tensorboard(logdir=os.path.join("outputs", "textual_inversion"), port=8888)
|
|
|
|
|