sgl-project--sglang/python/sglang/multimodal_gen/configs/sample/sampling_params.py

# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

# SPDX-License-Identifier: Apache-2.0
import argparse
import dataclasses
import hashlib
import json
import math
import os
import os.path
import re
import time
import unicodedata
import uuid
from dataclasses import dataclass, field
from enum import Enum, auto
from typing import TYPE_CHECKING, Any, ClassVar

from sglang.multimodal_gen.configs.post_training import RLRolloutArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.utils import StoreBoolean, expand_path_fields

logger = init_logger(__name__)

if TYPE_CHECKING:
    from sglang.multimodal_gen.runtime.server_args import ServerArgs


def _json_safe(obj: Any):
    """
    Recursively convert objects to JSON-serializable forms.
    - Enums -> their name
    - Callables -> stable module-qualified name
    - Sets/Tuples -> lists
    - Dicts/Lists -> recursively processed
    """
    if isinstance(obj, Enum):
        return obj.name
    if callable(obj):
        module = getattr(obj, "__module__", None)
        qualname = getattr(obj, "__qualname__", getattr(obj, "__name__", repr(obj)))
        return f"{module}.{qualname}" if module else qualname
    if isinstance(obj, dict):
        return {k: _json_safe(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple, set)):
        return [_json_safe(v) for v in obj]
    return obj


def generate_request_id() -> str:
    return str(uuid.uuid4())


def _sanitize_filename(name: str, replacement: str = "_", max_length: int = 150) -> str:
    """Create a filesystem- and ffmpeg-friendly filename.

    - Normalize to ASCII (drop accents and unsupported chars)
    - Replace spaces with underscores
    - Replace any char not in [A-Za-z0-9_.-] with replacement
    - Collapse multiple underscores
    - Trim leading/trailing dots/underscores and limit length
    """
    normalized = unicodedata.normalize("NFKD", name)
    ascii_name = normalized.encode("ascii", "ignore").decode("ascii")
    ascii_name = ascii_name.replace(" ", "_")
    ascii_name = re.sub(r"[^A-Za-z0-9._-]", replacement, ascii_name)
    ascii_name = re.sub(r"_+", "_", ascii_name).strip("._")
    if not ascii_name:
        ascii_name = "output"
    if max_length and len(ascii_name) > max_length:
        ascii_name = ascii_name[:max_length]
    return ascii_name


class DataType(Enum):
    IMAGE = auto()
    VIDEO = auto()
    MESH = auto()
    ACTION = auto()

    def get_default_extension(self) -> str:
        if self == DataType.IMAGE:
            return "png"
        if self == DataType.VIDEO:
            return "mp4"
        if self == DataType.ACTION:
            return "json"
        return "glb"


@dataclass
class SamplingParams:
    """
    Sampling parameters for generation.

    Dynamic batching compares these fields for compatibility, except fields
    marked with `batch_sig_exclude`.
    """

    data_type: DataType = DataType.VIDEO

    request_id: str | None = field(default=None, metadata={"batch_sig_exclude": True})

    # All fields below are copied from ForwardBatch

    # Image inputs
    image_path: str | list[str] | None = None

    # Video inputs (video-to-video conditioning)
    video_path: str | list[str] | None = None

    # Text inputs
    prompt: str | list[str] | None = field(
        default=None, metadata={"batch_sig_exclude": True}
    )
    negative_prompt: str = (
        "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
    )
    prompt_path: str | None = field(default=None, metadata={"batch_sig_exclude": True})
    output_path: str | None = field(default=None, metadata={"batch_sig_exclude": True})
    output_file_name: str | None = field(
        default=None, metadata={"batch_sig_exclude": True}
    )
    output_quality: str | None = "default"
    output_compression: int | None = None

    # Frame interpolation
    enable_frame_interpolation: bool = False
    frame_interpolation_exp: int = 1  # 1=2x, 2=4x
    frame_interpolation_scale: float = 1.0  # RIFE inference scale (0.5 for high-res)
    frame_interpolation_model_path: str | None = (
        None  # local dir or HF repo ID with flownet.pkl (default: elfgum/RIFE-4.22.lite)
    )

    # Upscaling
    enable_upscaling: bool = False
    upscaling_model_path: str | None = (
        None  # local .pth, HF repo ID, or repo_id:filename (default: ai-forever/Real-ESRGAN)
    )
    upscaling_scale: int = 4

    # Batch info
    num_outputs_per_prompt: int = 1
    seed: int | list[int] = field(default=42, metadata={"batch_sig_exclude": True})
    generator_device: str | None = None  # None means use the pipeline/model default

    # Original dimensions (before VAE scaling)
    num_frames: int = 1  # Default for image models
    num_frames_round_down: bool = (
        False  # Whether to round down num_frames if it's not divisible by num_gpus
    )

    # Subclasses can set these to provide model-specific default resolutions.
    # The base __post_init__ will apply them when height/width are not provided.
    _default_height: ClassVar[int | None] = None
    _default_width: ClassVar[int | None] = None

    height: int | None = None
    width: int | None = None
    fps: int = 24

    # Resolution validation
    supported_resolutions: list[tuple[int, int]] | None = field(
        default=None, metadata={"batch_sig_exclude": True}
    )  # None means all resolutions allowed

    # Output audio duration in seconds (models without an audio modality ignore this).
    sound_duration: float = 0.0

    # Denoising parameters
    num_inference_steps: int = None
    guidance_scale: float = 1.0
    guidance_scale_2: float = None
    true_cfg_scale: float = None  # for CFG vs guidance distillation (e.g., QwenImage)
    guidance_rescale: float = 0.0
    cfg_normalization: float | bool = 0.0
    boundary_ratio: float | None = None

    progressive_mode: str = "fullres"
    progressive_levels: int = 1
    progressive_delta: float = 0.01

    # TeaCache parameters
    enable_teacache: bool = False
    teacache_params: Any = (
        None  # TeaCacheParams or WanTeaCacheParams, set by model-specific subclass
    )

    # Profiling
    profile: bool = field(default=False, metadata={"batch_sig_exclude": True})
    num_profiled_timesteps: int = field(default=5, metadata={"batch_sig_exclude": True})
    profile_all_stages: bool = field(
        default=False, metadata={"batch_sig_exclude": True}
    )

    # Debugging
    debug: bool = field(default=False, metadata={"batch_sig_exclude": True})
    perf_dump_path: str | None = field(
        default=None, metadata={"batch_sig_exclude": True}
    )

    # Misc
    save_output: bool = True
    return_frames: bool = field(default=False, metadata={"batch_sig_exclude": True})
    rollout: bool = False
    rollout_sde_type: str = "sde"
    rollout_noise_level: float = 0.7
    rollout_log_prob_no_const: bool = False  # exclude constants in rollout logprob
    rollout_debug_mode: bool = (
        False  # return rollout debug tensors (intermediate states)
    )
    return_trajectory_latents: bool = False  # returns all latents for each timestep
    return_trajectory_decoded: bool = False  # returns decoded latents for each timestep
    rollout_return_denoising_env: bool = (
        False  # populate ``denoising_env`` (image/pos/neg kwargs, guidance) for RL replay
    )
    rollout_return_dit_trajectory: bool = (
        False  # per-step noisy latents + final latent + timesteps (RolloutDitTrajectory)
    )
    # 0-indexed denoising-loop step filters; None = all steps.
    rollout_sde_step_indices: list[int] | None = None
    rollout_return_step_indices: list[int] | None = None
    # if True, disallow user params to override subclass-defined protected fields
    no_override_protected_fields: bool = field(
        default=False, metadata={"batch_sig_exclude": True}
    )
    # whether to adjust num_frames for multi-GPU friendly splitting (default: True)
    adjust_frames: bool = True
    # if True, suppress verbose logging for this request
    suppress_logs: bool = field(default=False, metadata={"batch_sig_exclude": True})

    # return output file paths directly to client
    return_file_paths_only: bool = True
    enable_sequence_shard: bool | None = None
    diffusers_kwargs: dict | None = None
    max_sequence_length: int | None = None
    flow_shift: float | None = None

    # cosmos-related
    use_duration_template: bool | None = None
    use_resolution_template: bool | None = None
    use_system_prompt: bool | None = None
    use_guardrails: bool | None = None
    condition_inputs: dict[str, Any] = field(default_factory=dict)
    realtime_chunk_size: int | None = None

    # Prompt enhancement (ErnieImage)
    use_pe: bool | None = None

    def _set_output_file_ext(self):
        # add extension if needed
        output_extensions = (".mp4", ".jpg", ".png", ".webp", ".obj", ".glb", ".json")
        if not any(self.output_file_name.endswith(ext) for ext in output_extensions):
            self.output_file_name = (
                f"{self.output_file_name}.{self.data_type.get_default_extension()}"
            )

    def _set_output_file_name(self):
        # settle output_file_name
        if (
            self.output_file_name is None
            and self.prompt
            and isinstance(self.prompt, str)
        ):
            # generate a random filename
            # get a hash of current params
            params_dict = dataclasses.asdict(self)
            # Avoid recursion
            params_dict["output_file_name"] = ""

            # Convert to a stable JSON string
            params_str = json.dumps(_json_safe(params_dict), sort_keys=True)
            # Create a hash
            hasher = hashlib.sha256()
            hasher.update(params_str.encode("utf-8"))
            param_hash = hasher.hexdigest()[:8]

            timestamp = time.strftime("%Y%m%d-%H%M%S")
            base = f"{self.prompt[:100]}_{timestamp}_{param_hash}"
            self.output_file_name = base

        if self.output_file_name is None:
            timestamp = time.strftime("%Y%m%d-%H%M%S")
            self.output_file_name = f"output_{timestamp}"

        self.output_file_name = _sanitize_filename(self.output_file_name)

        # Ensure a proper extension is present
        self._set_output_file_ext()

    def __post_init__(self) -> None:
        assert self.num_frames >= 1

        if self.width is None and self._default_width is not None:
            self.width = self._default_width
        if self.height is None and self._default_height is not None:
            self.height = self._default_height

        # Handle output_quality to output_compression conversion
        if self.output_compression is None and self.output_quality is not None:
            self.output_compression = self._adjust_output_quality(
                self.output_quality, self.data_type
            )

        self._validate()

        # Allow env var to override num_inference_steps (for faster CI testing on AMD)
        env_steps = os.environ.get("SGLANG_TEST_NUM_INFERENCE_STEPS")
        if env_steps is not None and self.num_inference_steps is not None:
            self.num_inference_steps = int(env_steps)

    def build_request_extra(self) -> dict[str, Any]:
        """Return optional request-scoped extras for downstream pipeline stages."""
        extra = {}
        diffusers_kwargs = getattr(self, "diffusers_kwargs", None)
        if diffusers_kwargs:
            extra["diffusers_kwargs"] = diffusers_kwargs
        explicit_fields = getattr(self, "_explicit_fields", None)
        if explicit_fields is not None:
            extra["explicit_fields"] = sorted(explicit_fields)
        return extra

    def apply_request_extra(self, req: Any) -> None:
        """Merge request extras (model specific, e.g., LTX2.3) into an already-created pipeline request."""
        req.extra.update(self.build_request_extra())
        if self.condition_inputs:
            req.condition_inputs.update(self.condition_inputs)
        if self.realtime_chunk_size is not None:
            req.realtime_chunk_size = self.realtime_chunk_size

    def _adjust_output_quality(self, output_quality: str, data_type: DataType) -> int:
        """Convert output_quality string to compression level."""
        if data_type == DataType.ACTION:
            return 0
        output_quality_mapper = {"maximum": 100, "high": 90, "medium": 55, "low": 35}
        if output_quality == "default":
            return 50 if data_type == DataType.VIDEO else 75
        return output_quality_mapper.get(output_quality)

    def _validate(self):
        """
        check if the sampling params is correct by itself
        """
        if self.prompt_path and not self.prompt_path.endswith(".txt"):
            raise ValueError(
                f"prompt_path must be a txt file, got {self.prompt_path!r}"
            )

        # These are always required to be sane regardless of pipeline.
        if (
            not isinstance(self.num_outputs_per_prompt, int)
            or self.num_outputs_per_prompt <= 0
        ):
            raise ValueError(
                f"num_outputs_per_prompt must be a positive int, got {self.num_outputs_per_prompt!r}"
            )

        if isinstance(self.seed, list):
            if not self.seed:
                raise ValueError("seed list must not be empty")
            for seed in self.seed:
                if isinstance(seed, bool) or not isinstance(seed, int) or seed < 0:
                    raise ValueError(
                        f"seed list must contain non-negative ints, got {self.seed!r}"
                    )
        elif (
            isinstance(self.seed, bool)
            or not isinstance(self.seed, int)
            or self.seed < 0
        ):
            raise ValueError(
                f"seed must be a non-negative int or list of ints, got {self.seed!r}"
            )

        # Used by seconds() and video writer; fps <= 0 is always invalid.
        if not isinstance(self.fps, int) or self.fps <= 0:
            raise ValueError(f"fps must be a positive int, got {self.fps!r}")

        # num_frames is already asserted in __post_init__, but keep a friendly error here too
        # (e.g., when validation is triggered from other code paths).
        if not isinstance(self.num_frames, int) or self.num_frames <= 0:
            raise ValueError(
                f"num_frames must be a positive int, got {self.num_frames!r}"
            )

        if self.num_inference_steps is not None:
            if (
                not isinstance(self.num_inference_steps, int)
                or self.num_inference_steps <= 0
            ):
                raise ValueError(
                    f"num_inference_steps must be a positive int, got {self.num_inference_steps!r}"
                )

        if self.progressive_mode not in ("fullres", "dct", "dct_rewind"):
            raise ValueError(
                "progressive_mode must be one of 'fullres', 'dct', or "
                f"'dct_rewind', got {self.progressive_mode!r}"
            )
        if (
            isinstance(self.progressive_levels, bool)
            or not isinstance(self.progressive_levels, int)
            or self.progressive_levels <= 0
        ):
            raise ValueError(
                f"progressive_levels must be a positive int, got {self.progressive_levels!r}"
            )
        if (
            isinstance(self.progressive_delta, bool)
            or not isinstance(self.progressive_delta, (int, float))
            or not 0 < float(self.progressive_delta) < 1
        ):
            raise ValueError(
                f"progressive_delta must be in (0, 1), got {self.progressive_delta!r}"
            )

        # Numeric hyperparams should not be NaN/Inf and should be within basic ranges.
        # Note: bool is a subclass of int; reject it explicitly to avoid silent surprises.
        def _finite_non_negative_float(
            name: str, value: Any, allow_none: bool = True
        ) -> None:
            if value is None and allow_none:
                return
            if isinstance(value, bool) or not isinstance(value, (int, float)):
                raise ValueError(f"{name} must be a number, got {value!r}")
            if not math.isfinite(float(value)):
                raise ValueError(f"{name} must be finite, got {value!r}")
            if float(value) < 0.0:
                raise ValueError(f"{name} must be non-negative, got {value!r}")

        _finite_non_negative_float(
            "guidance_scale", self.guidance_scale, allow_none=True
        )
        _finite_non_negative_float(
            "guidance_scale_2", self.guidance_scale_2, allow_none=True
        )
        _finite_non_negative_float(
            "true_cfg_scale", self.true_cfg_scale, allow_none=True
        )
        _finite_non_negative_float(
            "guidance_rescale", self.guidance_rescale, allow_none=False
        )

        if self.cfg_normalization is None:
            self.cfg_normalization = 0.0
        elif isinstance(self.cfg_normalization, bool):
            self.cfg_normalization = 1.0 if self.cfg_normalization else 0.0

        if self.boundary_ratio is not None:
            if isinstance(self.boundary_ratio, bool) or not isinstance(
                self.boundary_ratio, (int, float)
            ):
                raise ValueError(
                    f"boundary_ratio must be a number, got {self.boundary_ratio!r}"
                )
            if not math.isfinite(float(self.boundary_ratio)):
                raise ValueError(
                    f"boundary_ratio must be finite, got {self.boundary_ratio!r}"
                )
            if not (0.0 <= float(self.boundary_ratio) <= 1.0):
                raise ValueError(
                    f"boundary_ratio must be within [0, 1], got {self.boundary_ratio!r}"
                )

        RLRolloutArgs.validate_sampling_params(self)

    def check_sampling_param(self):
        # Keep backward-compatibility for old call sites.
        self._validate()

    def _validate_with_pipeline_config(self, pipeline_config):
        """
        check if the sampling params is compatible and valid with server_args
        """
        task_type = pipeline_config.task_type
        if task_type.is_action_gen():
            return

        if task_type.requires_image_input():
            # requires image input
            if self.image_path is None:
                raise ValueError(
                    f"Served model with task type '{task_type.name}' requires an 'image_path' input, but none was provided"
                )

        if not task_type.accepts_image_input():
            # does not support image input
            if self.image_path is not None:
                raise ValueError(
                    f"input_reference is not supported for {task_type.name} models."
                )

    def _adjust(
        self,
        server_args,
    ):
        """
        final adjustment, called after merged with user params
        """
        expand_path_fields(self)

        # TODO: SamplingParams should not rely on ServerArgs
        pipeline_config = server_args.pipeline_config
        task_type = pipeline_config.task_type
        self.data_type = task_type.data_type()

        self._adjust_output_path(server_args)
        if task_type.is_action_gen():
            self._adjust_action_fields(server_args)
            return

        if task_type.is_mesh_gen():
            self._adjust_mesh_fields(server_args, pipeline_config)
            return

        if task_type.is_visual_gen():
            self._adjust_visual_fields(server_args, pipeline_config)

    def _adjust_output_path(self, server_args):
        if self.output_path is None:
            if server_args.output_path is not None:
                self.output_path = server_args.output_path
                logger.debug(
                    f"Overriding output_path with server configuration: {self.output_path}"
                )
            else:
                self.save_output = False

    def _adjust_action_fields(self, server_args):
        self.return_file_paths_only = False
        self.num_frames = 1
        self.adjust_frames = False
        if self.save_output and not server_args.comfyui_mode:
            self._set_output_file_name()

    def _adjust_mesh_fields(self, server_args, pipeline_config):
        if self.guidance_scale is None:
            try:
                from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import (
                    Hunyuan3D2PipelineConfig,
                )

                if isinstance(pipeline_config, Hunyuan3D2PipelineConfig):
                    self.guidance_scale = pipeline_config.guidance_scale
                else:
                    self.guidance_scale = 1.0
            except ImportError:
                self.guidance_scale = 1.0
        self.return_frames = False
        self.return_video = False
        self.num_frames = 1
        self.adjust_frames = False
        if self.save_output and not server_args.comfyui_mode:
            self._set_output_file_name()

    def _adjust_visual_fields(self, server_args, pipeline_config):
        if self.guidance_scale is None:
            self.guidance_scale = 1.0

        # Process negative prompt
        if self.negative_prompt is not None and not self.negative_prompt.isspace():
            # avoid stripping default negative prompt: ' ' for qwen-image
            self.negative_prompt = self.negative_prompt.strip()

        # Validate dimensions
        if self.num_frames <= 0:
            raise ValueError(
                f"height, width, and num_frames must be positive integers, got "
                f"height={self.height}, width={self.width}, "
                f"num_frames={self.num_frames}"
            )

        # Validate resolution against pipeline-specific supported resolutions
        if self.height is None and self.width is None:
            if self.supported_resolutions is not None:
                self.width, self.height = self.supported_resolutions[0]
                logger.info(
                    f"Resolution unspecified, using default: {self.supported_resolutions[0]}"
                )

        if self.height is not None and self.width is not None:
            if self.supported_resolutions is not None:
                if (self.width, self.height) not in self.supported_resolutions:
                    supported_str = ", ".join(
                        [f"{w}x{h}" for w, h in self.supported_resolutions]
                    )
                    error_msg = (
                        f"Unsupported resolution: {self.width}x{self.height}, output quality may suffer. "
                        f"Supported resolutions: {supported_str}"
                    )
                    logger.warning(error_msg)

        pipeline_name_lower = server_args.pipeline_config.__class__.__name__.lower()

        if (
            "wan" in pipeline_name_lower
            or "helios" in pipeline_name_lower
            or "joy" in pipeline_name_lower
            or "cosmos3" in pipeline_name_lower
        ) and (self.enable_sequence_shard is None or self.enable_sequence_shard):
            self.enable_sequence_shard = True
            logger.debug("Automatically enabled enable_sequence_shard")
        else:
            self.enable_sequence_shard = False

        if self.enable_sequence_shard:
            self.adjust_frames = False
            logger.info(
                "Sequence dimension shard is enabled, disabling frame adjustment for better performance"
            )

        if pipeline_config.task_type.is_image_gen():
            # settle num_frames
            if not server_args.pipeline_config.allow_set_num_frames():
                logger.debug("Setting `num_frames` to 1 for image generation model")
                self.num_frames = 1
        else:
            # mandatory frame adjusting logic, mod
            # NOTE: We must apply adjust_num_frames BEFORE the SP alignment logic below.
            # If we apply it after, adjust_num_frames might modify the frame count
            # and break the divisibility constraint (alignment) required by num_gpus.
            original_num_frames = self.num_frames
            self.num_frames = server_args.pipeline_config.adjust_num_frames(
                original_num_frames
            )
            logger.info(
                "Adjusting number of frames from %s to %s based on model",
                original_num_frames,
                self.num_frames,
            )

            if self.adjust_frames:
                # Adjust number of frames based on number of GPUs for video task
                use_temporal_scaling_frames = (
                    pipeline_config.vae_config.use_temporal_scaling_frames
                )
                num_frames = self.num_frames
                num_gpus = server_args.num_gpus
                temporal_scale_factor = (
                    pipeline_config.vae_config.arch_config.temporal_compression_ratio
                )

                if use_temporal_scaling_frames:
                    orig_latent_num_frames = (
                        num_frames - 1
                    ) // temporal_scale_factor + 1
                else:
                    orig_latent_num_frames = num_frames

                if orig_latent_num_frames % server_args.num_gpus != 0:
                    # Adjust latent frames to be divisible by number of GPUs
                    if self.num_frames_round_down:
                        # Ensure we have at least 1 batch per GPU
                        new_latent_num_frames = (
                            max(1, (orig_latent_num_frames // num_gpus)) * num_gpus
                        )
                    else:
                        new_latent_num_frames = (
                            math.ceil(orig_latent_num_frames / num_gpus) * num_gpus
                        )

                    if use_temporal_scaling_frames:
                        # Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor
                        new_num_frames = (
                            new_latent_num_frames - 1
                        ) * temporal_scale_factor + 1
                    else:
                        new_num_frames = new_latent_num_frames

                    logger.info(
                        "Adjusting number of frames from %s to %s based on number of GPUs (%s)",
                        self.num_frames,
                        new_num_frames,
                        server_args.num_gpus,
                    )
                    self.num_frames = new_num_frames

        if not server_args.comfyui_mode:
            self._set_output_file_name()

    @classmethod
    def from_pretrained(cls, model_path: str, **kwargs) -> "SamplingParams":
        from sglang.multimodal_gen.registry import get_model_info

        backend = kwargs.pop("backend", None)
        model_id = kwargs.pop("model_id", None)
        model_info = get_model_info(model_path, backend=backend, model_id=model_id)
        sampling_params: SamplingParams = model_info.sampling_param_cls(**kwargs)
        return sampling_params

    @staticmethod
    def from_user_sampling_params_args(
        model_path: str, server_args: "ServerArgs", *args, **kwargs
    ):
        pipeline_class_name = getattr(server_args, "pipeline_class_name", None)
        try:
            sampling_params = None
            if pipeline_class_name:
                from sglang.multimodal_gen.registry import get_pipeline_config_classes

                config_classes = get_pipeline_config_classes(pipeline_class_name)
                if config_classes is not None:
                    _, sampling_params_cls = config_classes
                    sampling_params = sampling_params_cls()

            if sampling_params is None:
                sampling_params = SamplingParams.from_pretrained(
                    model_path,
                    backend=server_args.backend,
                    model_id=server_args.model_id,
                )
        except (AttributeError, ValueError):
            # Handle safetensors files or other cases where model_index.json is not available
            # Use appropriate SamplingParams based on pipeline_class_name from registry
            if os.path.isfile(model_path) and model_path.endswith(".safetensors"):
                # Determine which sampling params to use based on pipeline_class_name
                from sglang.multimodal_gen.registry import get_pipeline_config_classes

                config_classes = (
                    get_pipeline_config_classes(pipeline_class_name)
                    if pipeline_class_name
                    else None
                )

                if config_classes is not None:
                    _, sampling_params_cls = config_classes
                    try:
                        sampling_params = sampling_params_cls()
                        logger.info(
                            f"Using {sampling_params_cls.__name__} for {pipeline_class_name} safetensors file (no model_index.json): %s",
                            model_path,
                        )
                    except Exception as import_error:
                        logger.warning(
                            f"Failed to instantiate {sampling_params_cls.__name__}: {import_error}. "
                            "Using default SamplingParams"
                        )
                        sampling_params = SamplingParams()
                else:
                    raise ValueError(
                        f"Could not get pipeline config classes for {pipeline_class_name}"
                    )
            else:
                # Re-raise if it's not a safetensors file issue
                raise

        user_kwargs = dict(kwargs)
        diffusers_kwargs = user_kwargs.pop("diffusers_kwargs", None)

        user_sampling_params = type(sampling_params)(*args, **user_kwargs)
        # TODO: refactor
        sampling_params._merge_with_user_params(
            user_sampling_params, explicit_fields=set(user_kwargs.keys())
        )
        sampling_params._explicit_fields = set(user_kwargs.keys())
        if diffusers_kwargs is not None:
            sampling_params.diffusers_kwargs = diffusers_kwargs
        sampling_params._adjust(server_args)

        sampling_params._validate_with_pipeline_config(server_args.pipeline_config)

        return sampling_params

    def output_size_str(self) -> str:
        return f"{self.width}x{self.height}"

    def seconds(self) -> float:
        return self.num_frames / self.fps

    @staticmethod
    def add_cli_args(parser: Any) -> Any:
        """Add CLI arguments for SamplingParam fields"""

        def add_argument(*name_or_flags, **kwargs):
            kwargs.setdefault("default", argparse.SUPPRESS)
            return parser.add_argument(*name_or_flags, **kwargs)

        add_argument("--data-type", type=str, nargs="+")
        add_argument(
            "--num-frames-round-down",
            action="store_true",
        )
        add_argument(
            "--enable-teacache",
            action="store_true",
        )

        # profiling
        add_argument(
            "--profile",
            action="store_true",
            help="Enable torch profiler for denoising stage",
        )
        add_argument(
            "--num-profiled-timesteps",
            type=int,
            help="Number of timesteps to profile after warmup",
        )
        add_argument(
            "--profile-all-stages",
            action="store_true",
            dest="profile_all_stages",
            help="Used with --profile, profile all pipeline stages",
        )

        add_argument(
            "--debug",
            action="store_true",
            help="",
        )

        # Progressive resolution growing (DCT spectral upsampling)
        add_argument(
            "--progressive-mode",
            type=str,
            dest="progressive_mode",
            choices=["fullres", "dct", "dct_rewind"],
            help="Progressive resolution mode. 'fullres' disables (default). "
            "'dct_rewind' uses DCT-II upsample + scheduler sigma rewind (recommended).",
        )
        add_argument(
            "--progressive-levels",
            type=int,
            dest="progressive_levels",
            help="Number of resolution halvings for progressive generation (default: 1).",
        )
        add_argument(
            "--progressive-delta",
            type=float,
            dest="progressive_delta",
            help="Noise-dominated tolerance δ for stage-transition thresholds (default: 0.01).",
        )

        add_argument(
            "--prompt",
            type=str,
            nargs="+",
            help="Text prompt(s) for generation. Use space-separated values for multiple prompts, e.g., --prompt 'prompt 1' 'prompt 2'",
        )
        add_argument(
            "--negative-prompt",
            type=str,
            help="Negative text prompt for generation",
        )
        add_argument(
            "--prompt-path",
            type=str,
            help="Path to a text file containing prompts (one per line)",
        )
        add_argument(
            "--output-file-name",
            type=str,
            help="Name of the output file",
        )
        add_argument(
            "--output-quality",
            type=str,
            help="Output quality setting (default, low, medium, high, maximum)",
        )
        add_argument(
            "--output-compression",
            type=int,
            help="Output compression level (0-100, higher means better quality but larger file size)",
        )
        add_argument(
            "--num-outputs-per-prompt",
            type=int,
            help="Number of outputs to generate per prompt",
        )
        add_argument(
            "--seed",
            type=int,
            nargs="+",
            help="Random seed for generation",
        )
        add_argument(
            "--generator-device",
            type=str,
            choices=["cuda", "musa", "cpu"],
            help="Device for random generator (cuda, musa or cpu). Default: use the model-specific setting.",
        )
        add_argument(
            "--num-frames",
            type=int,
            help="Number of frames to generate",
        )
        add_argument(
            "--sound-duration",
            type=float,
            help="Duration of generated audio in seconds; 0 disables audio output (audio-capable models only)",
        )
        add_argument(
            "--height",
            type=int,
            help="Height of generated output",
        )
        add_argument(
            "--width",
            type=int,
            help="Width of generated output",
        )
        # resolution shortcuts
        add_argument(
            "--4k",
            action="store_true",
            dest="resolution_4k",
            help="Set resolution to 4K (3840x2160)",
        )
        add_argument(
            "--2k",
            action="store_true",
            dest="resolution_2k",
            help="Set resolution to 2K (2560x1440)",
        )
        add_argument(
            "--1080p",
            action="store_true",
            dest="resolution_1080p",
            help="Set resolution to 1080p (1920x1080)",
        )
        add_argument(
            "--720p",
            action="store_true",
            dest="resolution_720p",
            help="Set resolution to 720p (1280x720)",
        )

        add_argument(
            "--fps",
            type=int,
            help="Frames per second for saved output",
        )
        add_argument(
            "--num-inference-steps",
            type=int,
            help="Number of denoising steps",
        )
        add_argument(
            "--guidance-scale",
            type=float,
            help="Classifier-free guidance scale",
        )
        add_argument(
            "--guidance-scale-2",
            type=float,
            dest="guidance_scale_2",
            help="Secondary guidance scale for dual-guidance models (e.g., Wan low-noise expert)",
        )
        add_argument(
            "--true-cfg-scale",
            type=float,
            dest="true_cfg_scale",
            help="True CFG scale for models that distinguish distilled guidance from standard CFG (e.g., Qwen-Image)",
        )
        add_argument(
            "--guidance-rescale",
            type=float,
            help="Guidance rescale factor",
        )
        add_argument(
            "--cfg-normalization",
            type=float,
            dest="cfg_normalization",
            help="CFG renormalization factor (for Z-Image). ",
        )
        add_argument(
            "--boundary-ratio",
            type=float,
            help="Boundary timestep ratio",
        )
        add_argument(
            "--save-output",
            action="store_true",
            help="Whether to save the output to disk",
        )
        add_argument(
            "--no-save-output",
            action="store_false",
            dest="save_output",
            help="Don't save the output to disk",
        )
        add_argument(
            "--return-frames",
            action="store_true",
            help="Whether to return the raw frames",
        )
        add_argument(
            "--image-path",
            "--image-paths",
            type=str,
            nargs="+",
            help=(
                "Path(s) to input image(s) for image-to-image / image-to-video "
                "generation. For multiple images, pass them space-separated "
                "(alias: --image-paths), e.g.: "
                '--image-paths "img1.png" "img2.png"'
            ),
        )
        add_argument(
            "--action",
            type=str,
            help=(
                "Action input. SANA-WM uses a WASD/IJKL DSL, e.g. "
                "'w-80,jw-40,w-40,lw-60,w-100'. Cosmos3 uses a JSON array of "
                "shape [T, D], e.g. '[[0.1, 0.2, ...], ...]'. Model-specific "
                "fields are ignored by other pipelines."
            ),
        )
        add_argument(
            "--translation-speed",
            "--translation_speed",
            type=float,
            dest="translation_speed",
            help="SANA-WM action DSL per-frame translation speed.",
        )
        add_argument(
            "--rotation-speed-deg",
            "--rotation_speed_deg",
            type=float,
            dest="rotation_speed_deg",
            help="SANA-WM action DSL per-frame rotation speed in degrees.",
        )
        add_argument(
            "--pitch-limit-deg",
            "--pitch_limit_deg",
            type=float,
            dest="pitch_limit_deg",
            help="SANA-WM action DSL absolute pitch clamp in degrees.",
        )
        add_argument(
            "--video-path",
            type=str,
            nargs="+",
            help=(
                "Path(s) to input video(s) for video-to-video generation. "
                "The first/last frames of the video become the conditioning "
                "frames for the generated output."
            ),
        )
        add_argument(
            "--action-mode",
            type=str,
            dest="action_mode",
            help=(
                "Cosmos3 action mode: 'forward_dynamics' (predict next frame "
                "from action), 'policy' (predict action from frame), or "
                "'inverse_dynamics' (predict action from two frames)."
            ),
        )
        add_argument(
            "--domain-id",
            type=int,
            dest="domain_id",
            help="Action embodiment domain ID (integer). Overrides --domain-name.",
        )
        add_argument(
            "--domain-name",
            type=str,
            dest="domain_name",
            help="Action embodiment domain name (e.g. 'av', 'camera_pose', 'umi').",
        )
        add_argument(
            "--raw-action-dim",
            type=int,
            dest="raw_action_dim",
            help=(
                "Number of active action dimensions to predict; remaining "
                "dimensions are zero-padded. Required for 'policy' and "
                "'inverse_dynamics' modes."
            ),
        )
        add_argument(
            "--action-fps",
            type=float,
            dest="action_fps",
            help=(
                "Frame rate used for action token temporal mRoPE positions. "
                "Defaults to the video fps when not set."
            ),
        )
        add_argument(
            "--moba-config-path",
            type=str,
            help="Path to a JSON file containing V-MoBA specific configurations.",
        )
        add_argument(
            "--return-trajectory-latents",
            action="store_true",
            help="Whether to return the trajectory",
        )

        # Rollout arguments
        RLRolloutArgs.add_cli_args(parser, add_argument=add_argument)

        add_argument(
            "--return-trajectory-decoded",
            action="store_true",
            help="Whether to return the decoded trajectory",
        )
        add_argument(
            "--diffusers-kwargs",
            type=str,
            help="JSON string of extra kwargs to pass to diffusers pipeline. "
            'Example: \'{"output_type": "latent", "clip_skip": 2}\'',
        )
        add_argument(
            "--no-override-protected-fields",
            action="store_true",
            help=(
                "If set, disallow user params to override fields defined in subclasses."
            ),
        )
        add_argument(
            "--adjust-frames",
            action=StoreBoolean,
            help=(
                "Enable/disable adjusting num_frames to evenly split latent frames across GPUs "
                "and satisfy model temporal constraints. If disabled, tokens might be padded for SP."
                "Default: true. Examples: --adjust-frames, --adjust-frames true, --adjust-frames false."
            ),
        )
        add_argument(
            "--return-file-paths-only",
            action=StoreBoolean,
            help="If set, output file will be saved early to get a performance boost, while output tensors will not be returned.",
        )
        add_argument(
            "--enable-sequence-shard",
            action=StoreBoolean,
            help="Enable sequence dimension shard with sequence parallelism.",
        )
        add_argument(
            "--enable-frame-interpolation",
            action="store_true",
            help="Enable post-generation frame interpolation using RIFE 4.22.lite.",
        )
        add_argument(
            "--frame-interpolation-exp",
            type=int,
            help="Frame interpolation exponent: 1=2x, 2=4x (default: 1).",
        )
        add_argument(
            "--frame-interpolation-scale",
            type=float,
            help="RIFE inference scale factor (default: 1.0; use 0.5 for high-res).",
        )
        add_argument(
            "--frame-interpolation-model-path",
            type=str,
            help="Local directory or HuggingFace repo ID containing RIFE flownet.pkl weights "
            "(default: elfgum/RIFE-4.22.lite, downloaded automatically). "
            "Only RIFE 4.22.lite architecture is supported; other RIFE versions or "
            "frame interpolation models are not compatible.",
        )
        add_argument(
            "--enable-upscaling",
            action="store_true",
            help="Enable post-generation upscaling using Real-ESRGAN.",
        )
        add_argument(
            "--upscaling-model-path",
            type=str,
            help="Local .pth file, HuggingFace repo ID, or repo_id:filename for Real-ESRGAN weights "
            "(default: ai-forever/Real-ESRGAN with RealESRGAN_x4.pth). "
            "Only RRDBNet (e.g. RealESRGAN_x4plus) and SRVGGNetCompact (e.g. realesr-animevideov3) "
            "architectures are supported; other super-resolution models are not compatible. "
            "Use 'repo_id:filename' to specify a custom weight file from a HF repo.",
        )
        add_argument(
            "--upscaling-scale",
            type=int,
            help="Upscaling factor (default: 4).",
        )
        return parser

    @classmethod
    def get_cli_args(cls, args: argparse.Namespace):
        # handle resolution shortcuts
        if hasattr(args, "resolution_4k") and args.resolution_4k:
            args.width = 3840
            args.height = 2160
        elif hasattr(args, "resolution_2k") and args.resolution_2k:
            args.width = 2560
            args.height = 1440
        elif hasattr(args, "resolution_1080p") and args.resolution_1080p:
            args.width = 1920
            args.height = 1080
        elif hasattr(args, "resolution_720p") and args.resolution_720p:
            args.width = 1280
            args.height = 720

        sampling_params_fields = {attr.name for attr in dataclasses.fields(cls)}
        args_attrs = set(vars(args).keys())
        attrs = sampling_params_fields & args_attrs
        cli_args = {
            attr: getattr(args, attr)
            for attr in attrs
            if hasattr(args, attr) and getattr(args, attr) is not None
        }
        if isinstance(cli_args.get("seed"), list) and len(cli_args["seed"]) == 1:
            cli_args["seed"] = cli_args["seed"][0]
        return cli_args

    def output_file_path(self):
        if self.output_path is None:
            return None
        return os.path.join(self.output_path, self.output_file_name)

    def _merge_with_user_params(
        self,
        user_params: "SamplingParams",
        explicit_fields: set[str] | None = None,
    ):
        """
        Merges parameters from a user-provided SamplingParams object.

        Args:
            explicit_fields: field names explicitly set by the user (e.g. from
                CLI kwargs). These are always treated as user-modified even when
                their value matches the base-class default.
        """
        if user_params is None:
            return

        predefined_fields = set(type(self).__annotations__.keys())

        # global switch: if True, allow overriding protected fields
        allow_override_protected = not user_params.no_override_protected_fields
        for field_info in dataclasses.fields(user_params):
            field_name = field_info.name
            user_value = getattr(user_params, field_name)
            if hasattr(SamplingParams, field_name):
                default_class_value = getattr(SamplingParams, field_name)
            elif field_info.default is not dataclasses.MISSING:
                default_class_value = field_info.default
            elif field_info.default_factory is not dataclasses.MISSING:
                default_class_value = field_info.default_factory()
            else:
                default_class_value = dataclasses.MISSING

            if explicit_fields is not None:
                is_user_modified = field_name in explicit_fields
            else:
                is_user_modified = user_value != default_class_value
            is_protected_field = field_name in predefined_fields
            if is_user_modified and (
                allow_override_protected or not is_protected_field
            ):
                setattr(self, field_name, user_value)
        if explicit_fields is not None:
            self._explicit_fields = set(explicit_fields)
        self.__post_init__()

    @property
    def n_tokens(self) -> int:
        # Calculate latent sizes
        if self.height and self.width:
            latents_size = [
                (self.num_frames - 1) // 4 + 1,
                self.height // 8,
                self.width // 8,
            ]
            n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
        else:
            n_tokens = -1
        return n_tokens


@dataclass
class CacheParams:
    cache_type: str = "none"