sgl-project--sglang/python/sglang/multimodal_gen/runtime/pipelines/ideogram.py

# SPDX-License-Identifier: Apache-2.0

import os
from dataclasses import dataclass
from functools import lru_cache
from typing import Any, cast

from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
    ComposedPipelineBase,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import (
    InputValidationStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ideogram import (
    Ideogram4DecodingStage,
    Ideogram4DenoisingStage,
    Ideogram4TextEncodingStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.denoising import (
    ProgressiveDenoisingStageRouter,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.ideogram import (
    Ideogram4ProgressiveDenoisingStage,
)
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
    maybe_download_model,
    verify_model_config_and_directory,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger

logger = init_logger(__name__)

_IDEOGRAM4_BASE_MODEL = "ideogram-ai/ideogram-4-fp8"
_IDEOGRAM4_NVFP4_COND_FILE = "diffusion_models/ideogram4_nvfp4_mixed.safetensors"
_IDEOGRAM4_NVFP4_UNCOND_FILE = (
    "diffusion_models/ideogram4_unconditional_nvfp4_mixed.safetensors"
)


@dataclass(frozen=True)
class Ideogram4Nvfp4ModelResolution:
    base_model_name: str
    base_model_path: str
    transformer_weights_path: str
    unconditional_transformer_weights_path: str | None


@lru_cache(maxsize=1)
def _resolve_ideogram4_base_model_path() -> str:
    return maybe_download_model(_IDEOGRAM4_BASE_MODEL, force_diffusers_model=True)


def _resolve_ideogram4_unconditional_transformer_weights_path(
    transformer_weights_path: str,
) -> str | None:
    if os.path.basename(transformer_weights_path) != os.path.basename(
        _IDEOGRAM4_NVFP4_COND_FILE
    ):
        return None
    return os.path.join(
        os.path.dirname(transformer_weights_path),
        os.path.basename(_IDEOGRAM4_NVFP4_UNCOND_FILE),
    )


def _resolve_ideogram4_nvfp4_transformer_weights_paths(
    server_args: ServerArgs, model_path: str
) -> tuple[str, str | None]:
    if server_args.transformer_weights_path is not None:
        transformer_weights_path = server_args.transformer_weights_path
        return (
            transformer_weights_path,
            _resolve_ideogram4_unconditional_transformer_weights_path(
                transformer_weights_path
            ),
        )

    local_nvfp4_path = maybe_download_model(
        model_path,
        allow_patterns=[
            _IDEOGRAM4_NVFP4_COND_FILE,
            _IDEOGRAM4_NVFP4_UNCOND_FILE,
        ],
    )
    return (
        os.path.join(local_nvfp4_path, _IDEOGRAM4_NVFP4_COND_FILE),
        os.path.join(local_nvfp4_path, _IDEOGRAM4_NVFP4_UNCOND_FILE),
    )


def resolve_ideogram4_nvfp4_model(
    server_args: ServerArgs, model_path: str
) -> Ideogram4Nvfp4ModelResolution:
    (
        transformer_weights_path,
        unconditional_transformer_weights_path,
    ) = _resolve_ideogram4_nvfp4_transformer_weights_paths(
        server_args,
        model_path,
    )
    return Ideogram4Nvfp4ModelResolution(
        base_model_name=_IDEOGRAM4_BASE_MODEL,
        base_model_path=_resolve_ideogram4_base_model_path(),
        transformer_weights_path=transformer_weights_path,
        unconditional_transformer_weights_path=unconditional_transformer_weights_path,
    )


class Ideogram4Pipeline(LoRAPipeline, ComposedPipelineBase):
    pipeline_name = "Ideogram4Pipeline"

    _required_config_modules = [
        "text_encoder",
        "tokenizer",
        "vae",
        "transformer",
        "unconditional_transformer",
        "scheduler",
    ]

    def _create_denoising_stage(self):
        transformer = self.get_module("transformer")
        unconditional_transformer = self.get_module("unconditional_transformer")
        return ProgressiveDenoisingStageRouter(
            standard_stage=Ideogram4DenoisingStage(
                transformer=transformer,
                unconditional_transformer=unconditional_transformer,
                pipeline=self,
            ),
            progressive_stage_factory=lambda: Ideogram4ProgressiveDenoisingStage(
                transformer=transformer,
                unconditional_transformer=unconditional_transformer,
                pipeline=self,
            ),
        )

    def create_pipeline_stages(self, server_args: ServerArgs):
        self.add_stage(InputValidationStage())
        self.add_stage_factory(
            RoleType.ENCODER,
            lambda: Ideogram4TextEncodingStage(
                text_encoder=self.get_module("text_encoder"),
                tokenizer=self.get_module("tokenizer"),
            ),
            "ideogram4_text_encoding_stage",
        )
        self.add_standard_latent_preparation_stage()
        self.add_stage_factory(
            RoleType.DENOISER,
            self._create_denoising_stage,
            "ideogram4_denoising_stage",
        )
        self.add_stage_factory(
            RoleType.DECODER,
            lambda: Ideogram4DecodingStage(vae=self.get_module("vae")),
            "ideogram4_decoding_stage",
        )


class Ideogram4Nvfp4Pipeline(Ideogram4Pipeline):
    pipeline_name = "Ideogram4Nvfp4Pipeline"
    _model_resolution: Ideogram4Nvfp4ModelResolution | None = None

    def _get_model_resolution(
        self,
        server_args: ServerArgs | None = None,
    ) -> Ideogram4Nvfp4ModelResolution:
        if self._model_resolution is None:
            if server_args is None:
                raise ValueError(
                    "server_args is required to resolve Ideogram4 NVFP4 paths"
                )
            self._model_resolution = resolve_ideogram4_nvfp4_model(
                server_args,
                self.model_path,
            )
        return self._model_resolution

    def _load_config(self) -> dict[str, Any]:
        model_resolution = self._get_model_resolution(self.server_args)
        logger.info("Model path: %s", self.model_path)
        logger.info(
            "Using base model '%s' at %s for config and non-transformer components",
            model_resolution.base_model_name,
            model_resolution.base_model_path,
        )
        config = verify_model_config_and_directory(model_resolution.base_model_path)
        return cast(dict[str, Any], config)

    def _resolve_component_path(
        self,
        server_args: ServerArgs,
        module_name: str,
        load_module_name: str,
    ) -> str:
        override_path = server_args.component_paths.get(module_name)
        if override_path is not None:
            return maybe_download_model(override_path)

        component_model_path = os.path.join(
            self._get_model_resolution(server_args).base_model_path,
            load_module_name,
        )
        logger.debug("Resolved component path: %s", component_model_path)
        return component_model_path

    def load_modules(
        self,
        server_args: ServerArgs,
        loaded_modules: dict | None = None,
    ) -> dict:
        model_resolution = self._get_model_resolution(server_args)
        server_args.transformer_weights_path = model_resolution.transformer_weights_path
        if model_resolution.unconditional_transformer_weights_path is not None:
            # The loader treats transformer_weights_path as the base DiT override.
            # Route the sibling unconditional DiT weights through the generic
            # per-component override map instead of hard-coding Ideogram there.
            component_transformer_weights_paths = dict(
                getattr(server_args, "component_transformer_weights_paths", {})
            )
            component_transformer_weights_paths.setdefault(
                "unconditional_transformer",
                model_resolution.unconditional_transformer_weights_path,
            )
            server_args.component_transformer_weights_paths = (
                component_transformer_weights_paths
            )
        logger.info(
            "NVFP4 transformer weights: %s",
            model_resolution.transformer_weights_path,
        )
        logger.info(
            "NVFP4 unconditional transformer weights: %s",
            server_args.component_transformer_weights_paths.get(
                "unconditional_transformer"
            ),
        )
        return super().load_modules(server_args, loaded_modules)


EntryClass = [Ideogram4Pipeline, Ideogram4Nvfp4Pipeline]