# 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 . # 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)