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

This commit is contained in:
Mick
2026-04-29 13:51:31 +08:00
committed by GitHub
parent 5a94088178
commit 71d2227a78
12 changed files with 222 additions and 66 deletions
@@ -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,
),
]
)
@@ -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
or batch.true_cfg_scale <= 1.0
cfg_scale is None
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
@@ -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,
)
@@ -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,
)
@@ -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,
@@ -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,
@@ -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_image, use_en_prompt=use_en_prompt, device=device
)
prompt = batch.prompt
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
@@ -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 = (
attention_mask.to(device=target_device)
if attention_mask is not None
else torch.ones(input_ids.shape[:2], device=target_device)
)
if has_postprocessed_attention_mask:
mask_to_store = (
postprocessed_attention_mask.to(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
@@ -14,10 +14,10 @@
"mean_abs_diff_threshold": 8.0
},
"flux_2_klein_image_t2i": {
"clip_threshold": 0.86,
"ssim_threshold": 0.60,
"psnr_threshold": 10.0,
"mean_abs_diff_threshold": 56.0
"clip_threshold": 0.83,
"ssim_threshold": 0.52,
"psnr_threshold": 9.3,
"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,
"psnr_threshold": 22.0,
"ssim_threshold": 0.89,
"psnr_threshold": 19.9,
"mean_abs_diff_threshold": 8.0
},
"qwen_image_t2i_cache_dit_enabled": {
"clip_threshold": 0.99,
"ssim_threshold": 0.86,
"psnr_threshold": 17.0,
"mean_abs_diff_threshold": 10.0
"clip_threshold": 0.98,
"ssim_threshold": 0.73,
"psnr_threshold": 13.5,
"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,
"psnr_threshold": 28.0,
"mean_abs_diff_threshold": 8.0
"ssim_threshold": 0.88,
"psnr_threshold": 19.5,
"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,
"ssim_threshold": 0.65,
"psnr_threshold": 13.0,
"mean_abs_diff_threshold": 26.0
"clip_threshold": 0.75,
"ssim_threshold": 0.53,
"psnr_threshold": 10.5,
"mean_abs_diff_threshold": 46.0
},
"qwen_image_edit_ti2i": {
"clip_threshold": 0.96,
"ssim_threshold": 0.95,
"psnr_threshold": 28.0,
"mean_abs_diff_threshold": 8.0
"ssim_threshold": 0.94,
"psnr_threshold": 25.5,
"mean_abs_diff_threshold": 10.0
},
"qwen_image_edit_2511_ti2i": {
"clip_threshold": 0.96,
"ssim_threshold": 0.95,
"psnr_threshold": 28.0,
"mean_abs_diff_threshold": 8.0
"psnr_threshold": 26.5,
"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,
"psnr_threshold": 26.0,
"ssim_threshold": 0.85,
"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,
@@ -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,
),
@@ -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,
+69 -16
View File
@@ -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)
return [
(filename, f"{SGL_TEST_FILES_CONSISTENCY_GT_BASE}/{filename}")
for filename in filenames
]
files = _find_remote_consistency_gt_files(
case_id, num_gpus, is_video, output_format
)
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:
url = f"{SGL_TEST_FILES_CONSISTENCY_GT_BASE}/{fn}"
remote_files = _find_remote_consistency_gt_files(
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)
fn = filenames[0]
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
return bool(
_find_remote_consistency_gt_files(case_id, num_gpus, is_video, output_format)
)
def extract_key_frames_from_video(