trainfortesting.py

import os
import argparse
import makeyourdance.utils.config as config
from accelerate import Accelerator
import accelerate
import transformers
from transformers import CLIPTextModel, CLIPTokenizer
import makeyourdance.utils.logger as logger
import wandb
import torch
from makeyourdance.utils.loader import Optimizer, Dataset, Dataloader
from diffusers.schedulers import EulerDiscreteScheduler
from diffusers.image_processor import VaeImageProcessor
from diffusers.models import AutoencoderKLTemporalDecoder, UNetSpatioTemporalConditionModel
# from makeyourdance.models.unet import UNetSpatioTemporalConditionModel
import torch.nn.functional as F

from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection

def main():
    accelerator = Accelerator()
    num_processes = accelerator.num_processes
    is_main_process = accelerator.is_main_process
    device = accelerator.device
    global_seed = 1 if not hasattr(cfgs, 'global_seed') else cfgs.global_seed
    accelerate.utils.set_seed(global_seed, device_specific=True)
    wandb_config = {
                    "batch_size": 8,
                    "attention_resolutions": [2, 4, 8],
                    "channel_mults":(1, 2, 4, 8),
                    "model_channels": [96],
                    "dropout": 0.,
                    "learn_rate": 0.0006,
                    "decay": 0.0005,
                    "num_head_channel": 32,
                }
    # if is_main_process and (not cfgs.is_debug) and cfgs.use_wandb:
    #     wandb.init(project='my-diff',
    #                entity='rightbrainai',
    #                config=wandb_config,
    #                dir=str(os.path.join(cfgs.work_dir)),
    #                )
    if is_main_process:
        os.makedirs(cfgs.work_dir, exist_ok=True)
        os.makedirs(f"{cfgs.work_dir}/samples", exist_ok=True)
        # os.makedirs(f"{cfgs.work_dir}/sanity_check", exist_ok=True)
        os.makedirs(f"{cfgs.work_dir}/checkpoints", exist_ok=True)

    accelerator.wait_for_everyone()
    noise_scheduler = EulerDiscreteScheduler()
    if is_main_process:
        logger.info('Loading pretrained model weights...')
    vae = AutoencoderKLTemporalDecoder.from_pretrained(cfgs.pretrain_path, subfolder='vae')
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(cfgs.pretrain_path, subfolder='image_encoder')
    feature_extractor = CLIPImageProcessor.from_pretrained(cfgs.pretrain_path, subfolder='feature_extractor')
    # if is_main_process:

    unet = UNetSpatioTemporalConditionModel.from_pretrained(cfgs.pretrain_path, subfolder='unet')

    vae.requires_grad_(False)
    image_encoder.requires_grad_(False)
    # feature_extractor.requires_grad_(False)
    unet.requires_grad_(True)

    # optimizer = Optimizer(cfgs, model.parameters())
    optimizer = torch.optim.AdamW(
        unet.parameters(),
        lr=0.0001,
    )
    enable_gc = True
    if enable_gc:
        unet.enable_gradient_checkpointing()
    unet.enable_xformers_memory_efficient_attention()
    vae.to(device)
    image_encoder.to(device)
    image_processor = VaeImageProcessor()
    # feature_extractor.to(device)


    height = 1024
    width = 576
    frame = 25
    steps = 1000
    batch_size = 1
    dataset = RandomTensorDataset(image_size=(3, height, width), video_size=(frame, 3, height, width))
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    unet.to(device)
    dataloader, unet, optimizer = (
        accelerator.prepare(dataloader, unet, optimizer)
    )
    fps = 6
    for step, batch in enumerate(dataloader):
        unet.train()
        image, video = batch

        image = image * 2.0 - 1.0
        image = _resize_with_antialiasing(image, (224, 224))
        image = (image + 1.0) / 2.0

        # Normalize the image with for CLIP input
        image = feature_extractor(
            images=image,
            do_normalize=True,
            do_center_crop=False,
            do_resize=False,
            do_rescale=False,
            return_tensors="pt",
        ).pixel_values
        image_embeddings = image_encoder(image).image_embeds
        image_embeddings = image_embeddings.unsqueeze(1)
        image = image_processor.preprocess(image, height=height, width=width)
        # duplicate image embeddings for each generation per prompt, using mps friendly method
        bs_embed, seq_len, _ = image_embeddings.shape
        image_embeddings = image_embeddings.repeat(1, 1, 1)
        image_embeddings = image_embeddings.view(bs_embed * 1, seq_len, -1)
        image = image.to(device=device, dtype=next(image_encoder.parameters()).dtype)
        # image = image_processor.preprocess(image, height=height, width=width)
        needs_upcasting = vae.dtype == torch.float16
        if needs_upcasting:
            vae.to(dtype=torch.float32)
        image_latents = vae.encode(image).latent_dist.mode()
        image_latents = image_latents.repeat(1, 1, 1, 1)
        image_latents = image_latents.to(image_embeddings.dtype)
        if needs_upcasting:
            vae.to(dtype=torch.float16)
        image_latents = image_latents.unsqueeze(1).repeat(1, frame, 1, 1, 1)
        add_time_ids = [fps, 127, 0.02]

        passed_add_embed_dim = unet.config.addition_time_embed_dim * len(add_time_ids)
        expected_add_embed_dim = unet.add_embedding.linear_1.in_features

        if expected_add_embed_dim != passed_add_embed_dim:
            raise ValueError(
                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
            )

        add_time_ids = torch.tensor([add_time_ids], dtype=image_embeddings.dtype)
        add_time_ids = add_time_ids.repeat(batch_size * 1, 1)

        added_time_ids = add_time_ids.to(device)
        num_channels_latents = unet.config.in_channels
        latent_model_input = torch.cat([image_latents, image_latents], dim=2)
        noise_pred = unet(
            latent_model_input,
            500,
            encoder_hidden_states=image_embeddings,
            added_time_ids=added_time_ids,
            return_dict=False,
        )[0]
        loss = F.mse_loss(noise_pred.float(), video.float(), reduction="mean")
        optimizer.zero_grad()
        accelerator.backward(loss)
        """ >>> gradient clipping >>> """
        torch.nn.utils.clip_grad_norm_(unet.parameters(), 1)
        """ <<< gradient clipping <<< """
        optimizer.step()
        logger.info('loss: {}'.format(loss))



from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import Resize
from PIL import Image
def _resize_with_antialiasing(input, size, interpolation="bicubic", align_corners=True):
    h, w = input.shape[-2:]
    factors = (h / size[0], w / size[1])

    # First, we have to determine sigma
    # Taken from skimage: https://github.com/scikit-image/scikit-image/blob/v0.19.2/skimage/transform/_warps.py#L171
    sigmas = (
        max((factors[0] - 1.0) / 2.0, 0.001),
        max((factors[1] - 1.0) / 2.0, 0.001),
    )

    # Now kernel size. Good results are for 3 sigma, but that is kind of slow. Pillow uses 1 sigma
    # https://github.com/python-pillow/Pillow/blob/master/src/libImaging/Resample.c#L206
    # But they do it in the 2 passes, which gives better results. Let's try 2 sigmas for now
    ks = int(max(2.0 * 2 * sigmas[0], 3)), int(max(2.0 * 2 * sigmas[1], 3))

    # Make sure it is odd
    if (ks[0] % 2) == 0:
        ks = ks[0] + 1, ks[1]

    if (ks[1] % 2) == 0:
        ks = ks[0], ks[1] + 1

    input = _gaussian_blur2d(input, ks, sigmas)

    output = torch.nn.functional.interpolate(input, size=size, mode=interpolation, align_corners=align_corners)
    return output


def _compute_padding(kernel_size):
    """Compute padding tuple."""
    # 4 or 6 ints:  (padding_left, padding_right,padding_top,padding_bottom)
    # https://pytorch.org/docs/stable/nn.html#torch.nn.functional.pad
    if len(kernel_size) < 2:
        raise AssertionError(kernel_size)
    computed = [k - 1 for k in kernel_size]

    # for even kernels we need to do asymmetric padding :(
    out_padding = 2 * len(kernel_size) * [0]

    for i in range(len(kernel_size)):
        computed_tmp = computed[-(i + 1)]

        pad_front = computed_tmp // 2
        pad_rear = computed_tmp - pad_front

        out_padding[2 * i + 0] = pad_front
        out_padding[2 * i + 1] = pad_rear

    return out_padding


def _filter2d(input, kernel):
    # prepare kernel
    b, c, h, w = input.shape
    tmp_kernel = kernel[:, None, ...].to(device=input.device, dtype=input.dtype)

    tmp_kernel = tmp_kernel.expand(-1, c, -1, -1)

    height, width = tmp_kernel.shape[-2:]

    padding_shape: list[int] = _compute_padding([height, width])
    input = torch.nn.functional.pad(input, padding_shape, mode="reflect")

    # kernel and input tensor reshape to align element-wise or batch-wise params
    tmp_kernel = tmp_kernel.reshape(-1, 1, height, width)
    input = input.view(-1, tmp_kernel.size(0), input.size(-2), input.size(-1))

    # convolve the tensor with the kernel.
    output = torch.nn.functional.conv2d(input, tmp_kernel, groups=tmp_kernel.size(0), padding=0, stride=1)

    out = output.view(b, c, h, w)
    return out


def _gaussian(window_size: int, sigma):
    if isinstance(sigma, float):
        sigma = torch.tensor([[sigma]])

    batch_size = sigma.shape[0]

    x = (torch.arange(window_size, device=sigma.device, dtype=sigma.dtype) - window_size // 2).expand(batch_size, -1)

    if window_size % 2 == 0:
        x = x + 0.5

    gauss = torch.exp(-x.pow(2.0) / (2 * sigma.pow(2.0)))

    return gauss / gauss.sum(-1, keepdim=True)


def _gaussian_blur2d(input, kernel_size, sigma):
    if isinstance(sigma, tuple):
        sigma = torch.tensor([sigma], dtype=input.dtype)
    else:
        sigma = sigma.to(dtype=input.dtype)

    ky, kx = int(kernel_size[0]), int(kernel_size[1])
    bs = sigma.shape[0]
    kernel_x = _gaussian(kx, sigma[:, 1].view(bs, 1))
    kernel_y = _gaussian(ky, sigma[:, 0].view(bs, 1))
    out_x = _filter2d(input, kernel_x[..., None, :])
    out = _filter2d(out_x, kernel_y[..., None])

    return out

class RandomTensorDataset(Dataset):
    def __init__(self, num_samples=1000, image_size=(3, 1024, 576), video_size=(25, 3, 1024, 576), transform=None):
        self.num_samples = num_samples
        self.image_size = image_size
        self.video_size = video_size
        self.transform = transform

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        image = torch.randn(self.image_size)
        video = torch.randn(self.video_size)
        if self.transform:
            image = self.transform(image)

        return image, video

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', type=str, default='configs/test.yaml',
                        help='JSON file for configuration')
    parser.add_argument('-w', '--wandb', type=str, default=True,
                        help='uploading the experiment to wandb')
    parser.add_argument('-gpu', '--gpu_ids', type=str, default=None)
    # parser.add_argument('-enable_wandb', action='store_true')
    args = parser.parse_args()
    cfgs, dict_cfgs = config.from_yaml(args.config)
    # os.environ['CUDA_VISIBLE_DEVICES'] = cfgs.gpus_id
    # Convert to NoneDict, which return None for missing key.
    # opt = Logger.dict_to_nonedict(opt)
    results_folder = cfgs.work_dir
    os.makedirs(results_folder, exist_ok=True)
    logger = logger.get_logger(name=cfgs.task_name, work_dir=cfgs.work_dir)


    main()
    if cfgs.use_wandb:
        wandb.finish()
    os._exit(0)