[Perf] GLM Image (#920)

Signed-off-by: JaredforReal <w13431838023@gmail.com>
Signed-off-by: Jared Wen <w13431838023@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/models/supported_models.md

@@ -41,6 +41,7 @@ th {
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
+|`GlmImageForConditionalGeneration` | GLM-Image | `zai-org/GLM-Image` |
 |`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |
 
 

examples/offline_inference/glm_image/README.md

@@ -0,0 +1,138 @@
+# GLM-Image Multistage End-to-End Inference
+
+This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture.
+
+## Architecture
+
+GLM-Image uses a 2-stage pipeline:
+
+```
+┌─────────────────────────────────────────────────────────────┐
+│                     GLM-Image Pipeline                       │
+├─────────────────────────────────────────────────────────────┤
+│                                                             │
+│  Stage 0 (AR Model)              Stage 1 (Diffusion)        │
+│  ┌─────────────────┐            ┌─────────────────────┐     │
+│  │ vLLM-optimized  │            │  GlmImagePipeline   │     │
+│  │ GlmImageFor     │  prior     │  ┌───────────────┐  │     │
+│  │ Conditional     │──tokens───►│  │ DiT Denoiser  │  │     │
+│  │ Generation      │            │  └───────────────┘  │     │
+│  │ (9B AR model)   │            │         │          │     │
+│  └─────────────────┘            │         ▼          │     │
+│         ▲                       │  ┌───────────────┐  │     │
+│         │                       │  │  VAE Decode   │──┼──► Image
+│    Text/Image                   │  └───────────────┘  │     │
+│      Input                      └─────────────────────┘     │
+│                                                             │
+└─────────────────────────────────────────────────────────────┘
+```
+
+## Features
+
+- **vLLM-optimized AR**: Uses PagedAttention and tensor parallelism for faster prior token generation
+- **Flexible deployment**: AR and Diffusion stages can run on different GPUs
+- **Text-to-Image**: Generate images from text descriptions
+- **Image-to-Image**: Edit existing images with text prompts
+
+## Usage
+
+### Text-to-Image
+
+```bash
+python end2end.py \
+    --model-path /path/to/glm-image \
+    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "A beautiful sunset over the ocean with sailing boats" \
+    --height 1024 \
+    --width 1024 \
+    --output output_t2i.png
+```
+
+### Image-to-Image (Image Editing)
+
+```bash
+python end2end.py \
+    --model-path /path/to/glm-image \
+    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "Transform this scene into a winter wonderland" \
+    --image input.png \
+    --output output_i2i.png
+```
+
+### With Custom Parameters
+
+```bash
+python end2end.py \
+    --model-path /path/to/glm-image \
+    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "A photorealistic cat sitting on a window sill" \
+    --height 1024 \
+    --width 1024 \
+    --num-inference-steps 50 \
+    --guidance-scale 1.5 \
+    --seed 42 \
+    --output output.png
+```
+
+## Shell Scripts
+
+### Run Text-to-Image
+
+```bash
+./run_t2i.sh
+```
+
+### Run Image-to-Image
+
+```bash
+./run_i2i.sh --image /path/to/input.png
+```
+
+## Stage Configuration
+
+The stage config (`glm_image.yaml`) defines:
+
+- **Stage 0 (AR)**: Uses `GPUARWorker` with vLLM engine
+
+  - Model: `GlmImageForConditionalGeneration`
+  - Output: `token_ids` (prior tokens)
+
+- **Stage 1 (Diffusion)**: Uses diffusion engine
+  - Model: `GlmImagePipeline`
+  - Output: Generated image
+
+See `vllm_omni/model_executor/stage_configs/glm_image.yaml` for full configuration.
+
+## Comparison with Single-Stage
+
+| Aspect      | Single-Stage (transformers) | Multistage (vLLM)   |
+| ----------- | --------------------------- | ------------------- |
+| AR Model    | transformers native         | vLLM PagedAttention |
+| Memory      | Higher (no KV cache opt)    | Lower (optimized)   |
+| Throughput  | Lower                       | Higher              |
+| Flexibility | Single GPU                  | Multi-GPU support   |
+
+## Troubleshooting
+
+### OOM Error
+
+Try reducing memory usage:
+
+```bash
+# In glm_image.yaml, adjust:
+gpu_memory_utilization: 0.5  # Reduce from 0.6
+```
+
+### Slow Initialization
+
+The first run loads model weights. Subsequent runs are faster:
+
+```bash
+--stage-init-timeout 900  # Increase timeout for slow storage
+```
+
+## Requirements
+
+- vLLM-Omni with GLM-Image support
+- CUDA-capable GPU (recommended: H100/A100 with 80GB)
+- GLM-Image model weights

examples/offline_inference/glm_image/end2end.py

@@ -0,0 +1,476 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+End-to-end offline inference example for GLM-Image with multistage architecture.
+
+This script tests the multistage pipeline where:
+- Stage 0 (AR): vLLM-optimized GlmImageForConditionalGeneration generates prior_token_ids
+- Stage 1 (Diffusion): GlmImagePipeline performs DiT denoising + VAE decode
+
+Usage (text-to-image):
+    python end2end.py \
+        --model-path /path/to/glm-image \
+        --config-path /path/to/glm_image.yaml \
+        --prompt "A beautiful sunset over the ocean" \
+        --output output_t2i.png
+
+Usage (image-to-image / image edit):
+    python end2end.py \
+        --model-path /path/to/glm-image \
+        --config-path /path/to/glm_image.yaml \
+        --prompt "Make it look like winter" \
+        --image input.png \
+        --output output_i2i.png
+
+Usage (with custom parameters):
+    python end2end.py \
+        --model-path /path/to/glm-image \
+        --config-path /path/to/glm_image.yaml \
+        --prompt "A cat sitting on a window sill" \
+        --height 1024 \
+        --width 1024 \
+        --num-inference-steps 50 \
+        --guidance-scale 1.5 \
+        --seed 42
+
+For more options, run:
+    python end2end.py --help
+"""
+
+import argparse
+import os
+import time
+from pathlib import Path
+
+from PIL import Image
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# Default stage config path (relative to vllm_omni package)
+DEFAULT_CONFIG_PATH = "vllm_omni/model_executor/stage_configs/glm_image.yaml"
+
+SEED = 42
+
+# GLM-Image special tokens
+GLM_IMAGE_EOS_TOKEN_ID = 16385  # eos_token_id from generation_config.json
+GLM_IMAGE_VISION_VOCAB_SIZE = 16512  # top_k should be vision_vocab_size
+
+
+def compute_max_tokens(height: int, width: int, factor: int = 32) -> int:
+    """
+    Compute max_new_tokens for GLM-Image AR generation.
+
+    GLM-Image generates tokens in this order for text-to-image:
+    1. Small preview image (half resolution in each dimension)
+    2. Large target image (full resolution)
+    3. EOS token
+
+    Args:
+        height: Target image height in pixels
+        width: Target image width in pixels
+        factor: Downsampling factor (32 for GLM-Image AR output)
+
+    Returns:
+        Total number of tokens to generate (small + large + EOS)
+    """
+    # Large image tokens (target resolution)
+    token_h = height // factor
+    token_w = width // factor
+    large_tokens = token_h * token_w
+
+    # Small preview tokens (half resolution in each dimension)
+    small_h = token_h // 2
+    small_w = token_w // 2
+    small_tokens = small_h * small_w
+
+    # Total: small + large + EOS
+    return small_tokens + large_tokens + 1
+
+
+def load_image(image_path: str) -> Image.Image:
+    """Load an image from file path."""
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+    return Image.open(image_path).convert("RGB")
+
+
+def save_image(image: Image.Image, output_path: str) -> None:
+    """Save an image to file path."""
+    output_dir = os.path.dirname(output_path)
+    if output_dir:
+        os.makedirs(output_dir, exist_ok=True)
+    image.save(output_path)
+    print(f"Image saved to: {output_path}")
+
+
+def build_prompt_for_t2i(
+    prompt: str,
+    height: int = 1024,
+    width: int = 1024,
+) -> dict:
+    """
+    Build prompt dict for text-to-image generation.
+
+    Args:
+        prompt: Text description for image generation
+        height: Target image height
+        width: Target image width
+
+    Returns:
+        Dict containing prompt and generation parameters
+    """
+    return {
+        "prompt": prompt,
+        "height": height,
+        "width": width,
+        # Pass target dimensions to AR processor for proper grid token generation
+        "mm_processor_kwargs": {
+            "target_h": height,
+            "target_w": width,
+        },
+    }
+
+
+def build_prompt_for_i2i(
+    prompt: str,
+    image: Image.Image,
+    height: int | None = None,
+    width: int | None = None,
+) -> dict:
+    """
+    Build prompt dict for image-to-image generation.
+
+    Args:
+        prompt: Text description for image editing
+        image: Source image for editing
+        height: Target image height (default: use source image size)
+        width: Target image width (default: use source image size)
+
+    Returns:
+        Dict containing prompt, image, and generation parameters
+    """
+    # Use source image dimensions if not specified
+    if height is None:
+        height = image.height
+    if width is None:
+        width = image.width
+
+    return {
+        "prompt": prompt,
+        "multi_modal_data": {
+            "image": image,
+        },
+        "height": height,
+        "width": width,
+        # Pass target dimensions to AR processor for proper grid token generation
+        "mm_processor_kwargs": {
+            "target_h": height,
+            "target_w": width,
+        },
+    }
+
+
+def main(args: argparse.Namespace) -> None:
+    """Main entry point for GLM-Image end-to-end inference."""
+    print("=" * 60)
+    print("GLM-Image Multistage End-to-End Inference")
+    print("=" * 60)
+
+    # Validate arguments
+    if not args.model_path:
+        raise ValueError("--model-path is required")
+
+    if not args.prompt:
+        raise ValueError("--prompt is required")
+
+    # Determine config path
+    config_path = args.config_path
+    if config_path is None:
+        # Try to find default config
+        if os.path.exists(DEFAULT_CONFIG_PATH):
+            config_path = DEFAULT_CONFIG_PATH
+        else:
+            # Try relative to script location
+            script_dir = Path(__file__).parent.parent.parent.parent
+            config_path = script_dir / "vllm_omni/model_executor/stage_configs/glm_image.yaml"
+            if not config_path.exists():
+                raise FileNotFoundError(
+                    f"Stage config not found. Please specify --config-path. Tried: {DEFAULT_CONFIG_PATH}"
+                )
+            config_path = str(config_path)
+
+    print(f"Model path: {args.model_path}")
+    print(f"Config path: {config_path}")
+    print(f"Prompt: {args.prompt}")
+
+    # Load source image for image-to-image mode
+    source_image = None
+    if args.image:
+        print(f"Source image: {args.image}")
+        source_image = load_image(args.image)
+        print(f"  Image size: {source_image.width}x{source_image.height}")
+
+    # Build prompt based on mode
+    if source_image is not None:
+        # Image-to-image mode
+        prompt_dict = build_prompt_for_i2i(
+            prompt=args.prompt,
+            image=source_image,
+            height=args.height,
+            width=args.width,
+        )
+        mode = "image-to-image"
+    else:
+        # Text-to-image mode
+        prompt_dict = build_prompt_for_t2i(
+            prompt=args.prompt,
+            height=args.height or 1024,
+            width=args.width or 1024,
+        )
+        mode = "text-to-image"
+
+    print(f"Mode: {mode}")
+    print(f"Target size: {prompt_dict.get('height', 1024)}x{prompt_dict.get('width', 1024)}")
+
+    # Add generation parameters to prompt
+    prompt_dict["seed"] = args.seed
+    prompt_dict["num_inference_steps"] = args.num_inference_steps
+    prompt_dict["guidance_scale"] = args.guidance_scale
+
+    if args.negative_prompt:
+        prompt_dict["negative_prompt"] = args.negative_prompt
+
+    # Initialize Omni with multistage config
+    print("\nInitializing Omni with multistage pipeline...")
+    start_time = time.time()
+
+    omni = Omni(
+        model=args.model_path,
+        stage_configs_path=config_path,
+        log_stats=args.enable_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    init_time = time.time() - start_time
+    print(f"Initialization completed in {init_time:.2f}s")
+
+    # Prepare prompts (support batch generation)
+    prompts = [prompt_dict for _ in range(args.num_prompts)]
+
+    # No explicit sampling_params for diffusion - parameters are in prompt_dict
+    # For multistage, the AR stage may need sampling params
+    from vllm import SamplingParams
+
+    # Compute max_tokens dynamically based on target image size
+    target_height = prompt_dict.get("height", 1024)
+    target_width = prompt_dict.get("width", 1024)
+    calculated_max_tokens = compute_max_tokens(target_height, target_width)
+
+    # Use calculated value unless user explicitly specified a different value
+    # Default args.max_tokens is 16384 (very large), so prefer calculated value
+    effective_max_tokens = calculated_max_tokens if args.max_tokens == 16384 else args.max_tokens
+
+    if args.verbose:
+        print(f"AR max_tokens: {effective_max_tokens} (calculated: {calculated_max_tokens}, arg: {args.max_tokens})")
+
+    # IMPORTANT: GLM-Image AR model requires these exact sampling parameters
+    # from generation_config.json for proper image token generation.
+    # - temperature=0.9, top_p=0.75, top_k=16512 (vision_vocab_size)
+    # - stop_token_ids=[16385] (eos_token_id) is CRITICAL to stop generation
+    ar_sampling_params = SamplingParams(
+        temperature=0.9,  # From generation_config.json
+        top_p=0.75,  # From generation_config.json
+        top_k=GLM_IMAGE_VISION_VOCAB_SIZE,  # 16512, vision vocabulary size
+        max_tokens=effective_max_tokens,
+        stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],  # 16385, CRITICAL for stopping
+        seed=args.seed,
+        detokenize=False,
+    )
+
+    # For diffusion stage, sampling_params contains diffusion-specific parameters
+    # These are passed as kwargs to the diffusion engine
+    diffusion_sampling_params = OmniDiffusionSamplingParams(
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        height=prompt_dict.get("height", 1024),
+        width=prompt_dict.get("width", 1024),
+        seed=args.seed,
+    )
+
+    # For multistage, we need sampling_params for each stage
+    # Stage 0 (AR): SamplingParams for vLLM
+    # Stage 1 (Diffusion): dict with diffusion kwargs
+    sampling_params_list = [ar_sampling_params, diffusion_sampling_params]
+
+    # Run generation
+    print(f"\nGenerating {args.num_prompts} image(s)...")
+    gen_start_time = time.time()
+
+    output_dir = os.path.dirname(args.output) if args.output else "outputs"
+    if output_dir:
+        os.makedirs(output_dir, exist_ok=True)
+
+    output_count = 0
+    for stage_outputs in omni.generate(prompts, sampling_params_list, py_generator=True):
+        if stage_outputs.final_output_type == "image":
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+
+                # Get generated images
+                images = output.images if hasattr(output, "images") else []
+                if not images and hasattr(output, "multimodal_output"):
+                    images = output.multimodal_output.get("images", [])
+
+                # Save each generated image
+                for idx, img in enumerate(images):
+                    if args.num_prompts == 1 and len(images) == 1:
+                        output_path = args.output
+                    else:
+                        base, ext = os.path.splitext(args.output)
+                        output_path = f"{base}_{request_id}_{idx}{ext}"
+
+                    if isinstance(img, Image.Image):
+                        save_image(img, output_path)
+                    else:
+                        print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
+
+                    output_count += 1
+
+        elif stage_outputs.final_output_type == "text":
+            # AR stage output (intermediate, for debugging)
+            if args.verbose:
+                for output in stage_outputs.request_output:
+                    print(f"AR output for request {output.request_id}:")
+                    print(f"  Token count: {len(output.outputs[0].token_ids)}")
+
+    gen_time = time.time() - gen_start_time
+    print(f"\nGeneration completed in {gen_time:.2f}s")
+    print(f"Generated {output_count} image(s)")
+
+    # Cleanup
+    omni.close()
+    print("\nDone!")
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="GLM-Image Multistage End-to-End Inference",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+
+    # Required arguments
+    parser.add_argument(
+        "--model-path",
+        type=str,
+        required=True,
+        help="Path to GLM-Image model directory or HuggingFace model ID",
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        required=True,
+        help="Text prompt for image generation",
+    )
+
+    # Optional arguments
+    parser.add_argument(
+        "--config-path",
+        type=str,
+        default=None,
+        help="Path to stage config YAML file (default: auto-detect)",
+    )
+    parser.add_argument(
+        "--image",
+        type=str,
+        default=None,
+        help="Path to source image for image-to-image mode",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="output_glm_image.png",
+        help="Output image path (default: output_glm_image.png)",
+    )
+    parser.add_argument(
+        "--negative-prompt",
+        type=str,
+        default=None,
+        help="Negative prompt for classifier-free guidance",
+    )
+
+    # Generation parameters
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=None,
+        help="Output image height (default: 1024 for t2i, source size for i2i)",
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=None,
+        help="Output image width (default: 1024 for t2i, source size for i2i)",
+    )
+    parser.add_argument(
+        "--num-inference-steps",
+        type=int,
+        default=50,
+        help="Number of diffusion denoising steps (default: 50)",
+    )
+    parser.add_argument(
+        "--guidance-scale",
+        type=float,
+        default=1.5,
+        help="Classifier-free guidance scale (default: 1.5)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=SEED,
+        help=f"Random seed for reproducibility (default: {SEED})",
+    )
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=16384,
+        help="Maximum tokens for AR generation (default: 16384)",
+    )
+
+    # Batch processing
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of images to generate (default: 1)",
+    )
+
+    # Runtime options
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable statistics logging",
+    )
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=600,
+        help="Timeout for stage initialization in seconds (default: 600)",
+    )
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="store_true",
+        default=False,
+        help="Enable verbose output",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/offline_inference/glm_image/run_i2i.sh

@@ -0,0 +1,93 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# Run GLM-Image image-to-image (editing) with multistage pipeline
+
+set -e
+
+# Default values
+MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}"
+CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}"
+PROMPT="${PROMPT:-Transform this image into an oil painting style}"
+INPUT_IMAGE=""
+OUTPUT="${OUTPUT:-output_i2i.png}"
+NUM_STEPS="${NUM_STEPS:-50}"
+GUIDANCE="${GUIDANCE:-1.5}"
+SEED="${SEED:-42}"
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --model-path)
+            MODEL_PATH="$2"
+            shift 2
+            ;;
+        --config-path)
+            CONFIG_PATH="$2"
+            shift 2
+            ;;
+        --prompt)
+            PROMPT="$2"
+            shift 2
+            ;;
+        --image)
+            INPUT_IMAGE="$2"
+            shift 2
+            ;;
+        --output)
+            OUTPUT="$2"
+            shift 2
+            ;;
+        --num-steps)
+            NUM_STEPS="$2"
+            shift 2
+            ;;
+        --guidance)
+            GUIDANCE="$2"
+            shift 2
+            ;;
+        --seed)
+            SEED="$2"
+            shift 2
+            ;;
+        *)
+            echo "Unknown option: $1"
+            exit 1
+            ;;
+    esac
+done
+
+# Check if input image is provided
+if [ -z "${INPUT_IMAGE}" ]; then
+    echo "Error: --image is required for image-to-image mode"
+    echo "Usage: ./run_i2i.sh --image /path/to/input.png [--prompt \"edit instruction\"]"
+    exit 1
+fi
+
+if [ ! -f "${INPUT_IMAGE}" ]; then
+    echo "Error: Input image not found: ${INPUT_IMAGE}"
+    exit 1
+fi
+
+echo "=============================================="
+echo "GLM-Image Image-to-Image Generation"
+echo "=============================================="
+echo "Model: ${MODEL_PATH}"
+echo "Config: ${CONFIG_PATH}"
+echo "Input: ${INPUT_IMAGE}"
+echo "Prompt: ${PROMPT}"
+echo "Output: ${OUTPUT}"
+echo "Steps: ${NUM_STEPS}"
+echo "Guidance: ${GUIDANCE}"
+echo "Seed: ${SEED}"
+echo "=============================================="
+
+python end2end.py \
+    --model-path "${MODEL_PATH}" \
+    --config-path "${CONFIG_PATH}" \
+    --prompt "${PROMPT}" \
+    --image "${INPUT_IMAGE}" \
+    --output "${OUTPUT}" \
+    --num-inference-steps "${NUM_STEPS}" \
+    --guidance-scale "${GUIDANCE}" \
+    --seed "${SEED}" \
+    --verbose

examples/offline_inference/glm_image/run_t2i.sh

@@ -0,0 +1,87 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# Run GLM-Image text-to-image generation with multistage pipeline
+
+set -e
+
+# Default values
+MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}"
+CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}"
+PROMPT="${PROMPT:-A beautiful mountain landscape with snow-capped peaks and a clear blue lake}"
+OUTPUT="${OUTPUT:-output_t2i.png}"
+HEIGHT="${HEIGHT:-1024}"
+WIDTH="${WIDTH:-1024}"
+NUM_STEPS="${NUM_STEPS:-50}"
+GUIDANCE="${GUIDANCE:-1.5}"
+SEED="${SEED:-42}"
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --model-path)
+            MODEL_PATH="$2"
+            shift 2
+            ;;
+        --config-path)
+            CONFIG_PATH="$2"
+            shift 2
+            ;;
+        --prompt)
+            PROMPT="$2"
+            shift 2
+            ;;
+        --output)
+            OUTPUT="$2"
+            shift 2
+            ;;
+        --height)
+            HEIGHT="$2"
+            shift 2
+            ;;
+        --width)
+            WIDTH="$2"
+            shift 2
+            ;;
+        --num-steps)
+            NUM_STEPS="$2"
+            shift 2
+            ;;
+        --guidance)
+            GUIDANCE="$2"
+            shift 2
+            ;;
+        --seed)
+            SEED="$2"
+            shift 2
+            ;;
+        *)
+            echo "Unknown option: $1"
+            exit 1
+            ;;
+    esac
+done
+
+echo "=============================================="
+echo "GLM-Image Text-to-Image Generation"
+echo "=============================================="
+echo "Model: ${MODEL_PATH}"
+echo "Config: ${CONFIG_PATH}"
+echo "Prompt: ${PROMPT}"
+echo "Output: ${OUTPUT}"
+echo "Size: ${WIDTH}x${HEIGHT}"
+echo "Steps: ${NUM_STEPS}"
+echo "Guidance: ${GUIDANCE}"
+echo "Seed: ${SEED}"
+echo "=============================================="
+
+python end2end.py \
+    --model-path "${MODEL_PATH}" \
+    --config-path "${CONFIG_PATH}" \
+    --prompt "${PROMPT}" \
+    --output "${OUTPUT}" \
+    --height "${HEIGHT}" \
+    --width "${WIDTH}" \
+    --num-inference-steps "${NUM_STEPS}" \
+    --guidance-scale "${GUIDANCE}" \
+    --seed "${SEED}" \
+    --verbose

examples/online_serving/glm_image/README.md

@@ -0,0 +1,247 @@
+# GLM-Image Online Serving
+
+This example demonstrates how to deploy GLM-Image for online image generation using vLLM-Omni.
+
+## 🛠️ Installation
+
+Please refer to [README.md](../../../README.md)
+
+## Run examples (GLM-Image)
+
+**Note**: These examples work with the default configuration on **2× NVIDIA A100 (80GB)** or equivalent. Stage 0 (AR) and Stage 1 (Diffusion) each use one GPU by default. For single-GPU setups, modify the stage configuration to share the same device.
+
+### Launch the Server
+
+```bash
+# Use default configuration
+vllm serve zai-org/GLM-Image --omni --port 8091
+```
+
+Or use the convenience script:
+
+```bash
+cd examples/online_serving/glm_image
+bash run_server.sh
+```
+
+If you have a custom stage configs file:
+
+```bash
+vllm serve zai-org/GLM-Image --omni --port 8091 --stage-configs-path /path/to/glm_image.yaml
+```
+
+### Send Requests
+
+Get into the glm_image folder:
+
+```bash
+cd examples/online_serving/glm_image
+```
+
+Send request via Python:
+
+```bash
+python openai_chat_client.py --prompt "A cute cat sitting on a window sill"
+```
+
+The Python client supports the following command-line arguments:
+
+- `--prompt` (or `-p`): Text prompt for generation (default: `A beautiful sunset over the ocean with sailing boats`)
+- `--output` (or `-o`): Output file path (default: `glm_image_output.png`)
+- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
+- `--image` (or `-i`): Input image path (for image-to-image editing)
+- `--height`: Image height in pixels (default: 1024)
+- `--width`: Image width in pixels (default: 1024)
+- `--steps`: Number of inference steps (default: 50)
+- `--guidance-scale`: Classifier-free guidance scale (default: 1.5)
+- `--seed`: Random seed (default: 42)
+- `--negative`: Negative prompt
+
+## Modality Control
+
+GLM-Image supports **text-to-image** and **image-to-image** modes.
+
+The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use the default configuration file: [`glm_image.yaml`](../../../vllm_omni/model_executor/stage_configs/glm_image.yaml)
+
+| Mode           | Input        | Output | Description                        |
+| -------------- | ------------ | ------ | ---------------------------------- |
+| Text-to-Image  | Text         | Image  | Generate images from text prompts  |
+| Image-to-Image | Image + Text | Image  | Edit images with text instructions |
+
+### Text-to-Image
+
+Generate images from text prompts:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "A photorealistic mountain landscape at sunset" \
+    --height 1024 \
+    --width 1024 \
+    --output landscape.png
+```
+
+**Using curl**
+
+```bash
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {"role": "user", "content": "A beautiful sunset over the ocean with sailing boats"}
+    ],
+    "extra_body": {
+      "height": 1024,
+      "width": 1024,
+      "num_inference_steps": 50,
+      "guidance_scale": 1.5,
+      "seed": 42
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+```
+
+Or use the script:
+
+```bash
+bash run_curl_text_to_image.sh "A futuristic city skyline at night"
+```
+
+### Image-to-Image (Image Editing)
+
+Edit images with text instructions:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "Convert this image to watercolor style" \
+    --image input.png \
+    --output watercolor.png
+```
+
+**Using curl**
+
+```bash
+IMG_B64=$(base64 < input.png | tr -d '\n')
+
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d @- <<EOF | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+{
+  "messages": [{
+    "role": "user",
+    "content": [
+      {"type": "text", "text": "Convert this image to watercolor style"},
+      {"type": "image_url", "image_url": {"url": "data:image/png;base64,'$IMG_B64'"}}
+    ]
+  }],
+  "extra_body": {
+    "height": 1024,
+    "width": 1024,
+    "num_inference_steps": 50,
+    "guidance_scale": 1.5,
+    "seed": 42
+  }
+}
+EOF
+```
+
+Or use the script:
+
+```bash
+bash run_curl_image_edit.sh input.png "Convert to watercolor style"
+```
+
+## Generation Parameters (extra_body)
+
+| Parameter             | Type  | Default | Description                         |
+| --------------------- | ----- | ------- | ----------------------------------- |
+| `height`              | int   | 1024    | Image height in pixels              |
+| `width`               | int   | 1024    | Image width in pixels               |
+| `num_inference_steps` | int   | 50      | Number of diffusion denoising steps |
+| `guidance_scale`      | float | 1.5     | Classifier-free guidance scale      |
+| `seed`                | int   | 42      | Random seed for reproducibility     |
+| `negative_prompt`     | str   | None    | Negative prompt                     |
+
+## Response Format
+
+```json
+{
+  "id": "chatcmpl-xxx",
+  "created": 1234567890,
+  "model": "zai-org/GLM-Image",
+  "choices": [
+    {
+      "index": 0,
+      "message": {
+        "role": "assistant",
+        "content": [
+          {
+            "type": "image_url",
+            "image_url": {
+              "url": "data:image/png;base64,..."
+            }
+          }
+        ]
+      },
+      "finish_reason": "stop"
+    }
+  ],
+  "usage": {}
+}
+```
+
+## Extract Image
+
+```bash
+# From a saved JSON response
+cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+```
+
+## Architecture
+
+GLM-Image uses a 2-stage multistage pipeline:
+
+```
+Stage 0 (AR Model)                Stage 1 (Diffusion)
+┌───────────────────┐            ┌─────────────────────┐
+│ vLLM-optimized    │  prior     │  GlmImagePipeline   │
+│ GlmImageFor       │──tokens──►│  ┌───────────────┐  │
+│ Conditional       │            │  │ DiT Denoiser  │  │
+│ Generation        │            │  └───────┬───────┘  │
+│ (9B AR model)     │            │          ▼          │
+└───────────────────┘            │  ┌───────────────┐  │
+        ▲                        │  │  VAE Decode   │──┼──► Image
+        │                        │  └───────────────┘  │
+   Text / Image                  └─────────────────────┘
+     Input
+```
+
+## VRAM Requirements
+
+| Stage             | VRAM                   |
+| :---------------- | :--------------------- |
+| Stage-0 (AR)      | **~18 GiB + KV Cache** |
+| Stage-1 (DiT+VAE) | **~20 GiB**            |
+| Total             | **~38 GiB + KV Cache** |
+
+## File Description
+
+| File                        | Description                           |
+| --------------------------- | ------------------------------------- |
+| `run_server.sh`             | Server startup script                 |
+| `run_curl_text_to_image.sh` | Text-to-image curl example            |
+| `run_curl_image_edit.sh`    | Image-to-image (editing) curl example |
+| `openai_chat_client.py`     | Python client (t2i + i2i)             |
+
+## FAQ
+
+- If you encounter OOM errors, adjust `gpu_memory_utilization` in the stage config:
+
+```yaml
+# In glm_image.yaml, reduce from default 0.6:
+gpu_memory_utilization: 0.5
+```
+
+- The first request may be slow due to model warmup. Subsequent requests will be faster.

examples/online_serving/glm_image/openai_chat_client.py

@@ -0,0 +1,172 @@
+#!/usr/bin/env python3
+"""
+GLM-Image OpenAI-compatible chat client for text-to-image and image-to-image.
+
+Usage:
+    # Text-to-image
+    python openai_chat_client.py --prompt "A cute cat" --output output.png
+
+    # Image-to-image (image editing)
+    python openai_chat_client.py --prompt "Convert to watercolor style" --image input.png --output output.png
+"""
+
+import argparse
+import base64
+from pathlib import Path
+
+import requests
+
+
+def generate_image(
+    prompt: str,
+    server_url: str = "http://localhost:8091",
+    image_path: str | None = None,
+    height: int = 1024,
+    width: int = 1024,
+    steps: int = 50,
+    guidance_scale: float = 1.5,
+    seed: int | None = None,
+    negative_prompt: str | None = None,
+) -> bytes | None:
+    """Generate or edit an image using the chat completions API.
+
+    Args:
+        prompt: Text description or editing instruction
+        server_url: Server URL
+        image_path: Path to input image (for image-to-image editing)
+        height: Image height in pixels
+        width: Image width in pixels
+        steps: Number of inference steps
+        guidance_scale: Classifier-free guidance scale
+        seed: Random seed for reproducibility
+        negative_prompt: Negative prompt
+
+    Returns:
+        Image bytes or None if failed
+    """
+    # Build message content
+    content: list[dict] = [{"type": "text", "text": prompt}]
+
+    if image_path:
+        img_path = Path(image_path)
+        if not img_path.exists():
+            print(f"Error: Image file not found: {image_path}")
+            return None
+        b64_data = base64.b64encode(img_path.read_bytes()).decode("utf-8")
+        suffix = img_path.suffix.lstrip(".").lower()
+        mime = {"jpg": "jpeg", "jpeg": "jpeg", "png": "png", "webp": "webp"}.get(suffix, "png")
+        content.append(
+            {
+                "type": "image_url",
+                "image_url": {"url": f"data:image/{mime};base64,{b64_data}"},
+            }
+        )
+
+    messages = [{"role": "user", "content": content}]
+
+    # Build request payload
+    extra_body: dict = {
+        "height": height,
+        "width": width,
+        "num_inference_steps": steps,
+        "guidance_scale": guidance_scale,
+    }
+    if seed is not None:
+        extra_body["seed"] = seed
+    if negative_prompt:
+        extra_body["negative_prompt"] = negative_prompt
+
+    payload = {"messages": messages, "extra_body": extra_body}
+
+    # Send request
+    try:
+        mode = "image-to-image" if image_path else "text-to-image"
+        print(f"Sending {mode} request to {server_url}...")
+        response = requests.post(
+            f"{server_url}/v1/chat/completions",
+            headers={"Content-Type": "application/json"},
+            json=payload,
+            timeout=600,
+        )
+        response.raise_for_status()
+        data = response.json()
+
+        # Extract image from response
+        choices = data.get("choices", [])
+        for choice in choices:
+            choice_content = choice.get("message", {}).get("content")
+            if isinstance(choice_content, list):
+                for item in choice_content:
+                    if isinstance(item, dict) and "image_url" in item:
+                        img_url = item["image_url"].get("url", "")
+                        if img_url.startswith("data:image"):
+                            _, b64 = img_url.split(",", 1)
+                            return base64.b64decode(b64)
+
+        print(f"Unexpected response format: {data}")
+        return None
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return None
+
+
+def main():
+    parser = argparse.ArgumentParser(description="GLM-Image chat client")
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        default="A beautiful sunset over the ocean with sailing boats",
+        help="Text prompt",
+    )
+    parser.add_argument("--output", "-o", default="glm_image_output.png", help="Output file")
+    parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
+
+    # Image-to-image
+    parser.add_argument(
+        "--image",
+        "-i",
+        type=str,
+        help="Input image path (for image-to-image editing)",
+    )
+
+    # Generation parameters
+    parser.add_argument("--height", type=int, default=1024, help="Image height")
+    parser.add_argument("--width", type=int, default=1024, help="Image width")
+    parser.add_argument("--steps", type=int, default=50, help="Inference steps")
+    parser.add_argument("--guidance-scale", type=float, default=1.5, help="CFG guidance scale")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    parser.add_argument("--negative", help="Negative prompt")
+
+    args = parser.parse_args()
+
+    mode = "image-to-image" if args.image else "text-to-image"
+    print(f"Mode: {mode}")
+    print(f"Prompt: {args.prompt}")
+    if args.image:
+        print(f"Input image: {args.image}")
+
+    image_bytes = generate_image(
+        prompt=args.prompt,
+        server_url=args.server,
+        image_path=args.image,
+        height=args.height,
+        width=args.width,
+        steps=args.steps,
+        guidance_scale=args.guidance_scale,
+        seed=args.seed,
+        negative_prompt=args.negative,
+    )
+
+    if image_bytes:
+        output_path = Path(args.output)
+        output_path.write_bytes(image_bytes)
+        print(f"Image saved to: {output_path}")
+        print(f"Size: {len(image_bytes) / 1024:.1f} KB")
+    else:
+        print("Failed to generate image")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/glm_image/run_curl_image_edit.sh

@@ -0,0 +1,61 @@
+#!/bin/bash
+# GLM-Image image-edit (image-to-image) curl example
+
+set -euo pipefail
+
+if [[ $# -lt 2 ]]; then
+  echo "Usage: $0 <input_image> \"<edit_prompt>\" [output_file]" >&2
+  exit 1
+fi
+
+INPUT_IMG=$1
+PROMPT=$2
+SERVER="${SERVER:-http://localhost:8091}"
+CURRENT_TIME=$(date +%Y%m%d%H%M%S)
+OUTPUT="${3:-glm_image_i2i_${CURRENT_TIME}.png}"
+
+if [[ ! -f "$INPUT_IMG" ]]; then
+  echo "Input image not found: $INPUT_IMG" >&2
+  exit 1
+fi
+
+# base64 encode (macOS uses -i, Linux uses -w0)
+if [[ "$(uname)" == "Darwin" ]]; then
+  IMG_B64=$(base64 < "$INPUT_IMG" | tr -d '\n')
+else
+  IMG_B64=$(base64 -w0 "$INPUT_IMG")
+fi
+
+REQUEST_JSON=$(
+  jq -n --arg prompt "$PROMPT" --arg img "$IMG_B64" '{
+    messages: [{
+      role: "user",
+      content: [
+        {"type": "text", "text": $prompt},
+        {"type": "image_url", "image_url": {"url": ("data:image/png;base64," + $img)}}
+      ]
+    }],
+    extra_body: {
+      height: 1024,
+      width: 1024,
+      num_inference_steps: 50,
+      guidance_scale: 1.5,
+      seed: 42
+    }
+  }'
+)
+
+echo "Generating edited image..."
+echo "Server: $SERVER"
+echo "Prompt: $PROMPT"
+echo "Input : $INPUT_IMG"
+echo "Output: $OUTPUT"
+
+curl -s "$SERVER/v1/chat/completions" \
+  -H "Content-Type: application/json" \
+  -d "$REQUEST_JSON" \
+  | jq -r '.choices[0].message.content[0].image_url.url' \
+  | cut -d',' -f2- \
+  | base64 -d > "$OUTPUT"
+
+echo "Image saved to: $OUTPUT"

examples/online_serving/glm_image/run_curl_text_to_image.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+# GLM-Image text-to-image curl example
+
+set -euo pipefail
+
+PROMPT="${1:-A beautiful sunset over the ocean with sailing boats}"
+SERVER="${SERVER:-http://localhost:8091}"
+OUTPUT="${OUTPUT:-glm_image_t2i_output.png}"
+
+echo "Generating image..."
+echo "Server: $SERVER"
+echo "Prompt: $PROMPT"
+echo "Output: $OUTPUT"
+
+curl -s "$SERVER/v1/chat/completions" \
+  -H "Content-Type: application/json" \
+  -d "{
+    \"messages\": [
+      {\"role\": \"user\", \"content\": \"$PROMPT\"}
+    ],
+    \"extra_body\": {
+      \"height\": 1024,
+      \"width\": 1024,
+      \"num_inference_steps\": 50,
+      \"guidance_scale\": 1.5,
+      \"seed\": 42
+    }
+  }" | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > "$OUTPUT"
+
+echo "Image saved to: $OUTPUT"

examples/online_serving/glm_image/run_server.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+# GLM-Image online serving startup script
+
+MODEL="${MODEL:-zai-org/GLM-Image}"
+PORT="${PORT:-8091}"
+
+echo "Starting GLM-Image server..."
+echo "Model: $MODEL"
+echo "Port: $PORT"
+
+vllm serve "$MODEL" --omni \
+    --port "$PORT"

vllm_omni/diffusion/models/glm_image/glm_image_transformer.py

@@ -17,7 +17,6 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.cache.base import CachedTransformer
 from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
 logger = init_logger(__name__)
 
@@ -354,8 +353,7 @@ class GlmImageAttention(nn.Module):
             nn.Dropout(0.0),
         )
 
-        # RoPE and attention
-        self.rope = RotaryEmbedding(is_neox_style=False)
+        # Attention
         self.attn = Attention(
             num_heads=num_heads,
             head_size=head_dim,
@@ -379,7 +377,6 @@ class GlmImageAttention(nn.Module):
             hidden_states: Image hidden states [B, img_seq_len, D]
             encoder_hidden_states: Text hidden states [B, text_seq_len, D]
             image_rotary_emb: Tuple of (cos, sin) for RoPE
-            attention_mask: Optional attention mask for text tokens
             kv_cache: Optional layer KV cache for image editing
             kv_cache_mode: Cache mode (WRITE, READ, SKIP)
 
@@ -407,16 +404,12 @@ class GlmImageAttention(nn.Module):
 
         # Apply RoPE only to image tokens (not text tokens)
         if image_rotary_emb is not None:
-            cos, sin = image_rotary_emb
-            cos = cos.to(query.dtype)
-            sin = sin.to(query.dtype)
             # Only apply RoPE to image part (after text_seq_length)
             query_img = query[:, text_seq_length:, :, :]
             key_img = key[:, text_seq_length:, :, :]
             from diffusers.models.embeddings import apply_rotary_emb
 
             query_img = apply_rotary_emb(query_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2)
-            # key_img = self.rope(key_img, cos, sin)
             key_img = apply_rotary_emb(key_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2)
             query = torch.cat([query[:, :text_seq_length, :, :], query_img], dim=1)
             key = torch.cat([key[:, :text_seq_length, :, :], key_img], dim=1)

vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py

@@ -4,7 +4,7 @@
 GlmImagePipeline implementation for vLLM-Omni.
 
 This pipeline implements GLM-Image text-to-image generation with:
-- AR stage: GlmImageForConditionalGeneration generates prior tokens
+- AR stage (vLLM): GLM-Image AR stage generates prior tokens
 - DiT stage: GlmImageTransformer2DModel performs diffusion denoising
 - VAE: AutoencoderKL decodes latents to images
 """
@@ -31,8 +31,6 @@ from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import (
     ByT5Tokenizer,
-    GlmImageForConditionalGeneration,
-    GlmImageProcessor,
     T5EncoderModel,
 )
 
@@ -244,13 +242,13 @@ class GlmImagePipeline(nn.Module):
     GLM-Image Pipeline for text-to-image and image-to-image generation.
 
     This pipeline integrates:
-    - AR model (GlmImageForConditionalGeneration): Generates prior image tokens
+    - AR stage (vLLM): Generates prior image tokens
     - Text encoder (T5EncoderModel): Encodes glyph/text embeddings
     - DiT model (GlmImageTransformer2DModel): Diffusion transformer
     - VAE (AutoencoderKL): Encodes/decodes images to/from latent space
 
     The pipeline flow:
-    1. AR generates prior_token_ids from text prompt
+    1. AR stage provides prior_token_ids (and optionally prior_token_image_ids)
     2. T5 encodes glyph text for text rendering
     3. DiT performs iterative denoising conditioned on prior tokens
     4. VAE decodes final latents to image
@@ -280,19 +278,6 @@ class GlmImagePipeline(nn.Module):
             model_path, subfolder="scheduler", local_files_only=True
         )
 
-        # Load AR model (vision_language_encoder)
-        logger.info("Loading GlmImageForConditionalGeneration (AR model)...")
-        self.vision_language_encoder = GlmImageForConditionalGeneration.from_pretrained(
-            model_path,
-            subfolder="vision_language_encoder",
-            local_files_only=True,
-            torch_dtype=torch.bfloat16,
-        ).to(self.device)
-        self.vision_language_encoder.eval()
-
-        # Load processor for AR model
-        self.processor = GlmImageProcessor.from_pretrained(model_path, subfolder="processor", local_files_only=True)
-
         # Load text encoder (T5 for glyph embeddings)
         logger.info("Loading T5EncoderModel (glyph encoder)...")
         self.text_encoder = T5EncoderModel.from_pretrained(
@@ -371,159 +356,6 @@ class GlmImagePipeline(nn.Module):
         if prompt is not None and not isinstance(prompt, (str, list)):
             raise ValueError(f"`prompt` must be of type `str` or `list` but is {type(prompt)}")
 
-    # ==================== AR Stage Methods ====================
-
-    @staticmethod
-    def _compute_generation_params(
-        image_grid_thw: torch.Tensor,
-        is_text_to_image: bool,
-    ) -> tuple[int, int, int, int]:
-        """
-        Compute AR generation parameters from image grid.
-
-        Args:
-            image_grid_thw: Image grid tensor of shape [N, 3] where each row is [t, h, w]
-            is_text_to_image: Whether this is text-to-image (vs image-to-image)
-
-        Returns:
-            Tuple of (max_new_tokens, large_image_start_offset, target_grid_h, target_grid_w)
-        """
-        grid_sizes = []
-        grid_hw = []
-
-        for i in range(image_grid_thw.shape[0]):
-            t, h, w = image_grid_thw[i].tolist()
-            grid_sizes.append(int(h * w))
-            grid_hw.append((int(h), int(w)))
-
-        if not is_text_to_image:
-            # Image-to-image: only generate target image tokens
-            max_new_tokens = grid_sizes[-1] + 1
-            large_image_start_offset = 0
-            target_grid_h, target_grid_w = grid_hw[-1]
-        else:
-            # Text-to-image: generate both small preview and large target
-            total_tokens = sum(grid_sizes)
-            max_new_tokens = total_tokens + 1
-            large_image_start_offset = sum(grid_sizes[1:])
-            target_grid_h, target_grid_w = grid_hw[0]
-
-        return max_new_tokens, large_image_start_offset, target_grid_h, target_grid_w
-
-    @staticmethod
-    def _extract_large_image_tokens(
-        outputs: torch.Tensor, input_length: int, large_image_start_offset: int, large_image_tokens: int
-    ) -> torch.Tensor:
-        """Extract large image tokens from AR output."""
-        generated_tokens = outputs[0][input_length:]
-        large_image_start = large_image_start_offset
-        large_image_end = large_image_start + large_image_tokens
-        return generated_tokens[large_image_start:large_image_end]
-
-    @staticmethod
-    def _upsample_token_ids(token_ids: torch.Tensor, token_h: int, token_w: int) -> torch.Tensor:
-        """Upsample token IDs by 2x using nearest neighbor interpolation."""
-        token_ids = token_ids.view(1, 1, token_h, token_w)
-        token_ids = torch.nn.functional.interpolate(token_ids.float(), scale_factor=2, mode="nearest").to(
-            dtype=torch.long
-        )
-        token_ids = token_ids.view(1, -1)
-        return token_ids
-
-    @torch.inference_mode()
-    def generate_prior_tokens(
-        self,
-        prompt: str,
-        height: int,
-        width: int,
-        image: list[PIL.Image.Image] | None = None,
-        factor: int = 32,
-    ) -> tuple[torch.Tensor, list[torch.Tensor] | None]:
-        """
-        Generate prior tokens using the AR model.
-
-        Args:
-            prompt: Text prompt for generation
-            height: Target image height
-            width: Target image width
-            image: Optional condition images for image-to-image
-            factor: Token factor (default 32)
-
-        Returns:
-            Tuple of (prior_token_ids, prior_token_image_ids)
-            prior_token_image_ids is a list of tensors, one per condition image
-        """
-        device = self.vision_language_encoder.device
-        height = (height // factor) * factor
-        width = (width // factor) * factor
-        is_text_to_image = image is None or len(image) == 0
-
-        # Build message content
-        content = []
-        if image is not None:
-            for img in image:
-                content.append({"type": "image", "image": img})
-        content.append({"type": "text", "text": prompt})
-        messages = [{"role": "user", "content": content}]
-
-        # Apply chat template - processor will handle target dimensions and build grid
-        inputs = self.processor.apply_chat_template(
-            messages,
-            tokenize=True,
-            target_h=height,
-            target_w=width,
-            return_dict=True,
-            return_tensors="pt",
-        ).to(device)
-
-        image_grid_thw = inputs.get("image_grid_thw")
-
-        # Compute generation parameters from the full grid
-        max_new_tokens, large_image_offset, token_h, token_w = self._compute_generation_params(
-            image_grid_thw=image_grid_thw, is_text_to_image=is_text_to_image
-        )
-
-        # Process condition images if provided
-        # Use image_grid_thw[:-1] to exclude the target image grid (last entry)
-        prior_token_image_ids = None
-        if image is not None and image_grid_thw is not None and len(image_grid_thw) > 1:
-            # Get features only for condition images (exclude target image grid)
-            condition_grid = image_grid_thw[:-1]
-            prior_token_image_embed = self.vision_language_encoder.get_image_features(
-                inputs["pixel_values"], condition_grid
-            ).pooler_output
-            prior_token_image_embed = torch.cat(prior_token_image_embed, dim=0)
-            flat_prior_token_image_ids = self.vision_language_encoder.get_image_tokens(
-                prior_token_image_embed, condition_grid
-            )
-            # Split by image grid sizes and convert to list
-            split_sizes = (condition_grid.prod(dim=-1)).tolist()
-            prior_token_image_ids_list = torch.split(flat_prior_token_image_ids, split_sizes, dim=0)
-            # Convert to list with upsampling
-            prior_token_image_ids = []
-            for i, token_ids in enumerate(prior_token_image_ids_list):
-                grid_t, grid_h, grid_w = condition_grid[i].tolist()
-                token_ids = token_ids.view(1, -1)
-                # Upsample 2x (from d32 to d64)
-                token_ids_upsampled = self._upsample_token_ids(token_ids, grid_h, grid_w)
-                prior_token_image_ids.append(token_ids_upsampled)
-
-        # Generate with AR model
-        outputs = self.vision_language_encoder.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            do_sample=True,
-        )
-
-        # Extract and upsample tokens
-        large_image_tokens = token_h * token_w
-        prior_token_ids_d32 = self._extract_large_image_tokens(
-            outputs, inputs["input_ids"].shape[-1], large_image_offset, large_image_tokens
-        )
-        prior_token_ids = self._upsample_token_ids(prior_token_ids_d32, token_h, token_w)
-
-        return prior_token_ids, prior_token_image_ids
-
     # ==================== Text Encoding Methods ====================
 
     def get_glyph_texts(self, prompt: str | list[str]) -> list[str]:
@@ -803,11 +635,14 @@ class GlmImagePipeline(nn.Module):
         # Process each condition image through transformer to populate KV cache
         for condition_image, condition_prior_token_id in zip(condition_images, prior_token_image_ids):
             condition_image = condition_image.to(device=self.device, dtype=prompt_embeds.dtype)
+            # Move prior token ids to device (may come from CPU due to serialization)
+            if isinstance(condition_prior_token_id, torch.Tensor):
+                condition_prior_token_id = condition_prior_token_id.to(device=self.device)
 
             # Encode condition image to latent space
             # Use argmax (mode) for deterministic encoding of condition images
             condition_latent = retrieve_latents(
-                self.vae.encode(condition_image), generator=generator, sample_mode="argmax"
+                self.vae.encode(condition_image.unsqueeze(0)), generator=generator, sample_mode="argmax"
             )
             condition_latent = (condition_latent - latents_mean) / latents_std
 
@@ -846,6 +681,17 @@ class GlmImagePipeline(nn.Module):
         first_prompt = req.prompts[0]
         prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
 
+        # NOTE: DiffusionEngine does an internal warmup "dummy run" during
+        # initialization. That request has no request_id and does not carry
+        # Stage-0 (AR) outputs via req.extra. For GLM-Image, we allow that
+        # specific warmup request to proceed by synthesizing minimal prior
+        # tokens, while still raising a clear error for real requests.
+        is_dummy_warmup = (
+            not getattr(req, "request_id", None)
+            and prompt == "dummy run"
+            and (req.sampling_params.num_inference_steps == 1)
+        )
+
         # Get pre-computed prompt embeddings if provided
         if isinstance(first_prompt, str):
             prompt_embeds = None
@@ -856,20 +702,18 @@ class GlmImagePipeline(nn.Module):
 
         # Get condition images for Image Edit mode
         # Use pre-processed images from pre_process_func
-        preprocessed_images = (
-            None
-            if isinstance(first_prompt, str)
-            else [first_prompt.get("additional_information", {}).get("preprocessed_image")]
-        )
-        condition_images = (
-            None
-            if isinstance(first_prompt, str)
-            else first_prompt.get("additional_information", {}).get("prompt_image")
-        )
+        preprocessed_images = None
+        if not isinstance(first_prompt, str):
+            img = first_prompt.get("additional_information", {}).get("preprocessed_image")
+            if img is not None:
+                preprocessed_images = [img]
+
         img_height = req.sampling_params.height
         img_width = req.sampling_params.width
 
-        is_image_edit = preprocessed_images is not None
+        # Warmup may include a dummy image for image-input-capable models.
+        # Treat that as t2i warmup to avoid requiring i2i-only KV-cache inputs.
+        is_image_edit = (preprocessed_images is not None) and (not is_dummy_warmup)
 
         # Use image dimensions as default if available
         height = req.sampling_params.height or img_height or self.default_sample_size * self.vae_scale_factor
@@ -877,46 +721,67 @@ class GlmImagePipeline(nn.Module):
         num_inference_steps = req.sampling_params.num_inference_steps or 50
         guidance_scale = req.sampling_params.guidance_scale or 1.5
 
-        # 0. Validate inputs
         self.check_inputs(prompt=prompt, height=height, width=width, prompt_embeds=prompt_embeds)
 
         batch_size = 1
         do_classifier_free_guidance = guidance_scale > 1.0
 
-        # Set seed if provided
         generator = None
         if req.sampling_params.seed is not None:
             generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed)
 
         # 1. Get prior tokens - either from external source (multistage) or generate internally
         # Check if prior_token_ids are provided externally (from AR stage in multistage mode)
+        # First try sampling_params.extra_args (for direct API usage)
+        # Then try prompt["extra"] (for multistage ar2diffusion output)
         external_prior_tokens = req.sampling_params.extra_args.get("prior_token_ids")
         external_prior_image_ids = req.sampling_params.extra_args.get("prior_token_image_ids")
+        if external_prior_tokens is None and isinstance(first_prompt, dict):
+            prompt_extra = first_prompt.get("extra", {})
+            external_prior_tokens = prompt_extra.get("prior_token_ids")
+            external_prior_image_ids = prompt_extra.get("prior_token_image_ids")
 
         if external_prior_tokens is not None:
-            # Multistage mode: use externally provided prior tokens from vLLM AR stage
-            logger.info("Using externally provided prior tokens from AR stage...")
             prior_token_id = external_prior_tokens
             if isinstance(prior_token_id, list):
                 prior_token_id = torch.tensor(prior_token_id, dtype=torch.long, device=self.device)
             elif isinstance(prior_token_id, torch.Tensor):
                 prior_token_id = prior_token_id.to(device=self.device, dtype=torch.long)
-            # Ensure shape is [1, num_tokens] for batch processing
             if prior_token_id.dim() == 1:
                 prior_token_id = prior_token_id.unsqueeze(0)
-            prior_token_image_ids = external_prior_image_ids
+            prior_token_image_ids = None
+            if external_prior_image_ids is not None:
+                if isinstance(external_prior_image_ids, torch.Tensor):
+                    prior_token_image_ids = [external_prior_image_ids]
+                else:
+                    prior_token_image_ids = list(external_prior_image_ids)
+
+                normalized: list[torch.Tensor] = []
+                for item in prior_token_image_ids:
+                    if isinstance(item, torch.Tensor):
+                        tensor_item = item.to(device=self.device, dtype=torch.long)
+                    else:
+                        tensor_item = torch.tensor(item, dtype=torch.long, device=self.device)
+                    if tensor_item.dim() == 1:
+                        tensor_item = tensor_item.unsqueeze(0)
+                    normalized.append(tensor_item)
+                prior_token_image_ids = normalized
+        elif is_dummy_warmup:
+            # Synthesize a minimal, shape-correct prior token sequence.
+            # The diffusion transformer expects prior_token_id shape:
+            #   [B, (H_lat/p)*(W_lat/p)] where H_lat=H/vae_scale_factor.
+            h_lat = height // self.vae_scale_factor
+            w_lat = width // self.vae_scale_factor
+            seq_len = (h_lat * w_lat) // (self._patch_size**2)
+            prior_token_id = torch.zeros((1, seq_len), dtype=torch.long, device=self.device)
+            prior_token_image_ids = None
         else:
-            # Single-stage mode: generate prior tokens with internal AR model
-            logger.info("Generating prior tokens with AR model...")
-            prior_token_id, prior_token_image_ids = self.generate_prior_tokens(
-                prompt=prompt,
-                image=condition_images,
-                height=height,
-                width=width,
+            raise ValueError(
+                "GLM-Image diffusion pipeline expects prior tokens from the vLLM AR stage. "
+                "Pass them via req.extra['prior_token_ids'] (and for i2i, req.extra['prior_token_image_ids'])."
             )
 
         # 2. Encode prompt for glyph embeddings
-        logger.info("Encoding prompt...")
         prompt_embeds, negative_prompt_embeds = self.encode_prompt(
             prompt,
             do_classifier_free_guidance=do_classifier_free_guidance,
@@ -928,15 +793,17 @@ class GlmImagePipeline(nn.Module):
 
         # 3. Prepare KV cache for Image Edit mode
         kv_caches = None
+        if is_image_edit and prior_token_image_ids is None:
+            raise ValueError(
+                "Image edit (i2i) requires req.extra['prior_token_image_ids'] from the AR stage to build KV cache."
+            )
         if is_image_edit and prior_token_image_ids is not None:
-            logger.info("Preparing KV cache for Image Edit mode...")
             kv_caches = self._prepare_condition_image_kv_cache(
                 condition_images=preprocessed_images,
                 prior_token_image_ids=prior_token_image_ids,
                 prompt_embeds=prompt_embeds,
                 generator=generator,
             )
-            # Switch to read mode for denoising
             kv_caches.set_mode("read")
 
         # 4. Prepare latents
@@ -971,8 +838,7 @@ class GlmImagePipeline(nn.Module):
         target_size = torch.tensor([[height, width]], dtype=prompt_embeds.dtype, device=self.device)
         crop_coords = torch.zeros((1, 2), dtype=prompt_embeds.dtype, device=self.device)
 
-        # 7. Denoising loop with CFG-parallel support
-        logger.info(f"Starting denoising loop with {num_inference_steps} steps...")
+        # 7. Denoising loop
         latents = self.diffuse(
             latents=latents,
             prior_token_id=prior_token_id,
@@ -987,7 +853,6 @@ class GlmImagePipeline(nn.Module):
         )
 
         # 8. VAE decode
-        logger.info("Decoding latents with VAE...")
         latents = latents.to(self.vae.dtype)
         latents_mean = (
             torch.tensor(self.vae.config.latents_mean)
@@ -1002,8 +867,6 @@ class GlmImagePipeline(nn.Module):
         latents = latents * latents_std + latents_mean
         image = self.vae.decode(latents, return_dict=False, generator=generator)[0]
 
-        # 9. Leave post-process to vllm-omni pipeline
-
         return DiffusionOutput(output=image)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm_omni/entrypoints/omni_diffusion.py

@@ -47,8 +47,9 @@ class OmniDiffusion:
         if engine_input_source is not None:
             self.od_config.omni_kv_config.setdefault("engine_input_source", engine_input_source)
 
+        # Detect model class and load config
         # Diffusers-style models expose `model_index.json` with `_class_name`.
-        # Non-diffusers models (e.g. Bagel, NextStep) only have `config.json`,
+        # Non-diffusers models (e.g. Bagel, NextStep, GLM-Image) only have `config.json`,
         # so we fall back to reading that and mapping model_type manually.
         try:
             config_dict = get_hf_file_to_dict(
@@ -71,6 +72,7 @@ class OmniDiffusion:
             if cfg is None:
                 raise ValueError(f"Could not find config.json or model_index.json for model {od_config.model}")
 
+            # Map model_type or architecture to pipeline class
             model_type = cfg.get("model_type")
             architectures = cfg.get("architectures") or []
             # Bagel/NextStep models don't have a model_index.json, so we set the pipeline class name manually
@@ -83,10 +85,17 @@ class OmniDiffusion:
                     od_config.model_class_name = "NextStep11Pipeline"
                 od_config.tf_model_config = TransformerConfig()
                 od_config.update_multimodal_support()
+            elif model_type == "glm-image" or "GlmImageForConditionalGeneration" in architectures:
+                pipeline_class = "GlmImagePipeline"
             elif architectures and len(architectures) == 1:
-                od_config.model_class_name = architectures[0]
-            else:
-                raise
+                pipeline_class = architectures[0]
+
+            if pipeline_class is None:
+                raise ValueError(f"Unknown model type: {model_type}, architectures: {architectures}")
+
+            od_config.model_class_name = pipeline_class
+            od_config.tf_model_config = TransformerConfig()
+            od_config.update_multimodal_support()
 
         self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config)
 

vllm_omni/entrypoints/omni_stage.py

@@ -44,6 +44,7 @@ from vllm_omni.entrypoints.omni_llm import OmniLLM
 from vllm_omni.entrypoints.stage_utils import (
     SHUTDOWN_TASK,
     OmniStageTaskType,
+    _resolve_model_tokenizer_paths,
     _to_dict,
     is_profiler_task,
     maybe_dump_to_shm,
@@ -251,7 +252,7 @@ class OmniStage:
     """
 
     def __init__(self, stage_config: Any, stage_init_timeout: int = 300):
-        logger.info(f"[OmniStage] stage_config: {stage_config}")
+        logger.debug(f"[OmniStage] stage_config: {stage_config}")
         self.stage_config = stage_config
         self.engine = None
         self.async_engine = None
@@ -708,6 +709,9 @@ def _stage_worker(
     if stage_type != "diffusion":
         _resolve_worker_cls(engine_args)
 
+    # Handle non-standard model directory structures (e.g., tokenizer in root, model in subdir)
+    model = _resolve_model_tokenizer_paths(model, engine_args)
+
     # Resolve ZMQ queue endpoints if needed
     zmq_ctx = None
     if isinstance(in_q, str) or isinstance(out_q, str):
@@ -933,8 +937,9 @@ def _stage_worker(
             _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
 
             batch_request_ids.append(rid)
-            if isinstance(ein, (str, dict)):
-                # Types like OmniTextPrompt, TextPrompt are TypedDict, essentially dict and enters this branch
+
+            if isinstance(ein, (dict, str)):
+                # For diffusion stage-0, ein might be a string prompt directly
                 batch_engine_inputs.append(ein)
             elif isinstance(ein, Sequence):
                 batch_engine_inputs.extend(ein)
@@ -1108,6 +1113,9 @@ async def _stage_worker_async(
     final_output = stage_payload.get("final_output", False)
     final_output_type = stage_payload.get("final_output_type", None)
 
+    # Handle non-standard model directory structures (e.g., tokenizer in root, model in subdir)
+    model = _resolve_model_tokenizer_paths(model, engine_args)
+
     if stage_type != "diffusion":
         _resolve_worker_cls(engine_args)
 

vllm_omni/entrypoints/openai/serving_chat.py

@@ -274,6 +274,80 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             output_modalities if output_modalities is not None else self.engine_client.output_modalities
         )
 
+        # Omni multistage image generation: Stage-0 (AR) should receive a clean
+        # text prompt (and optional conditioning image/size) so the model's own
+        # processor can construct the correct inputs.
+        # If we pass pre-tokenized chat-template ids, GLM-Image can become
+        # effectively unconditioned and produce nonsense images.
+        if request.modalities and ("image" in request.modalities):
+            try:
+                messages_as_dicts: list[dict[str, Any]] = []
+                for msg in request.messages:
+                    if hasattr(msg, "model_dump"):
+                        messages_as_dicts.append(msg.model_dump())
+                    elif isinstance(msg, dict):
+                        messages_as_dicts.append(msg)
+                    else:
+                        messages_as_dicts.append(
+                            {
+                                "role": getattr(msg, "role", "user"),
+                                "content": getattr(msg, "content", ""),
+                            }
+                        )
+                extracted_prompt, reference_images = self._extract_diffusion_prompt_and_images(messages_as_dicts)
+                if not extracted_prompt:
+                    return self.create_error_response("No text prompt found in messages")
+
+                extra_body = getattr(request, "extra_body", None) or {}
+                height = extra_body.get("height")
+                width = extra_body.get("width")
+                if "size" in extra_body:
+                    try:
+                        size_str = extra_body["size"]
+                        if isinstance(size_str, str) and "x" in size_str.lower():
+                            w, h = size_str.lower().split("x")
+                            width, height = int(w), int(h)
+                    except Exception:
+                        pass
+                negative_prompt = extra_body.get("negative_prompt")
+
+                engine_prompt_image: dict[str, Any] | None = None
+                if reference_images:
+                    # Best-effort decode first reference image for i2i.
+                    try:
+                        img_bytes = base64.b64decode(reference_images[0])
+                        img = Image.open(BytesIO(img_bytes))
+                        engine_prompt_image = {"image": img}
+                    except Exception:
+                        engine_prompt_image = None
+
+                # Override the prompts produced by chat-template preprocessing.
+                tprompt: OmniTextPrompt = {"prompt": extracted_prompt}
+                if negative_prompt is not None:
+                    tprompt["negative_prompt"] = negative_prompt
+                # GLM-Image's _call_hf_processor expects target_h/target_w in mm_processor_kwargs
+                mm_processor_kwargs: dict[str, Any] = {}
+                if height is not None:
+                    mm_processor_kwargs["target_h"] = height
+                if width is not None:
+                    mm_processor_kwargs["target_w"] = width
+                if mm_processor_kwargs:
+                    tprompt["mm_processor_kwargs"] = mm_processor_kwargs
+                if engine_prompt_image is not None:
+                    tprompt["multi_modal_data"] = engine_prompt_image
+
+                engine_prompts = [tprompt]
+                # Store height/width for applying to diffusion stage sampling params later
+                _image_gen_height = height
+                _image_gen_width = width
+            except Exception as e:
+                logger.warning("Failed to build image-generation prompt for omni multistage: %s", e)
+                _image_gen_height = None
+                _image_gen_width = None
+        else:
+            _image_gen_height = None
+            _image_gen_width = None
+
         # Schedule the request and get the result generator.
         generators: list[AsyncGenerator[RequestOutput, None]] = []
         try:
@@ -284,6 +358,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     # Use standard OpenAI API parameters for comprehension stage
                     sampling_params_list = self._build_sampling_params_list_from_request(request)
 
+                # Apply user-specified height/width to diffusion stage(s) for image generation
+                if _image_gen_height is not None or _image_gen_width is not None:
+                    for idx, sp in enumerate(sampling_params_list):
+                        # Diffusion stages typically have height/width attributes
+                        if hasattr(sp, "height") and _image_gen_height is not None:
+                            sp.height = _image_gen_height
+                        if hasattr(sp, "width") and _image_gen_width is not None:
+                            sp.width = _image_gen_width
+
                 self._log_inputs(
                     request_id,
                     engine_prompt,
@@ -420,6 +503,35 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 request=request
             )
 
+        # Preserve a clean text prompt for downstream stages (e.g., GLM-Image diffusion).
+        # For /v1/chat/completions, `request_prompt` is often the rendered chat template.
+        # Diffusion models generally want the raw user caption instead.
+        output_modalities = getattr(self.engine_client, "output_modalities", None)
+        if output_modalities and ("image" in output_modalities):
+            messages_as_dicts: list[dict[str, Any]] = []
+            for msg in messages:
+                if hasattr(msg, "model_dump"):
+                    messages_as_dicts.append(msg.model_dump())
+                elif isinstance(msg, dict):
+                    messages_as_dicts.append(msg)
+                else:
+                    messages_as_dicts.append(
+                        {
+                            "role": getattr(msg, "role", "user"),
+                            "content": getattr(msg, "content", ""),
+                        }
+                    )
+            extracted_prompt, _ = self._extract_diffusion_prompt_and_images(messages_as_dicts)
+            if extracted_prompt:
+                engine_prompt["prompt"] = extracted_prompt
+
+        mm_processor_kwargs = getattr(request, "mm_processor_kwargs", None)
+        if mm_processor_kwargs is not None:
+            engine_prompt["mm_processor_kwargs"] = mm_processor_kwargs
+
+        if hasattr(request, "cache_salt") and request.cache_salt is not None:
+            engine_prompt["cache_salt"] = request.cache_salt
+
         return conversation, [engine_prompt]
 
     def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[SamplingParams]:

vllm_omni/entrypoints/stage_utils.py

@@ -302,3 +302,41 @@ def _to_dict(x: Any) -> dict[str, Any]:
             return dict(x)
         except Exception:
             return {}
+
+
+def _resolve_model_tokenizer_paths(
+    model: str,
+    engine_args: dict,
+) -> str:
+    """Resolve model and tokenizer paths for non-standard directory structures.
+
+    Some models (e.g., GLM-Image) have tokenizer in root and model in subdirectory.
+    This function handles model_subdir and tokenizer_subdir engine_args.
+
+    Args:
+        model: Base model path
+        engine_args: Engine arguments (modified in-place to remove subdir args
+            and set tokenizer if needed)
+
+    Returns:
+        Resolved model path (may be subdirectory of original)
+    """
+    import os
+
+    model_subdir = engine_args.pop("model_subdir", None)
+    tokenizer_subdir = engine_args.pop("tokenizer_subdir", None)
+    base_model_path = model
+
+    if model_subdir:
+        model = os.path.join(model, model_subdir)
+        logger.info(f"Using model subdirectory: {model}")
+
+    if tokenizer_subdir is not None:
+        tokenizer_path = os.path.join(base_model_path, tokenizer_subdir) if tokenizer_subdir else base_model_path
+        engine_args["tokenizer"] = tokenizer_path
+        logger.info(f"Using tokenizer from: {tokenizer_path}")
+    elif model_subdir and "tokenizer" not in engine_args:
+        engine_args["tokenizer"] = base_model_path
+        logger.info(f"Using tokenizer from base model path: {base_model_path}")
+
+    return model

vllm_omni/inputs/preprocess.py

@@ -33,23 +33,41 @@ class OmniInputPreprocessor(InputPreprocessor):
         *,
         mm_uuids: MultiModalUUIDDict | None = None,
     ) -> OmniTokenInputs | MultiModalInputs:
+        """Process text prompts with support for mm_processor_kwargs.
+
+        Extends base class to support mm_processor_kwargs without multi_modal_data.
+        This is needed for models like GLM-Image where text-to-image generation
+        requires processor kwargs (target_h, target_w) to format the prompt.
+        """
         prompt_text = parsed_content["prompt"]
+        mm_processor_kwargs = parsed_content.get("mm_processor_kwargs") or {}
 
         inputs: OmniTokenInputs | MultiModalInputs
         if multi_modal_data := parsed_content.get("multi_modal_data"):
             inputs = self._process_multimodal(
                 prompt_text,
                 multi_modal_data,
-                parsed_content.get("mm_processor_kwargs") or {},
+                mm_processor_kwargs,
                 tokenization_kwargs=tokenization_kwargs,
                 mm_uuids=mm_uuids,
             )
+            # Preserve prompt_embeds and additional_information
             prompt_embeds = parsed_content.get("prompt_embeds")
             if prompt_embeds is not None:
                 inputs["prompt_embeds"] = prompt_embeds
             additional_information = parsed_content.get("additional_information")
             if additional_information is not None:
                 inputs["additional_information"] = additional_information
+        elif mm_processor_kwargs:
+            # Support mm_processor_kwargs without multi_modal_data
+            # Use case: GLM-Image text-to-image needs processor to generate grid tokens
+            inputs = self._process_multimodal(
+                prompt_text,
+                {},  # Empty multi_modal_data
+                mm_processor_kwargs,
+                tokenization_kwargs=tokenization_kwargs,
+                mm_uuids=mm_uuids,
+            )
         else:
             prompt_token_ids = self._tokenize_prompt(
                 prompt_text,
@@ -77,8 +95,10 @@ class OmniInputPreprocessor(InputPreprocessor):
         prompt_embeds = parsed_content.get("prompt_embeds")
         additional_information = parsed_content.get("additional_information")
 
+        multi_modal_data = parsed_content.get("multi_modal_data")
+
         inputs: OmniTokenInputs | MultiModalInputs
-        if multi_modal_data := parsed_content.get("multi_modal_data"):
+        if multi_modal_data:
             inputs = self._process_multimodal(
                 prompt_token_ids,
                 multi_modal_data,

vllm_omni/model_executor/models/glm_image/__init__.py

@@ -0,0 +1,3 @@
+from .glm_image_ar import GlmImageForConditionalGeneration
+
+__all__ = ["GlmImageForConditionalGeneration"]

vllm_omni/model_executor/models/registry.py

@@ -63,6 +63,11 @@ _OMNI_MODELS = {
         "qwen3_tts_code2wav",
         "Qwen3TTSCode2Wav",
     ),
+    "GlmImageForConditionalGeneration": (
+        "glm_image",
+        "glm_image_ar",
+        "GlmImageForConditionalGeneration",
+    ),
 }
 
 

vllm_omni/model_executor/stage_configs/glm_image.yaml

@@ -0,0 +1,80 @@
+# 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"
+      max_batch_size: 1
+      requires_multimodal_data: true # Required for i2i mode to receive source images
+    engine_args:
+      model_stage: ar
+      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.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.6
+      enforce_eager: false
+      trust_remote_code: true
+      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: 1281 # For 1024x1024: small(16x16=256) + large(32x32=1024) + EOS(1)
+      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
+      max_batch_size: 1
+      requires_multimodal_data: true # Required for i2i mode to pass condition images
+    engine_args:
+      model_stage: dit
+      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

vllm_omni/model_executor/stage_configs/glm_image_muilticonnector.yaml

@@ -0,0 +1,93 @@
+# Stage config for running GLM-Image with 2-stage architecture (MultiConnector version)
+# Stage 0: AR Model (vLLM implementation) - generates prior_token_ids
+# Stage 1: Diffusion (DiT + VAE) - denoising and image decoding
+#
+# This config uses OmniConnectors for inter-stage communication,
+# enabling efficient tensor transfer between stages on different processes/nodes.
+
+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"
+      max_batch_size: 1
+      requires_multimodal_data: true # Required for i2i mode to receive source images
+    engine_args:
+      model_stage: ar
+      model_arch: GlmImageForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.6
+      enforce_eager: false
+      trust_remote_code: true
+      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
+      hf_config_name: vision_language_encoder # Subfolder in model path
+    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: 1281 # For 1024x1024: small(16x16=256) + large(32x32=1024) + EOS(1)
+      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" # Use separate GPU for diffusion
+      max_batch_size: 1
+      requires_multimodal_data: true # Required for i2i mode to pass condition images
+    engine_args:
+      model_stage: dit
+      # 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:
+      seed: 42
+      num_inference_steps: 50
+      guidance_scale: 1.5
+      height: 1024
+      width: 1024
+
+# Top-level runtime config with MultiConnector support
+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
+      to: 1
+      window_size: -1
+
+# OmniConnector configuration for efficient inter-stage tensor transfer
+connectors:
+  - type: tensor_transfer
+    source_stage: 0
+    target_stage: 1
+    # Transfer prior_token_ids efficiently between stages
+    fields:
+      - name: prior_token_ids
+        dtype: int64
+      - name: prior_token_image_ids
+        dtype: int64
+        optional: true

vllm_omni/model_executor/stage_input_processors/glm_image.py

@@ -0,0 +1,265 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage input processor for GLM-Image: AR → Diffusion transition."""
+
+from typing import Any
+
+import torch
+from vllm.inputs import TextPrompt
+from vllm.logger import init_logger
+
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+logger = init_logger(__name__)
+
+
+def _upsample_token_ids(token_ids: torch.Tensor, token_h: int, token_w: int) -> torch.Tensor:
+    """Upsample token IDs by 2x using nearest neighbor interpolation.
+
+    GLM-Image AR model generates tokens at 32x downsampling, but DiT expects
+    16x downsampling, so we need to upsample by 2x.
+
+    Args:
+        token_ids: Prior token IDs of shape [num_tokens]
+        token_h: Height in token space (at 32x downsampling)
+        token_w: Width in token space (at 32x downsampling)
+
+    Returns:
+        Upsampled token IDs of shape [num_tokens * 4]
+    """
+    token_ids = token_ids.view(1, 1, token_h, token_w)
+    token_ids = torch.nn.functional.interpolate(token_ids.float(), scale_factor=2, mode="nearest").to(dtype=torch.long)
+    token_ids = token_ids.view(-1)
+    return token_ids
+
+
+def _parse_generated_tokens(
+    token_ids: list[int],
+    height: int,
+    width: int,
+    factor: int = 32,
+    is_i2i: bool = False,
+) -> tuple[torch.Tensor, int, int]:
+    """Parse AR-generated tokens to extract prior_token_ids.
+
+    Args:
+        token_ids: Generated token IDs from AR model
+        height: Target image height
+        width: Target image width
+        factor: Downsampling factor (default 32)
+        is_i2i: Whether this is image-to-image mode. In i2i mode, the AR model
+                generates only large image tokens (no small preview tokens).
+    """
+    # Calculate token dimensions for target image
+    token_h = height // factor
+    token_w = width // factor
+    large_image_tokens = token_h * token_w
+
+    # Calculate small preview image dimensions (used in text-to-image)
+    small_token_h = token_h // 2
+    small_token_w = token_w // 2
+    small_image_tokens = small_token_h * small_token_w
+
+    token_tensor = torch.tensor(token_ids, dtype=torch.long)
+
+    # Remove EOS token (16385) from the end if present
+    eos_token_id = 16385
+    if len(token_ids) > 0 and token_ids[-1] == eos_token_id:
+        token_tensor = token_tensor[:-1]
+
+    actual_tokens = len(token_tensor)
+
+    logger.debug(
+        f"[_parse_generated_tokens] height={height}, width={width}, "
+        f"token_h={token_h}, token_w={token_w}, "
+        f"large_image_tokens={large_image_tokens}, small_image_tokens={small_image_tokens}, "
+        f"actual_tokens={actual_tokens}"
+    )
+
+    if is_i2i:
+        # Image-to-image mode: check if AR generated small+large tokens (like t2i) or just large tokens
+        # Some AR models output small+large even in i2i mode
+        if actual_tokens >= small_image_tokens + large_image_tokens:
+            # AR generated full t2i-style output, extract large tokens after small
+            large_start = small_image_tokens
+            large_end = large_start + large_image_tokens
+            prior_token_ids_d32 = token_tensor[large_start:large_end]
+            actual_h, actual_w = token_h, token_w
+        else:
+            # AR generated only large tokens (pure i2i output)
+            prior_token_ids_d32 = token_tensor[:large_image_tokens]
+            actual_h, actual_w = token_h, token_w
+    elif actual_tokens >= small_image_tokens + large_image_tokens:
+        # Text-to-image: extract large image tokens after small image tokens
+        large_start = small_image_tokens
+        large_end = large_start + large_image_tokens
+        prior_token_ids_d32 = token_tensor[large_start:large_end]
+        actual_h, actual_w = token_h, token_w
+    elif actual_tokens >= large_image_tokens:
+        # Image-to-image: large image tokens are at the beginning
+        prior_token_ids_d32 = token_tensor[:large_image_tokens]
+        actual_h, actual_w = token_h, token_w
+    else:
+        # Insufficient tokens - try to infer the actual grid size
+        import math
+
+        for scale in [1, 2, 4]:
+            test_h = token_h // scale
+            test_w = token_w // scale
+            test_small_h = test_h // 2
+            test_small_w = test_w // 2
+            test_large = test_h * test_w
+            test_small = test_small_h * test_small_w
+
+            if actual_tokens >= test_small + test_large:
+                prior_token_ids_d32 = token_tensor[test_small : test_small + test_large]
+                actual_h, actual_w = test_h, test_w
+                height = test_h * factor
+                width = test_w * factor
+                logger.warning(f"Adjusted grid to {test_h}x{test_w}, output will be {height}x{width}")
+                break
+            elif actual_tokens >= test_large:
+                prior_token_ids_d32 = token_tensor[:test_large]
+                actual_h, actual_w = test_h, test_w
+                height = test_h * factor
+                width = test_w * factor
+                logger.warning(f"Adjusted grid to {test_h}x{test_w}, output will be {height}x{width}")
+                break
+        else:
+            sqrt_tokens = int(math.sqrt(actual_tokens))
+            actual_h = actual_w = sqrt_tokens
+            usable_tokens = sqrt_tokens * sqrt_tokens
+            prior_token_ids_d32 = token_tensor[:usable_tokens]
+            height = sqrt_tokens * factor
+            width = sqrt_tokens * factor
+            logger.error(f"Grid pattern mismatch. Using {sqrt_tokens}x{sqrt_tokens}, output: {height}x{width}")
+
+    # Upsample from 32x to 16x
+    prior_token_ids = _upsample_token_ids(prior_token_ids_d32, actual_h, actual_w)
+
+    return prior_token_ids, height, width
+
+
+def ar2diffusion(
+    stage_list: list[Any],
+    engine_input_source: list[int],
+    prompt: OmniTokensPrompt | TextPrompt | list | None = None,
+    requires_multimodal_data: bool = False,
+) -> list[dict[str, Any]]:
+    """Process AR stage outputs to create Diffusion stage inputs."""
+    if not engine_input_source:
+        raise ValueError("engine_input_source cannot be empty")
+
+    source_stage_id = engine_input_source[0]
+    if source_stage_id >= len(stage_list):
+        raise IndexError(f"Invalid stage_id: {source_stage_id}")
+
+    if stage_list[source_stage_id].engine_outputs is None:
+        raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
+
+    ar_outputs = stage_list[source_stage_id].engine_outputs
+    diffusion_inputs = []
+
+    # Normalize prompt to list
+    if not isinstance(prompt, list):
+        prompt = [prompt] if prompt is not None else [{}]
+
+    for i, ar_output in enumerate(ar_outputs):
+        output = ar_output.outputs[0]
+        generated_token_ids = output.token_ids
+
+        # Get original prompt info
+        original_prompt = prompt[i] if i < len(prompt) else {}
+        if isinstance(original_prompt, dict):
+            pass
+        elif hasattr(original_prompt, "_asdict"):
+            original_prompt = original_prompt._asdict()
+        elif hasattr(original_prompt, "__dict__"):
+            original_prompt = vars(original_prompt)
+        else:
+            original_prompt = {}
+
+        height = original_prompt.get("height", 1024)
+        width = original_prompt.get("width", 1024)
+        text_prompt = original_prompt.get("prompt", "")
+
+        # Detect i2i mode first by checking if multimodal_output contains prior_token_image_ids
+        is_i2i = False
+        if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output:
+            mm_output = ar_output.multimodal_output
+            if isinstance(mm_output, dict) and mm_output.get("prior_token_image_ids") is not None:
+                is_i2i = True
+
+        # Parse and upsample prior tokens
+        prior_token_ids, pixel_h, pixel_w = _parse_generated_tokens(generated_token_ids, height, width, is_i2i=is_i2i)
+
+        # Get prior_token_image_ids from AR model output (for i2i mode)
+        # This contains VQ-VAE tokens from input image, used for KV cache conditioning
+        # NOTE: multimodal_output is attached to ar_output (RequestOutput), NOT output (CompletionOutput)
+        prior_token_image_ids = None
+
+        # Check ar_output (RequestOutput) for multimodal_output - this is the correct location
+        if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output:
+            mm_output = ar_output.multimodal_output
+            if isinstance(mm_output, dict):
+                raw_prior_image_ids = mm_output.get("prior_token_image_ids")
+                if raw_prior_image_ids is not None:
+                    # Handle different formats:
+                    # 1. Single tensor -> wrap in list
+                    # 2. List of tensors -> use as-is
+                    # 3. List of Python lists (from serialization) -> convert to tensors
+                    if isinstance(raw_prior_image_ids, torch.Tensor):
+                        prior_token_image_ids = [raw_prior_image_ids]
+                    elif isinstance(raw_prior_image_ids, list):
+                        # Check if elements are tensors or Python lists
+                        if raw_prior_image_ids and isinstance(raw_prior_image_ids[0], torch.Tensor):
+                            prior_token_image_ids = raw_prior_image_ids
+                        elif raw_prior_image_ids and isinstance(raw_prior_image_ids[0], list):
+                            # Convert Python lists back to tensors
+                            prior_token_image_ids = [torch.tensor(ids, dtype=torch.long) for ids in raw_prior_image_ids]
+                        else:
+                            logger.warning(
+                                f"[ar2diffusion] Request {i}: unexpected prior_token_image_ids format: "
+                                f"{type(raw_prior_image_ids[0]) if raw_prior_image_ids else 'empty'}"
+                            )
+        else:
+            # Fallback: also check output (CompletionOutput) in case of different vLLM versions
+            if hasattr(output, "multimodal_output") and output.multimodal_output:
+                mm_output = output.multimodal_output
+                logger.debug(f"[ar2diffusion] Request {i}: found multimodal_output on CompletionOutput (fallback)")
+                if isinstance(mm_output, dict):
+                    raw_prior_image_ids = mm_output.get("prior_token_image_ids")
+                    if raw_prior_image_ids is not None:
+                        if isinstance(raw_prior_image_ids, torch.Tensor):
+                            prior_token_image_ids = [raw_prior_image_ids]
+                        elif isinstance(raw_prior_image_ids, list):
+                            prior_token_image_ids = raw_prior_image_ids
+
+        diffusion_input = {
+            "prompt": text_prompt,
+            "height": pixel_h,
+            "width": pixel_w,
+            "extra": {
+                "prior_token_ids": prior_token_ids,
+                "prior_token_image_ids": prior_token_image_ids,
+            },
+        }
+
+        if requires_multimodal_data:
+            mm_data = original_prompt.get("multi_modal_data")
+            if mm_data:
+                pil_image = mm_data.get("image")
+                if pil_image is None:
+                    # Try "images" (plural) as fallback
+                    images = mm_data.get("images")
+                    if images:
+                        pil_image = images[0] if isinstance(images, list) else images
+                diffusion_input["pil_image"] = pil_image
+
+        for key in ["seed", "num_inference_steps", "guidance_scale", "negative_prompt"]:
+            if key in original_prompt:
+                diffusion_input[key] = original_prompt[key]
+
+        diffusion_inputs.append(diffusion_input)
+
+    return diffusion_inputs

vllm_omni/patch.py

@@ -17,6 +17,32 @@ from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.layers.rotary_embedding import OmniMRotaryEmbedding
 from vllm_omni.request import OmniRequest
 
+# =============================================================================
+# Patch GlmImageTextConfig to expose mrope_section in rope_parameters
+# =============================================================================
+# GLM-Image uses M-RoPE with mrope_section: [8, 12, 12], but transformers'
+# implementation doesn't expose it in rope_parameters. vLLM's uses_mrope
+# detection relies on "mrope_section" being present in rope_parameters.
+# This patch ensures proper M-RoPE detection for GLM-Image.
+try:
+    from transformers.models.glm_image.configuration_glm_image import GlmImageTextConfig
+
+    _original_glm_image_text_config_init = GlmImageTextConfig.__init__
+
+    def _patched_glm_image_text_config_init(self, *args, **kwargs):
+        _original_glm_image_text_config_init(self, *args, **kwargs)
+        # Ensure rope_parameters exists and contains mrope_section
+        if self.rope_parameters is None:
+            self.rope_parameters = {}
+        if isinstance(self.rope_parameters, dict) and "mrope_section" not in self.rope_parameters:
+            # GLM-Image uses mrope_section: [8, 12, 12] for T/H/W dimensions
+            self.rope_parameters["mrope_section"] = [8, 12, 12]
+
+    GlmImageTextConfig.__init__ = _patched_glm_image_text_config_init
+except ImportError:
+    # GlmImageTextConfig not available, skip patching
+    pass
+
 # Extend RequestStatus enum with omni-specific statuses
 if not hasattr(RequestStatus, "WAITING_FOR_CHUNK"):
     # The value - 1 is intentionally chosen to ensure it is treated

vllm_omni/worker/gpu_model_runner.py

@@ -133,6 +133,8 @@ class OmniGPUModelRunner(GPUModelRunner):
                 use_audio_in_video = bool(use_audio_in_video_value.item())
 
         if supports_mrope(self.get_model()):
+            # Model implements SupportsMRoPE interface
+            # Pass all extracted metadata; models use what they need via **kwargs
             req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions(
                 req_state.prompt_token_ids,
                 mm_features=req_state.mm_features,
@@ -154,6 +156,71 @@ class OmniGPUModelRunner(GPUModelRunner):
                 use_audio_in_video=use_audio_in_video,
             )
 
+    def _calc_mrope_positions(self, scheduler_output: "SchedulerOutput"):
+        """Calculate M-RoPE positions for scheduled tokens.
+
+        Delegates to the upstream implementation first, then applies a fixup
+        pass for models that pre-compute 2D spatial decode positions (e.g.
+        GLM-Image).  This avoids duplicating the full upstream method while
+        still supporting non-linear decode position patterns.
+
+        Models opt-in by declaring ``precomputed_mrope_decode = True`` as a
+        class attribute.  When set, ``get_mrope_input_positions`` is expected
+        to return positions covering **both** prefill and decode tokens.
+        """
+        # Run upstream logic (handles prompt positions + linear decode fallback)
+        super()._calc_mrope_positions(scheduler_output)
+
+        # Only run the fixup if the model pre-computes decode M-RoPE positions
+        if not getattr(self.get_model(), "precomputed_mrope_decode", False):
+            return
+
+        self._fixup_precomputed_mrope_decode_positions(scheduler_output)
+
+    def _fixup_precomputed_mrope_decode_positions(self, scheduler_output: "SchedulerOutput") -> None:
+        """Overwrite linear decode M-RoPE positions with pre-computed ones.
+
+        For image-generation models (like GLM-Image) that output tokens in 2D
+        grid order, ``get_mrope_input_positions`` returns positions for the
+        full sequence (prefill + decode).  The upstream runner only uses the
+        prefill portion and falls back to linear increments for decode.  This
+        method patches the decode slice with the correct pre-computed values.
+        """
+        from vllm.utils import length_from_prompt_token_ids_or_embeds
+
+        mrope_pos_ptr = 0
+        for index, req_id in enumerate(self.input_batch.req_ids):
+            req = self.requests[req_id]
+            assert req.mrope_positions is not None
+
+            num_computed_tokens = self.input_batch.num_computed_tokens_cpu[index]
+            num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
+            num_prompt_tokens = length_from_prompt_token_ids_or_embeds(req.prompt_token_ids, req.prompt_embeds)
+
+            if num_computed_tokens + num_scheduled_tokens > num_prompt_tokens:
+                prompt_part_len = max(0, num_prompt_tokens - num_computed_tokens)
+                completion_part_len = max(0, num_scheduled_tokens - prompt_part_len)
+            else:
+                prompt_part_len = num_scheduled_tokens
+                completion_part_len = 0
+
+            mrope_pos_ptr += prompt_part_len
+
+            if completion_part_len > 0:
+                dst_start = mrope_pos_ptr
+                decode_start = num_computed_tokens + prompt_part_len
+                decode_end = decode_start + completion_part_len
+                total_precomputed = req.mrope_positions.shape[1]
+
+                if decode_end <= total_precomputed:
+                    # Overwrite the linear positions written by upstream with
+                    # the correct pre-computed 2D spatial positions.
+                    self.mrope_positions.cpu[:, dst_start : dst_start + completion_part_len] = req.mrope_positions[
+                        :, decode_start:decode_end
+                    ]
+
+                mrope_pos_ptr += completion_part_len
+
     def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
         """Update the cached states and the persistent batch with the scheduler
         output.