# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.

from .configuration_ministral_dlm import MinistralDLMConfig


class NemotronLabsDiffusionImageConfig(MinistralDLMConfig):
    model_type = "nemotron_labs_diffusion_image"

    def __init__(
        self,
        max_position_embeddings=262144,
        rope_parameters=None,
        rope_scaling=None,
        **kwargs,
    ):
        # Newer transformers standardize RoPE params during PretrainedConfig init.
        # Make these attributes available early and normalize legacy rope_scaling.
        if rope_parameters is None and rope_scaling is not None:
            rope_parameters = dict(rope_scaling)

        self.max_position_embeddings = max_position_embeddings
        if rope_parameters is not None:
            self.rope_parameters = rope_parameters
        if rope_scaling is not None:
            self.rope_scaling = rope_scaling

        super().__init__(
            max_position_embeddings=max_position_embeddings,
            rope_parameters=rope_parameters,
            rope_scaling=rope_scaling,
            **kwargs,
        )

        if getattr(self, "rope_parameters", None) is None and getattr(self, "rope_scaling", None) is not None:
            self.rope_parameters = dict(self.rope_scaling)