docs/finetrainers-src-codebase/examples/inference/inference.py

import argparse
import json
import os
import time
import traceback
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import datasets.distributed
import torch
import wandb
from diffusers.hooks import HookRegistry, ModelHook
from diffusers.utils import export_to_video

from finetrainers import data, get_logger, logging, parallel, patches, utils
from finetrainers.args import AttentionProviderInference
from finetrainers.config import ModelType
from finetrainers.models import ModelSpecification, attention_provider
from finetrainers.models.cogvideox import CogVideoXModelSpecification
from finetrainers.models.cogview4 import CogView4ModelSpecification
from finetrainers.models.flux import FluxModelSpecification
from finetrainers.models.wan import WanModelSpecification
from finetrainers.parallel import ParallelBackendEnum
from finetrainers.state import ParallelBackendType
from finetrainers.utils import ArgsConfigMixin


logger = get_logger()


def main():
    try:
        import multiprocessing

        multiprocessing.set_start_method("fork")
    except Exception as e:
        logger.error(
            f'Failed to set multiprocessing start method to "fork". This can lead to poor performance, high memory usage, or crashes. '
            f"See: https://pytorch.org/docs/stable/notes/multiprocessing.html\n"
            f"Error: {e}"
        )

    try:
        args = BaseArgs()
        args.parse_args()

        model_specification_cls = get_model_specifiction_cls(args.model_name, args.inference_type)
        model_specification = model_specification_cls(
            pretrained_model_name_or_path=args.pretrained_model_name_or_path,
            tokenizer_id=args.tokenizer_id,
            tokenizer_2_id=args.tokenizer_2_id,
            tokenizer_3_id=args.tokenizer_3_id,
            text_encoder_id=args.text_encoder_id,
            text_encoder_2_id=args.text_encoder_2_id,
            text_encoder_3_id=args.text_encoder_3_id,
            transformer_id=args.transformer_id,
            vae_id=args.vae_id,
            text_encoder_dtype=args.text_encoder_dtype,
            text_encoder_2_dtype=args.text_encoder_2_dtype,
            text_encoder_3_dtype=args.text_encoder_3_dtype,
            transformer_dtype=args.transformer_dtype,
            vae_dtype=args.vae_dtype,
            revision=args.revision,
            cache_dir=args.cache_dir,
        )

        inferencer = Inference(args, model_specification)
        inferencer.run()

    except KeyboardInterrupt:
        logger.info("Received keyboard interrupt. Exiting...")
    except Exception as e:
        logger.error(f"An error occurred during training: {e}")
        logger.error(traceback.format_exc())


class InferenceType(str, Enum):
    TEXT_TO_IMAGE = "text_to_image"
    TEXT_TO_VIDEO = "text_to_video"
    IMAGE_TO_VIDEO = "image_to_video"


# We do a union because every ArgsConfigMixin registered to BaseArgs can be looked up using the `__getattribute__` override
BaseArgsType = Union[
    "BaseArgs", "ParallelArgs", "ModelArgs", "InferenceArgs", "AttentionProviderArgs", "TorchConfigArgs"
]

_DTYPE_MAP = {
    "bf16": torch.bfloat16,
    "fp16": torch.float16,
    "fp32": torch.float32,
    "float8_e4m3fn": torch.float8_e4m3fn,
    "float8_e5m2": torch.float8_e5m2,
}


SUPPORTED_MODEL_CONFIGS = {
    ModelType.COGVIDEOX: {
        InferenceType.TEXT_TO_VIDEO: CogVideoXModelSpecification,
    },
    ModelType.COGVIEW4: {
        InferenceType.TEXT_TO_IMAGE: CogView4ModelSpecification,
    },
    ModelType.FLUX: {
        InferenceType.TEXT_TO_IMAGE: FluxModelSpecification,
    },
    ModelType.WAN: {
        InferenceType.TEXT_TO_VIDEO: WanModelSpecification,
        InferenceType.IMAGE_TO_VIDEO: WanModelSpecification,
    },
}


def get_model_specifiction_cls(model_name: str, inference_type: InferenceType) -> ModelSpecification:
    """
    Get the model specification class for the given model name and inference type.
    """
    if model_name not in SUPPORTED_MODEL_CONFIGS:
        raise ValueError(
            f"Model {model_name} not supported. Supported models are: {list(SUPPORTED_MODEL_CONFIGS.keys())}"
        )
    if inference_type not in SUPPORTED_MODEL_CONFIGS[model_name]:
        raise ValueError(
            f"Inference type {inference_type} not supported for model {model_name}. Supported inference types are: {list(SUPPORTED_MODEL_CONFIGS[model_name].keys())}"
        )
    return SUPPORTED_MODEL_CONFIGS[model_name][inference_type]


@dataclass
class State:
    # Parallel state
    parallel_backend: ParallelBackendType = None

    # Training state
    generator: torch.Generator = None


class Inference:
    def __init__(self, args: BaseArgsType, model_specification: ModelSpecification):
        self.args = args
        self.model_specification = model_specification
        self.state = State()

        self.pipeline = None
        self.dataset = None
        self.dataloader = None

        self._init_distributed()
        self._init_config_options()

        patches.perform_patches_for_inference(args, self.state.parallel_backend)

    def run(self) -> None:
        try:
            self._prepare_pipeline()
            self._prepare_distributed()
            self._prepare_dataset()
            self._inference()
        except Exception as e:
            logger.error(f"Error during inference: {e}")
            self.state.parallel_backend.destroy()
            raise e

    def _prepare_pipeline(self) -> None:
        logger.info("Initializing pipeline")

        transformer = self.model_specification.load_diffusion_models()["transformer"]
        self.pipeline = self.model_specification.load_pipeline(
            transformer=transformer,
            enable_slicing=self.args.enable_slicing,
            enable_tiling=self.args.enable_tiling,
            enable_model_cpu_offload=False,  # TODO(aryan): handle model/sequential/group offloading
            training=False,
        )

    def _prepare_distributed(self) -> None:
        parallel_backend = self.state.parallel_backend
        cp_mesh = parallel_backend.get_mesh("cp") if parallel_backend.context_parallel_enabled else None

        if parallel_backend.context_parallel_enabled:
            cp_mesh = parallel_backend.get_mesh()["cp"]
            parallel_backend.apply_context_parallel(self.pipeline.transformer, cp_mesh)

        registry = HookRegistry.check_if_exists_or_initialize(self.pipeline.transformer)
        hook = AttentionProviderHook(
            self.args.attn_provider, cp_mesh, self.args.cp_rotate_method, self.args.cp_reduce_precision
        )
        registry.register_hook(hook, "attn_provider")

        self._maybe_torch_compile()

        self._init_logging()
        self._init_trackers()
        self._init_directories()

    def _prepare_dataset(self) -> None:
        logger.info("Preparing dataset for inference")
        parallel_backend = self.state.parallel_backend

        dp_mesh = None
        if parallel_backend.data_replication_enabled:
            dp_mesh = parallel_backend.get_mesh("dp_replicate")
        if dp_mesh is not None:
            local_rank, dp_world_size = dp_mesh.get_local_rank(), dp_mesh.size()
        else:
            local_rank, dp_world_size = 0, 1

        dataset = data.ValidationDataset(self.args.dataset_file)
        dataset._data = datasets.distributed.split_dataset_by_node(dataset._data, local_rank, dp_world_size)
        dataloader = data.DPDataLoader(
            local_rank,
            dataset,
            batch_size=1,
            num_workers=0,  # TODO(aryan): handle dataloader_num_workers
            collate_fn=lambda items: items,
        )

        self.dataset = dataset
        self.dataloader = dataloader

    def _inference(self) -> None:
        parallel_backend = self.state.parallel_backend
        seed = self.args.seed if self.args.seed is not None else 0
        generator = torch.Generator(device=parallel_backend.device).manual_seed(seed)

        if parallel_backend._dp_degree > 1:
            dp_mesh = parallel_backend.get_mesh("dp")
            dp_local_rank, dp_world_size = dp_mesh.get_local_rank(), dp_mesh.size()
        else:
            dp_mesh = None
            dp_local_rank, dp_world_size = parallel_backend.local_rank, 1

        self.pipeline.to(parallel_backend.device)

        memory_statistics = utils.get_memory_statistics()
        logger.info(f"Memory before inference start: {json.dumps(memory_statistics, indent=4)}")

        data_iterator = iter(self.dataloader)
        main_process_prompts_to_filenames = {}  # Used to save model card
        all_processes_artifacts = []  # Used to gather artifacts from all processes

        while True:
            inference_data = next(data_iterator, None)
            if inference_data is None:
                break

            inference_data = inference_data[0]
            with torch.inference_mode():
                inference_artifacts = self.model_specification.validation(
                    pipeline=self.pipeline, generator=generator, **inference_data
                )

            if dp_local_rank != 0:
                continue

            PROMPT = inference_data["prompt"]
            IMAGE = inference_data.get("image", None)
            VIDEO = inference_data.get("video", None)
            EXPORT_FPS = inference_data.get("export_fps", 30)

            # 2.1. If there are any initial images or videos, they will be logged to keep track of them as
            # conditioning for generation.
            prompt_filename = utils.string_to_filename(PROMPT)[:25]
            artifacts = {
                "input_image": data.ImageArtifact(value=IMAGE),
                "input_video": data.VideoArtifact(value=VIDEO),
            }

            # 2.2. Track the artifacts generated from inference
            for i, inference_artifact in enumerate(inference_artifacts):
                if inference_artifact.value is None:
                    continue
                artifacts.update({f"artifact_{i}": inference_artifact})

            # 2.3. Save the artifacts to the output directory and create appropriate logging objects
            # TODO(aryan): Currently, we only support WandB so we've hardcoded it here. Needs to be revisited.
            for index, (key, artifact) in enumerate(list(artifacts.items())):
                assert isinstance(artifact, (data.ImageArtifact, data.VideoArtifact))
                if artifact.value is None:
                    continue

                time_, rank, ext = int(time.time()), parallel_backend.rank, artifact.file_extension
                filename = f"inference-{rank}-{index}-{prompt_filename}-{time_}.{ext}"
                output_filename = os.path.join(self.args.output_dir, filename)

                if parallel_backend.is_main_process and ext in ["mp4", "jpg", "jpeg", "png"]:
                    main_process_prompts_to_filenames[PROMPT] = filename

                if isinstance(artifact, data.ImageArtifact):
                    artifact.value.save(output_filename)
                    all_processes_artifacts.append(wandb.Image(output_filename, caption=PROMPT))
                elif isinstance(artifact, data.VideoArtifact):
                    export_to_video(artifact.value, output_filename, fps=EXPORT_FPS)
                    all_processes_artifacts.append(wandb.Video(output_filename, caption=PROMPT))

        # 3. Cleanup & log artifacts
        parallel_backend.wait_for_everyone()

        memory_statistics = utils.get_memory_statistics()
        logger.info(f"Memory after inference end: {json.dumps(memory_statistics, indent=4)}")

        # Gather artifacts from all processes. We also need to flatten them since each process returns a list of artifacts.
        all_artifacts = [None] * dp_world_size
        if dp_world_size > 1:
            torch.distributed.all_gather_object(all_artifacts, all_processes_artifacts)
        else:
            all_artifacts = [all_processes_artifacts]
        all_artifacts = [artifact for artifacts in all_artifacts for artifact in artifacts]

        if parallel_backend.is_main_process:
            tracker_key = "inference"
            artifact_log_dict = {}

            image_artifacts = [artifact for artifact in all_artifacts if isinstance(artifact, wandb.Image)]
            if len(image_artifacts) > 0:
                artifact_log_dict["images"] = image_artifacts
            video_artifacts = [artifact for artifact in all_artifacts if isinstance(artifact, wandb.Video)]
            if len(video_artifacts) > 0:
                artifact_log_dict["videos"] = video_artifacts
            parallel_backend.log({tracker_key: artifact_log_dict}, step=0)

        parallel_backend.wait_for_everyone()

    def _init_distributed(self) -> None:
        world_size = int(os.environ.get("WORLD_SIZE", torch.cuda.device_count()))

        # TODO(aryan): handle other backends
        backend_cls: parallel.ParallelBackendType = parallel.get_parallel_backend_cls(self.args.parallel_backend)
        self.state.parallel_backend = backend_cls(
            world_size=world_size,
            pp_degree=self.args.pp_degree,
            dp_degree=self.args.dp_degree,
            dp_shards=self.args.dp_shards,
            cp_degree=self.args.cp_degree,
            tp_degree=self.args.tp_degree,
            backend="nccl",
            timeout=self.args.init_timeout,
            logging_dir=self.args.logging_dir,
            output_dir=self.args.output_dir,
        )

        if self.args.seed is not None:
            self.state.parallel_backend.enable_determinism(self.args.seed)

    def _init_logging(self) -> None:
        logging._set_parallel_backend(self.state.parallel_backend)
        logging.set_dependency_log_level(self.args.verbose, self.state.parallel_backend.is_local_main_process)
        logger.info("Initialized Finetrainers")

    def _init_trackers(self) -> None:
        # TODO(aryan): handle multiple trackers
        trackers = [self.args.report_to]
        experiment_name = self.args.tracker_name or "finetrainers-inference"
        self.state.parallel_backend.initialize_trackers(
            trackers, experiment_name=experiment_name, config=self.args.to_dict(), log_dir=self.args.logging_dir
        )

    def _init_directories(self) -> None:
        if self.state.parallel_backend.is_main_process:
            self.args.output_dir = Path(self.args.output_dir)
            self.args.output_dir.mkdir(parents=True, exist_ok=True)

    def _init_config_options(self) -> None:
        # Enable TF32 for faster training on Ampere GPUs: https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
        if self.args.allow_tf32 and torch.cuda.is_available():
            torch.backends.cuda.matmul.allow_tf32 = True
        torch.set_float32_matmul_precision(self.args.float32_matmul_precision)

    def _maybe_torch_compile(self):
        for model_name, compile_scope in zip(self.args.compile_modules, self.args.compile_scopes):
            model = getattr(self.pipeline, model_name, None)
            if model is not None:
                logger.info(f"Applying torch.compile to '{model_name}' with scope '{compile_scope}'.")
                compiled_model = utils.apply_compile(model, compile_scope)
                setattr(self.pipeline, model_name, compiled_model)


class AttentionProviderHook(ModelHook):
    def __init__(
        self,
        provider: str,
        mesh: Optional[torch.distributed.device_mesh.DeviceMesh] = None,
        rotate_method: str = "allgather",
        reduce_precision: bool = False,
    ):
        super().__init__()
        self.provider = provider
        self.mesh = mesh
        self.rotate_method = rotate_method
        self.convert_to_fp32 = not reduce_precision

    def new_forward(self, module: torch.nn.Module, *args, **kwargs) -> Any:
        with attention_provider(
            self.provider, mesh=self.mesh, convert_to_fp32=self.convert_to_fp32, rotate_method=self.rotate_method
        ):
            return self.fn_ref.original_forward(*args, **kwargs)


class ParallelArgs(ArgsConfigMixin):
    """
    Args:
        parallel_backend (`str`, defaults to "accelerate"):
            The parallel backend to use for inference. Choose between ['accelerate', 'ptd'].
        pp_degree (`int`, defaults to 1):
            The degree of pipeline parallelism.
        dp_degree (`int`, defaults to 1):
            The degree of data parallelism (number of model replicas).
        dp_shards (`int`, defaults to 1):
            The number of data parallel shards (number of model partitions).
        cp_degree (`int`, defaults to 1):
            The degree of context parallelism.
    """

    parallel_backend: ParallelBackendEnum = ParallelBackendEnum.ACCELERATE
    pp_degree: int = 1
    dp_degree: int = 1
    dp_shards: int = 1
    cp_degree: int = 1
    tp_degree: int = 1
    cp_rotate_method: str = "allgather"
    cp_reduce_precision: bool = False

    def add_args(self, parser: argparse.ArgumentParser) -> None:
        parser.add_argument("--parallel_backend", type=str, default="accelerate", choices=["accelerate", "ptd"])
        parser.add_argument("--pp_degree", type=int, default=1)
        parser.add_argument("--dp_degree", type=int, default=1)
        parser.add_argument("--dp_shards", type=int, default=1)
        parser.add_argument("--cp_degree", type=int, default=1)
        parser.add_argument("--tp_degree", type=int, default=1)
        parser.add_argument("--cp_rotate_method", type=str, default="allgather", choices=["allgather", "alltoall"])
        parser.add_argument("--cp_reduce_precision", action="store_true")

    def map_args(self, argparse_args: argparse.Namespace, mapped_args: "BaseArgs"):
        mapped_args.parallel_backend = argparse_args.parallel_backend
        mapped_args.pp_degree = argparse_args.pp_degree
        mapped_args.dp_degree = argparse_args.dp_degree
        mapped_args.dp_shards = argparse_args.dp_shards
        mapped_args.cp_degree = argparse_args.cp_degree
        mapped_args.tp_degree = argparse_args.tp_degree
        mapped_args.cp_rotate_method = argparse_args.cp_rotate_method
        mapped_args.cp_reduce_precision = argparse_args.cp_reduce_precision

    def validate_args(self, args: "BaseArgs"):
        if args.parallel_backend != "ptd":
            raise ValueError("Only 'ptd' parallel backend is supported for now.")
        if any(x > 1 for x in [args.pp_degree, args.dp_degree, args.dp_shards, args.tp_degree]):
            raise ValueError("Parallel degrees must be 1 except for `cp_degree` for now.")


class ModelArgs(ArgsConfigMixin):
    """
    Args:
        model_name (`str`):
            Name of model to train.
        pretrained_model_name_or_path (`str`):
            Path to pretrained model or model identifier from https://huggingface.co/models. The model should be
            loadable based on specified `model_name`.
        revision (`str`, defaults to `None`):
            If provided, the model will be loaded from a specific branch of the model repository.
        variant (`str`, defaults to `None`):
            Variant of model weights to use. Some models provide weight variants, such as `fp16`, to reduce disk
            storage requirements.
        cache_dir (`str`, defaults to `None`):
            The directory where the downloaded models and datasets will be stored, or loaded from.
        tokenizer_id (`str`, defaults to `None`):
            Identifier for the tokenizer model. This is useful when using a different tokenizer than the default from `pretrained_model_name_or_path`.
        tokenizer_2_id (`str`, defaults to `None`):
            Identifier for the second tokenizer model. This is useful when using a different tokenizer than the default from `pretrained_model_name_or_path`.
        tokenizer_3_id (`str`, defaults to `None`):
            Identifier for the third tokenizer model. This is useful when using a different tokenizer than the default from `pretrained_model_name_or_path`.
        text_encoder_id (`str`, defaults to `None`):
            Identifier for the text encoder model. This is useful when using a different text encoder than the default from `pretrained_model_name_or_path`.
        text_encoder_2_id (`str`, defaults to `None`):
            Identifier for the second text encoder model. This is useful when using a different text encoder than the default from `pretrained_model_name_or_path`.
        text_encoder_3_id (`str`, defaults to `None`):
            Identifier for the third text encoder model. This is useful when using a different text encoder than the default from `pretrained_model_name_or_path`.
        transformer_id (`str`, defaults to `None`):
            Identifier for the transformer model. This is useful when using a different transformer model than the default from `pretrained_model_name_or_path`.
        vae_id (`str`, defaults to `None`):
            Identifier for the VAE model. This is useful when using a different VAE model than the default from `pretrained_model_name_or_path`.
        text_encoder_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
            Data type for the text encoder when generating text embeddings.
        text_encoder_2_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
            Data type for the text encoder 2 when generating text embeddings.
        text_encoder_3_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
            Data type for the text encoder 3 when generating text embeddings.
        transformer_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
            Data type for the transformer model.
        vae_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
            Data type for the VAE model.
        layerwise_upcasting_modules (`List[str]`, defaults to `[]`):
            Modules that should have fp8 storage weights but higher precision computation. Choose between ['transformer'].
        layerwise_upcasting_storage_dtype (`torch.dtype`, defaults to `float8_e4m3fn`):
            Data type for the layerwise upcasting storage. Choose between ['float8_e4m3fn', 'float8_e5m2'].
        layerwise_upcasting_skip_modules_pattern (`List[str]`, defaults to `["patch_embed", "pos_embed", "x_embedder", "context_embedder", "^proj_in$", "^proj_out$", "norm"]`):
            Modules to skip for layerwise upcasting. Layers such as normalization and modulation, when casted to fp8 precision
            naively (as done in layerwise upcasting), can lead to poorer training and inference quality. We skip these layers
            by default, and recommend adding more layers to the default list based on the model architecture.
        enable_slicing (`bool`, defaults to `False`):
            Whether to enable VAE slicing.
        enable_tiling (`bool`, defaults to `False`):
            Whether to enable VAE tiling.
    """

    model_name: str = None
    pretrained_model_name_or_path: str = None
    revision: Optional[str] = None
    variant: Optional[str] = None
    cache_dir: Optional[str] = None
    tokenizer_id: Optional[str] = None
    tokenizer_2_id: Optional[str] = None
    tokenizer_3_id: Optional[str] = None
    text_encoder_id: Optional[str] = None
    text_encoder_2_id: Optional[str] = None
    text_encoder_3_id: Optional[str] = None
    transformer_id: Optional[str] = None
    vae_id: Optional[str] = None
    text_encoder_dtype: torch.dtype = torch.bfloat16
    text_encoder_2_dtype: torch.dtype = torch.bfloat16
    text_encoder_3_dtype: torch.dtype = torch.bfloat16
    transformer_dtype: torch.dtype = torch.bfloat16
    vae_dtype: torch.dtype = torch.bfloat16
    layerwise_upcasting_modules: List[str] = []
    layerwise_upcasting_storage_dtype: torch.dtype = torch.float8_e4m3fn
    # fmt: off
    layerwise_upcasting_skip_modules_pattern: List[str] = ["patch_embed", "pos_embed", "x_embedder", "context_embedder", "time_embed", "^proj_in$", "^proj_out$", "norm"]
    # fmt: on
    enable_slicing: bool = False
    enable_tiling: bool = False

    def add_args(self, parser: argparse.ArgumentParser) -> None:
        parser.add_argument(
            "--model_name", type=str, required=True, choices=[x.value for x in ModelType.__members__.values()]
        )
        parser.add_argument("--pretrained_model_name_or_path", type=str, required=True)
        parser.add_argument("--revision", type=str, default=None, required=False)
        parser.add_argument("--variant", type=str, default=None)
        parser.add_argument("--cache_dir", type=str, default=None)
        parser.add_argument("--tokenizer_id", type=str, default=None)
        parser.add_argument("--tokenizer_2_id", type=str, default=None)
        parser.add_argument("--tokenizer_3_id", type=str, default=None)
        parser.add_argument("--text_encoder_id", type=str, default=None)
        parser.add_argument("--text_encoder_2_id", type=str, default=None)
        parser.add_argument("--text_encoder_3_id", type=str, default=None)
        parser.add_argument("--transformer_id", type=str, default=None)
        parser.add_argument("--vae_id", type=str, default=None)
        parser.add_argument("--text_encoder_dtype", type=str, default="bf16")
        parser.add_argument("--text_encoder_2_dtype", type=str, default="bf16")
        parser.add_argument("--text_encoder_3_dtype", type=str, default="bf16")
        parser.add_argument("--transformer_dtype", type=str, default="bf16")
        parser.add_argument("--vae_dtype", type=str, default="bf16")
        parser.add_argument("--layerwise_upcasting_modules", type=str, default=[], nargs="+", choices=["transformer"])
        parser.add_argument(
            "--layerwise_upcasting_storage_dtype",
            type=str,
            default="float8_e4m3fn",
            choices=["float8_e4m3fn", "float8_e5m2"],
        )
        parser.add_argument(
            "--layerwise_upcasting_skip_modules_pattern",
            type=str,
            default=["patch_embed", "pos_embed", "x_embedder", "context_embedder", "^proj_in$", "^proj_out$", "norm"],
            nargs="+",
        )
        parser.add_argument("--enable_slicing", action="store_true")
        parser.add_argument("--enable_tiling", action="store_true")

    def map_args(self, argparse_args: argparse.Namespace, mapped_args: "BaseArgs"):
        mapped_args.model_name = argparse_args.model_name
        mapped_args.pretrained_model_name_or_path = argparse_args.pretrained_model_name_or_path
        mapped_args.revision = argparse_args.revision
        mapped_args.variant = argparse_args.variant
        mapped_args.cache_dir = argparse_args.cache_dir
        mapped_args.tokenizer_id = argparse_args.tokenizer_id
        mapped_args.tokenizer_2_id = argparse_args.tokenizer_2_id
        mapped_args.tokenizer_3_id = argparse_args.tokenizer_3_id
        mapped_args.text_encoder_id = argparse_args.text_encoder_id
        mapped_args.text_encoder_2_id = argparse_args.text_encoder_2_id
        mapped_args.text_encoder_3_id = argparse_args.text_encoder_3_id
        mapped_args.transformer_id = argparse_args.transformer_id
        mapped_args.vae_id = argparse_args.vae_id
        mapped_args.text_encoder_dtype = _DTYPE_MAP[argparse_args.text_encoder_dtype]
        mapped_args.text_encoder_2_dtype = _DTYPE_MAP[argparse_args.text_encoder_2_dtype]
        mapped_args.text_encoder_3_dtype = _DTYPE_MAP[argparse_args.text_encoder_3_dtype]
        mapped_args.transformer_dtype = _DTYPE_MAP[argparse_args.transformer_dtype]
        mapped_args.vae_dtype = _DTYPE_MAP[argparse_args.vae_dtype]
        mapped_args.layerwise_upcasting_modules = argparse_args.layerwise_upcasting_modules
        mapped_args.layerwise_upcasting_storage_dtype = _DTYPE_MAP[argparse_args.layerwise_upcasting_storage_dtype]
        mapped_args.layerwise_upcasting_skip_modules_pattern = argparse_args.layerwise_upcasting_skip_modules_pattern
        mapped_args.enable_slicing = argparse_args.enable_slicing
        mapped_args.enable_tiling = argparse_args.enable_tiling

    def validate_args(self, args: "BaseArgs"):
        pass


class InferenceArgs(ArgsConfigMixin):
    """
    Args:
        inference_type (`str`):
            The type of inference to run. Choose between ['text_to_video'].
        dataset_file (`str`, defaults to `None`):
            Path to a CSV/JSON/PARQUET/ARROW file containing information for inference. The file must contain atleast the
            "caption" column. Other columns such as "image_path" and "video_path" can be provided too. If provided, "image_path"
            will be used to load a PIL.Image.Image and set the "image" key in the sample dictionary. Similarly, "video_path"
            will be used to load a List[PIL.Image.Image] and set the "video" key in the sample dictionary.
            The dataset file may contain other attributes such as:
                - "height" and "width" and "num_frames": Resolution
                - "num_inference_steps": Number of inference steps
                - "guidance_scale": Classifier-free Guidance Scale
                - ... (any number of additional attributes can be provided. The ModelSpecification::validate method will be
                invoked with the sample dictionary to validate the sample.)
    """

    inference_type: InferenceType = InferenceType.TEXT_TO_VIDEO
    dataset_file: str = None

    def add_args(self, parser: argparse.ArgumentParser) -> None:
        parser.add_argument(
            "--inference_type",
            type=str,
            default=InferenceType.TEXT_TO_VIDEO.value,
            choices=[x.value for x in InferenceType.__members__.values()],
        )
        parser.add_argument("--dataset_file", type=str, required=True)

    def map_args(self, argparse_args: argparse.Namespace, mapped_args: "BaseArgs"):
        mapped_args.inference_type = InferenceType(argparse_args.inference_type)
        mapped_args.dataset_file = argparse_args.dataset_file

    def validate_args(self, args: "BaseArgs"):
        pass


class AttentionProviderArgs(ArgsConfigMixin):
    """
    Args:
        attn_provider (`str`, defaults to "native"):
            The attention provider to use for inference. Choose between ['flash', 'flash_varlen', 'flex', 'native', '_native_cudnn', '_native_efficient', '_native_flash', '_native_math', 'sage', 'sage_varlen', '_sage_qk_int8_pv_fp8_cuda', '_sage_qk_int8_pv_fp8_cuda_sm90', '_sage_qk_int8_pv_fp16_cuda', '_sage_qk_int8_pv_fp16_triton', 'xformers'].
    """

    attn_provider: AttentionProviderInference = "native"
    # attn_provider_specialized_modules: List[str] = []

    def add_args(self, parser: argparse.ArgumentParser) -> None:
        # fmt: off
        parser.add_argument("--attn_provider", type=str, default="native", choices=["flash", "flash_varlen", "flex", "native", "_native_cudnn", "_native_efficient", "_native_flash", "_native_math", "sage", "sage_varlen", "_sage_qk_int8_pv_fp8_cuda", "_sage_qk_int8_pv_fp8_cuda_sm90", "_sage_qk_int8_pv_fp16_cuda", "_sage_qk_int8_pv_fp16_triton", "xformers"])
        # fmt: on

    def map_args(self, argparse_args: argparse.Namespace, mapped_args: "BaseArgs"):
        mapped_args.attn_provider = argparse_args.attn_provider

    def validate_args(self, args: "BaseArgs"):
        pass


class TorchConfigArgs(ArgsConfigMixin):
    """
    Args:
        compile_modules (`List[str]`, defaults to `[]`):
            Modules that should be regionally compiled with `torch.compile`.
        compile_scopes (`str`, defaults to `None`):
            The scope of compilation for each `--compile_modules`. Choose between ['regional', 'full']. Must have the same length as
            `--compile_modules`. If `None`, will default to `regional` for all modules.
        allow_tf32 (`bool`, defaults to `False`):
            Whether or not to allow the use of TF32 matmul on compatible hardware.
        float32_matmul_precision (`str`, defaults to `highest`):
            The precision to use for float32 matmul. Choose between ['highest', 'high', 'medium'].
    """

    compile_modules: List[str] = []
    compile_scopes: List[str] = None
    allow_tf32: bool = False
    float32_matmul_precision: str = "highest"

    def add_args(self, parser: argparse.ArgumentParser) -> None:
        parser.add_argument("--compile_modules", type=str, default=[], nargs="+")
        parser.add_argument("--compile_scopes", type=str, default=None, nargs="+")
        parser.add_argument("--allow_tf32", action="store_true")
        parser.add_argument(
            "--float32_matmul_precision",
            type=str,
            default="highest",
            choices=["highest", "high", "medium"],
            help="The precision to use for float32 matmul. Choose between ['highest', 'high', 'medium'].",
        )

    def map_args(self, argparse_args: argparse.Namespace, mapped_args: "BaseArgs"):
        compile_scopes = argparse_args.compile_scopes
        if len(argparse_args.compile_modules) > 0:
            if compile_scopes is None:
                compile_scopes = "regional"
            if isinstance(compile_scopes, list) and len(compile_scopes) == 1:
                compile_scopes = compile_scopes[0]
            if isinstance(compile_scopes, str):
                compile_scopes = [compile_scopes] * len(argparse_args.compile_modules)
        else:
            compile_scopes = []

        mapped_args.compile_modules = argparse_args.compile_modules
        mapped_args.compile_scopes = compile_scopes
        mapped_args.allow_tf32 = argparse_args.allow_tf32
        mapped_args.float32_matmul_precision = argparse_args.float32_matmul_precision

    def validate_args(self, args: "BaseArgs"):
        if len(args.compile_modules) > 0:
            assert len(args.compile_modules) == len(args.compile_scopes) and all(
                x in ["regional", "full"] for x in args.compile_scopes
            ), (
                "Compile modules and compile scopes must be of the same length and compile scopes must be either 'regional' or 'full'"
            )


class MiscellaneousArgs(ArgsConfigMixin):
    """
    Args:
        seed (`int`, defaults to `None`):
            Random seed for reproducibility under same initialization conditions.
        tracker_name (`str`, defaults to `finetrainers`):
            Name of the tracker/project to use for logging inference metrics.
        output_dir (`str`, defaults to `None`):
            The directory where the model checkpoints and logs will be stored.
        logging_dir (`str`, defaults to `logs`):
            The directory where the logs will be stored.
        logging_steps (`int`, defaults to `1`):
            Inference logs will be tracked every `logging_steps` steps.
        nccl_timeout (`int`, defaults to `1800`):
            Timeout for the NCCL communication.
        report_to (`str`, defaults to `wandb`):
            The name of the logger to use for logging inference metrics. Choose between ['wandb'].
        verbose (`int`, defaults to `1`):
            Whether or not to print verbose logs.
                - 0: Diffusers/Transformers warning logging on local main process only
                - 1: Diffusers/Transformers info logging on local main process only
                - 2: Diffusers/Transformers debug logging on local main process only
                - 3: Diffusers/Transformers debug logging on all processes
    """

    seed: Optional[int] = None
    tracker_name: str = "finetrainers-inference"
    output_dir: str = None
    logging_dir: Optional[str] = "logs"
    init_timeout: int = 300  # 5 minutes
    nccl_timeout: int = 600  # 10 minutes
    report_to: str = "wandb"
    verbose: int = 1

    def add_args(self, parser: argparse.ArgumentParser) -> None:
        parser.add_argument("--seed", type=int, default=None)
        parser.add_argument("--tracker_name", type=str, default="finetrainers")
        parser.add_argument("--output_dir", type=str, default="finetrainers-inference")
        parser.add_argument("--logging_dir", type=str, default="logs")
        parser.add_argument("--init_timeout", type=int, default=300)
        parser.add_argument("--nccl_timeout", type=int, default=600)
        parser.add_argument("--report_to", type=str, default="none", choices=["none", "wandb"])
        parser.add_argument("--verbose", type=int, default=0, choices=[0, 1, 2, 3])

    def map_args(self, argparse_args: argparse.Namespace, mapped_args: "BaseArgs"):
        mapped_args.seed = argparse_args.seed
        mapped_args.tracker_name = argparse_args.tracker_name
        mapped_args.output_dir = argparse_args.output_dir
        mapped_args.logging_dir = argparse_args.logging_dir
        mapped_args.init_timeout = argparse_args.init_timeout
        mapped_args.nccl_timeout = argparse_args.nccl_timeout
        mapped_args.report_to = argparse_args.report_to
        mapped_args.verbose = argparse_args.verbose

    def validate_args(self, args: "BaseArgs"):
        pass


class BaseArgs:
    """The arguments for the finetrainers inference script."""

    parallel_args = ParallelArgs()
    model_args = ModelArgs()
    inference_args = InferenceArgs()
    attention_provider_args = AttentionProviderArgs()
    torch_config_args = TorchConfigArgs()
    miscellaneous_args = MiscellaneousArgs()

    _registered_config_mixins: List[ArgsConfigMixin] = []
    _arg_group_map: Dict[str, ArgsConfigMixin] = {}

    def __init__(self):
        self._arg_group_map: Dict[str, ArgsConfigMixin] = {
            "parallel_args": self.parallel_args,
            "model_args": self.model_args,
            "inference_args": self.inference_args,
            "attention_provider_args": self.attention_provider_args,
            "torch_config_args": self.torch_config_args,
            "miscellaneous_args": self.miscellaneous_args,
        }

        for arg_config_mixin in self._arg_group_map.values():
            self.register_args(arg_config_mixin)

    def to_dict(self) -> Dict[str, Any]:
        arguments_to_dict = {}
        for config_mixin in self._registered_config_mixins:
            arguments_to_dict[config_mixin.__class__.__name__] = config_mixin.to_dict()

        return arguments_to_dict

    def register_args(self, config: ArgsConfigMixin) -> None:
        if not hasattr(self, "_extended_add_arguments"):
            self._extended_add_arguments = []
        self._extended_add_arguments.append((config.add_args, config.validate_args, config.map_args))
        self._registered_config_mixins.append(config)

    def parse_args(self):
        parser = argparse.ArgumentParser()

        for extended_add_arg_fns in getattr(self, "_extended_add_arguments", []):
            add_fn, _, _ = extended_add_arg_fns
            add_fn(parser)

        args, remaining_args = parser.parse_known_args()
        logger.debug(f"Remaining unparsed arguments: {remaining_args}")

        for extended_add_arg_fns in getattr(self, "_extended_add_arguments", []):
            _, _, map_fn = extended_add_arg_fns
            map_fn(args, self)

        for extended_add_arg_fns in getattr(self, "_extended_add_arguments", []):
            _, validate_fn, _ = extended_add_arg_fns
            validate_fn(self)

        return self

    def __getattribute__(self, name: str):
        try:
            return object.__getattribute__(self, name)
        except AttributeError:
            for arg_group in self._arg_group_map.values():
                if hasattr(arg_group, name):
                    return getattr(arg_group, name)
            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")

    def __setattr__(self, name: str, value: Any):
        if name in self.__dict__:
            object.__setattr__(self, name, value)
            return
        for arg_group in self._arg_group_map.values():
            if hasattr(arg_group, name):
                setattr(arg_group, name, value)
                return
        object.__setattr__(self, name, value)


if __name__ == "__main__":
    main()