[Bugfix] Fix GLM-Image output dimensions and image edit pipeline

[JaredforReal]

Purpose

Fixes several issues in the GLM-Image serving and AR-to-diffusion pipeline:

• Pass output dimensions to AR stage: The generate_images and edit_images endpoints now inject target_h / target_w into mm_processor_kwargs so the AR stage generates tokens for the correct grid size.
• Fix small preview token calculation: Replaces the incorrect token_h // 2 / token_w // 2 formula with sqrt(ratio) * (factor // 2), matching GlmImageProcessor._build_prompt_with_target_shape. This was producing wrong preview grids for non-square images.
• Normalize img2img modality key: GlmImageDataParser.parse_mm_data maps "img2img" → "image" so downstream code (mm_hashes, kwargs merging) sees a single consistent modality key.
• Robust dimension resolution in ar2diffusion: Prefers mm_processor_kwargs over top-level prompt fields, with integer coercion and informative logging.

Test plan

• Verify text-to-image generation produces correct output sizes for square (e.g. 1024x1024) and non-square requests
• Verify image editing (/v1/images/edits) works end-to-end
• Confirm small preview grid dimensions match expected values

Benchmark

Terminal 1

vllm serve zai-org/GLM-Image --port 8091 --host 0.0.0.0 --served-model-name glm-image --omni --enable-diffusion-pipeline-profiler

Terminal 2 for T2I mode

python benchmarks/diffusion/diffusion_benchmark_service.py \
    --model glm-image \
    --backend vllm-omni \ 
    --dataset vbench \
    --task t2i \
    --num-prompts 20 \
    --height 1024 --width 1024 \    
    --num-inference-steps 50 \
    --max-concurrency 1 \
    --warmup-requests 1 \  
    --output-file glm_image_t2i_1024_baseline.json 

================= Serving Benchmark Result =================
Backend:                                 vllm-omni
Model:                                   /workspace/GLM-Image
Dataset:                                 vbench
Task:                                    t2i
--------------------------------------------------
Benchmark duration (s):                  726.87
Request rate:                            inf
Max request concurrency:                 1
Successful requests:                     20/20
--------------------------------------------------
Request throughput (req/s):              0.03
Latency Mean (s):                        36.3433
Latency Median (s):                      36.2490
Latency P99 (s):                         37.1899
Latency P95 (s):                         37.1436
--------------------------------------------------
Peak Memory Max (MB):                    23420.00
Peak Memory Mean (MB):                   23420.00
Peak Memory Median (MB):                 23420.00
--------------------------------------------------
Stage Durations Mean (s):
  GlmImagePipeline.text_encoder.forward: 0.0142
  GlmImagePipeline.diffuse:              13.8316
  GlmImagePipeline.vae.decode:           0.1907

============================================================

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[Copilot]

Pull request overview

Fixes GLM-Image image sizing inconsistencies across the serving layer, AR token generation/parsing, and AR→diffusion handoff so non-square and edit requests resolve to the correct output grid dimensions.

Changes:

• Inject mm_processor_kwargs.target_h/target_w (and fallback height/width) into image generation/edit prompts in the OpenAI API server.
• Align small preview token grid calculation in _parse_generated_tokens with GlmImageProcessor._build_prompt_with_target_shape.
• Normalize img2img multimodal inputs to image within the GLM-Image AR data parser and improve dimension resolution in ar2diffusion.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
vllm_omni/model_executor/stage_input_processors/glm_image.py	Updates preview grid math and makes ar2diffusion prefer mm_processor_kwargs target dims with coercion + logging.
vllm_omni/model_executor/models/glm_image/glm_image_ar.py	Normalizes img2img modality key to image before parsing multimodal inputs.
vllm_omni/entrypoints/openai/api_server.py	Passes requested output dimensions down to the AR stage via mm_processor_kwargs for generations and edits.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 23445c1481

ℹ️ 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".

[scyiwei1986]

I use this mr on 0.18.0, but i got some wired output.

with the folloing changes.
/vllm-workspace/vllm-omni/vllm_omni/model_executor/models/glm_image/glm_image_ar.py
44c44
< from vllm.inputs.data import MultiModalDataDict

from vllm.inputs import MultiModalDataDict

I use offline mode to run the test .

python end2end.py --model-path /opt/data/weights/modelscope/GLM-Image --config-path ../config/glm_image_graph.yaml --prompt "a dog sitting on the table" --height 1024 --width 1920 --num-inference-steps 20 --output big_output_e2e.png > e2e.log 2>&1 &

glm_image_graph.yaml
# Stage config for running GLM-Image with 2-stage architecture
# Stage 0: AR Model (vLLM implementation) - generates prior_token_ids
# Stage 1: Diffusion (DiT + VAE) - denoising and image decoding

stage_args:
  # Stage 0: AR Model (GlmImageForConditionalGeneration)
  # This stage uses the vLLM-optimized AR model to generate prior tokens
  # for conditioning the diffusion process.
  - stage_id: 0
    stage_type: llm
    runtime:
      process: true
      devices: "0"
      requires_multimodal_data: true # Required for i2i mode to receive source images
    engine_args:
      model_stage: ar
      max_num_seqs: 1
      model_arch: GlmImageForConditionalGeneration
      model_subdir: vision_language_encoder # AR model config.json is in this subdirectory
      tokenizer_subdir: processor # Use processor's tokenizer (not ByT5 from tokenizer/)
      worker_cls: vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker
      #worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
      gpu_memory_utilization: 0.7
      async_scheduling: true
      enforce_eager: false
      trust_remote_code: true
      compilation_config: 
        cudagraph_capture_sizes: [1,2]
        cudagraph_mode: "FULL_DECODE_ONLY"
      engine_output_type: token_ids # Output prior_token_ids for diffusion stage
      distributed_executor_backend: "mp"
      enable_prefix_caching: false
      max_num_batched_tokens: 32768
    final_output: false # AR is not the final output
    is_comprehension: true
    default_sampling_params:
      temperature: 0.9 # From model's generation_config.json
      top_p: 0.75 # From model's generation_config.json
      top_k: 16512 # vision_vocab_size from generation_config.json
      max_tokens: 3601 # Large enough for 1920x1024 (needs 2401) and some margin
      # For reference: 1024x1024 needs 1281, 1920x1024 needs 2401, 2048x2048 needs 5121
      stop_token_ids: [16385] # eos_token_id from generation_config.json
      seed: 42
      detokenize: false

  # Stage 1: Diffusion (DiT + VAE)
  # This stage receives prior_token_ids from AR and performs denoising + VAE decode
  - stage_id: 1
    stage_type: diffusion
    runtime:
      process: true
      devices: "1" # Can use different GPU, or same GPU if memory allows
      requires_multimodal_data: true # Required for i2i mode to pass condition images
    engine_args:
      model_stage: dit
      max_num_seqs: 1
      model_arch: GlmImagePipeline # Required for diffusion model class resolution
      # Diffusion-specific parameters
      num_gpus: 1
      enforce_eager: true
      trust_remote_code: true
      distributed_executor_backend: "mp"
    engine_input_source: [0] # Input from AR stage
    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion
    final_output: true
    final_output_type: image
    default_sampling_params:
      # Diffusion-specific parameters only (no LLM params like temperature/top_p/top_k)
      seed: 42
      num_inference_steps: 50
      guidance_scale: 1.5
      height: 1024
      width: 1024

# Top-level runtime config
runtime:
  enabled: true
  defaults:
    window_size: -1 # Trigger downstream only after full upstream completion
    max_inflight: 1 # Process serially within each stage

  edges:
    - from: 0 # AR → Diffusion: trigger after AR completes
      to: 1
      window_size: -1

[JaredforReal]

@scyiwei1986 don't set any max_tokens in the config.yaml, it would be calculated internally.

and i try ur prompt:

curl -s http://172.18.67.228:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [
      {"role": "user", "content": "a dog sitting on the table"}    
    ],
    "extra_body": {
      "height": 1024,
      "width": 1920,
      "num_inference_steps": 20,
      "true_cfg_scale": 1.5,
      "seed": 42
    }
  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > dog.png

I got:

and offline end2end.py
CUDA_VISIBLE_DEVICES=1,2 python end2end.py --model-path /workspace/GLM-Image --prompt "a dog sitting on the table" --height 1024 --width 1924 --output dog.png --num-inference-steps 20

I got:

Works fine on H100 GPU, but I don't have any NPU machine, so we might need help from the community

[hsliuustc0106]

lgtm