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

[diffusion] fix images/edit with 2 images (#17520)

Browse Source 
Signed-off-by: Chi McIsaac <chixie.mcisaac@gmail.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Chi McIsaac
2026-03-05 03:56:39 -05:00
committed by GitHub
parent dbc896f204
commit c36de62bfc
4 changed files with 129 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/multimodal_gen/configs/pipeline_configs/flux.py
View File

@@ -509,12 +509,6 @@ class Flux2PipelineConfig(FluxPipelineConfig):
def postprocess_image_latent(self, latent_condition, batch):
batch_size = batch.batch_size
# get image_latent_ids right after scale & shift
image_latent_ids = _prepare_image_ids([latent_condition])
image_latent_ids = image_latent_ids.repeat(batch_size, 1, 1)
image_latent_ids = image_latent_ids.to(get_local_torch_device())
batch.condition_image_latent_ids = image_latent_ids
# latent: (1, 128, 32, 32)
packed = self.maybe_pack_latents(
latent_condition, None, batch
@@ -526,6 +520,11 @@ class Flux2PipelineConfig(FluxPipelineConfig):
image_latents = image_latents.repeat(batch_size, 1, 1)
return image_latents
def prepare_condition_image_latent_ids(self, image_latents, batch):
image_latent_ids = _prepare_image_ids(image_latents)
image_latent_ids = image_latent_ids.repeat(batch.batch_size, 1, 1)
batch.condition_image_latent_ids = image_latent_ids.to(get_local_torch_device())
def get_freqs_cis(self, prompt_embeds, width, height, device, rotary_emb, batch):
txt_ids = _prepare_text_ids(prompt_embeds).to(device=device)

9
python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py
View File

@@ -225,6 +225,10 @@ class ImageVAEEncodingStage(PipelineStage):
images = [images]
all_image_latents = []
prepare_condition_image_latent_ids = getattr(
server_args.pipeline_config, "prepare_condition_image_latent_ids", None
)
condition_latents = [] if callable(prepare_condition_image_latent_ids) else None
for image in images:
image = self.preprocess(
image,
@@ -309,12 +313,17 @@ class ImageVAEEncodingStage(PipelineStage):
latent_condition -= shift_factor
latent_condition = latent_condition * scaling_factor
if condition_latents is not None:
condition_latents.append(latent_condition)
image_latent = server_args.pipeline_config.postprocess_image_latent(
latent_condition, batch
)
all_image_latents.append(image_latent)
batch.image_latent = torch.cat(all_image_latents, dim=1)
if condition_latents is not None:
prepare_condition_image_latent_ids(condition_latents, batch)
self.offload_model()
return batch

66
python/sglang/multimodal_gen/test/server/perf_baselines.json
View File

@@ -447,6 +447,72 @@
"expected_avg_denoise_ms": 1077.26,
"expected_median_denoise_ms": 1096.5
},
"flux_2_ti2i_multi_image_cache_dit": {
"stages_ms": {
"ImageVAEEncodingStage": 282.83,
"DenoisingStage": 26936.93,
"DecodingStage": 129.33,
"TextEncodingStage": 737.01,
"LatentPreparationStage": 0.84,
"TimestepPreparationStage": 20.57,
"InputValidationStage": 84.29
},
"denoise_step_ms": {
"0": 1846.55,
"1": 232.6,
"2": 1621.05,
"3": 1606.1,
"4": 1441.78,
"5": 59.2,
"6": 60.32,
"7": 235.36,
"8": 1443.21,
"9": 59.71,
"10": 59.79,
"11": 236.12,
"12": 1439.11,
"13": 59.97,
"14": 60.46,
"15": 239.29,
"16": 1441.32,
"17": 60.76,
"18": 60.68,
"19": 240.16,
"20": 1442.46,
"21": 60.14,
"22": 61.1,
"23": 239.26,
"24": 1443.28,
"25": 59.41,
"26": 60.74,
"27": 238.02,
"28": 1444.38,
"29": 59.09,
"30": 59.12,
"31": 241.69,
"32": 1441.99,
"33": 60.44,
"34": 61.93,
"35": 241.0,
"36": 1443.56,
"37": 60.55,
"38": 61.07,
"39": 238.11,
"40": 1443.33,
"41": 59.08,
"42": 60.43,
"43": 239.23,
"44": 1444.53,
"45": 61.77,
"46": 61.65,
"47": 239.84,
"48": 1443.89,
"49": 59.92
},
"expected_e2e_ms": 28591.93,
"expected_avg_denoise_ms": 533.42,
"expected_median_denoise_ms": 235.74
},
"flux_image_t2i_2_gpus": {
"stages_ms": {
"InputValidationStage": 0.03,

65
python/sglang/multimodal_gen/test/server/testcase_configs.py
View File

@@ -321,6 +321,13 @@ MULTI_IMAGE_TI2I_sampling_params = DiffusionSamplingParams(
],
direct_url_test=True,
)
MULTI_IMAGE_TI2I_UPLOAD_sampling_params = DiffusionSamplingParams(
prompt="The magician bear is on the left, the alchemist bear is on the right, facing each other in the central park square.",
image_path=[
"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_1.jpg",
"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/edit2509/edit2509_2.jpg",
],
)
MULTI_FRAME_I2I_sampling_params = DiffusionSamplingParams(
prompt="a high quality, cute halloween themed illustration, consistent style and lighting",
image_path=[
@@ -699,22 +706,6 @@ TWO_GPU_CASES_A = [
),
]
# Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536.
if not current_platform.is_hip():
TWO_GPU_CASES_A.append(
DiffusionTestCase(
"turbo_wan2_2_i2v_a14b_2gpu",
DiffusionServerArgs(
model_path="IPostYellow/TurboWan2.2-I2V-A14B-Diffusers",
modality="video",
custom_validator="video",
num_gpus=2,
tp_size=2,
),
TURBOWAN_I2V_sampling_params,
)
)
TWO_GPU_CASES_B = [
DiffusionTestCase(
"wan2_1_i2v_14b_480P_2gpu",
@@ -800,6 +791,48 @@ TWO_GPU_CASES_B = [
),
]
if not current_platform.is_hip():
# Flux2 multi-image edit with cache-dit, regression test
ONE_GPU_CASES_B.append(
DiffusionTestCase(
"flux_2_ti2i_multi_image_cache_dit",
DiffusionServerArgs(
model_path="black-forest-labs/FLUX.2-dev",
modality="image",
enable_cache_dit=True,
),
MULTI_IMAGE_TI2I_UPLOAD_sampling_params,
)
)
# Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536.
ONE_GPU_CASES_B.append(
DiffusionTestCase(
"turbo_wan2_1_t2v_1.3b",
DiffusionServerArgs(
model_path="IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers",
modality="video",
custom_validator="video",
),
DiffusionSamplingParams(
prompt=T2V_PROMPT,
),
)
)
# Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536.
TWO_GPU_CASES_A.append(
DiffusionTestCase(
"turbo_wan2_2_i2v_a14b_2gpu",
DiffusionServerArgs(
model_path="IPostYellow/TurboWan2.2-I2V-A14B-Diffusers",
modality="video",
custom_validator="video",
num_gpus=2,
tp_size=2,
),
TURBOWAN_I2V_sampling_params,
)
)
# Load global configuration
BASELINE_CONFIG = BaselineConfig.load(
Path(__file__).with_name("perf_baselines.json")