# Copyright 2024 NVIDIA CORPORATION & AFFILIATES # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # SPDX-License-Identifier: Apache-2.0 import os import random import re from collections import OrderedDict import numpy as np import torch from accelerate.state import DistributedType from diffusion.utils.logger import get_root_logger from tools.download import find_model def save_checkpoint( work_dir, epoch, model, accelerator=None, model_ema=None, optimizer=None, lr_scheduler=None, generator=torch.Generator(device="cpu").manual_seed(42), keep_last=False, step=None, saved_info=None, add_symlink=False, add_suffix=None, ): if accelerator is not None and accelerator.distributed_type == DistributedType.FSDP: return save_checkpoint_fsdp( work_dir=work_dir, epoch=epoch, model=model, accelerator=accelerator, lr_scheduler=lr_scheduler, generator=generator, keep_last=keep_last, step=step, saved_info=saved_info, add_symlink=add_symlink, add_suffix=add_suffix, ) else: return save_checkpoint_ddp( work_dir=work_dir, epoch=epoch, model=model, model_ema=model_ema, optimizer=optimizer, lr_scheduler=lr_scheduler, generator=generator, keep_last=keep_last, step=step, saved_info=saved_info, add_symlink=add_symlink, add_suffix=add_suffix, ) def save_checkpoint_ddp( work_dir, epoch, model, model_ema=None, optimizer=None, lr_scheduler=None, generator=torch.Generator(device="cpu").manual_seed(42), keep_last=False, step=None, saved_info=None, add_symlink=False, add_suffix=None, ): os.makedirs(work_dir, exist_ok=True) state_dict = dict(state_dict=model.state_dict()) if model_ema is not None: state_dict["state_dict_ema"] = model_ema.state_dict() if optimizer is not None: state_dict["optimizer"] = optimizer.state_dict() if lr_scheduler is not None: state_dict["scheduler"] = lr_scheduler.state_dict() if epoch is not None: state_dict["epoch"] = epoch file_path = os.path.join(work_dir, f"epoch_{epoch}.pth") if step is not None: state_dict["step"] = step file_path = file_path.split(".pth")[0] + f"_step_{step}.pth" # Save additional information from saved_info dict if saved_info is not None: for key, value in saved_info.items(): if value is not None: state_dict[key] = value if add_suffix is not None: file_path = file_path.replace(".pth", f"_{add_suffix}.pth") rng_state = { "torch": torch.get_rng_state(), "torch_cuda": torch.cuda.get_rng_state_all(), "numpy": np.random.get_state(), "python": random.getstate(), "generator": generator.get_state(), } state_dict["rng_state"] = rng_state logger = get_root_logger() torch.save(state_dict, file_path) logger.info(f"Saved checkpoint of epoch {epoch} to {file_path.format(epoch)}.") if keep_last: for i in range(epoch): previous_ckgt = file_path.format(i) if os.path.exists(previous_ckgt): os.remove(previous_ckgt) if add_symlink: link_path = os.path.join(os.path.dirname(file_path), "latest.pth") if os.path.exists(link_path) or os.path.islink(link_path): os.remove(link_path) os.symlink(os.path.abspath(file_path), link_path) return file_path def save_checkpoint_fsdp( work_dir, epoch, model=None, accelerator=None, lr_scheduler=None, generator=torch.Generator(device="cpu").manual_seed(42), keep_last=False, step=None, saved_info=None, add_symlink=False, add_suffix=None, ): """FSDP checkpoint save function, sharding""" logger = get_root_logger() checkpoint_dir = os.path.join(work_dir, f"epoch_{epoch}") if step is not None: checkpoint_dir = checkpoint_dir + f"_step_{step}" if add_suffix is not None: checkpoint_dir = checkpoint_dir + f"_{add_suffix}" os.makedirs(checkpoint_dir, exist_ok=True) model_dir = os.path.join(checkpoint_dir, "model") os.makedirs(model_dir, exist_ok=True) accelerator.save_state(model_dir) merged_state_dict = None if model is not None: try: merged_state_dict = accelerator.get_state_dict(model) except Exception as exc: logger.warning(f"Unable to export merged FSDP checkpoint for inference: {exc}") if accelerator.is_main_process: metadata = dict() rng_state = { "torch": torch.get_rng_state(), "torch_cuda": torch.cuda.get_rng_state_all(), "numpy": np.random.get_state(), "python": random.getstate(), "generator": generator.get_state(), } metadata["rng_state"] = rng_state if lr_scheduler is not None: metadata["scheduler"] = lr_scheduler.state_dict() if epoch is not None: metadata["epoch"] = epoch if step is not None: metadata["step"] = step # Save additional information from saved_info dict if saved_info is not None: for key, value in saved_info.items(): if value is not None: metadata[key] = value torch.save(metadata, os.path.join(checkpoint_dir, "metadata.pth")) if keep_last: checkpoints = sorted( [d for d in os.listdir(work_dir) if os.path.isdir(os.path.join(work_dir, d)) and d.startswith("epoch_")] ) for old_ckpt in checkpoints[:-1]: old_path = os.path.join(work_dir, old_ckpt) if os.path.exists(old_path): import shutil shutil.rmtree(old_path) if add_symlink: link_path = os.path.join(work_dir, "latest.pth") if os.path.exists(link_path) or os.path.islink(link_path): os.remove(link_path) os.symlink(os.path.abspath(checkpoint_dir), link_path) logger.info(f"Saved checkpoint to {checkpoint_dir}") model_link_path = checkpoint_dir + ".pth" if merged_state_dict is not None: torch.save({"state_dict": merged_state_dict}, model_link_path) else: fsdp_model_path = os.path.join(model_dir, "pytorch_model_fsdp.bin") if os.path.exists(fsdp_model_path): state_dict = torch.load(fsdp_model_path, map_location="cpu") torch.save({"state_dict": state_dict}, model_link_path) accelerator.wait_for_everyone() return checkpoint_dir def load_checkpoint( checkpoint, model, model_ema=None, optimizer=None, lr_scheduler=None, load_ema=False, resume_optimizer=True, resume_lr_scheduler=True, null_embed_path=None, FSDP=False, remove_state_dict_keys=None, ): if FSDP: return load_checkpoint_fsdp( checkpoint=checkpoint, model=model, remove_state_dict_keys=remove_state_dict_keys, ) else: return load_checkpoint_ddp( checkpoint=checkpoint, model=model, model_ema=model_ema, optimizer=optimizer, lr_scheduler=lr_scheduler, load_ema=load_ema, resume_optimizer=resume_optimizer, resume_lr_scheduler=resume_lr_scheduler, null_embed_path=null_embed_path, remove_state_dict_keys=remove_state_dict_keys, ) def load_checkpoint_ddp( checkpoint, model, model_ema=None, optimizer=None, lr_scheduler=None, load_ema=False, resume_optimizer=True, resume_lr_scheduler=True, null_embed_path=None, remove_state_dict_keys=None, ): assert isinstance(checkpoint, str) logger = get_root_logger() ckpt_file = checkpoint checkpoint = find_model(ckpt_file) if remove_state_dict_keys is None: remove_state_dict_keys = [] remove_state_dict_keys.extend(["pos_embed", "base_model.pos_embed", "model.pos_embed"]) for key in remove_state_dict_keys: if key in checkpoint["state_dict"]: del checkpoint["state_dict"][key] if "state_dict_ema" in checkpoint and key in checkpoint["state_dict_ema"]: del checkpoint["state_dict_ema"][key] if load_ema: state_dict = checkpoint["state_dict_ema"] else: state_dict = checkpoint.get("state_dict", checkpoint) # to be compatible with the official checkpoint null_embed = torch.load(null_embed_path, map_location="cpu") state_dict["y_embedder.y_embedding"] = null_embed["uncond_prompt_embeds"][0] rng_state = checkpoint.get("rng_state", None) def load_ckpt_with_auto_reshape(model, state_dict, strict=False): new_state_dict = OrderedDict() for k, v in model.state_dict().items(): if k in state_dict: # auto reshape missing dimensions (e.g. [dim,dim2,1,1] -> [dim,dim2,1,1,1]) if state_dict[k].dim() < v.dim(): new_shape = state_dict[k].shape + (1,) * (v.dim() - state_dict[k].dim()) new_state_dict[k] = state_dict[k].reshape(*new_shape) else: new_state_dict[k] = state_dict[k] else: print(f"Warning: Missing key {k} in checkpoint") missing, unexpect = model.load_state_dict(new_state_dict, strict=strict) return missing, unexpect missing, unexpect = load_ckpt_with_auto_reshape(model, state_dict, strict=False) if model_ema is not None: model_ema.load_state_dict(checkpoint["state_dict_ema"], strict=False) if optimizer is not None and resume_optimizer: optimizer.load_state_dict(checkpoint["optimizer"]) if lr_scheduler is not None and resume_lr_scheduler: lr_scheduler.load_state_dict(checkpoint["scheduler"]) epoch = 0 # Load saved_info dictionary containing video_step, image_step, etc. saved_info = {} known_saved_keys = ["video_step", "image_step"] # Add more keys as needed for key in known_saved_keys: value = checkpoint.get(key, None) if value is not None: saved_info[key] = value if optimizer is not None and resume_optimizer: epoch_match = re.search(r"epoch_(\d+)", ckpt_file) epoch = int(epoch_match.group(1)) if epoch_match else 0 logger.info( f"Resume checkpoint of epoch {epoch} from {ckpt_file}. Load ema: {load_ema}, " f"resume optimizer: {resume_optimizer}, resume lr scheduler: {resume_lr_scheduler}." ) return epoch, missing, unexpect, rng_state, saved_info logger.info(f"Load checkpoint from {ckpt_file}. Load ema: {load_ema}.") return epoch, missing, unexpect, None, saved_info def load_checkpoint_fsdp( checkpoint, model, remove_state_dict_keys=None, ): assert isinstance(checkpoint, str) logger = get_root_logger() # 1 load model if ".pth" in checkpoint: state_dict_model = find_model(checkpoint) state_dict_model = state_dict_model.get("state_dict", state_dict_model) metadata = {} else: if os.path.isfile(checkpoint): checkpoint = os.path.dirname(checkpoint) assert os.path.isdir(checkpoint), f"Checkpoint directory {checkpoint} does not exist!" state_dict_model = find_model(os.path.join(checkpoint, "model", "pytorch_model_fsdp.bin")) # Load metadata to get video_step and image_step try: metadata = torch.load(os.path.join(checkpoint, "metadata.pth"), map_location="cpu") except: metadata = {} if remove_state_dict_keys is None: remove_state_dict_keys = [] remove_state_dict_keys.extend(["pos_embed", "base_model.pos_embed", "model.pos_embed"]) for key in remove_state_dict_keys: if key in state_dict_model: del state_dict_model[key] def load_ckpt_with_auto_reshape(model, state_dict, strict=False): new_state_dict = OrderedDict() for k, v in model.state_dict().items(): if k in state_dict: # auto reshape missing dimensions (e.g. [dim,dim2,1,1] -> [dim,dim2,1,1,1]) if state_dict[k].dim() < v.dim(): new_shape = state_dict[k].shape + (1,) * (v.dim() - state_dict[k].dim()) new_state_dict[k] = state_dict[k].reshape(*new_shape) else: new_state_dict[k] = state_dict[k] else: print(f"Warning: Missing key {k} in checkpoint") missing, unexpect = model.load_state_dict(new_state_dict, strict=strict) return missing, unexpect missing, unexpect = load_ckpt_with_auto_reshape(model, state_dict_model, strict=False) # missing, unexpect = model.load_state_dict(state_dict_model, strict=False) logger.info(f"Load checkpoint of {checkpoint}.") # Load saved_info dictionary containing video_step, image_step, etc. saved_info = {} known_saved_keys = ["video_step", "image_step"] # Add more keys as needed for key in known_saved_keys: value = metadata.get(key, None) if value is not None: saved_info[key] = value return None, missing, unexpect, None, saved_info