import yaml
import json
import torch
import wandb
import torchvision
import numpy as np
from torch import nn
from tqdm import tqdm
from abc import abstractmethod
from fractions import Fraction
import matplotlib.pyplot as plt
from dataclasses import dataclass
from torch.distributed import barrier
from torch.utils.data import DataLoader
from gdf import GDF
from gdf import AdaptiveLossWeight
from core import WarpCore
from core.data import setup_webdataset_path, MultiGetter, MultiFilter, Bucketeer
from core.utils import EXPECTED, EXPECTED_TRAIN, update_weights_ema, create_folder_if_necessary
import webdataset as wds
from webdataset.handlers import warn_and_continue
import transformers
transformers.utils.logging.set_verbosity_error()
class DataCore(WarpCore):
@dataclass(frozen=True)
class Config(WarpCore.Config):
image_size: int = EXPECTED_TRAIN
webdataset_path: str = EXPECTED_TRAIN
grad_accum_steps: int = EXPECTED_TRAIN
batch_size: int = EXPECTED_TRAIN
multi_aspect_ratio: list = None
captions_getter: list = None
dataset_filters: list = None
bucketeer_random_ratio: float = 0.05
@dataclass(frozen=True)
class Extras(WarpCore.Extras):
transforms: torchvision.transforms.Compose = EXPECTED
clip_preprocess: torchvision.transforms.Compose = EXPECTED
@dataclass(frozen=True)
class Models(WarpCore.Models):
tokenizer: nn.Module = EXPECTED
text_model: nn.Module = EXPECTED
image_model: nn.Module = None
config: Config
def webdataset_path(self):
if isinstance(self.config.webdataset_path, str) and (self.config.webdataset_path.strip().startswith(
'pipe:') or self.config.webdataset_path.strip().startswith('file:')):
return self.config.webdataset_path
else:
dataset_path = self.config.webdataset_path
if isinstance(self.config.webdataset_path, str) and self.config.webdataset_path.strip().endswith('.yml'):
with open(self.config.webdataset_path, 'r', encoding='utf-8') as file:
dataset_path = yaml.safe_load(file)
return setup_webdataset_path(dataset_path, cache_path=f"{self.config.experiment_id}_webdataset_cache.yml")
def webdataset_preprocessors(self, extras: Extras):
def identity(x):
if isinstance(x, bytes):
x = x.decode('utf-8')
return x
# CUSTOM CAPTIONS GETTER -----
def get_caption(oc, c, p_og=0.05): # cog_contexual, cog_caption
if p_og > 0 and np.random.rand() < p_og and len(oc) > 0:
return identity(oc)
else:
return identity(c)
captions_getter = MultiGetter(rules={
('old_caption', 'caption'): lambda oc, c: get_caption(json.loads(oc)['og_caption'], c, p_og=0.05)
})
return [
('jpg;png',
torchvision.transforms.ToTensor() if self.config.multi_aspect_ratio is not None else extras.transforms,
'images'),
('txt', identity, 'captions') if self.config.captions_getter is None else (
self.config.captions_getter[0], eval(self.config.captions_getter[1]), 'captions'),
]
def setup_data(self, extras: Extras) -> WarpCore.Data:
# SETUP DATASET
dataset_path = self.webdataset_path()
preprocessors = self.webdataset_preprocessors(extras)
handler = warn_and_continue
dataset = wds.WebDataset(
dataset_path, resampled=True, handler=handler
).select(
MultiFilter(rules={
f[0]: eval(f[1]) for f in self.config.dataset_filters
}) if self.config.dataset_filters is not None else lambda _: True
).shuffle(690, handler=handler).decode(
"pilrgb", handler=handler
).to_tuple(
*[p[0] for p in preprocessors], handler=handler
).map_tuple(
*[p[1] for p in preprocessors], handler=handler
).map(lambda x: {p[2]: x[i] for i, p in enumerate(preprocessors)})
def identity(x):
return x
# SETUP DATALOADER
real_batch_size = self.config.batch_size // (self.world_size * self.config.grad_accum_steps)
dataloader = DataLoader(
dataset, batch_size=real_batch_size, num_workers=8, pin_memory=True,
collate_fn=identity if self.config.multi_aspect_ratio is not None else None
)
if self.is_main_node:
print(f"Training with batch size {self.config.batch_size} ({real_batch_size}/GPU)")
if self.config.multi_aspect_ratio is not None:
aspect_ratios = [float(Fraction(f)) for f in self.config.multi_aspect_ratio]
dataloader_iterator = Bucketeer(dataloader, density=self.config.image_size ** 2, factor=32,
ratios=aspect_ratios, p_random_ratio=self.config.bucketeer_random_ratio,
interpolate_nearest=False) # , use_smartcrop=True)
else:
dataloader_iterator = iter(dataloader)
return self.Data(dataset=dataset, dataloader=dataloader, iterator=dataloader_iterator)
def get_conditions(self, batch: dict, models: Models, extras: Extras, is_eval=False, is_unconditional=False,
eval_image_embeds=False, return_fields=None):
if return_fields is None:
return_fields = ['clip_text', 'clip_text_pooled', 'clip_img']
captions = batch.get('captions', None)
images = batch.get('images', None)
batch_size = len(captions)
text_embeddings = None
text_pooled_embeddings = None
if 'clip_text' in return_fields or 'clip_text_pooled' in return_fields:
if is_eval:
if is_unconditional:
captions_unpooled = ["" for _ in range(batch_size)]
else:
captions_unpooled = captions
else:
rand_idx = np.random.rand(batch_size) > 0.05
captions_unpooled = [str(c) if keep else "" for c, keep in zip(captions, rand_idx)]
clip_tokens_unpooled = models.tokenizer(captions_unpooled, truncation=True, padding="max_length",
max_length=models.tokenizer.model_max_length,
return_tensors="pt").to(self.device)
text_encoder_output = models.text_model(**clip_tokens_unpooled, output_hidden_states=True)
if 'clip_text' in return_fields:
text_embeddings = text_encoder_output.hidden_states[-1]
if 'clip_text_pooled' in return_fields:
text_pooled_embeddings = text_encoder_output.text_embeds.unsqueeze(1)
image_embeddings = None
if 'clip_img' in return_fields:
image_embeddings = torch.zeros(batch_size, 768, device=self.device)
if images is not None:
images = images.to(self.device)
if is_eval:
if not is_unconditional and eval_image_embeds:
image_embeddings = models.image_model(extras.clip_preprocess(images)).image_embeds
else:
rand_idx = np.random.rand(batch_size) > 0.9
if any(rand_idx):
image_embeddings[rand_idx] = models.image_model(extras.clip_preprocess(images[rand_idx])).image_embeds
image_embeddings = image_embeddings.unsqueeze(1)
return {
'clip_text': text_embeddings,
'clip_text_pooled': text_pooled_embeddings,
'clip_img': image_embeddings
}
class TrainingCore(DataCore, WarpCore):
@dataclass(frozen=True)
class Config(DataCore.Config, WarpCore.Config):
updates: int = EXPECTED_TRAIN
backup_every: int = EXPECTED_TRAIN
save_every: int = EXPECTED_TRAIN
# EMA UPDATE
ema_start_iters: int = None
ema_iters: int = None
ema_beta: float = None
use_fsdp: bool = None
@dataclass() # not frozen, means that fields are mutable. Doesn't support EXPECTED
class Info(WarpCore.Info):
ema_loss: float = None
adaptive_loss: dict = None
@dataclass(frozen=True)
class Models(WarpCore.Models):
generator: nn.Module = EXPECTED
generator_ema: nn.Module = None # optional
@dataclass(frozen=True)
class Optimizers(WarpCore.Optimizers):
generator: any = EXPECTED
@dataclass(frozen=True)
class Extras(WarpCore.Extras):
gdf: GDF = EXPECTED
sampling_configs: dict = EXPECTED
info: Info
config: Config
@abstractmethod
def forward_pass(self, data: WarpCore.Data, extras: WarpCore.Extras, models: Models):
raise NotImplementedError("This method needs to be overriden")
@abstractmethod
def backward_pass(self, update, loss, loss_adjusted, models: Models, optimizers: Optimizers,
schedulers: WarpCore.Schedulers):
raise NotImplementedError("This method needs to be overriden")
@abstractmethod
def models_to_save(self) -> list:
raise NotImplementedError("This method needs to be overriden")
@abstractmethod
def encode_latents(self, batch: dict, models: Models, extras: Extras) -> torch.Tensor:
raise NotImplementedError("This method needs to be overriden")
@abstractmethod
def decode_latents(self, latents: torch.Tensor, batch: dict, models: Models, extras: Extras) -> torch.Tensor:
raise NotImplementedError("This method needs to be overriden")
def train(self, data: WarpCore.Data, extras: WarpCore.Extras, models: Models, optimizers: Optimizers,
schedulers: WarpCore.Schedulers):
start_iter = self.info.iter + 1
max_iters = self.config.updates * self.config.grad_accum_steps
if self.is_main_node:
print(f"STARTING AT STEP: {start_iter}/{max_iters}")
pbar = tqdm(range(start_iter, max_iters + 1)) if self.is_main_node else range(start_iter,
max_iters + 1) # <--- DDP
if 'generator' in self.models_to_save():
models.generator.train()
for i in pbar:
# FORWARD PASS
loss, loss_adjusted = self.forward_pass(data, extras, models)
# # BACKWARD PASS
grad_norm = self.backward_pass(
i % self.config.grad_accum_steps == 0 or i == max_iters, loss, loss_adjusted,
models, optimizers, schedulers
)
self.info.iter = i
# UPDATE EMA
if models.generator_ema is not None and i % self.config.ema_iters == 0:
update_weights_ema(
models.generator_ema, models.generator,
beta=(self.config.ema_beta if i > self.config.ema_start_iters else 0)
)
# UPDATE LOSS METRICS
self.info.ema_loss = loss.mean().item() if self.info.ema_loss is None else self.info.ema_loss * 0.99 + loss.mean().item() * 0.01
if self.is_main_node and self.config.wandb_project is not None and np.isnan(loss.mean().item()) or np.isnan(
grad_norm.item()):
wandb.alert(
title=f"NaN value encountered in training run {self.info.wandb_run_id}",
text=f"Loss {loss.mean().item()} - Grad Norm {grad_norm.item()}. Run {self.info.wandb_run_id}",
wait_duration=60 * 30
)
if self.is_main_node:
logs = {
'loss': self.info.ema_loss,
'raw_loss': loss.mean().item(),
'grad_norm': grad_norm.item(),
'lr': optimizers.generator.param_groups[0]['lr'] if optimizers.generator is not None else 0,
'total_steps': self.info.total_steps,
}
pbar.set_postfix(logs)
if self.config.wandb_project is not None:
wandb.log(logs)
if i == 1 or i % (self.config.save_every * self.config.grad_accum_steps) == 0 or i == max_iters:
# SAVE AND CHECKPOINT STUFF
if np.isnan(loss.mean().item()):
if self.is_main_node and self.config.wandb_project is not None:
tqdm.write("Skipping sampling & checkpoint because the loss is NaN")
wandb.alert(title=f"Skipping sampling & checkpoint for training run {self.config.wandb_run_id}",
text=f"Skipping sampling & checkpoint at {self.info.total_steps} for training run {self.info.wandb_run_id} iters because loss is NaN")
else:
if isinstance(extras.gdf.loss_weight, AdaptiveLossWeight):
self.info.adaptive_loss = {
'bucket_ranges': extras.gdf.loss_weight.bucket_ranges.tolist(),
'bucket_losses': extras.gdf.loss_weight.bucket_losses.tolist(),
}
self.save_checkpoints(models, optimizers)
if self.is_main_node:
create_folder_if_necessary(f'{self.config.output_path}/{self.config.experiment_id}/')
self.sample(models, data, extras)
def save_checkpoints(self, models: Models, optimizers: Optimizers, suffix=None):
barrier()
suffix = '' if suffix is None else suffix
self.save_info(self.info, suffix=suffix)
models_dict = models.to_dict()
optimizers_dict = optimizers.to_dict()
for key in self.models_to_save():
model = models_dict[key]
if model is not None:
self.save_model(model, f"{key}{suffix}", is_fsdp=self.config.use_fsdp)
for key in optimizers_dict:
optimizer = optimizers_dict[key]
if optimizer is not None:
self.save_optimizer(optimizer, f'{key}_optim{suffix}',
fsdp_model=models_dict[key] if self.config.use_fsdp else None)
if suffix == '' and self.info.total_steps > 1 and self.info.total_steps % self.config.backup_every == 0:
self.save_checkpoints(models, optimizers, suffix=f"_{self.info.total_steps // 1000}k")
torch.cuda.empty_cache()
def sample(self, models: Models, data: WarpCore.Data, extras: Extras):
if 'generator' in self.models_to_save():
models.generator.eval()
with torch.no_grad():
batch = next(data.iterator)
conditions = self.get_conditions(batch, models, extras, is_eval=True, is_unconditional=False, eval_image_embeds=False)
unconditions = self.get_conditions(batch, models, extras, is_eval=True, is_unconditional=True, eval_image_embeds=False)
latents = self.encode_latents(batch, models, extras)
noised, _, _, logSNR, noise_cond, _ = extras.gdf.diffuse(latents, shift=1, loss_shift=1)
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
pred = models.generator(noised, noise_cond, **conditions)
pred = extras.gdf.undiffuse(noised, logSNR, pred)[0]
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
*_, (sampled, _, _) = extras.gdf.sample(
models.generator, conditions,
latents.shape, unconditions, device=self.device, **extras.sampling_configs
)
if models.generator_ema is not None:
*_, (sampled_ema, _, _) = extras.gdf.sample(
models.generator_ema, conditions,
latents.shape, unconditions, device=self.device, **extras.sampling_configs
)
else:
sampled_ema = sampled
if self.is_main_node:
noised_images = torch.cat(
[self.decode_latents(noised[i:i + 1], batch, models, extras) for i in range(len(noised))], dim=0)
pred_images = torch.cat(
[self.decode_latents(pred[i:i + 1], batch, models, extras) for i in range(len(pred))], dim=0)
sampled_images = torch.cat(
[self.decode_latents(sampled[i:i + 1], batch, models, extras) for i in range(len(sampled))], dim=0)
sampled_images_ema = torch.cat(
[self.decode_latents(sampled_ema[i:i + 1], batch, models, extras) for i in range(len(sampled_ema))],
dim=0)
images = batch['images']
if images.size(-1) != noised_images.size(-1) or images.size(-2) != noised_images.size(-2):
images = nn.functional.interpolate(images, size=noised_images.shape[-2:], mode='bicubic')
collage_img = torch.cat([
torch.cat([i for i in images.cpu()], dim=-1),
torch.cat([i for i in noised_images.cpu()], dim=-1),
torch.cat([i for i in pred_images.cpu()], dim=-1),
torch.cat([i for i in sampled_images.cpu()], dim=-1),
torch.cat([i for i in sampled_images_ema.cpu()], dim=-1),
], dim=-2)
torchvision.utils.save_image(collage_img, f'{self.config.output_path}/{self.config.experiment_id}/{self.info.total_steps:06d}.jpg')
torchvision.utils.save_image(collage_img, f'{self.config.experiment_id}_latest_output.jpg')
captions = batch['captions']
if self.config.wandb_project is not None:
log_data = [
[captions[i]] + [wandb.Image(sampled_images[i])] + [wandb.Image(sampled_images_ema[i])] + [
wandb.Image(images[i])] for i in range(len(images))]
log_table = wandb.Table(data=log_data, columns=["Captions", "Sampled", "Sampled EMA", "Orig"])
wandb.log({"Log": log_table})
if isinstance(extras.gdf.loss_weight, AdaptiveLossWeight):
plt.plot(extras.gdf.loss_weight.bucket_ranges, extras.gdf.loss_weight.bucket_losses[:-1])
plt.ylabel('Raw Loss')
plt.ylabel('LogSNR')
wandb.log({"Loss/LogSRN": plt})
if 'generator' in self.models_to_save():
models.generator.train()