train.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
from makeyourdance.models.unet import UNetSpatioTemporalConditionModel
import torch.nn.functional as F
from makeyourdance.pipeline.svdpipeline import _resize_with_antialiasing
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 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:
        logger.info('loading custom unet...')

    unet = UNetSpatioTemporalConditionModel.custom_load(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, trainable_parameters.parameters())
    optimizer = torch.optim.AdamW(
        unet,
        lr=0.0001,
    )
    enable_gc = False
    if enable_gc:
        unet.enable_gradient_checkpointing()
    unet.enable_xformers_memory_efficient_attention()
    vae.to(device)
    image_encoder.to(device)
    scheduler = Scheduler()
    # feature_extractor.to(device)



    steps = 1000
    batch_size = 1
    dataset = RandomTensorDataset()
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    unet.to(device)
    dataloader, unet, optimizer = (
        accelerator.prepare(dataloader, unet, optimizer, scheduler)
    )
    fps = 6
    for step, batch in enumerate(dataloader):
        unet.train()
        image, video = batch

        image = image.to(device=device, dtype=next(image_encoder.parameters()).dtype)
        image_embeddings = image_encoder(image).image_embeds
        image_embeddings = image_embeddings.unsqueeze(1)

        # 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)
        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, 25, 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([video, 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


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='config/pretrain.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)