train_c_controlnet.py

import torch
import torchvision
from torch import nn, optim
from transformers import AutoTokenizer, CLIPTextModelWithProjection, CLIPVisionModelWithProjection
from warmup_scheduler import GradualWarmupScheduler

import sys
import os
import wandb
from dataclasses import dataclass

from gdf import GDF, EpsilonTarget, CosineSchedule
from gdf import VPScaler, CosineTNoiseCond, DDPMSampler, P2LossWeight
from torchtools.transforms import SmartCrop

from modules import EfficientNetEncoder
from modules import StageC
from modules import ResBlock, AttnBlock, TimestepBlock, FeedForwardBlock
from modules import Previewer
from modules import ControlNet, ControlNetDeliverer
from modules import controlnet_filters

from train.base import DataCore, TrainingCore

from core import WarpCore
from core.utils import EXPECTED, EXPECTED_TRAIN, load_or_fail

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy
from torch.distributed.fsdp.wrap import ModuleWrapPolicy
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
import functools
from accelerate import init_empty_weights
from accelerate.utils import set_module_tensor_to_device
from contextlib import contextmanager


class WurstCore(TrainingCore, DataCore, WarpCore):
    @dataclass(frozen=True)
    class Config(TrainingCore.Config, DataCore.Config, WarpCore.Config):
        # TRAINING PARAMS
        lr: float = EXPECTED_TRAIN
        warmup_updates: int = EXPECTED_TRAIN
        offset_noise: float = None
        dtype: str = None

        # MODEL VERSION
        model_version: str = EXPECTED  # 3.6B or 1B
        clip_image_model_name: str = 'openai/clip-vit-large-patch14'
        clip_text_model_name: str = 'laion/CLIP-ViT-bigG-14-laion2B-39B-b160k'

        # CHECKPOINT PATHS
        effnet_checkpoint_path: str = EXPECTED
        previewer_checkpoint_path: str = EXPECTED
        generator_checkpoint_path: str = None
        controlnet_checkpoint_path: str = None

        # controlnet settings
        controlnet_blocks: list = EXPECTED
        controlnet_filter: str = EXPECTED
        controlnet_filter_params: dict = None
        controlnet_bottleneck_mode: str = None

    @dataclass(frozen=True)
    class Models(TrainingCore.Models, DataCore.Models, WarpCore.Models):
        effnet: nn.Module = EXPECTED
        previewer: nn.Module = EXPECTED
        controlnet: nn.Module = EXPECTED

    @dataclass(frozen=True)
    class Schedulers(WarpCore.Schedulers):
        controlnet: any = None

    @dataclass(frozen=True)
    class Extras(TrainingCore.Extras, DataCore.Extras, WarpCore.Extras):
        gdf: GDF = EXPECTED
        sampling_configs: dict = EXPECTED
        effnet_preprocess: torchvision.transforms.Compose = EXPECTED
        controlnet_filter: controlnet_filters.BaseFilter = EXPECTED

    # @dataclass() # not frozen, means that fields are mutable. Doesn't support EXPECTED
    # class Info(WarpCore.Info):
    #     ema_loss: float = None

    @dataclass(frozen=True)
    class Optimizers(TrainingCore.Optimizers, WarpCore.Optimizers):
        generator: any = None
        controlnet: any = EXPECTED

    info: TrainingCore.Info
    config: Config

    def setup_extras_pre(self) -> Extras:
        gdf = GDF(
            schedule=CosineSchedule(clamp_range=[0.0001, 0.9999]),
            input_scaler=VPScaler(), target=EpsilonTarget(),
            noise_cond=CosineTNoiseCond(),
            loss_weight=P2LossWeight(),
            offset_noise=self.config.offset_noise if self.config.offset_noise is not None else 0.0
        )
        sampling_configs = {"cfg": 5, "sampler": DDPMSampler(gdf), "shift": 1, "timesteps": 20}

        effnet_preprocess = torchvision.transforms.Compose([
            torchvision.transforms.Normalize(
                mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)
            )
        ])

        clip_preprocess = torchvision.transforms.Compose([
            torchvision.transforms.Resize(224, interpolation=torchvision.transforms.InterpolationMode.BICUBIC),
            torchvision.transforms.CenterCrop(224),
            torchvision.transforms.Normalize(
                mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711)
            )
        ])

        if self.config.training:
            transforms = torchvision.transforms.Compose([
                torchvision.transforms.ToTensor(),
                torchvision.transforms.Resize(self.config.image_size, interpolation=torchvision.transforms.InterpolationMode.BILINEAR, antialias=True),
                SmartCrop(self.config.image_size, randomize_p=0.3, randomize_q=0.2)
            ])
        else:
            transforms = None

        controlnet_filter = getattr(controlnet_filters, self.config.controlnet_filter)(
            self.device,
            **(self.config.controlnet_filter_params if self.config.controlnet_filter_params is not None else {})
        )

        return self.Extras(
            gdf=gdf,
            sampling_configs=sampling_configs,
            transforms=transforms,
            effnet_preprocess=effnet_preprocess,
            clip_preprocess=clip_preprocess,
            controlnet_filter=controlnet_filter
        )

    def get_cnet(self, batch: dict, models: Models, extras: Extras, cnet_input=None, **kwargs):
        images = batch['images']
        with torch.no_grad():
            if cnet_input is None:
                cnet_input = extras.controlnet_filter(images, **kwargs)
            if isinstance(cnet_input, tuple):
                cnet_input, cnet_input_preview = cnet_input
            else:
                cnet_input_preview = cnet_input
            cnet_input, cnet_input_preview = cnet_input.to(self.device), cnet_input_preview.to(self.device)
        cnet = models.controlnet(cnet_input)
        return cnet, cnet_input_preview

    def get_conditions(self, batch: dict, models: Models, extras: Extras, is_eval=False, is_unconditional=False,
                       eval_image_embeds=False, return_fields=None):
        with torch.no_grad():
            conditions = super().get_conditions(
                batch, models, extras, is_eval, is_unconditional,
                eval_image_embeds, return_fields=return_fields or ['clip_text', 'clip_text_pooled', 'clip_img']
            )
        return conditions

    def setup_models(self, extras: Extras) -> Models:
        dtype = getattr(torch, self.config.dtype) if self.config.dtype else torch.float32

        # EfficientNet encoder
        effnet = EfficientNetEncoder().to(self.device)
        effnet_checkpoint = load_or_fail(self.config.effnet_checkpoint_path)
        effnet.load_state_dict(effnet_checkpoint if 'state_dict' not in effnet_checkpoint else effnet_checkpoint['state_dict'])
        effnet.eval().requires_grad_(False)
        del effnet_checkpoint

        # Previewer
        previewer = Previewer().to(self.device)
        previewer_checkpoint = load_or_fail(self.config.previewer_checkpoint_path)
        previewer.load_state_dict(previewer_checkpoint if 'state_dict' not in previewer_checkpoint else previewer_checkpoint['state_dict'])
        previewer.eval().requires_grad_(False)
        del previewer_checkpoint

        @contextmanager
        def dummy_context():
            yield None

        loading_context = dummy_context if self.config.training else init_empty_weights

        with loading_context():
            # Diffusion models
            if self.config.model_version == '3.6B':
                generator = StageC()
            elif self.config.model_version == '1B':
                generator = StageC(c_cond=1536, c_hidden=[1536, 1536], nhead=[24, 24], blocks=[[4, 12], [12, 4]])
            else:
                raise ValueError(f"Unknown model version {self.config.model_version}")
                
        if self.config.generator_checkpoint_path is not None:
            if loading_context is dummy_context:
                generator.load_state_dict(load_or_fail(self.config.generator_checkpoint_path))
            else:
                for param_name, param in load_or_fail(self.config.generator_checkpoint_path).items():
                    set_module_tensor_to_device(generator, param_name, "cpu", value=param)
        generator = generator.to(dtype).to(self.device)
        generator = self.load_model(generator, 'generator')

        # if self.config.use_fsdp:
        #     fsdp_auto_wrap_policy = ModuleWrapPolicy([ResBlock, AttnBlock, TimestepBlock, FeedForwardBlock])
        #     generator = FSDP(generator, **self.fsdp_defaults, auto_wrap_policy=fsdp_auto_wrap_policy, device_id=self.device)

        # CLIP encoders
        tokenizer = AutoTokenizer.from_pretrained(self.config.clip_text_model_name)
        text_model = CLIPTextModelWithProjection.from_pretrained(self.config.clip_text_model_name).requires_grad_(False).to(dtype).to(self.device)
        image_model = CLIPVisionModelWithProjection.from_pretrained(self.config.clip_image_model_name).requires_grad_(False).to(dtype).to(self.device)

        # ControlNet
        controlnet = ControlNet(
            c_in=extras.controlnet_filter.num_channels(),
            proj_blocks=self.config.controlnet_blocks,
            bottleneck_mode=self.config.controlnet_bottleneck_mode
        )

        if self.config.controlnet_checkpoint_path is not None:
            controlnet_checkpoint = load_or_fail(self.config.controlnet_checkpoint_path)
            controlnet.load_state_dict(controlnet_checkpoint if 'state_dict' not in controlnet_checkpoint else controlnet_checkpoint['state_dict'])
        controlnet = controlnet.to(dtype).to(self.device)

        controlnet = self.load_model(controlnet, 'controlnet')
        controlnet.backbone.eval().requires_grad_(True)

        if self.config.use_fsdp:
            fsdp_auto_wrap_policy = functools.partial(size_based_auto_wrap_policy, min_num_params=3000)
            controlnet = FSDP(controlnet, **self.fsdp_defaults, auto_wrap_policy=fsdp_auto_wrap_policy, device_id=self.device)

        return self.Models(
            effnet=effnet, previewer=previewer,
            generator=generator, generator_ema=None,
            controlnet=controlnet,
            tokenizer=tokenizer, text_model=text_model, image_model=image_model
        )

    def setup_optimizers(self, extras: Extras, models: Models) -> Optimizers:
        optimizer = optim.AdamW(models.controlnet.parameters(), lr=self.config.lr)  # , eps=1e-7, betas=(0.9, 0.95))
        optimizer = self.load_optimizer(optimizer, 'controlnet_optim',
                                        fsdp_model=models.controlnet if self.config.use_fsdp else None)
        return self.Optimizers(generator=None, controlnet=optimizer)

    def setup_schedulers(self, extras: Extras, models: Models, optimizers: Optimizers) -> Schedulers:
        scheduler = GradualWarmupScheduler(optimizers.controlnet, multiplier=1, total_epoch=self.config.warmup_updates)
        scheduler.last_epoch = self.info.total_steps
        return self.Schedulers(controlnet=scheduler)

    def forward_pass(self, data: WarpCore.Data, extras: Extras, models: Models):
        batch = next(data.iterator)

        cnet, _ = self.get_cnet(batch, models, extras)
        conditions = {**self.get_conditions(batch, models, extras), 'cnet': cnet}
        with torch.no_grad():
            latents = self.encode_latents(batch, models, extras)
            noised, noise, target, logSNR, noise_cond, loss_weight = extras.gdf.diffuse(latents, shift=1, loss_shift=1)

        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
            pred = models.generator(noised, noise_cond, **conditions)
            loss = nn.functional.mse_loss(pred, target, reduction='none').mean(dim=[1, 2, 3])
            loss_adjusted = (loss * loss_weight).mean() / self.config.grad_accum_steps

        return loss, loss_adjusted

    def backward_pass(self, update, loss, loss_adjusted, models: Models, optimizers: Optimizers,
                      schedulers: Schedulers):
        if update:
            loss_adjusted.backward()
            grad_norm = nn.utils.clip_grad_norm_(models.controlnet.parameters(), 1.0)
            optimizers_dict = optimizers.to_dict()
            for k in optimizers_dict:
                if optimizers_dict[k] is not None and k != 'training':
                    optimizers_dict[k].step()
            schedulers_dict = schedulers.to_dict()
            for k in schedulers_dict:
                if k != 'training':
                    schedulers_dict[k].step()
            for k in optimizers_dict:
                if optimizers_dict[k] is not None and k != 'training':
                    optimizers_dict[k].zero_grad(set_to_none=True)
            self.info.total_steps += 1
        else:
            loss_adjusted.backward()
            grad_norm = torch.tensor(0.0).to(self.device)

        return grad_norm

    def models_to_save(self):
        return ['controlnet']  # ['generator', 'generator_ema']

    # LATENT ENCODING & PROCESSING ----------
    def encode_latents(self, batch: dict, models: Models, extras: Extras) -> torch.Tensor:
        images = batch['images'].to(self.device)
        return models.effnet(extras.effnet_preprocess(images))

    def decode_latents(self, latents: torch.Tensor, batch: dict, models: Models, extras: Extras) -> torch.Tensor:
        return models.previewer(latents)

    def sample(self, models: Models, data: WarpCore.Data, extras: Extras):
        models.controlnet.eval()
        with torch.no_grad():
            batch = next(data.iterator)

            cnet, cnet_input = self.get_cnet(batch, models, extras)
            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)
            conditions, unconditions = {**conditions, 'cnet': cnet}, {**unconditions, 'cnet': cnet}

            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')
                    cnet_input = nn.functional.interpolate(cnet_input, size=noised_images.shape[-2:], mode='bicubic')
                    if cnet_input.size(1) == 1:
                        cnet_input = cnet_input.repeat(1, 3, 1, 1)
                    elif cnet_input.size(1) > 3:
                        cnet_input = cnet_input[:, :3]

                collage_img = torch.cat([
                    torch.cat([i for i in images.cpu()], dim=-1),
                    torch.cat([i for i in cnet_input.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(cnet_input[i])] + [wandb.Image(images[i])] for i in range(len(images))]
                    log_table = wandb.Table(data=log_data,
                                            columns=["Captions", "Sampled", "Sampled EMA", "Cnet", "Orig"])
                    wandb.log({"Log": log_table})
            models.controlnet.train()
            models.controlnet.backbone.eval()


if __name__ == '__main__':
    print("Launching Script")
    warpcore = WurstCore(
        config_file_path=sys.argv[1] if len(sys.argv) > 1 else None,
        device=torch.device(int(os.environ.get("SLURM_LOCALID")))
    )
    warpcore.fsdp_defaults['sharding_strategy'] = ShardingStrategy.NO_SHARD

    # RUN TRAINING
    warpcore()