dominic1021/LatentSync

1# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15from tqdm.auto import tqdm16import os, argparse, datetime, math17import logging18from omegaconf import OmegaConf19import shutil20 21from latentsync.data.syncnet_dataset import SyncNetDataset22from latentsync.models.syncnet import SyncNet23from latentsync.models.syncnet_wav2lip import SyncNetWav2Lip24from latentsync.utils.util import gather_loss, plot_loss_chart25from accelerate.utils import set_seed26 27import torch28from diffusers import AutoencoderKL29from diffusers.utils.logging import get_logger30from einops import rearrange31import torch.distributed as dist32from torch.nn.parallel import DistributedDataParallel as DDP33from torch.utils.data.distributed import DistributedSampler34from latentsync.utils.util import init_dist, cosine_loss35 36logger = get_logger(__name__)37 38 39def main(config):40    # Initialize distributed training41    local_rank = init_dist()42    global_rank = dist.get_rank()43    num_processes = dist.get_world_size()44    is_main_process = global_rank == 045 46    seed = config.run.seed + global_rank47    set_seed(seed)48 49    # Logging folder50    folder_name = "train" + datetime.datetime.now().strftime(f"-%Y_%m_%d-%H:%M:%S")51    output_dir = os.path.join(config.data.train_output_dir, folder_name)52 53    # Make one log on every process with the configuration for debugging.54    logging.basicConfig(55        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",56        datefmt="%m/%d/%Y %H:%M:%S",57        level=logging.INFO,58    )59 60    # Handle the output folder creation61    if is_main_process:62        os.makedirs(output_dir, exist_ok=True)63        os.makedirs(f"{output_dir}/checkpoints", exist_ok=True)64        os.makedirs(f"{output_dir}/loss_charts", exist_ok=True)65        shutil.copy(config.config_path, output_dir)66 67    device = torch.device(local_rank)68 69    if config.data.latent_space:70        vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)71        vae.requires_grad_(False)72        vae.to(device)73    else:74        vae = None75 76    # Dataset and Dataloader setup77    train_dataset = SyncNetDataset(config.data.train_data_dir, config.data.train_fileslist, config)78    val_dataset = SyncNetDataset(config.data.val_data_dir, config.data.val_fileslist, config)79 80    train_distributed_sampler = DistributedSampler(81        train_dataset,82        num_replicas=num_processes,83        rank=global_rank,84        shuffle=True,85        seed=config.run.seed,86    )87 88    # DataLoaders creation:89    train_dataloader = torch.utils.data.DataLoader(90        train_dataset,91        batch_size=config.data.batch_size,92        shuffle=False,93        sampler=train_distributed_sampler,94        num_workers=config.data.num_workers,95        pin_memory=False,96        drop_last=True,97        worker_init_fn=train_dataset.worker_init_fn,98    )99    100    num_samples_limit = 640101 102    val_batch_size = min(103        num_samples_limit // config.data.num_frames, config.data.batch_size104    )  # limit batch size to avoid CUDA OOM105 106    val_dataloader = torch.utils.data.DataLoader(107        val_dataset,108        batch_size=val_batch_size,109        shuffle=False,110        num_workers=config.data.num_workers,111        pin_memory=False,112        drop_last=False,113        worker_init_fn=val_dataset.worker_init_fn,114    )115 116    # Model117    syncnet = SyncNet(OmegaConf.to_container(config.model)).to(device)118    # syncnet = SyncNetWav2Lip().to(device)119 120    optimizer = torch.optim.AdamW(121        list(filter(lambda p: p.requires_grad, syncnet.parameters())), lr=config.optimizer.lr122    )123 124    if config.ckpt.resume_ckpt_path != "":125        if is_main_process:126            logger.info(f"Load checkpoint from: {config.ckpt.resume_ckpt_path}")127        ckpt = torch.load(config.ckpt.resume_ckpt_path, map_location=device)128 129        syncnet.load_state_dict(ckpt["state_dict"])130        global_step = ckpt["global_step"]131        train_step_list = ckpt["train_step_list"]132        train_loss_list = ckpt["train_loss_list"]133        val_step_list = ckpt["val_step_list"]134        val_loss_list = ckpt["val_loss_list"]135    else:136        global_step = 0137        train_step_list = []138        train_loss_list = []139        val_step_list = []140        val_loss_list = []141 142    # DDP wrapper143    syncnet = DDP(syncnet, device_ids=[local_rank], output_device=local_rank)144 145    num_update_steps_per_epoch = math.ceil(len(train_dataloader))146    num_train_epochs = math.ceil(config.run.max_train_steps / num_update_steps_per_epoch)147    # validation_steps = int(config.ckpt.save_ckpt_steps // 5)148    # validation_steps = 100149 150    if is_main_process:151        logger.info("***** Running training *****")152        logger.info(f"  Num examples = {len(train_dataset)}")153        logger.info(f"  Num Epochs = {num_train_epochs}")154        logger.info(f"  Instantaneous batch size per device = {config.data.batch_size}")155        logger.info(f"  Total train batch size (w. parallel & distributed) = {config.data.batch_size * num_processes}")156        logger.info(f"  Total optimization steps = {config.run.max_train_steps}")157 158    first_epoch = global_step // num_update_steps_per_epoch159    num_val_batches = config.data.num_val_samples // (num_processes * config.data.batch_size)160 161    # Only show the progress bar once on each machine.162    progress_bar = tqdm(163        range(0, config.run.max_train_steps), initial=global_step, desc="Steps", disable=not is_main_process164    )165 166    # Support mixed-precision training167    scaler = torch.cuda.amp.GradScaler() if config.run.mixed_precision_training else None168 169    for epoch in range(first_epoch, num_train_epochs):170        train_dataloader.sampler.set_epoch(epoch)171        syncnet.train()172 173        for step, batch in enumerate(train_dataloader):174            ### >>>> Training >>>> ###175 176            frames = batch["frames"].to(device, dtype=torch.float16)177            audio_samples = batch["audio_samples"].to(device, dtype=torch.float16)178            y = batch["y"].to(device, dtype=torch.float32)179 180            if config.data.latent_space:181                max_batch_size = (182                    num_samples_limit // config.data.num_frames183                )  # due to the limited cuda memory, we split the input frames into parts184                if frames.shape[0] > max_batch_size:185                    assert (186                        frames.shape[0] % max_batch_size == 0187                    ), f"max_batch_size {max_batch_size} should be divisible by batch_size {frames.shape[0]}"188                    frames_part_results = []189                    for i in range(0, frames.shape[0], max_batch_size):190                        frames_part = frames[i : i + max_batch_size]191                        frames_part = rearrange(frames_part, "b f c h w -> (b f) c h w")192                        with torch.no_grad():193                            frames_part = vae.encode(frames_part).latent_dist.sample() * 0.18215194                        frames_part_results.append(frames_part)195                    frames = torch.cat(frames_part_results, dim=0)196                else:197                    frames = rearrange(frames, "b f c h w -> (b f) c h w")198                    with torch.no_grad():199                        frames = vae.encode(frames).latent_dist.sample() * 0.18215200 201                frames = rearrange(frames, "(b f) c h w -> b (f c) h w", f=config.data.num_frames)202            else:203                frames = rearrange(frames, "b f c h w -> b (f c) h w")204 205            if config.data.lower_half:206                height = frames.shape[2]207                frames = frames[:, :, height // 2 :, :]208 209            # audio_embeds = wav2vec_encoder(audio_samples).last_hidden_state210 211            # Mixed-precision training212            with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=config.run.mixed_precision_training):213                vision_embeds, audio_embeds = syncnet(frames, audio_samples)214 215            loss = cosine_loss(vision_embeds.float(), audio_embeds.float(), y).mean()216 217            optimizer.zero_grad()218 219            # Backpropagate220            if config.run.mixed_precision_training:221                scaler.scale(loss).backward()222                """ >>> gradient clipping >>> """223                scaler.unscale_(optimizer)224                torch.nn.utils.clip_grad_norm_(syncnet.parameters(), config.optimizer.max_grad_norm)225                """ <<< gradient clipping <<< """226                scaler.step(optimizer)227                scaler.update()228            else:229                loss.backward()230                """ >>> gradient clipping >>> """231                torch.nn.utils.clip_grad_norm_(syncnet.parameters(), config.optimizer.max_grad_norm)232                """ <<< gradient clipping <<< """233                optimizer.step()234 235            progress_bar.update(1)236            global_step += 1237 238            global_average_loss = gather_loss(loss, device)239            train_step_list.append(global_step)240            train_loss_list.append(global_average_loss)241 242            if is_main_process and global_step % config.run.validation_steps == 0:243                logger.info(f"Validation at step {global_step}")244                val_loss = validation(245                    val_dataloader,246                    device,247                    syncnet,248                    cosine_loss,249                    config.data.latent_space,250                    config.data.lower_half,251                    vae,252                    num_val_batches,253                )254                val_step_list.append(global_step)255                val_loss_list.append(val_loss)256                logger.info(f"Validation loss at step {global_step} is {val_loss:0.3f}")257 258            if is_main_process and global_step % config.ckpt.save_ckpt_steps == 0:259                checkpoint_save_path = os.path.join(output_dir, f"checkpoints/checkpoint-{global_step}.pt")260                torch.save(261                    {262                        "state_dict": syncnet.module.state_dict(),  # to unwrap DDP263                        "global_step": global_step,264                        "train_step_list": train_step_list,265                        "train_loss_list": train_loss_list,266                        "val_step_list": val_step_list,267                        "val_loss_list": val_loss_list,268                    },269                    checkpoint_save_path,270                )271                logger.info(f"Saved checkpoint to {checkpoint_save_path}")272                plot_loss_chart(273                    os.path.join(output_dir, f"loss_charts/loss_chart-{global_step}.png"),274                    ("Train loss", train_step_list, train_loss_list),275                    ("Val loss", val_step_list, val_loss_list),276                )277 278            progress_bar.set_postfix({"step_loss": global_average_loss})279            if global_step >= config.run.max_train_steps:280                break281 282    progress_bar.close()283    dist.destroy_process_group()284 285 286@torch.no_grad()287def validation(val_dataloader, device, syncnet, cosine_loss, latent_space, lower_half, vae, num_val_batches):288    syncnet.eval()289 290    losses = []291    val_step = 0292    while True:293        for step, batch in enumerate(val_dataloader):294            ### >>>> Validation >>>> ###295 296            frames = batch["frames"].to(device, dtype=torch.float16)297            audio_samples = batch["audio_samples"].to(device, dtype=torch.float16)298            y = batch["y"].to(device, dtype=torch.float32)299 300            if latent_space:301                num_frames = frames.shape[1]302                frames = rearrange(frames, "b f c h w -> (b f) c h w")303                frames = vae.encode(frames).latent_dist.sample() * 0.18215304                frames = rearrange(frames, "(b f) c h w -> b (f c) h w", f=num_frames)305            else:306                frames = rearrange(frames, "b f c h w -> b (f c) h w")307 308            if lower_half:309                height = frames.shape[2]310                frames = frames[:, :, height // 2 :, :]311 312            with torch.autocast(device_type="cuda", dtype=torch.float16):313                vision_embeds, audio_embeds = syncnet(frames, audio_samples)314 315            loss = cosine_loss(vision_embeds.float(), audio_embeds.float(), y).mean()316 317            losses.append(loss.item())318 319            val_step += 1320            if val_step > num_val_batches:321                syncnet.train()322                if len(losses) == 0:323                    raise RuntimeError("No validation data")324                return sum(losses) / len(losses)325 326 327if __name__ == "__main__":328    parser = argparse.ArgumentParser(description="Code to train the expert lip-sync discriminator")329    parser.add_argument("--config_path", type=str, default="configs/syncnet/syncnet_16_vae.yaml")330    args = parser.parse_args()331 332    # Load a configuration file333    config = OmegaConf.load(args.config_path)334    config.config_path = args.config_path335 336    main(config)337