[Diffusion][SenseNova] support SenseNova-U1.5-8B-MoT (#36606)
Co-authored-by: wuyuefeng <wuyuefeng@noreply.gitcode.com> Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
Yuefeng Wu
co-authored by
wuyuefeng
Xiaoyu Zhang
ronnie_zheng
wuyuefeng
Xiaoyu Zhang
ronnie_zheng
parent
78da625190
commit
daf66f6670
@@ -53,6 +53,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConf
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.sana_video import (
|
||||
SanaVideoPipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.sensenova_u1 import (
|
||||
SenseNovaU1PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.stablediffusion3 import (
|
||||
StableDiffusion3PipelineConfig,
|
||||
)
|
||||
@@ -83,6 +86,7 @@ __all__ = [
|
||||
"PipelineConfig",
|
||||
"SanaPipelineConfig",
|
||||
"SanaVideoPipelineConfig",
|
||||
"SenseNovaU1PipelineConfig",
|
||||
"SlidingTileAttnConfig",
|
||||
"MOVAPipelineConfig",
|
||||
"Pi05PipelineConfig",
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from dataclasses import dataclass
|
||||
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.base import (
|
||||
ModelTaskType,
|
||||
PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
|
||||
ModelDeploymentConfig,
|
||||
)
|
||||
|
||||
|
||||
def _is_runtime_option_requested(value) -> bool:
|
||||
if value is None:
|
||||
return False
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, (dict, list, tuple, set)):
|
||||
return bool(value)
|
||||
return True
|
||||
|
||||
|
||||
def _is_arg_explicitly_set(server_args, option: str) -> bool:
|
||||
is_explicit = getattr(server_args, "is_arg_explicitly_set", None)
|
||||
if callable(is_explicit):
|
||||
return is_explicit(option)
|
||||
return _is_runtime_option_requested(getattr(server_args, option, None))
|
||||
|
||||
|
||||
def _component_residency_requests_offload(value) -> bool:
|
||||
if not _is_runtime_option_requested(value):
|
||||
return False
|
||||
if isinstance(value, dict):
|
||||
values = value.values()
|
||||
elif isinstance(value, str):
|
||||
values = value.split(",")
|
||||
else:
|
||||
values = value
|
||||
|
||||
for raw_value in values:
|
||||
mode = str(raw_value).split("=", 1)[-1].strip().replace("_", "-").lower()
|
||||
if mode in ("component-offload", "layerwise-offload"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _set_compatible_runtime_defaults(server_args) -> None:
|
||||
compatible_defaults = {
|
||||
"component_residency": None,
|
||||
"cpu_offload_components": None,
|
||||
"dit_cpu_offload": False,
|
||||
"text_encoder_cpu_offload": False,
|
||||
"image_encoder_cpu_offload": False,
|
||||
"vae_cpu_offload": False,
|
||||
"dit_layerwise_offload": False,
|
||||
"layerwise_offload_components": None,
|
||||
"quantization": None,
|
||||
"quantization_ignored_layers": None,
|
||||
"transformer_weights_path": None,
|
||||
"component_paths": {},
|
||||
"component_weights_paths": {},
|
||||
"component_quantizations": {},
|
||||
"component_quantization_ignored_layers": {},
|
||||
"component_precisions": {},
|
||||
"attention_backend": None,
|
||||
"component_attention_backends": {},
|
||||
"attention_backend_config": None,
|
||||
}
|
||||
for option, value in compatible_defaults.items():
|
||||
if not _is_arg_explicitly_set(server_args, option):
|
||||
setattr(server_args, option, value)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SenseNovaU1PipelineConfig(PipelineConfig):
|
||||
"""Native SenseNova-U1 text-to-image pipeline configuration."""
|
||||
|
||||
task_type: ModelTaskType = ModelTaskType.T2I
|
||||
model_precision: str = "bf16"
|
||||
should_use_guidance: bool = True
|
||||
supports_cfg_parallel: bool = False
|
||||
|
||||
def supports_dynamic_batching(self):
|
||||
return False
|
||||
|
||||
def supports_disaggregation(self) -> bool:
|
||||
return False
|
||||
|
||||
def supports_sequential_multi_output_inference(self):
|
||||
return True
|
||||
|
||||
def validate_server_args(self, server_args) -> None:
|
||||
if server_args.num_gpus != 1:
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline currently supports num_gpus=1. "
|
||||
"Native tensor/pipeline parallelism is not implemented yet."
|
||||
)
|
||||
if getattr(server_args, "enable_torch_compile", False):
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline does not support torch.compile yet. "
|
||||
"Please omit --enable-torch-compile."
|
||||
)
|
||||
if getattr(server_args, "lora_path", None):
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline does not support LoRA adapters yet. "
|
||||
"Please omit --lora-path."
|
||||
)
|
||||
_set_compatible_runtime_defaults(server_args)
|
||||
if _is_arg_explicitly_set(
|
||||
server_args, "component_residency"
|
||||
) and _component_residency_requests_offload(
|
||||
getattr(server_args, "component_residency", None)
|
||||
):
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline does not support component residency "
|
||||
"offload modes yet. Please omit --component-residency."
|
||||
)
|
||||
unsupported_runtime_options = {
|
||||
"cpu_offload_components": "CPU offload",
|
||||
"dit_cpu_offload": "DiT CPU offload",
|
||||
"text_encoder_cpu_offload": "text encoder CPU offload",
|
||||
"image_encoder_cpu_offload": "image encoder CPU offload",
|
||||
"vae_cpu_offload": "VAE CPU offload",
|
||||
"dit_layerwise_offload": "DiT layerwise offload",
|
||||
"layerwise_offload_components": "layerwise offload",
|
||||
"quantization": "quantization",
|
||||
"quantization_ignored_layers": "quantization ignored layers",
|
||||
"transformer_weights_path": "pre-quantized transformer weights",
|
||||
"component_paths": "component path overrides",
|
||||
"component_weights_paths": "component weight path overrides",
|
||||
"component_quantizations": "component quantization",
|
||||
"component_quantization_ignored_layers": (
|
||||
"component quantization ignored layers"
|
||||
),
|
||||
"component_precisions": "component precision overrides",
|
||||
}
|
||||
for option, description in unsupported_runtime_options.items():
|
||||
if _is_arg_explicitly_set(
|
||||
server_args, option
|
||||
) and _is_runtime_option_requested(getattr(server_args, option, None)):
|
||||
raise ValueError(
|
||||
f"SenseNovaU1Pipeline does not support {description} yet. "
|
||||
f"Please omit --{option.replace('_', '-')}."
|
||||
)
|
||||
if _is_arg_explicitly_set(
|
||||
server_args, "attention_backend"
|
||||
) and _is_runtime_option_requested(
|
||||
getattr(server_args, "attention_backend", None)
|
||||
):
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline does not support custom attention backends yet. "
|
||||
"Please omit --attention-backend."
|
||||
)
|
||||
if _is_arg_explicitly_set(
|
||||
server_args, "component_attention_backends"
|
||||
) and _is_runtime_option_requested(
|
||||
getattr(server_args, "component_attention_backends", None)
|
||||
):
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline does not support component attention backends yet. "
|
||||
"Please omit --component-attention-backends."
|
||||
)
|
||||
if _is_arg_explicitly_set(
|
||||
server_args, "attention_backend_config"
|
||||
) and _is_runtime_option_requested(
|
||||
getattr(server_args, "attention_backend_config", None)
|
||||
):
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline does not support attention backend config yet. "
|
||||
"Please omit --attention-backend-config."
|
||||
)
|
||||
|
||||
def get_model_deployment_config(self) -> ModelDeploymentConfig:
|
||||
return ModelDeploymentConfig(
|
||||
speed_mode_enable_torch_compile_by_default=False,
|
||||
keep_resident_min_available_gb=80,
|
||||
auto_enable_cfg_parallel=False,
|
||||
supports_cfg_parallel=False,
|
||||
)
|
||||
@@ -10,6 +10,9 @@ from sglang.multimodal_gen.configs.sample.lingbot_video_moe import (
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.pi05 import Pi05SamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.sensenova_u1 import (
|
||||
SenseNovaU1SamplingParams,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"SamplingParams",
|
||||
@@ -18,4 +21,5 @@ __all__ = [
|
||||
"Ideogram4SamplingParams",
|
||||
"Pi05SamplingParams",
|
||||
"LingBotVideoMoESamplingParams",
|
||||
"SenseNovaU1SamplingParams",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from sglang.multimodal_gen.configs.sample.sampling_params import (
|
||||
DataType,
|
||||
SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sensenova_u1 import (
|
||||
DEFAULT_CFG_INTERVAL,
|
||||
DEFAULT_CFG_NORM,
|
||||
DEFAULT_ENABLE_TIMESTEP_SHIFT,
|
||||
DEFAULT_T_EPS,
|
||||
DEFAULT_THINK_MODE,
|
||||
DEFAULT_TIMESTEP_SHIFT,
|
||||
SENSENOVA_U1_CFG_NORM_CHOICES,
|
||||
SENSENOVA_U1_REQUEST_EXTRA_KEY,
|
||||
SENSENOVA_U1_RESOLUTION_ALIGNMENT,
|
||||
)
|
||||
|
||||
_PUBLIC_OVERRIDE_FIELDS = {
|
||||
"prompt",
|
||||
"prompt_path",
|
||||
"height",
|
||||
"width",
|
||||
"num_inference_steps",
|
||||
"guidance_scale",
|
||||
"num_outputs_per_prompt",
|
||||
"seed",
|
||||
"save_output",
|
||||
"output_path",
|
||||
"output_file_name",
|
||||
"output_quality",
|
||||
"output_compression",
|
||||
"quality",
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class SenseNovaU1SamplingParams(SamplingParams):
|
||||
data_type: DataType = field(default=DataType.IMAGE, init=False)
|
||||
height: int = 2048
|
||||
width: int = 2048
|
||||
num_frames: int = 1
|
||||
fps: int = 1
|
||||
num_inference_steps: int = 50
|
||||
guidance_scale: float = 4.0
|
||||
cfg_norm: str = DEFAULT_CFG_NORM
|
||||
timestep_shift: float = DEFAULT_TIMESTEP_SHIFT
|
||||
enable_timestep_shift: bool = DEFAULT_ENABLE_TIMESTEP_SHIFT
|
||||
cfg_interval: tuple[float, float] = DEFAULT_CFG_INTERVAL
|
||||
t_eps: float = DEFAULT_T_EPS
|
||||
think_mode: bool = DEFAULT_THINK_MODE
|
||||
negative_prompt: None = field(default=None, init=False)
|
||||
|
||||
@classmethod
|
||||
def supported_override_fields(cls) -> set[str]:
|
||||
return set(_PUBLIC_OVERRIDE_FIELDS)
|
||||
|
||||
@classmethod
|
||||
def get_cli_args(cls, args):
|
||||
cli_args = super().get_cli_args(args)
|
||||
return {
|
||||
key: value
|
||||
for key, value in cli_args.items()
|
||||
if key in _PUBLIC_OVERRIDE_FIELDS
|
||||
}
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if isinstance(self.cfg_interval, list):
|
||||
self.cfg_interval = tuple(float(x) for x in self.cfg_interval)
|
||||
super().__post_init__()
|
||||
|
||||
def _validate(self) -> None:
|
||||
super()._validate()
|
||||
if (
|
||||
self.width % SENSENOVA_U1_RESOLUTION_ALIGNMENT != 0
|
||||
or self.height % SENSENOVA_U1_RESOLUTION_ALIGNMENT != 0
|
||||
):
|
||||
raise ValueError(
|
||||
"SenseNova-U1 requires width and height to be divisible by "
|
||||
f"{SENSENOVA_U1_RESOLUTION_ALIGNMENT}, got "
|
||||
f"{self.width}x{self.height}."
|
||||
)
|
||||
if self.num_frames != 1:
|
||||
raise ValueError(
|
||||
f"SenseNova-U1 is an image model and requires num_frames=1, got {self.num_frames}."
|
||||
)
|
||||
if self.cfg_norm not in SENSENOVA_U1_CFG_NORM_CHOICES:
|
||||
raise ValueError(
|
||||
f"cfg_norm must be one of {SENSENOVA_U1_CFG_NORM_CHOICES}, "
|
||||
f"got {self.cfg_norm!r}"
|
||||
)
|
||||
if len(self.cfg_interval) != 2:
|
||||
raise ValueError("cfg_interval must contain exactly two values")
|
||||
start, end = self.cfg_interval
|
||||
if not 0.0 <= float(start) <= float(end) <= 1.0:
|
||||
raise ValueError(
|
||||
f"cfg_interval must satisfy 0 <= start <= end <= 1, got {self.cfg_interval!r}"
|
||||
)
|
||||
|
||||
def build_request_extra(self) -> dict[str, Any]:
|
||||
extra = super().build_request_extra()
|
||||
extra[SENSENOVA_U1_REQUEST_EXTRA_KEY] = {
|
||||
"cfg_norm": self.cfg_norm,
|
||||
"timestep_shift": self.timestep_shift,
|
||||
"enable_timestep_shift": self.enable_timestep_shift,
|
||||
"cfg_interval": tuple(self.cfg_interval),
|
||||
"t_eps": self.t_eps,
|
||||
"think_mode": self.think_mode,
|
||||
}
|
||||
return extra
|
||||
@@ -0,0 +1,55 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Shared constants for the native SenseNova-U1 integration."""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
SENSENOVA_U1_REQUEST_EXTRA_KEY = "sensenova_u1"
|
||||
|
||||
SENSENOVA_U1_MODEL_IDS = {
|
||||
"sensenova/sensenova-u1.5-8b-mot",
|
||||
}
|
||||
SENSENOVA_U1_ADAPTER_ONLY_MODEL_IDS = {
|
||||
"sensenova/sensenova-u1.5-8b-mot-loras",
|
||||
}
|
||||
|
||||
SENSENOVA_U1_CFG_NORM_CHOICES = (
|
||||
"none",
|
||||
"global",
|
||||
"channel",
|
||||
"cfg_zero_star",
|
||||
)
|
||||
SENSENOVA_U1_RESOLUTION_ALIGNMENT = 32
|
||||
|
||||
DEFAULT_CFG_NORM = "none"
|
||||
DEFAULT_TIMESTEP_SHIFT = 3.0
|
||||
DEFAULT_ENABLE_TIMESTEP_SHIFT = True
|
||||
DEFAULT_CFG_INTERVAL = (0.0, 1.0)
|
||||
DEFAULT_T_EPS = 0.02
|
||||
DEFAULT_THINK_MODE = False
|
||||
|
||||
|
||||
def is_sensenova_u1_model(model_path: str) -> bool:
|
||||
"""Identify SenseNova-U1 Hub IDs and local base checkpoints."""
|
||||
if os.path.isdir(model_path):
|
||||
config_path = os.path.join(model_path, "config.json")
|
||||
try:
|
||||
with open(config_path) as config_file:
|
||||
config = json.load(config_file)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return False
|
||||
|
||||
if not isinstance(config, dict):
|
||||
return False
|
||||
architectures = config.get("architectures", [])
|
||||
return (
|
||||
config.get("model_type") == "neo_chat"
|
||||
and isinstance(architectures, list)
|
||||
and "NEOChatModel" in architectures
|
||||
)
|
||||
|
||||
return model_path.rstrip("/").lower() in SENSENOVA_U1_MODEL_IDS
|
||||
|
||||
|
||||
def is_sensenova_u1_adapter_only_model(model_path: str) -> bool:
|
||||
return model_path.rstrip("/").lower() in SENSENOVA_U1_ADAPTER_ONLY_MODEL_IDS
|
||||
@@ -101,6 +101,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.sana_video import (
|
||||
SanaVideoPipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import SanaWMPipelineConfig
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.sensenova_u1 import (
|
||||
SenseNovaU1PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.stablediffusion3 import (
|
||||
StableDiffusion3PipelineConfig,
|
||||
)
|
||||
@@ -181,6 +184,9 @@ from sglang.multimodal_gen.configs.sample.qwenimage import (
|
||||
from sglang.multimodal_gen.configs.sample.sana import SanaSamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.sana_video import SanaVideoSamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.sana_wm import SanaWMSamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.sensenova_u1 import (
|
||||
SenseNovaU1SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.stablediffusion3 import (
|
||||
StableDiffusion3SamplingParams,
|
||||
)
|
||||
@@ -200,6 +206,11 @@ from sglang.multimodal_gen.configs.sample.zimage import (
|
||||
ZImageSamplingParams,
|
||||
ZImageTurboSamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sensenova_u1 import (
|
||||
SENSENOVA_U1_MODEL_IDS,
|
||||
is_sensenova_u1_adapter_only_model,
|
||||
is_sensenova_u1_model,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
@@ -458,17 +469,24 @@ def has_registered_diffusion_model_path(model_path: str) -> bool:
|
||||
_ensure_registry_initialized()
|
||||
all_model_hf_paths = sorted(_MODEL_HF_PATH_TO_NAME.keys(), key=len, reverse=True)
|
||||
|
||||
if is_sensenova_u1_model(model_path):
|
||||
return True
|
||||
|
||||
if model_path in _MODEL_HF_PATH_TO_NAME:
|
||||
return True
|
||||
|
||||
model_short_name = get_model_short_name(model_path.lower())
|
||||
for registered_model_hf_id in all_model_hf_paths:
|
||||
if registered_model_hf_id.lower() in SENSENOVA_U1_MODEL_IDS:
|
||||
continue
|
||||
registered_model_name = get_model_short_name(registered_model_hf_id.lower())
|
||||
if registered_model_name in model_short_name:
|
||||
return True
|
||||
|
||||
normalized_model_path = _normalize_hf_cache_path(model_path)
|
||||
for registered_model_hf_id in all_model_hf_paths:
|
||||
if registered_model_hf_id.lower() in SENSENOVA_U1_MODEL_IDS:
|
||||
continue
|
||||
cache_repo_fragment = (
|
||||
f"models--{registered_model_hf_id.lower().replace('/', '--')}"
|
||||
)
|
||||
@@ -502,6 +520,13 @@ def _get_config_info(
|
||||
"falling back to automatic detection."
|
||||
)
|
||||
|
||||
# SenseNova Hub IDs require an exact match, while local checkpoints are
|
||||
# identified from their config metadata rather than their directory name.
|
||||
if is_sensenova_u1_model(model_path):
|
||||
for registered_hf_id in all_model_hf_paths:
|
||||
if registered_hf_id.lower() in SENSENOVA_U1_MODEL_IDS:
|
||||
return _CONFIG_REGISTRY.get(_MODEL_HF_PATH_TO_NAME[registered_hf_id])
|
||||
|
||||
# 1. Exact match
|
||||
if model_path in _MODEL_HF_PATH_TO_NAME:
|
||||
model_id = _MODEL_HF_PATH_TO_NAME[model_path]
|
||||
@@ -511,6 +536,8 @@ def _get_config_info(
|
||||
# 2. Partial match: find the best (longest) match against all registered model hf paths.
|
||||
model_short_name = get_model_short_name(model_path.lower())
|
||||
for registered_model_hf_id in all_model_hf_paths:
|
||||
if registered_model_hf_id.lower() in SENSENOVA_U1_MODEL_IDS:
|
||||
continue
|
||||
registered_model_name = get_model_short_name(registered_model_hf_id.lower())
|
||||
|
||||
if registered_model_name in model_short_name:
|
||||
@@ -529,6 +556,8 @@ def _get_config_info(
|
||||
# -> models--black-forest-labs--flux.2-dev-nvfp4 (to match with cache_repo_fragment)
|
||||
normalized_model_path = _normalize_hf_cache_path(model_path)
|
||||
for registered_model_hf_id in all_model_hf_paths:
|
||||
if registered_model_hf_id.lower() in SENSENOVA_U1_MODEL_IDS:
|
||||
continue
|
||||
cache_repo_fragment = (
|
||||
f"models--{registered_model_hf_id.lower().replace('/', '--')}"
|
||||
)
|
||||
@@ -661,6 +690,16 @@ def get_model_info(
|
||||
elif isinstance(backend, str):
|
||||
backend = Backend.from_string(backend)
|
||||
|
||||
if is_sensenova_u1_adapter_only_model(model_path):
|
||||
logger.error(
|
||||
"SenseNova-U1 adapter-only checkpoint '%s' does not contain base "
|
||||
"model weights or config. SenseNova-U1 adapters are not supported "
|
||||
"yet; use the base checkpoint 'sensenova/SenseNova-U1.5-8B-MoT' "
|
||||
"directly.",
|
||||
model_path,
|
||||
)
|
||||
return None
|
||||
|
||||
# Handle explicit diffusers backend
|
||||
if backend == Backend.DIFFUSERS:
|
||||
logger.info(
|
||||
@@ -980,6 +1019,13 @@ def _register_configs():
|
||||
)
|
||||
],
|
||||
)
|
||||
register_configs(
|
||||
sampling_param_cls=SenseNovaU1SamplingParams,
|
||||
pipeline_config_cls=SenseNovaU1PipelineConfig,
|
||||
hf_model_paths=[
|
||||
"sensenova/SenseNova-U1.5-8B-MoT",
|
||||
],
|
||||
)
|
||||
register_configs(
|
||||
sampling_param_cls=FastH3SamplingParams,
|
||||
pipeline_config_cls=FastH3PipelineConfig,
|
||||
@@ -1392,6 +1438,9 @@ def is_known_non_diffusers_multimodal_model(model_path: str) -> bool:
|
||||
|
||||
def get_non_diffusers_pipeline_name(model_path: str) -> Optional[str]:
|
||||
"""Get the pipeline name for a known non-diffusers model."""
|
||||
if is_sensenova_u1_model(model_path):
|
||||
return "SenseNovaU1Pipeline"
|
||||
|
||||
normalized_model_path = _normalize_hf_cache_path(model_path)
|
||||
model_short_name = get_model_short_name(normalized_model_path)
|
||||
for pattern, pipeline_name in KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS.items():
|
||||
|
||||
@@ -163,9 +163,12 @@ def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None
|
||||
# respect config file by overriding args with args parsed from it
|
||||
if config_file:
|
||||
config_args = ServerArgs.load_config_file(config_file) or {}
|
||||
sampling_param_fields = {
|
||||
field.name for field in dataclasses.fields(sampling_params_cls)
|
||||
}
|
||||
if hasattr(sampling_params_cls, "supported_override_fields"):
|
||||
sampling_param_fields = sampling_params_cls.supported_override_fields()
|
||||
else:
|
||||
sampling_param_fields = {
|
||||
field.name for field in dataclasses.fields(sampling_params_cls)
|
||||
}
|
||||
sampling_params_kwargs.update(
|
||||
{
|
||||
key: value
|
||||
|
||||
@@ -18,7 +18,6 @@ import tempfile
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from copy import copy
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, List, Optional, Sequence, Union
|
||||
|
||||
@@ -41,6 +40,12 @@ from sglang.multimodal_gen.configs.sample.sampling_params import (
|
||||
DataType,
|
||||
SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.request_utils import (
|
||||
expand_request_outputs as expand_request_outputs,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.request_utils import (
|
||||
normalize_output_seeds as normalize_output_seeds,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_logger
|
||||
@@ -214,129 +219,6 @@ class MaterializedOutput:
|
||||
fps: int = 0
|
||||
|
||||
|
||||
def normalize_output_seeds(
|
||||
seed: int | list[int],
|
||||
*,
|
||||
num_outputs_per_prompt: int,
|
||||
num_prompts: int = 1,
|
||||
prompt_index: int = 0,
|
||||
) -> list[int]:
|
||||
"""
|
||||
return a list of seed with size equal to `num_outputs_per_prompt`
|
||||
"""
|
||||
if num_outputs_per_prompt <= 0:
|
||||
raise ValueError(
|
||||
f"num_outputs_per_prompt must be positive, got {num_outputs_per_prompt}"
|
||||
)
|
||||
|
||||
if isinstance(seed, list):
|
||||
seeds = [int(item) for item in seed]
|
||||
total_outputs = num_outputs_per_prompt * num_prompts
|
||||
if len(seeds) == num_outputs_per_prompt:
|
||||
return seeds
|
||||
if len(seeds) == total_outputs:
|
||||
start = prompt_index * num_outputs_per_prompt
|
||||
return seeds[start : start + num_outputs_per_prompt]
|
||||
raise ValueError(
|
||||
"seed list length must match num_outputs_per_prompt "
|
||||
f"({num_outputs_per_prompt}) or total outputs ({total_outputs}), "
|
||||
f"got {len(seeds)}"
|
||||
)
|
||||
|
||||
base_seed = int(seed)
|
||||
return [base_seed + i for i in range(num_outputs_per_prompt)]
|
||||
|
||||
|
||||
def _with_output_index_suffix(output_file_name: str, output_index: int) -> str:
|
||||
base, ext = os.path.splitext(output_file_name)
|
||||
return f"{base}_{output_index}{ext}"
|
||||
|
||||
|
||||
def _copy_trace_ctx_for_output(req: Req, request_id: str | None, output_index: int):
|
||||
trace_ctx = req.trace_ctx
|
||||
if output_index == 0 or not trace_ctx.tracing_enable:
|
||||
return trace_ctx
|
||||
|
||||
output_trace_ctx = TraceReqContext(
|
||||
rid=request_id,
|
||||
module_name=trace_ctx.module_name,
|
||||
external_trace_header=trace_ctx.external_trace_header,
|
||||
)
|
||||
output_trace_ctx.trace_req_start()
|
||||
return output_trace_ctx
|
||||
|
||||
|
||||
def _copy_req_for_output(
|
||||
req: Req,
|
||||
*,
|
||||
request_id: str | None,
|
||||
output_index: int,
|
||||
) -> Req:
|
||||
"""Create a lightweight per-output ``Req`` without deep-copying tensors."""
|
||||
output_req = copy(req)
|
||||
output_req.sampling_params = copy(req.sampling_params)
|
||||
output_req.extra = dict(req.extra)
|
||||
output_req.condition_inputs = dict(req.condition_inputs)
|
||||
output_req.trace_ctx = _copy_trace_ctx_for_output(req, request_id, output_index)
|
||||
return output_req
|
||||
|
||||
|
||||
def expand_request_outputs(
|
||||
req: Req,
|
||||
*,
|
||||
num_prompts: int = 1,
|
||||
prompt_index: int = 0,
|
||||
) -> list[Req]:
|
||||
"""
|
||||
Expand a req to a list with size equal to `num_prompts`
|
||||
"""
|
||||
num_outputs = int(req.num_outputs_per_prompt)
|
||||
# each req must has different seed
|
||||
seeds = normalize_output_seeds(
|
||||
req.seed,
|
||||
num_outputs_per_prompt=num_outputs,
|
||||
num_prompts=num_prompts,
|
||||
prompt_index=prompt_index,
|
||||
)
|
||||
|
||||
if num_outputs == 1:
|
||||
req.seed = seeds[0]
|
||||
req.seeds = None
|
||||
req.generator = None
|
||||
req.sampling_params.refresh_request_extra_after_output_expansion(req)
|
||||
return [req]
|
||||
|
||||
expanded: list[Req] = []
|
||||
for output_index, seed in enumerate(seeds):
|
||||
output_request_id = (
|
||||
f"{req.request_id}:{output_index}" if req.request_id is not None else None
|
||||
)
|
||||
output_req = _copy_req_for_output(
|
||||
req, request_id=output_request_id, output_index=output_index
|
||||
)
|
||||
output_req.seed = seed
|
||||
output_req.num_outputs_per_prompt = 1
|
||||
output_req.seeds = None
|
||||
output_req.generator = None
|
||||
output_req.extra["parent_request_id"] = req.request_id
|
||||
output_req.extra["output_index"] = output_index
|
||||
|
||||
if output_request_id is not None:
|
||||
output_req.request_id = output_request_id
|
||||
|
||||
if req.output_file_name:
|
||||
output_req.output_file_name = _with_output_index_suffix(
|
||||
req.output_file_name, output_index
|
||||
)
|
||||
output_req.sampling_params.refresh_request_extra_after_output_expansion(
|
||||
output_req
|
||||
)
|
||||
output_req.validate()
|
||||
expanded.append(output_req)
|
||||
|
||||
return expanded
|
||||
|
||||
|
||||
def _normalize_audio_to_numpy(audio: Any) -> np.ndarray | None:
|
||||
"""Convert audio (torch / numpy) into a float32 numpy array in [-1, 1], best-effort."""
|
||||
if audio is None:
|
||||
|
||||
@@ -1371,13 +1371,15 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
"""
|
||||
merge batched output
|
||||
"""
|
||||
if parts.output_file_paths:
|
||||
merged.output_file_paths = parts.output_file_paths
|
||||
if any(metrics is not None for metrics in parts.metrics_list):
|
||||
merged.metrics_list = parts.metrics_list
|
||||
merged.metrics = next(
|
||||
metrics for metrics in parts.metrics_list if metrics is not None
|
||||
)
|
||||
if merged.error is not None:
|
||||
return
|
||||
if parts.output_file_paths:
|
||||
merged.output_file_paths = parts.output_file_paths
|
||||
if parts.tensor_outputs:
|
||||
merged.output = torch.cat(parts.tensor_outputs, dim=0)
|
||||
elif parts.list_outputs:
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""SenseNova-U1 native model registration for multimodal generation."""
|
||||
|
||||
from sglang.multimodal_gen.runtime.models.sensenova_u1.neo_unify import ( # noqa: F401
|
||||
NEOChatConfig,
|
||||
NEOChatModel,
|
||||
NEOLLMConfig,
|
||||
NEOMoELLMConfig,
|
||||
NEOVisionConfig,
|
||||
NEOVisionModel,
|
||||
register,
|
||||
)
|
||||
|
||||
register()
|
||||
|
||||
__all__ = [
|
||||
"NEOChatConfig",
|
||||
"NEOChatModel",
|
||||
"NEOLLMConfig",
|
||||
"NEOMoELLMConfig",
|
||||
"NEOVisionConfig",
|
||||
"NEOVisionModel",
|
||||
"register",
|
||||
]
|
||||
@@ -0,0 +1,42 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.models import ( # noqa: F401
|
||||
sensenova_u1 as _sensenova_u1,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
|
||||
|
||||
|
||||
def load_model_and_tokenizer(
|
||||
model_path: str,
|
||||
server_args: ServerArgs,
|
||||
) -> dict[str, Any]:
|
||||
dtype = PRECISION_TO_TYPE.get(
|
||||
server_args.pipeline_config.model_precision, torch.bfloat16
|
||||
)
|
||||
model_kwargs: dict[str, Any] = {"torch_dtype": dtype}
|
||||
if server_args.trust_remote_code:
|
||||
model_kwargs["trust_remote_code"] = True
|
||||
if server_args.revision is not None:
|
||||
model_kwargs["revision"] = server_args.revision
|
||||
|
||||
tokenizer_kwargs: dict[str, Any] = {}
|
||||
if server_args.trust_remote_code:
|
||||
tokenizer_kwargs["trust_remote_code"] = True
|
||||
if server_args.revision is not None:
|
||||
tokenizer_kwargs["revision"] = server_args.revision
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, **tokenizer_kwargs)
|
||||
model = AutoModel.from_pretrained(model_path, **model_kwargs).eval()
|
||||
device = get_local_torch_device()
|
||||
current_platform.set_device(device)
|
||||
model = model.to(device)
|
||||
return {"model": model, "tokenizer": tokenizer}
|
||||
@@ -0,0 +1,8 @@
|
||||
# SenseNova-U1 Vendored Model Files
|
||||
|
||||
The Python files in this directory are adapted from:
|
||||
|
||||
- Repository: https://github.com/OpenSenseNova/SenseNova-U1
|
||||
- Commit: 2f42002f9b819506c9deb44599f0809e30252aba
|
||||
|
||||
They were modified for SGLang multimodal_gen native integration.
|
||||
@@ -0,0 +1,57 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from .configuration_neo_chat import NEOChatConfig, NEOLLMConfig, NEOMoELLMConfig
|
||||
from .configuration_neo_vit import NEOVisionConfig
|
||||
from .modeling_neo_chat import NEOChatModel
|
||||
from .modeling_neo_vit import NEOVisionModel
|
||||
from .modeling_qwen3 import _HAS_FLASH_ATTN as has_flash_attn
|
||||
from .modeling_qwen3 import (
|
||||
Qwen3ForCausalLM,
|
||||
effective_attn_backend,
|
||||
get_attn_backend,
|
||||
set_attn_backend,
|
||||
)
|
||||
from .modeling_qwen3_moe import Qwen3MoeForCausalLM
|
||||
|
||||
__all__ = [
|
||||
"NEOChatConfig",
|
||||
"NEOLLMConfig",
|
||||
"NEOMoELLMConfig",
|
||||
"NEOVisionConfig",
|
||||
"NEOChatModel",
|
||||
"NEOVisionModel",
|
||||
"Qwen3ForCausalLM",
|
||||
"Qwen3MoeForCausalLM",
|
||||
"register",
|
||||
"set_attn_backend",
|
||||
"get_attn_backend",
|
||||
"effective_attn_backend",
|
||||
"has_flash_attn",
|
||||
]
|
||||
|
||||
|
||||
_REGISTERED = False
|
||||
|
||||
|
||||
def register() -> None:
|
||||
"""Register NEO-Unify types with ``transformers.Auto*``.
|
||||
|
||||
After calling this (or simply ``import sensenova_u1``), users can load a
|
||||
SenseNova-U1 checkpoint via plain ``AutoConfig.from_pretrained`` /
|
||||
``AutoModel.from_pretrained``.
|
||||
"""
|
||||
global _REGISTERED
|
||||
if _REGISTERED:
|
||||
return
|
||||
|
||||
from transformers import AutoConfig, AutoModel
|
||||
|
||||
AutoConfig.register("neo_vision", NEOVisionConfig, exist_ok=True)
|
||||
AutoConfig.register("neo_chat", NEOChatConfig, exist_ok=True)
|
||||
|
||||
AutoModel.register(NEOVisionConfig, NEOVisionModel, exist_ok=True)
|
||||
AutoModel.register(NEOChatConfig, NEOChatModel, exist_ok=True)
|
||||
|
||||
_REGISTERED = True
|
||||
+217
@@ -0,0 +1,217 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
import copy
|
||||
|
||||
from transformers import Qwen3Config, Qwen3MoeConfig
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.utils import logging
|
||||
|
||||
from .configuration_neo_vit import NEOVisionConfig
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
def _restore_legacy_rope_theta(config) -> None:
|
||||
"""Expose the v4 rope attribute expected by the vendored model code."""
|
||||
if hasattr(config, "rope_theta"):
|
||||
return
|
||||
rope_parameters = getattr(config, "rope_parameters", None) or {}
|
||||
config.rope_theta = float(rope_parameters.get("rope_theta", 10000.0))
|
||||
|
||||
|
||||
class NEOLLMConfig(Qwen3Config):
|
||||
"""Config for the dense Qwen3 backbone used by NEO-Unify.
|
||||
|
||||
Extends ``Qwen3Config`` with two extra rope knobs used by the spatial
|
||||
(height/width) rotary axes that are layered on top of the temporal one.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, rope_theta_hw=10000.0, max_position_embeddings_hw=10000, **kwargs
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
_restore_legacy_rope_theta(self)
|
||||
self.rope_theta_hw = rope_theta_hw
|
||||
self.max_position_embeddings_hw = max_position_embeddings_hw
|
||||
|
||||
|
||||
class NEOMoELLMConfig(Qwen3MoeConfig):
|
||||
"""Config for the Qwen3-MoE backbone used by NEO-Unify.
|
||||
|
||||
Extends ``Qwen3MoeConfig`` with the same ``rope_theta_hw`` /
|
||||
``max_position_embeddings_hw`` extras as :class:`NEOLLMConfig`, and adds a
|
||||
*generation-path* MoE branch alongside the standard understanding-path one.
|
||||
In the A3B unified model every decoder layer carries two parallel sparse
|
||||
MoE blocks routed by the per-token ``image_gen_indicators`` mask:
|
||||
|
||||
* ``mlp`` - sparse MoE for the understanding path
|
||||
(``num_experts`` experts, ``num_experts_per_tok`` active,
|
||||
expert width ``moe_intermediate_size``).
|
||||
* ``mlp_mot_gen`` - sparse MoE for the image generation path
|
||||
(``gen_num_experts`` experts, ``gen_num_experts_per_tok``
|
||||
active, expert width ``gen_moe_intermediate_size``).
|
||||
|
||||
Each gen-path knob falls back to its understanding-path counterpart when
|
||||
unset, so vanilla single-MoE configs keep working without changes.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rope_theta_hw=10000.0,
|
||||
max_position_embeddings_hw=10000,
|
||||
gen_num_experts=None,
|
||||
gen_num_experts_per_tok=None,
|
||||
gen_moe_intermediate_size=None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
_restore_legacy_rope_theta(self)
|
||||
self.rope_theta_hw = rope_theta_hw
|
||||
self.max_position_embeddings_hw = max_position_embeddings_hw
|
||||
|
||||
# Generation-path MoE knobs default to the understanding-path values
|
||||
# so legacy single-MoE configs (where both branches share the same
|
||||
# router width / expert count) keep working unchanged.
|
||||
self.gen_num_experts = (
|
||||
int(gen_num_experts)
|
||||
if gen_num_experts is not None
|
||||
else int(self.num_experts)
|
||||
)
|
||||
self.gen_num_experts_per_tok = (
|
||||
int(gen_num_experts_per_tok)
|
||||
if gen_num_experts_per_tok is not None
|
||||
else int(self.num_experts_per_tok)
|
||||
)
|
||||
self.gen_moe_intermediate_size = (
|
||||
int(gen_moe_intermediate_size)
|
||||
if gen_moe_intermediate_size is not None
|
||||
else int(self.moe_intermediate_size)
|
||||
)
|
||||
|
||||
# ``Qwen3Attention`` (used by NEO-Unify MoE layers) reads
|
||||
# ``config.layer_types[layer_idx]`` to decide between ``"full_attention"``
|
||||
# and ``"sliding_attention"``. Older / vanilla ``Qwen3MoeConfig`` does
|
||||
# not populate that field, so we backfill it here mirroring the dense
|
||||
# ``Qwen3Config`` behaviour: sliding-attention layers start at
|
||||
# ``max_window_layers`` when ``use_sliding_window`` is enabled.
|
||||
existing = getattr(self, "layer_types", None)
|
||||
if not existing or len(existing) != self.num_hidden_layers:
|
||||
use_swa = (
|
||||
bool(getattr(self, "use_sliding_window", False))
|
||||
and getattr(self, "sliding_window", None) is not None
|
||||
)
|
||||
max_window_layers = int(getattr(self, "max_window_layers", 0) or 0)
|
||||
self.layer_types = [
|
||||
(
|
||||
"sliding_attention"
|
||||
if (use_swa and i >= max_window_layers)
|
||||
else "full_attention"
|
||||
)
|
||||
for i in range(self.num_hidden_layers)
|
||||
]
|
||||
|
||||
|
||||
def _is_moe_llm_config(llm_config) -> bool:
|
||||
"""Detect whether an ``llm_config`` (dict or object) targets a MoE backbone.
|
||||
|
||||
Order of checks: explicit ``model_type``, ``architectures`` entry that
|
||||
contains ``MoE/MoeForCausalLM``, or presence of MoE-specific keys
|
||||
(``num_experts``).
|
||||
"""
|
||||
if isinstance(llm_config, dict):
|
||||
model_type = llm_config.get("model_type", "")
|
||||
archs = llm_config.get("architectures") or []
|
||||
has_num_experts = "num_experts" in llm_config
|
||||
else:
|
||||
model_type = getattr(llm_config, "model_type", "")
|
||||
archs = getattr(llm_config, "architectures", None) or []
|
||||
has_num_experts = hasattr(llm_config, "num_experts")
|
||||
|
||||
if isinstance(model_type, str) and "moe" in model_type.lower():
|
||||
return True
|
||||
for arch in archs:
|
||||
arch_str = str(arch)
|
||||
if "Moe" in arch_str or "MoE" in arch_str:
|
||||
return True
|
||||
return (
|
||||
bool(has_num_experts)
|
||||
and getattr(llm_config, "num_experts", 0)
|
||||
and int(getattr(llm_config, "num_experts", 0)) > 1
|
||||
)
|
||||
|
||||
|
||||
def _build_llm_config(llm_config):
|
||||
"""Instantiate the right LLM config object from a dict or pre-built config."""
|
||||
if isinstance(llm_config, dict):
|
||||
if _is_moe_llm_config(llm_config):
|
||||
return NEOMoELLMConfig(**llm_config)
|
||||
return NEOLLMConfig(**llm_config)
|
||||
return llm_config
|
||||
|
||||
|
||||
class NEOChatConfig(PretrainedConfig):
|
||||
model_type = "neo_chat"
|
||||
is_composition = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vision_config=None,
|
||||
llm_config=None,
|
||||
use_backbone_lora=0,
|
||||
use_llm_lora=0,
|
||||
downsample_ratio=0.5,
|
||||
template=None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
if vision_config is None:
|
||||
vision_config = {"architectures": ["NEOVisionModel"]}
|
||||
logger.info(
|
||||
"vision_config is None. Initializing the NEOVisionConfig with default values."
|
||||
)
|
||||
|
||||
if llm_config is None:
|
||||
llm_config = {"architectures": ["Qwen3ForCausalLM"]}
|
||||
logger.info(
|
||||
"llm_config is None. Initializing the LlamaConfig config with default values (`LlamaConfig`)."
|
||||
)
|
||||
assert "architectures" in llm_config, (
|
||||
"Should specify architecture in llm_config"
|
||||
)
|
||||
|
||||
if isinstance(vision_config, dict):
|
||||
self.vision_config = NEOVisionConfig(**vision_config)
|
||||
else:
|
||||
self.vision_config = vision_config
|
||||
|
||||
self.llm_config = _build_llm_config(llm_config)
|
||||
|
||||
self.use_backbone_lora = use_backbone_lora
|
||||
self.use_llm_lora = use_llm_lora
|
||||
self.downsample_ratio = downsample_ratio
|
||||
self.template = template
|
||||
self.tie_word_embeddings = self.llm_config.tie_word_embeddings
|
||||
|
||||
@property
|
||||
def is_moe_llm(self) -> bool:
|
||||
"""Convenience flag so callers can switch between dense / MoE LLM."""
|
||||
return isinstance(self.llm_config, NEOMoELLMConfig)
|
||||
|
||||
def to_dict(self):
|
||||
"""
|
||||
Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
|
||||
|
||||
Returns:
|
||||
`Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
|
||||
"""
|
||||
output = copy.deepcopy(self.__dict__)
|
||||
output["vision_config"] = self.vision_config.to_dict()
|
||||
output["llm_config"] = self.llm_config.to_dict()
|
||||
output["model_type"] = self.__class__.model_type
|
||||
output["use_backbone_lora"] = self.use_backbone_lora
|
||||
output["use_llm_lora"] = self.use_llm_lora
|
||||
output["downsample_ratio"] = self.downsample_ratio
|
||||
output["template"] = self.template
|
||||
|
||||
return output
|
||||
+70
@@ -0,0 +1,70 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, Union
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.utils import logging
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
def _as_singleton_tuple(value: Any) -> tuple[Any, ...]:
|
||||
if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
|
||||
if len(value) == 1:
|
||||
return _as_singleton_tuple(value[0])
|
||||
return tuple(value)
|
||||
return (value,)
|
||||
|
||||
|
||||
class NEOVisionConfig(PretrainedConfig):
|
||||
model_type = "neo_vision"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_channels=3,
|
||||
patch_size=16,
|
||||
hidden_size=1024,
|
||||
llm_hidden_size=2048,
|
||||
downsample_ratio=0.5,
|
||||
rope_theta_vision=10000.0,
|
||||
max_position_embeddings_vision=10000,
|
||||
min_pixels=65536,
|
||||
max_pixels=4194304,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.llm_hidden_size = _as_singleton_tuple(llm_hidden_size)
|
||||
self.downsample_ratio = _as_singleton_tuple(downsample_ratio)
|
||||
self.rope_theta_vision = rope_theta_vision
|
||||
self.max_position_embeddings_vision = max_position_embeddings_vision
|
||||
self.num_channels = num_channels
|
||||
self.patch_size = patch_size
|
||||
self.min_pixels = min_pixels
|
||||
self.max_pixels = max_pixels
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
|
||||
) -> "PretrainedConfig":
|
||||
config_dict, kwargs = cls.get_config_dict(
|
||||
pretrained_model_name_or_path, **kwargs
|
||||
)
|
||||
|
||||
if "vision_config" in config_dict:
|
||||
config_dict = config_dict["vision_config"]
|
||||
|
||||
if (
|
||||
"model_type" in config_dict
|
||||
and hasattr(cls, "model_type")
|
||||
and config_dict["model_type"] != cls.model_type
|
||||
):
|
||||
logger.warning(
|
||||
f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
|
||||
f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
|
||||
)
|
||||
|
||||
return cls.from_dict(config_dict, **kwargs)
|
||||
@@ -0,0 +1,420 @@
|
||||
"""
|
||||
Conversation prompt templates.
|
||||
|
||||
We kindly request that you import fastchat instead of copying this file if you wish to use it.
|
||||
If you have changes in mind, please contribute back so the community can benefit collectively and continue to maintain these valuable templates.
|
||||
|
||||
Modified from https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
|
||||
"""
|
||||
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
import dataclasses
|
||||
from enum import IntEnum, auto
|
||||
from typing import Dict, List, Tuple, Union
|
||||
|
||||
|
||||
class SeparatorStyle(IntEnum):
|
||||
"""Separator styles."""
|
||||
|
||||
ADD_COLON_SINGLE = auto()
|
||||
ADD_COLON_TWO = auto()
|
||||
ADD_COLON_SPACE_SINGLE = auto()
|
||||
NO_COLON_SINGLE = auto()
|
||||
NO_COLON_TWO = auto()
|
||||
ADD_NEW_LINE_SINGLE = auto()
|
||||
LLAMA2 = auto()
|
||||
CHATGLM = auto()
|
||||
CHATML = auto()
|
||||
CHATINTERN = auto()
|
||||
DOLLY = auto()
|
||||
RWKV = auto()
|
||||
PHOENIX = auto()
|
||||
ROBIN = auto()
|
||||
FALCON_CHAT = auto()
|
||||
CHATGLM3 = auto()
|
||||
INTERNVL_ZH = auto()
|
||||
MPT = auto()
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class Conversation:
|
||||
"""A class that manages prompt templates and keeps all conversation history."""
|
||||
|
||||
# The name of this template
|
||||
name: str
|
||||
# The template of the system prompt
|
||||
system_template: str = "{system_message}"
|
||||
# The system message
|
||||
system_message: str = ""
|
||||
# The names of two roles
|
||||
roles: Tuple[str] = ("USER", "ASSISTANT")
|
||||
# All messages. Each item is (role, message).
|
||||
messages: List[List[str]] = ()
|
||||
# The number of few shot examples
|
||||
offset: int = 0
|
||||
# The separator style and configurations
|
||||
sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SINGLE
|
||||
sep: str = "\n"
|
||||
sep2: str = None
|
||||
# Stop criteria (the default one is EOS token)
|
||||
stop_str: Union[str, List[str]] = None
|
||||
# Stops generation if meeting any token in this list
|
||||
stop_token_ids: List[int] = None
|
||||
|
||||
def get_prompt(self) -> str:
|
||||
"""Get the prompt for generation."""
|
||||
if self.system_message is not None and self.system_message != "":
|
||||
system_prompt = self.system_template.format(
|
||||
system_message=self.system_message
|
||||
)
|
||||
else:
|
||||
system_prompt = ""
|
||||
|
||||
if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + ": " + message + self.sep
|
||||
else:
|
||||
ret += role + ":"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
|
||||
seps = [self.sep, self.sep2]
|
||||
ret = "" if system_prompt == "" else system_prompt + seps[0]
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if message:
|
||||
ret += role + ": " + message + seps[i % 2]
|
||||
else:
|
||||
ret += role + ":"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + ": " + message + self.sep
|
||||
else:
|
||||
ret += role + ": " # must be end with a space
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.ADD_NEW_LINE_SINGLE:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + "\n" + message + self.sep
|
||||
else:
|
||||
ret += role + "\n"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
|
||||
ret = system_prompt
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + message + self.sep
|
||||
else:
|
||||
ret += role
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.NO_COLON_TWO:
|
||||
seps = [self.sep, self.sep2]
|
||||
ret = system_prompt
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if message:
|
||||
ret += role + message + seps[i % 2]
|
||||
else:
|
||||
ret += role
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.RWKV:
|
||||
ret = system_prompt
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if message:
|
||||
ret += (
|
||||
role
|
||||
+ ": "
|
||||
+ message.replace("\r\n", "\n").replace("\n\n", "\n")
|
||||
)
|
||||
ret += "\n\n"
|
||||
else:
|
||||
ret += role + ":"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.LLAMA2:
|
||||
seps = [self.sep, self.sep2]
|
||||
ret = system_prompt if system_prompt != "" else "[INST] "
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
tag = self.roles[i % 2]
|
||||
if message:
|
||||
if i == 0:
|
||||
ret += message + " "
|
||||
else:
|
||||
ret += tag + " " + message + seps[i % 2]
|
||||
else:
|
||||
ret += tag
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.CHATGLM:
|
||||
# source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
|
||||
# source2: https://huggingface.co/THUDM/chatglm2-6b/blob/e186c891cf64310ac66ef10a87e6635fa6c2a579/modeling_chatglm.py#L926
|
||||
round_add_n = 1 if self.name == "chatglm2" else 0
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if i % 2 == 0:
|
||||
ret += f"[Round {i // 2 + round_add_n}]{self.sep}"
|
||||
|
||||
if message:
|
||||
ret += f"{role}：{message}{self.sep}"
|
||||
else:
|
||||
ret += f"{role}："
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.CHATML:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep + "\n"
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + "\n" + message + self.sep + "\n"
|
||||
else:
|
||||
ret += role + "\n"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.CHATGLM3:
|
||||
ret = system_prompt
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + "\n" + " " + message
|
||||
else:
|
||||
ret += role
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.CHATINTERN:
|
||||
# source: https://huggingface.co/internlm/internlm-chat-7b-8k/blob/bd546fa984b4b0b86958f56bf37f94aa75ab8831/modeling_internlm.py#L771
|
||||
seps = [self.sep, self.sep2]
|
||||
ret = system_prompt
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
# if i % 2 == 0:
|
||||
# ret += "<s>"
|
||||
if message:
|
||||
ret += role + ":" + message + seps[i % 2] + "\n"
|
||||
else:
|
||||
ret += role + ":"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.DOLLY:
|
||||
seps = [self.sep, self.sep2]
|
||||
ret = system_prompt
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if message:
|
||||
ret += role + ":\n" + message + seps[i % 2]
|
||||
if i % 2 == 1:
|
||||
ret += "\n\n"
|
||||
else:
|
||||
ret += role + ":\n"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.PHOENIX:
|
||||
ret = system_prompt
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + ": " + "<s>" + message + "</s>"
|
||||
else:
|
||||
ret += role + ": " + "<s>"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.ROBIN:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + ":\n" + message + self.sep
|
||||
else:
|
||||
ret += role + ":\n"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.FALCON_CHAT:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + ": " + message + self.sep
|
||||
else:
|
||||
ret += role + ":"
|
||||
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.INTERNVL_ZH:
|
||||
seps = [self.sep, self.sep2]
|
||||
ret = "" if system_prompt == "" else self.system_message + seps[0]
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if message:
|
||||
ret += role + ": " + message + seps[i % 2]
|
||||
else:
|
||||
ret += role + ":"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.MPT:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if message:
|
||||
if type(message) is tuple:
|
||||
message, _, _ = message
|
||||
ret += role + message + self.sep
|
||||
else:
|
||||
if i != len(self.messages) and message is not None:
|
||||
ret += role + self.sep
|
||||
else:
|
||||
ret += role
|
||||
return ret
|
||||
else:
|
||||
raise ValueError(f"Invalid style: {self.sep_style}")
|
||||
|
||||
def set_system_message(self, system_message: str):
|
||||
"""Set the system message."""
|
||||
self.system_message = system_message
|
||||
|
||||
def append_message(self, role: str, message: str):
|
||||
"""Append a new message."""
|
||||
self.messages.append([role, message])
|
||||
|
||||
def update_last_message(self, message: str):
|
||||
"""Update the last output.
|
||||
|
||||
The last message is typically set to be None when constructing the prompt,
|
||||
so we need to update it in-place after getting the response from a model.
|
||||
"""
|
||||
self.messages[-1][1] = message
|
||||
|
||||
def to_gradio_chatbot(self):
|
||||
"""Convert the conversation to gradio chatbot format."""
|
||||
ret = []
|
||||
for i, (role, msg) in enumerate(self.messages[self.offset :]):
|
||||
if i % 2 == 0:
|
||||
ret.append([msg, None])
|
||||
else:
|
||||
ret[-1][-1] = msg
|
||||
return ret
|
||||
|
||||
def to_openai_api_messages(self):
|
||||
"""Convert the conversation to OpenAI chat completion format."""
|
||||
ret = [{"role": "system", "content": self.system_message}]
|
||||
|
||||
for i, (_, msg) in enumerate(self.messages[self.offset :]):
|
||||
if i % 2 == 0:
|
||||
ret.append({"role": "user", "content": msg})
|
||||
else:
|
||||
if msg is not None:
|
||||
ret.append({"role": "assistant", "content": msg})
|
||||
return ret
|
||||
|
||||
def copy(self):
|
||||
return Conversation(
|
||||
name=self.name,
|
||||
system_template=self.system_template,
|
||||
system_message=self.system_message,
|
||||
roles=self.roles,
|
||||
messages=[[x, y] for x, y in self.messages],
|
||||
offset=self.offset,
|
||||
sep_style=self.sep_style,
|
||||
sep=self.sep,
|
||||
sep2=self.sep2,
|
||||
stop_str=self.stop_str,
|
||||
stop_token_ids=self.stop_token_ids,
|
||||
)
|
||||
|
||||
def dict(self):
|
||||
return {
|
||||
"template_name": self.name,
|
||||
"system_message": self.system_message,
|
||||
"roles": self.roles,
|
||||
"messages": self.messages,
|
||||
"offset": self.offset,
|
||||
}
|
||||
|
||||
|
||||
# A global registry for all conversation templates
|
||||
conv_templates: Dict[str, Conversation] = {}
|
||||
|
||||
|
||||
def register_conv_template(template: Conversation, override: bool = False):
|
||||
"""Register a new conversation template."""
|
||||
if not override:
|
||||
assert template.name not in conv_templates, (
|
||||
f"{template.name} has been registered."
|
||||
)
|
||||
|
||||
conv_templates[template.name] = template
|
||||
|
||||
|
||||
def get_conv_template(name: str) -> Conversation:
|
||||
"""Get a conversation template."""
|
||||
return conv_templates[name].copy()
|
||||
|
||||
|
||||
# Both Hermes-2 and neo1_0-chat are chatml-format conversation templates. The difference
|
||||
# is that during training, the preprocessing function for the Hermes-2 template doesn't add
|
||||
# <s> at the beginning of the tokenized sequence, while the neo1_0-chat template does.
|
||||
# Therefore, they are completely equivalent during inference.
|
||||
|
||||
# These Unicode-escaped strings preserve the exact Chinese prompts from the
|
||||
# pinned upstream source while complying with the multimodal_gen source lint.
|
||||
# Upstream prompt retained for benchmark compatibility:
|
||||
_INTERNVL_SYSTEM_MESSAGE = (
|
||||
"\u4f60\u662f\u7531\u4e0a\u6d77\u4eba\u5de5\u667a\u80fd\u5b9e\u9a8c\u5ba4"
|
||||
"\u8054\u5408\u5546\u6c64\u79d1\u6280\u5f00\u53d1\u7684\u4e66\u751f\u591a"
|
||||
"\u6a21\u6001\u5927\u6a21\u578b\uff0c\u82f1\u6587\u540d\u53ebInternVL, "
|
||||
"\u662f\u4e00\u4e2a\u6709\u7528\u65e0\u5bb3\u7684\u4eba\u5de5\u667a\u80fd"
|
||||
"\u52a9\u624b\u3002"
|
||||
)
|
||||
# Newer upstream prompt, intentionally not used by the first three templates:
|
||||
_INTERNVL2_5_SYSTEM_MESSAGE = (
|
||||
"\u4f60\u662f\u4e66\u751f\xb7\u4e07\u8c61\uff0c\u82f1\u6587\u540d\u662f"
|
||||
"InternVL\uff0c\u662f\u7531\u4e0a\u6d77\u4eba\u5de5\u667a\u80fd\u5b9e\u9a8c"
|
||||
"\u5ba4\u3001\u6e05\u534e\u5927\u5b66\u53ca\u591a\u5bb6\u5408\u4f5c\u5355"
|
||||
"\u4f4d\u8054\u5408\u5f00\u53d1\u7684\u591a\u6a21\u6001\u5927\u8bed\u8a00"
|
||||
"\u6a21\u578b\u3002"
|
||||
)
|
||||
|
||||
register_conv_template(
|
||||
Conversation(
|
||||
name="Hermes-2",
|
||||
system_template="<|im_start|>system\n{system_message}",
|
||||
# note: The new system prompt was not used here to avoid changes in benchmark performance.
|
||||
# system_message=_INTERNVL2_5_SYSTEM_MESSAGE
|
||||
system_message=_INTERNVL_SYSTEM_MESSAGE,
|
||||
roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
|
||||
sep_style=SeparatorStyle.MPT,
|
||||
sep="<|im_end|>",
|
||||
stop_str="<|endoftext|>",
|
||||
)
|
||||
)
|
||||
|
||||
register_conv_template(
|
||||
Conversation(
|
||||
name="internlm2-chat",
|
||||
system_template="<|im_start|>system\n{system_message}",
|
||||
# note: The new system prompt was not used here to avoid changes in benchmark performance.
|
||||
# system_message=_INTERNVL2_5_SYSTEM_MESSAGE
|
||||
system_message=_INTERNVL_SYSTEM_MESSAGE,
|
||||
roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
|
||||
sep_style=SeparatorStyle.MPT,
|
||||
sep="<|im_end|>",
|
||||
)
|
||||
)
|
||||
|
||||
register_conv_template(
|
||||
Conversation(
|
||||
name="phi3-chat",
|
||||
system_template="<|system|>\n{system_message}",
|
||||
# note: The new system prompt was not used here to avoid changes in benchmark performance.
|
||||
# system_message=_INTERNVL2_5_SYSTEM_MESSAGE
|
||||
system_message=_INTERNVL_SYSTEM_MESSAGE,
|
||||
roles=("<|user|>\n", "<|assistant|>\n"),
|
||||
sep_style=SeparatorStyle.MPT,
|
||||
sep="<|end|>",
|
||||
)
|
||||
)
|
||||
|
||||
register_conv_template(
|
||||
Conversation(
|
||||
name="internvl2_5",
|
||||
system_template="<|im_start|>system\n{system_message}",
|
||||
system_message=_INTERNVL2_5_SYSTEM_MESSAGE,
|
||||
roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
|
||||
sep_style=SeparatorStyle.MPT,
|
||||
sep="<|im_end|>\n",
|
||||
)
|
||||
)
|
||||
|
||||
register_conv_template(
|
||||
Conversation(
|
||||
name="neo1_0",
|
||||
system_template="<|im_start|>system\n{system_message}",
|
||||
system_message="",
|
||||
roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
|
||||
sep_style=SeparatorStyle.MPT,
|
||||
sep="<|im_end|>\n",
|
||||
)
|
||||
)
|
||||
+648
@@ -0,0 +1,648 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
import logging
|
||||
import math
|
||||
from functools import lru_cache
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def modulate(x, shift, scale=None):
|
||||
if shift is None:
|
||||
return x * (1 + scale)
|
||||
return x * (1 + scale) + shift
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-5):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||
return output * self.weight
|
||||
|
||||
|
||||
class TimestepEmbedder(nn.Module):
|
||||
"""
|
||||
Embeds scalar timesteps into vector representations.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size, frequency_embedding_size=256):
|
||||
super().__init__()
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, hidden_size, bias=True),
|
||||
)
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
|
||||
@staticmethod
|
||||
def timestep_embedding(t: torch.Tensor, dim: int, max_period: float = 10000.0):
|
||||
"""
|
||||
Create sinusoidal timestep embeddings.
|
||||
:param t: a 1-D Tensor of N indices, one per batch element. These may be fractional.
|
||||
:param dim: the dimension of the output.
|
||||
:param max_period: controls the minimum frequency of the embeddings.
|
||||
:return: an (N, D) Tensor of positional embeddings.
|
||||
"""
|
||||
# https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
|
||||
half = dim // 2
|
||||
freqs = torch.exp(
|
||||
-math.log(max_period)
|
||||
* torch.arange(start=0, end=half, dtype=torch.float32)
|
||||
/ half
|
||||
).to(device=t.device)
|
||||
args = t[:, None].float() * freqs[None]
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
if dim % 2:
|
||||
embedding = torch.cat(
|
||||
[embedding, torch.zeros_like(embedding[:, :1])], dim=-1
|
||||
)
|
||||
return embedding
|
||||
|
||||
def forward(self, t):
|
||||
t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
|
||||
t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype))
|
||||
return t_emb
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(self, channels, mlp_ratio=1.0):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.intermediate_size = int(channels * mlp_ratio)
|
||||
|
||||
self.in_ln = nn.LayerNorm(self.channels, eps=1e-6)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(self.channels, self.intermediate_size),
|
||||
nn.SiLU(),
|
||||
nn.Linear(self.intermediate_size, self.channels),
|
||||
)
|
||||
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(), nn.Linear(channels, 3 * channels, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x, y):
|
||||
shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(y).chunk(3, dim=-1)
|
||||
h = modulate(self.in_ln(x), shift_mlp, scale_mlp)
|
||||
h = self.mlp(h)
|
||||
return x + gate_mlp * h
|
||||
|
||||
|
||||
# class FinalLayer(nn.Module):
|
||||
|
||||
# def __init__(self, model_channels, out_channels):
|
||||
# super().__init__()
|
||||
# self.norm_final = nn.LayerNorm(model_channels, elementwise_affine=False, eps=1e-6)
|
||||
# self.linear = nn.Linear(model_channels, out_channels, bias=True)
|
||||
# self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(model_channels, 2 * model_channels, bias=True))
|
||||
|
||||
# def forward(self, x, c):
|
||||
# shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
|
||||
# x = modulate(self.norm_final(x), shift, scale)
|
||||
# x = self.linear(x)
|
||||
# return x
|
||||
|
||||
# class SimpleMLPAdaLN(nn.Module):
|
||||
|
||||
# def __init__(self, input_dim, out_dim, dim=1536, layers=12, mlp_ratio=1.0):
|
||||
# super().__init__()
|
||||
# self.input_dim = input_dim
|
||||
# self.out_dim = out_dim
|
||||
# self.dim = dim
|
||||
# self.layers = layers
|
||||
# self.mlp_ratio = mlp_ratio
|
||||
|
||||
# self.time_embed = TimestepEmbedder(dim)
|
||||
# self.input_proj = nn.Linear(input_dim, dim)
|
||||
|
||||
# res_blocks = []
|
||||
# for _ in range(layers):
|
||||
# res_blocks.append(ResBlock(dim, mlp_ratio))
|
||||
# self.res_blocks = nn.ModuleList(res_blocks)
|
||||
|
||||
# self.final_layer = FinalLayer(dim, out_dim)
|
||||
|
||||
# self.grad_checkpointing = False
|
||||
|
||||
# self.initialize_weights()
|
||||
|
||||
# def initialize_weights(self):
|
||||
# def _basic_init(module):
|
||||
# if isinstance(module, nn.Linear):
|
||||
# torch.nn.init.xavier_uniform_(module.weight)
|
||||
# if module.bias is not None:
|
||||
# nn.init.constant_(module.bias, 0)
|
||||
|
||||
# self.apply(_basic_init)
|
||||
|
||||
# # Initialize timestep embedding MLP
|
||||
# nn.init.normal_(self.time_embed.mlp[0].weight, std=0.02)
|
||||
# nn.init.normal_(self.time_embed.mlp[2].weight, std=0.02)
|
||||
|
||||
# # Zero-out adaLN modulation layers
|
||||
# for block in self.res_blocks:
|
||||
# nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
|
||||
# nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
|
||||
|
||||
# # Zero-out output layers
|
||||
# nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
|
||||
# nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
|
||||
# nn.init.constant_(self.final_layer.linear.weight, 0)
|
||||
# nn.init.constant_(self.final_layer.linear.bias, 0)
|
||||
|
||||
# def forward(self, x, t):
|
||||
# """
|
||||
# x.shape = (bsz, input_dim)
|
||||
# t.shape = (bsz,)
|
||||
# """
|
||||
|
||||
# x = self.input_proj(x)
|
||||
# t = self.time_embed(t)
|
||||
|
||||
# y = t
|
||||
|
||||
# for block in self.res_blocks:
|
||||
# if self.grad_checkpointing and self.training:
|
||||
# x = checkpoint(block, x, y, use_reentrant=True)
|
||||
# else:
|
||||
# x = block(x, y)
|
||||
|
||||
# return self.final_layer(x, y)
|
||||
|
||||
|
||||
class FlowMatchingHead(nn.Module):
|
||||
def __init__(self, input_dim, out_dim, dim=1536, layers=12, mlp_ratio=1.0):
|
||||
super(FlowMatchingHead, self).__init__()
|
||||
self.net = SimpleMLPAdaLN(
|
||||
input_dim=input_dim,
|
||||
out_dim=out_dim,
|
||||
dim=dim,
|
||||
layers=layers,
|
||||
mlp_ratio=mlp_ratio,
|
||||
)
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return self.net.input_proj.weight.dtype
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
return self.net.input_proj.weight.device
|
||||
|
||||
def forward(self, x, t):
|
||||
x = self.net(x, t)
|
||||
return x
|
||||
|
||||
|
||||
def precompute_freqs_cis_2d(
|
||||
dim: int, height: int, width: int, theta: float = 10000.0, scale=16.0
|
||||
):
|
||||
# assert H * H == end
|
||||
# flat_patch_pos = torch.linspace(-1, 1, end) # N = end
|
||||
x_pos = torch.linspace(0, scale, width)
|
||||
y_pos = torch.linspace(0, scale, height)
|
||||
y_pos, x_pos = torch.meshgrid(y_pos, x_pos, indexing="ij")
|
||||
y_pos = y_pos.reshape(-1)
|
||||
x_pos = x_pos.reshape(-1)
|
||||
freqs = 1.0 / (
|
||||
theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)
|
||||
) # Hc/4
|
||||
x_freqs = torch.outer(x_pos, freqs).float() # N Hc/4
|
||||
y_freqs = torch.outer(y_pos, freqs).float() # N Hc/4
|
||||
x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs)
|
||||
y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs)
|
||||
freqs_cis = torch.cat(
|
||||
[x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1
|
||||
) # N,Hc/4,2
|
||||
freqs_cis = freqs_cis.reshape(height * width, -1)
|
||||
return freqs_cis
|
||||
|
||||
|
||||
class NerfEmbedder(nn.Module):
|
||||
def __init__(self, in_channels, hidden_size_input, max_freqs):
|
||||
super().__init__()
|
||||
self.max_freqs = max_freqs
|
||||
self.hidden_size_input = hidden_size_input
|
||||
self.embedder = nn.Sequential(
|
||||
nn.Linear(in_channels + max_freqs**2, hidden_size_input, bias=True),
|
||||
)
|
||||
|
||||
@lru_cache
|
||||
def fetch_pos(self, patch_size, device, dtype):
|
||||
pos = precompute_freqs_cis_2d(
|
||||
self.max_freqs**2 * 2, patch_size, patch_size
|
||||
).real
|
||||
pos = pos[None, :, :].to(device=device, dtype=dtype)
|
||||
return pos
|
||||
|
||||
def forward(self, inputs):
|
||||
B, P2, C = inputs.shape
|
||||
patch_size = int(P2**0.5)
|
||||
device = inputs.device
|
||||
dtype = inputs.dtype
|
||||
dct = self.fetch_pos(patch_size, device, dtype)
|
||||
dct = dct.repeat(B, 1, 1)
|
||||
inputs = torch.cat([inputs, dct], dim=-1)
|
||||
inputs = self.embedder(inputs)
|
||||
return inputs
|
||||
|
||||
|
||||
class SimpleMLPAdaLN(nn.Module):
|
||||
"""
|
||||
The MLP for Diffusion Loss.
|
||||
:param in_channels: channels in the input Tensor.
|
||||
:param model_channels: base channel count for the model.
|
||||
:param out_channels: channels in the output Tensor.
|
||||
:param z_channels: channels in the condition.
|
||||
:param num_res_blocks: number of residual blocks per downsample.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
model_channels,
|
||||
out_channels,
|
||||
z_channels,
|
||||
num_res_blocks,
|
||||
patch_size,
|
||||
grad_checkpointing=False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.model_channels = model_channels
|
||||
self.out_channels = out_channels
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.grad_checkpointing = grad_checkpointing
|
||||
self.patch_size = patch_size
|
||||
|
||||
self.cond_embed = nn.Linear(z_channels, patch_size**2 * model_channels)
|
||||
|
||||
self.input_proj = nn.Linear(in_channels, model_channels)
|
||||
|
||||
res_blocks = []
|
||||
for i in range(num_res_blocks):
|
||||
res_blocks.append(
|
||||
ResBlock(
|
||||
model_channels,
|
||||
)
|
||||
)
|
||||
|
||||
self.res_blocks = nn.ModuleList(res_blocks)
|
||||
self.final_layer = FinalLayer(model_channels, out_channels)
|
||||
|
||||
self.initialize_weights()
|
||||
|
||||
def initialize_weights(self):
|
||||
def _basic_init(module):
|
||||
if isinstance(module, nn.Linear):
|
||||
torch.nn.init.xavier_uniform_(module.weight)
|
||||
if module.bias is not None:
|
||||
nn.init.constant_(module.bias, 0)
|
||||
|
||||
self.apply(_basic_init)
|
||||
|
||||
# Zero-out adaLN modulation layers
|
||||
for block in self.res_blocks:
|
||||
nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
|
||||
nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
|
||||
|
||||
# Zero-out output layers
|
||||
nn.init.constant_(self.final_layer.linear.weight, 0)
|
||||
nn.init.constant_(self.final_layer.linear.bias, 0)
|
||||
|
||||
def forward(self, x, c):
|
||||
"""
|
||||
Apply the model to an input batch.
|
||||
:param x: an [N x C] Tensor of inputs.
|
||||
:param t: a 1-D batch of timesteps.
|
||||
:param c: conditioning from AR transformer.
|
||||
:return: an [N x C] Tensor of outputs.
|
||||
"""
|
||||
x = self.input_proj(x)
|
||||
c = self.cond_embed(c)
|
||||
|
||||
y = c.reshape(-1, self.patch_size**2, self.model_channels)
|
||||
|
||||
for block in self.res_blocks:
|
||||
x = block(x, y)
|
||||
|
||||
return self.final_layer(x)
|
||||
|
||||
|
||||
class FinalLayer(nn.Module):
|
||||
"""
|
||||
The final layer adopted from DiT.
|
||||
"""
|
||||
|
||||
def __init__(self, model_channels, out_channels):
|
||||
super().__init__()
|
||||
self.norm_final = nn.LayerNorm(
|
||||
model_channels, elementwise_affine=False, eps=1e-6
|
||||
)
|
||||
self.linear = nn.Linear(model_channels, out_channels, bias=True)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.norm_final(x)
|
||||
x = self.linear(x)
|
||||
return x
|
||||
|
||||
|
||||
#################################################################################
|
||||
# Sine/Cosine Positional Embedding Functions #
|
||||
#################################################################################
|
||||
# https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
|
||||
|
||||
|
||||
def get_2d_sincos_pos_embed(
|
||||
embed_dim, grid_size, cls_token=False, extra_tokens=0, pe_interpolation=1.0
|
||||
):
|
||||
"""
|
||||
grid_size: int of the grid height and width
|
||||
return:
|
||||
pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
|
||||
"""
|
||||
grid_h = np.arange(grid_size, dtype=np.float32) / pe_interpolation
|
||||
grid_w = np.arange(grid_size, dtype=np.float32) / pe_interpolation
|
||||
grid = np.meshgrid(grid_w, grid_h) # here w goes first
|
||||
grid = np.stack(grid, axis=0)
|
||||
|
||||
grid = grid.reshape([2, 1, grid_size, grid_size])
|
||||
pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
|
||||
if cls_token and extra_tokens > 0:
|
||||
pos_embed = np.concatenate(
|
||||
[np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0
|
||||
)
|
||||
return pos_embed
|
||||
|
||||
|
||||
def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
|
||||
assert embed_dim % 2 == 0
|
||||
|
||||
# use half of dimensions to encode grid_h
|
||||
emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)
|
||||
emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)
|
||||
|
||||
emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
|
||||
return emb
|
||||
|
||||
|
||||
def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
|
||||
"""
|
||||
embed_dim: output dimension for each position
|
||||
pos: a list of positions to be encoded: size (M,)
|
||||
out: (M, D)
|
||||
"""
|
||||
assert embed_dim % 2 == 0
|
||||
omega = np.arange(embed_dim // 2, dtype=np.float64)
|
||||
omega /= embed_dim / 2.0
|
||||
omega = 1.0 / 10000**omega # (D/2,)
|
||||
|
||||
pos = pos.reshape(-1) # (M,)
|
||||
out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product
|
||||
|
||||
emb_sin = np.sin(out) # (M, D/2)
|
||||
emb_cos = np.cos(out) # (M, D/2)
|
||||
|
||||
emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)
|
||||
return emb
|
||||
|
||||
|
||||
# --------------------------------------------------------
|
||||
# Interpolate position embeddings for high-resolution
|
||||
# References:
|
||||
# DeiT: https://github.com/facebookresearch/deit
|
||||
# --------------------------------------------------------
|
||||
def interpolate_pos_embed(
|
||||
model_path, pe_key: str = "gen_pos_embed", new_len: int = 4096
|
||||
):
|
||||
state_dict = torch.load(model_path, map_location="cpu")
|
||||
|
||||
pos_embed_1d = state_dict[pe_key]
|
||||
_, ori_len, embed_dim = pos_embed_1d.shape
|
||||
|
||||
ori_size = int(ori_len**0.5)
|
||||
new_size = int(new_len**0.5)
|
||||
|
||||
if ori_size != new_size:
|
||||
logger.info(
|
||||
"Position interpolate from %dx%d to %dx%d"
|
||||
% (ori_size, ori_size, new_size, new_size)
|
||||
)
|
||||
pos_embed_2d = pos_embed_1d.reshape(-1, ori_size, ori_size, embed_dim).permute(
|
||||
0, 3, 1, 2
|
||||
)
|
||||
pos_embed_2d = torch.nn.functional.interpolate(
|
||||
pos_embed_2d, size=(new_size, new_size), mode="bicubic", align_corners=False
|
||||
)
|
||||
pos_embed_1d = pos_embed_2d.permute(0, 2, 3, 1).flatten(1, 2)
|
||||
state_dict[pe_key] = pos_embed_1d
|
||||
|
||||
torch.save(state_dict, model_path)
|
||||
|
||||
|
||||
class PositionEmbedding(nn.Module):
|
||||
def __init__(self, max_num_patch_per_side, hidden_size):
|
||||
super().__init__()
|
||||
self.max_num_patch_per_side = max_num_patch_per_side
|
||||
self.hidden_size = hidden_size
|
||||
self.pos_embed = nn.Parameter(
|
||||
torch.zeros(max_num_patch_per_side**2, hidden_size), requires_grad=False
|
||||
)
|
||||
self._init_weights()
|
||||
|
||||
def _init_weights(self):
|
||||
# Initialize (and freeze) pos_embed by sin-cos embedding:
|
||||
pos_embed = get_2d_sincos_pos_embed(
|
||||
self.hidden_size, self.max_num_patch_per_side
|
||||
)
|
||||
self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float())
|
||||
|
||||
def forward(self, position_ids):
|
||||
return self.pos_embed[position_ids]
|
||||
|
||||
|
||||
class ResidualConvBlock(nn.Module):
|
||||
def __init__(self, channels: int):
|
||||
super().__init__()
|
||||
self.block = nn.Sequential(
|
||||
nn.Conv2d(channels, channels, kernel_size=3, padding=1),
|
||||
nn.SiLU(),
|
||||
nn.Conv2d(channels, channels, kernel_size=3, padding=1),
|
||||
)
|
||||
nn.init.zeros_(self.block[2].weight)
|
||||
nn.init.zeros_(self.block[2].bias)
|
||||
|
||||
def forward(self, x):
|
||||
return x + self.block(x)
|
||||
|
||||
|
||||
class PostConvSmoother(nn.Module):
|
||||
def __init__(self, in_channels=3, hidden_channels=64, num_blocks=3):
|
||||
super().__init__()
|
||||
self.in_proj = nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1)
|
||||
self.blocks = nn.Sequential(
|
||||
*[ResidualConvBlock(hidden_channels) for _ in range(num_blocks)]
|
||||
)
|
||||
self.out_proj = nn.Conv2d(hidden_channels, in_channels, kernel_size=1)
|
||||
|
||||
nn.init.zeros_(self.out_proj.weight)
|
||||
nn.init.zeros_(self.out_proj.bias)
|
||||
|
||||
def forward(self, x):
|
||||
h = self.in_proj(x)
|
||||
h = self.blocks(h)
|
||||
return x + self.out_proj(h)
|
||||
|
||||
|
||||
class ProgressiveConvDecoder(nn.Module):
|
||||
def __init__(self, hidden_dim=4096, out_channels=3):
|
||||
super().__init__()
|
||||
|
||||
# self.proj = nn.Linear(hidden_dim, 1024)
|
||||
# self.act = nn.SiLU()
|
||||
|
||||
self.up_blocks = nn.ModuleList(
|
||||
[
|
||||
nn.Sequential(
|
||||
nn.Upsample(scale_factor=2, mode="nearest"),
|
||||
nn.Conv2d(hidden_dim, 512, kernel_size=3, padding=1),
|
||||
nn.GroupNorm(32, 512),
|
||||
nn.SiLU(),
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.Upsample(scale_factor=2, mode="nearest"),
|
||||
nn.Conv2d(512, 256, kernel_size=3, padding=1),
|
||||
nn.GroupNorm(32, 256),
|
||||
nn.SiLU(),
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.Upsample(scale_factor=2, mode="nearest"),
|
||||
nn.Conv2d(256, 64, kernel_size=3, padding=1),
|
||||
nn.GroupNorm(32, 64),
|
||||
nn.SiLU(),
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.Upsample(scale_factor=2, mode="nearest"),
|
||||
nn.Conv2d(64, 32, kernel_size=3, padding=1),
|
||||
nn.GroupNorm(16, 32),
|
||||
nn.SiLU(),
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.Upsample(scale_factor=2, mode="nearest"),
|
||||
nn.Conv2d(32, 16, kernel_size=3, padding=1),
|
||||
nn.SiLU(),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
self.out_conv = nn.Conv2d(16, out_channels, kernel_size=3, padding=1)
|
||||
|
||||
def forward(self, x_2d):
|
||||
# B, C, H, W = x_2d.shape
|
||||
# x = x_2d.permute(0, 2, 3, 1).contiguous() # (B, H, W, C)
|
||||
# x = self.proj(x)
|
||||
# x = self.act(x)
|
||||
# x = x.permute(0, 3, 1, 2).contiguous() # (B, 512, H, W)
|
||||
x = x_2d
|
||||
for block in self.up_blocks:
|
||||
x = block(x)
|
||||
|
||||
out = self.out_conv(x)
|
||||
return out
|
||||
|
||||
|
||||
class PatchDecoder_postps(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# layer 1: H/32 -> H/8 (4x upscale)
|
||||
|
||||
self.conv1 = nn.Conv2d(4096, 4096, kernel_size=3, padding=1)
|
||||
self.ps1 = nn.PixelShuffle(4)
|
||||
self.act1 = nn.GELU()
|
||||
|
||||
# layer 2: H/8 -> H (8x upscale)
|
||||
self.conv2 = nn.Conv2d(256, 192, kernel_size=3, padding=1)
|
||||
self.ps2 = nn.PixelShuffle(8)
|
||||
|
||||
def forward(self, x):
|
||||
# x shape: [B, 4096, H/32, W/32]
|
||||
x = self.ps1(self.act1(self.conv1(x))) # -> [B, 256, H/8, W/8]
|
||||
x = self.ps2(self.conv2(x)) # -> [B, 3, H, W]
|
||||
return x
|
||||
|
||||
|
||||
class PatchDecoder_preps(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# layer 1: H/32 -> H/16 (2x upscale)
|
||||
self.ps1 = nn.PixelShuffle(2)
|
||||
self.conv1 = nn.Conv2d(1024, 1024, kernel_size=3, padding=1)
|
||||
self.act1 = nn.GELU()
|
||||
|
||||
# layer 2: H/16 -> H/8 (2x upscale)
|
||||
self.ps2 = nn.PixelShuffle(2)
|
||||
self.conv2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
|
||||
self.act2 = nn.GELU()
|
||||
|
||||
# layer 3: H/8 -> H (8x upscale)
|
||||
self.ps3 = nn.PixelShuffle(8)
|
||||
self.conv3 = nn.Conv2d(4, 3, kernel_size=3, padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
# x shape: [B, 4096, H/32, W/32]
|
||||
x = self.act1(self.conv1(self.ps1((x)))) # -> [B, 256, H/16, W/16]
|
||||
x = self.act2(self.conv2(self.ps2((x)))) # -> [B, 256, H/8, W/8]
|
||||
x = self.conv3(self.ps3((x))) # -> [B, 3, H, W]
|
||||
return x
|
||||
|
||||
|
||||
class PatchDecoder_preps1(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# layer 1: H/32 -> H/16 (2x upscale)
|
||||
self.ps1 = nn.PixelShuffle(2)
|
||||
self.conv1 = nn.Conv2d(1024, 1024, kernel_size=3, padding=1)
|
||||
self.act1 = nn.GELU()
|
||||
|
||||
# layer 2: H/16 -> H/8 (2x upscale)
|
||||
self.ps2 = nn.PixelShuffle(2)
|
||||
self.conv2 = nn.Conv2d(256, 192, kernel_size=3, padding=1)
|
||||
|
||||
# layer 3: H/8 -> H (8x upscale)
|
||||
self.ps3 = nn.PixelShuffle(8)
|
||||
|
||||
def forward(self, x):
|
||||
# x shape: [B, 4096, H/32, W/32]
|
||||
x = self.act1(self.conv1(self.ps1((x)))) # -> [B, 256, H/16, W/16]
|
||||
x = self.ps3(self.conv2(self.ps2((x)))) # -> [B, 256, H/8, W/8]
|
||||
return x
|
||||
|
||||
|
||||
class ConvDecoder(nn.Module):
|
||||
def __init__(self, input_dim=4096, hidden_dim=1024):
|
||||
super().__init__()
|
||||
# layer 1: H/32 -> H/16 (2x upscale)
|
||||
self.ps1 = nn.PixelShuffle(2)
|
||||
self.conv1 = nn.Conv2d(input_dim // 4, hidden_dim, kernel_size=3, padding=1)
|
||||
self.act1 = nn.GELU()
|
||||
|
||||
# layer 2: H/16 -> H/8 (2x upscale)
|
||||
self.ps2 = nn.PixelShuffle(2)
|
||||
self.conv2 = nn.Conv2d(hidden_dim // 4, 192, kernel_size=3, padding=1)
|
||||
|
||||
# layer 3: H/8 -> H (8x upscale)
|
||||
self.ps3 = nn.PixelShuffle(8)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.act1(self.conv1(self.ps1((x))))
|
||||
x = self.ps3(self.conv2(self.ps2((x))))
|
||||
return x
|
||||
+2739
File diff suppressed because it is too large
Load Diff
+285
@@ -0,0 +1,285 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
from torch import nn
|
||||
from transformers.modeling_outputs import BaseModelOutputWithPooling
|
||||
from transformers.modeling_utils import PreTrainedModel
|
||||
|
||||
from .configuration_neo_vit import NEOVisionConfig
|
||||
|
||||
|
||||
def precompute_rope_freqs_sincos(
|
||||
dim: int, max_position: int, base: float = 10000.0, device=None
|
||||
):
|
||||
"""Precompute 1D RoPE cosine and sine values."""
|
||||
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device).float() / dim))
|
||||
t = torch.arange(max_position, device=device).type_as(inv_freq)
|
||||
freqs = torch.outer(t, inv_freq)
|
||||
return torch.cos(freqs), torch.sin(freqs)
|
||||
|
||||
|
||||
def build_abs_positions_from_grid_hw(grid_hw: torch.Tensor, device=None):
|
||||
"""
|
||||
Compute patch coordinates (x, y)
|
||||
|
||||
Args:
|
||||
grid_hw: (B, 2) tensor representing (H, W) per image
|
||||
"""
|
||||
device = grid_hw.device
|
||||
B = grid_hw.shape[0]
|
||||
|
||||
# Get the number of patches per image
|
||||
H = grid_hw[:, 0]
|
||||
W = grid_hw[:, 1]
|
||||
N = H * W
|
||||
N_total = N.sum()
|
||||
|
||||
# Create the batch index for each patch (B x patch count)
|
||||
patch_to_sample = torch.repeat_interleave(
|
||||
torch.arange(B, device=device), N
|
||||
) # (N_total,)
|
||||
|
||||
# Generate intra-image patch index (row-major order)
|
||||
patch_id_within_image = torch.arange(N_total, device=device)
|
||||
patch_id_within_image = (
|
||||
patch_id_within_image
|
||||
- torch.cumsum(torch.cat([torch.tensor([0], device=device), N[:-1]]), dim=0)[
|
||||
patch_to_sample
|
||||
]
|
||||
)
|
||||
|
||||
# Get H/W for each patch according to its image
|
||||
W_per_patch = W[patch_to_sample]
|
||||
abs_x = patch_id_within_image % W_per_patch
|
||||
abs_y = patch_id_within_image // W_per_patch
|
||||
|
||||
return abs_x, abs_y
|
||||
|
||||
|
||||
def apply_rotary_emb_1d(
|
||||
x: torch.Tensor,
|
||||
cos_cached: torch.Tensor,
|
||||
sin_cached: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
):
|
||||
"""Apply 1D RoPE to part of the input tensor."""
|
||||
# x: (..., seq_len, dim_part)
|
||||
# positions: (..., seq_len)
|
||||
# cos_cached: (max_pos, dim_part / 2)
|
||||
|
||||
cos = cos_cached[positions] # Shape: (positions.shape, dim_part / 2)
|
||||
sin = sin_cached[positions] # Shape: (positions.shape, dim_part / 2)
|
||||
|
||||
x1 = x[..., 0::2]
|
||||
x2 = x[..., 1::2]
|
||||
|
||||
rotated_x1 = x1 * cos - x2 * sin
|
||||
rotated_x2 = x1 * sin + x2 * cos
|
||||
|
||||
x_rotated = torch.empty_like(x)
|
||||
x_rotated[..., 0::2] = rotated_x1
|
||||
x_rotated[..., 1::2] = rotated_x2
|
||||
return x_rotated
|
||||
|
||||
|
||||
def apply_2d_rotary_pos_emb(
|
||||
x: torch.Tensor,
|
||||
cos_cached_x: torch.Tensor,
|
||||
sin_cached_x: torch.Tensor,
|
||||
cos_cached_y: torch.Tensor,
|
||||
sin_cached_y: torch.Tensor,
|
||||
abs_positions_x: torch.Tensor,
|
||||
abs_positions_y: torch.Tensor,
|
||||
):
|
||||
"""Apply 2D RoPE to input tensor x."""
|
||||
dim = x.shape[-1]
|
||||
dim_half = dim // 2
|
||||
|
||||
# Use the first half of the embedding for one RoPE direction and the second
|
||||
# half for the other direction. The split order must stay consistent.
|
||||
x_part_1 = x[..., :dim_half]
|
||||
x_part_2 = x[..., dim_half:]
|
||||
|
||||
# Apply rotations associated with abs_positions_x to x_part_1.
|
||||
rotated_part_1 = apply_rotary_emb_1d(
|
||||
x_part_1, cos_cached_x, sin_cached_x, abs_positions_x
|
||||
)
|
||||
# Apply rotations associated with abs_positions_y to x_part_2.
|
||||
rotated_part_2 = apply_rotary_emb_1d(
|
||||
x_part_2, cos_cached_y, sin_cached_y, abs_positions_y
|
||||
)
|
||||
|
||||
# Concatenate them back in the same order used by the split.
|
||||
return torch.cat((rotated_part_1, rotated_part_2), dim=-1)
|
||||
|
||||
|
||||
class NEOVisionEmbeddings(nn.Module):
|
||||
"""
|
||||
Embedding Module for Vision.
|
||||
"""
|
||||
|
||||
def __init__(self, config: NEOVisionConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.embed_dim = config.hidden_size
|
||||
self.llm_embed_dim = config.llm_hidden_size[0]
|
||||
self.downsample_factor = int(1 / config.downsample_ratio[0])
|
||||
self.patch_size = config.patch_size
|
||||
|
||||
self.patch_embedding = nn.Conv2d(
|
||||
in_channels=config.num_channels,
|
||||
out_channels=self.embed_dim,
|
||||
kernel_size=self.patch_size,
|
||||
stride=self.patch_size,
|
||||
)
|
||||
self.dense_embedding = nn.Conv2d(
|
||||
in_channels=self.embed_dim,
|
||||
out_channels=self.llm_embed_dim,
|
||||
kernel_size=self.downsample_factor,
|
||||
stride=self.downsample_factor,
|
||||
)
|
||||
self.gelu = nn.GELU()
|
||||
|
||||
self.rope_dim_part = self.embed_dim // 2
|
||||
self.max_position_embeddings_vision = config.max_position_embeddings_vision
|
||||
self.rope_theta_vision = config.rope_theta_vision
|
||||
|
||||
# These deterministic caches are not checkpoint state. In
|
||||
# Transformers 5, ``from_pretrained`` constructs models on the meta
|
||||
# device; tensors computed here would later be materialized as
|
||||
# uninitialized memory because persistent=False buffers are absent
|
||||
# from the checkpoint. Build them lazily on the first real device.
|
||||
self.register_buffer("cos_cached_x", None, persistent=False)
|
||||
self.register_buffer("sin_cached_x", None, persistent=False)
|
||||
self.register_buffer("cos_cached_y", None, persistent=False)
|
||||
self.register_buffer("sin_cached_y", None, persistent=False)
|
||||
|
||||
def _ensure_rope_cache(self, device: torch.device) -> None:
|
||||
if self.cos_cached_x is not None and self.cos_cached_x.device == device:
|
||||
return
|
||||
|
||||
cos, sin = precompute_rope_freqs_sincos(
|
||||
self.rope_dim_part,
|
||||
self.max_position_embeddings_vision,
|
||||
base=self.rope_theta_vision,
|
||||
device=device,
|
||||
)
|
||||
self.cos_cached_x = cos
|
||||
self.sin_cached_x = sin
|
||||
self.cos_cached_y = cos.clone()
|
||||
self.sin_cached_y = sin.clone()
|
||||
|
||||
def _apply_2d_rotary_pos_emb(self, patch_embeds, grid_hw):
|
||||
"""
|
||||
Apply 2D Rotary Position Embedding to the patch embeddings.
|
||||
"""
|
||||
abs_pos_x, abs_pos_y = build_abs_positions_from_grid_hw(
|
||||
grid_hw, device=patch_embeds.device
|
||||
)
|
||||
embeddings = apply_2d_rotary_pos_emb(
|
||||
patch_embeds.to(
|
||||
torch.float32
|
||||
), # RoPE calculations are often more stable in float32
|
||||
self.cos_cached_x,
|
||||
self.sin_cached_x,
|
||||
self.cos_cached_y,
|
||||
self.sin_cached_y,
|
||||
abs_pos_x,
|
||||
abs_pos_y,
|
||||
).to(self.patch_embedding.weight.dtype)
|
||||
return embeddings
|
||||
|
||||
def forward(self, pixel_values: torch.FloatTensor, grid_hw=None) -> torch.Tensor:
|
||||
|
||||
pixel_values = pixel_values.view( #
|
||||
-1,
|
||||
3,
|
||||
self.patch_size,
|
||||
self.patch_size,
|
||||
) # [28072, 768] -> [28072, 3, 16, 16]
|
||||
patch_embeds = self.gelu(self.patch_embedding(pixel_values)).view(
|
||||
-1, self.embed_dim
|
||||
)
|
||||
self._ensure_rope_cache(patch_embeds.device)
|
||||
patch_embeds = self._apply_2d_rotary_pos_emb(
|
||||
patch_embeds, grid_hw
|
||||
) # [28072, 1024]
|
||||
assert (grid_hw[:, 0] * grid_hw[:, 1]).sum() == patch_embeds.shape[0]
|
||||
|
||||
patches_list = []
|
||||
cur_position = 0
|
||||
for i in range(grid_hw.shape[0]):
|
||||
h, w = grid_hw[i]
|
||||
patches_per_img = (
|
||||
patch_embeds[cur_position : cur_position + h * w]
|
||||
.view(h, w, -1)
|
||||
.unsqueeze(0)
|
||||
)
|
||||
patches_per_img = self.dense_embedding(patches_per_img.permute(0, 3, 1, 2))
|
||||
patches_per_img = patches_per_img.permute(0, 2, 3, 1)
|
||||
patches_list.append(patches_per_img.view(-1, patches_per_img.shape[-1]))
|
||||
cur_position += h * w
|
||||
|
||||
embeddings = torch.cat(
|
||||
patches_list, dim=0
|
||||
) # (N_total // downsample_factor**2, C)
|
||||
|
||||
assert cur_position == patch_embeds.shape[0]
|
||||
assert embeddings.shape[0] == int(
|
||||
patch_embeds.shape[0] / self.downsample_factor**2
|
||||
)
|
||||
|
||||
return embeddings
|
||||
|
||||
|
||||
class NEOVisionModel(PreTrainedModel):
|
||||
main_input_name = "pixel_values"
|
||||
_supports_flash_attn_2 = True
|
||||
supports_gradient_checkpointing = True
|
||||
config_class = NEOVisionConfig
|
||||
# support transformers 4.51.+
|
||||
_tp_plan = ""
|
||||
|
||||
def __init__(self, config: NEOVisionConfig):
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
|
||||
self.embeddings = NEOVisionEmbeddings(config)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
pixel_values: Optional[torch.FloatTensor] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
pixel_embeds: Optional[torch.FloatTensor] = None,
|
||||
grid_hw: Optional[torch.Tensor] = None,
|
||||
) -> Union[Tuple, BaseModelOutputWithPooling]:
|
||||
output_hidden_states = (
|
||||
output_hidden_states
|
||||
if output_hidden_states is not None
|
||||
else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = (
|
||||
return_dict if return_dict is not None else self.config.use_return_dict
|
||||
)
|
||||
|
||||
if pixel_values is None and pixel_embeds is None:
|
||||
raise ValueError("You have to specify pixel_values or pixel_embeds")
|
||||
|
||||
if pixel_embeds is not None:
|
||||
hidden_states = pixel_embeds
|
||||
else:
|
||||
assert pixel_values.dim() == 2, (
|
||||
f"pixel_values must be 2D for native resolution, got: {pixel_values.dim()}"
|
||||
)
|
||||
hidden_states = self.embeddings(pixel_values, grid_hw=grid_hw)
|
||||
|
||||
return BaseModelOutputWithPooling(
|
||||
last_hidden_state=hidden_states,
|
||||
pooler_output=None,
|
||||
hidden_states=None,
|
||||
attentions=None,
|
||||
)
|
||||
+1540
File diff suppressed because it is too large
Load Diff
+619
@@ -0,0 +1,619 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from transformers.cache_utils import Cache, DynamicCache
|
||||
from transformers.generation import GenerationMixin
|
||||
from transformers.masking_utils import create_causal_mask
|
||||
from transformers.modeling_layers import GradientCheckpointingLayer
|
||||
from transformers.modeling_outputs import (
|
||||
BaseModelOutputWithPast,
|
||||
CausalLMOutputWithPast,
|
||||
)
|
||||
from transformers.modeling_utils import PreTrainedModel
|
||||
from transformers.processing_utils import Unpack
|
||||
from transformers.utils import TransformersKwargs, can_return_tuple
|
||||
from transformers.utils.deprecation import deprecate_kwarg
|
||||
|
||||
from .configuration_neo_chat import NEOMoELLMConfig
|
||||
from .modeling_qwen3 import (
|
||||
Qwen3Attention,
|
||||
Qwen3RMSNorm,
|
||||
create_block_causal_mask,
|
||||
)
|
||||
from .transformers_compat import (
|
||||
causal_mask_kwargs,
|
||||
model_input_compat,
|
||||
tied_weights_keys,
|
||||
)
|
||||
|
||||
|
||||
class Qwen3MoeMLP(nn.Module):
|
||||
"""Single expert FFN. Same structure as :class:`Qwen3MLP` but the
|
||||
intermediate size is parameterised so it can be ``moe_intermediate_size``
|
||||
(per-expert) for experts and ``intermediate_size`` for any dense fallback.
|
||||
"""
|
||||
|
||||
def __init__(self, config, intermediate_size: Optional[int] = None):
|
||||
super().__init__()
|
||||
from transformers.activations import ACT2FN
|
||||
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.intermediate_size = (
|
||||
intermediate_size
|
||||
if intermediate_size is not None
|
||||
else config.intermediate_size
|
||||
)
|
||||
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
||||
self.act_fn = ACT2FN[config.hidden_act]
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class Qwen3MoeSparseMoeBlock(nn.Module):
|
||||
"""Top-k softmax-routed MoE block matching HuggingFace's Qwen3-MoE layout.
|
||||
|
||||
Parameter names (``gate.weight``, ``experts.{i}.gate_proj/up_proj/down_proj``)
|
||||
are kept identical so converted A3B checkpoints load directly via the
|
||||
``mlp.*`` / ``mlp_mot_gen.*`` keys. The block is parameterised explicitly
|
||||
so the same class can serve both the understanding branch (``num_experts``
|
||||
experts, top-k = ``num_experts_per_tok``, width ``moe_intermediate_size``)
|
||||
and the image-generation branch (``gen_num_experts`` etc.).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: NEOMoELLMConfig,
|
||||
num_experts: Optional[int] = None,
|
||||
num_experts_per_tok: Optional[int] = None,
|
||||
moe_intermediate_size: Optional[int] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.num_experts = (
|
||||
int(num_experts) if num_experts is not None else int(config.num_experts)
|
||||
)
|
||||
self.top_k = int(
|
||||
num_experts_per_tok
|
||||
if num_experts_per_tok is not None
|
||||
else config.num_experts_per_tok
|
||||
)
|
||||
self.norm_topk_prob = bool(getattr(config, "norm_topk_prob", True))
|
||||
self.hidden_size = config.hidden_size
|
||||
|
||||
expert_intermediate_size = int(
|
||||
moe_intermediate_size
|
||||
if moe_intermediate_size is not None
|
||||
else config.moe_intermediate_size
|
||||
)
|
||||
|
||||
self.gate = nn.Linear(config.hidden_size, self.num_experts, bias=False)
|
||||
self.experts = nn.ModuleList(
|
||||
[
|
||||
Qwen3MoeMLP(config, intermediate_size=expert_intermediate_size)
|
||||
for _ in range(self.num_experts)
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
orig_shape = hidden_states.shape
|
||||
hidden_dim = orig_shape[-1]
|
||||
flat = hidden_states.view(-1, hidden_dim)
|
||||
n_tokens = flat.shape[0]
|
||||
|
||||
router_logits = self.gate(flat)
|
||||
routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float32)
|
||||
routing_weights, selected_experts = torch.topk(
|
||||
routing_weights, self.top_k, dim=-1
|
||||
)
|
||||
if self.norm_topk_prob:
|
||||
routing_weights = routing_weights / routing_weights.sum(
|
||||
dim=-1, keepdim=True
|
||||
)
|
||||
routing_weights = routing_weights.to(flat.dtype)
|
||||
|
||||
output = torch.zeros(
|
||||
(n_tokens, hidden_dim), dtype=flat.dtype, device=flat.device
|
||||
)
|
||||
# (num_experts, top_k, num_tokens)
|
||||
expert_mask = F.one_hot(selected_experts, num_classes=self.num_experts).permute(
|
||||
2, 1, 0
|
||||
)
|
||||
|
||||
for expert_idx in range(self.num_experts):
|
||||
idx, top_x = torch.where(expert_mask[expert_idx])
|
||||
if top_x.numel() == 0:
|
||||
continue
|
||||
expert_layer = self.experts[expert_idx]
|
||||
current_state = flat.index_select(0, top_x)
|
||||
current_out = (
|
||||
expert_layer(current_state) * routing_weights[top_x, idx, None]
|
||||
)
|
||||
output.index_add_(0, top_x, current_out.to(flat.dtype))
|
||||
|
||||
return output.view(*orig_shape)
|
||||
|
||||
|
||||
class Qwen3MoeDecoderLayer(GradientCheckpointingLayer):
|
||||
"""A Qwen3-MoE decoder block with the NEO-Unify two-branch structure.
|
||||
|
||||
Mirrors ``Qwen3DecoderLayer`` from :mod:`modeling_qwen3` but uses sparse
|
||||
MoE blocks on *both* branches:
|
||||
|
||||
* ``self.mlp`` - understanding-path MoE
|
||||
(``num_experts`` / ``num_experts_per_tok`` /
|
||||
``moe_intermediate_size``)
|
||||
* ``self.mlp_mot_gen`` - image-generation-path MoE
|
||||
(``gen_num_experts`` / ``gen_num_experts_per_tok`` /
|
||||
``gen_moe_intermediate_size``)
|
||||
|
||||
Layers listed in ``mlp_only_layers`` or those not aligned with
|
||||
``decoder_sparse_step`` fall back to a dense :class:`Qwen3MoeMLP` on the
|
||||
understanding branch (matching upstream Qwen3-MoE), while the
|
||||
generation branch still uses a sparse MoE.
|
||||
"""
|
||||
|
||||
def __init__(self, config: NEOMoELLMConfig, layer_idx: int):
|
||||
super().__init__()
|
||||
self.hidden_size = config.hidden_size
|
||||
self.self_attn = Qwen3Attention(config=config, layer_idx=layer_idx)
|
||||
|
||||
mlp_only_layers = list(getattr(config, "mlp_only_layers", []) or [])
|
||||
decoder_sparse_step = int(getattr(config, "decoder_sparse_step", 1) or 1)
|
||||
is_sparse = (
|
||||
int(config.num_experts) > 0
|
||||
and layer_idx not in mlp_only_layers
|
||||
and (layer_idx + 1) % decoder_sparse_step == 0
|
||||
)
|
||||
|
||||
if is_sparse:
|
||||
self.mlp = Qwen3MoeSparseMoeBlock(
|
||||
config,
|
||||
num_experts=config.num_experts,
|
||||
num_experts_per_tok=config.num_experts_per_tok,
|
||||
moe_intermediate_size=config.moe_intermediate_size,
|
||||
)
|
||||
else:
|
||||
self.mlp = Qwen3MoeMLP(config, intermediate_size=config.intermediate_size)
|
||||
|
||||
# Image-generation branch: in the A3B checkpoint this is *also* a sparse
|
||||
# MoE block (``gen_num_experts`` experts, typically smaller than the und
|
||||
# branch's ``num_experts``). ``NEOMoELLMConfig`` defaults the gen-path
|
||||
# knobs to their und-path counterparts so legacy single-pool configs
|
||||
# keep working.
|
||||
self.mlp_mot_gen = Qwen3MoeSparseMoeBlock(
|
||||
config,
|
||||
num_experts=getattr(config, "gen_num_experts", config.num_experts),
|
||||
num_experts_per_tok=getattr(
|
||||
config, "gen_num_experts_per_tok", config.num_experts_per_tok
|
||||
),
|
||||
moe_intermediate_size=getattr(
|
||||
config, "gen_moe_intermediate_size", config.moe_intermediate_size
|
||||
),
|
||||
)
|
||||
|
||||
self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.input_layernorm_mot_gen = Qwen3RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
self.post_attention_layernorm = Qwen3RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
self.post_attention_layernorm_mot_gen = Qwen3RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
self.attention_type = config.layer_types[layer_idx]
|
||||
|
||||
def forward_und(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
image_gen_indicators: torch.Tensor,
|
||||
exist_non_image_gen_tokens: bool,
|
||||
exist_image_gen_tokens: bool,
|
||||
indexes: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
use_cache: Optional[bool] = False,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> torch.Tensor:
|
||||
residual = hidden_states
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
hidden_states, _ = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
image_gen_indicators=image_gen_indicators,
|
||||
exist_non_image_gen_tokens=exist_non_image_gen_tokens,
|
||||
exist_image_gen_tokens=exist_image_gen_tokens,
|
||||
indexes=indexes,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
residual = hidden_states
|
||||
hidden_states = self.post_attention_layernorm(hidden_states)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
return hidden_states
|
||||
|
||||
def forward_gen(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
image_gen_indicators: torch.Tensor,
|
||||
exist_non_image_gen_tokens: bool,
|
||||
exist_image_gen_tokens: bool,
|
||||
indexes: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
use_cache: Optional[bool] = False,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> torch.Tensor:
|
||||
residual = hidden_states
|
||||
hidden_states = self.input_layernorm_mot_gen(hidden_states)
|
||||
hidden_states, _ = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
image_gen_indicators=image_gen_indicators,
|
||||
exist_non_image_gen_tokens=exist_non_image_gen_tokens,
|
||||
exist_image_gen_tokens=exist_image_gen_tokens,
|
||||
indexes=indexes,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
residual = hidden_states
|
||||
hidden_states = self.post_attention_layernorm_mot_gen(hidden_states)
|
||||
hidden_states = self.mlp_mot_gen(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
return hidden_states
|
||||
|
||||
@deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
image_gen_indicators: torch.Tensor,
|
||||
exist_non_image_gen_tokens: bool,
|
||||
exist_image_gen_tokens: bool,
|
||||
indexes: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
use_cache: Optional[bool] = False,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> torch.Tensor:
|
||||
if exist_non_image_gen_tokens and not exist_image_gen_tokens:
|
||||
return self.forward_und(
|
||||
hidden_states,
|
||||
image_gen_indicators,
|
||||
exist_non_image_gen_tokens,
|
||||
exist_image_gen_tokens,
|
||||
indexes,
|
||||
attention_mask,
|
||||
position_ids,
|
||||
past_key_values,
|
||||
use_cache,
|
||||
cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
if not exist_non_image_gen_tokens and exist_image_gen_tokens:
|
||||
return self.forward_gen(
|
||||
hidden_states,
|
||||
image_gen_indicators,
|
||||
exist_non_image_gen_tokens,
|
||||
exist_image_gen_tokens,
|
||||
indexes,
|
||||
attention_mask,
|
||||
position_ids,
|
||||
past_key_values,
|
||||
use_cache,
|
||||
cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Mixed und/gen path — see the NOTE in Qwen3Attention.forward (modeling_qwen3.py).
|
||||
raise NotImplementedError(
|
||||
"Mixed und/gen decoder-layer forward is not yet validated (issue #207). "
|
||||
"Split the sequence at token-type boundaries and use forward_und / forward_gen."
|
||||
)
|
||||
|
||||
# Mixed batch: dispatch tokens per branch then merge back. Matches the
|
||||
# dense ``Qwen3DecoderLayer.forward`` mixed-path implementation.
|
||||
residual = hidden_states
|
||||
|
||||
_hidden_states = hidden_states.new_zeros(hidden_states.shape)
|
||||
if exist_non_image_gen_tokens:
|
||||
_hidden_states[~image_gen_indicators] = self.input_layernorm(
|
||||
hidden_states[~image_gen_indicators]
|
||||
)
|
||||
if exist_image_gen_tokens:
|
||||
_hidden_states[image_gen_indicators] = self.input_layernorm_mot_gen(
|
||||
hidden_states[image_gen_indicators]
|
||||
)
|
||||
hidden_states = _hidden_states
|
||||
|
||||
hidden_states, _ = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
image_gen_indicators=image_gen_indicators,
|
||||
exist_non_image_gen_tokens=exist_non_image_gen_tokens,
|
||||
exist_image_gen_tokens=exist_image_gen_tokens,
|
||||
indexes=indexes,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
residual = hidden_states
|
||||
|
||||
_hidden_states = hidden_states.new_zeros(hidden_states.shape)
|
||||
if exist_non_image_gen_tokens:
|
||||
und_hidden = self.post_attention_layernorm(
|
||||
hidden_states[~image_gen_indicators]
|
||||
)
|
||||
# MoE expects a 3D input (batch, seq, hidden); promote then squeeze.
|
||||
if und_hidden.dim() == 2:
|
||||
und_hidden = und_hidden.unsqueeze(0)
|
||||
_hidden_states[~image_gen_indicators] = self.mlp(und_hidden).squeeze(0)
|
||||
else:
|
||||
_hidden_states[~image_gen_indicators] = self.mlp(und_hidden)
|
||||
if exist_image_gen_tokens:
|
||||
_hidden_states[image_gen_indicators] = self.mlp_mot_gen(
|
||||
self.post_attention_layernorm_mot_gen(
|
||||
hidden_states[image_gen_indicators]
|
||||
)
|
||||
)
|
||||
|
||||
hidden_states = _hidden_states
|
||||
hidden_states = residual + hidden_states
|
||||
return hidden_states
|
||||
|
||||
|
||||
class Qwen3MoePreTrainedModel(PreTrainedModel):
|
||||
config: NEOMoELLMConfig
|
||||
base_model_prefix = "model"
|
||||
supports_gradient_checkpointing = True
|
||||
_no_split_modules = ["Qwen3MoeDecoderLayer"]
|
||||
_skip_keys_device_placement = ["past_key_values"]
|
||||
_supports_flash_attn = True
|
||||
_supports_sdpa = True
|
||||
_supports_flex_attn = True
|
||||
|
||||
_can_compile_fullgraph = False # MoE routing has data-dependent control flow.
|
||||
_supports_attention_backend = True
|
||||
_can_record_outputs = {
|
||||
"hidden_states": Qwen3MoeDecoderLayer,
|
||||
"attentions": Qwen3Attention,
|
||||
}
|
||||
|
||||
|
||||
class Qwen3MoeModel(Qwen3MoePreTrainedModel):
|
||||
def __init__(self, config: NEOMoELLMConfig):
|
||||
super().__init__(config)
|
||||
self.padding_idx = config.pad_token_id
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.embed_tokens = nn.Embedding(
|
||||
config.vocab_size, config.hidden_size, self.padding_idx
|
||||
)
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
Qwen3MoeDecoderLayer(config, layer_idx)
|
||||
for layer_idx in range(config.num_hidden_layers)
|
||||
]
|
||||
)
|
||||
self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.norm_mot_gen = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
self.has_sliding_layers = "sliding_attention" in self.config.layer_types
|
||||
self.current_index = -1
|
||||
|
||||
self.post_init()
|
||||
|
||||
@model_input_compat
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
image_gen_indicators: Optional[torch.Tensor] = None,
|
||||
indexes: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> BaseModelOutputWithPast:
|
||||
if image_gen_indicators is None:
|
||||
exist_non_image_gen_tokens = True
|
||||
exist_image_gen_tokens = False
|
||||
else:
|
||||
# Convert the CUDA reductions once before the decoder loop. If the
|
||||
# scalar tensors reach every layer, each Python branch can force a
|
||||
# host-device synchronization and collapse async weight prefetch.
|
||||
exist_non_image_gen_tokens = bool((~image_gen_indicators).any().item())
|
||||
exist_image_gen_tokens = bool(image_gen_indicators.any().item())
|
||||
|
||||
if (input_ids is None) ^ (inputs_embeds is not None):
|
||||
raise ValueError(
|
||||
"You must specify exactly one of input_ids or inputs_embeds"
|
||||
)
|
||||
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.embed_tokens(input_ids)
|
||||
|
||||
if use_cache and past_key_values is None:
|
||||
past_key_values = DynamicCache(config=self.config)
|
||||
|
||||
if cache_position is None:
|
||||
past_seen_tokens = (
|
||||
past_key_values.get_seq_length() if past_key_values is not None else 0
|
||||
)
|
||||
cache_position = torch.arange(
|
||||
past_seen_tokens,
|
||||
past_seen_tokens + inputs_embeds.shape[1],
|
||||
device=inputs_embeds.device,
|
||||
)
|
||||
|
||||
if position_ids is None:
|
||||
position_ids = cache_position.unsqueeze(0)
|
||||
|
||||
if not isinstance(causal_mask_mapping := attention_mask, dict):
|
||||
if input_ids is not None:
|
||||
mask_kwargs = causal_mask_kwargs(
|
||||
create_causal_mask,
|
||||
config=self.config,
|
||||
inputs_embeds=inputs_embeds,
|
||||
attention_mask=attention_mask,
|
||||
cache_position=cache_position,
|
||||
past_key_values=past_key_values,
|
||||
position_ids=position_ids,
|
||||
)
|
||||
causal_mask_mapping = {
|
||||
"full_attention": create_causal_mask(**mask_kwargs),
|
||||
}
|
||||
self.current_index += 1
|
||||
indexes = torch.LongTensor([[self.current_index], [0], [0]]).to(
|
||||
input_ids.device
|
||||
)
|
||||
else:
|
||||
causal_mask_mapping = {
|
||||
"full_attention": create_block_causal_mask(indexes[0]),
|
||||
}
|
||||
self.current_index = indexes[0].max()
|
||||
else:
|
||||
self.current_index = indexes[0].max()
|
||||
|
||||
hidden_states = inputs_embeds
|
||||
|
||||
for decoder_layer in self.layers[: self.config.num_hidden_layers]:
|
||||
hidden_states = decoder_layer(
|
||||
hidden_states,
|
||||
image_gen_indicators=image_gen_indicators,
|
||||
exist_non_image_gen_tokens=exist_non_image_gen_tokens,
|
||||
exist_image_gen_tokens=exist_image_gen_tokens,
|
||||
indexes=indexes,
|
||||
attention_mask=causal_mask_mapping[decoder_layer.attention_type],
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if not exist_image_gen_tokens:
|
||||
hidden_states = self.norm(hidden_states)
|
||||
elif not exist_non_image_gen_tokens:
|
||||
hidden_states = self.norm_mot_gen(hidden_states)
|
||||
else:
|
||||
_hidden_states = hidden_states.new_zeros(hidden_states.shape)
|
||||
_hidden_states[~image_gen_indicators] = self.norm(
|
||||
hidden_states[~image_gen_indicators]
|
||||
)
|
||||
_hidden_states[image_gen_indicators] = self.norm_mot_gen(
|
||||
hidden_states[image_gen_indicators]
|
||||
)
|
||||
hidden_states = _hidden_states
|
||||
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state=hidden_states,
|
||||
past_key_values=past_key_values if use_cache else None,
|
||||
)
|
||||
|
||||
|
||||
class Qwen3MoeForCausalLM(Qwen3MoePreTrainedModel, GenerationMixin):
|
||||
_tied_weights_keys = tied_weights_keys(
|
||||
"lm_head.weight", "model.embed_tokens.weight"
|
||||
)
|
||||
_tp_plan = {"lm_head": "colwise_rep"}
|
||||
_pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
|
||||
|
||||
def __init__(self, config: NEOMoELLMConfig):
|
||||
super().__init__(config)
|
||||
self.model = Qwen3MoeModel(config)
|
||||
self.vocab_size = config.vocab_size
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
self.post_init()
|
||||
|
||||
@can_return_tuple
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
indexes: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
logits_to_keep: Union[int, torch.Tensor] = 0,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> CausalLMOutputWithPast:
|
||||
outputs: BaseModelOutputWithPast = self.model(
|
||||
input_ids=input_ids,
|
||||
indexes=indexes,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
hidden_states = outputs.last_hidden_state
|
||||
slice_indices = (
|
||||
slice(-logits_to_keep, None)
|
||||
if isinstance(logits_to_keep, int)
|
||||
else logits_to_keep
|
||||
)
|
||||
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
loss = self.loss_function(
|
||||
logits=logits,
|
||||
labels=labels,
|
||||
vocab_size=self.config.vocab_size,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Qwen3MoeForCausalLM",
|
||||
"Qwen3MoeModel",
|
||||
"Qwen3MoePreTrainedModel",
|
||||
"Qwen3MoeDecoderLayer",
|
||||
"Qwen3MoeSparseMoeBlock",
|
||||
"Qwen3MoeMLP",
|
||||
]
|
||||
+71
@@ -0,0 +1,71 @@
|
||||
"""Small compatibility seams for the supported Transformers 4/5 window."""
|
||||
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from functools import lru_cache
|
||||
from typing import Any, Callable
|
||||
|
||||
import torch
|
||||
from packaging.version import Version
|
||||
|
||||
try:
|
||||
from transformers.utils.generic import merge_with_config_defaults
|
||||
from transformers.utils.output_capturing import capture_outputs
|
||||
|
||||
def model_input_compat(func: Callable[..., Any]) -> Callable[..., Any]:
|
||||
return merge_with_config_defaults(capture_outputs(func))
|
||||
|
||||
except ImportError: # Transformers 4.x
|
||||
pass
|
||||
|
||||
|
||||
@lru_cache(maxsize=None)
|
||||
def _parameter_names(callable_: Callable[..., Any]) -> frozenset[str]:
|
||||
return frozenset(inspect.signature(callable_).parameters)
|
||||
|
||||
|
||||
def causal_mask_kwargs(
|
||||
mask_factory: Callable[..., Any],
|
||||
*,
|
||||
config: Any,
|
||||
inputs_embeds: torch.Tensor,
|
||||
attention_mask: torch.Tensor | None,
|
||||
cache_position: torch.Tensor,
|
||||
past_key_values: Any,
|
||||
position_ids: torch.Tensor,
|
||||
) -> dict[str, Any]:
|
||||
"""Build arguments accepted by the installed ``create_causal_mask``.
|
||||
|
||||
Transformers 4.57 uses ``input_embeds`` and ``cache_position`` while newer
|
||||
5.x releases use ``inputs_embeds`` and derive the cache position internally.
|
||||
"""
|
||||
parameters = _parameter_names(mask_factory)
|
||||
embedding_parameter = (
|
||||
"inputs_embeds" if "inputs_embeds" in parameters else "input_embeds"
|
||||
)
|
||||
candidates = {
|
||||
"config": config,
|
||||
embedding_parameter: inputs_embeds,
|
||||
"attention_mask": attention_mask,
|
||||
"cache_position": cache_position,
|
||||
"past_key_values": past_key_values,
|
||||
"position_ids": position_ids,
|
||||
}
|
||||
return {name: value for name, value in candidates.items() if name in parameters}
|
||||
|
||||
|
||||
def pretrained_dtype_kwargs(dtype: torch.dtype) -> dict[str, torch.dtype]:
|
||||
"""Use the public dtype keyword supported throughout Transformers 4.57+."""
|
||||
return {"dtype": dtype}
|
||||
|
||||
|
||||
def tied_weights_keys(output_key: str, input_key: str) -> list[str] | dict[str, str]:
|
||||
"""Return the `_tied_weights_keys` shape expected by Transformers 4 or 5."""
|
||||
import transformers
|
||||
|
||||
if Version(transformers.__version__).major >= 5:
|
||||
return {output_key: input_key}
|
||||
return [output_key]
|
||||
@@ -0,0 +1,165 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torchvision.transforms as T
|
||||
from PIL import Image
|
||||
|
||||
SYSTEM_MESSAGE_FOR_GEN = (
|
||||
"You are an image generation and editing assistant that accurately understands and executes "
|
||||
"user intent.\n\nYou support two modes:\n\n1. Think Mode:\nIf the task requires reasoning, you "
|
||||
"MUST start with a <think></think> block. Put all reasoning inside the block using plain text. "
|
||||
"DO NOT include any image tags. Keep it reasonable and directly useful for producing the final "
|
||||
"image.\n\n2. Non-Think Mode:\nIf no reasoning is needed, directly produce the final image.\n\n"
|
||||
"Task Types:\n\nA. Text-to-Image Generation:\n"
|
||||
"- Generate a high-quality image based on the user's description.\n"
|
||||
"- Ensure visual clarity, semantic consistency, and completeness.\n"
|
||||
"- DO NOT introduce elements that contradict or override the user's intent.\n\n"
|
||||
"B. Image Editing:\n"
|
||||
"- Use the provided image(s) as input or reference for modification or transformation.\n"
|
||||
"- The result can be an edited image or a new image based on the reference(s).\n"
|
||||
"- Preserve all unspecified attributes unless explicitly changed.\n\n"
|
||||
"General Rules:\n"
|
||||
"- For any visible text in the image, follow the language specified for the rendered text in "
|
||||
"the user's description, not the language of the prompt. If no language is specified, use the "
|
||||
"user's input language."
|
||||
)
|
||||
|
||||
IMAGENET_MEAN = (0.485, 0.456, 0.406)
|
||||
IMAGENET_STD = (0.229, 0.224, 0.225)
|
||||
|
||||
|
||||
def round_by_factor(number: float, factor: int) -> int:
|
||||
"""Returns the closest integer to `number` that is divisible by `factor`."""
|
||||
return round(number / factor) * factor
|
||||
|
||||
|
||||
def ceil_by_factor(number: float, factor: int) -> int:
|
||||
"""Returns the smallest integer >= `number` that is divisible by `factor`."""
|
||||
return math.ceil(number / factor) * factor
|
||||
|
||||
|
||||
def floor_by_factor(number: float, factor: int) -> int:
|
||||
"""Returns the largest integer <= `number` that is divisible by `factor`."""
|
||||
return math.floor(number / factor) * factor
|
||||
|
||||
|
||||
def smart_resize(
|
||||
height: int,
|
||||
width: int,
|
||||
factor: int = 32,
|
||||
min_pixels: int = 65536,
|
||||
max_pixels: int = 4194304,
|
||||
) -> tuple[int, int]:
|
||||
"""Rescale so that H/W are divisible by `factor` and total pixels ∈ [min, max].
|
||||
|
||||
Copied from https://github.com/QwenLM/Qwen2.5-VL/blob/main/qwen-vl-utils/src/qwen_vl_utils/vision_process.py#L60
|
||||
"""
|
||||
if max(height, width) / min(height, width) > 200:
|
||||
raise ValueError(
|
||||
f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
|
||||
)
|
||||
h_bar = max(factor, round_by_factor(height, factor))
|
||||
w_bar = max(factor, round_by_factor(width, factor))
|
||||
if h_bar * w_bar > max_pixels:
|
||||
beta = math.sqrt((height * width) / max_pixels)
|
||||
h_bar = max(factor, floor_by_factor(height / beta, factor))
|
||||
w_bar = max(factor, floor_by_factor(width / beta, factor))
|
||||
elif h_bar * w_bar < min_pixels:
|
||||
beta = math.sqrt(min_pixels / (height * width))
|
||||
h_bar = ceil_by_factor(height * beta, factor)
|
||||
w_bar = ceil_by_factor(width * beta, factor)
|
||||
return h_bar, w_bar
|
||||
|
||||
|
||||
def dynamic_preprocess_native_resolution(
|
||||
image: Image.Image,
|
||||
size_factor: int = 32,
|
||||
min_pixels: int = 65536,
|
||||
max_pixels: int = 4194304,
|
||||
**_kwargs,
|
||||
) -> Image.Image:
|
||||
width, height = image.size
|
||||
resized_height, resized_width = smart_resize(
|
||||
height,
|
||||
width,
|
||||
factor=size_factor,
|
||||
min_pixels=min_pixels,
|
||||
max_pixels=max_pixels,
|
||||
)
|
||||
return image.resize((resized_width, resized_height))
|
||||
|
||||
|
||||
def preprocess_pixel_values(pixel_values: torch.Tensor, patch_size: int = 16):
|
||||
c, h, w = pixel_values.shape
|
||||
grid_h = h // patch_size
|
||||
grid_w = w // patch_size
|
||||
|
||||
flatten_pixel_values = (
|
||||
pixel_values.view(c, grid_h, patch_size, grid_w, patch_size)
|
||||
.permute(1, 3, 0, 2, 4) # [grid_h, grid_w, c, patch_size, patch_size]
|
||||
.reshape(grid_h * grid_w, c * patch_size**2)
|
||||
)
|
||||
|
||||
grid_hw = torch.tensor([[grid_h, grid_w]], device=pixel_values.device)
|
||||
return flatten_pixel_values, grid_hw
|
||||
|
||||
|
||||
def get_contrasting_background(image: Image.Image):
|
||||
"""Return a background color for RGBA->RGB conversion, or ``None`` to use default.
|
||||
|
||||
The original Neo_Unify implementation computed a contrasting background
|
||||
from the alpha channel. For this open-source release we fall back to a
|
||||
plain white background; callers that need the smarter behavior can override
|
||||
this function.
|
||||
"""
|
||||
del image
|
||||
return (255, 255, 255)
|
||||
|
||||
|
||||
def load_image_native(
|
||||
image,
|
||||
patch_size: int = 16,
|
||||
downsample_ratio: float = 0.5,
|
||||
min_pixels: int = 65536,
|
||||
max_pixels: int = 4194304,
|
||||
upscale: bool = False,
|
||||
):
|
||||
"""Load and preprocess an image: RGB convert, smart-resize, normalize, patchify."""
|
||||
if not isinstance(image, Image.Image):
|
||||
image = Image.open(image)
|
||||
if image.mode == "RGBA":
|
||||
bg_color = get_contrasting_background(image)
|
||||
if bg_color:
|
||||
background = Image.new("RGB", image.size, bg_color)
|
||||
background.paste(image, mask=image.split()[3])
|
||||
image = background.convert("RGB")
|
||||
else:
|
||||
image = image.convert("RGB")
|
||||
else:
|
||||
image = image.convert("RGB")
|
||||
|
||||
if upscale:
|
||||
image = image.resize((image.width * 2, image.height * 2), Image.BILINEAR)
|
||||
|
||||
transform = T.Compose(
|
||||
[
|
||||
T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
|
||||
T.ToTensor(),
|
||||
T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
|
||||
]
|
||||
)
|
||||
|
||||
new_image = dynamic_preprocess_native_resolution(
|
||||
image,
|
||||
size_factor=int(patch_size // downsample_ratio),
|
||||
min_pixels=min_pixels,
|
||||
max_pixels=max_pixels,
|
||||
)
|
||||
pixel_values, grid_hw = preprocess_pixel_values(
|
||||
transform(new_image).to(torch.float32), patch_size=patch_size
|
||||
)
|
||||
return pixel_values, grid_hw
|
||||
@@ -0,0 +1,73 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.sensenova_u1 import (
|
||||
SenseNovaU1PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.sensenova_u1 import (
|
||||
SenseNovaU1SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
|
||||
from sglang.multimodal_gen.runtime.models.sensenova_u1.loader import (
|
||||
load_model_and_tokenizer,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages import InputValidationStage
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sensenova_u1 import (
|
||||
SenseNovaU1GenerationStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class SenseNovaU1Pipeline(ComposedPipelineBase):
|
||||
pipeline_name = "SenseNovaU1Pipeline"
|
||||
pipeline_config_cls = SenseNovaU1PipelineConfig
|
||||
sampling_params_cls = SenseNovaU1SamplingParams
|
||||
_required_config_modules: list[str] = []
|
||||
|
||||
def validate_disagg_role(self, role: RoleType) -> None:
|
||||
if role != RoleType.MONOLITHIC:
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline only supports monolithic deployment; "
|
||||
f"disaggregation role {role.value!r} is not supported"
|
||||
)
|
||||
|
||||
def load_modules(
|
||||
self,
|
||||
server_args: ServerArgs,
|
||||
loaded_modules: dict[str, torch.nn.Module] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if loaded_modules is not None and {"model", "tokenizer"} <= set(loaded_modules):
|
||||
return loaded_modules
|
||||
|
||||
if server_args.num_gpus != 1:
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline currently supports num_gpus=1. "
|
||||
"Native tensor/pipeline parallelism is not implemented yet."
|
||||
)
|
||||
modules = load_model_and_tokenizer(self.model_path, server_args)
|
||||
logger.info("Loaded SenseNova-U1 model from %s", self.model_path)
|
||||
return modules
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs) -> None:
|
||||
del server_args
|
||||
self.add_stage(InputValidationStage())
|
||||
self.add_stage(
|
||||
SenseNovaU1GenerationStage(
|
||||
model=self.get_module("model"),
|
||||
tokenizer=self.get_module("tokenizer"),
|
||||
),
|
||||
"sensenova_u1_generation_stage",
|
||||
)
|
||||
|
||||
|
||||
EntryClass = SenseNovaU1Pipeline
|
||||
@@ -0,0 +1,135 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Shared helpers for expanding one request into per-output requests."""
|
||||
|
||||
import os
|
||||
from copy import copy, deepcopy
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
from sglang.srt.observability.trace import TraceReqContext
|
||||
|
||||
|
||||
def normalize_output_seeds(
|
||||
seed: int | list[int],
|
||||
*,
|
||||
num_outputs_per_prompt: int,
|
||||
num_prompts: int = 1,
|
||||
prompt_index: int = 0,
|
||||
) -> list[int]:
|
||||
"""Return the seeds for one prompt's outputs."""
|
||||
if num_outputs_per_prompt <= 0:
|
||||
raise ValueError(
|
||||
f"num_outputs_per_prompt must be positive, got {num_outputs_per_prompt}"
|
||||
)
|
||||
|
||||
if isinstance(seed, list):
|
||||
seeds = [int(item) for item in seed]
|
||||
total_outputs = num_outputs_per_prompt * num_prompts
|
||||
if len(seeds) == num_outputs_per_prompt:
|
||||
return seeds
|
||||
if len(seeds) == total_outputs:
|
||||
start = prompt_index * num_outputs_per_prompt
|
||||
return seeds[start : start + num_outputs_per_prompt]
|
||||
raise ValueError(
|
||||
"seed list length must match num_outputs_per_prompt "
|
||||
f"({num_outputs_per_prompt}) or total outputs ({total_outputs}), "
|
||||
f"got {len(seeds)}"
|
||||
)
|
||||
|
||||
base_seed = int(seed)
|
||||
return [base_seed + i for i in range(num_outputs_per_prompt)]
|
||||
|
||||
|
||||
def _with_output_index_suffix(output_file_name: str, output_index: int) -> str:
|
||||
base, ext = os.path.splitext(output_file_name)
|
||||
return f"{base}_{output_index}{ext}"
|
||||
|
||||
|
||||
def _trace_ctx_for_output(
|
||||
req: Req,
|
||||
request_id: str | None,
|
||||
output_index: int,
|
||||
*,
|
||||
reuse_parent_trace_ctx: bool,
|
||||
):
|
||||
trace_ctx = req.trace_ctx
|
||||
if reuse_parent_trace_ctx or output_index == 0 or not trace_ctx.tracing_enable:
|
||||
return trace_ctx
|
||||
|
||||
output_trace_ctx = TraceReqContext(
|
||||
rid=request_id,
|
||||
module_name=trace_ctx.module_name,
|
||||
external_trace_header=trace_ctx.external_trace_header,
|
||||
)
|
||||
output_trace_ctx.trace_req_start()
|
||||
return output_trace_ctx
|
||||
|
||||
|
||||
def expand_request_outputs(
|
||||
req: Req,
|
||||
*,
|
||||
num_prompts: int = 1,
|
||||
prompt_index: int = 0,
|
||||
reuse_parent_trace_ctx: bool = False,
|
||||
preserve_parent_metrics: bool = False,
|
||||
) -> list[Req]:
|
||||
"""Expand one request into independent per-output requests.
|
||||
|
||||
Entry points use separate trace roots because they own and finish every
|
||||
expanded request scope. Sequential pipeline execution reuses the parent
|
||||
context because the executor owns only the parent request trace lifecycle.
|
||||
"""
|
||||
num_outputs = int(req.num_outputs_per_prompt)
|
||||
seeds = normalize_output_seeds(
|
||||
req.seed,
|
||||
num_outputs_per_prompt=num_outputs,
|
||||
num_prompts=num_prompts,
|
||||
prompt_index=prompt_index,
|
||||
)
|
||||
|
||||
if num_outputs == 1:
|
||||
req.seed = seeds[0]
|
||||
req.seeds = None
|
||||
req.generator = None
|
||||
req.sampling_params.refresh_request_extra_after_output_expansion(req)
|
||||
return [req]
|
||||
|
||||
expanded: list[Req] = []
|
||||
for output_index, seed in enumerate(seeds):
|
||||
output_request_id = (
|
||||
f"{req.request_id}:{output_index}" if req.request_id is not None else None
|
||||
)
|
||||
output_metrics = deepcopy(req.metrics) if preserve_parent_metrics else None
|
||||
output_req = copy(req)
|
||||
output_req.sampling_params = copy(req.sampling_params)
|
||||
output_req.extra = dict(req.extra)
|
||||
output_req.condition_inputs = dict(req.condition_inputs)
|
||||
output_req.trace_ctx = _trace_ctx_for_output(
|
||||
req,
|
||||
output_request_id,
|
||||
output_index,
|
||||
reuse_parent_trace_ctx=reuse_parent_trace_ctx,
|
||||
)
|
||||
output_req.seed = seed
|
||||
output_req.num_outputs_per_prompt = 1
|
||||
output_req.seeds = None
|
||||
output_req.generator = None
|
||||
output_req.extra["parent_request_id"] = req.request_id
|
||||
output_req.extra["output_index"] = output_index
|
||||
|
||||
if output_request_id is not None:
|
||||
output_req.request_id = output_request_id
|
||||
|
||||
if req.output_file_name:
|
||||
output_req.output_file_name = _with_output_index_suffix(
|
||||
req.output_file_name, output_index
|
||||
)
|
||||
output_req.sampling_params.refresh_request_extra_after_output_expansion(
|
||||
output_req
|
||||
)
|
||||
output_req.validate()
|
||||
if output_metrics is not None:
|
||||
output_req.metrics = output_metrics
|
||||
output_req.metrics.request_id = output_req.request_id
|
||||
expanded.append(output_req)
|
||||
|
||||
return expanded
|
||||
@@ -5,6 +5,8 @@
|
||||
Input validation stage for diffusion pipelines.
|
||||
"""
|
||||
|
||||
from typing import Iterator
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision.transforms.functional as TF
|
||||
@@ -13,6 +15,9 @@ from PIL import Image
|
||||
from sglang.multimodal_gen.configs.pipeline_configs import WanI2V480PConfig
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.request_utils import (
|
||||
expand_request_outputs,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
|
||||
@@ -48,6 +53,24 @@ class InputValidationStage(PipelineStage):
|
||||
super().__init__()
|
||||
self.vae_image_processor = vae_image_processor
|
||||
|
||||
def iter_sequential_requests(
|
||||
self, batch: Req, server_args: ServerArgs
|
||||
) -> Iterator[Req]:
|
||||
if not server_args.pipeline_config.supports_sequential_multi_output_inference():
|
||||
return iter((batch,))
|
||||
|
||||
num_outputs = max(1, int(batch.num_outputs_per_prompt or 1))
|
||||
if num_outputs == 1:
|
||||
return iter((batch,))
|
||||
|
||||
return iter(
|
||||
expand_request_outputs(
|
||||
batch,
|
||||
reuse_parent_trace_ctx=True,
|
||||
preserve_parent_metrics=True,
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _calculate_dimensions_from_area(
|
||||
max_area: float, aspect_ratio: float, mod_value: int
|
||||
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Model-specific helpers and stages for SenseNova-U1."""
|
||||
|
||||
from .stages import SenseNovaU1GenerationStage
|
||||
|
||||
__all__ = ["SenseNovaU1GenerationStage"]
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Pipeline lifecycle stages for the native SenseNova-U1 implementation."""
|
||||
|
||||
from .generation import SenseNovaU1GenerationStage
|
||||
|
||||
__all__ = ["SenseNovaU1GenerationStage"]
|
||||
+104
@@ -0,0 +1,104 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.configs.sensenova_u1 import (
|
||||
DEFAULT_CFG_INTERVAL,
|
||||
DEFAULT_CFG_NORM,
|
||||
DEFAULT_ENABLE_TIMESTEP_SHIFT,
|
||||
DEFAULT_T_EPS,
|
||||
DEFAULT_THINK_MODE,
|
||||
DEFAULT_TIMESTEP_SHIFT,
|
||||
SENSENOVA_U1_REQUEST_EXTRA_KEY,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
|
||||
OutputBatch,
|
||||
Req,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
|
||||
|
||||
def _denorm_sensenova_output(x: torch.Tensor) -> torch.Tensor:
|
||||
"""Convert SenseNova's normalized image tensor from [-1, 1] to [0, 1]."""
|
||||
return ((x.float() + 1.0) * 0.5).clamp(0, 1)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SenseNovaU1GenerationOptions:
|
||||
cfg_norm: str = DEFAULT_CFG_NORM
|
||||
timestep_shift: float = DEFAULT_TIMESTEP_SHIFT
|
||||
enable_timestep_shift: bool = DEFAULT_ENABLE_TIMESTEP_SHIFT
|
||||
cfg_interval: tuple[float, float] = DEFAULT_CFG_INTERVAL
|
||||
t_eps: float = DEFAULT_T_EPS
|
||||
think_mode: bool = DEFAULT_THINK_MODE
|
||||
|
||||
@classmethod
|
||||
def from_batch(cls, batch: Req) -> SenseNovaU1GenerationOptions:
|
||||
extra = batch.extra.get(SENSENOVA_U1_REQUEST_EXTRA_KEY, {})
|
||||
return cls(
|
||||
cfg_norm=extra.get("cfg_norm", DEFAULT_CFG_NORM),
|
||||
timestep_shift=float(extra.get("timestep_shift", DEFAULT_TIMESTEP_SHIFT)),
|
||||
enable_timestep_shift=bool(
|
||||
extra.get("enable_timestep_shift", DEFAULT_ENABLE_TIMESTEP_SHIFT)
|
||||
),
|
||||
cfg_interval=tuple(extra.get("cfg_interval", DEFAULT_CFG_INTERVAL)),
|
||||
t_eps=float(extra.get("t_eps", DEFAULT_T_EPS)),
|
||||
think_mode=bool(extra.get("think_mode", DEFAULT_THINK_MODE)),
|
||||
)
|
||||
|
||||
|
||||
class SenseNovaU1GenerationStage(PipelineStage):
|
||||
def __init__(self, model: torch.nn.Module, tokenizer: Any):
|
||||
super().__init__()
|
||||
self.model = model
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
@property
|
||||
def role_affinity(self) -> RoleType:
|
||||
return RoleType.DENOISER
|
||||
|
||||
def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch:
|
||||
del server_args
|
||||
options = SenseNovaU1GenerationOptions.from_batch(batch)
|
||||
if int(batch.num_outputs_per_prompt) != 1:
|
||||
raise ValueError(
|
||||
"SenseNova-U1 expects output expansion before generation; "
|
||||
f"got num_outputs_per_prompt={batch.num_outputs_per_prompt}."
|
||||
)
|
||||
seed = batch.seed[0] if isinstance(batch.seed, list) else int(batch.seed)
|
||||
|
||||
out = self.model.t2i_generate(
|
||||
self.tokenizer,
|
||||
batch.prompt,
|
||||
image_size=(int(batch.width), int(batch.height)),
|
||||
cfg_scale=float(batch.guidance_scale),
|
||||
cfg_norm=options.cfg_norm,
|
||||
timestep_shift=options.timestep_shift,
|
||||
enable_timestep_shift=options.enable_timestep_shift,
|
||||
cfg_interval=options.cfg_interval,
|
||||
num_steps=int(batch.num_inference_steps),
|
||||
batch_size=1,
|
||||
t_eps=options.t_eps,
|
||||
think_mode=options.think_mode,
|
||||
seed=seed,
|
||||
)
|
||||
think_text = None
|
||||
if options.think_mode:
|
||||
images, think_text = out
|
||||
else:
|
||||
images = out
|
||||
|
||||
images = _denorm_sensenova_output(images)
|
||||
samples = [sample.contiguous() for sample in images]
|
||||
usage = {"think_text": think_text} if think_text is not None else None
|
||||
return OutputBatch(
|
||||
output=samples,
|
||||
metrics=batch.metrics,
|
||||
usage=usage,
|
||||
)
|
||||
@@ -112,6 +112,27 @@ def is_ltx2_two_stage_pipeline_name(pipeline_class_name: str | None) -> bool:
|
||||
return pipeline_class_name in LTX2_TWO_STAGE_PIPELINE_NAMES
|
||||
|
||||
|
||||
def _infer_direct_constructor_explicit_arg_names(server_args) -> set[str]:
|
||||
explicit_arg_names: set[str] = set()
|
||||
for attr in dataclasses.fields(server_args):
|
||||
if not attr.init or attr.name == "_explicit_arg_names":
|
||||
continue
|
||||
|
||||
value = getattr(server_args, attr.name)
|
||||
if attr.default is not dataclasses.MISSING:
|
||||
default = attr.default
|
||||
elif attr.default_factory is not dataclasses.MISSING:
|
||||
default = attr.default_factory()
|
||||
else:
|
||||
explicit_arg_names.add(attr.name)
|
||||
continue
|
||||
|
||||
if value != default:
|
||||
explicit_arg_names.add(attr.name)
|
||||
|
||||
return explicit_arg_names
|
||||
|
||||
|
||||
def _normalize_component_precisions(value: object) -> dict[str, str]:
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError("component_precisions must be a mapping")
|
||||
@@ -1864,11 +1885,17 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
raise ValueError(f"Could not parse attention backend config: {config_str}")
|
||||
|
||||
def __post_init__(self):
|
||||
if not self._explicit_arg_names:
|
||||
self._explicit_arg_names = _infer_direct_constructor_explicit_arg_names(
|
||||
self
|
||||
)
|
||||
|
||||
# configure logger before use
|
||||
configure_logger(server_args=self)
|
||||
|
||||
component_paths: dict[str, str] = {}
|
||||
component_weights_paths = dict(self.component_weights_paths)
|
||||
migrated_component_weight_path = False
|
||||
for component, path in self.component_paths.items():
|
||||
if not is_explicit_weight_file_reference(path):
|
||||
component_paths[component] = path
|
||||
@@ -1880,6 +1907,11 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
f"{existing!r} and {path!r}"
|
||||
)
|
||||
component_weights_paths[component] = path
|
||||
migrated_component_weight_path = True
|
||||
if migrated_component_weight_path and self.is_arg_explicitly_set(
|
||||
"component_paths"
|
||||
):
|
||||
self._explicit_arg_names.add("component_weights_paths")
|
||||
self.component_paths = component_paths
|
||||
self.component_weights_paths = component_weights_paths
|
||||
self.component_precisions = _normalize_component_precisions(
|
||||
|
||||
@@ -10,6 +10,9 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import (
|
||||
InputValidationStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import (
|
||||
LatentPreparationStage,
|
||||
)
|
||||
@@ -131,6 +134,55 @@ class TestMultiOutputGrouping(unittest.TestCase):
|
||||
["rid:0", "rid:1"],
|
||||
)
|
||||
|
||||
def test_sequential_stage_matches_entrypoint_output_expansion(self):
|
||||
def make_req():
|
||||
req = Req(
|
||||
sampling_params=SamplingParams(
|
||||
request_id="rid",
|
||||
prompt="p",
|
||||
output_path="/tmp",
|
||||
output_file_name="image.png",
|
||||
num_outputs_per_prompt=2,
|
||||
seed=[100, 101],
|
||||
)
|
||||
)
|
||||
return req
|
||||
|
||||
entrypoint_outputs = expand_request_outputs(make_req())
|
||||
sequential_parent = make_req()
|
||||
sequential_outputs = list(
|
||||
InputValidationStage().iter_sequential_requests(
|
||||
sequential_parent,
|
||||
SimpleNamespace(
|
||||
pipeline_config=SimpleNamespace(
|
||||
supports_sequential_multi_output_inference=lambda: True
|
||||
)
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
def expansion_signature(req):
|
||||
return (
|
||||
req.request_id,
|
||||
req.seed,
|
||||
req.num_outputs_per_prompt,
|
||||
req.output_file_name,
|
||||
req.extra["parent_request_id"],
|
||||
req.extra["output_index"],
|
||||
req.metrics.request_id,
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
[expansion_signature(req) for req in sequential_outputs],
|
||||
[expansion_signature(req) for req in entrypoint_outputs],
|
||||
)
|
||||
self.assertTrue(
|
||||
all(
|
||||
req.trace_ctx is sequential_parent.trace_ctx
|
||||
for req in sequential_outputs
|
||||
)
|
||||
)
|
||||
|
||||
def test_split_batched_latents_uses_original_batched_tensor(self):
|
||||
stage = LatentPreparationStage.__new__(LatentPreparationStage)
|
||||
src = Req(sampling_params=SamplingParams(prompt="p"))
|
||||
|
||||
@@ -0,0 +1,851 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import asyncio
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.sensenova_u1 import (
|
||||
SenseNovaU1PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.sampling_params import DataType
|
||||
from sglang.multimodal_gen.configs.sample.sensenova_u1 import (
|
||||
SenseNovaU1SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sensenova_u1 import (
|
||||
SENSENOVA_U1_REQUEST_EXTRA_KEY,
|
||||
)
|
||||
from sglang.multimodal_gen.registry import (
|
||||
_get_config_info,
|
||||
get_model_info,
|
||||
get_non_diffusers_pipeline_name,
|
||||
is_registered_diffusion_model_path,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
|
||||
process_generation_batch,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker
|
||||
from sglang.multimodal_gen.runtime.models.sensenova_u1.neo_unify.configuration_neo_vit import (
|
||||
NEOVisionConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.sensenova_u1.neo_unify.conversation import (
|
||||
get_conv_template,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.sensenova_u1.neo_unify.modeling_neo_chat import (
|
||||
_randn_with_seed,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import (
|
||||
PipelineExecutor,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import (
|
||||
InputValidationStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sensenova_u1 import (
|
||||
SenseNovaU1GenerationStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.multimodal_gen.runtime.server_args.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.perf_logger import MemorySnapshot
|
||||
|
||||
|
||||
class _FakeSenseNovaModel:
|
||||
def __init__(self):
|
||||
self.call_kwargs = None
|
||||
|
||||
def t2i_generate(self, tokenizer, prompt, **kwargs):
|
||||
self.call_kwargs = {"tokenizer": tokenizer, "prompt": prompt, **kwargs}
|
||||
return torch.tensor(
|
||||
[
|
||||
[
|
||||
[[-1.0, 0.0], [0.5, 1.0]],
|
||||
[[-1.0, 0.0], [0.5, 1.0]],
|
||||
[[-1.0, 0.0], [0.5, 1.0]],
|
||||
]
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class _RecordingTraceContext:
|
||||
tracing_enable = True
|
||||
|
||||
def __init__(self):
|
||||
self.finish_count = 0
|
||||
self.started_slices = []
|
||||
self.finished_slices = []
|
||||
|
||||
def trace_req_finish(self):
|
||||
self.finish_count += 1
|
||||
|
||||
def trace_slice_start(self, name, level=0):
|
||||
self.started_slices.append((name, level))
|
||||
|
||||
def trace_slice_end(self, name, level=0, **kwargs):
|
||||
self.finished_slices.append((name, level))
|
||||
|
||||
|
||||
class _SequentialTestExecutor(PipelineExecutor):
|
||||
def __init__(self, server_args, *, fail=False, fail_request_ids=None):
|
||||
super().__init__(server_args)
|
||||
self.fail = fail
|
||||
self.fail_request_ids = set(fail_request_ids or [])
|
||||
self.executed_requests = []
|
||||
|
||||
def execute_group(self, stages, batches, server_args):
|
||||
for batch in batches:
|
||||
batch.metrics.record_stage("InputValidationStage", 0.125)
|
||||
batch.metrics.record_memory_snapshot(
|
||||
"after_validation",
|
||||
MemorySnapshot(
|
||||
allocated_mb=100.0,
|
||||
reserved_mb=200.0,
|
||||
peak_allocated_mb=300.0,
|
||||
peak_reserved_mb=400.0,
|
||||
),
|
||||
)
|
||||
return batches
|
||||
|
||||
def execute(self, stages, batch, server_args):
|
||||
self.executed_requests.append(batch)
|
||||
if self.fail or batch.request_id in self.fail_request_ids:
|
||||
raise RuntimeError(f"generation failed for {batch.request_id}")
|
||||
return OutputBatch(
|
||||
output_file_paths=[batch.output_file_name],
|
||||
metrics=batch.metrics,
|
||||
)
|
||||
|
||||
|
||||
class _SequentialTestPipeline:
|
||||
def __init__(self, server_args, *, fail=False, fail_request_ids=None):
|
||||
self.input_stage = InputValidationStage()
|
||||
self.executor = _SequentialTestExecutor(
|
||||
server_args,
|
||||
fail=fail,
|
||||
fail_request_ids=fail_request_ids,
|
||||
)
|
||||
|
||||
def forward_batch_sequentially(self, batches, server_args):
|
||||
return self.executor.execute_group_sequentially(
|
||||
[self.input_stage, object()], batches, server_args
|
||||
)
|
||||
|
||||
|
||||
class _WorkerBackedSchedulerClient:
|
||||
def __init__(self, worker):
|
||||
self.worker = worker
|
||||
|
||||
async def forward(self, batches):
|
||||
return next(self.worker.execute_forward_sequentially(batches))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("template_name", "expected_system_message"),
|
||||
[
|
||||
(
|
||||
"Hermes-2",
|
||||
"\u4f60\u662f\u7531\u4e0a\u6d77\u4eba\u5de5\u667a\u80fd\u5b9e\u9a8c\u5ba4\u8054\u5408\u5546\u6c64\u79d1\u6280\u5f00\u53d1\u7684\u4e66\u751f\u591a\u6a21\u6001\u5927\u6a21\u578b\uff0c\u82f1\u6587\u540d\u53ebInternVL, \u662f\u4e00\u4e2a\u6709\u7528\u65e0\u5bb3\u7684\u4eba\u5de5\u667a\u80fd\u52a9\u624b\u3002",
|
||||
),
|
||||
(
|
||||
"internlm2-chat",
|
||||
"\u4f60\u662f\u7531\u4e0a\u6d77\u4eba\u5de5\u667a\u80fd\u5b9e\u9a8c\u5ba4\u8054\u5408\u5546\u6c64\u79d1\u6280\u5f00\u53d1\u7684\u4e66\u751f\u591a\u6a21\u6001\u5927\u6a21\u578b\uff0c\u82f1\u6587\u540d\u53ebInternVL, \u662f\u4e00\u4e2a\u6709\u7528\u65e0\u5bb3\u7684\u4eba\u5de5\u667a\u80fd\u52a9\u624b\u3002",
|
||||
),
|
||||
(
|
||||
"phi3-chat",
|
||||
"\u4f60\u662f\u7531\u4e0a\u6d77\u4eba\u5de5\u667a\u80fd\u5b9e\u9a8c\u5ba4\u8054\u5408\u5546\u6c64\u79d1\u6280\u5f00\u53d1\u7684\u4e66\u751f\u591a\u6a21\u6001\u5927\u6a21\u578b\uff0c\u82f1\u6587\u540d\u53ebInternVL, \u662f\u4e00\u4e2a\u6709\u7528\u65e0\u5bb3\u7684\u4eba\u5de5\u667a\u80fd\u52a9\u624b\u3002",
|
||||
),
|
||||
(
|
||||
"internvl2_5",
|
||||
"\u4f60\u662f\u4e66\u751f\xb7\u4e07\u8c61\uff0c\u82f1\u6587\u540d\u662fInternVL\uff0c\u662f\u7531\u4e0a\u6d77\u4eba\u5de5\u667a\u80fd\u5b9e\u9a8c\u5ba4\u3001\u6e05\u534e\u5927\u5b66\u53ca\u591a\u5bb6\u5408\u4f5c\u5355\u4f4d\u8054\u5408\u5f00\u53d1\u7684\u591a\u6a21\u6001\u5927\u8bed\u8a00\u6a21\u578b\u3002",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_sensenova_u1_conversation_preserves_upstream_system_prompt(
|
||||
template_name, expected_system_message
|
||||
):
|
||||
assert get_conv_template(template_name).system_message == expected_system_message
|
||||
|
||||
|
||||
def _force_generator_fallback(monkeypatch, device_type):
|
||||
original_generator = torch.Generator
|
||||
|
||||
def unsupported_device_generator(device="cpu"):
|
||||
if torch.device(device).type == device_type:
|
||||
raise RuntimeError(f"Generator is unsupported on {device_type}")
|
||||
return original_generator(device)
|
||||
|
||||
monkeypatch.setattr(torch, "Generator", unsupported_device_generator)
|
||||
|
||||
|
||||
def test_sensenova_u1_randn_fallback_preserves_cpu_rng(monkeypatch):
|
||||
_force_generator_fallback(monkeypatch, "cpu")
|
||||
rng_state = torch.get_rng_state().clone()
|
||||
|
||||
first = _randn_with_seed((2, 3), device="cpu", dtype=torch.float32, seed=17)
|
||||
second = _randn_with_seed((2, 3), device="cpu", dtype=torch.float32, seed=17)
|
||||
|
||||
assert torch.equal(first, second)
|
||||
assert torch.equal(torch.get_rng_state(), rng_state)
|
||||
|
||||
|
||||
def test_sensenova_u1_randn_fallback_preserves_device_rng(monkeypatch):
|
||||
device_type = current_platform.device_type
|
||||
if not device_type or device_type == "cpu":
|
||||
pytest.skip("No accelerator is available")
|
||||
|
||||
device = torch.device(device_type, 0)
|
||||
device_module = torch.get_device_module(device)
|
||||
if not device_module.is_available():
|
||||
pytest.skip(f"{device_type} is not available")
|
||||
|
||||
_force_generator_fallback(monkeypatch, device_type)
|
||||
cpu_rng_state = torch.get_rng_state().clone()
|
||||
device_rng_state = device_module.get_rng_state(device).clone()
|
||||
|
||||
first = _randn_with_seed((2, 3), device=device, dtype=torch.float32, seed=17)
|
||||
second = _randn_with_seed((2, 3), device=device, dtype=torch.float32, seed=17)
|
||||
|
||||
assert torch.equal(first, second)
|
||||
assert torch.equal(torch.get_rng_state(), cpu_rng_state)
|
||||
assert torch.equal(device_module.get_rng_state(device), device_rng_state)
|
||||
|
||||
|
||||
def test_sensenova_u1_registry_resolves_local_and_hf_paths(tmp_path):
|
||||
_get_config_info.cache_clear()
|
||||
get_model_info.cache_clear()
|
||||
|
||||
local_path = tmp_path / "checkpoint-revision-abc123"
|
||||
local_path.mkdir()
|
||||
(local_path / "config.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"architectures": ["NEOChatModel"],
|
||||
"model_type": "neo_chat",
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
assert is_registered_diffusion_model_path(str(local_path))
|
||||
assert get_non_diffusers_pipeline_name(str(local_path)) == "SenseNovaU1Pipeline"
|
||||
|
||||
local_model_info = get_model_info(str(local_path))
|
||||
assert local_model_info is not None
|
||||
assert local_model_info.pipeline_config_cls is SenseNovaU1PipelineConfig
|
||||
assert local_model_info.sampling_param_cls is SenseNovaU1SamplingParams
|
||||
|
||||
model_info = get_model_info("sensenova/SenseNova-U1.5-8B-MoT")
|
||||
assert model_info is not None
|
||||
assert model_info.pipeline_config_cls is SenseNovaU1PipelineConfig
|
||||
assert model_info.sampling_param_cls is SenseNovaU1SamplingParams
|
||||
|
||||
modelscope_id = "SenseNova/SenseNova-U1.5-8B-MoT"
|
||||
assert is_registered_diffusion_model_path(modelscope_id)
|
||||
assert get_non_diffusers_pipeline_name(modelscope_id) == "SenseNovaU1Pipeline"
|
||||
assert get_model_info(modelscope_id) is not None
|
||||
get_model_info.cache_clear()
|
||||
|
||||
|
||||
def test_sensenova_u1_registry_requires_exact_hub_id(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"sglang.multimodal_gen.registry.maybe_download_model_index",
|
||||
lambda _: {},
|
||||
)
|
||||
_get_config_info.cache_clear()
|
||||
get_model_info.cache_clear()
|
||||
|
||||
unrelated_repo = "acme/SenseNova-U1.5-8B-MoT"
|
||||
assert not is_registered_diffusion_model_path(unrelated_repo)
|
||||
assert get_non_diffusers_pipeline_name(unrelated_repo) is None
|
||||
assert _get_config_info(unrelated_repo) is None
|
||||
|
||||
_get_config_info.cache_clear()
|
||||
get_model_info.cache_clear()
|
||||
|
||||
|
||||
def test_sensenova_u1_registry_does_not_route_lora_only_repositories(tmp_path):
|
||||
lora_repo = "sensenova/SenseNova-U1.5-8B-MoT-LoRA"
|
||||
lora_path = tmp_path / "SenseNova-U1.5-8B-MoT-LoRA"
|
||||
lora_path.mkdir()
|
||||
(lora_path / "adapter_config.json").write_text("{}")
|
||||
|
||||
assert get_non_diffusers_pipeline_name(lora_repo) is None
|
||||
assert get_non_diffusers_pipeline_name(str(lora_path)) is None
|
||||
assert not is_registered_diffusion_model_path(lora_repo)
|
||||
assert not is_registered_diffusion_model_path(str(lora_path))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["auto", "sglang", "diffusers"])
|
||||
def test_sensenova_u1_known_adapter_only_repo_rejected_before_backend_resolution(
|
||||
monkeypatch, backend
|
||||
):
|
||||
def fail_model_index_download(_):
|
||||
raise AssertionError("adapter-only repo should not download model_index")
|
||||
|
||||
def fail_diffusers_resolution(**_kwargs):
|
||||
raise AssertionError("adapter-only repo should not resolve diffusers info")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"sglang.multimodal_gen.registry.maybe_download_model_index",
|
||||
fail_model_index_download,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"sglang.multimodal_gen.registry._get_diffusers_model_info",
|
||||
fail_diffusers_resolution,
|
||||
)
|
||||
get_model_info.cache_clear()
|
||||
|
||||
loras_repo = "sensenova/SenseNova-U1.5-8B-MoT-LoRAs"
|
||||
|
||||
assert get_non_diffusers_pipeline_name(loras_repo) is None
|
||||
assert get_model_info(loras_repo, backend=backend) is None
|
||||
get_model_info.cache_clear()
|
||||
|
||||
|
||||
def test_sensenova_u1_sampling_params_keep_private_defaults_internal():
|
||||
params = SenseNovaU1SamplingParams(prompt="hello", width=2304, height=4096)
|
||||
|
||||
assert params.guidance_scale == 4.0
|
||||
assert params.num_inference_steps == 50
|
||||
assert params.num_outputs_per_prompt == 1
|
||||
assert params.cfg_norm == "none"
|
||||
assert params.timestep_shift == 3.0
|
||||
|
||||
extra = params.build_request_extra()[SENSENOVA_U1_REQUEST_EXTRA_KEY]
|
||||
assert extra == {
|
||||
"cfg_norm": "none",
|
||||
"timestep_shift": 3.0,
|
||||
"enable_timestep_shift": True,
|
||||
"cfg_interval": (0.0, 1.0),
|
||||
"t_eps": 0.02,
|
||||
"think_mode": False,
|
||||
}
|
||||
|
||||
|
||||
def test_sensenova_u1_rejects_unaligned_resolution():
|
||||
with pytest.raises(ValueError, match="divisible by 32"):
|
||||
SenseNovaU1SamplingParams(width=2160, height=3840)
|
||||
|
||||
|
||||
def test_sensenova_u1_accepts_openai_image_api_num_frames():
|
||||
params = SenseNovaU1SamplingParams(
|
||||
prompt="hello",
|
||||
width=2048,
|
||||
height=2048,
|
||||
num_frames=1,
|
||||
)
|
||||
|
||||
assert params.num_frames == 1
|
||||
assert params.data_type == DataType.IMAGE
|
||||
|
||||
|
||||
def test_sensenova_u1_scheduler_capabilities():
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
|
||||
assert not config.supports_dynamic_batching()
|
||||
assert config.supports_sequential_multi_output_inference()
|
||||
|
||||
|
||||
def test_sensenova_u1_rejects_multi_gpu_during_arg_validation():
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
|
||||
with pytest.raises(ValueError, match="num_gpus=1"):
|
||||
config.validate_server_args(
|
||||
SimpleNamespace(
|
||||
num_gpus=2,
|
||||
enable_torch_compile=False,
|
||||
lora_path=None,
|
||||
attention_backend=None,
|
||||
component_attention_backends={},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_sensenova_u1_clears_auto_tuned_runtime_defaults():
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
args = SimpleNamespace(
|
||||
num_gpus=1,
|
||||
enable_torch_compile=False,
|
||||
lora_path=None,
|
||||
component_residency={"transformer": "layerwise-offload"},
|
||||
cpu_offload_components=["transformer"],
|
||||
dit_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
image_encoder_cpu_offload=True,
|
||||
vae_cpu_offload=True,
|
||||
dit_layerwise_offload=True,
|
||||
layerwise_offload_components=["transformer"],
|
||||
quantization=None,
|
||||
quantization_ignored_layers=None,
|
||||
transformer_weights_path=None,
|
||||
component_paths={"model": "/tmp/component"},
|
||||
component_weights_paths={"model": "/tmp/model.safetensors"},
|
||||
component_quantizations={},
|
||||
component_quantization_ignored_layers={},
|
||||
component_precisions={},
|
||||
attention_backend="aiter",
|
||||
component_attention_backends={"text_encoder": "torch_sdpa"},
|
||||
attention_backend_config={"foo": "bar"},
|
||||
is_arg_explicitly_set=lambda _name: False,
|
||||
)
|
||||
|
||||
config.validate_server_args(args)
|
||||
|
||||
assert args.component_residency is None
|
||||
assert args.cpu_offload_components is None
|
||||
assert args.dit_cpu_offload is False
|
||||
assert args.text_encoder_cpu_offload is False
|
||||
assert args.image_encoder_cpu_offload is False
|
||||
assert args.vae_cpu_offload is False
|
||||
assert args.dit_layerwise_offload is False
|
||||
assert args.layerwise_offload_components is None
|
||||
assert args.component_paths == {}
|
||||
assert args.component_weights_paths == {}
|
||||
assert args.attention_backend is None
|
||||
assert args.component_attention_backends == {}
|
||||
assert args.attention_backend_config is None
|
||||
|
||||
|
||||
def test_sensenova_u1_allows_explicit_resident_component_residency():
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
args = SimpleNamespace(
|
||||
num_gpus=1,
|
||||
enable_torch_compile=False,
|
||||
lora_path=None,
|
||||
component_residency={"transformer": "resident"},
|
||||
cpu_offload_components=None,
|
||||
dit_cpu_offload=False,
|
||||
text_encoder_cpu_offload=False,
|
||||
image_encoder_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
dit_layerwise_offload=False,
|
||||
layerwise_offload_components=None,
|
||||
quantization=None,
|
||||
quantization_ignored_layers=None,
|
||||
transformer_weights_path=None,
|
||||
component_paths={},
|
||||
component_weights_paths={},
|
||||
component_quantizations={},
|
||||
component_quantization_ignored_layers={},
|
||||
component_precisions={},
|
||||
attention_backend=None,
|
||||
component_attention_backends={},
|
||||
attention_backend_config={},
|
||||
is_arg_explicitly_set=lambda name: name == "component_residency",
|
||||
)
|
||||
|
||||
config.validate_server_args(args)
|
||||
|
||||
assert args.component_residency == {"transformer": "resident"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("override", "expected"),
|
||||
[
|
||||
({"enable_torch_compile": True}, "torch.compile"),
|
||||
({"lora_path": "sensenova/SenseNova-U1.5-8B-MoT-LoRAs"}, "LoRA adapters"),
|
||||
(
|
||||
{"component_residency": {"transformer": "component-offload"}},
|
||||
"component residency offload",
|
||||
),
|
||||
({"cpu_offload_components": ["transformer"]}, "CPU offload"),
|
||||
({"dit_cpu_offload": True}, "DiT CPU offload"),
|
||||
({"text_encoder_cpu_offload": True}, "text encoder CPU offload"),
|
||||
({"image_encoder_cpu_offload": True}, "image encoder CPU offload"),
|
||||
({"vae_cpu_offload": True}, "VAE CPU offload"),
|
||||
({"dit_layerwise_offload": True}, "DiT layerwise offload"),
|
||||
({"layerwise_offload_components": ["transformer"]}, "layerwise offload"),
|
||||
({"quantization": "fp8"}, "quantization"),
|
||||
({"quantization_ignored_layers": ["foo"]}, "quantization ignored layers"),
|
||||
(
|
||||
{"transformer_weights_path": "/tmp/transformer.safetensors"},
|
||||
"pre-quantized transformer weights",
|
||||
),
|
||||
({"component_paths": {"model": "/tmp/component"}}, "component path overrides"),
|
||||
(
|
||||
{"component_weights_paths": {"model": "/tmp/model.safetensors"}},
|
||||
"component weight path overrides",
|
||||
),
|
||||
({"component_quantizations": {"transformer": "fp8"}}, "component quantization"),
|
||||
(
|
||||
{"component_quantization_ignored_layers": {"transformer": ["foo"]}},
|
||||
"component quantization ignored layers",
|
||||
),
|
||||
({"component_precisions": {"transformer": "fp16"}}, "component precision"),
|
||||
({"attention_backend": "fa"}, "custom attention backends"),
|
||||
(
|
||||
{"component_attention_backends": {"text_encoder": "torch_sdpa"}},
|
||||
"component attention backends",
|
||||
),
|
||||
({"attention_backend_config": {"foo": "bar"}}, "attention backend config"),
|
||||
],
|
||||
)
|
||||
def test_sensenova_u1_rejects_unsupported_runtime_modes(override, expected):
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
args = {
|
||||
"num_gpus": 1,
|
||||
"enable_torch_compile": False,
|
||||
"lora_path": None,
|
||||
"component_residency": None,
|
||||
"cpu_offload_components": None,
|
||||
"dit_cpu_offload": None,
|
||||
"text_encoder_cpu_offload": None,
|
||||
"image_encoder_cpu_offload": None,
|
||||
"vae_cpu_offload": False,
|
||||
"dit_layerwise_offload": None,
|
||||
"layerwise_offload_components": None,
|
||||
"quantization": None,
|
||||
"quantization_ignored_layers": None,
|
||||
"transformer_weights_path": None,
|
||||
"component_paths": {},
|
||||
"component_weights_paths": {},
|
||||
"component_quantizations": {},
|
||||
"component_quantization_ignored_layers": {},
|
||||
"component_precisions": {},
|
||||
"attention_backend": None,
|
||||
"component_attention_backends": {},
|
||||
"attention_backend_config": {},
|
||||
}
|
||||
args.update(override)
|
||||
|
||||
with pytest.raises(ValueError, match=expected):
|
||||
config.validate_server_args(SimpleNamespace(**args))
|
||||
|
||||
|
||||
def test_sensenova_u1_rejects_direct_server_args_quantization():
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
|
||||
with pytest.raises(ValueError, match="quantization"):
|
||||
ServerArgs(
|
||||
model_path="sensenova/SenseNova-U1.5-8B-MoT",
|
||||
pipeline_config=config,
|
||||
quantization="fp8",
|
||||
)
|
||||
|
||||
|
||||
def test_sensenova_u1_rejects_file_valued_component_paths(tmp_path):
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
|
||||
with pytest.raises(ValueError, match="component weight path overrides"):
|
||||
ServerArgs(
|
||||
model_path="sensenova/SenseNova-U1.5-8B-MoT",
|
||||
pipeline_config=config,
|
||||
component_paths={"model": str(tmp_path / "model.safetensors")},
|
||||
)
|
||||
|
||||
|
||||
def test_sensenova_u1_vision_config_round_trips_sequence_fields(tmp_path):
|
||||
config = NEOVisionConfig(llm_hidden_size=2048, downsample_ratio=0.5)
|
||||
config.save_pretrained(tmp_path)
|
||||
|
||||
loaded = NEOVisionConfig.from_pretrained(tmp_path)
|
||||
|
||||
assert loaded.llm_hidden_size == (2048,)
|
||||
assert loaded.downsample_ratio == (0.5,)
|
||||
|
||||
|
||||
def test_sensenova_u1_vision_config_normalizes_nested_singletons():
|
||||
config = NEOVisionConfig(llm_hidden_size=[[2048]], downsample_ratio=[[0.5]])
|
||||
|
||||
assert config.llm_hidden_size == (2048,)
|
||||
assert config.downsample_ratio == (0.5,)
|
||||
|
||||
|
||||
def test_sensenova_u1_rejects_video_frame_count():
|
||||
with pytest.raises(ValueError, match="num_frames=1"):
|
||||
SenseNovaU1SamplingParams(width=2048, height=2048, num_frames=2)
|
||||
|
||||
|
||||
def test_sensenova_u1_cli_args_expose_only_sglang_compatible_fields():
|
||||
args = SimpleNamespace(
|
||||
prompt="hello",
|
||||
width=2304,
|
||||
height=4096,
|
||||
guidance_scale=4.5,
|
||||
num_inference_steps=30,
|
||||
num_outputs_per_prompt=2,
|
||||
cfg_norm="global",
|
||||
timestep_shift=9.0,
|
||||
think_mode=True,
|
||||
)
|
||||
|
||||
cli_args = SenseNovaU1SamplingParams.get_cli_args(args)
|
||||
|
||||
assert cli_args["prompt"] == "hello"
|
||||
assert cli_args["width"] == 2304
|
||||
assert cli_args["height"] == 4096
|
||||
assert cli_args["guidance_scale"] == 4.5
|
||||
assert cli_args["num_inference_steps"] == 30
|
||||
assert cli_args["num_outputs_per_prompt"] == 2
|
||||
assert "cfg_norm" not in cli_args
|
||||
assert "timestep_shift" not in cli_args
|
||||
assert "think_mode" not in cli_args
|
||||
|
||||
|
||||
def test_sensenova_u1_generation_stage_uses_sglang_params_and_single_model_batch():
|
||||
sampling = SenseNovaU1SamplingParams(
|
||||
prompt="a mountain lake",
|
||||
width=2304,
|
||||
height=4096,
|
||||
guidance_scale=4.5,
|
||||
num_inference_steps=30,
|
||||
seed=123,
|
||||
)
|
||||
batch = SimpleNamespace(
|
||||
prompt=sampling.prompt,
|
||||
width=sampling.width,
|
||||
height=sampling.height,
|
||||
guidance_scale=sampling.guidance_scale,
|
||||
num_inference_steps=sampling.num_inference_steps,
|
||||
seed=sampling.seed,
|
||||
num_outputs_per_prompt=sampling.num_outputs_per_prompt,
|
||||
extra=sampling.build_request_extra(),
|
||||
metrics=None,
|
||||
)
|
||||
model = _FakeSenseNovaModel()
|
||||
stage = SenseNovaU1GenerationStage(model=model, tokenizer="tok")
|
||||
|
||||
output = stage.forward(batch, server_args=SimpleNamespace())
|
||||
|
||||
assert len(output.output) == 1
|
||||
assert torch.allclose(
|
||||
output.output[0],
|
||||
torch.tensor(
|
||||
[
|
||||
[[0.0, 0.5], [0.75, 1.0]],
|
||||
[[0.0, 0.5], [0.75, 1.0]],
|
||||
[[0.0, 0.5], [0.75, 1.0]],
|
||||
]
|
||||
),
|
||||
)
|
||||
assert model.call_kwargs["tokenizer"] == "tok"
|
||||
assert model.call_kwargs["prompt"] == "a mountain lake"
|
||||
assert model.call_kwargs["image_size"] == (2304, 4096)
|
||||
assert model.call_kwargs["cfg_scale"] == 4.5
|
||||
assert model.call_kwargs["num_steps"] == 30
|
||||
assert model.call_kwargs["batch_size"] == 1
|
||||
assert model.call_kwargs["seed"] == 123
|
||||
|
||||
|
||||
def test_sensenova_u1_multi_output_request_expands_before_generation_stage():
|
||||
sampling = SenseNovaU1SamplingParams(
|
||||
prompt="a mountain lake",
|
||||
width=2304,
|
||||
height=4096,
|
||||
num_outputs_per_prompt=2,
|
||||
)
|
||||
batch = Req(
|
||||
request_id="req-0",
|
||||
prompt=sampling.prompt,
|
||||
width=sampling.width,
|
||||
height=sampling.height,
|
||||
guidance_scale=sampling.guidance_scale,
|
||||
num_inference_steps=sampling.num_inference_steps,
|
||||
seed=42,
|
||||
sampling_params=sampling,
|
||||
extra=sampling.build_request_extra(),
|
||||
output_file_name="sample.png",
|
||||
)
|
||||
server_args = SimpleNamespace(pipeline_config=SenseNovaU1PipelineConfig())
|
||||
input_stage = InputValidationStage()
|
||||
stage = SenseNovaU1GenerationStage(model=_FakeSenseNovaModel(), tokenizer="tok")
|
||||
batch.metrics.record_stage("InputValidationStage", 0.125)
|
||||
batch.metrics.record_memory_snapshot(
|
||||
"after_validation",
|
||||
MemorySnapshot(
|
||||
allocated_mb=100.0,
|
||||
reserved_mb=200.0,
|
||||
peak_allocated_mb=300.0,
|
||||
peak_reserved_mb=400.0,
|
||||
),
|
||||
)
|
||||
|
||||
expanded = list(input_stage.iter_sequential_requests(batch, server_args))
|
||||
|
||||
assert [req.num_outputs_per_prompt for req in expanded] == [1, 1]
|
||||
assert [req.seed for req in expanded] == [42, 43]
|
||||
assert [req.request_id for req in expanded] == ["req-0:0", "req-0:1"]
|
||||
assert [req.output_file_name for req in expanded] == [
|
||||
"sample_0.png",
|
||||
"sample_1.png",
|
||||
]
|
||||
assert [req.metrics.request_id for req in expanded] == ["req-0:0", "req-0:1"]
|
||||
assert all(req.trace_ctx is batch.trace_ctx for req in expanded)
|
||||
assert all(req.metrics is not batch.metrics for req in expanded)
|
||||
assert expanded[0].metrics is not expanded[1].metrics
|
||||
assert all(
|
||||
req.metrics.stages == {"InputValidationStage": 125.0} for req in expanded
|
||||
)
|
||||
assert all(
|
||||
req.metrics.memory_snapshots["after_validation"].peak_reserved_mb == 400.0
|
||||
for req in expanded
|
||||
)
|
||||
assert (
|
||||
expanded[0].metrics.memory_snapshots["after_validation"]
|
||||
is not expanded[1].metrics.memory_snapshots["after_validation"]
|
||||
)
|
||||
|
||||
expanded[0].metrics.record_stage("child-only", 0.5)
|
||||
expanded[0].metrics.memory_snapshots["after_validation"].peak_reserved_mb = 999.0
|
||||
assert "child-only" not in expanded[1].metrics.stages
|
||||
assert "child-only" not in batch.metrics.stages
|
||||
assert (
|
||||
expanded[1].metrics.memory_snapshots["after_validation"].peak_reserved_mb
|
||||
== 400.0
|
||||
)
|
||||
assert batch.metrics.memory_snapshots["after_validation"].peak_reserved_mb == 400.0
|
||||
|
||||
for req in expanded:
|
||||
output = stage.forward(req, server_args=SimpleNamespace())
|
||||
assert len(output.output) == 1
|
||||
|
||||
|
||||
def test_sensenova_u1_multi_output_rejects_short_seed_list():
|
||||
sampling = SenseNovaU1SamplingParams(
|
||||
prompt="a mountain lake",
|
||||
width=2304,
|
||||
height=4096,
|
||||
num_outputs_per_prompt=2,
|
||||
seed=[7],
|
||||
)
|
||||
batch = Req(
|
||||
request_id="req-0",
|
||||
prompt=sampling.prompt,
|
||||
width=sampling.width,
|
||||
height=sampling.height,
|
||||
guidance_scale=sampling.guidance_scale,
|
||||
num_inference_steps=sampling.num_inference_steps,
|
||||
seed=sampling.seed,
|
||||
sampling_params=sampling,
|
||||
extra=sampling.build_request_extra(),
|
||||
output_file_name="sample.png",
|
||||
)
|
||||
server_args = SimpleNamespace(pipeline_config=SenseNovaU1PipelineConfig())
|
||||
|
||||
with pytest.raises(ValueError, match="seed list length"):
|
||||
list(InputValidationStage().iter_sequential_requests(batch, server_args))
|
||||
|
||||
|
||||
def _make_sensenova_u1_sequential_entrypoint(*, fail=False, fail_request_ids=None):
|
||||
sampling = SenseNovaU1SamplingParams(
|
||||
prompt="a mountain lake",
|
||||
width=2304,
|
||||
height=4096,
|
||||
num_outputs_per_prompt=2,
|
||||
save_output=False,
|
||||
suppress_logs=True,
|
||||
)
|
||||
trace_ctx = _RecordingTraceContext()
|
||||
batch = Req(
|
||||
request_id="req-0",
|
||||
prompt=sampling.prompt,
|
||||
width=sampling.width,
|
||||
height=sampling.height,
|
||||
seed=42,
|
||||
sampling_params=sampling,
|
||||
extra=sampling.build_request_extra(),
|
||||
output_file_name="sample.png",
|
||||
trace_ctx=trace_ctx,
|
||||
)
|
||||
server_args = SimpleNamespace(pipeline_config=SenseNovaU1PipelineConfig())
|
||||
pipeline = _SequentialTestPipeline(
|
||||
server_args, fail=fail, fail_request_ids=fail_request_ids
|
||||
)
|
||||
worker = GPUWorker.__new__(GPUWorker)
|
||||
worker.pipeline = pipeline
|
||||
worker.server_args = server_args
|
||||
worker.is_output_rank = True
|
||||
worker._runtime_peak_reserved_mb = 0.0
|
||||
worker._release_warmup_pool_before_serving = False
|
||||
worker._realtime_sessions = SimpleNamespace(attach=lambda _req: None)
|
||||
return batch, trace_ctx, pipeline.executor, _WorkerBackedSchedulerClient(worker)
|
||||
|
||||
|
||||
def _force_cpu_entrypoint(monkeypatch):
|
||||
monkeypatch.setattr(current_platform, "is_cpu", lambda: True)
|
||||
monkeypatch.setattr(current_platform, "is_mps", lambda: False)
|
||||
monkeypatch.setattr(current_platform, "is_npu", lambda: False)
|
||||
monkeypatch.setattr(
|
||||
"sglang.multimodal_gen.runtime.entrypoints.openai.utils.get_global_server_args",
|
||||
lambda: SimpleNamespace(batching_max_size=1),
|
||||
)
|
||||
|
||||
|
||||
def test_sensenova_u1_multi_output_entrypoint_success(monkeypatch):
|
||||
_force_cpu_entrypoint(monkeypatch)
|
||||
batch, trace_ctx, executor, scheduler_client = (
|
||||
_make_sensenova_u1_sequential_entrypoint()
|
||||
)
|
||||
|
||||
paths, result = asyncio.run(process_generation_batch(scheduler_client, batch))
|
||||
|
||||
assert paths == ["sample_0.png", "sample_1.png"]
|
||||
assert result.error is None
|
||||
assert [req.request_id for req in executor.executed_requests] == [
|
||||
"req-0:0",
|
||||
"req-0:1",
|
||||
]
|
||||
assert [req.seed for req in executor.executed_requests] == [42, 43]
|
||||
assert result.metrics_list is not None
|
||||
assert [metrics.request_id for metrics in result.metrics_list] == [
|
||||
"req-0:0",
|
||||
"req-0:1",
|
||||
]
|
||||
assert all(
|
||||
"InputValidationStage" in metrics.stages
|
||||
and "PipelineExecutor.sequential_wait" in metrics.stages
|
||||
and metrics.memory_snapshots["after_validation"].peak_reserved_mb == 400.0
|
||||
for metrics in result.metrics_list
|
||||
)
|
||||
assert all(req.trace_ctx is trace_ctx for req in executor.executed_requests)
|
||||
assert trace_ctx.started_slices == [("gpu_forward", 2)]
|
||||
assert trace_ctx.finished_slices == [("gpu_forward", 2)]
|
||||
assert trace_ctx.finish_count == 1
|
||||
|
||||
|
||||
def test_sensenova_u1_multi_output_entrypoint_failure(monkeypatch):
|
||||
_force_cpu_entrypoint(monkeypatch)
|
||||
batch, trace_ctx, executor, scheduler_client = (
|
||||
_make_sensenova_u1_sequential_entrypoint(fail=True)
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="generation failed for req-0:0"):
|
||||
asyncio.run(process_generation_batch(scheduler_client, batch))
|
||||
|
||||
assert [req.request_id for req in executor.executed_requests] == [
|
||||
"req-0:0",
|
||||
"req-0:1",
|
||||
]
|
||||
assert all(
|
||||
"InputValidationStage" in req.metrics.stages
|
||||
and "PipelineExecutor.sequential_wait" in req.metrics.stages
|
||||
and req.metrics.memory_snapshots["after_validation"].peak_reserved_mb == 400.0
|
||||
for req in executor.executed_requests
|
||||
)
|
||||
assert all(req.trace_ctx is trace_ctx for req in executor.executed_requests)
|
||||
assert trace_ctx.started_slices == [("gpu_forward", 2)]
|
||||
assert trace_ctx.finished_slices == [("gpu_forward", 2)]
|
||||
assert trace_ctx.finish_count == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failed_request_id", ["req-0:0", "req-0:1"])
|
||||
def test_sensenova_u1_multi_output_entrypoint_mixed_failure_fails_parent(
|
||||
monkeypatch, failed_request_id
|
||||
):
|
||||
_force_cpu_entrypoint(monkeypatch)
|
||||
batch, trace_ctx, executor, scheduler_client = (
|
||||
_make_sensenova_u1_sequential_entrypoint(fail_request_ids={failed_request_id})
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
RuntimeError, match=f"generation failed for {failed_request_id}"
|
||||
):
|
||||
asyncio.run(process_generation_batch(scheduler_client, batch))
|
||||
|
||||
assert [req.request_id for req in executor.executed_requests] == [
|
||||
"req-0:0",
|
||||
"req-0:1",
|
||||
]
|
||||
assert all(req.trace_ctx is trace_ctx for req in executor.executed_requests)
|
||||
assert trace_ctx.started_slices == [("gpu_forward", 2)]
|
||||
assert trace_ctx.finished_slices == [("gpu_forward", 2)]
|
||||
assert trace_ctx.finish_count == 1
|
||||
Reference in New Issue
Block a user