[Bugfix]: Fix LoRA loading failure for modules with numeric indices (e.g., to_out.0 in Diffusion Transformers)

[Wang-Shengyuan]

#35734

• vLLM version: latest main
• PyTorch version: 2.10.0
• Python version: 3.10

Describe the bug

When loading a PEFT LoRA checkpoint that targets modules inside nn.Sequential or nn.ModuleList (e.g., to_out.0), vLLM has two related failures:

1. Validation failure – check_unexpected_modules raises ValueError, incorrectly treating the module as unexpected.
2. Silent weight drop – Even if validation is bypassed, _get_lora_layer_weights cannot match the LoRA weight key to_out.0 to the vLLM model's module name to_out, so the LoRA weights are silently ignored.

Error message (failure 1)

ValueError: While loading /path/to/lora, expected target modules in
{'to_out', 'to_q', 'to_k', 'to_v'} but received
['transformer_blocks.0.attn.to_out.0', 'transformer_blocks.1.attn.to_out.0', ...].
Please verify that the loaded LoRA module is correct

Root cause

Validation (lora_model.py, check_unexpected_modules):

# Current code extracts the last dot-segment as the module suffix:
elif module_name.rsplit(".", 1)[-1] not in expected_lora_modules:
    unexpected_modules.append(module_name)

For transformer_blocks.0.attn.to_out.0, rsplit(".", 1)[-1] returns "0", which is the numeric ModuleList index — not the actual module name to_out.

Weight lookup (model_manager.py, _get_lora_layer_weights):

return lora_model.get_lora(org_module_name)

The lookup is exact-match only. When the LoRA checkpoint stores weights under to_out.0 but the vLLM model registers the layer as to_out, the lookup fails and returns None.

How to reproduce

# 1. Train a LoRA with PEFT targeting a ModuleList member
from peft import LoraConfig
config = LoraConfig(target_modules=["to_k", "to_q", "to_v", "to_out"])
# PEFT resolves to_out → to_out.0 (first child of ModuleList)
# Saved checkpoint keys: base_model.model.*.to_out.0.lora_A.weight

# 2. Load with vLLM → ValueError

Affected models

Any model where LoRA targets a layer inside nn.Sequential / nn.ModuleList:

Model	Module structure
Qwen-Image-Edit (Diffusion Transformer)	attn.to_out = ModuleList([Linear, Dropout])
Flux	Same attention pattern
Any HuggingFace model with sequential projections	nn.Sequential(Linear, ...)

Model structure reference

transformer_blocks.0.attn.to_out: ModuleList
  ├── [0]: Linear    ← PEFT targets this → key becomes to_out.0
  └── [1]: Dropout

Expected behavior

LoRA should load and be applied successfully. When the module suffix is a numeric index, vLLM should resolve it to the parent module name.

Workaround

Manually rename to_out.0 → to_out in the saved checkpoint weight keys — error-prone and inconvenient.

Root Cause

When PEFT trains a LoRA on a layer inside nn.ModuleList, the weight key includes a numeric index (e.g., to_out.0.lora_A.weight). After parse_fine_tuned_lora_name, the module name becomes transformer_blocks.0.attn.to_out.0.

Problem 1 – Validation:

# Before: extracts "0" as the suffix → not in expected modules → rejects
module_name.rsplit(".", 1)[-1]  # "0"

Problem 2 – Weight lookup:

# Before: exact match only → key "to_out.0" ≠ model module "to_out" → None
lora_model.get_lora(org_module_name)

Changes

vllm/lora/lora_model.py

1. _get_effective_module_suffix (new static method) — Extracts the effective module suffix. When the last segment is purely numeric, returns the parent segment instead.

@staticmethod
def _get_effective_module_suffix(module_name: str) -> str:
    parts = module_name.rsplit(".", 2)
    suffix = parts[-1]
    if suffix.isdigit() and len(parts) > 1:
        suffix = parts[-2]
    return suffix

2. get_lora_by_indexed_name (new method) — Searches for LoRA weights stored under {module_name}.{digit}. Used as a fallback when exact-match lookup fails.

def get_lora_by_indexed_name(self, module_name: str) -> LoRALayerWeights | None:
    prefix = module_name + "."
    for lora_name in self.loras:
        if lora_name.startswith(prefix):
            trailing = lora_name[len(prefix):]
            if trailing.isdigit():
                return self.loras[lora_name]
    return None

3. check_unexpected_modules — Replaced inline suffix extraction with _get_effective_module_suffix.

# Before
elif module_name.rsplit(".", 1)[-1] not in expected_lora_modules:
    unexpected_modules.append(module_name)

# After
else:
    suffix = cls._get_effective_module_suffix(module_name)
    if suffix not in expected_lora_modules:
        unexpected_modules.append(module_name)

vllm/lora/model_manager.py

4. _get_lora_layer_weights — Added indexed-name fallback after exact match, following the same pattern as the existing pooling-model fallback.

result = lora_model.get_lora(org_module_name)
if result is not None:
    return result

# Fallback for indexed modules
result = lora_model.get_lora_by_indexed_name(org_module_name)
if result is not None:
    logger.info_once(
        "LoRA weights for module '%s' matched using indexed-name "
        "fallback (e.g., '%s.0'). This typically occurs with PEFT "
        "checkpoints targeting nn.Sequential/nn.ModuleList members.",
        org_module_name, org_module_name,
    )
return result

[gemini-code-assist]

Code Review

This pull request addresses a bug that caused LoRA loading to fail for modules with numeric indices, such as those found in nn.Sequential or nn.ModuleList. The fix is well-structured, addressing both the validation and weight lookup failures. A new helper function correctly extracts the effective module suffix, and a fallback mechanism is added to locate LoRA weights for indexed modules. The changes are accompanied by a comprehensive set of new tests that cover the new logic and various edge cases. The implementation appears correct and effectively resolves the issue.

[mergify]

Hi @Wang-Shengyuan, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[jeejeelee]

Could you please provide the reproduce script?

[Wang-Shengyuan]

Thank you for your reply!

The reproduce script can be like

"""
LoRA training sanity check script using vLLM-Omni
Supports batch inference acceleration
"""
import argparse
import gc
import json
import logging
import os
import time
from pathlib import Path
from typing import Optional, List, Dict, Any

import torch
from datasets import load_dataset
from PIL import Image, ImageDraw, ImageFont
from torch.utils.data import Dataset
from tqdm import tqdm

from vllm.lora.request import LoRARequest
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.lora.utils import stable_lora_int_id
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform

logging.basicConfig(
    format="[%(asctime)s] %(levelname)s %(filename)s:%(lineno)d: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)


class HuggingFaceImageEditDataset(Dataset):
    """Image editing dataset loaded from HuggingFace Hub"""

    def __init__(
        self,
        repo_id: str,
        split: str = "train",
        num_samples: Optional[int] = None,
    ):
        self.repo_id = repo_id
        self.split = split

        try:
            token = os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get("HF_TOKEN")
            kwargs = {"token": token} if token else {}
            self.hf_dataset = load_dataset(repo_id, split=split, **kwargs)
            logger.info(f"Loaded dataset from HuggingFace: {repo_id}, split={split}, samples={len(self.hf_dataset)}")
        except Exception as e:
            raise RuntimeError(f"Failed to load dataset {repo_id} from HuggingFace: {e}")

        if num_samples is not None and num_samples > 0:
            self.hf_dataset = self.hf_dataset.select(range(min(num_samples, len(self.hf_dataset))))
            logger.info(f"Limited to {len(self.hf_dataset)} samples")

    def __len__(self):
        return len(self.hf_dataset)

    def __getitem__(self, idx):
        sample = self.hf_dataset[idx]

        control_img = sample["control_images"][0] if isinstance(sample["control_images"], list) else sample["control_images"]
        if not isinstance(control_img, Image.Image):
            control_img = Image.open(control_img)
        control_img = control_img.convert("RGB")

        target_img = sample.get("target_image")
        if target_img is not None and not isinstance(target_img, Image.Image):
            target_img = Image.open(target_img)
        target_img = target_img.convert("RGB") if target_img is not None else None

        prompt = sample.get("prompt", "")

        return {
            "control_image": control_img,
            "target_image": target_img,
            "prompt": prompt,
            "index": idx,
        }


def add_title_to_image(image: Image.Image, title: str, font_size: int = 36) -> Image.Image:
    """Add a title banner above the image"""
    title_height = int(font_size * 1.8)
    new_image = Image.new("RGB", (image.width, image.height + title_height), color="white")
    new_image.paste(image, (0, title_height))

    draw = ImageDraw.Draw(new_image)
    font = None
    font_paths = [
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
        "/System/Library/Fonts/Helvetica.ttc",
        "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
    ]
    for font_path in font_paths:
        try:
            if os.path.exists(font_path):
                font = ImageFont.truetype(font_path, font_size)
                break
        except:
            continue
    if font is None:
        font = ImageFont.load_default()

    bbox = draw.textbbox((0, 0), title, font=font)
    text_width = bbox[2] - bbox[0]
    text_height = bbox[3] - bbox[1]
    text_x = max((image.width - text_width) // 2, 8)
    text_y = max((title_height - text_height) // 2, 2)
    draw.text((text_x, text_y), title, fill="black", font=font)

    return new_image


def resize_and_pad_image(image: Image.Image, target_w: int, target_h: int, bg_color: str = "white") -> Image.Image:
    """Scale image proportionally and pad to target size"""
    if image.mode != "RGB":
        image = image.convert("RGB")

    scale = min(target_w / image.width, target_h / image.height)
    new_w = max(1, int(round(image.width * scale)))
    new_h = max(1, int(round(image.height * scale)))

    resized = image.resize((new_w, new_h), Image.Resampling.LANCZOS)
    canvas = Image.new("RGB", (target_w, target_h), color=bg_color)
    paste_x = (target_w - new_w) // 2
    paste_y = (target_h - new_h) // 2
    canvas.paste(resized, (paste_x, paste_y))
    return canvas


def validate_lora_path(lora_path: str) -> bool:
    """Validate LoRA path for PEFT format"""
    if not os.path.exists(lora_path):
        logger.error(f"LoRA path does not exist: {lora_path}")
        return False

    adapter_config = os.path.join(lora_path, "adapter_config.json")
    if not os.path.exists(adapter_config):
        logger.error(f"Missing adapter_config.json: {adapter_config}")
        logger.error("vLLM-Omni requires PEFT-format LoRA weights")
        files = os.listdir(lora_path)
        logger.info(f"LoRA directory contents: {files}")
        return False

    has_weights = any(
        f.endswith(('.safetensors', '.bin'))
        for f in os.listdir(lora_path)
    )
    if not has_weights:
        logger.error("Missing LoRA weight files (.safetensors or .bin)")
        return False

    logger.info(f"✅ LoRA path validated: {lora_path}")
    return True


def create_omni_pipeline(
    model: str,
    lora_path: Optional[str] = None,
    lora_scale: float = 1.0,
    cache_backend: Optional[str] = None,
    tensor_parallel_size: int = 1,
    enforce_eager: bool = False,
    enable_cpu_offload: bool = False,
) -> Omni:
    """Create vLLM-Omni inference pipeline"""
    logger.info(f"Loading model: {model}")

    if lora_path:
        if not validate_lora_path(lora_path):
            raise ValueError(f"Invalid LoRA path: {lora_path}")
        logger.info(f"Loading LoRA: {lora_path}")

    parallel_config = DiffusionParallelConfig(
        ulysses_degree=1,
        ring_degree=1,
        cfg_parallel_size=1,
        tensor_parallel_size=tensor_parallel_size,
    )

    cache_config = None
    if cache_backend == "cache_dit":
        cache_config = {
            "Fn_compute_blocks": 1,
            "Bn_compute_blocks": 0,
            "max_warmup_steps": 4,
            "residual_diff_threshold": 0.24,
            "max_continuous_cached_steps": 3,
            "enable_taylorseer": False,
            "taylorseer_order": 1,
            "scm_steps_mask_policy": None,
            "scm_steps_policy": "dynamic",
        }
    elif cache_backend == "tea_cache":
        cache_config = {"rel_l1_thresh": 0.2}

    omni = Omni(
        model=model,
        lora_path=lora_path,
        lora_scale=lora_scale,
        cache_backend=cache_backend,
        cache_config=cache_config,
        parallel_config=parallel_config,
        enforce_eager=enforce_eager,
        enable_cpu_offload=enable_cpu_offload,
    )
    logger.info("Pipeline loaded successfully")
    return omni


def release_omni_pipeline(omni: Omni):
    """Release pipeline memory"""
    try:
        omni.close()
        logger.info("Omni pipeline closed")
    except Exception as e:
        logger.warning(f"Error closing Omni pipeline: {e}")

    del omni

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

    gc.collect()
    time.sleep(3)

    logger.info("Pipeline memory released")


def batch_inference_with_omni(
    omni: Omni,
    samples: List[Dict[str, Any]],
    negative_prompt: Optional[str] = None,
    num_inference_steps: int = 20,
    cfg_scale: float = 4.0,
    guidance_scale: float = 1.0,
    seed: int = 42,
    lora_path: Optional[str] = None,
    lora_scale: float = 1.0,
) -> List[Optional[Image.Image]]:
    """
    Batch inference using vLLM-Omni's batching capability.

    Note: lora_request must be passed in sampling_params to activate LoRA.
    Passing lora_path to Omni() only pre-loads the adapter;
    each inference request must activate it via lora_request.
    """
    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)

    requests = []
    for sample in samples:
        requests.append({
            "prompt": sample["prompt"],
            "negative_prompt": negative_prompt,
            "multi_modal_data": {"image": sample["control_image"]},
        })

    lora_request = None
    if lora_path:
        lora_request = LoRARequest(
            lora_name="inference_lora",
            lora_int_id=stable_lora_int_id(lora_path),
            lora_path=lora_path,
        )
        logger.info(f"LoRA request created: path={lora_path}, scale={lora_scale}")

    sampling_params = OmniDiffusionSamplingParams(
        generator=generator,
        true_cfg_scale=cfg_scale,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        num_outputs_per_prompt=1,
        lora_request=lora_request,
        lora_scale=lora_scale,
    )

    outputs = omni.generate(requests, sampling_params)

    results = []
    for output in outputs:
        try:
            if hasattr(output, "request_output") and output.request_output:
                req_out = output.request_output[0]
                if isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images") and req_out.images:
                    results.append(req_out.images[0])
                else:
                    results.append(None)
            else:
                results.append(None)
        except Exception as e:
            logger.warning(f"Failed to parse output: {e}")
            results.append(None)

    return results


def inference_with_omni(
    omni: Omni,
    control_image: Image.Image,
    prompt: str,
    negative_prompt: Optional[str] = None,
    num_inference_steps: int = 20,
    cfg_scale: float = 4.0,
    guidance_scale: float = 1.0,
    seed: int = 42,
    lora_path: Optional[str] = None,
    lora_scale: float = 1.0,
) -> Image.Image:
    """Single-sample inference (kept for compatibility)"""
    results = batch_inference_with_omni(
        omni=omni,
        samples=[{"control_image": control_image, "prompt": prompt}],
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        cfg_scale=cfg_scale,
        guidance_scale=guidance_scale,
        seed=seed,
        lora_path=lora_path,
        lora_scale=lora_scale,
    )

    if not results or results[0] is None:
        raise ValueError("No output generated from omni.generate()")

    return results[0]


def run_inference_stage(
    omni: Omni,
    dataset: HuggingFaceImageEditDataset,
    output_dir: str,
    stage_name: str,
    args: argparse.Namespace,
    batch_size: int = 1,
    lora_path: Optional[str] = None,
    lora_scale: float = 1.0,
) -> Dict[int, Optional[Image.Image]]:
    """
    Run an inference stage with optional batch processing.

    Args:
        lora_path: LoRA weight path; must be passed per-request to activate LoRA.
        lora_scale: LoRA weight scaling factor.
    """
    results = {}
    num_samples = len(dataset)

    if batch_size > 1:
        # Batch inference mode
        for batch_start in tqdm(range(0, num_samples, batch_size), desc=f"{stage_name} inference"):
            batch_end = min(batch_start + batch_size, num_samples)
            batch_samples = []
            batch_indices = []

            for idx in range(batch_start, batch_end):
                sample = dataset[idx]
                batch_samples.append({
                    "control_image": sample["control_image"],
                    "prompt": sample["prompt"],
                })
                batch_indices.append(sample["index"])

            try:
                batch_results = batch_inference_with_omni(
                    omni=omni,
                    samples=batch_samples,
                    negative_prompt=args.negative_prompt,
                    num_inference_steps=args.num_inference_steps,
                    cfg_scale=args.cfg_scale,
                    guidance_scale=args.guidance_scale,
                    seed=args.seed,
                    lora_path=lora_path,
                    lora_scale=lora_scale,
                )

                for sample_idx, result in zip(batch_indices, batch_results):
                    if result is not None:
                        save_path = os.path.join(output_dir, f"sample_{sample_idx:04d}_{stage_name.lower()}.png")
                        result.save(save_path)
                        results[sample_idx] = result
                        logger.info(f"✅ {stage_name} sample {sample_idx} done")
                    else:
                        results[sample_idx] = None
                        logger.warning(f"⚠️ {stage_name} sample {sample_idx} returned empty result")

            except Exception as e:
                logger.error(f"❌ {stage_name} batch {batch_start}-{batch_end} failed: {e}")
                for sample_idx in batch_indices:
                    results[sample_idx] = None
    else:
        # Sequential inference mode
        for idx in tqdm(range(num_samples), desc=f"{stage_name} inference"):
            sample = dataset[idx]
            sample_idx = sample["index"]

            try:
                result = inference_with_omni(
                    omni=omni,
                    control_image=sample["control_image"],
                    prompt=sample["prompt"],
                    negative_prompt=args.negative_prompt,
                    num_inference_steps=args.num_inference_steps,
                    cfg_scale=args.cfg_scale,
                    guidance_scale=args.guidance_scale,
                    seed=args.seed,
                    lora_path=lora_path,
                    lora_scale=lora_scale,
                )
                save_path = os.path.join(output_dir, f"sample_{sample_idx:04d}_{stage_name.lower()}.png")
                result.save(save_path)
                results[sample_idx] = result
                logger.info(f"✅ {stage_name} sample {idx + 1}/{num_samples} done")
            except Exception as e:
                logger.error(f"❌ {stage_name} sample {idx + 1}/{num_samples} failed: {e}")
                results[sample_idx] = None

    return results


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="LoRA training sanity check using vLLM-Omni")

    # Model arguments
    parser.add_argument("--base_model", type=str, required=True, help="Base model path")
    parser.add_argument("--lora_path", type=str, default=None, help="LoRA weight path (optional)")
    parser.add_argument("--lora_scale", type=float, default=1.0, help="LoRA weight scaling factor")

    # Dataset arguments
    parser.add_argument("--hf_repo_id", type=str, required=True, help="HuggingFace dataset repo ID")
    parser.add_argument("--hf_split", type=str, default="train", help="HuggingFace dataset split")
    parser.add_argument("--num_samples", type=int, default=None, help="Limit number of inference samples")

    # Output arguments
    parser.add_argument("--output_dir", type=str, default="./sanity_check_results_vllm", help="Output directory")

    # Inference arguments
    parser.add_argument("--num_inference_steps", type=int, default=20, help="Number of inference steps")
    parser.add_argument("--cfg_scale", type=float, default=4.0, help="CFG scale")
    parser.add_argument("--guidance_scale", type=float, default=1.0, help="Guidance scale")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")
    parser.add_argument("--negative_prompt", type=str, default=None, help="Negative prompt")
    parser.add_argument("--batch_size", type=int, default=1, help="Batch size for inference acceleration")

    # vLLM-Omni specific arguments
    parser.add_argument("--cache_backend", type=str, default=None, choices=["cache_dit", "tea_cache"], help="Cache backend for acceleration")
    parser.add_argument("--tensor_parallel_size", type=int, default=1, help="Tensor parallel size")
    parser.add_argument("--enforce_eager", action="store_true", help="Disable torch.compile")
    parser.add_argument("--enable_cpu_offload", action="store_true", help="Enable CPU offload")

    # Visualization arguments
    parser.add_argument("--viz_font_size", type=int, default=36, help="Font size for comparison image titles")
    parser.add_argument("--viz_max_side", type=int, default=768, help="Max side length for each panel in comparison image")

    return parser.parse_args()


def main():
    args = parse_args()

    # Create output directories
    os.makedirs(args.output_dir, exist_ok=True)
    base_output_dir = os.path.join(args.output_dir, "base_model")
    lora_output_dir = os.path.join(args.output_dir, "lora_model")
    comparison_dir = os.path.join(args.output_dir, "comparison")
    os.makedirs(base_output_dir, exist_ok=True)
    os.makedirs(comparison_dir, exist_ok=True)
    if args.lora_path:
        os.makedirs(lora_output_dir, exist_ok=True)

    # Validate LoRA path
    if args.lora_path and not validate_lora_path(args.lora_path):
        logger.error("LoRA path validation failed, please check file format")
        return

    # Save config
    config_info = {
        "base_model": args.base_model,
        "lora_path": args.lora_path,
        "dataset": args.hf_repo_id,
        "split": args.hf_split,
        "num_inference_steps": args.num_inference_steps,
        "cfg_scale": args.cfg_scale,
        "guidance_scale": args.guidance_scale,
        "seed": args.seed,
        "cache_backend": args.cache_backend,
        "tensor_parallel_size": args.tensor_parallel_size,
        "batch_size": args.batch_size,
    }
    with open(os.path.join(args.output_dir, "config.json"), "w", encoding="utf-8") as f:
        json.dump(config_info, f, indent=2, ensure_ascii=False)

    # Load dataset
    logger.info("Loading dataset...")
    dataset = HuggingFaceImageEditDataset(
        repo_id=args.hf_repo_id,
        split=args.hf_split,
        num_samples=args.num_samples,
    )

    # Pre-collect all sample info
    sample_info = {}
    for idx in range(len(dataset)):
        sample = dataset[idx]
        sample_info[sample["index"]] = {
            "control_image": sample["control_image"],
            "target_image": sample["target_image"],
            "prompt": sample["prompt"],
        }

    # ========== Stage 1: Base model inference ==========
    logger.info("=" * 60)
    logger.info("Stage 1: Loading base model and running inference...")
    logger.info("=" * 60)

    omni_base = create_omni_pipeline(
        model=args.base_model,
        lora_path=None,
        cache_backend=args.cache_backend,
        tensor_parallel_size=args.tensor_parallel_size,
        enforce_eager=args.enforce_eager,
        enable_cpu_offload=args.enable_cpu_offload,
    )

    base_results = run_inference_stage(
        omni=omni_base,
        dataset=dataset,
        output_dir=base_output_dir,
        stage_name="Base",
        args=args,
        batch_size=args.batch_size,
    )

    logger.info("Releasing base model memory...")
    release_omni_pipeline(omni_base)
    omni_base = None

    # ========== Stage 2: LoRA model inference ==========
    lora_results = {}
    if args.lora_path:
        logger.info("=" * 60)
        logger.info("Stage 2: Loading LoRA model and running inference...")
        logger.info("=" * 60)

        omni_lora = create_omni_pipeline(
            model=args.base_model,
            lora_path=args.lora_path,
            lora_scale=args.lora_scale,
            cache_backend=args.cache_backend,
            tensor_parallel_size=args.tensor_parallel_size,
            enforce_eager=args.enforce_eager,
            enable_cpu_offload=args.enable_cpu_offload,
        )

        lora_results = run_inference_stage(
            omni=omni_lora,
            dataset=dataset,
            output_dir=lora_output_dir,
            stage_name="LoRA",
            args=args,
            batch_size=args.batch_size,
            lora_path=args.lora_path,  # Key: pass lora_path per-request to activate LoRA
            lora_scale=args.lora_scale,
        )

        logger.info("Releasing LoRA model memory...")
        release_omni_pipeline(omni_lora)
        omni_lora = None

    # ========== Stage 3: Generate comparison images ==========
    logger.info("=" * 60)
    logger.info("Stage 3: Generating comparison images...")
    logger.info("=" * 60)

    results = []
    for idx, sample_idx in enumerate(sorted(sample_info.keys()), 1):
        info = sample_info[sample_idx]
        control_image = info["control_image"]
        target_image = info["target_image"]
        prompt = info["prompt"]
        base_result = base_results.get(sample_idx)
        lora_result = lora_results.get(sample_idx) if args.lora_path else None

        if base_result is None:
            continue

        control_path = os.path.join(comparison_dir, f"sample_{sample_idx:04d}_control.png")
        control_image.save(control_path)

        target_path = None
        if target_image is not None:
            target_path = os.path.join(comparison_dir, f"sample_{sample_idx:04d}_target.png")
            target_image.save(target_path)

        ref_w, ref_h = control_image.width, control_image.height
        ref_max_side = max(ref_w, ref_h)
        scale = min(1.0, args.viz_max_side / ref_max_side) if ref_max_side > 0 else 1.0
        tile_w = max(256, int(round(ref_w * scale)))
        tile_h = max(256, int(round(ref_h * scale)))

        panels = []

        control_norm = resize_and_pad_image(control_image, tile_w, tile_h)
        panels.append(add_title_to_image(control_norm, "Control", font_size=args.viz_font_size))

        base_norm = resize_and_pad_image(base_result, tile_w, tile_h)
        panels.append(add_title_to_image(base_norm, "Base Model", font_size=args.viz_font_size))

        if lora_result is not None:
            lora_norm = resize_and_pad_image(lora_result, tile_w, tile_h)
            panels.append(add_title_to_image(lora_norm, "LoRA Model", font_size=args.viz_font_size))

        if target_image is not None:
            target_norm = resize_and_pad_image(target_image, tile_w, tile_h)
            panels.append(add_title_to_image(target_norm, "Target", font_size=args.viz_font_size))

        panel_w, panel_h = panels[0].width, panels[0].height
        comparison_image = Image.new("RGB", (panel_w * len(panels), panel_h), color="white")
        for i, panel in enumerate(panels):
            comparison_image.paste(panel, (i * panel_w, 0))

        comparison_path = os.path.join(comparison_dir, f"sample_{sample_idx:04d}_comparison.png")
        comparison_image.save(comparison_path)

        result_info = {
            "sample_index": sample_idx,
            "prompt": prompt,
            "control_path": control_path,
            "base_path": os.path.join(base_output_dir, f"sample_{sample_idx:04d}_base.png"),
            "lora_path": os.path.join(lora_output_dir, f"sample_{sample_idx:04d}_lora.png") if args.lora_path else None,
            "target_path": target_path,
            "comparison_path": comparison_path,
            "status": "success",
        }
        results.append(result_info)
        logger.info(f"✅ Comparison {idx}/{len(sample_info)} generated")

    successful = len(results)
    summary = {
        "total_samples": len(dataset),
        "successful_samples": successful,
        "failed_samples": len(dataset) - successful,
        "results": results,
    }

    with open(os.path.join(args.output_dir, "summary.json"), "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)

    logger.info("=" * 60)
    logger.info("Sanity check complete!")
    logger.info(f"Total samples: {len(dataset)}")
    logger.info(f"Successful: {successful}")
    logger.info(f"Failed: {len(dataset) - successful}")
    logger.info(f"Results saved to: {args.output_dir}")
    logger.info("=" * 60)


if __name__ == "__main__":
    main()

with configs

python vllm_infer.py \
    --base_model my-path-to/Qwen/Qwen-Image-Edit-2509 \
    --lora_path my-path-to-a-lora-peft-adapter \
    --hf_repo_id TsienDragon/face_segmentation_20 \
    --hf_split train \
    --output_dir ./output/ \
    --num_samples 10 \
    --batch_size 4 \
    --num_inference_steps 20 \
    --cfg_scale 4.0 \
    --seed 42

[Wang-Shengyuan]

The lora adapter is trained with peft in a standard practice. The architecture of the adapter is as follows (adapter_config.json).

{
  "alora_invocation_tokens": null,
  "alpha_pattern": {},
  "arrow_config": null,
  "auto_mapping": null,
  "base_model_name_or_path": null,
  "bias": "none",
  "corda_config": null,
  "ensure_weight_tying": false,
  "eva_config": null,
  "exclude_modules": null,
  "fan_in_fan_out": false,
  "inference_mode": false,
  "init_lora_weights": "gaussian",
  "layer_replication": null,
  "layers_pattern": null,
  "layers_to_transform": null,
  "loftq_config": {},
  "lora_alpha": 32,
  "lora_bias": false,
  "lora_dropout": 0.0,
  "megatron_config": null,
  "megatron_core": "megatron.core",
  "modules_to_save": null,
  "peft_type": "LORA",
  "peft_version": "0.18.0",
  "qalora_group_size": 16,
  "r": 32,
  "rank_pattern": {},
  "revision": null,
  "target_modules": [
    "to_v",
    "to_out.0",
    "to_q",
    "to_k"
  ],
  "target_parameters": null,
  "task_type": null,
  "trainable_token_indices": null,
  "use_dora": false,
  "use_qalora": false,
  "use_rslora": false
}

[mergify]

Hi @Wang-Shengyuan, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[github-actions]

This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you!