[diffusion] CI: update ground truth with official output (#23714)

python/sglang/multimodal_gen/configs/pipeline_configs/flux.py

@@ -77,7 +77,13 @@ class FluxPipelineConfig(ImagePipelineConfig):
                 return_overflowing_tokens=False,
                 return_length=False,
             ),
-            None,
+            dict(
+                max_length=512,
+                padding="max_length",
+                truncation=True,
+                return_overflowing_tokens=False,
+                return_length=False,
+            ),
         ]
     )
 

python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py

@@ -42,13 +42,19 @@ def qwen_image_preprocess_text(prompt):
     return txt
 
 
-def qwen_image_postprocess_text(outputs, _text_inputs, drop_idx=34):
+def qwen_image_postprocess_text(
+    outputs, _text_inputs, drop_idx=34, return_attention_mask=False
+):
     # squeeze the batch dim
     hidden_states = outputs.hidden_states[-1]
     split_hidden_states = _extract_masked_hidden(
         hidden_states, _text_inputs.attention_mask
     )
     split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+    attn_mask_list = [
+        torch.ones(e.size(0), dtype=torch.long, device=e.device)
+        for e in split_hidden_states
+    ]
     max_seq_len = max([e.size(0) for e in split_hidden_states])
     prompt_embeds = torch.stack(
         [
@@ -56,6 +62,16 @@ def qwen_image_postprocess_text(outputs, _text_inputs, drop_idx=34):
             for u in split_hidden_states
         ]
     )
+    if return_attention_mask:
+        encoder_attention_mask = torch.stack(
+            [
+                torch.cat([u, u.new_zeros(max_seq_len - u.size(0))])
+                for u in attn_mask_list
+            ]
+        )
+        if encoder_attention_mask.all():
+            return prompt_embeds, None
+        return prompt_embeds, encoder_attention_mask
     return prompt_embeds
 
 
@@ -187,9 +203,14 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig
         # Qwen-Image follows the official diffusers true-CFG behavior:
         # after combining cond/uncond with true_cfg_scale, match the per-token norm
         # back to the conditional branch.
+        cfg_scale = (
+            batch.true_cfg_scale
+            if batch.true_cfg_scale is not None
+            else batch.guidance_scale
+        )
         if (
-            batch.true_cfg_scale is None
+            cfg_scale is None
-            or batch.true_cfg_scale <= 1.0
+            or cfg_scale <= 1.0
             or not batch.do_classifier_free_guidance
         ):
             return noise_pred
@@ -586,6 +607,16 @@ class QwenImageLayeredPipelineConfig(QwenImageEditPipelineConfig):
     resolution: int = 640
     vae_precision: str = "bf16"
 
+    def postprocess_cfg_noise(
+        self,
+        batch,
+        noise_pred: torch.Tensor,
+        noise_pred_cond: torch.Tensor,
+    ) -> torch.Tensor:
+        if not batch.cfg_normalize:
+            return noise_pred
+        return super().postprocess_cfg_noise(batch, noise_pred, noise_pred_cond)
+
     def _prepare_edit_cond_kwargs(
         self, batch, prompt_embeds, rotary_emb, device, dtype
     ):
@@ -597,7 +628,7 @@ class QwenImageLayeredPipelineConfig(QwenImageEditPipelineConfig):
         vae_scale_factor = self.get_vae_scale_factor()
 
         img_shapes = batch.img_shapes
-        txt_seq_lens = batch.txt_seq_lens
+        txt_seq_lens = [prompt_embeds[0].shape[1]]
 
         freqs_cis = QwenImageEditPlusPipelineConfig.get_freqs_cis(
             img_shapes, txt_seq_lens, rotary_emb, device, dtype

python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py

@@ -326,10 +326,13 @@ class TokenizerLoader(ComponentLoader):
         ):
             return AutoProcessor.from_pretrained(component_model_path)
 
+        # Qwen-Image's model_index declares Qwen2Tokenizer; using the fast class
+        # changes text preprocessing and shifts official GT comparisons.
+        use_fast = self.component_architecture != "Qwen2Tokenizer"
         return AutoTokenizer.from_pretrained(
             component_model_path,
             padding_side="right",
-            use_fast=True,
+            use_fast=use_fast,
         )
 
 

python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

@@ -616,6 +616,9 @@ class QwenImageCrossAttention(nn.Module):
         **cross_attention_kwargs,
     ):
         seq_len_txt = encoder_hidden_states.shape[1]
+        attn_mask = cross_attention_kwargs.get("attn_mask")
+        if attn_mask is None:
+            attn_mask = cross_attention_kwargs.get("attention_mask")
 
         img_query, img_key, img_value, txt_query, txt_key, txt_value = (
             _get_qkv_projections(self, hidden_states, encoder_hidden_states)
@@ -680,6 +683,7 @@ class QwenImageCrossAttention(nn.Module):
             joint_query,
             joint_key,
             joint_value,
+            attn_mask=attn_mask,
             num_replicated_prefix=seq_len_txt,
         )
 
@@ -1223,6 +1227,8 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
 
         if isinstance(encoder_hidden_states, list):
             encoder_hidden_states = encoder_hidden_states[0]
+        if isinstance(encoder_hidden_states_mask, list):
+            encoder_hidden_states_mask = encoder_hidden_states_mask[0]
 
         hidden_states = self.img_in(hidden_states)
 
@@ -1238,6 +1244,21 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
         encoder_hidden_states = self.txt_norm(encoder_hidden_states)
         encoder_hidden_states = self.txt_in(encoder_hidden_states)
 
+        block_attention_kwargs = attention_kwargs.copy() if attention_kwargs else {}
+        if encoder_hidden_states_mask is not None:
+            encoder_hidden_states_mask = encoder_hidden_states_mask.to(
+                device=hidden_states.device, dtype=torch.bool
+            )
+            batch_size, image_seq_len = hidden_states.shape[:2]
+            image_mask = torch.ones(
+                (batch_size, image_seq_len),
+                dtype=torch.bool,
+                device=hidden_states.device,
+            )
+            block_attention_kwargs["attn_mask"] = torch.cat(
+                [encoder_hidden_states_mask, image_mask], dim=1
+            )
+
         temb = self.time_text_embed(timestep, hidden_states, additional_t_cond)
 
         temb_img_silu = F.silu(temb)
@@ -1257,7 +1278,7 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
                 temb_img_silu=temb_img_silu,
                 temb_txt_silu=temb_txt_silu,
                 image_rotary_emb=image_rotary_emb,
-                joint_attention_kwargs=attention_kwargs,
+                joint_attention_kwargs=block_attention_kwargs,
                 modulate_index=modulate_index,
             )
 

python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py

@@ -431,7 +431,6 @@ class Qwen2_5_VLTextModel(nn.Module):
 
         # It may already have been prepared by e.g. `generate`
         if not isinstance(causal_mask_mapping := attention_mask, dict):
-            # Prepare mask arguments
             mask_kwargs = {
                 "config": self.config,
                 "inputs_embeds": inputs_embeds,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -621,6 +621,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
             {
                 "encoder_hidden_states_2": batch.clip_embedding_pos,
                 "encoder_attention_mask": batch.prompt_attention_mask,
+                "encoder_hidden_states_mask": batch.prompt_attention_mask,
             }
             | server_args.pipeline_config.prepare_pos_cond_kwargs(
                 batch,
@@ -641,6 +642,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
                 {
                     "encoder_hidden_states_2": batch.clip_embedding_neg,
                     "encoder_attention_mask": batch.negative_attention_mask,
+                    "encoder_hidden_states_mask": batch.negative_attention_mask,
                 }
                 | server_args.pipeline_config.prepare_neg_cond_kwargs(
                     batch,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image_layered.py

@@ -293,6 +293,9 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
         prompt_embeds_mask = prompt_embeds_mask.view(num_images_per_prompt, seq_len)
 
+        if prompt_embeds_mask is not None and prompt_embeds_mask.all():
+            prompt_embeds_mask = None
+
         return prompt_embeds, prompt_embeds_mask
 
     # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline._encode_vae_image
@@ -412,7 +415,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         batch: Req,
         server_args: ServerArgs,
     ) -> Req:
-        use_en_prompt = True
+        use_en_prompt = batch.use_en_prompt
         device = get_local_torch_device()
         layers = batch.num_frames
         num_inference_steps = batch.num_inference_steps
@@ -443,9 +446,11 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         image = image.unsqueeze(2)
         image = image.to(dtype=torch.bfloat16)
 
-        prompt = self.get_image_caption(
+        prompt = batch.prompt
-            prompt_image, use_en_prompt=use_en_prompt, device=device
+        if not prompt or prompt.isspace():
-        )
+            prompt = self.get_image_caption(
+                prompt_image, use_en_prompt=use_en_prompt, device=device
+            )
 
         prompt_embeds, prompt_embeds_mask = self.encode_prompt(
             prompt=prompt,
@@ -514,9 +519,9 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         is_rgb = torch.tensor([0]).to(device=device, dtype=torch.long)
 
         batch.prompt_embeds = [prompt_embeds]
-        batch.prompt_embeds_mask = [prompt_embeds_mask]
+        batch.prompt_attention_mask = [prompt_embeds_mask]
         batch.negative_prompt_embeds = [negative_prompt_embeds]
-        batch.negative_prompt_embeds_mask = [negative_prompt_embeds_mask]
+        batch.negative_attention_mask = [negative_prompt_embeds_mask]
         batch.latents = latents
         batch.image_latent = image_latents
         batch.timesteps = timesteps

python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py

@@ -317,7 +317,14 @@ class TextEncodingStage(PipelineStage):
             if "pipeline_config" in postprocess_sig.parameters:
                 # required by models like LTX
                 postprocess_kwargs["pipeline_config"] = server_args.pipeline_config
+            if "return_attention_mask" in postprocess_sig.parameters:
+                postprocess_kwargs["return_attention_mask"] = return_attention_mask
             prompt_embeds = postprocess_func(outputs, text_inputs, **postprocess_kwargs)
+            has_postprocessed_attention_mask = False
+            postprocessed_attention_mask = None
+            if isinstance(prompt_embeds, tuple):
+                prompt_embeds, postprocessed_attention_mask = prompt_embeds
+                has_postprocessed_attention_mask = True
             if dtype is not None:
                 prompt_embeds = prompt_embeds.to(device=target_device, dtype=dtype)
             else:
@@ -332,11 +339,18 @@ class TextEncodingStage(PipelineStage):
                 pooled_embeds_list.append(pooled_output.to(device=target_device))
 
             if return_attention_mask:
-                mask_to_store = (
+                if has_postprocessed_attention_mask:
-                    attention_mask.to(device=target_device)
+                    mask_to_store = (
-                    if attention_mask is not None
+                        postprocessed_attention_mask.to(device=target_device)
-                    else torch.ones(input_ids.shape[:2], device=target_device)
+                        if postprocessed_attention_mask is not None
-                )
+                        else None
+                    )
+                elif attention_mask is not None:
+                    mask_to_store = attention_mask.to(device=target_device)
+                else:
+                    mask_to_store = torch.ones(
+                        input_ids.shape[:2], device=target_device
+                    )
                 attn_masks_list.append(mask_to_store)
 
         # Shape results according to return_type

python/sglang/multimodal_gen/test/server/consistency_threshold.json

@@ -14,10 +14,10 @@
             "mean_abs_diff_threshold": 8.0
         },
         "flux_2_klein_image_t2i": {
-            "clip_threshold": 0.86,
+            "clip_threshold": 0.83,
-            "ssim_threshold": 0.60,
+            "ssim_threshold": 0.52,
-            "psnr_threshold": 10.0,
+            "psnr_threshold": 9.3,
-            "mean_abs_diff_threshold": 56.0
+            "mean_abs_diff_threshold": 68.0
         },
         "zimage_image_t2i": {
             "clip_threshold": 0.92,
@@ -27,21 +27,33 @@
         },
         "zimage_image_t2i_multi_lora": {
             "clip_threshold": 0.92,
-            "ssim_threshold": 0.92,
+            "ssim_threshold": 0.89,
-            "psnr_threshold": 22.0,
+            "psnr_threshold": 19.9,
             "mean_abs_diff_threshold": 8.0
         },
         "qwen_image_t2i_cache_dit_enabled": {
-            "clip_threshold": 0.99,
+            "clip_threshold": 0.98,
-            "ssim_threshold": 0.86,
+            "ssim_threshold": 0.73,
-            "psnr_threshold": 17.0,
+            "psnr_threshold": 13.5,
-            "mean_abs_diff_threshold": 10.0
+            "mean_abs_diff_threshold": 26.0
+        },
+        "qwen_image_t2i_2_gpus": {
+            "clip_threshold": 0.98,
+            "ssim_threshold": 0.79,
+            "psnr_threshold": 14.3,
+            "mean_abs_diff_threshold": 20.0
         },
         "flux_2_image_t2i": {
+            "clip_threshold": 0.98,
+            "ssim_threshold": 0.86,
+            "psnr_threshold": 15.0,
+            "mean_abs_diff_threshold": 13.0
+        },
+        "flux_2_ti2i": {
             "clip_threshold": 0.96,
-            "ssim_threshold": 0.95,
+            "ssim_threshold": 0.88,
-            "psnr_threshold": 28.0,
+            "psnr_threshold": 19.5,
-            "mean_abs_diff_threshold": 8.0
+            "mean_abs_diff_threshold": 13.5
         },
         "layerwise_offload": {
             "clip_threshold": 0.94,
@@ -56,22 +68,22 @@
             "mean_abs_diff_threshold": 9.5
         },
         "qwen_image_edit_2509_ti2i": {
-            "clip_threshold": 0.92,
+            "clip_threshold": 0.75,
-            "ssim_threshold": 0.65,
+            "ssim_threshold": 0.53,
-            "psnr_threshold": 13.0,
+            "psnr_threshold": 10.5,
-            "mean_abs_diff_threshold": 26.0
+            "mean_abs_diff_threshold": 46.0
         },
         "qwen_image_edit_ti2i": {
             "clip_threshold": 0.96,
-            "ssim_threshold": 0.95,
+            "ssim_threshold": 0.94,
-            "psnr_threshold": 28.0,
+            "psnr_threshold": 25.5,
-            "mean_abs_diff_threshold": 8.0
+            "mean_abs_diff_threshold": 10.0
         },
         "qwen_image_edit_2511_ti2i": {
             "clip_threshold": 0.96,
             "ssim_threshold": 0.95,
-            "psnr_threshold": 28.0,
+            "psnr_threshold": 26.5,
-            "mean_abs_diff_threshold": 8.0
+            "mean_abs_diff_threshold": 10.0
         },
         "qwen_image_layered_i2i": {
             "clip_threshold": 0.92,
@@ -87,10 +99,22 @@
         },
         "wan2_1_t2v_1.3b": {
             "clip_threshold": 0.94,
-            "ssim_threshold": 0.94,
+            "ssim_threshold": 0.85,
-            "psnr_threshold": 26.0,
+            "psnr_threshold": 25.0,
             "mean_abs_diff_threshold": 8.0
         },
+        "ltx_2.3_one_stage_ti2v": {
+            "clip_threshold": 0.57,
+            "ssim_threshold": 0.42,
+            "psnr_threshold": 9.0,
+            "mean_abs_diff_threshold": 57.0
+        },
+        "ltx_2.3_two_stage_t2v_2gpus": {
+            "clip_threshold": 0.78,
+            "ssim_threshold": 0.20,
+            "psnr_threshold": 12.7,
+            "mean_abs_diff_threshold": 49.0
+        },
         "wan2_1_t2v_1.3b_teacache_enabled": {
             "clip_threshold": 0.93,
             "ssim_threshold": 0.92,
@@ -121,12 +145,6 @@
             "psnr_threshold": 9.5,
             "mean_abs_diff_threshold": 46.0
         },
-        "zimage_image_t2i_multi_lora": {
-            "clip_threshold": 0.90,
-            "ssim_threshold": 0.52,
-            "psnr_threshold": 9.5,
-            "mean_abs_diff_threshold": 46.0
-        },
         "fsdp-inference": {
             "clip_threshold": 0.92,
             "ssim_threshold": 0.90,

python/sglang/multimodal_gen/test/server/gpu_cases.py

@@ -331,6 +331,10 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [
             extras=[
                 "--pipeline-class-name LTX2TwoStageHQPipeline --ltx2-two-stage-device-mode snapshot"
             ],
+            env_vars={
+                "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
+                "SGLANG_LTX2_SNAPSHOT_RELEASE_EMPTY_CACHE": "true",
+            },
         ),
         T2I_sampling_params,
     ),

python/sglang/multimodal_gen/test/server/perf_baselines.json

@@ -812,7 +812,7 @@
                 "46": 42.62,
                 "47": 5.33,
                 "48": 189.41,
-                "49": 43.03
+                "49": 177.01
             },
             "expected_e2e_ms": 4956.71,
             "expected_avg_denoise_ms": 85.38,

python/sglang/multimodal_gen/test/test_utils.py

@@ -31,7 +31,14 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_BASE = "https://raw.githubusercontent.com/sglang-bot/sglang-ci-data/main/diffusion-ci/consistency_gt/official_generated"
+SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE = "https://raw.githubusercontent.com/sglang-bot/sglang-ci-data/main/diffusion-ci/consistency_gt/sglang_generated"
 SGL_TEST_FILES_CONSISTENCY_GT_BASE = "https://raw.githubusercontent.com/sglang-bot/sglang-ci-data/main/diffusion-ci/consistency_gt"
+SGL_TEST_FILES_CONSISTENCY_GT_BASES = (
+    SGL_TEST_FILES_OFFICIAL_CONSISTENCY_GT_BASE,
+    SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE,
+    SGL_TEST_FILES_CONSISTENCY_GT_BASE,
+)
 CONSISTENCY_THRESHOLD_JSON_PATH = (
     Path(__file__).resolve().parent / "server" / "consistency_threshold.json"
 )
@@ -873,11 +880,55 @@ def get_consistency_gt_remote_files(
     case_id: str, num_gpus: int, is_video: bool, output_format: str | None = None
 ) -> list[tuple[str, str]]:
     """Return GT filenames with their remote raw URLs."""
-    filenames = _consistency_gt_filenames(case_id, num_gpus, is_video, output_format)
+    files = _find_remote_consistency_gt_files(
-    return [
+        case_id, num_gpus, is_video, output_format
-        (filename, f"{SGL_TEST_FILES_CONSISTENCY_GT_BASE}/{filename}")
+    )
-        for filename in filenames
+    if files:
-    ]
+        return files
+
+    return _remote_consistency_gt_candidates(
+        SGL_TEST_FILES_CONSISTENCY_GT_BASE, case_id, num_gpus, is_video, output_format
+    )
+
+
+def _remote_consistency_gt_candidates(
+    base_url: str,
+    case_id: str,
+    num_gpus: int,
+    is_video: bool,
+    output_format: str | None = None,
+) -> list[tuple[str, str]]:
+    filenames = get_consistency_gt_candidates(
+        case_id, num_gpus, is_video, output_format
+    )
+    return [(filename, f"{base_url}/{filename}") for filename in filenames]
+
+
+def _remote_file_exists(url: str) -> bool:
+    try:
+        return requests.head(url, timeout=10, allow_redirects=True).status_code == 200
+    except requests.RequestException:
+        return False
+
+
+def _find_remote_consistency_gt_files(
+    case_id: str,
+    num_gpus: int,
+    is_video: bool,
+    output_format: str | None = None,
+) -> list[tuple[str, str]]:
+    for base_url in SGL_TEST_FILES_CONSISTENCY_GT_BASES:
+        candidates = _remote_consistency_gt_candidates(
+            base_url, case_id, num_gpus, is_video, output_format
+        )
+        if is_video:
+            if all(_remote_file_exists(url) for _, url in candidates):
+                return candidates
+        else:
+            for filename, url in candidates:
+                if _remote_file_exists(url):
+                    return [(filename, url)]
+    return []
 
 
 def _get_consistency_gt_dir() -> Path | None:
@@ -942,13 +993,20 @@ def load_consistency_gt(
             images.append(np.array(Image.open(path).convert("RGB")))
         logger.info(f"Loaded {len(images)} GT images for {case_id} from {gt_dir}")
     else:
-        for fn in filenames:
+        remote_files = _find_remote_consistency_gt_files(
-            url = f"{SGL_TEST_FILES_CONSISTENCY_GT_BASE}/{fn}"
+            case_id, num_gpus, is_video, output_format
+        )
+        if not remote_files:
+            raise FileNotFoundError(
+                f"GT image not found for {case_id}. Tried: {', '.join(filenames)}"
+            )
+        for _, url in remote_files:
             resp = requests.get(url, timeout=30)
             if resp.status_code != 200:
                 raise FileNotFoundError(f"GT image not found: {url}")
             images.append(np.array(Image.open(io.BytesIO(resp.content)).convert("RGB")))
-        logger.info(f"Loaded {len(images)} GT images for {case_id} from sglang-ci-data")
+        source_dir = remote_files[0][1].rsplit("/", 1)[0]
+        logger.info(f"Loaded {len(images)} GT images for {case_id} from {source_dir}")
 
     embeddings = [compute_clip_embedding(arr) for arr in images]
     loaded_gt = LoadedConsistencyGT(images=images, embeddings=embeddings)
@@ -987,14 +1045,9 @@ def gt_exists(
             return all((gt_dir / c).exists() for c in candidates)
         return any((gt_dir / c).exists() for c in candidates)
 
-    filenames = _consistency_gt_filenames(case_id, num_gpus, is_video, output_format)
+    return bool(
-    fn = filenames[0]
+        _find_remote_consistency_gt_files(case_id, num_gpus, is_video, output_format)
-    url = f"{SGL_TEST_FILES_CONSISTENCY_GT_BASE}/{fn}"
+    )
-    try:
-        r = requests.head(url, timeout=10)
-        return r.status_code == 200
-    except Exception:
-        return False
 
 
 def extract_key_frames_from_video(