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

# SPDX-License-Identifier: Apache-2.0
"""ErnieImage text-to-image pipeline."""

import json
import os

from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
    ComposedPipelineBase,
)
from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline
from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import (
    InputValidationStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ernie_image_pe import (
    PromptEnhancementStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import (
    TextEncodingStage,
)
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
    maybe_download_model,
    maybe_download_model_index,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger

logger = init_logger(__name__)


class ErnieImagePipeline(LoRAPipeline, ComposedPipelineBase):

    pipeline_name = "ErnieImagePipeline"

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

    def _has_pe_in_model_index(self, server_args) -> bool:
        try:
            model_index = maybe_download_model_index(server_args.model_path)
            return "pe" in model_index and model_index["pe"] is not None
        except Exception:
            return False

    def _read_tokenizer_model_max_length(self, model_path: str):
        """Read model_max_length from tokenizer/tokenizer_config.json.

        Supports both local paths and HuggingFace Hub model IDs.
        Returns None if the value cannot be determined.
        """
        tokenizer_config_subpath = os.path.join("tokenizer", "tokenizer_config.json")

        # Local path
        if os.path.exists(model_path):
            config_path = os.path.join(model_path, tokenizer_config_subpath)
            if os.path.exists(config_path):
                with open(config_path, encoding="utf-8") as f:
                    config = json.load(f)
                return config.get("model_max_length")
            return None

        # Remote HuggingFace Hub model ID
        try:
            import tempfile

            from huggingface_hub import hf_hub_download

            with tempfile.TemporaryDirectory() as tmp_dir:
                config_path = hf_hub_download(
                    repo_id=model_path,
                    filename=tokenizer_config_subpath,
                    local_dir=tmp_dir,
                )
                with open(config_path, encoding="utf-8") as f:
                    config = json.load(f)
                return config.get("model_max_length")
        except Exception as e:
            logger.warning(
                "Failed to read tokenizer_config.json from %s: %s", model_path, e
            )
            return None

    def _resolve_pe_tokenizer_path(self, model_path: str, server_args) -> str:
        """Resolve the directory that contains the PE tokenizer files."""
        pe_component_path = server_args.component_paths.get(
            "pe", os.path.join(model_path, "pe")
        )
        if os.path.exists(os.path.join(pe_component_path, "tokenizer_config.json")):
            return pe_component_path
        pe_tokenizer_dir = os.path.join(model_path, "pe_tokenizer")
        if os.path.exists(os.path.join(pe_tokenizer_dir, "tokenizer_config.json")):
            return pe_tokenizer_dir
        return pe_component_path

    def _read_pe_model_max_length(self, model_path: str, server_args) -> int | None:
        # If model_path is a Hub ID, download the full model first (or use cache)
        # so that pe/tokenizer_config.json is available locally.
        if not os.path.exists(model_path):
            try:
                model_path = maybe_download_model(
                    model_path, force_diffusers_model=True
                )
            except Exception as e:
                logger.warning(
                    "Failed to download model to read pe/tokenizer_config.json: %s", e
                )
                return None

        tokenizer_path = self._resolve_pe_tokenizer_path(model_path, server_args)
        config_path = os.path.join(tokenizer_path, "tokenizer_config.json")
        if os.path.exists(config_path):
            try:
                with open(config_path, encoding="utf-8") as f:
                    config = json.load(f)
                val = config.get("model_max_length")
                if val is not None:
                    return int(val)
            except Exception as e:
                logger.warning(
                    "Failed to read tokenizer_config.json from %s: %s",
                    tokenizer_path,
                    e,
                )
        return None

    def load_modules(self, server_args, loaded_modules=None):
        has_pe = self._has_pe_in_model_index(server_args)
        if has_pe:
            if "pe" not in self._required_config_modules:
                self._required_config_modules.insert(0, "pe")
            logger.info("PE model detected in model_index.json, will load PE module.")

        pipeline_config = server_args.pipeline_config

        # --- Text encoder max_length ---
        text_model_max_length = self._read_tokenizer_model_max_length(
            server_args.model_path
        )
        if text_model_max_length is not None:
            # 1. Update arch_config.text_len so the model knows the true sequence length
            if (
                hasattr(pipeline_config, "text_encoder_configs")
                and pipeline_config.text_encoder_configs
            ):
                arch_config = pipeline_config.text_encoder_configs[0].arch_config
                arch_config.text_len = text_model_max_length
                arch_config.tokenizer_kwargs["max_length"] = text_model_max_length
            # 2. Update text_encoder_extra_args used by TextEncodingStage tokenization
            if (
                hasattr(pipeline_config, "text_encoder_extra_args")
                and pipeline_config.text_encoder_extra_args
            ):
                pipeline_config.text_encoder_extra_args[0][
                    "max_length"
                ] = text_model_max_length
            logger.info(
                "Set text encoder model_max_length=%d from tokenizer/tokenizer_config.json",
                text_model_max_length,
            )
        else:
            logger.warning(
                "Could not read model_max_length from tokenizer/tokenizer_config.json, "
                "text encoder will use the default text_len from arch config."
            )

        # --- PE model_max_length ---
        if has_pe:
            pe_model_max_length = self._read_pe_model_max_length(
                server_args.model_path, server_args
            )
            if pe_model_max_length is not None:
                pipeline_config.pe_model_max_length = pe_model_max_length
                logger.info(
                    "Set PE model_max_length=%d from pe/tokenizer_config.json",
                    pe_model_max_length,
                )
            else:
                raise RuntimeError(
                    "PE model is present but 'model_max_length' could not be read from "
                    "pe/tokenizer_config.json. Please ensure the PE component directory "
                    "contains a valid tokenizer_config.json with a 'model_max_length' field."
                )

        return super().load_modules(server_args, loaded_modules)

    def create_pipeline_stages(self, server_args):
        self.add_stage(InputValidationStage())

        pe_model = self.get_module("pe")
        if pe_model is not None:
            pe_tokenizer = getattr(pe_model, "pe_tokenizer", None)
            if pe_tokenizer is None:
                from transformers import AutoTokenizer

                pe_tokenizer_path = self._resolve_pe_tokenizer_path(
                    self.model_path, server_args
                )
                logger.warning(
                    "pe_tokenizer not found on pe_model (%s), loading from %s",
                    type(pe_model).__name__,
                    pe_tokenizer_path,
                )
                pe_tokenizer = AutoTokenizer.from_pretrained(
                    pe_tokenizer_path,
                    trust_remote_code=server_args.trust_remote_code,
                )
            self.add_stage(
                PromptEnhancementStage(
                    pe_model=pe_model,
                    pe_tokenizer=pe_tokenizer,
                ),
                "prompt_enhancement_stage",
            )

        self.add_stage(
            TextEncodingStage(
                text_encoders=[self.get_module("text_encoder")],
                tokenizers=[self.get_module("tokenizer")],
            ),
            "prompt_encoding_stage_primary",
        )

        self.add_standard_timestep_preparation_stage()
        self.add_standard_latent_preparation_stage()
        self.add_standard_denoising_stage()
        self.add_standard_decoding_stage()


EntryClass = ErnieImagePipeline