# SPDX-License-Identifier: Apache-2.0
"""
Generic sampling parameters for diffusers backend.

This module provides generic sampling parameters that work with any diffusers pipeline.
"""

from dataclasses import dataclass, field
from typing import Any, ClassVar

from sglang.multimodal_gen.configs.sample.sampling_params import (
    DataType,
    SamplingParams,
)


@dataclass
class DiffusersGenericSamplingParams(SamplingParams):
    """
    Generic sampling parameters for diffusers backend.

    These parameters cover the most common options across different diffusers pipelines.
    The diffusers pipeline will use whichever parameters it supports.

    For pipeline-specific parameters, use `diffusers_kwargs` dict which will be
    passed directly to the diffusers pipeline call.
    """

    _default_height: ClassVar[int] = 1024
    _default_width: ClassVar[int] = 1024

    # Override defaults with more conservative values that work across pipelines
    num_frames: int = 1  # default to image generation
    height: int = 1024
    width: int = 1024
    num_inference_steps: int = 30
    guidance_scale: float = 7.5
    negative_prompt: str = ""

    # extra kwargs to pass directly to the diffusers pipeline
    # example: {"output_type": "latent", "return_dict": False}
    diffusers_kwargs: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if self.num_frames > 1:
            self.data_type = DataType.VIDEO
        else:
            self.data_type = DataType.IMAGE

        super().__post_init__()