lipforcing/methods/model.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from abc import abstractmethod
from typing import Dict, Any, Optional, TYPE_CHECKING, Callable, List
import contextlib
import os

import torch

from lipforcing.configs.opt import get_scheduler
from lipforcing.utils import instantiate
from lipforcing.utils.distributed import synchronize, world_size
from lipforcing.utils.io_utils import s3_load
import lipforcing.utils.logging_utils as logger
import lipforcing.utils.basic_utils as basic_utils
from lipforcing.utils.distributed import is_rank0

if TYPE_CHECKING:
    from lipforcing.configs.config import BaseModelConfig
    from lipforcing.networks.network import FastGenNetwork


class FastGenModel(torch.nn.Module):
    def __init__(self, config: BaseModelConfig):
        """FastGenModel class for implementing training interface for all lipforcing networks.

        Args:
            config (BaseModelConfig): The configuration for the model
        """
        super().__init__()
        self.config = config

        # device
        self.device = torch.device(config.device)
        if self.device.type not in ["cuda", "cpu"]:
            raise ValueError(f"FastGenModel models only support cuda and cpu devices, got {self.device.type}")

        # precision and autocast
        self.set_precision(
            precision=self.config.precision,
            precision_amp=self.config.precision_amp,
            precision_amp_infer=self.config.precision_amp_infer,
            precision_amp_enc=self.config.precision_amp_enc,
            precision_fsdp=self.config.precision_fsdp,
        )

        # input shape
        self.input_shape = config.input_shape
        logger.info(f"Input shape is {self.input_shape}.")

        # define the name of the EMA networks to use
        use_ema = config.use_ema
        if isinstance(use_ema, bool):
            use_ema = ["ema"] if use_ema else []
        if not all(isinstance(name, str) and name.startswith("ema") for name in use_ema):
            raise ValueError(f"use_ema must be a bool or a list of strings starting with `ema`, got {use_ema}.")
        self.use_ema = use_ema

        # instantiate all necessary nets and submodules
        self.build_model()

    def _setup_ema(self):
        """Initialize EMA networks. Only call during build_model(), before checkpoint loading."""
        for name in self.use_ema:
            if not hasattr(self, name):
                logger.info(f"Initializing EMA network {name}")
                ema = instantiate(self.config.net)
                ema.eval().requires_grad_(False)
                setattr(self, name, ema)
            else:
                logger.warning(
                    f"EMA network {name} already exists, skipping initialization. "
                    "This is expected if loading pretrained network weights"
                )
            ema = getattr(self, name)
            # Only rank 0 loads weights if using meta initialization (non-rank-0 has meta tensors in self.net)
            if (not self.config.fsdp_meta_init) or is_rank0():
                net_load_info = ema.load_state_dict(self.net.state_dict(), strict=False)
                logger.success(f"Loaded EMA network {name}. Loading info: {net_load_info}")
            # Broadcast EMA weights from rank 0 to all other ranks when using meta init
            if world_size() > 1 and self.config.fsdp_meta_init:
                ema.to(device=self.device)
                for param in ema.parameters():
                    torch.distributed.broadcast(param.data, src=0)
                for buffer in ema.buffers():
                    torch.distributed.broadcast(buffer.data, src=0)
        synchronize()

    def _get_meta_init_context(self, fsdp_meta_init: bool = None):
        """Get context manager for FSDP meta initialization.

        When fsdp_meta_init is enabled, non-rank-0 processes use meta device
        for memory-efficient loading. Rank 0 loads weights normally, then
        FSDP syncs weights to other ranks via sync_module_states.

        Args:
            fsdp_meta_init: Whether to use meta initialization. If None, uses
                self.config.fsdp_meta_init.
        """
        if fsdp_meta_init is None:
            fsdp_meta_init = self.config.fsdp_meta_init
        use_meta = fsdp_meta_init and not is_rank0()
        if use_meta:
            return torch.device("meta")
        return contextlib.nullcontext()

    def set_precision(
        self,
        precision: str = "float32",
        precision_amp: str | None = None,
        precision_amp_infer: str | None = None,
        precision_amp_enc: str | None = None,
        precision_fsdp: str | None = None,
    ):
        """Set the model/data precision and automatic mixed precision (AMP) precision for training and inference.

        All precision arguments are strings that are mapped to torch dtypes according to PRECISION_MAP:
            "float16" -> torch.float16
            "bfloat16" -> torch.bfloat16
            "float32" -> torch.float32
            "float64" -> torch.float64

        Note that the precision of the time steps is handled in the noise scheduler (defaulting to float64 for numerical stability).

        Args:
            precision: Precision for model/optimizer states and data. Recommended to be float32 if precision_amp is not None.
            precision_amp: Precision for AMP during training. If None or equal to precision, AMP is disabled during training.
            precision_amp_infer: Precision for AMP during inference. If None or equal to precision, AMP is disabled during inference.
            precision_amp_enc: Precision for AMP en-/decoder (e.g., for VAEs or text encoders).
                If None or equal to precision, AMP is disabled during en-/decoding.
            precision_fsdp: Precision for FSDP2 parameter storage and gradient reduction.
                If None, defaults to `precision`.
        """

        # precision for model/optimizer states and data
        self.precision = basic_utils.PRECISION_MAP[precision]

        # precision for FSDP2 parameter storage and gradient reduction (defaults to model precision)
        self.precision_fsdp = (
            basic_utils.PRECISION_MAP[precision_fsdp] if precision_fsdp is not None else self.precision
        )

        # precision for AMP training
        if precision_amp is None or precision_amp == precision:
            # AMP is disabled during training
            self.precision_amp = None
        else:
            self.precision_amp = basic_utils.PRECISION_MAP[precision_amp]
            if self.precision != torch.float32:
                logger.warning(
                    f"Autocast to {self.precision_amp} is enabled and model and data are cast to {self.precision}. "
                    f"It is recommended to set `config.model.precision` to `float32`."
                )

        # precision for AMP inference
        if precision_amp_infer is None or precision_amp_infer == precision:
            # AMP is disabled during inference
            self.precision_amp_infer = None
        else:
            self.precision_amp_infer = basic_utils.PRECISION_MAP[precision_amp_infer]

        # precision for AMP en-/decoder (e.g., for VAEs or text encoders)
        if precision_amp_enc is None or precision_amp_enc == precision:
            # AMP is disabled during en-/decoding
            self.precision_amp_enc = None
        else:
            self.precision_amp_enc = basic_utils.PRECISION_MAP[precision_amp_enc]

        logger.critical(
            f"Model and data precision: {self.precision}. AMP training precision: {self.precision_amp}. "
            f"AMP en-/decoder precision: {self.precision_amp_enc}. AMP inference precision: {self.precision_amp_infer}. "
            f"FSDP precision: {self.precision_fsdp}."
        )

    @property
    def teacher_config(self) -> dict:
        teacher_config = self.config.net
        if self.config.teacher is not None:
            logger.critical("Using teacher config (usually due to teacher architecture being different from student)")
            teacher_config = self.config.teacher
        return teacher_config

    def build_teacher(self):
        # instantiate the teacher
        logger.info("Instantiating the teacher")
        meta_init_teacher = self.config.add_teacher_to_fsdp_dict and self.config.fsdp_meta_init
        logger.info(
            f"build_teacher: add_teacher_to_fsdp_dict={self.config.add_teacher_to_fsdp_dict}, "
            f"fsdp_meta_init={self.config.fsdp_meta_init}, meta_init_teacher={meta_init_teacher}"
        )
        with self._get_meta_init_context(meta_init_teacher):
            self.teacher = instantiate(self.teacher_config)
        logger.info(
            f"Teacher guidance scale set to {self.config.guidance_scale} (skip-layer guidance: {self.config.skip_layers})"
        )

        # load pre-trained teacher model
        model_path = self.config.pretrained_model_path
        if model_path is not None and len(model_path) > 0:
            FastGenModel._load_pretrained_model(self.teacher, model_path, fsdp_meta_init=meta_init_teacher)
        self.teacher.eval().requires_grad_(False)
        synchronize()

    def load_student_weights_and_ema(self):
        # path to an external network ckpt different from teacher (e.g. pretrained kd, pretrained self-forcing, etc.)
        pretrained_student_net_path = self.config.pretrained_student_net_path
        has_student_path = pretrained_student_net_path is not None and len(pretrained_student_net_path) > 0
        # path to the pretrained teacher model ckpt
        pretrained_model_path = self.config.pretrained_model_path
        has_model_path = pretrained_model_path is not None and len(pretrained_model_path) > 0

        if self.config.load_student_weights:
            logger.info("Loading student weights")
            if has_student_path:
                FastGenModel._load_pretrained_model(
                    self.net, pretrained_student_net_path, fsdp_meta_init=self.config.fsdp_meta_init
                )
            elif has_model_path:
                if getattr(self, "teacher", None) is not None:
                    logger.info("Loading student weights from teacher weights")
                    # initialize the consistency network with the teacher weights
                    # Only rank 0 loads weights if using meta initialization
                    if (not self.config.fsdp_meta_init) or is_rank0():
                        net_load_info = self.net.load_state_dict(self.teacher.state_dict(), strict=False)
                        logger.success(f"Net initializing info: {net_load_info}")
                else:
                    FastGenModel._load_pretrained_model(
                        self.net, pretrained_model_path, fsdp_meta_init=self.config.fsdp_meta_init
                    )
            else:
                logger.warning(
                    "No student weights specified. This might be intended if the student initialization already "
                    "loads pretrained weights (e.g., from diffusers)."
                )

            if has_student_path or has_model_path:
                synchronize()

        elif has_student_path:
            logger.warning("Ignoring `pretrained_student_net_path` since `load_student_weights` is False.")
        elif has_model_path and getattr(self, "teacher", None) is None:
            logger.warning("Ignoring `pretrained_model_path` since `load_student_weights` is False.")

        # load EMA weights
        self._setup_ema()

    def build_model(self):
        # instantiate the generator network
        logger.info("Instantiating the generator network")
        with self._get_meta_init_context():
            self.net = instantiate(self.config.net)
        no_grad_params = [n for n, p in self.net.named_parameters() if not p.requires_grad]
        if any(no_grad_params):
            logger.warning(
                f"The `requires_grad` attribute of these parameters is `False` at initialization and will be set to `True`: {no_grad_params}"
            )
        self.net.train().requires_grad_(True)

        # initialize the preprocessors if they exist, only in the net model
        # this is useful for models that require specific preprocessing. e.g. SD model for image / text encoding
        if hasattr(self.net, "init_preprocessors") and self.config.enable_preprocessors:
            self.net.init_preprocessors()

    def on_train_begin(self, is_fsdp=False):
        self._is_fsdp = is_fsdp  # Store for later use (e.g., to skip EMA during inference)
        ctx = dict(dtype=self.precision, device=self.device)

        if is_fsdp:
            # Cast fsdp_dict modules to precision_fsdp before FSDP wrapping (when AMP is disabled).
            # This sets the parameter storage dtype that FSDP will preserve for shards and gradient reduction.
            #
            # Skip on ranks where parameters are on the meta device (fsdp_meta_init=True,
            # non-rank-0): the cast on a meta tensor is logically a no-op, but in practice
            # iterating modules and casting them appears to materialize tensors during the
            # bf16->fp32 conversion, defeating meta-init's memory savings and pushing the
            # cgroup over its limit.  After FSDP wrap, sync_module_states broadcasts rank-0's
            # post-cast fp32 weights into the per-rank shards, which is what we want anyway.
            for net_name, net in self.fsdp_dict.items():
                has_meta_params = any(p.is_meta for p in net.parameters())
                if has_meta_params:
                    logger.info(
                        f"Skipping pre-FSDP cast for {net_name}: parameters on meta "
                        f"device (will be filled by FSDP sync_module_states broadcast)"
                    )
                else:
                    logger.debug(f"Casting {net_name} to dtype={self.precision_fsdp} (pre-FSDP).")
                    net.to(dtype=self.precision_fsdp)
                # Synchronize on every rank, even those that skipped the cast,
                # so the per-iteration barrier (which the .to() side enters
                # via synchronize() below) is matched by the meta side too.
                # Otherwise rank 0 calls synchronize() and meta ranks don't,
                # causing a NCCL ALLREDUCE deadlock.
                synchronize()

            # Move EMA networks as they aren't handled by FSDP
            for net_name, net in self.ema_dict.items():
                logger.debug(f"Starting moving EMA {net_name} to device: {self.device}.")
                net.to(device=self.device)
                synchronize()
                logger.debug(f"Completed moving EMA {net_name} to device: {self.device}.")
        else:
            # If no FSDP, we need to manually handle casting and device management
            for net_name, net in self.fsdp_dict.items():
                logger.debug(f"Starting moving {net_name} to context: {ctx}.")
                net.to(**ctx)
                synchronize()
                logger.debug(f"Completed moving {net_name} to context: {ctx}.")

        # Handle teacher separately if it's not in the FSDP dict
        if getattr(self, "teacher", None) is not None:
            fsdp_dict_keys = list(self.fsdp_dict.keys())
            logger.info(
                f"Teacher check: add_teacher_to_fsdp_dict={self.config.add_teacher_to_fsdp_dict}, "
                f"fsdp_dict keys={fsdp_dict_keys}, teacher in fsdp_dict={'teacher' in fsdp_dict_keys}"
            )
            if "teacher" not in self.fsdp_dict:
                # No gradients for teacher, can put in lower precision
                logger.info(f"Started converting teacher to context: {ctx}.")
                self.teacher.to(**ctx)
                synchronize()

        # For networks that don't need gradients, we always manually handle casting and device management
        if hasattr(self.net, "init_preprocessors") and self.config.enable_preprocessors:
            logger.debug(f"Starting moving preprocessors to context: {ctx}.")
            if hasattr(self.net, "vae"):
                self.net.vae.to(**ctx)
                synchronize()
            if hasattr(self.net, "text_encoder"):
                self.net.text_encoder.to(**ctx)
                synchronize()
            if hasattr(self.net, "image_encoder"):
                self.net.image_encoder.to(**ctx)
                synchronize()
            logger.debug(f"Completed moving preprocessors to context: {ctx}.")

        synchronize()

    def gen_data_from_net(
        self,
        input_student: torch.Tensor,
        t_student: torch.Tensor,
        condition: Optional[Any] = None,
    ) -> torch.Tensor:
        gen_data = self.net(input_student, t_student, condition=condition, fwd_pred_type="x0")
        return gen_data

    @classmethod
    def _student_sample_loop(
        cls,
        net: FastGenNetwork,
        x: torch.Tensor,
        t_list: torch.Tensor,
        condition: Any = None,
        student_sample_type: str = "sde",
        **kwargs,
    ) -> torch.Tensor:
        """
        Sample loop for the student network.

        Args:
            net: The FastGenNetwork network
            x: The latents to start from
            t_list: Timesteps to sample
            condition: Optional conditioning information
            student_sample_type: Type of student multistep sampling

        Returns:
            The sampled data
        """
        batch_size = x.shape[0]

        # Check if network has custom conditioning preservation hooks
        # This allows video I2V/v2w models to handle conditioning without
        # complicating this generic loop with model-specific logic
        has_preserve_hook = hasattr(net, "preserve_conditioning")

        x_pred = x
        for t_cur, t_next in zip(t_list[:-1], t_list[1:]):
            # Forward pass to get x0 prediction
            t_batch = t_cur.expand(batch_size)
            x_pred = net(x, t_batch, condition=condition, fwd_pred_type="x0")

            # Allow network to preserve conditioning frames
            if has_preserve_hook:
                x_pred = net.preserve_conditioning(x_pred, condition)

            # One step reverse process
            if t_next > 0:
                t_next_batch = t_next.expand(batch_size)
                if student_sample_type == "sde":
                    eps_infer = torch.randn_like(x_pred)
                elif student_sample_type == "ode":
                    eps_infer = net.noise_scheduler.x0_to_eps(xt=x, x0=x_pred, t=t_batch)
                else:
                    raise NotImplementedError(
                        f"student_sample_type must be one of 'sde', 'ode' but got {student_sample_type}"
                    )
                x = net.noise_scheduler.forward_process(x_pred, eps_infer, t_next_batch)

                # Preserve conditioning frames after adding noise
                if has_preserve_hook:
                    x = net.preserve_conditioning(x, condition)

        return x_pred

    @classmethod
    def generator_fn(
        cls,
        net: FastGenNetwork,
        noise: torch.Tensor,
        student_sample_steps: int = 1,
        t_list: Optional[List[float]] = None,
        data: torch.Tensor = None,
        precision_amp: Optional[torch.dtype] = None,
        **kwargs,
    ) -> torch.Tensor:
        """
        Single-step or multistep generation with the distilled network.

        Args:
            net: The FastGenNetwork network
            noise: Pure noise to start from (zero-mean, unit-variance Gaussian)
            student_sample_steps: Number of student diffusion steps
            t_list: Timesteps to sample (defaults to None: use noise_schedule.get_t_list() instead)
            data (torch.Tensor, optional): Additional data to add to initial latents.
                Useful for inpainting or other conditional tasks. Defaults to None.
            precision_amp (torch.dtype, optional): If not None, uses autocast with this dtype for inference.
            **kwargs: Additional keyword arguments passed to the network.

        Returns:
            Generated sample from the distilled single-step or multistep student.
        """
        with basic_utils.inference_mode(net, precision_amp=precision_amp, device_type=noise.device.type):
            # Default timestep schedule
            if t_list is None:
                t_list = net.noise_scheduler.get_t_list(sample_steps=student_sample_steps, device=noise.device)
            else:
                assert (
                    len(t_list) - 1 == student_sample_steps
                ), f"t_list length (excluding zero) != student_sample_steps: {len(t_list) - 1} != {student_sample_steps}"
                t_list = torch.tensor(t_list, device=noise.device, dtype=net.noise_scheduler.t_precision)
            assert t_list[-1].item() == 0, "t_list[-1] must be zero"

            # Initialize with noise scaling
            latents = net.noise_scheduler.latents(noise=noise, t_init=t_list[0])

            # Add optional data (e.g., for inpainting)
            if data is not None:
                latents = latents + data

            # Multistep sampling loop
            return cls._student_sample_loop(net, latents, t_list=t_list, **kwargs).to(dtype=noise.dtype)

    def sample(self, net: FastGenNetwork, noise: torch.Tensor, **kwargs) -> torch.Tensor:
        assert hasattr(net, "sample")
        with basic_utils.inference_mode(net, precision_amp=self.precision_amp_infer):
            return net.sample(
                noise,
                guidance_scale=self.config.guidance_scale,
                **kwargs,
            ).to(dtype=noise.dtype)

    def _prepare_training_data(self, data: Dict[str, Any]) -> tuple[torch.Tensor, Any, Any]:
        """Prepare training data and conditions from input data dict.

        Args:
            data: Data dict containing real data, conditions, etc.

        Returns:
            tuple of (real_data, condition, neg_condition)
        """
        real_data = data["real"]
        if getattr(self.net, "is_vid2vid", False):
            # handle vid2vid
            vid_context = data["vid_context"]  # this is processed in trainer.py
            condition = {
                "text_embeds": data["condition"],
                "vid_context": vid_context,
            }
            neg_condition = {
                "text_embeds": data["neg_condition"],
                "vid_context": vid_context,
            }
        elif getattr(self.net, "is_i2v", False):
            # handle i2v (WanI2V style)
            first_frame_cond = data["first_frame_cond"]  # this is processed in trainer.py
            condition = {
                "text_embeds": data["condition"],
                "first_frame_cond": first_frame_cond,
            }
            neg_condition = {
                "text_embeds": data["neg_condition"],
                "first_frame_cond": first_frame_cond,
            }
            if hasattr(self.net, "image_encoder"):
                condition["encoder_hidden_states_image"] = data["encoder_hidden_states_image"]
                neg_condition["encoder_hidden_states_image"] = data["encoder_hidden_states_image"]
        elif getattr(self.net, "is_video2world", False):
            # handle video2world (Cosmos style)
            conditioning_latents = data["conditioning_latents"]  # this is processed in trainer.py
            condition_mask = data["condition_mask"]  # this is processed in trainer.py
            condition = {
                "text_embeds": data["condition"],
                "conditioning_latents": conditioning_latents,
                "condition_mask": condition_mask,
            }
            neg_condition = {
                "text_embeds": data["neg_condition"],
                "conditioning_latents": conditioning_latents,
                "condition_mask": condition_mask,
            }
        else:
            # handle other cases
            condition = data["condition"]
            neg_condition = data["neg_condition"]

        return real_data, condition, neg_condition

    @abstractmethod
    def _get_outputs(
        self,
        gen_data: torch.Tensor,
        input_student: torch.Tensor = None,
        condition: Any = None,
    ) -> Dict[str, torch.Tensor | Callable]:
        """
        Get model outputs as a dictionary of tensors.
        """

    @abstractmethod
    def single_train_step(
        self, data: Dict[str, Any], iteration: int
    ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]:
        """
        Single training step for the model.

        Args:
            data (Dict[str, Any]): Data dict for the current iteration.
            iteration (int): Current training iteration

        Returns:
            loss_map (dict[str, torch.Tensor]): Dictionary containing the loss values
            outputs (dict[str, torch.Tensor]): Dictionary containing the network output
        """

    def init_optimizers(self):
        """Initialize optimizers, lr_schedulers and grad_scalers"""
        # instantiate the optimizer for the generator network
        self.net_optimizer = instantiate(self.config.net_optimizer, model=self.net)
        # instantiate the lr scheduler for the generator network
        self.net_lr_scheduler = get_scheduler(self.net_optimizer, self.config.net_scheduler)
        # instantiate the gradient scaler (only fp16 needs grad_scaler)
        grad_scaler_required = self.precision == torch.float16 or self.precision_amp == torch.float16
        grad_scaler_enabled = self.config.grad_scaler_enabled and grad_scaler_required
        if grad_scaler_required:
            if grad_scaler_enabled:
                logger.info(
                    f"Grad scaler enabled with init scale {self.config.grad_scaler_init_scale} and growth interval {self.config.grad_scaler_growth_interval}."
                )
            else:
                logger.warning(
                    f"Grad scaler disabled but recommended when using float16 precision (precision={self.precision}, precision_amp={self.precision_amp})."
                )
        self.grad_scaler = torch.amp.GradScaler(
            init_scale=self.config.grad_scaler_init_scale,
            growth_interval=self.config.grad_scaler_growth_interval,
            enabled=grad_scaler_enabled,
        )

    def get_optimizers(self, iteration: int) -> list[torch.optim.Optimizer]:
        """
        Get the optimizers for the current iteration
        Args:
            iteration (int): The current training iteration
        """
        return [self.net_optimizer]

    def get_lr_schedulers(self, iteration: int) -> list[torch.optim.lr_scheduler]:
        """
        Get the lr schedulers for the current iteration
        Args:
            iteration (int): The current training iteration
        """
        return [self.net_lr_scheduler]

    def optimizers_zero_grad(self, iteration: int) -> None:
        """
        Zero the gradients of the optimizers based on the iteration
        """
        for optimizer in self.get_optimizers(iteration):
            optimizer.zero_grad(set_to_none=True)

    def should_use_grad_scaler(self, optimizer: torch.optim.Optimizer) -> bool:
        """
        Check if grad_scaler should be used for the given optimizer.

        GradScaler only works with float32 gradients. When model weights are in
        FP16/BF16 (e.g., with FSDP2 mixed precision), gradients are also in that
        dtype and grad_scaler cannot be used.

        Args:
            optimizer: The optimizer to check

        Returns:
            True if grad_scaler should be used, False otherwise
        """
        if not self.grad_scaler.is_enabled():
            return False

        # Check if any gradient is not float32
        for param_group in optimizer.param_groups:
            for param in param_group["params"]:
                if param.grad is not None and param.grad.dtype != torch.float32:
                    return False
        return True

    def optimizers_schedulers_step(self, iteration: int) -> None:
        """
        Step the optimizer and scheduler step based on the iteration,
        and gradient scaler is also updated
        """
        for optimizer in self.get_optimizers(iteration):
            if self.should_use_grad_scaler(optimizer):
                self.grad_scaler.step(optimizer)
                self.grad_scaler.update()
            else:
                optimizer.step()

        for scheduler in self.get_lr_schedulers(iteration):
            scheduler.step()

    @staticmethod
    def _load_pretrained_model(
        model: torch.nn.Module,
        pretrained_model_path: str,
        device: Optional[torch.device] = "cpu",
        fsdp_meta_init: bool = False,
    ) -> None:
        """
        Load the pre-trained model from the given path
        Args:
            model (torch.nn.Module): The model to load
            pretrained_model_path (str): The path to the pretrained model
            device (Optional[torch.device]): The device to load the model on
            fsdp_meta_init (bool): Whether to use meta initialization for FSDP
        """
        # Only rank-0 loads weights if using meta initialization
        if (not fsdp_meta_init) or is_rank0():
            logger.info(f"Loading the pretrained diffusion model from {pretrained_model_path}")
            if pretrained_model_path.startswith("s3://"):
                key = pretrained_model_path.split("/")[-1]
                local_path = os.path.join(
                    os.environ.get("LIPFORCING_OUTPUT_ROOT", "outputs"), "model", key.split("/")[-1]
                )
                if os.path.exists(local_path):
                    logger.info(f"Model already exists at {local_path}, loading from local cache")
                    model_dict = torch.load(local_path, weights_only=True, map_location=device)
                else:
                    model_dict = torch.load(s3_load(pretrained_model_path), weights_only=True, map_location=device)
                    os.makedirs(os.path.dirname(local_path), exist_ok=True)
                    torch.save(model_dict, local_path)
            else:
                assert os.path.isfile(pretrained_model_path), f"{pretrained_model_path} is not a valid file"
                model_dict = torch.load(pretrained_model_path, weights_only=True, map_location=device)

            for k, v in model_dict.items():
                if isinstance(v, torch.Tensor) and v.ndim == 0:
                    # since FSDP2 cannot handle 0-dim. tensors, we adapted all network definitions to use 1-dim.
                    # tensors with numel equal to 1
                    model_dict[k] = v.unsqueeze(0)
                    logger.debug(f"Changed {k} from 0-dim. tensor to 1-dim. tensor with numel equal to 1.")

            model_load_info = model.load_state_dict(model_dict, strict=False)
            torch.cuda.empty_cache()
            logger.success(f"Model loading info: {model_load_info}")
        synchronize()

    def autocast(self):
        """Return the autocast context manager for training"""
        return torch.autocast(
            device_type=self.device.type,
            dtype=self.precision_amp,
            enabled=self.precision_amp is not None,
        )

    @property
    def ema_dict(self):
        """Return dict containing all EMA networks"""
        return {name: getattr(self, name) for name in self.use_ema}

    @property
    def net_inference(self):
        """Return the network to use for inference.

        Uses EMA network when available and not using FSDP.
        TODO: When FSDP is enabled, EMA networks are not wrapped and have dtype issues,
        so we fall back to the main network which is properly FSDP-wrapped.
        Note that inference of the EMA network is possible using the scripts in scripts/inference/.
        """
        use_ema_for_inference = self.use_ema and not getattr(self, "_is_fsdp", False)
        return getattr(self, self.use_ema[0]) if use_ema_for_inference else self.net

    @property
    def fsdp_dict(self):
        """Return dict containing all networks to be sharded.

        By default, this is the same as the model dict.
        If the model has a teacher and add_teacher_to_fsdp_dict is True, the teacher is added to the dict.
        """
        model_dict = self.model_dict
        if getattr(self, "teacher", None) is not None and self.config.add_teacher_to_fsdp_dict:
            model_dict["teacher"] = self.teacher
        return model_dict

    @property
    def model_dict(self):
        """Return the model dict containing the student and EMA networks"""
        return torch.nn.ModuleDict({"net": self.net, **self.ema_dict})

    @property
    def optimizer_dict(self):
        """Return a dict containing all the optimizers"""
        return {
            "net": self.net_optimizer,
        }

    @property
    def scheduler_dict(self):
        """Return a dict containing all the lr schedulers"""
        return {
            "net": self.net_lr_scheduler,
        }