721 lines
38 KiB
Python
721 lines
38 KiB
Python
import os
|
|
import random
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import accelerate
|
|
import diffusers
|
|
import numpy as np
|
|
import pyrallis
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torch.utils.checkpoint
|
|
import transformers
|
|
from accelerate import Accelerator
|
|
# matplotlib.use('Agg') # Set the backend to non-interactive (Agg)
|
|
from accelerate.logging import get_logger
|
|
from accelerate.state import AcceleratorState
|
|
from accelerate.utils import ProjectConfiguration, set_seed
|
|
from diffusers import PriorTransformer, DDPMScheduler
|
|
from diffusers import UNet2DConditionModel, KandinskyV22Pipeline
|
|
from diffusers.loaders import AttnProcsLayers
|
|
from diffusers.models.attention_processor import LoRAAttnProcessor
|
|
from diffusers.training_utils import EMAModel
|
|
from packaging import version
|
|
from torchvision import transforms
|
|
from tqdm import tqdm
|
|
from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection
|
|
from transformers.utils import ContextManagers
|
|
|
|
from model.pipeline_pops import pOpsPipeline
|
|
from model import pops_utils
|
|
from training.dataset import TexturedDataset, InstructDataset, PairsDataset, \
|
|
SceneDataset, UnionDataset, ClothesDataset
|
|
from training.train_config import TrainConfig
|
|
from utils import vis_utils
|
|
|
|
logger = get_logger(__name__, log_level="INFO")
|
|
|
|
|
|
class Coach:
|
|
def __init__(self, config: TrainConfig):
|
|
self.cfg = config
|
|
self.cfg.output_dir.mkdir(exist_ok=True, parents=True)
|
|
(self.cfg.output_dir / 'cfg.yaml').write_text(pyrallis.dump(self.cfg))
|
|
(self.cfg.output_dir / 'run.sh').write_text(f'python {Path(__file__).name} {" ".join(sys.argv)}')
|
|
|
|
self.logging_dir = self.cfg.output_dir / 'logs'
|
|
accelerator_project_config = ProjectConfiguration(
|
|
total_limit=2, project_dir=self.cfg.output_dir, logging_dir=self.logging_dir
|
|
)
|
|
self.accelerator = Accelerator(
|
|
gradient_accumulation_steps=self.cfg.gradient_accumulation_steps,
|
|
mixed_precision=self.cfg.mixed_precision,
|
|
log_with=self.cfg.report_to,
|
|
project_config=accelerator_project_config,
|
|
)
|
|
logger.info(self.accelerator.state, main_process_only=False)
|
|
if self.accelerator.is_local_main_process:
|
|
transformers.utils.logging.set_verbosity_warning()
|
|
diffusers.utils.logging.set_verbosity_info()
|
|
else:
|
|
transformers.utils.logging.set_verbosity_error()
|
|
diffusers.utils.logging.set_verbosity_error()
|
|
|
|
if self.cfg.seed is not None:
|
|
set_seed(self.cfg.seed)
|
|
|
|
if self.accelerator.is_main_process:
|
|
self.logging_dir.mkdir(exist_ok=True, parents=True)
|
|
|
|
self.noise_scheduler = DDPMScheduler(beta_schedule='squaredcos_cap_v2', prediction_type='sample')
|
|
self.image_processor = CLIPImageProcessor.from_pretrained(self.cfg.image_processor_path,
|
|
subfolder='image_processor')
|
|
self.tokenizer = CLIPTokenizer.from_pretrained(self.cfg.tokenizer_path, subfolder='tokenizer')
|
|
|
|
def deepspeed_zero_init_disabled_context_manager():
|
|
"""
|
|
returns either a context list that includes one that will disable zero.Init or an empty context list
|
|
"""
|
|
deepspeed_plugin = AcceleratorState().deepspeed_plugin if accelerate.state.is_initialized() else None
|
|
if deepspeed_plugin is None:
|
|
return []
|
|
|
|
return [deepspeed_plugin.zero3_init_context_manager(enable=False)]
|
|
|
|
self.weight_dtype = torch.float32
|
|
if self.accelerator.mixed_precision == "fp16":
|
|
self.weight_dtype = torch.float16
|
|
elif self.accelerator.mixed_precision == "bf16":
|
|
self.weight_dtype = torch.bfloat16
|
|
with ContextManagers(deepspeed_zero_init_disabled_context_manager()):
|
|
self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(self.cfg.pretrained_image_encoder,
|
|
subfolder='image_encoder',
|
|
torch_dtype=self.weight_dtype).eval()
|
|
self.text_encoder = CLIPTextModelWithProjection.from_pretrained(self.cfg.text_encoder_path,
|
|
subfolder='text_encoder',
|
|
torch_dtype=self.weight_dtype).eval()
|
|
print('args.pretrained_prior_path =', self.cfg.pretrained_prior_path)
|
|
self.prior = PriorTransformer.from_pretrained(
|
|
self.cfg.pretrained_prior_path, subfolder="prior"
|
|
)
|
|
|
|
self.clip_mean = self.prior.clip_mean.clone()
|
|
self.clip_std = self.prior.clip_std.clone()
|
|
|
|
# Freeze text_encoder and image_encoder
|
|
self.text_encoder.requires_grad_(False)
|
|
self.image_encoder.requires_grad_(False)
|
|
|
|
# Load full model for vis
|
|
self.unet = UNet2DConditionModel.from_pretrained('kandinsky-community/kandinsky-2-2-decoder',
|
|
subfolder='unet').to(torch.float16).to(self.accelerator.device)
|
|
self.prior_pipeline = pOpsPipeline.from_pretrained(
|
|
'kandinsky-community/kandinsky-2-2-prior',
|
|
prior=self.prior,
|
|
image_encoder=self.image_encoder,
|
|
torch_dtype=torch.float16)
|
|
self.prior_pipeline = self.prior_pipeline.to(self.accelerator.device)
|
|
self.decoder = KandinskyV22Pipeline.from_pretrained('kandinsky-community/kandinsky-2-2-decoder', unet=self.unet,
|
|
torch_dtype=torch.float16)
|
|
self.decoder = self.decoder.to(self.accelerator.device)
|
|
|
|
if self.cfg.lora_rank is not None:
|
|
self.prior.requires_grad_(False)
|
|
self.prior.to(self.accelerator.device, dtype=self.weight_dtype)
|
|
lora_attn_procs = {}
|
|
for name in self.prior.attn_processors.keys():
|
|
lora_attn_procs[name] = LoRAAttnProcessor(hidden_size=2048, rank=self.cfg.lora_rank)
|
|
|
|
self.prior.set_attn_processor(lora_attn_procs)
|
|
|
|
self.lora_layers = AttnProcsLayers(self.prior.attn_processors)
|
|
|
|
# Create EMA for the prior.
|
|
if self.cfg.use_ema:
|
|
ema_prior = PriorTransformer.from_pretrained(
|
|
self.cfg.pretrained_prior_path, subfolder="prior"
|
|
)
|
|
ema_prior = EMAModel(ema_prior.parameters(), model_cls=PriorTransformer, model_config=ema_prior.config)
|
|
ema_prior.to(self.accelerator.device)
|
|
|
|
# `accelerate` 0.16.0 will have better support for customized saving
|
|
if version.parse(accelerate.__version__) >= version.parse("0.16.0"):
|
|
# create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
|
|
def save_model_hook(models, weights, output_dir):
|
|
if self.cfg.use_ema:
|
|
ema_prior.save_pretrained(os.path.join(output_dir, "prior_ema"))
|
|
|
|
for i, model in enumerate(models):
|
|
model.save_pretrained(os.path.join(output_dir, "prior"))
|
|
|
|
# make sure to pop weight so that corresponding model is not saved again
|
|
weights.pop()
|
|
|
|
def load_model_hook(models, input_dir):
|
|
if self.cfg.use_ema:
|
|
load_model = EMAModel.from_pretrained(os.path.join(input_dir, "prior_ema"), PriorTransformer)
|
|
ema_prior.load_state_dict(load_model.state_dict())
|
|
ema_prior.to(self.accelerator.device)
|
|
del load_model
|
|
|
|
for i in range(len(models)):
|
|
# pop models so that they are not loaded again
|
|
model = models.pop()
|
|
|
|
# load diffusers style into model
|
|
load_model = PriorTransformer.from_pretrained(input_dir, subfolder="prior")
|
|
model.register_to_config(**load_model.config)
|
|
|
|
model.load_state_dict(load_model.state_dict())
|
|
del load_model
|
|
|
|
self.accelerator.register_save_state_pre_hook(save_model_hook)
|
|
self.accelerator.register_load_state_pre_hook(load_model_hook)
|
|
|
|
if self.cfg.allow_tf32:
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
|
|
if self.cfg.use_8bit_adam:
|
|
try:
|
|
import bitsandbytes as bnb
|
|
except ImportError:
|
|
raise ImportError(
|
|
"Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
|
|
)
|
|
|
|
optimizer_cls = bnb.optim.AdamW8bit
|
|
else:
|
|
optimizer_cls = torch.optim.AdamW
|
|
|
|
if self.cfg.lora_rank is not None:
|
|
params_to_optimize = list(self.lora_layers.parameters())
|
|
else:
|
|
if self.cfg.training_mode == 'full':
|
|
params_to_optimize = list(self.prior.parameters())
|
|
elif '_layers' in self.cfg.training_mode:
|
|
layers_to_train = int(self.cfg.training_mode.split('_')[0])
|
|
params_to_optimize = []
|
|
self.prior.norm_out.requires_grad_(False)
|
|
self.prior.proj_to_clip_embeddings.requires_grad_(False)
|
|
self.prior.time_proj.requires_grad_(False)
|
|
self.prior.time_embedding.requires_grad_(False)
|
|
self.prior.transformer_blocks[layers_to_train:].requires_grad_(False)
|
|
self.prior.clip_mean.requires_grad_(False)
|
|
self.prior.clip_std.requires_grad_(False)
|
|
# Populate params_to_optimize with all params that require grad and print its name
|
|
for name, param in self.prior.named_parameters():
|
|
if param.requires_grad:
|
|
print(name)
|
|
params_to_optimize.append(param)
|
|
else:
|
|
raise ValueError(f"Invalid training mode: {self.cfg.training_mode}")
|
|
# Count the number of parameters
|
|
num_params = sum(p.numel() for p in params_to_optimize)
|
|
# Pretty print the number of parameters
|
|
print(f"Number of parameters: {num_params}")
|
|
|
|
self.optimizer = optimizer_cls(
|
|
params_to_optimize,
|
|
lr=self.cfg.lr,
|
|
betas=(self.cfg.adam_beta1, self.cfg.adam_beta2),
|
|
weight_decay=self.cfg.adam_weight_decay,
|
|
eps=self.cfg.adam_epsilon,
|
|
)
|
|
|
|
self.train_dataloader, self.validation_dataloader = self.get_dataloaders()
|
|
|
|
if self.cfg.lora_rank is not None:
|
|
self.lora_layers, self.optimizer, self.train_dataloader = self.accelerator.prepare(
|
|
self.lora_layers, self.optimizer, self.train_dataloader
|
|
)
|
|
else:
|
|
self.prior, self.optimizer, self.train_dataloader = self.accelerator.prepare(
|
|
self.prior, self.optimizer, self.train_dataloader
|
|
)
|
|
|
|
self.image_encoder.to(self.accelerator.device, dtype=self.weight_dtype)
|
|
self.text_encoder.to(self.accelerator.device, dtype=self.weight_dtype)
|
|
self.clip_mean = self.clip_mean.to(self.weight_dtype).to(self.accelerator.device)
|
|
self.clip_std = self.clip_std.to(self.weight_dtype).to(self.accelerator.device)
|
|
|
|
self.train_step = 0 if self.cfg.resume_from_step is None else self.cfg.resume_from_step
|
|
print(self.train_step)
|
|
self.device = "cuda"
|
|
|
|
if self.cfg.resume_from_path is not None:
|
|
prior_state_dict = torch.load(self.cfg.resume_from_path, map_location=self.device)
|
|
msg = self.prior.load_state_dict(prior_state_dict, strict=False)
|
|
print(msg)
|
|
|
|
def save_model(self, save_path):
|
|
save_path.parent.mkdir(exist_ok=True, parents=True)
|
|
prior_state_dict = self.prior.state_dict()
|
|
del prior_state_dict["clip_mean"]
|
|
del prior_state_dict["clip_std"]
|
|
torch.save(prior_state_dict, save_path)
|
|
|
|
if self.cfg.lora_rank is not None:
|
|
lora_state_dict = self.lora_layers.state_dict()
|
|
torch.save(lora_state_dict, save_path.parent / "lora.ckpt")
|
|
|
|
|
|
def unnormalize_and_pil(self, tensor):
|
|
unnormed = tensor * torch.tensor(self.image_processor.image_std).view(3, 1, 1).to(
|
|
tensor.device) + torch.tensor(self.image_processor.image_mean).view(3, 1, 1).to(tensor.device)
|
|
return transforms.ToPILImage()(unnormed)
|
|
|
|
def save_images(self, image, conds, input_embeds, hidden_states, target_embeds, label='', save_path=''):
|
|
self.prior.eval()
|
|
input_images = []
|
|
captions = []
|
|
for i in range(len(conds)):
|
|
pil_image = self.unnormalize_and_pil(conds[i]).resize((self.cfg.img_size, self.cfg.img_size))
|
|
# Check if all zeros
|
|
if self.cfg.mode == 'clothes' and np.all(np.array(pil_image) > 250):
|
|
continue
|
|
input_images.append(pil_image)
|
|
captions.append("Condition")
|
|
if image is not None:
|
|
input_images.append(self.unnormalize_and_pil(image).resize((self.cfg.img_size, self.cfg.img_size)))
|
|
captions.append(f"Target {label}")
|
|
|
|
seeds = range(2)
|
|
output_images = []
|
|
embebds_to_vis = []
|
|
embeds_captions = []
|
|
if self.cfg.mode not in ['hybrid']:
|
|
embebds_to_vis += [target_embeds]
|
|
embeds_captions += ["Target Reconstruct" if image is not None else "Source Reconstruct"]
|
|
if self.cfg.mode not in ['instruct', 'colorization', 'clothes']: # ,'histogram']:
|
|
for alpha in [0.5]:
|
|
embebds_to_vis.append(alpha * hidden_states[:, 0] + (1 - alpha) * hidden_states[:, 1])
|
|
embeds_captions.append(f"{alpha:.2f}A + {1 - alpha:.2f}B")
|
|
for embs in embebds_to_vis:
|
|
embs_unnormed = (embs * self.clip_std) + self.clip_mean
|
|
zero_embeds = self.prior_pipeline.get_zero_embed(embs.shape[0], device=embs.device)
|
|
direct_from_emb = self.decoder(image_embeds=embs_unnormed, negative_image_embeds=zero_embeds,
|
|
num_inference_steps=50, height=512,
|
|
width=512, guidance_scale=4).images
|
|
output_images = output_images + direct_from_emb
|
|
captions += embeds_captions
|
|
|
|
for seed in seeds:
|
|
negative_input_embeds = torch.zeros_like(input_embeds)
|
|
negative_hidden_states = torch.zeros_like(hidden_states)
|
|
for scale in [1, 4]:
|
|
img_emb = self.prior_pipeline(input_embeds=input_embeds, input_hidden_states=hidden_states,
|
|
negative_input_embeds=negative_input_embeds,
|
|
negative_input_hidden_states=negative_hidden_states,
|
|
num_inference_steps=25,
|
|
num_images_per_prompt=1,
|
|
guidance_scale=scale,
|
|
generator=torch.Generator(device="cuda").manual_seed(seed))
|
|
|
|
negative_emb = img_emb.negative_image_embeds
|
|
for seed_2 in range(1):
|
|
images = self.decoder(image_embeds=img_emb.image_embeds, negative_image_embeds=negative_emb,
|
|
num_inference_steps=50, height=512,
|
|
width=512, guidance_scale=4,
|
|
generator=torch.Generator(device="cuda").manual_seed(seed_2)).images
|
|
# torch.Generator(device="cuda").manual_seed(s) for s in range(4)
|
|
output_images += images
|
|
captions.append(f"prior_s {seed}, cfg {scale}, unet_s {seed_2}")
|
|
|
|
all_images = input_images + output_images
|
|
gen_images = vis_utils.create_table_plot(images=all_images, captions=captions)
|
|
gen_images.save(save_path)
|
|
self.prior.train()
|
|
|
|
def get_dataloaders(self) -> torch.utils.data.DataLoader:
|
|
sampler_train = None
|
|
if self.cfg.mode == 'texture':
|
|
dataset_path = self.cfg.dataset_path
|
|
if not isinstance(self.cfg.dataset_path, list):
|
|
dataset_path = [self.cfg.dataset_path]
|
|
datasets = []
|
|
for path in dataset_path:
|
|
datasets.append(TexturedDataset(
|
|
dataset_dir=path,
|
|
randoms_dir=self.cfg.randoms_dir,
|
|
image_processor=self.image_processor,
|
|
tokenizer=self.tokenizer
|
|
))
|
|
dataset = torch.utils.data.ConcatDataset(datasets)
|
|
dataset_weights = []
|
|
for single_dataset in datasets:
|
|
dataset_weights.extend([len(dataset) / len(single_dataset)] * len(single_dataset))
|
|
sampler_train = torch.utils.data.WeightedRandomSampler(weights=dataset_weights,
|
|
num_samples=len(dataset_weights))
|
|
|
|
validation_dataset = TexturedDataset(
|
|
dataset_dir=self.cfg.val_dataset_path,
|
|
multi_dirs=False,
|
|
textures_dir=self.cfg.textures_dir,
|
|
image_processor=self.image_processor,
|
|
tokenizer=self.tokenizer
|
|
)
|
|
elif self.cfg.mode == 'instruct':
|
|
dataset = InstructDataset(
|
|
dataset_dir=self.cfg.dataset_path,
|
|
image_processor=self.image_processor,
|
|
tokenizer=self.tokenizer,
|
|
)
|
|
validation_dataset = InstructDataset(dataset_dir=self.cfg.val_dataset_path,
|
|
image_processor=self.image_processor, tokenizer=self.tokenizer)
|
|
elif self.cfg.mode == 'clothes':
|
|
dataset = ClothesDataset(
|
|
image_processor=self.image_processor,
|
|
tokenizer=self.tokenizer,
|
|
)
|
|
validation_dataset = ClothesDataset(
|
|
image_processor=self.image_processor,
|
|
tokenizer=self.tokenizer,
|
|
val_mode=True
|
|
)
|
|
elif self.cfg.mode == 'union':
|
|
dataset_path = self.cfg.dataset_path
|
|
if not isinstance(self.cfg.dataset_path, list):
|
|
dataset_path = [self.cfg.dataset_path]
|
|
datasets = []
|
|
for path in dataset_path:
|
|
datasets.append(UnionDataset(dataset_dir=path,
|
|
random_dir=self.cfg.randoms_dir,
|
|
image_processor=self.image_processor,
|
|
tokenizer=self.tokenizer))
|
|
dataset = torch.utils.data.ConcatDataset(datasets)
|
|
dataset_weights = []
|
|
for single_dataset in datasets:
|
|
dataset_weights.extend([len(dataset) / len(single_dataset)] * len(single_dataset))
|
|
sampler_train = torch.utils.data.WeightedRandomSampler(weights=dataset_weights,
|
|
num_samples=len(dataset_weights))
|
|
|
|
validation_dataset = PairsDataset(dataset_dir=self.cfg.val_dataset_path,
|
|
image_processor=self.image_processor, tokenizer=self.tokenizer)
|
|
elif self.cfg.mode == 'scene':
|
|
dataset_path = self.cfg.dataset_path
|
|
if not isinstance(self.cfg.dataset_path, list):
|
|
dataset_path = [self.cfg.dataset_path]
|
|
datasets = []
|
|
for path in dataset_path:
|
|
datasets.append(SceneDataset(
|
|
dataset_dir=path,
|
|
background_dir=self.cfg.backgrounds_dir,
|
|
random_dir=self.cfg.randoms_dir,
|
|
image_processor=self.image_processor,
|
|
tokenizer=self.tokenizer
|
|
))
|
|
dataset = torch.utils.data.ConcatDataset(datasets)
|
|
dataset_weights = []
|
|
for single_dataset in datasets:
|
|
dataset_weights.extend([len(dataset) / len(single_dataset)] * len(single_dataset))
|
|
sampler_train = torch.utils.data.WeightedRandomSampler(weights=dataset_weights,
|
|
num_samples=len(dataset_weights))
|
|
validation_dataset = TexturedDataset(
|
|
dataset_dir=self.cfg.val_dataset_path,
|
|
multi_dirs=False,
|
|
textures_dir=self.cfg.textures_dir,
|
|
image_processor=self.image_processor,
|
|
tokenizer=self.tokenizer
|
|
)
|
|
else:
|
|
raise ValueError(f"Invalid mode: {self.cfg.mode}")
|
|
train_dataloader = torch.utils.data.DataLoader(
|
|
dataset, batch_size=self.cfg.train_batch_size, shuffle=sampler_train is None,
|
|
num_workers=self.cfg.num_workers,
|
|
sampler=sampler_train
|
|
)
|
|
|
|
validation_dataloader = torch.utils.data.DataLoader(
|
|
validation_dataset, batch_size=1, shuffle=True, num_workers=self.cfg.num_workers
|
|
)
|
|
return train_dataloader, validation_dataloader
|
|
|
|
def train(self):
|
|
pbar = tqdm(range(self.train_step, self.cfg.max_train_steps + 1))
|
|
while self.train_step < self.cfg.max_train_steps:
|
|
train_loss = 0.0
|
|
self.prior.train()
|
|
|
|
for sample_idx, batch in enumerate(self.train_dataloader):
|
|
# self.optimizer.zero_grad()
|
|
with self.accelerator.accumulate(self.prior):
|
|
image, cond = batch
|
|
|
|
image = image.to(self.weight_dtype).to(self.accelerator.device)
|
|
if 'source_a' in cond:
|
|
cond['source_a'] = cond['source_a'].to(self.weight_dtype).to(self.accelerator.device)
|
|
if 'source_b' in cond:
|
|
cond['source_b'] = cond['source_b'].to(self.weight_dtype).to(self.accelerator.device)
|
|
if 'crops' in cond:
|
|
for crop_ind in range(len(cond['crops'])):
|
|
cond['crops'][crop_ind] = cond['crops'][crop_ind].to(self.weight_dtype).to(
|
|
self.accelerator.device)
|
|
for key in cond.keys():
|
|
if isinstance(cond[key], torch.Tensor):
|
|
cond[key] = cond[key].to(self.accelerator.device)
|
|
|
|
with torch.no_grad():
|
|
text_mask = cond['mask']
|
|
text_encoder_output = self.text_encoder(cond['input_ids'])
|
|
text_encoder_hidden_states = text_encoder_output.last_hidden_state
|
|
# Take only the unmasked tokens
|
|
text_encoder_concat = None
|
|
|
|
if self.cfg.mode == 'instruct':
|
|
text_target = self.text_encoder(cond['condition_ids']).text_embeds
|
|
text_encoder_concat = text_encoder_hidden_states[:, :text_mask.sum().item()]
|
|
|
|
image_embeds = self.image_encoder(image).image_embeds
|
|
|
|
# Sample noise that we'll add to the image_embeds
|
|
noise = torch.randn_like(image_embeds)
|
|
bsz = image_embeds.shape[0]
|
|
# Sample a random timestep for each image
|
|
timesteps = torch.randint(0, self.noise_scheduler.config.num_train_timesteps, (bsz,),
|
|
device=image_embeds.device)
|
|
timesteps = timesteps.long()
|
|
image_embeds = (image_embeds - self.clip_mean) / self.clip_std
|
|
noisy_latents = self.noise_scheduler.add_noise(image_embeds, noise, timesteps)
|
|
|
|
target = image_embeds
|
|
|
|
input_image_embeds, input_hidden_state = self.preprocess_inputs(cond=cond,
|
|
text_encoder_concat=text_encoder_concat)
|
|
|
|
# Create the input sequence for the prior
|
|
loss = 0
|
|
image_feat_seq = torch.zeros_like(text_encoder_hidden_states)
|
|
image_feat_seq[:, :input_hidden_state.shape[1]] = input_hidden_state
|
|
image_txt_mask = torch.zeros_like(text_mask)
|
|
image_txt_mask[:, :input_hidden_state.shape[1]] = 1
|
|
|
|
# Predict the noise residual and compute loss
|
|
model_pred = self.prior(
|
|
noisy_latents,
|
|
timestep=timesteps,
|
|
proj_embedding=input_image_embeds,
|
|
encoder_hidden_states=image_feat_seq,
|
|
attention_mask=image_txt_mask,
|
|
).predicted_image_embedding
|
|
|
|
# Calculate the loss
|
|
if self.cfg.mode == 'instruct':
|
|
prior_loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
|
|
normed_pred = model_pred.float() * self.clip_std.detach() + self.clip_mean.detach()
|
|
normed_pred = F.normalize(normed_pred, p=2, dim=-1)
|
|
normed_text_target = F.normalize(text_target.float(), p=2, dim=-1)
|
|
clip_loss = 1 - (normed_pred * normed_text_target).sum(dim=-1)
|
|
prior_loss = prior_loss + self.cfg.clip_strength * clip_loss
|
|
else:
|
|
prior_loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
|
|
|
|
loss += prior_loss
|
|
# Gather the losses across all processes for logging (if we use distributed training).
|
|
avg_loss = self.accelerator.gather(loss.repeat(self.cfg.train_batch_size)).mean()
|
|
train_loss += avg_loss.item() / self.cfg.gradient_accumulation_steps
|
|
|
|
# Backprop
|
|
self.accelerator.backward(loss)
|
|
if self.accelerator.sync_gradients:
|
|
self.accelerator.clip_grad_norm_(self.prior.parameters(), self.cfg.max_grad_norm)
|
|
self.optimizer.step()
|
|
self.optimizer.zero_grad()
|
|
|
|
# Checks if the accelerator has performed an optimization step behind the scenes
|
|
if self.accelerator.sync_gradients:
|
|
if self.cfg.use_ema:
|
|
self.ema_prior.step(self.prior.parameters())
|
|
pbar.update(1)
|
|
self.train_step += 1
|
|
train_loss = 0.0
|
|
|
|
if self.train_step % self.cfg.checkpointing_steps == 1:
|
|
if self.accelerator.is_main_process:
|
|
save_path = self.cfg.output_dir / f"learned_prior.pth"
|
|
self.save_model(save_path)
|
|
logger.info(f"Saved state to {save_path}")
|
|
pbar.set_postfix(**{"loss": loss.cpu().detach().item()})
|
|
|
|
if self.cfg.log_image_frequency > 0 and (
|
|
self.train_step % self.cfg.log_image_frequency == 1) or self.train_step < 10:
|
|
image_save_path = self.cfg.output_dir / 'images' / f"{self.train_step}_step_images.jpg"
|
|
image_save_path.parent.mkdir(exist_ok=True, parents=True)
|
|
# Apply the full diffusion process
|
|
conds_list = []
|
|
if 'source_a' in cond:
|
|
conds_list.append(cond['source_a'][0])
|
|
if 'source_b' in cond:
|
|
conds_list.append(cond['source_b'][0])
|
|
if 'crops' in cond:
|
|
for crop_ind in range(len(cond['crops'])):
|
|
conds_list.append(cond['crops'][crop_ind][0])
|
|
|
|
label = ''
|
|
if 'text_condition' in cond:
|
|
label += cond['text_condition'][0]
|
|
self.save_images(image=image[0], conds=conds_list,
|
|
label=label,
|
|
input_embeds=input_image_embeds[0].unsqueeze(0),
|
|
hidden_states=input_hidden_state[0].unsqueeze(0),
|
|
target_embeds=image_embeds[0].unsqueeze(0),
|
|
save_path=image_save_path)
|
|
|
|
if self.cfg.log_validation > 0 and (self.train_step % self.cfg.log_validation == 10):
|
|
# Run validation
|
|
self.log_validation()
|
|
|
|
if self.train_step >= self.cfg.max_train_steps:
|
|
break
|
|
|
|
self.train_dataloader, self.validation_dataloader = self.get_dataloaders()
|
|
pbar.close()
|
|
|
|
def preprocess_inputs(self, cond, text_encoder_concat=None):
|
|
|
|
# Specific pre-processing for each mode
|
|
if self.cfg.mode == 'instruct':
|
|
input_image_embeds, input_hidden_state = pops_utils.preprocess(cond['source_a'], None,
|
|
self.image_encoder,
|
|
self.clip_mean.detach(),
|
|
self.clip_std.detach(),
|
|
concat_hidden_states=text_encoder_concat)
|
|
# For CFG
|
|
if random.random() < 0.2:
|
|
input_hidden_state = torch.zeros_like(input_hidden_state)
|
|
for b_ind in range(cond['source_a'].shape[0]):
|
|
cond['source_a'][b_ind] = cond['source_a'][b_ind] * 0
|
|
elif self.cfg.mode == 'clothes':
|
|
input_image_embeds, input_hidden_state = pops_utils.preprocess(None, None,
|
|
self.image_encoder,
|
|
self.clip_mean.detach(),
|
|
self.clip_std.detach(),
|
|
image_list=cond['crops'])
|
|
if random.random() < 0.2:
|
|
input_hidden_state = torch.zeros_like(input_hidden_state)
|
|
# For vis purposes zero all images
|
|
cond['crops'] = [torch.zeros_like(cond['crops'][0])]
|
|
else:
|
|
if self.cfg.mode == 'texture':
|
|
should_drop_cond = []
|
|
for b_ind in range(cond['source_a'].shape[0]):
|
|
# For CFG and random data
|
|
if 'is_random' in cond and cond['is_random'][b_ind]:
|
|
should_drop_a = random.random() < 0.5
|
|
should_drop_b = True
|
|
else:
|
|
# Randomly dropping one of the conditions to make the model use both?
|
|
should_drop_a = random.random() < 0.25
|
|
should_drop_b = random.random() < 0.25
|
|
|
|
if should_drop_a:
|
|
cond['source_a'][b_ind] = cond['source_a'][b_ind] * 0
|
|
if should_drop_b:
|
|
cond['source_b'][b_ind] = cond['source_b'][b_ind] * 0
|
|
should_drop_cond.append((should_drop_a, should_drop_b))
|
|
elif self.cfg.mode == 'union':
|
|
should_drop_cond = []
|
|
for b_ind in range(cond['source_a'].shape[0]):
|
|
should_drop_a, should_drop_b = False, False
|
|
if ('is_random' in cond and cond['is_random'][b_ind]) or random.random() < 0.2:
|
|
should_drop_a = True
|
|
should_drop_b = True
|
|
|
|
if should_drop_a:
|
|
cond['source_a'][b_ind] = cond['source_a'][b_ind] * 0
|
|
if should_drop_b:
|
|
cond['source_b'][b_ind] = cond['source_b'][b_ind] * 0
|
|
should_drop_cond.append((should_drop_a, should_drop_b))
|
|
elif self.cfg.mode == 'scene':
|
|
should_drop_cond = []
|
|
for b_ind in range(cond['source_a'].shape[0]):
|
|
should_drop_a, should_drop_b = False, False
|
|
if ('is_random' in cond and cond['is_random'][b_ind]) or (random.random() < 0.1):
|
|
should_drop_a = True
|
|
should_drop_b = True
|
|
|
|
if should_drop_a:
|
|
cond['source_a'][b_ind] = cond['source_a'][b_ind] * 0
|
|
if should_drop_b:
|
|
cond['source_b'][b_ind] = cond['source_b'][b_ind] * 0
|
|
should_drop_cond.append((should_drop_a, should_drop_b))
|
|
else:
|
|
raise ValueError(f"Invalid mode: {self.cfg.mode}")
|
|
input_image_embeds, input_hidden_state = pops_utils.preprocess(cond['source_a'], cond['source_b'],
|
|
self.image_encoder,
|
|
self.clip_mean.detach(),
|
|
self.clip_std.detach(),
|
|
should_drop_cond=should_drop_cond)
|
|
|
|
return input_image_embeds, input_hidden_state
|
|
|
|
def log_validation(self):
|
|
for sample_idx, batch in tqdm(enumerate(self.validation_dataloader)):
|
|
image, cond = batch
|
|
image = image.to(self.weight_dtype).to(self.accelerator.device)
|
|
if 'source_a' in cond:
|
|
cond['source_a'] = cond['source_a'].to(self.weight_dtype).to(self.accelerator.device)
|
|
if 'source_b' in cond:
|
|
cond['source_b'] = cond['source_b'].to(self.weight_dtype).to(self.accelerator.device)
|
|
if 'crops' in cond:
|
|
for crop_ind in range(len(cond['crops'])):
|
|
cond['crops'][crop_ind] = cond['crops'][crop_ind].to(self.weight_dtype).to(
|
|
self.accelerator.device)
|
|
for key in cond.keys():
|
|
if isinstance(cond[key], torch.Tensor):
|
|
cond[key] = cond[key].to(self.accelerator.device)
|
|
|
|
text_mask = cond['mask']
|
|
|
|
with torch.no_grad():
|
|
target_embeds = self.image_encoder(image).image_embeds
|
|
target_embeds = (target_embeds - self.clip_mean) / self.clip_std
|
|
|
|
text_encoder_output = self.text_encoder(cond['input_ids'])
|
|
text_encoder_hidden_states = text_encoder_output.last_hidden_state
|
|
# Take only the unmasked tokens
|
|
if self.cfg.mode == 'instruct':
|
|
text_encoder_concat = text_encoder_hidden_states[:, :text_mask.sum().item()]
|
|
input_image_embeds, input_hidden_state = pops_utils.preprocess(cond['source_a'],
|
|
None,
|
|
self.image_encoder,
|
|
self.clip_mean.detach(),
|
|
self.clip_std.detach(),
|
|
concat_hidden_states=text_encoder_concat)
|
|
elif self.cfg.mode == 'clothes':
|
|
input_image_embeds, input_hidden_state = pops_utils.preprocess(None, None,
|
|
self.image_encoder,
|
|
self.clip_mean.detach(),
|
|
self.clip_std.detach(),
|
|
image_list=cond['crops'])
|
|
|
|
else:
|
|
input_image_embeds, input_hidden_state = pops_utils.preprocess(cond['source_a'],
|
|
cond['source_b'],
|
|
self.image_encoder,
|
|
self.clip_mean.detach(),
|
|
self.clip_std.detach())
|
|
|
|
image_save_path = self.cfg.output_dir / 'val_images' / f"{self.train_step}_step_{sample_idx}_images.jpg"
|
|
image_save_path.parent.mkdir(exist_ok=True, parents=True)
|
|
|
|
save_target_image = image[0] if self.cfg.textures_dir is None else None
|
|
conds_list = []
|
|
if 'source_a' in cond:
|
|
conds_list.append(cond['source_a'][0])
|
|
if 'source_b' in cond:
|
|
conds_list.append(cond['source_b'][0])
|
|
if 'crops' in cond:
|
|
for crop_ind in range(len(cond['crops'])):
|
|
conds_list.append(cond['crops'][crop_ind][0])
|
|
# Apply the full diffusion process
|
|
self.save_images(image=save_target_image, conds=conds_list,
|
|
label=cond['text_condition'][0] if 'text_condition' in cond else '',
|
|
input_embeds=input_image_embeds[0].unsqueeze(0),
|
|
hidden_states=input_hidden_state[0].unsqueeze(0),
|
|
target_embeds=target_embeds[0].unsqueeze(0),
|
|
save_path=image_save_path)
|
|
|
|
if sample_idx == self.cfg.n_val_images:
|
|
break
|