base.py

import os
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, dtype=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, dtype=torch.float32 if dtype == None else dtype)
            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
        optimizer_type: str = None
        stochastic_rounding: 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):
        if 'SLURM_LOCALID' in os.environ:
            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()