utils.py

"""Training state, logging metrics, and sample-writing helpers."""

from collections import defaultdict, deque
import logging
from pathlib import Path

import torch
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torchvision.utils import save_image
from tqdm import tqdm

from libs.uvit_t2i_vq import UViT


LOGGER = logging.getLogger("styledrop")


def update_ema(destination, source, rate):
    source_parameters = dict(source.named_parameters())
    for name, destination_parameter in destination.named_parameters():
        source_parameter = source_parameters[name]
        if "adapter" in name:
            destination_parameter.data.copy_(source_parameter.data)
        else:
            destination_parameter.data.mul_(rate).add_(
                source_parameter.data,
                alpha=1 - rate,
            )


class TrainState:
    def __init__(
        self,
        optimizer,
        lr_scheduler,
        step,
        nnet,
        nnet_ema,
    ):
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler
        self.step = step
        self.nnet = nnet
        self.nnet_ema = nnet_ema

    def ema_update(self, rate=0.9999):
        update_ema(self.nnet_ema, self.nnet, rate)

    def save(self, path):
        """Save only the trained adapter and its step number."""
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)
        torch.save(self.step, path / "step.pth")
        torch.save(
            self.nnet.adapter.state_dict(),
            path / "adapter.pth",
        )

    @staticmethod
    def _merge_checkpoint(model, checkpoint):
        """Keep initialized values for adapter keys absent from the base model."""
        current_state = model.state_dict()
        return {
            name: checkpoint.get(name, value).clone()
            for name, value in current_state.items()
        }

    def _load_base_checkpoint(self, path):
        path = Path(path)
        LOGGER.info("Loading base checkpoint from %s", path)
        self.step = torch.load(path / "step.pth", map_location="cpu")

        for name in ("nnet", "nnet_ema"):
            model = getattr(self, name)
            checkpoint = torch.load(
                path / f"{name}.pth",
                map_location="cpu",
            )
            model.load_state_dict(
                self._merge_checkpoint(model, checkpoint)
            )

    def _load_adapter(self, path):
        path = Path(path)
        LOGGER.info("Loading adapter from %s", path)
        adapter = torch.load(path, map_location="cpu")
        self.nnet.adapter.load_state_dict(adapter)
        self.nnet_ema.adapter.load_state_dict(adapter)

    def resume(
        self,
        checkpoint_root,
        adapter_path=None,
        step=None,
        required=False,
    ):
        checkpoint_root = Path(checkpoint_root)
        if not checkpoint_root.exists():
            if required:
                raise FileNotFoundError(
                    "Base MUSE checkpoint was not found at "
                    f"{checkpoint_root}. See README.md for setup."
                )
            return False

        if checkpoint_root.suffix == ".ckpt":
            checkpoint_path = checkpoint_root
        else:
            checkpoint_steps = [
                int(path.stem)
                for path in checkpoint_root.iterdir()
                if path.suffix == ".ckpt" and path.stem.isdigit()
            ]
            if step is None:
                if not checkpoint_steps:
                    if required:
                        raise FileNotFoundError(
                            f"No *.ckpt directories found in {checkpoint_root}"
                        )
                    return False
                step = max(checkpoint_steps)
            checkpoint_path = checkpoint_root / f"{step}.ckpt"

        self._load_base_checkpoint(checkpoint_path)
        if adapter_path:
            self._load_adapter(adapter_path)
        return True

    def to(self, device):
        self.nnet.to(device)
        self.nnet_ema.to(device)

    def freeze(self):
        """Freeze the denoiser except for its StyleDrop adapter."""
        self.nnet.requires_grad_(False)
        self.nnet.adapter.requires_grad_(True)
        self.nnet_ema.requires_grad_(False)


def initialize_train_state(config, device):
    nnet = UViT(**config.nnet)
    nnet_ema = UViT(**config.nnet)
    nnet_ema.eval()
    LOGGER.info(
        "Denoiser has %s parameters",
        f"{sum(parameter.numel() for parameter in nnet.parameters()):,}",
    )

    optimizer = AdamW(
        nnet.adapter.parameters(),
        **config.optimizer,
    )

    warmup_steps = config.lr_scheduler.warmup_steps

    def lr_multiplier(step):
        if warmup_steps > 0:
            return min(step / warmup_steps, 1.0)
        return 1.0

    lr_scheduler = LambdaLR(optimizer, lr_multiplier)
    state = TrainState(
        optimizer=optimizer,
        lr_scheduler=lr_scheduler,
        step=0,
        nnet=nnet,
        nnet_ema=nnet_ema,
    )
    state.ema_update(rate=0)
    state.to(device)
    return state


def _batch_sizes(n_samples, batch_size):
    full_batches, remainder = divmod(n_samples, batch_size)
    sizes = [batch_size] * full_batches
    if remainder:
        sizes.append(remainder)
    return sizes


def sample_to_directory(
    accelerator,
    path,
    n_samples,
    mini_batch_size,
    sample_fn,
    unpreprocess_fn,
):
    if path:
        Path(path).mkdir(parents=True, exist_ok=True)

    sample_index = 0
    distributed_batch_size = (
        mini_batch_size * accelerator.num_processes
    )
    for current_batch_size in tqdm(
        _batch_sizes(n_samples, distributed_batch_size),
        disable=not accelerator.is_main_process,
        desc="sampling",
    ):
        samples = unpreprocess_fn(sample_fn(mini_batch_size))
        samples = accelerator.gather(samples.contiguous())[
            :current_batch_size
        ]
        if accelerator.is_main_process:
            for sample in samples:
                save_image(
                    sample,
                    Path(path) / f"{sample_index}.png",
                )
                sample_index += 1


class SmoothedValue:
    def __init__(
        self,
        window_size=20,
        fmt="{median:.4f} ({global_avg:.4f})",
    ):
        self.deque = deque(maxlen=window_size)
        self.total = 0.0
        self.count = 0
        self.fmt = fmt

    def update(self, value, n=1):
        self.deque.append(value)
        self.count += n
        self.total += value * n

    @property
    def median(self):
        return torch.tensor(list(self.deque)).median().item()

    @property
    def avg(self):
        return torch.tensor(
            list(self.deque),
            dtype=torch.float32,
        ).mean().item()

    @property
    def global_avg(self):
        return self.total / self.count

    @property
    def max(self):
        return max(self.deque)

    @property
    def value(self):
        return self.deque[-1]

    def __str__(self):
        return self.fmt.format(
            median=self.median,
            avg=self.avg,
            global_avg=self.global_avg,
            max=self.max,
            value=self.value,
        )


class MetricLogger:
    def __init__(self, delimiter=" "):
        self.meters = defaultdict(SmoothedValue)
        self.delimiter = delimiter

    def update(self, **values):
        for name, value in values.items():
            if isinstance(value, torch.Tensor):
                value = value.item()
            if not isinstance(value, (float, int)):
                raise TypeError(
                    f"Metric {name!r} is not numeric: {type(value)}"
                )
            self.meters[name].update(value)

    def __getattr__(self, name):
        if name in self.meters:
            return self.meters[name]
        raise AttributeError(
            f"{type(self).__name__!s} has no attribute {name!r}"
        )

    def __str__(self):
        return self.delimiter.join(
            f"{name}: {meter}"
            for name, meter in self.meters.items()
        )


def get_grad_norm_(parameters, norm_type=2.0):
    parameters = [
        parameter
        for parameter in parameters
        if parameter.grad is not None
    ]
    if not parameters:
        return torch.tensor(0.0)

    norm_type = float(norm_type)
    device = parameters[0].grad.device
    if norm_type == float("inf"):
        return max(
            parameter.grad.detach().abs().max().to(device)
            for parameter in parameters
        )

    return torch.norm(
        torch.stack(
            [
                torch.norm(
                    parameter.grad.detach(),
                    norm_type,
                ).to(device)
                for parameter in parameters
            ]
        ),
        norm_type,
    )