diffusion/utils/config.py

# 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 json
from dataclasses import asdict, dataclass, field
from typing import Any, Dict, List, Optional, Tuple, Union


@dataclass
class BaseConfig:
    def get(self, attribute_name, default=None):
        return getattr(self, attribute_name, default)

    def pop(self, attribute_name, default=None):
        if hasattr(self, attribute_name):
            value = getattr(self, attribute_name)
            delattr(self, attribute_name)
            return value
        else:
            return default

    def __str__(self):
        return json.dumps(asdict(self), indent=4)


@dataclass
class DataConfig(BaseConfig):
    data_dir: List[str] = field(default_factory=list)
    caption_proportion: Dict[str, int] = field(default_factory=lambda: {"prompt": 1})
    external_caption_suffixes: List[str] = field(default_factory=list)
    external_clipscore_suffixes: List[str] = field(default_factory=list)
    caption_selection_type: str = (
        "clipscore"  # clipscore: use $external_clipscore_suffixes, proportion: use $caption_proportion
    )
    clip_thr_temperature: float = 1.0
    clip_thr: float = 0.0
    del_img_clip_thr: float = 0.0
    sort_dataset: bool = False
    load_text_feat: bool = False
    load_vae_feat: bool = False
    aspect_ratio_type: str = "ASPECT_RATIO_1024"
    transform: str = "default_train"
    type: str = "SanaWebDatasetMS"
    image_size: int = 512
    hq_only: bool = False
    valid_num: int = 0
    data: Any = None
    num_frames: int = 81
    extra: Any = None


@dataclass
class VideoDataConfig(DataConfig):
    data_dir: Dict[str, str] = field(default_factory=lambda: {"video_toy_data: data/video_toy_data"})
    aspect_ratio_type: str = "ASPECT_RATIO_VIDEO_256_MS"
    external_data_filter: Dict[str, Dict[str, Dict[str, float]]] = field(default_factory=lambda: {})
    motion_score_file_thres: Dict[str, Optional[float]] = field(default_factory=dict)
    motion_score_cal_type: str = "average"  # average, max
    target_fps: int = 16
    resample_fps: bool = True
    shuffle_dataset: bool = False
    vae_cache_dir: Optional[str] = None
    json_cache_dir: Optional[str] = None
    load_first_frame: bool = False


@dataclass
class ModelConfig(BaseConfig):
    model: str = "SanaMS_600M_P1_D28"
    teacher: Optional[str] = None
    image_size: int = 512
    mixed_precision: str = "fp16"  # ['fp16', 'fp32', 'bf16']
    fp32_attention: bool = True
    load_from: Optional[str] = None
    discriminator_model: Optional[str] = None
    teacher_model: Optional[str] = None
    teacher_model_weight_dtype: Optional[str] = None
    resume_from: Optional[Union[Dict[str, Any], str]] = field(
        default_factory=lambda: {
            "checkpoint": None,
            "load_ema": False,
            "resume_lr_scheduler": True,
            "resume_optimizer": True,
        }
    )
    aspect_ratio_type: Optional[str] = None
    multi_scale: bool = True
    pe_interpolation: float = 1.0
    micro_condition: bool = False
    attn_type: str = "linear"
    autocast_linear_attn: bool = False
    ffn_type: str = "glumbconv"
    mlp_acts: List[Optional[str]] = field(default_factory=lambda: ["silu", "silu", None])
    mlp_ratio: float = 2.5
    use_pe: bool = False
    pos_embed_type: str = "sincos"  # "sincos", "flux_rope", "wan_rope"
    qk_norm: bool = False
    class_dropout_prob: float = 0.0
    linear_head_dim: int = 32
    cross_norm: bool = False
    cross_attn_type: str = "flash"
    logvar: bool = False
    cfg_scale: int = 4
    cfg_embed: bool = False
    cfg_embed_scale: float = 1.0
    guidance_type: str = "classifier-free"
    pag_applied_layers: List[int] = field(default_factory=lambda: [8])
    # for ladd
    ladd_multi_scale: bool = True
    head_block_ids: Optional[List[int]] = None
    extra: Any = None


@dataclass
class ModelVideoConfig(ModelConfig):
    # stage1
    remove_state_dict_keys: Optional[List[str]] = None
    # stage2
    rope_fhw_dim: Optional[Tuple[int, int, int]] = None
    t_kernel_size: int = 3
    flash_attn_window_count: Optional[List[int]] = None
    pack_latents: bool = False
    additional_inchannels: int = 0
    encode_image_prompt_embeds: bool = False
    # stage3
    cross_attn_image_embeds: bool = False
    image_latent_mode: str = "video_zero"
    # chunkcasual
    chunk_index: Optional[List[int]] = None
    # Hybrid-attention video models
    softmax_ratio: Optional[float] = 0.0
    softmax_layer_indices: Optional[List[int]] = None
    softmax_attn_type: str = "V2VGatedSoftmaxAttention"
    softmax_head_dim: int = 256
    attn_res_block_size: int = 8


@dataclass
class AEConfig(BaseConfig):
    vae_type: str = "AutoencoderDC"
    vae_pretrained: str = "mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers"
    weight_dtype: str = "float32"
    scale_factor: float = 0.41407
    scaling_factor: Optional[Union[float, List[float]]] = None  # for st-dc-ae
    vae_latent_dim: int = 32
    vae_downsample_rate: int = 32
    sample_posterior: bool = True
    vae_stride: Optional[List[int]] = None
    use_causal_encode: bool = False
    if_cache: bool = False
    cache_dir: Optional[str] = None
    # Framewise / tiling fields used by LTX2VAE_diffusers for long-video decode.
    use_framewise_encoding: bool = False
    use_framewise_decoding: bool = False
    tile_sample_stride_num_frames: int = 64
    tile_sample_min_num_frames: int = 96
    extra: Any = None


@dataclass
class TextEncoderConfig(BaseConfig):
    text_encoder_name: str = "gemma-2-2b-it"
    caption_channels: int = 2304
    y_norm: bool = True
    y_norm_scale_factor: float = 1.0
    model_max_length: int = 300
    chi_prompt: List[Optional[str]] = field(default_factory=lambda: [])
    extra: Any = None


@dataclass
class ImageEncoderConfig(BaseConfig):
    image_encoder_name: Optional[str] = None
    image_encoder_path: Optional[str] = None
    weight_dtype: Optional[str] = "bf16"


@dataclass
class SchedulerConfig(BaseConfig):
    train_sampling_steps: int = 1000
    predict_flow_v: bool = True
    noise_schedule: str = "linear_flow"
    pred_sigma: bool = False
    learn_sigma: bool = True
    vis_sampler: str = "flow_dpm-solver"
    flow_shift: float = 1.0
    inference_flow_shift: Optional[float] = None
    # logit-normal timestep
    weighting_scheme: Optional[str] = "logit_normal"
    weighting_scheme_discriminator: Optional[str] = "logit_normal_trigflow"
    add_noise_timesteps: List[float] = field(default_factory=lambda: [1.57080])
    logit_mean: float = 0.0
    logit_std: float = 1.0
    logit_mean_discriminator: float = 0.0
    logit_std_discriminator: float = 1.0
    mode_scale: float = 1.29
    sigma_data: float = 1.0
    p_low: Optional[float] = None
    p_high: Optional[float] = None
    timestep_norm_scale_factor: float = 1.0
    pretrain_timestep_norm_scale_factor: float = 1.0
    discrete_norm_timestep: bool = False
    extra: Any = None


@dataclass
class TrainingConfig(BaseConfig):
    num_workers: int = 4
    seed: int = 42
    train_batch_size: int = 32
    train_batch_size_image: int = 32
    early_stop_hours: float = 100
    num_epochs: int = 100
    gradient_accumulation_steps: int = 1
    grad_checkpointing: bool = False
    gradient_clip: float = 1.0
    gc_step: int = 1
    optimizer: Dict[str, Any] = field(
        default_factory=lambda: {"eps": 1.0e-10, "lr": 0.0001, "type": "AdamW", "weight_decay": 0.03}
    )
    optimizer_D: Dict[str, Any] = field(
        default_factory=lambda: {"eps": 1.0e-10, "lr": 0.0001, "type": "AdamW", "weight_decay": 0.03}
    )
    load_from_optimizer: bool = False
    load_from_lr_scheduler: bool = False
    resume_lr_scheduler: bool = True
    lr_schedule: str = "constant"
    lr_schedule_args: Dict[str, int] = field(default_factory=lambda: {"num_warmup_steps": 500})
    auto_lr: Optional[Dict[str, str]] = field(default_factory=lambda: {"rule": "sqrt"})
    ema_rate: float = 0.9999
    eval_batch_size: int = 16
    use_fsdp: bool = False
    fsdp_version: int = 1
    cp_size: int = 0
    use_flash_attn: bool = False
    eval_sampling_steps: int = 250
    lora_rank: int = 4
    log_interval: int = 50
    mask_type: str = "null"
    mask_loss_coef: float = 0.0
    load_mask_index: bool = False
    snr_loss: bool = False
    real_prompt_ratio: float = 1.0
    save_image_epochs: int = 1
    save_model_epochs: int = 1
    save_model_steps: int = 1000000
    visualize: bool = False
    null_embed_root: str = "output/pretrained_models/"
    valid_prompt_embed_root: str = "output/tmp_embed/"
    validation_prompts: List[str] = field(
        default_factory=lambda: [
            "dog",
            "portrait photo of a girl, photograph, highly detailed face, depth of field",
            "Self-portrait oil painting, a beautiful cyborg with golden hair, 8k",
            "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
            "A photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece",
        ]
    )
    local_save_vis: bool = False
    deterministic_validation: bool = True
    online_metric: bool = False
    eval_metric_step: int = 5000
    online_metric_dir: str = "metric_helper"
    work_dir: str = "output/"
    skip_step: int = 0
    loss_type: str = "huber"
    huber_c: float = 0.001
    num_ddim_timesteps: int = 50
    w_max: float = 15.0
    w_min: float = 3.0
    ema_decay: float = 0.95
    debug_nan: bool = False
    ema_update: bool = False
    ema_rate: float = 0.9999
    weight_loss: bool = True
    tangent_warmup_steps: int = 10000
    scm_cfg_scale: Union[float, List[float]] = field(default_factory=lambda: [1.0])
    cfg_interval: Optional[List[float]] = None
    scm_logvar_loss: bool = True
    norm_invariant_to_spatial_dim: bool = True
    norm_same_as_512_scale: bool = False
    g_norm_constant: float = 0.1
    g_norm_r: float = 1.0
    show_gradient: bool = False
    lr_scale: Optional[Dict[str, List[str]]] = None
    # for ladd
    adv_lambda: float = 1.0
    scm_loss: bool = True
    scm_lambda: float = 1.0
    loss_scale: float = 1.0
    r1_penalty: bool = False
    r1_penalty_weight: float = 1.0e-5
    diff_timesteps_D: bool = True
    # for adversarial loss
    suffix_checkpoints: Optional[str] = "disc"
    misaligned_pairs_D: bool = False
    discriminator_loss: str = "cross entropy"
    largest_timestep: float = 1.57080
    train_largest_timestep: bool = False
    largest_timestep_prob: float = 0.5
    extra: Any = None


@dataclass
class TrainVideoConfig(TrainingConfig):
    validation_images: Optional[List[str]] = None
    image_prior_type: Optional[str] = None  # [flux-siglip2
    joint_training_interval: int = 50
    timestep_weight: bool = False
    noise_multiplier: Optional[float] = 0.0
    ltx_image_condition_prob: float = 0.0  # for ltx, the image condition is used for the first frame
    chunk_sampling_strategy: str = "uniform"  # uniform, incremental
    same_timestep_prob: float = (
        0.0  # for incremental sampling, the probability of using the same timestep for all chunks
    )
    # temporal coherence loss for video training
    temporal_coherence_loss: bool = False
    temporal_coherence_weight: float = 0.0


@dataclass
class ControlNetConfig(BaseConfig):
    control_signal_type: str = "scribble"
    validation_scribble_maps: List[str] = field(
        default_factory=lambda: [
            "output/tmp_embed/controlnet/dog_scribble_thickness_3.jpg",
            "output/tmp_embed/controlnet/girl_scribble_thickness_3.jpg",
            "output/tmp_embed/controlnet/cyborg_scribble_thickness_3.jpg",
            "output/tmp_embed/controlnet/Astronaut_scribble_thickness_3.jpg",
            "output/tmp_embed/controlnet/mountain_scribble_thickness_3.jpg",
        ]
    )


@dataclass
class ModelGrowthConfig(BaseConfig):
    """Model growth configuration for initializing larger models from smaller ones"""

    pretrained_ckpt_path: str = ""
    init_strategy: str = "constant"  # ['cyclic', 'block_expand', 'progressive', 'interpolation', 'random', 'constant']
    init_params: Dict[str, Any] = field(
        default_factory=lambda: {
            "expand_ratio": 3,
            "noise_scale": 0.01,
        }
    )
    source_num_layers: int = 20
    target_num_layers: int = 30
    extra: Any = None


@dataclass
class SanaConfig(BaseConfig):
    data: DataConfig
    model: ModelConfig
    vae: AEConfig
    text_encoder: TextEncoderConfig
    scheduler: SchedulerConfig
    train: TrainingConfig
    controlnet: Optional[ControlNetConfig] = None
    model_growth: Optional[ModelGrowthConfig] = None
    work_dir: str = "output/"
    resume_from: Optional[str] = None
    load_from: Optional[str] = None
    debug: bool = False
    caching: bool = False
    report_to: str = "wandb"
    tracker_project_name: str = "sana-video-baseline"
    name: str = "baseline"
    loss_report_name: str = "loss"


@dataclass
class WanTextEncoderConfig(BaseConfig):
    t5_model: str = "umt5_xxl"
    t5_dtype: str = "bfloat16"
    text_len: int = 512
    t5_checkpoint: str = "checkpoints/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth"
    t5_tokenizer: str = "google/umt5-xxl"
    extra: Any = None
    caption_channels: int = 4096


@dataclass
class DistillConfig(BaseConfig):
    pass


@dataclass
class SanaVideoConfig(BaseConfig):
    data: VideoDataConfig
    model: ModelVideoConfig
    vae: AEConfig
    text_encoder: TextEncoderConfig
    scheduler: SchedulerConfig
    train: TrainVideoConfig
    image_data: Optional[DataConfig] = None
    image_encoder: Optional[ImageEncoderConfig] = field(default_factory=lambda: {})
    model_growth: Optional[ModelGrowthConfig] = None
    text_encoder_wan: Optional[WanTextEncoderConfig] = None
    work_dir: str = "output/"
    resume_from: Optional[str] = None
    load_from: Optional[str] = None
    debug: bool = False
    caching: bool = False
    report_to: str = "wandb"
    tracker_project_name: str = "sana-video"
    name: str = "baseline"
    loss_report_name: str = "loss"
    task: str = "t2v"  # t2v or ti2v
    distill: Optional[DistillConfig] = None


@dataclass
class SanaVideoStage1Config(BaseConfig):
    data: DataConfig
    model: ModelVideoConfig
    vae: AEConfig
    text_encoder: TextEncoderConfig
    scheduler: SchedulerConfig
    train: TrainVideoConfig
    model_growth: Optional[ModelGrowthConfig] = None
    work_dir: str = "output/"
    resume_from: Optional[str] = None
    load_from: Optional[str] = None
    debug: bool = False
    caching: bool = False
    report_to: str = "wandb"
    tracker_project_name: str = "sana-video"
    name: str = "baseline"
    loss_report_name: str = "loss"
    task: str = "t2v"  # t2v or ti2v or df


def model_init_config(config: SanaConfig, latent_size: int = 32):

    pred_sigma = getattr(config.scheduler, "pred_sigma", True)
    learn_sigma = getattr(config.scheduler, "learn_sigma", True) and pred_sigma
    return {
        "input_size": latent_size,
        "pe_interpolation": config.model.pe_interpolation,
        "config": config,
        "model_max_length": config.text_encoder.model_max_length,
        "qk_norm": config.model.qk_norm,
        "micro_condition": config.model.micro_condition,
        "caption_channels": config.text_encoder.caption_channels,
        "class_dropout_prob": config.model.class_dropout_prob,
        "y_norm": config.text_encoder.y_norm,
        "attn_type": config.model.attn_type,
        "ffn_type": config.model.ffn_type,
        "mlp_ratio": config.model.mlp_ratio,
        "mlp_acts": list(config.model.mlp_acts),
        "in_channels": config.vae.vae_latent_dim,
        "y_norm_scale_factor": config.text_encoder.y_norm_scale_factor,
        "use_pe": config.model.use_pe,
        "pos_embed_type": config.model.pos_embed_type,
        "linear_head_dim": config.model.linear_head_dim,
        "pred_sigma": pred_sigma,
        "learn_sigma": learn_sigma,
        "cross_norm": config.model.cross_norm,
        "cross_attn_type": config.model.cross_attn_type,
        "timestep_norm_scale_factor": config.scheduler.timestep_norm_scale_factor,
        "discrete_norm_timestep": config.scheduler.discrete_norm_timestep,
    }


def model_video_init_config(config: SanaVideoConfig, latent_size: int = 32):

    pred_sigma = getattr(config.scheduler, "pred_sigma", True)
    learn_sigma = getattr(config.scheduler, "learn_sigma", True) and pred_sigma
    if config.model.model.startswith("SanaVideo2_"):
        return {
            "input_size": latent_size,
            "config": config,
            "model_max_length": config.text_encoder.model_max_length,
            "qk_norm": config.model.qk_norm,
            "caption_channels": config.text_encoder.caption_channels,
            "class_dropout_prob": config.model.class_dropout_prob,
            "y_norm": config.text_encoder.y_norm,
            "ffn_type": config.model.ffn_type,
            "mlp_ratio": config.model.mlp_ratio,
            "in_channels": config.vae.vae_latent_dim,
            "y_norm_scale_factor": config.text_encoder.y_norm_scale_factor,
            "use_pe": config.model.use_pe,
            "pos_embed_type": config.model.pos_embed_type,
            "rope_fhw_dim": config.model.rope_fhw_dim,
            "linear_head_dim": config.model.linear_head_dim,
            "softmax_head_dim": config.model.softmax_head_dim,
            "softmax_ratio": config.model.softmax_ratio,
            "pred_sigma": pred_sigma,
            "cross_norm": config.model.cross_norm,
            "attn_res_block_size": config.model.attn_res_block_size,
            "timestep_norm_scale_factor": config.scheduler.timestep_norm_scale_factor,
        }
    return {
        "input_size": latent_size,
        "pe_interpolation": config.model.pe_interpolation,
        "config": config,
        "model_max_length": config.text_encoder.model_max_length,
        "qk_norm": config.model.qk_norm,
        "micro_condition": config.model.micro_condition,
        "caption_channels": config.text_encoder.caption_channels,
        "class_dropout_prob": config.model.class_dropout_prob,
        "y_norm": config.text_encoder.y_norm,
        "attn_type": config.model.attn_type,
        "ffn_type": config.model.ffn_type,
        "mlp_ratio": config.model.mlp_ratio,
        "mlp_acts": list(config.model.mlp_acts),
        "in_channels": config.vae.vae_latent_dim,
        "use_pe": config.model.use_pe,
        "pos_embed_type": config.model.pos_embed_type,
        "rope_fhw_dim": config.model.rope_fhw_dim,
        "linear_head_dim": config.model.linear_head_dim,
        "pred_sigma": pred_sigma,
        "learn_sigma": learn_sigma,
        "cross_norm": config.model.cross_norm,
        "cross_attn_type": config.model.cross_attn_type,
        "cross_attn_image_embeds": config.model.cross_attn_image_embeds,
        "t_kernel_size": config.model.t_kernel_size,
        "flash_attn_window_count": config.model.flash_attn_window_count,
        "pack_latents": config.model.pack_latents,
    }


def model_v2v_init_config(config: SanaVideoConfig, latent_size: int = 32):
    """Build the shared SANA-Streaming V2V model arguments.

    V2V concatenates the noisy target and clean source latents along the
    channel dimension. Keeping this helper in the public config module makes
    training and inference use the same doubled-input model definition.
    """

    model_kwargs = model_video_init_config(config, latent_size=latent_size)
    model_kwargs.update(
        additional_inchannels=config.model.additional_inchannels,
        timestep_norm_scale_factor=config.scheduler.timestep_norm_scale_factor,
        softmax_ratio=config.model.softmax_ratio,
        softmax_layer_indices=config.model.softmax_layer_indices,
        softmax_attn_type=config.model.softmax_attn_type,
    )
    return model_kwargs