import json
import os
import uuid
from copy import deepcopy
from typing import Any

from modules.util.config.BaseConfig import BaseConfig
from modules.util.config.ConceptConfig import ConceptConfig
from modules.util.config.SampleConfig import SampleConfig
from modules.util.config.SecretsConfig import SecretsConfig
from modules.util.enum.AlignPropLoss import AlignPropLoss
from modules.util.enum.AttentionMechanism import AttentionMechanism
from modules.util.enum.ConfigPart import ConfigPart
from modules.util.enum.DataType import DataType
from modules.util.enum.EMAMode import EMAMode
from modules.util.enum.GradientCheckpointingMethod import GradientCheckpointingMethod
from modules.util.enum.ImageFormat import ImageFormat
from modules.util.enum.LearningRateScaler import LearningRateScaler
from modules.util.enum.LearningRateScheduler import LearningRateScheduler
from modules.util.enum.LossScaler import LossScaler
from modules.util.enum.LossWeight import LossWeight
from modules.util.enum.ModelFormat import ModelFormat
from modules.util.enum.ModelType import ModelType, PeftType
from modules.util.enum.Optimizer import Optimizer
from modules.util.enum.TimestepDistribution import TimestepDistribution
from modules.util.enum.TimeUnit import TimeUnit
from modules.util.enum.TrainingMethod import TrainingMethod
from modules.util.ModelNames import EmbeddingName, ModelNames
from modules.util.ModelWeightDtypes import ModelWeightDtypes
from modules.util.torch_util import default_device


class TrainOptimizerConfig(BaseConfig):
    optimizer: Optimizer
    adam_w_mode: bool
    alpha: float
    amsgrad: bool
    beta1: float
    beta2: float
    beta3: float
    bias_correction: bool
    block_wise: bool
    capturable: bool
    centered: bool
    clip_threshold: float
    d0: float
    d_coef: float
    dampening: float
    decay_rate: float
    decouple: bool
    differentiable: bool
    eps: float
    eps2: float
    foreach: bool
    fsdp_in_use: bool
    fused: bool
    fused_back_pass: bool
    growth_rate: float
    initial_accumulator_value: int
    is_paged: bool
    log_every: int
    lr_decay: float
    max_unorm: float
    maximize: bool
    min_8bit_size: int
    momentum: float
    nesterov: bool
    no_prox: bool
    optim_bits: int
    percentile_clipping: float
    r: float
    relative_step: bool
    safeguard_warmup: bool
    scale_parameter: bool
    stochastic_rounding: bool
    use_bias_correction: bool
    use_triton: bool
    warmup_init: bool
    weight_decay: float
    weight_lr_power: float
    decoupled_decay: bool
    fixed_decay: bool
    weight_decouple: bool
    rectify: bool
    degenerated_to_sgd: bool
    k: int
    xi: float
    n_sma_threshold: int
    ams_bound: bool
    r: float
    adanorm: bool
    adam_debias: bool
    slice_p: int
    cautious: bool

    def __init__(self, data: list[(str, Any, type, bool)]):
        super().__init__(data)

    @staticmethod
    def default_values():
        data = []

        # name, default value, data type, nullable
        data.append(("optimizer", Optimizer.ADAMW, Optimizer, False))
        data.append(("adam_w_mode", False, bool, False))
        data.append(("alpha", None, float, True))
        data.append(("amsgrad", False, bool, False))
        data.append(("beta1", None, float, True))
        data.append(("beta2", None, float, True))
        data.append(("beta3", None, float, True))
        data.append(("bias_correction", False, bool, False))
        data.append(("block_wise", False, bool, False))
        data.append(("capturable", False, bool, False))
        data.append(("centered", False, bool, False))
        data.append(("clip_threshold", None, float, True))
        data.append(("d0", None, float, True))
        data.append(("d_coef", None, float, True))
        data.append(("dampening", None, float, True))
        data.append(("decay_rate", None, float, True))
        data.append(("decouple", False, bool, False))
        data.append(("differentiable", False, bool, False))
        data.append(("eps", None, float, True))
        data.append(("eps2", None, float, True))
        data.append(("foreach", False, bool, True))  # Disabled, because it uses too much VRAM
        data.append(("fsdp_in_use", False, bool, False))
        data.append(("fused", False, bool, False))
        data.append(("fused_back_pass", False, bool, False))
        data.append(("growth_rate", None, float, True))
        data.append(("initial_accumulator_value", None, int, True))
        data.append(("is_paged", False, bool, False))
        data.append(("log_every", None, int, True))
        data.append(("lr_decay", None, float, True))
        data.append(("max_unorm", None, float, True))
        data.append(("maximize", False, bool, False))
        data.append(("min_8bit_size", None, int, True))
        data.append(("momentum", None, float, True))
        data.append(("nesterov", False, bool, False))
        data.append(("no_prox", False, bool, False))
        data.append(("optim_bits", None, int, True))
        data.append(("percentile_clipping", None, float, True))
        data.append(("r", None, float, True))
        data.append(("relative_step", False, bool, False))
        data.append(("safeguard_warmup", False, bool, False))
        data.append(("scale_parameter", False, bool, False))
        data.append(("stochastic_rounding", True, bool, False))
        data.append(("use_bias_correction", False, bool, False))
        data.append(("use_triton", False, bool, False))
        data.append(("warmup_init", False, bool, False))
        data.append(("weight_decay", None, float, True))
        data.append(("weight_lr_power", None, float, True))
        data.append(("decoupled_decay", False, bool, False))
        data.append(("fixed_decay", False, bool, False))
        data.append(("rectify", False, bool, False))
        data.append(("degenerated_to_sgd", False, bool, False))
        data.append(("k", None, int, True))
        data.append(("xi", None, float, True))
        data.append(("n_sma_threshold", None, int, True))
        data.append(("ams_bound", False, bool, False))
        data.append(("r", None, float, True))
        data.append(("adanorm", False, bool, False))
        data.append(("adam_debias", False, bool, False))
        data.append(("slice_p", None, int, True))
        data.append(("cautious", False, bool, False))

        return TrainOptimizerConfig(data)


class TrainModelPartConfig(BaseConfig):
    model_name: str
    include: bool
    train: bool
    stop_training_after: int
    stop_training_after_unit: TimeUnit
    learning_rate: float
    weight_dtype: DataType
    dropout_probability: float
    train_embedding: bool
    attention_mask: bool
    guidance_scale: float

    def __init__(self, data: list[(str, Any, type, bool)]):
        super().__init__(data)

    @staticmethod
    def default_values():
        data = []

        # name, default value, data type, nullable
        data.append(("model_name", "", str, False))
        data.append(("include", True, bool, False))
        data.append(("train", True, bool, False))
        data.append(("stop_training_after", None, int, True))
        data.append(("stop_training_after_unit", TimeUnit.NEVER, TimeUnit, False))
        data.append(("learning_rate", None, float, True))
        data.append(("weight_dtype", DataType.NONE, DataType, False))
        data.append(("dropout_probability", 0.0, float, False))
        data.append(("train_embedding", True, bool, False))
        data.append(("attention_mask", False, bool, False))
        data.append(("guidance_scale", 1.0, float, False))

        return TrainModelPartConfig(data)


class TrainEmbeddingConfig(BaseConfig):
    uuid: str
    model_name: str
    placeholder: str
    train: bool
    stop_training_after: int
    stop_training_after_unit: TimeUnit
    token_count: int
    initial_embedding_text: str

    def __init__(self, data: list[(str, Any, type, bool)]):
        super().__init__(data)

    @staticmethod
    def default_values():
        data = []

        # name, default value, data type, nullable
        data.append(("uuid", str(uuid.uuid4()), str, False))
        data.append(("model_name", "", str, False))
        data.append(("placeholder", "<embedding>", str, False))
        data.append(("train", True, bool, False))
        data.append(("stop_training_after", None, int, True))
        data.append(("stop_training_after_unit", TimeUnit.NEVER, TimeUnit, False))
        data.append(("token_count", 1, int, False))
        data.append(("initial_embedding_text", "*", str, False))

        return TrainEmbeddingConfig(data)


class TrainConfig(BaseConfig):
    training_method: TrainingMethod
    model_type: ModelType
    debug_mode: bool
    debug_dir: str
    workspace_dir: str
    cache_dir: str
    tensorboard: bool
    tensorboard_expose: bool
    tensorboard_port: str
    validation: bool
    validate_after: float
    validate_after_unit: TimeUnit
    continue_last_backup: bool
    include_train_config: ConfigPart

    # model settings
    base_model_name: str
    weight_dtype: DataType
    output_dtype: DataType
    output_model_format: ModelFormat
    output_model_destination: str
    gradient_checkpointing: GradientCheckpointingMethod
    enable_async_offloading: bool
    enable_activation_offloading: bool
    layer_offload_fraction: float
    force_circular_padding: bool

    # data settings
    concept_file_name: str
    concepts: list[ConceptConfig]
    aspect_ratio_bucketing: bool
    latent_caching: bool
    clear_cache_before_training: bool

    # training settings
    learning_rate_scheduler: LearningRateScheduler
    custom_learning_rate_scheduler: str | None
    # Dict keys are literally called "key" and "value"; not a tuple because
    # of restrictions with ConfigList.
    scheduler_params: list[dict[str, str]]
    learning_rate: float
    learning_rate_warmup_steps: float
    learning_rate_cycles: float
    epochs: int
    batch_size: int
    gradient_accumulation_steps: int
    ema: EMAMode
    ema_decay: float
    ema_update_step_interval: int
    dataloader_threads: int
    train_device: str
    temp_device: str
    train_dtype: DataType
    fallback_train_dtype: DataType
    enable_autocast_cache: bool
    only_cache: bool
    resolution: str
    attention_mechanism: AttentionMechanism
    align_prop: bool
    align_prop_probability: float
    align_prop_loss: AlignPropLoss
    align_prop_weight: float
    align_prop_steps: int
    align_prop_truncate_steps: float
    align_prop_cfg_scale: float
    mse_strength: float
    mae_strength: float
    log_cosh_strength: float
    vb_loss_strength: float
    loss_weight_fn: LossWeight
    loss_weight_strength: float
    dropout_probability: float
    loss_scaler: LossScaler
    learning_rate_scaler: LearningRateScaler
    clip_grad_norm: float

    # noise
    offset_noise_weight: float
    perturbation_noise_weight: float
    rescale_noise_scheduler_to_zero_terminal_snr: bool
    force_v_prediction: bool
    force_epsilon_prediction: bool
    timestep_distribution: TimestepDistribution
    min_noising_strength: float
    max_noising_strength: float

    noising_weight: float
    noising_bias: float

    # unet
    unet: TrainModelPartConfig

    # prior
    prior: TrainModelPartConfig

    # text encoder
    text_encoder: TrainModelPartConfig
    text_encoder_layer_skip: int

    # text encoder 2
    text_encoder_2: TrainModelPartConfig
    text_encoder_2_layer_skip: int

    # text encoder 3
    text_encoder_3: TrainModelPartConfig
    text_encoder_3_layer_skip: int

    # vae
    vae: TrainModelPartConfig

    # effnet encoder
    effnet_encoder: TrainModelPartConfig

    # decoder
    decoder: TrainModelPartConfig

    # decoder text encoder
    decoder_text_encoder: TrainModelPartConfig

    # decoder vqgan
    decoder_vqgan: TrainModelPartConfig

    # masked training
    masked_training: bool
    unmasked_probability: float
    unmasked_weight: float
    normalize_masked_area_loss: bool

    # embedding
    embedding_learning_rate: float
    preserve_embedding_norm: bool
    embedding: TrainEmbeddingConfig
    additional_embeddings: list[TrainEmbeddingConfig]
    embedding_weight_dtype: DataType

    # lora
    peft_type: PeftType
    lora_model_name: str
    lora_rank: int
    lora_alpha: float
    lora_decompose: bool
    lora_decompose_norm_epsilon: bool
    lora_weight_dtype: DataType
    lora_layers: str  # comma-separated
    lora_layer_preset: str
    bundle_additional_embeddings: bool

    # optimizer
    optimizer: TrainOptimizerConfig
    optimizer_defaults: dict[str, TrainOptimizerConfig]

    # sample settings
    sample_definition_file_name: str
    samples: list[SampleConfig]
    sample_after: float
    sample_after_unit: TimeUnit
    sample_image_format: ImageFormat
    samples_to_tensorboard: bool
    non_ema_sampling: bool

    # backup settings
    backup_after: float
    backup_after_unit: TimeUnit
    rolling_backup: bool
    rolling_backup_count: int
    backup_before_save: bool
    save_every: int
    save_every_unit: TimeUnit
    save_skip_first: int
    save_filename_prefix: str

    # secrets - not saved into config file
    secrets: SecretsConfig

    def __init__(self, data: list[(str, Any, type, bool)]):
        super().__init__(
            data,
            config_version=6,
            config_migrations={
                0: self.__migration_0,
                1: self.__migration_1,
                2: self.__migration_2,
                3: self.__migration_3,
                4: self.__migration_4,
                5: self.__migration_5,
            }
        )

    def __migration_0(self, data: dict) -> dict:
        optimizer_settings = {}
        migrated_data = {}
        for key, value in data.items():
            # move optimizer settings to sub object
            if key == 'optimizer':
                optimizer_settings['optimizer'] = value
            elif key.startswith('optimizer'):
                optimizer_settings[key.removeprefix('optimizer_')] = value
            else:
                migrated_data[key] = value

        if 'optimizer' in optimizer_settings:
            migrated_data['optimizer'] = optimizer_settings
            migrated_data['optimizer_defaults'] = {
                optimizer_settings['optimizer']: deepcopy(optimizer_settings)
            }

        return migrated_data

    def __migration_1(self, data: dict) -> dict:
        migrated_data = {
            "unet": {},
            "prior": {},
            "text_encoder": {},
            "text_encoder_2": {},
            "vae": {},
            "effnet_encoder": {},
            "decoder": {},
            "decoder_text_encoder": {},
            "decoder_vqgan": {},
            "embeddings": [{}],
        }

        for key, value in data.items():
            if key == "train_unet":
                migrated_data["unet"]["train"] = value
            elif key == "train_unet_epochs":
                migrated_data["unet"]["stop_training_after"] = value
                migrated_data["unet"]["stop_training_after_unit"] = TimeUnit.EPOCH
            elif key == "unet_learning_rate":
                migrated_data["unet"]["learning_rate"] = value
            elif key == "unet_weight_dtype":
                migrated_data["unet"]["weight_dtype"] = value

            elif key == "train_prior":
                migrated_data["prior"]["train"] = value
            elif key == "prior_model_name":
                migrated_data["prior"]["model_name"] = value
            elif key == "train_prior_epochs":
                migrated_data["prior"]["stop_training_after"] = value
                migrated_data["prior"]["stop_training_after_unit"] = TimeUnit.EPOCH
            elif key == "prior_learning_rate":
                migrated_data["prior"]["learning_rate"] = value
            elif key == "prior_weight_dtype":
                migrated_data["prior"]["weight_dtype"] = value

            elif key == "train_text_encoder":
                migrated_data["text_encoder"]["train"] = value
            elif key == "train_text_encoder_epochs":
                migrated_data["text_encoder"]["stop_training_after"] = value
                migrated_data["text_encoder"]["stop_training_after_unit"] = TimeUnit.EPOCH
            elif key == "text_encoder_learning_rate":
                migrated_data["text_encoder"]["learning_rate"] = value
            elif key == "text_encoder_weight_dtype":
                migrated_data["text_encoder"]["weight_dtype"] = value

            elif key == "train_text_encoder_2":
                migrated_data["text_encoder_2"]["train"] = value
            elif key == "train_text_encoder_2_epochs":
                migrated_data["text_encoder_2"]["stop_training_after"] = value
                migrated_data["text_encoder_2"]["stop_training_after_unit"] = TimeUnit.EPOCH
            elif key == "text_encoder_2_learning_rate":
                migrated_data["text_encoder_2"]["learning_rate"] = value
            elif key == "text_encoder_2_weight_dtype":
                migrated_data["text_encoder_2"]["weight_dtype"] = value

            elif key == "vae_model_name":
                migrated_data["vae"]["model_name"] = value
            elif key == "vae_weight_dtype":
                migrated_data["vae"]["weight_dtype"] = value

            elif key == "effnet_encoder_model_name":
                migrated_data["effnet_encoder"]["model_name"] = value
            elif key == "effnet_encoder_weight_dtype":
                migrated_data["effnet_encoder"]["weight_dtype"] = value

            elif key == "decoder_model_name":
                migrated_data["decoder"]["model_name"] = value
            elif key == "decoder_weight_dtype":
                migrated_data["decoder"]["weight_dtype"] = value

            elif key == "decoder_text_encoder_weight_dtype":
                migrated_data["decoder_text_encoder"]["weight_dtype"] = value

            elif key == "decoder_vqgan_weight_dtype":
                migrated_data["decoder_vqgan"]["weight_dtype"] = value

            elif key == "embedding_model_names" and len(value) > 0:
                migrated_data["embeddings"][0]["model_name"] = value[0]
            elif key == "token_count":
                migrated_data["embeddings"][0]["token_count"] = value
            elif key == "initial_embedding_text":
                migrated_data["embeddings"][0]["initial_embedding_text"] = value

            else:
                migrated_data[key] = value

        return migrated_data

    def __migration_2(self, data: dict) -> dict:
        migrated_data = data.copy()
        min_snr_gamma = migrated_data.pop("min_snr_gamma", 0.0)
        model_type = ModelType(migrated_data.get("model_type", ModelType.STABLE_DIFFUSION_15))
        if min_snr_gamma:
            migrated_data["loss_weight_fn"] = LossWeight.MIN_SNR_GAMMA
            migrated_data["loss_weight_strength"] = min_snr_gamma
        elif model_type.is_wuerstchen():
            migrated_data["loss_weight_fn"] = LossWeight.P2
            migrated_data["loss_weight_strength"] = 1.0

        return migrated_data

    def __migration_3(self, data: dict) -> dict:
        migrated_data = data.copy()

        noising_weight = migrated_data.pop("noising_weight", 0.0)
        noising_bias = migrated_data.pop("noising_bias", 0.5)

        if noising_weight != 0:
            migrated_data["timestep_distribution"] = TimestepDistribution.SIGMOID
            migrated_data["noising_weight"] = noising_weight
            migrated_data["noising_bias"] = noising_bias - 0.5
        else:
            migrated_data["timestep_distribution"] = TimestepDistribution.UNIFORM
            migrated_data["noising_weight"] = 0.0
            migrated_data["noising_bias"] = 0.0

        return migrated_data

    def __migration_4(self, data: dict) -> dict:
        migrated_data = data.copy()

        gradient_checkpointing = migrated_data.pop("gradient_checkpointing", True)

        if gradient_checkpointing:
            migrated_data["gradient_checkpointing"] = GradientCheckpointingMethod.ON
        else:
            migrated_data["gradient_checkpointing"] = GradientCheckpointingMethod.OFF

        return migrated_data

    def __migration_5(self, data: dict) -> dict:
        migrated_data = data.copy()

        if "save_after" in migrated_data:
            migrated_data["save_every"] = migrated_data.pop("save_after")
        if "save_after_unit" in migrated_data:
            migrated_data["save_every_unit"] = migrated_data.pop("save_after_unit")

        return migrated_data

    def weight_dtypes(self) -> ModelWeightDtypes:
        return ModelWeightDtypes(
            self.train_dtype,
            self.fallback_train_dtype,
            self.weight_dtype if self.unet.weight_dtype == DataType.NONE else self.unet.weight_dtype,
            self.weight_dtype if self.prior.weight_dtype == DataType.NONE else self.prior.weight_dtype,
            self.weight_dtype if self.text_encoder.weight_dtype == DataType.NONE else self.text_encoder.weight_dtype,
            self.weight_dtype if self.text_encoder_2.weight_dtype == DataType.NONE else self.text_encoder_2.weight_dtype,
            self.weight_dtype if self.text_encoder_3.weight_dtype == DataType.NONE else self.text_encoder_3.weight_dtype,
            self.weight_dtype if self.vae.weight_dtype == DataType.NONE else self.vae.weight_dtype,
            self.weight_dtype if self.effnet_encoder.weight_dtype == DataType.NONE else self.effnet_encoder.weight_dtype,
            self.weight_dtype if self.decoder.weight_dtype == DataType.NONE else self.decoder.weight_dtype,
            self.weight_dtype if self.decoder_text_encoder.weight_dtype == DataType.NONE else self.decoder_text_encoder.weight_dtype,
            self.weight_dtype if self.decoder_vqgan.weight_dtype == DataType.NONE else self.decoder_vqgan.weight_dtype,
            self.weight_dtype if self.lora_weight_dtype == DataType.NONE else self.lora_weight_dtype,
            self.weight_dtype if self.embedding_weight_dtype == DataType.NONE else self.embedding_weight_dtype,
        )

    def model_names(self) -> ModelNames:
        return ModelNames(
            base_model=self.base_model_name,
            prior_model=self.prior.model_name,
            effnet_encoder_model=self.effnet_encoder.model_name,
            decoder_model=self.decoder.model_name,
            vae_model=self.vae.model_name,
            lora=self.lora_model_name,
            embedding=EmbeddingName(self.embedding.uuid, self.embedding.model_name),
            additional_embeddings=[EmbeddingName(embedding.uuid, embedding.model_name) for embedding in
                                   self.additional_embeddings],
            include_text_encoder=self.text_encoder.include,
            include_text_encoder_2=self.text_encoder_2.include,
            include_text_encoder_3=self.text_encoder_3.include,
        )

    def train_any_embedding(self) -> bool:
        return self.training_method == TrainingMethod.EMBEDDING \
            or any(embedding.train for embedding in self.additional_embeddings)

    def train_text_encoder_or_embedding(self) -> bool:
        return (self.text_encoder.train and self.training_method != TrainingMethod.EMBEDDING) \
            or ((self.text_encoder.train_embedding or not self.model_type.has_multiple_text_encoders())
                and self.train_any_embedding())

    def train_text_encoder_2_or_embedding(self) -> bool:
        return (self.text_encoder_2.train and self.training_method != TrainingMethod.EMBEDDING) \
            or ((self.text_encoder_2.train_embedding or not self.model_type.has_multiple_text_encoders())
                and self.train_any_embedding())

    def train_text_encoder_3_or_embedding(self) -> bool:
        return (self.text_encoder_3.train and self.training_method != TrainingMethod.EMBEDDING) \
            or ((self.text_encoder_3.train_embedding or not self.model_type.has_multiple_text_encoders())
                and self.train_any_embedding())

    def get_last_backup_path(self) -> str | None:
        backups_path = os.path.join(self.workspace_dir, "backup")
        if os.path.exists(backups_path):
            backup_paths = sorted(
                [path for path in os.listdir(backups_path) if
                 os.path.isdir(os.path.join(backups_path, path))],
                reverse=True,
            )

            if backup_paths:
                last_backup_path = backup_paths[0]
                return os.path.join(backups_path, last_backup_path)

        return None

    def to_settings_dict(self) -> dict:
        config = TrainConfig.default_values().from_dict(self.to_dict())

        config.concepts = None
        config.samples = None

        return config.to_dict()

    def to_pack_dict(self) -> dict:
        config = TrainConfig.default_values().from_dict(self.to_dict())

        if config.concepts is None:
            with open(config.concept_file_name, 'r') as f:
                concepts = json.load(f)
                for i in range(len(concepts)):
                    concepts[i] = ConceptConfig.default_values().from_dict(concepts[i])
                config.concepts = concepts

        if config.samples is None:
            with open(config.sample_definition_file_name, 'r') as f:
                samples = json.load(f)
                for i in range(len(samples)):
                    samples[i] = SampleConfig.default_values().from_dict(samples[i])
                config.samples = samples

        return config.to_dict()

    def to_unpacked_config(self) -> 'TrainConfig':
        config = TrainConfig.default_values().from_dict(self.to_dict())
        config.concepts = None
        config.samples = None
        return config

    @staticmethod
    def default_values() -> 'TrainConfig':
        data = []

        # name, default value, data type, nullable

        # general settings
        data.append(("training_method", TrainingMethod.FINE_TUNE, TrainingMethod, False))
        data.append(("model_type", ModelType.STABLE_DIFFUSION_15, ModelType, False))
        data.append(("debug_mode", False, bool, False))
        data.append(("debug_dir", "debug", str, False))
        data.append(("workspace_dir", "workspace/run", str, False))
        data.append(("cache_dir", "workspace-cache/run", str, False))
        data.append(("tensorboard", True, bool, False))
        data.append(("tensorboard_expose", False, bool, False))
        data.append(("tensorboard_port", 6006, int, False))
        data.append(("validation", False, bool, False))
        data.append(("validate_after", 1, int, False))
        data.append(("validate_after_unit", TimeUnit.EPOCH, TimeUnit, False))
        data.append(("continue_last_backup", False, bool, False))
        data.append(("include_train_config", ConfigPart.NONE, ConfigPart, False))

        # model settings
        data.append(("base_model_name", "stable-diffusion-v1-5/stable-diffusion-v1-5", str, False))
        data.append(("weight_dtype", DataType.FLOAT_32, DataType, False))
        data.append(("output_dtype", DataType.FLOAT_32, DataType, False))
        data.append(("output_model_format", ModelFormat.SAFETENSORS, ModelFormat, False))
        data.append(("output_model_destination", "models/model.safetensors", str, False))
        data.append(("gradient_checkpointing", GradientCheckpointingMethod.ON, GradientCheckpointingMethod, False))
        data.append(("enable_async_offloading", True, bool, False))
        data.append(("enable_activation_offloading", True, bool, False))
        data.append(("layer_offload_fraction", 0.0, float, False))
        data.append(("force_circular_padding", False, bool, False))

        # data settings
        data.append(("concept_file_name", "training_concepts/concepts.json", str, False))
        data.append(("concepts", None, list[ConceptConfig], True))
        data.append(("aspect_ratio_bucketing", True, bool, False))
        data.append(("latent_caching", True, bool, False))
        data.append(("clear_cache_before_training", True, bool, False))

        # training settings
        data.append(("learning_rate_scheduler", LearningRateScheduler.CONSTANT, LearningRateScheduler, False))
        data.append(("custom_learning_rate_scheduler", None, str, True))
        data.append(("scheduler_params", [], list[dict[str, str]], True))
        data.append(("learning_rate", 3e-6, float, False))
        data.append(("learning_rate_warmup_steps", 200.0, float, False))
        data.append(("learning_rate_cycles", 1.0, float, False))
        data.append(("epochs", 100, int, False))
        data.append(("batch_size", 1, int, False))
        data.append(("gradient_accumulation_steps", 1, int, False))
        data.append(("ema", EMAMode.OFF, EMAMode, False))
        data.append(("ema_decay", 0.999, float, False))
        data.append(("ema_update_step_interval", 5, int, False))
        data.append(("dataloader_threads", 2, int, False))
        data.append(("train_device", default_device.type, str, False))
        data.append(("temp_device", "cpu", str, False))
        data.append(("train_dtype", DataType.FLOAT_16, DataType, False))
        data.append(("fallback_train_dtype", DataType.BFLOAT_16, DataType, False))
        data.append(("enable_autocast_cache", True, bool, False))
        data.append(("only_cache", False, bool, False))
        data.append(("resolution", "512", str, False))
        data.append(("attention_mechanism", AttentionMechanism.XFORMERS, AttentionMechanism, False))
        data.append(("align_prop", False, bool, False))
        data.append(("align_prop_probability", 0.1, float, False))
        data.append(("align_prop_loss", AlignPropLoss.AESTHETIC, AlignPropLoss, False))
        data.append(("align_prop_weight", 0.01, float, False))
        data.append(("align_prop_steps", 20, int, False))
        data.append(("align_prop_truncate_steps", 0.5, float, False))
        data.append(("align_prop_cfg_scale", 7.0, float, False))
        data.append(("mse_strength", 1.0, float, False))
        data.append(("mae_strength", 0.0, float, False))
        data.append(("log_cosh_strength", 0.0, float, False))
        data.append(("vb_loss_strength", 1.0, float, False))
        data.append(("loss_weight_fn", LossWeight.CONSTANT, LossWeight, False))
        data.append(("loss_weight_strength", 5.0, float, False))
        data.append(("dropout_probability", 0.0, float, False))
        data.append(("loss_scaler", LossScaler.NONE, LossScaler, False))
        data.append(("learning_rate_scaler", LearningRateScaler.NONE, LearningRateScaler, False))
        data.append(("clip_grad_norm", 1.0, float, True))

        # noise
        data.append(("offset_noise_weight", 0.0, float, False))
        data.append(("perturbation_noise_weight", 0.0, float, False))
        data.append(("rescale_noise_scheduler_to_zero_terminal_snr", False, bool, False))
        data.append(("force_v_prediction", False, bool, False))
        data.append(("force_epsilon_prediction", False, bool, False))
        data.append(("min_noising_strength", 0.0, float, False))
        data.append(("max_noising_strength", 1.0, float, False))
        data.append(("timestep_distribution", TimestepDistribution.UNIFORM, TimestepDistribution, False))
        data.append(("noising_weight", 0.0, float, False))
        data.append(("noising_bias", 0.0, float, False))

        # unet
        unet = TrainModelPartConfig.default_values()
        unet.train = True
        unet.stop_training_after = 0
        unet.learning_rate = None
        unet.weight_dtype = DataType.NONE
        data.append(("unet", unet, TrainModelPartConfig, False))

        # prior
        prior = TrainModelPartConfig.default_values()
        prior.model_name = ""
        prior.train = True
        prior.stop_training_after = 0
        prior.learning_rate = None
        prior.weight_dtype = DataType.NONE
        data.append(("prior", prior, TrainModelPartConfig, False))

        # text encoder
        text_encoder = TrainModelPartConfig.default_values()
        text_encoder.train = True
        text_encoder.stop_training_after = 30
        text_encoder.stop_training_after_unit = TimeUnit.EPOCH
        text_encoder.learning_rate = None
        text_encoder.weight_dtype = DataType.NONE
        data.append(("text_encoder", text_encoder, TrainModelPartConfig, False))
        data.append(("text_encoder_layer_skip", 0, int, False))

        # text encoder 2
        text_encoder_2 = TrainModelPartConfig.default_values()
        text_encoder_2.train = True
        text_encoder_2.stop_training_after = 30
        text_encoder_2.stop_training_after_unit = TimeUnit.EPOCH
        text_encoder_2.learning_rate = None
        text_encoder_2.weight_dtype = DataType.NONE
        data.append(("text_encoder_2", text_encoder_2, TrainModelPartConfig, False))
        data.append(("text_encoder_2_layer_skip", 0, int, False))

        # text encoder 3
        text_encoder_3 = TrainModelPartConfig.default_values()
        text_encoder_3.train = True
        text_encoder_3.stop_training_after = 30
        text_encoder_3.stop_training_after_unit = TimeUnit.EPOCH
        text_encoder_3.learning_rate = None
        text_encoder_3.weight_dtype = DataType.NONE
        data.append(("text_encoder_3", text_encoder_3, TrainModelPartConfig, False))
        data.append(("text_encoder_3_layer_skip", 0, int, False))

        # vae
        vae = TrainModelPartConfig.default_values()
        vae.model_name = ""
        vae.weight_dtype = DataType.FLOAT_32
        data.append(("vae", vae, TrainModelPartConfig, False))

        # effnet encoder
        effnet_encoder = TrainModelPartConfig.default_values()
        effnet_encoder.model_name = ""
        effnet_encoder.weight_dtype = DataType.NONE
        data.append(("effnet_encoder", effnet_encoder, TrainModelPartConfig, False))

        # decoder
        decoder = TrainModelPartConfig.default_values()
        decoder.model_name = ""
        decoder.weight_dtype = DataType.NONE
        data.append(("decoder", decoder, TrainModelPartConfig, False))

        # decoder text encoder
        decoder_text_encoder = TrainModelPartConfig.default_values()
        decoder_text_encoder.weight_dtype = DataType.NONE
        data.append(("decoder_text_encoder", decoder_text_encoder, TrainModelPartConfig, False))

        # decoder vqgan
        decoder_vqgan = TrainModelPartConfig.default_values()
        decoder_vqgan.weight_dtype = DataType.NONE
        data.append(("decoder_vqgan", decoder_vqgan, TrainModelPartConfig, False))

        # masked training
        data.append(("masked_training", False, bool, False))
        data.append(("unmasked_probability", 0.1, float, False))
        data.append(("unmasked_weight", 0.1, float, False))
        data.append(("normalize_masked_area_loss", False, bool, False))

        # embedding
        data.append(("embedding_learning_rate", None, float, True))
        data.append(("preserve_embedding_norm", False, bool, False))
        data.append(("embedding", TrainEmbeddingConfig.default_values(), TrainEmbeddingConfig, False))
        data.append(("additional_embeddings", [], list[TrainEmbeddingConfig], False))
        data.append(("embedding_weight_dtype", DataType.FLOAT_32, DataType, False))

        # lora
        data.append(("peft_type", PeftType.LORA, PeftType, False))
        data.append(("lora_model_name", "", str, False))
        data.append(("lora_rank", 16, int, False))
        data.append(("lora_alpha", 1.0, float, False))
        data.append(("lora_decompose", False, bool, False))
        data.append(("lora_decompose_norm_epsilon", True, bool, False))
        data.append(("lora_weight_dtype", DataType.FLOAT_32, DataType, False))
        data.append(("lora_layers", "", str, False))
        data.append(("lora_layer_preset", None, str, True))
        data.append(("bundle_additional_embeddings", True, bool, False))

        # optimizer
        data.append(("optimizer", TrainOptimizerConfig.default_values(), TrainOptimizerConfig, False))
        data.append(("optimizer_defaults", {}, dict[str, TrainOptimizerConfig], False))

        # sample settings
        data.append(("sample_definition_file_name", "training_samples/samples.json", str, False))
        data.append(("samples", None, list[SampleConfig], True))
        data.append(("sample_after", 10, int, False))
        data.append(("sample_after_unit", TimeUnit.MINUTE, TimeUnit, False))
        data.append(("sample_image_format", ImageFormat.JPG, ImageFormat, False))
        data.append(("samples_to_tensorboard", True, bool, False))
        data.append(("non_ema_sampling", True, bool, False))

        # backup settings
        data.append(("backup_after", 30, int, False))
        data.append(("backup_after_unit", TimeUnit.MINUTE, TimeUnit, False))
        data.append(("rolling_backup", False, bool, False))
        data.append(("rolling_backup_count", 3, int, False))
        data.append(("backup_before_save", True, bool, False))
        data.append(("save_every", 0, int, False))
        data.append(("save_every_unit", TimeUnit.NEVER, TimeUnit, False))
        data.append(("save_skip_first", 0, int, False))
        data.append(("save_filename_prefix", "", str, False))

        # secrets
        secrets = SecretsConfig.default_values()
        data.append(("secrets", secrets, SecretsConfig, False))

        return TrainConfig(data)
