minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] model: support fal Ideogram V4 Fast and Instant (#31177)

Browse Source
This commit is contained in:
Mick
2026-07-14 19:51:35 +08:00
committed by GitHub
parent ee000f6734
commit 43241b7f3f
17 changed files with 644 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/models/dits/__init__.py
+5 -1
View File
@@ -4,7 +4,10 @@ from sglang.multimodal_gen.configs.models.dits.cosmos3video import Cosmos3VideoC
from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig
from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConfig
from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
from sglang.multimodal_gen.configs.models.dits.ideogram import Ideogram4DiTConfig
from sglang.multimodal_gen.configs.models.dits.ideogram import (
Ideogram4DistilledDiTConfig,
Ideogram4DiTConfig,
)
from sglang.multimodal_gen.configs.models.dits.lingbot_world import (
LingBotWorldVideoConfig,
)
@@ -21,6 +24,7 @@ __all__ = [
"HeliosConfig",
"HunyuanVideoConfig",
"Ideogram4DiTConfig",
"Ideogram4DistilledDiTConfig",
"LingBotWorldVideoConfig",
"LongLive2VideoConfig",
"WanVideoConfig",
python/sglang/multimodal_gen/configs/models/dits/ideogram.py
+30
View File
@@ -19,6 +19,26 @@ class Ideogram4DiTArchConfig(DiTArchConfig):
num_attention_heads: int = 18
num_layers: int = 34
rope_theta: int = 5_000_000
param_names_mapping: dict = field(
default_factory=lambda: {
r"^(layers\.\d+\.attention)\.to_q\.(.*)$": (
r"\1.qkv.\2",
0,
3,
),
r"^(layers\.\d+\.attention)\.to_k\.(.*)$": (
r"\1.qkv.\2",
1,
3,
),
r"^(layers\.\d+\.attention)\.to_v\.(.*)$": (
r"\1.qkv.\2",
2,
3,
),
r"^(layers\.\d+\.attention)\.to_out\.0\.(.*)$": r"\1.o.\2",
}
)
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_layer])
_supported_attention_backends: set[AttentionBackendEnum] = field(
default_factory=lambda: {
@@ -37,3 +57,13 @@ class Ideogram4DiTArchConfig(DiTArchConfig):
class Ideogram4DiTConfig(DiTConfig):
arch_config: DiTArchConfig = field(default_factory=Ideogram4DiTArchConfig)
prefix: str = "ideogram4"
# The official FP8 checkpoint stores row-wise FP8 weights without a
# quantization_config, so its native loader intentionally defaults to the
# dedicated weight-only FP8 linears. Distilled fal checkpoints instead
# store floating-point weights and must use ordinary TP-aware linears.
use_weight_only_fp8_linears: bool = True
@dataclass
class Ideogram4DistilledDiTConfig(Ideogram4DiTConfig):
use_weight_only_fp8_linears: bool = False
python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py
+2
View File
@@ -29,6 +29,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import (
Hunyuan3D2PipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.ideogram import (
Ideogram4DistilledPipelineConfig,
Ideogram4PipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.lingbot_world import (
@@ -64,6 +65,7 @@ __all__ = [
"FastHunyuanConfig",
"Hunyuan3D2PipelineConfig",
"Ideogram4PipelineConfig",
"Ideogram4DistilledPipelineConfig",
"FluxPipelineConfig",
"Flux2PipelineConfig",
"Flux2KleinPipelineConfig",
python/sglang/multimodal_gen/configs/pipeline_configs/ideogram.py
+9 -1
View File
@@ -3,7 +3,10 @@
from dataclasses import dataclass, field
from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig
from sglang.multimodal_gen.configs.models.dits.ideogram import Ideogram4DiTConfig
from sglang.multimodal_gen.configs.models.dits.ideogram import (
Ideogram4DistilledDiTConfig,
Ideogram4DiTConfig,
)
from sglang.multimodal_gen.configs.models.encoders.ideogram import (
Ideogram4TextEncoderConfig,
)
@@ -300,3 +303,8 @@ class Ideogram4PipelineConfig(ImagePipelineConfig):
grid_h = batch.height // patch
grid_w = batch.width // patch
return (batch_size, grid_h * grid_w, self.dit_config.arch_config.in_channels)
@dataclass
class Ideogram4DistilledPipelineConfig(Ideogram4PipelineConfig):
dit_config: DiTConfig = field(default_factory=Ideogram4DistilledDiTConfig)
python/sglang/multimodal_gen/configs/sample/ideogram.py
+22
View File
@@ -26,6 +26,18 @@ IDEOGRAM4_PRESETS: dict[str, dict[str, object]] = {
"mu": 0.5,
"std": 1.75,
},
"V4_FAST_20": {
"num_steps": 20,
"guidance_schedule": (1.0,) * 20,
"mu": 0.0,
"std": 1.75,
},
"V4_INSTANT_8": {
"num_steps": 8,
"guidance_schedule": (1.0,) * 8,
"mu": 0.0,
"std": 1.75,
},
}
@@ -76,3 +88,13 @@ class Ideogram4SamplingParams(SamplingParams):
self.num_inference_steps = preset_steps
self.guidance_scale = float(preset_cfg["guidance_schedule"][-1])
super().__post_init__()
@dataclass
class Ideogram4FastSamplingParams(Ideogram4SamplingParams):
preset: str = "V4_FAST_20"
@dataclass
class Ideogram4InstantSamplingParams(Ideogram4SamplingParams):
preset: str = "V4_INSTANT_8"
python/sglang/multimodal_gen/registry.py
+16 -1
View File
@@ -59,6 +59,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import (
Hunyuan3D2PipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.ideogram import (
Ideogram4DistilledPipelineConfig,
Ideogram4PipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.joy_echo import (
@@ -119,7 +120,11 @@ from sglang.multimodal_gen.configs.sample.hunyuan import (
HunyuanSamplingParams,
)
from sglang.multimodal_gen.configs.sample.hunyuan3d import Hunyuan3DSamplingParams
from sglang.multimodal_gen.configs.sample.ideogram import Ideogram4SamplingParams
from sglang.multimodal_gen.configs.sample.ideogram import (
Ideogram4FastSamplingParams,
Ideogram4InstantSamplingParams,
Ideogram4SamplingParams,
)
from sglang.multimodal_gen.configs.sample.joy_echo import JoyEchoSamplingParams
from sglang.multimodal_gen.configs.sample.joy_image import (
JoyImageEditSamplingParams,
@@ -1112,6 +1117,16 @@ def _register_configs():
)
# Ideogram 4
register_configs(
sampling_param_cls=Ideogram4FastSamplingParams,
pipeline_config_cls=Ideogram4DistilledPipelineConfig,
hf_model_paths=["fal/ideogram-v4-fast"],
)
register_configs(
sampling_param_cls=Ideogram4InstantSamplingParams,
pipeline_config_cls=Ideogram4DistilledPipelineConfig,
hf_model_paths=["fal/ideogram-v4-instant"],
)
register_configs(
sampling_param_cls=Ideogram4SamplingParams,
pipeline_config_cls=Ideogram4PipelineConfig,
python/sglang/multimodal_gen/runtime/models/dits/ideogram.py
+36 -5
View File
@@ -36,6 +36,9 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
Qwen3VLTextRotaryEmbedding,
qwen3_apply_rotary_pos_emb,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
LayerwiseOffloadableModuleMixin,
)
from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT
OUTPUT_IMAGE_INDICATOR = 2
@@ -83,10 +86,11 @@ def _linear(
quant_config: QuantizationConfig | None = None,
prefix: str = "",
gather_output: bool = True,
use_weight_only_fp8_linears: bool = True,
):
tp_size = _tp_size()
use_column_parallel = tp_size > 1 and out_features % tp_size == 0
if quant_config is None:
if quant_config is None and use_weight_only_fp8_linears:
if use_column_parallel:
return WeightOnlyFP8ColumnParallelLinear(
in_features,
@@ -119,13 +123,14 @@ def _merged_column_linear(
bias: bool = True,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
use_weight_only_fp8_linears: bool = True,
):
tp_size = _tp_size()
use_column_parallel = tp_size > 1 and all(
output_size % tp_size == 0 for output_size in output_sizes
)
out_features = sum(output_sizes)
if quant_config is None:
if quant_config is None and use_weight_only_fp8_linears:
if use_column_parallel:
return WeightOnlyFP8MergedColumnParallelLinear(
in_features,
@@ -158,10 +163,11 @@ def _row_linear(
bias: bool = True,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
use_weight_only_fp8_linears: bool = True,
):
tp_size = _tp_size()
use_row_parallel = tp_size > 1 and in_features % tp_size == 0
if quant_config is None:
if quant_config is None and use_weight_only_fp8_linears:
if use_row_parallel:
return WeightOnlyFP8RowParallelLinear(
in_features,
@@ -197,6 +203,7 @@ class Ideogram4Attention(nn.Module):
supported_attention_backends,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
use_weight_only_fp8_linears: bool = True,
) -> None:
super().__init__()
self.hidden_size = hidden_size
@@ -211,6 +218,7 @@ class Ideogram4Attention(nn.Module):
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.qkv",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
self.norm_q = Ideogram4RMSNorm(self.head_dim, eps=eps)
self.norm_k = Ideogram4RMSNorm(self.head_dim, eps=eps)
@@ -228,6 +236,7 @@ class Ideogram4Attention(nn.Module):
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.o",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
def forward(self, x, cos, sin, attn_mask, attn_mask_meta):
@@ -251,6 +260,7 @@ class Ideogram4MLP(nn.Module):
hidden_dim: int,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
use_weight_only_fp8_linears: bool = True,
) -> None:
super().__init__()
self.w1 = _linear(
@@ -260,6 +270,7 @@ class Ideogram4MLP(nn.Module):
quant_config=quant_config,
prefix=f"{prefix}.w1",
gather_output=False,
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
self.w2 = _row_linear(
hidden_dim,
@@ -267,6 +278,7 @@ class Ideogram4MLP(nn.Module):
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.w2",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
self.w3 = _linear(
dim,
@@ -275,6 +287,7 @@ class Ideogram4MLP(nn.Module):
quant_config=quant_config,
prefix=f"{prefix}.w3",
gather_output=False,
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
def forward(self, x):
@@ -292,6 +305,7 @@ class Ideogram4TransformerBlock(nn.Module):
supported_attention_backends,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
use_weight_only_fp8_linears: bool = True,
):
super().__init__()
self.attention = Ideogram4Attention(
@@ -301,12 +315,14 @@ class Ideogram4TransformerBlock(nn.Module):
supported_attention_backends=supported_attention_backends,
quant_config=quant_config,
prefix=f"{prefix}.attention",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
self.feed_forward = Ideogram4MLP(
hidden_size,
intermediate_size,
quant_config=quant_config,
prefix=f"{prefix}.feed_forward",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
self.attention_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
self.ffn_norm1 = Ideogram4RMSNorm(hidden_size, eps=norm_eps)
@@ -318,6 +334,7 @@ class Ideogram4TransformerBlock(nn.Module):
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.adaln_modulation",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
def forward(self, x, cos, sin, adaln_input, attn_mask, attn_mask_meta):
@@ -359,6 +376,7 @@ class Ideogram4EmbedScalar(nn.Module):
input_range: tuple[float, float],
quant_config: QuantizationConfig | None = None,
prefix: str = "",
use_weight_only_fp8_linears: bool = True,
) -> None:
super().__init__()
self.dim = dim
@@ -369,6 +387,7 @@ class Ideogram4EmbedScalar(nn.Module):
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.mlp_in",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
self.mlp_out = _linear(
dim,
@@ -376,6 +395,7 @@ class Ideogram4EmbedScalar(nn.Module):
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.mlp_out",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
def forward(self, x):
@@ -394,6 +414,7 @@ class Ideogram4FinalLayer(nn.Module):
adaln_dim: int,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
use_weight_only_fp8_linears: bool = True,
) -> None:
super().__init__()
self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False)
@@ -403,6 +424,7 @@ class Ideogram4FinalLayer(nn.Module):
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.linear",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
self.adaln_modulation = _linear(
adaln_dim,
@@ -410,6 +432,7 @@ class Ideogram4FinalLayer(nn.Module):
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.adaln_modulation",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
def forward(self, x, c):
@@ -417,14 +440,15 @@ class Ideogram4FinalLayer(nn.Module):
return self.linear(self.norm_final(x) * scale)
class Ideogram4Transformer2DModel(BaseDiT):
class Ideogram4Transformer2DModel(BaseDiT, LayerwiseOffloadableModuleMixin):
_repeated_blocks = ["Ideogram4TransformerBlock"]
layer_names = ["layers"]
_fsdp_shard_conditions = Ideogram4DiTConfig().arch_config._fsdp_shard_conditions
_compile_conditions = Ideogram4DiTConfig().arch_config._compile_conditions
_supported_attention_backends = (
Ideogram4DiTConfig().arch_config._supported_attention_backends
)
param_names_mapping = {}
param_names_mapping = Ideogram4DiTConfig().arch_config.param_names_mapping
reverse_param_names_mapping = {}
def __init__(
@@ -436,6 +460,7 @@ class Ideogram4Transformer2DModel(BaseDiT):
) -> None:
super().__init__(config, hf_config, **kwargs)
cfg = config.arch_config
use_weight_only_fp8_linears = config.use_weight_only_fp8_linears
self._supported_attention_backends = cfg._supported_attention_backends
hidden_size = cfg.num_attention_heads * cfg.attention_head_dim
self.hidden_size = hidden_size
@@ -447,6 +472,7 @@ class Ideogram4Transformer2DModel(BaseDiT):
bias=True,
quant_config=quant_config,
prefix="input_proj",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
self.llm_cond_norm = Ideogram4RMSNorm(cfg.llm_features_dim, eps=1e-6)
self.llm_cond_proj = _linear(
@@ -455,12 +481,14 @@ class Ideogram4Transformer2DModel(BaseDiT):
bias=True,
quant_config=quant_config,
prefix="llm_cond_proj",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
self.t_embedding = Ideogram4EmbedScalar(
hidden_size,
input_range=(0.0, 1.0),
quant_config=quant_config,
prefix="t_embedding",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
self.adaln_proj = _linear(
hidden_size,
@@ -468,6 +496,7 @@ class Ideogram4Transformer2DModel(BaseDiT):
bias=True,
quant_config=quant_config,
prefix="adaln_proj",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
self.embed_image_indicator = nn.Embedding(2, hidden_size)
self.rotary_emb = Qwen3VLTextRotaryEmbedding(
@@ -486,6 +515,7 @@ class Ideogram4Transformer2DModel(BaseDiT):
supported_attention_backends=self._supported_attention_backends,
quant_config=quant_config,
prefix=f"layers.{i}",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
for i in range(cfg.num_layers)
]
@@ -496,6 +526,7 @@ class Ideogram4Transformer2DModel(BaseDiT):
adaln_dim=cfg.adaln_dim,
quant_config=quant_config,
prefix="final_layer",
use_weight_only_fp8_linears=use_weight_only_fp8_linears,
)
def post_load_weights(self) -> None:
python/sglang/multimodal_gen/runtime/models/encoders/ideogram.py
+5 -1
View File
@@ -51,7 +51,11 @@ class IdeogramQwen3VLTextEncoder(TextEncoder):
text_config,
quant_config=quant_config,
use_weight_only_fp8=self._uses_weight_only_fp8,
use_tensor_parallel=True,
# bitsandbytes 4-bit quant states can be sliced safely for output
# (column-parallel) shards, but not for the row-parallel input
# shards used by attention/MLP output projections. Replicate the
# relatively small NF4 text encoder while keeping DiT TP enabled.
use_tensor_parallel=not self._uses_bitsandbytes_4bit,
)
@torch.no_grad()
python/sglang/multimodal_gen/runtime/pipelines/ideogram.py
+109 -1
View File
@@ -1,10 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
import json
import os
from dataclasses import dataclass
from functools import lru_cache
from typing import Any, cast
from huggingface_hub import hf_hub_download, snapshot_download
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
@@ -34,6 +37,15 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
_IDEOGRAM4_BASE_MODEL = "ideogram-ai/ideogram-4-fp8"
_IDEOGRAM4_DISTILLED_COMPONENTS_MODEL = "ideogram-ai/ideogram-4-nf4-diffusers"
_IDEOGRAM4_DISTILLED_COMPONENTS_REVISION = "1874bc70267ba2c823a7239e1d70dd308c8d64dc"
_IDEOGRAM4_DISTILLED_COMPONENT_PATTERNS = [
"model_index.json",
"scheduler/*",
"text_encoder/*",
"tokenizer/*",
"vae/*",
]
_IDEOGRAM4_NVFP4_COND_FILE = "diffusion_models/ideogram4_nvfp4_mixed.safetensors"
_IDEOGRAM4_NVFP4_UNCOND_FILE = (
"diffusion_models/ideogram4_unconditional_nvfp4_mixed.safetensors"
@@ -53,6 +65,20 @@ def _resolve_ideogram4_base_model_path() -> str:
return maybe_download_model(_IDEOGRAM4_BASE_MODEL, force_diffusers_model=True)
@lru_cache(maxsize=1)
def _resolve_ideogram4_distilled_components_path() -> str:
# fal's model cards explicitly reference the NF4 Diffusers repository for
# these shared components. Download only those components: its conditional
# and unconditional base transformers are unused by distilled variants.
return snapshot_download(
repo_id=_IDEOGRAM4_DISTILLED_COMPONENTS_MODEL,
revision=_IDEOGRAM4_DISTILLED_COMPONENTS_REVISION,
allow_patterns=_IDEOGRAM4_DISTILLED_COMPONENT_PATTERNS,
ignore_patterns=["*.onnx", "*.msgpack"],
max_workers=8,
)
def _resolve_ideogram4_unconditional_transformer_weights_path(
transformer_weights_path: str,
) -> str | None:
@@ -241,4 +267,86 @@ class Ideogram4Nvfp4Pipeline(Ideogram4Pipeline):
return super().load_modules(server_args, loaded_modules)
EntryClass = [Ideogram4Pipeline, Ideogram4Nvfp4Pipeline]
class Ideogram4DistilledPipeline(Ideogram4Pipeline):
_required_config_modules = [
"text_encoder",
"tokenizer",
"vae",
"transformer",
"scheduler",
]
_distilled_transformer_path: str | None = None
def _load_config(self) -> dict[str, Any]:
logger.info(
"Using '%s' for distilled config and non-transformer components",
_IDEOGRAM4_DISTILLED_COMPONENTS_MODEL,
)
model_index_path = hf_hub_download(
repo_id=_IDEOGRAM4_DISTILLED_COMPONENTS_MODEL,
filename="model_index.json",
revision=_IDEOGRAM4_DISTILLED_COMPONENTS_REVISION,
)
with open(model_index_path, encoding="utf-8") as model_index_file:
return cast(dict[str, Any], json.load(model_index_file))
def _resolve_distilled_transformer_path(self) -> str:
if self._distilled_transformer_path is None:
model_path = (
self.model_path
if os.path.exists(self.model_path)
else snapshot_download(
repo_id=self.model_path,
allow_patterns=["transformer/*"],
ignore_patterns=["*.onnx", "*.msgpack"],
max_workers=8,
)
)
self._distilled_transformer_path = os.path.join(model_path, "transformer")
return self._distilled_transformer_path
def _resolve_component_path(
self,
server_args: ServerArgs,
module_name: str,
load_module_name: str,
) -> str:
override_path = server_args.component_paths.get(module_name)
if override_path is not None:
return maybe_download_model(override_path)
if module_name == "transformer":
return self._resolve_distilled_transformer_path()
return os.path.join(
_resolve_ideogram4_distilled_components_path(), load_module_name
)
def _create_denoising_stage(self):
transformer = self.get_module("transformer")
return ProgressiveDenoisingStageRouter(
standard_stage=Ideogram4DenoisingStage(
transformer=transformer,
unconditional_transformer=None,
pipeline=self,
),
progressive_stage_factory=lambda: Ideogram4ProgressiveDenoisingStage(
transformer=transformer,
unconditional_transformer=None,
pipeline=self,
),
)
class Ideogram4FastPipeline(Ideogram4DistilledPipeline):
pipeline_name = "Ideogram4FastPipeline"
class Ideogram4InstantPipeline(Ideogram4DistilledPipeline):
pipeline_name = "Ideogram4InstantPipeline"
EntryClass = [
Ideogram4Pipeline,
Ideogram4Nvfp4Pipeline,
Ideogram4FastPipeline,
Ideogram4InstantPipeline,
]
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ideogram.py
+28 -22
View File
@@ -285,6 +285,8 @@ class Ideogram4DenoisingStage(DenoisingStage):
def _dual_transformer_execution_mode(
self,
) -> DualTransformerExecutionMode | None:
if self.unconditional_transformer is None:
return None
return DualTransformerExecutionMode.PAIRED_PER_STEP
def _cache_dit_secondary_uses_primary_config(self) -> bool:
@@ -457,30 +459,34 @@ class Ideogram4DenoisingStage(DenoisingStage):
)
pos_v = pos_out[:, max_text_tokens : max_text_tokens + num_image_tokens]
self._manage_unconditional_transformer_use_site(batch)
with set_forward_context(
current_timestep=i,
attn_metadata=step.attn_metadata,
forward_batch=batch,
):
neg_v = self._run_ideogram_transformer(
self.unconditional_transformer,
dict(
llm_features=ctx.extra["ideogram4_neg_llm_features"],
x=z,
t=t,
position_ids=ctx.extra["ideogram4_neg_position_ids"],
segment_ids=ctx.extra["ideogram4_neg_segment_ids"],
indicator=ctx.extra["ideogram4_neg_indicator"],
attn_mask=ctx.extra["ideogram4_neg_attn_mask"],
attn_mask_meta=ctx.extra["ideogram4_neg_attn_mask_meta"],
),
)
neg_v = None
if self.unconditional_transformer is not None:
self._manage_unconditional_transformer_use_site(batch)
with set_forward_context(
current_timestep=i,
attn_metadata=step.attn_metadata,
forward_batch=batch,
):
neg_v = self._run_ideogram_transformer(
self.unconditional_transformer,
dict(
llm_features=ctx.extra["ideogram4_neg_llm_features"],
x=z,
t=t,
position_ids=ctx.extra["ideogram4_neg_position_ids"],
segment_ids=ctx.extra["ideogram4_neg_segment_ids"],
indicator=ctx.extra["ideogram4_neg_indicator"],
attn_mask=ctx.extra["ideogram4_neg_attn_mask"],
attn_mask_meta=ctx.extra["ideogram4_neg_attn_mask_meta"],
),
)
with maybe_nvtx_range("scheduler_step", use_nvtx):
velocity = (
guidance_schedule[i] * pos_v + (1.0 - guidance_schedule[i]) * neg_v
)
velocity = pos_v
if neg_v is not None:
velocity = (
guidance_schedule[i] * pos_v + (1.0 - guidance_schedule[i]) * neg_v
)
ctx.latents = z + velocity * schedule_deltas[i]
python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/ideogram.py
+4 -3
View File
@@ -153,9 +153,10 @@ class Ideogram4ProgressiveDenoisingStage(
"""Progressive-resolution denoising stage for Ideogram 4.
Inherits the progressive loop from ProgressiveDenoisingStage and the
Ideogram-specific dual-transformer forward pass from Ideogram4DenoisingStage
via MRO. __init__ calls DenoisingStage directly to avoid cooperative-init
incompatibility between the two parent signatures.
Ideogram-specific forward pass from Ideogram4DenoisingStage via MRO. The
base checkpoint uses two transformers, while distilled checkpoints pass no
unconditional transformer. __init__ calls DenoisingStage directly to avoid
cooperative-init incompatibility between the two parent signatures.
MRO for method resolution:
Ideogram4ProgressiveDenoisingStage
python/sglang/multimodal_gen/runtime/server_args/server_args.py
+4
View File
@@ -129,10 +129,14 @@ DEFAULT_BCG_TEXT_BUCKETS = (64, 128, 256, 512, 1024)
BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS = frozenset(
{
"comfy-org/ideogram-4",
"fal/ideogram-v4-fast",
"fal/ideogram-v4-instant",
"glm-image",
"ideogram-4",
"ideogram-4-fp8",
"ideogram-4-nf4",
"ideogram-v4-fast",
"ideogram-v4-instant",
"ideogram-ai/ideogram-4-fp8",
"ideogram-ai/ideogram-4-nf4",
"qwen/qwen-image",
python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py
+8
View File
@@ -251,6 +251,14 @@ class TestDiffusionBCGPadding(unittest.TestCase):
"comfy-org/ideogram-4",
BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS,
)
self.assertIn(
"fal/ideogram-v4-fast",
BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS,
)
self.assertIn(
"fal/ideogram-v4-instant",
BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS,
)
self.assertIn(
"Ideogram4PipelineConfig",
BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS,
python/sglang/multimodal_gen/test/unit/test_ideogram4.py
+255 -1
View File
@@ -9,15 +9,21 @@ import torch
import torch.nn.functional as F
from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig
from sglang.multimodal_gen.configs.models.dits.ideogram import Ideogram4DiTConfig
from sglang.multimodal_gen.configs.models.dits.ideogram import (
Ideogram4DistilledDiTConfig,
Ideogram4DiTConfig,
)
from sglang.multimodal_gen.configs.models.encoders.ideogram import (
Ideogram4TextEncoderConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.ideogram import (
Ideogram4DistilledPipelineConfig,
Ideogram4PipelineConfig,
)
from sglang.multimodal_gen.configs.sample.ideogram import (
IDEOGRAM4_PRESETS,
Ideogram4FastSamplingParams,
Ideogram4InstantSamplingParams,
Ideogram4SamplingParams,
)
from sglang.multimodal_gen.registry import _get_config_info, get_model_info
@@ -51,8 +57,17 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader i
from sglang.multimodal_gen.runtime.loader.fsdp_load import (
load_model_from_full_model_state_dict,
)
from sglang.multimodal_gen.runtime.loader.utils import (
get_param_names_mapping,
hf_to_custom_state_dict,
)
from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
LayerwiseOffloadableModuleMixin,
)
from sglang.multimodal_gen.runtime.models.dits.ideogram import (
Ideogram4ColumnParallelLinear,
Ideogram4MergedColumnParallelLinear,
Ideogram4RowParallelLinear,
Ideogram4Transformer2DModel,
)
@@ -60,6 +75,8 @@ from sglang.multimodal_gen.runtime.models.encoders.ideogram import (
IdeogramQwen3VLTextEncoder,
)
from sglang.multimodal_gen.runtime.pipelines.ideogram import (
Ideogram4FastPipeline,
_resolve_ideogram4_distilled_components_path,
_resolve_ideogram4_unconditional_transformer_weights_path,
)
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
@@ -173,6 +190,12 @@ def _fake_ideogram_pipeline(transformer, unconditional_transformer):
class TestIdeogram4(unittest.TestCase):
def test_ideogram_dit_supports_layerwise_offload(self):
self.assertTrue(
issubclass(Ideogram4Transformer2DModel, LayerwiseOffloadableModuleMixin)
)
self.assertEqual(Ideogram4Transformer2DModel.layer_names, ["layers"])
def test_registry_resolves_model_index_class_name(self):
get_model_info.cache_clear()
_get_config_info.cache_clear()
@@ -206,6 +229,77 @@ class TestIdeogram4(unittest.TestCase):
self.assertIs(info.pipeline_config_cls, Ideogram4PipelineConfig)
self.assertIs(info.sampling_param_cls, Ideogram4SamplingParams)
def test_registry_resolves_fal_distilled_repos_to_native_pipelines(self):
get_model_info.cache_clear()
_get_config_info.cache_clear()
fast = get_model_info("fal/ideogram-v4-fast", backend="sglang")
instant = get_model_info("fal/ideogram-v4-instant", backend="sglang")
self.assertEqual(fast.pipeline_cls.__name__, "Ideogram4FastPipeline")
self.assertIs(fast.pipeline_config_cls, Ideogram4DistilledPipelineConfig)
self.assertIs(fast.sampling_param_cls, Ideogram4FastSamplingParams)
self.assertEqual(instant.pipeline_cls.__name__, "Ideogram4InstantPipeline")
self.assertIs(instant.pipeline_config_cls, Ideogram4DistilledPipelineConfig)
self.assertIs(instant.sampling_param_cls, Ideogram4InstantSamplingParams)
def test_fal_distilled_pipeline_resolves_component_only_repo(self):
pipeline = object.__new__(Ideogram4FastPipeline)
pipeline.model_path = "fal/ideogram-v4-fast"
pipeline._distilled_transformer_path = None
server_args = SimpleNamespace(component_paths={})
with (
patch(
"sglang.multimodal_gen.runtime.pipelines.ideogram.snapshot_download",
return_value="/cache/fast",
) as download,
patch(
"sglang.multimodal_gen.runtime.pipelines.ideogram._resolve_ideogram4_distilled_components_path",
return_value="/cache/components",
),
):
transformer = pipeline._resolve_component_path(
server_args, "transformer", "transformer"
)
vae = pipeline._resolve_component_path(server_args, "vae", "vae")
self.assertEqual(transformer, "/cache/fast/transformer")
self.assertEqual(vae, "/cache/components/vae")
self.assertNotIn("unconditional_transformer", pipeline._required_config_modules)
download.assert_called_once_with(
repo_id="fal/ideogram-v4-fast",
allow_patterns=["transformer/*"],
ignore_patterns=["*.onnx", "*.msgpack"],
max_workers=8,
)
def test_fal_distilled_pipeline_downloads_pinned_shared_components(self):
_resolve_ideogram4_distilled_components_path.cache_clear()
try:
with patch(
"sglang.multimodal_gen.runtime.pipelines.ideogram.snapshot_download",
return_value="/cache/components",
) as download:
path = _resolve_ideogram4_distilled_components_path()
finally:
_resolve_ideogram4_distilled_components_path.cache_clear()
self.assertEqual(path, "/cache/components")
download.assert_called_once_with(
repo_id="ideogram-ai/ideogram-4-nf4-diffusers",
revision="1874bc70267ba2c823a7239e1d70dd308c8d64dc",
allow_patterns=[
"model_index.json",
"scheduler/*",
"text_encoder/*",
"tokenizer/*",
"vae/*",
],
ignore_patterns=["*.onnx", "*.msgpack"],
max_workers=8,
)
def test_registry_resolves_official_nf4_repo_to_native_pipeline(self):
get_model_info.cache_clear()
_get_config_info.cache_clear()
@@ -348,6 +442,17 @@ class TestIdeogram4(unittest.TestCase):
with self.assertRaisesRegex(ValueError, "Unknown Ideogram 4 preset"):
Ideogram4SamplingParams(preset="V4_FAST")
def test_ideogram_distilled_sampling_defaults(self):
fast = Ideogram4FastSamplingParams()
instant = Ideogram4InstantSamplingParams()
self.assertEqual(fast.preset, "V4_FAST_20")
self.assertEqual(fast.num_inference_steps, 20)
self.assertEqual(fast.guidance_scale, 1.0)
self.assertEqual(instant.preset, "V4_INSTANT_8")
self.assertEqual(instant.num_inference_steps, 8)
self.assertEqual(instant.guidance_scale, 1.0)
def test_ideogram_sampling_params_merge_recomputes_preset_fields(self):
target = Ideogram4SamplingParams()
user = Ideogram4SamplingParams(
@@ -450,6 +555,25 @@ class TestIdeogram4(unittest.TestCase):
)
self.assertTrue(all(use.target_dtype is None for use in uses))
def test_ideogram_distilled_denoiser_uses_one_transformer(self):
import sglang.multimodal_gen.runtime.server_args as server_args_module
prev_args = server_args_module._global_server_args
try:
set_global_server_args(_fake_server_args())
transformer = FakeIdeogramTransformer()
stage = Ideogram4DenoisingStage(
transformer=transformer,
unconditional_transformer=None,
pipeline=_fake_ideogram_pipeline(transformer, None),
)
uses = stage.component_uses(_fake_server_args(), "stage")
finally:
set_global_server_args(prev_args)
self.assertEqual([use.component_name for use in uses], ["transformer"])
self.assertIsNone(stage._dual_transformer_execution_mode())
def test_ideogram_stages_inherit_common_stage_bases(self):
import sglang.multimodal_gen.runtime.server_args as server_args_module
@@ -664,6 +788,99 @@ class TestIdeogram4(unittest.TestCase):
)
self.assertEqual(state["layers.0.attention.qkv.weight"].dtype, FP8_WEIGHT_DTYPE)
def test_distilled_ideogram_dit_uses_unquantized_linears(self):
import sglang.multimodal_gen.runtime.server_args as server_args_module
prev_args = server_args_module._global_server_args
try:
set_global_server_args(
SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
)
with patch(
"sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size",
return_value=1,
):
with torch.device("meta"):
model = Ideogram4Transformer2DModel(
Ideogram4DistilledDiTConfig(), {}
)
finally:
set_global_server_args(prev_args)
self.assertIsInstance(model.input_proj.quant_method, UnquantizedLinearMethod)
self.assertIsInstance(
model.layers[0].attention.qkv.quant_method, UnquantizedLinearMethod
)
state = model.state_dict()
self.assertEqual(tuple(state["input_proj.weight"].shape), (4608, 128))
self.assertNotIn("input_proj.weight_scale", state)
self.assertNotIn("layers.0.attention.qkv.weight_scale", state)
def test_distilled_ideogram_maps_diffusers_attention_weights(self):
mapping = get_param_names_mapping(
Ideogram4DistilledDiTConfig().arch_config.param_names_mapping
)
weights = [
("layers.0.attention.to_q.weight", torch.full((2, 2), 1.0)),
("layers.0.attention.to_k.weight", torch.full((2, 2), 2.0)),
("layers.0.attention.to_v.weight", torch.full((2, 2), 3.0)),
("layers.0.attention.to_out.0.weight", torch.full((2, 2), 4.0)),
]
mapped, _ = hf_to_custom_state_dict(iter(weights), mapping)
torch.testing.assert_close(
mapped["layers.0.attention.qkv.weight"],
torch.cat([weight for _, weight in weights[:3]], dim=0),
)
torch.testing.assert_close(mapped["layers.0.attention.o.weight"], weights[3][1])
def test_distilled_ideogram_dit_uses_tp_unquantized_linears(self):
import sglang.multimodal_gen.runtime.server_args as server_args_module
fake_tp_group = SimpleNamespace(world_size=2, rank_in_group=1)
prev_args = server_args_module._global_server_args
try:
set_global_server_args(
SimpleNamespace(attention_backend="torch_sdpa", comfyui_mode=False)
)
with (
patch(
"sglang.multimodal_gen.runtime.models.dits.ideogram.model_parallel_is_initialized",
return_value=True,
),
patch(
"sglang.multimodal_gen.runtime.models.dits.ideogram.get_tp_world_size",
return_value=2,
),
patch(
"sglang.multimodal_gen.runtime.layers.linear.get_tp_group",
return_value=fake_tp_group,
),
patch(
"sglang.multimodal_gen.runtime.layers.attention.layer.get_ring_parallel_world_size",
return_value=1,
),
):
with torch.device("meta"):
model = Ideogram4Transformer2DModel(
Ideogram4DistilledDiTConfig(), {}
)
finally:
set_global_server_args(prev_args)
self.assertIsInstance(model.input_proj, Ideogram4ColumnParallelLinear)
self.assertIsInstance(
model.layers[0].attention.qkv, Ideogram4MergedColumnParallelLinear
)
self.assertIsInstance(model.layers[0].attention.o, Ideogram4RowParallelLinear)
self.assertIsInstance(model.input_proj.quant_method, UnquantizedLinearMethod)
self.assertEqual(tuple(model.input_proj.weight.shape), (2304, 128))
self.assertEqual(
tuple(model.layers[0].attention.qkv.weight.shape), (6912, 4608)
)
self.assertEqual(tuple(model.layers[0].attention.o.weight.shape), (4608, 2304))
def test_ideogram_dit_uses_tp_fp8_linears_when_tp_is_initialized(self):
import sglang.multimodal_gen.runtime.server_args as server_args_module
@@ -1070,6 +1287,43 @@ class TestIdeogram4(unittest.TestCase):
self.assertTrue(layer.mlp.down_proj.input_is_parallel)
self.assertTrue(layer.mlp.down_proj.reduce_results)
def test_ideogram_nf4_text_encoder_is_replicated_under_tp(self):
config = Ideogram4TextEncoderConfig()
config.update_model_arch(
{
"quantization_config": {
"quant_method": "bitsandbytes",
"load_in_4bit": True,
"bnb_4bit_quant_type": "nf4",
}
}
)
config.arch_config.text_config = Qwen3VLTextConfig(
vocab_size=32,
hidden_size=16,
intermediate_size=32,
num_hidden_layers=1,
num_attention_heads=2,
num_key_value_heads=2,
head_dim=8,
max_position_embeddings=64,
pad_token_id=0,
)
with (
patch(
"sglang.multimodal_gen.runtime.models.encoders.ideogram.BitsAndBytesConfig.from_config",
return_value=object(),
),
patch(
"sglang.multimodal_gen.runtime.models.encoders.ideogram.Qwen3VLTextModel"
) as text_model,
):
text_model.return_value = torch.nn.Identity()
IdeogramQwen3VLTextEncoder(config)
self.assertFalse(text_model.call_args.kwargs["use_tensor_parallel"])
def test_denoise_and_decode_shape_check(self):
import sglang.multimodal_gen.runtime.server_args as server_args_module
python/sglang/utils.py
+4
View File
@@ -38,6 +38,10 @@ KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS: dict[str, str] = {
"pi0.5": "Pi05Pipeline",
"hunyuan3d": "Hunyuan3D2Pipeline",
"flux.2-dev-nvfp4": "Flux2NvfpPipeline",
"fal/ideogram-v4-fast": "Ideogram4FastPipeline",
"fal--ideogram-v4-fast": "Ideogram4FastPipeline",
"fal/ideogram-v4-instant": "Ideogram4InstantPipeline",
"fal--ideogram-v4-instant": "Ideogram4InstantPipeline",
"comfy-org/ideogram-4": "Ideogram4Nvfp4Pipeline",
"comfy-org--ideogram-4": "Ideogram4Nvfp4Pipeline",
}