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()