[Bugfix] Fix config misalignment between offline and online diffusion inference (Wan2.2, Qwen-Image series)

[SamitHuang]

Summary

Fixes config arguments that diverge between offline (Omni.generate) and online (serving chat / API) inference for diffusion models, causing different generation results for the same prompt+seed. Affected models: Wan2.2 (T2V/I2V), Qwen-Image, Qwen-Image-Edit, and Qwen-Image-Layered.

Root Causes & Fixes

1. guidance_scale_2 leaked the 0.0 sentinel value — OmniDiffusionRequest.__post_init__ auto-filled guidance_scale_2 from guidance_scale before resolving the 0.0 sentinel back to 1.0. Online requests that omitted guidance_scale ended up with guidance_scale_2 = 0.0, which disabled CFG entirely on the low-noise stage in Wan2.2 (and silently altered quality for Qwen-Image models). Fix: resolve the sentinel first, then auto-fill guidance_scale_2.

2. num_inference_steps hardcoded to 50 in serving_chat.py — The chat endpoint forced 50 steps regardless of the pipeline's own default (Wan2.2 uses 40, others may differ). Fix: change the OmniDiffusionSamplingParams.num_inference_steps default to None (sentinel) so each pipeline's forward() applies its own model-specific default when the caller does not specify a value.

3. Redundant guidance_scale_provided flag in AsyncOmniDiffusion — AsyncOmniDiffusion.generate() manually set guidance_scale_provided before OmniDiffusionRequest.__post_init__ ran, which then overwrote it anyway. Fix: remove the redundant pre-set; __post_init__ now handles it correctly and consistently for both offline and online paths.

4. cfg_scale not mapped to true_cfg_scale in online serving — The offline script uses --cfg-scale which maps to true_cfg_scale in OmniDiffusionSamplingParams, but the online serving_chat.py handler only read extra_body.get("true_cfg_scale"), silently ignoring the cfg_scale key sent by clients. Fix: accept cfg_scale as an alias for true_cfg_scale.

5. layers and resolution not passed through in online serving — Qwen-Image-Layered parameters from extra_body were not forwarded to OmniDiffusionSamplingParams. Fix: add passthrough for both.

6. RGB→RGBA conversion missing for Qwen-Image-Layered — The model's VAE expects 4-channel (RGBA) input, but online serving decodes user-uploaded images as RGB. Fix: add automatic RGB→RGBA conversion in both the preprocessing function and the fallback path in forward().

Files Changed

File	Change
vllm_omni/diffusion/request.py	Reorder __post_init__: resolve guidance_scale sentinel → set CFG flag → auto-fill guidance_scale_2
vllm_omni/inputs/data.py	Change num_inference_steps default from 50 to None
vllm_omni/entrypoints/openai/serving_chat.py	Remove hardcoded defaults; accept cfg_scale alias; pass layers/resolution
vllm_omni/entrypoints/async_omni_diffusion.py	Remove redundant guidance_scale_provided pre-set
vllm_omni/diffusion/worker/diffusion_model_runner.py	Guard cache refresh against num_inference_steps=None
vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py	Auto-convert RGB→RGBA input images

Test Plan

For each model, offline and online inference were run with identical parameters (same seed, prompt, num_inference_steps, guidance scales). The raw output tensors (before any video/image encoding) were saved and compared at the numpy level to verify pixel-level identity.

Test Parameters

Model	Prompt	Seed	Steps	cfg_scale	guidance_scale	Extra
Wan2.2 (T2V)	"a cat walking..."	42	40	-	5.0	480p, 17 frames
Qwen-Image	"a beautiful landscape..."	42	50	-	1.0	1024×1024
Qwen-Image-Edit	"edit the image..."	42	50	4.0	1.0	1024×1024
Qwen-Image-Layered	"a rabbit"	0	50	4.0	-	4 layers, 640 resolution, rabbit.png input

Test Scripts & Commands

Wan2.2, Qwen-Image, Qwen-Image-Edit

• Offline: Custom E2E test script e2e_test_config_alignment.py (not committed) calling Omni.generate() directly with explicit OmniDiffusionSamplingParams.
• Online: Same script starting vllm serve <model> --omni, then sending requests via requests.post() to /v1/chat/completions.
• Output images saved to: e2e_test_outputs/ (e.g., qwen-image_offline.png, qwen-image_online.png, wan22_offline_raw.npy, wan22_online_raw.npy).

Qwen-Image-Layered

Offline command:

cd examples && CUDA_VISIBLE_DEVICES=7 python offline_inference/image_to_image/image_edit.py \
    --model "Qwen/Qwen-Image-Layered" \
    --image rabbit.png \
    --prompt "a rabbit" \
    --output "e2e_layered_offline" \
    --num-inference-steps 50 \
    --cfg-scale 4.0 \
    --seed 0 \
    --layers 4 \
    --color-format "RGBA" \
    --enforce-eager

Online server:

CUDA_VISIBLE_DEVICES=3 vllm serve "Qwen/Qwen-Image-Layered" \
    --omni --port 8092 --enforce-eager

Online request:

import base64, requests

with open('examples/rabbit.png', 'rb') as f:
    img_b64 = base64.b64encode(f.read()).decode()

payload = {
    'model': 'Qwen/Qwen-Image-Layered',
    'messages': [{
        'role': 'user',
        'content': [
            {'type': 'image_url', 'image_url': {'url': f'data:image/png;base64,{img_b64}'}},
            {'type': 'text', 'text': 'a rabbit'}
        ]
    }],
    'extra_body': {
        'num_inference_steps': 50,
        'cfg_scale': 4.0,
        'seed': 0,
        'layers': 4,
    }
}

resp = requests.post('http://localhost:8092/v1/chat/completions', json=payload, timeout=600)

Output images saved to:

• Offline: examples/e2e_layered_offline_0.png to e2e_layered_offline_3.png (4 RGBA layers, 864×480)
• Online: raw decoded tensor saved as /tmp/e2e_layered_online/raw_decoded.npy (the online API postprocessing for layered output has a pre-existing bug — images are compared at the raw tensor level before postprocessing)

Offline:

Online:

Comparison:

import numpy as np
offline = np.load('/tmp/e2e_layered_offline/raw_decoded.npy')
online  = np.load('/tmp/e2e_layered_online/raw_decoded.npy')
print(f'Shapes: {offline.shape} vs {online.shape}')
print(f'Pixel-identical: {np.array_equal(offline, online)}')
# Shapes: (4, 4, 480, 864) vs (4, 4, 480, 864)
# Pixel-identical: True

Test Results

All four models produce pixel-identical output between offline and online inference:

Model                   Offline Shape              Online Shape               Pixel-Identical
─────────────────────── ────────────────────────── ────────────────────────── ───────────────
Wan2.2 (T2V)            (1, 3, 17, 480, 832)       (1, 3, 17, 480, 832)       ✅ True
Qwen-Image              (1024, 1024, 3)            (1024, 1024, 3)            ✅ True
Qwen-Image-Edit         (1024, 1024, 3)            (1024, 1024, 3)            ✅ True
Qwen-Image-Layered      (4, 4, 480, 864)           (4, 4, 480, 864)           ✅ True

Qwen-Image-Layered: Additional Fixes Verified

The online serving path for Qwen-Image-Layered had three additional issues fixed in this PR:

1. cfg_scale → true_cfg_scale mapping: Before fix, raw_true_cfg_scale=None (ignored); after fix, raw_true_cfg_scale=4.0 (correctly mapped).
2. RGB→RGBA conversion: Before fix, the online request crashed with expected input to have 4 channels, but got 3 channels; after fix, RGB images are automatically converted to RGBA.
3. layers/resolution passthrough: These parameters are now correctly forwarded from extra_body to the pipeline.

Server log comparison (effective params in pipeline's forward()):

# Before fix (online):
true_cfg_scale=4.0, raw_true_cfg_scale=None, guidance_scale_provided=False

# After fix (online):
true_cfg_scale=4.0, raw_true_cfg_scale=4.0, guidance_scale_provided=False

# Offline reference:
true_cfg_scale=4.0, raw_true_cfg_scale=4.0, guidance_scale_provided=True

Note: guidance_scale_provided differs (True offline vs False online) because the offline script explicitly passes --guidance-scale 1.0 while the online request omits it. This does not affect the output because Qwen-Image-Layered is not a guidance-distilled model — the guidance_scale value is unused when transformer.guidance_embeds is False.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 140ef4afc3

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Restore a concrete default for num_inference_steps

Changing OmniDiffusionSamplingParams.num_inference_steps to None breaks pipelines that still require an explicit step count. In DreamIDOmniPipeline.forward, the value is read directly and passed into get_scheduler_time_steps without a fallback, and FlowUniPCMultistepScheduler.set_timesteps asserts that num_inference_steps is not None; requests that omit this field now fail at runtime instead of using the previous default behavior.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Always refresh cache state even when steps are omitted

The new num_inference_steps is not None gate prevents cache_backend.refresh(...) from running for requests that rely on pipeline defaults (now common after this commit). That leaves per-request cache state stale across generations: for example, TeaCacheBackend.refresh is the hook reset path and is documented/implemented as required before each generation. With cache enabled and omitted step count, subsequent requests can reuse prior residual/counter state and produce incorrect outputs.

Useful? React with 👍 / 👎.

[hsliuustc0106]

also for qwen-image-layered

[SamitHuang]

also for qwen-image-layered

already covered by the same fixed for qwen-image-edit

[david6666666]

also for qwen-image-layered

@bjf-frz qwen-image-layered args should be alignmented, please check after this pr.thanks

[ZJY0516]

LGTM

[hsliuustc0106]

@SamitHuang please also attach the result for qwen-image-layered

[david6666666]

@bjf-frz PTAL thanks, and RGBA issue appears to remain unresolved.

[SamitHuang]

@bjf-frz PTAL thanks, and RGBA issue appears to remain unresolved.

addressed

[bjf-frz]

lgtm

[david6666666]

LGTM

[hsliuustc0106]

@ZJY0516 @SamitHuang remember to add tests to confirm these online and offline alignments