[vLLM] add vLLM offline inference (#47)

* [vLLM] add vLLM offline inference

* update x2i tokenizer_path

* add copyright

* refactor requirements

* rm requirements.txt

* update readme

Author: zhaoyinglia

README.md

@@ -64,10 +64,10 @@ Emu3.5 Team, BAAI
 ### Environment Setup
 
 ```bash
-# Python 3.10 or higher is required.
+# Requires Python 3.12 or higher.
 git clone https://github.com/baaivision/Emu3.5
 cd Emu3.5
-pip install -r requirements.txt
+pip install -r requirements/transformers.txt
 pip install flash_attn==2.8.3 --no-build-isolation
 ```
 ### Configuration
@@ -112,6 +112,44 @@ CUDA_VISIBLE_DEVICES=0,1 python inference.py --cfg configs/example_config_visual
 
 Protobuf outputs are written to `outputs/<exp_name>/proto/`. For better throughput, we recommend ≥2 GPUs.
 
+
+### Run Inference with vLLM
+
+#### vLLM Enviroment Setup
+
+1. [Optional Recommendation] Use a virtual environment
+```bash
+conda create -n Emu3p5 python=3.12
+```
+
+2. Install vLLM and apply the patch files.
+```bash
+# Requires Python 3.12 or higher.
+# Recommended: CUDA 12.8.
+pip install -r requirements/vllm.txt
+pip install flash_attn==2.8.3 --no-build-isolation
+
+cd Emu3.5
+python src/patch/apply.py
+```
+
+#### Example Configurations by Task
+
+```bash
+# 🖼️ Text-to-Image (T2I) task
+CUDA_VISIBLE_DEVICES=0,1 python inference_vllm.py --cfg configs/example_config_t2i.py
+
+# 🔄 Any-to-Image (X2I) task
+CUDA_VISIBLE_DEVICES=0,1 python inference_vllm.py --cfg configs/example_config_x2i.py
+
+# 🎯 Visual Guidance task
+CUDA_VISIBLE_DEVICES=0,1 python inference_vllm.py --cfg configs/example_config_visual_guidance.py
+
+# 📖 Visual Narrative task
+CUDA_VISIBLE_DEVICES=0,1 python inference_vllm.py --cfg configs/example_config_visual_narrative.py
+```
+
+
 ### Visualize Protobuf Outputs
 
 To visualize generated protobuf files (--video: Generate video visualizations for interleaved output):

configs/example_config_x2i.py

@@ -8,7 +8,7 @@
 model_path = "path_to_emu3.5_model" # download from hf
 vq_path = "path_to_vq_model" # download from hf
 
-tokenizer_path = "path_to_tokenizer"
+tokenizer_path = "./src/tokenizer_emu3_ibq"
 vq_type = "ibq"
 
 task_type = "x2i"

inference_vllm.py

@@ -0,0 +1,155 @@
+# Copyright 2025 BAAI. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import os
+import torch
+
+import importlib as imp
+import os.path as osp
+
+from pathlib import Path
+from PIL import Image
+from tqdm import tqdm
+
+from src.utils.model_utils import build_emu3p5_vllm
+from src.utils.vllm_generation_utils import generate
+from src.utils.generation_utils import multimodal_decode
+from src.utils.painting_utils import ProtoWriter
+from src.utils.input_utils import build_image
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--cfg", default="", type=str)
+    parser.add_argument("--tensor-parallel-size", default=2, type=int)
+    parser.add_argument("--gpu-memory-utilization", default=0.7, type=float)
+    parser.add_argument("--seed", default=6666, type=int)
+    args = parser.parse_args()
+    return args
+
+
+def inference(
+    cfg,
+    model,
+    tokenizer,
+    vq_model,
+):
+    save_path = cfg.save_path
+
+    os.makedirs(save_path, exist_ok=True)
+    os.makedirs(f"{save_path}/proto", exist_ok=True)
+    proto_writer = ProtoWriter()
+
+    for name, question in tqdm(cfg.prompts, total=len(cfg.prompts)):
+        if osp.exists(f"{save_path}/proto/{name}.pb"):
+            print(f"[WARNING] Result already exists, skipping {name}", flush=True)
+            continue
+
+        torch.cuda.empty_cache()
+
+        reference_image = None
+        if not isinstance(question, str):
+            if isinstance(question["reference_image"], list):
+                print(f"[INFO] {len(question['reference_image'])} reference images are provided")
+                reference_image = []
+                for img in question["reference_image"]:
+                    reference_image.append(Image.open(img).convert("RGB"))
+            else:
+                print (f"[INFO] 1 reference image is provided")
+                reference_image = Image.open(question["reference_image"]).convert("RGB")
+            question = question["prompt"]
+        else:
+            print(f"[INFO] No reference image is provided")
+
+        proto_writer.clear()
+        proto_writer.extend([["question", question]])
+        if reference_image is not None:
+            if isinstance(reference_image, list):
+                for idx, img in enumerate(reference_image):
+                    proto_writer.extend([[f"reference_image", img]])
+            else:
+                proto_writer.extend([["reference_image", reference_image]])
+
+        success = True
+        prompt = cfg.template.format(question=question)
+
+        print(f"[INFO] Handling prompt: {prompt}")
+        if reference_image is not None:
+            if isinstance(reference_image, list):
+                image_str = ""
+                for img in reference_image:
+                    image_str += build_image(img, cfg, tokenizer, vq_model)
+            else:
+                image_str = build_image(reference_image, cfg, tokenizer, vq_model)
+            prompt = prompt.replace("<|IMAGE|>", image_str)
+            unc_prompt = cfg.unc_prompt.replace("<|IMAGE|>", image_str)
+        else:
+            unc_prompt = cfg.unc_prompt
+
+        input_ids = tokenizer.encode(prompt, return_tensors="pt", add_special_tokens=False)
+
+        if input_ids[0, 0] != cfg.special_token_ids["BOS"]:
+            BOS = torch.Tensor([[cfg.special_token_ids["BOS"]]], dtype=input_ids.dtype)
+            input_ids = torch.cat([BOS, input_ids], dim=1)
+
+        unconditional_ids = tokenizer.encode(unc_prompt, return_tensors="pt", add_special_tokens=False)
+
+        for result_tokens in generate(cfg, model, tokenizer, input_ids, unconditional_ids):
+            try:
+                print(f"{result_tokens.shape=}")
+                result = tokenizer.decode(result_tokens, skip_special_tokens=False)
+                mm_out = multimodal_decode(result, tokenizer, vq_model)
+                proto_writer.extend(mm_out)
+            except Exception as e:
+                success = False
+                print(f"[ERROR] Failed to generate token sequence: {e}")
+                break
+
+        if not success:
+            continue
+
+        proto_writer.save(f"{save_path}/proto/{name}.pb")
+
+
+def main():
+    args = parse_args()
+    cfg_name = Path(args.cfg).stem
+    cfg_package = Path(args.cfg).parent.__str__().replace("/", ".")
+    cfg = imp.import_module(f".{cfg_name}", package=cfg_package)
+
+    if isinstance(cfg.prompts, dict):
+        cfg.prompts = [(n, p) for n, p in cfg.prompts.items()]
+    else:
+        cfg.prompts = [(f"{idx:03d}", p) for idx, p in enumerate(cfg.prompts)]
+
+    cfg.prompts = [(n, p) for n, p in cfg.prompts if not osp.exists(f"{cfg.save_path}/proto/{n}.pb")]
+    cfg.num_prompts = len(cfg.prompts)
+
+    model, tokenizer, vq_model = build_emu3p5_vllm(
+        cfg.model_path,
+        cfg.tokenizer_path,
+        cfg.vq_path,
+        vq_type=cfg.vq_type,
+        vq_device=cfg.vq_device,
+        seed=cfg.seed,
+        tensor_parallel_size=args.tensor_parallel_size,
+        gpu_memory_utilization=args.gpu_memory_utilization,
+        **getattr(cfg, "diffusion_decoder_kwargs", {}),
+    )
+    print(f"[INFO] Model loaded successfully")
+    cfg.special_token_ids = {}
+    for k, v in cfg.special_tokens.items():
+        cfg.special_token_ids[k] = tokenizer.encode(v)[0]
+
+    inference(
+        cfg=cfg,
+        model=model,
+        tokenizer=tokenizer,
+        vq_model=vq_model,
+    )
+    print(f"[INFO] Inference finished")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt requirements/common.txtrequirements.txt renamed to requirements/common.txt

@@ -1,8 +1,3 @@
-torch>=2.6.0
-torchvision>=0.15.0
-torchaudio>=2.0.0
-transformers==4.48.2
-accelerate>=0.20.0
 pillow>=9.0.0
 numpy>=1.21.0
 tqdm>=4.64.0
@@ -11,4 +6,4 @@ tiktoken>=0.12.0
 imageio==2.37.0
 imageio-ffmpeg==0.6.0
 omegaconf==2.3.0
-gradio==5.49.1
+gradio==5.49.1

requirements/transformers.txt

@@ -0,0 +1,7 @@
+-r common.txt
+
+torch>=2.6.0
+torchvision>=0.15.0
+torchaudio>=2.0.0
+transformers==4.48.2
+accelerate>=0.20.0

requirements/vllm.txt

@@ -0,0 +1,3 @@
+-r common.txt
+
+vllm==0.11.0; python_version > '3.11' # torch==2.8.0

src/patch/apply.py

@@ -0,0 +1,140 @@
+# -*- coding: utf-8 -*-
+# Copyright 2025 BAAI. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+import os
+import subprocess
+import sys
+import shutil
+from pathlib import Path
+
+REQUIRED_VLLM_VERSION = "0.11.0"
+
+
+def get_vllm_site():
+    try:
+        import vllm
+        if getattr(vllm, "__version__", None) != REQUIRED_VLLM_VERSION:
+            print(f"[FATAL] vLLM version must be {REQUIRED_VLLM_VERSION}, "
+                  f"but found {vllm.__version__}. Aborting.")
+            sys.exit(10)
+        print(f"[INFO] vLLM version verified: {vllm.__version__}")
+        return Path(vllm.__file__).parent
+    except ImportError:
+        print("[ERROR] vllm is not installed. Please run: pip install vllm==0.11.0")
+        sys.exit(1)
+
+
+def run_patch(patch_file, site_dir, dry_run=False):
+    cmd = ["patch", f"-p2"]
+    if dry_run:
+        cmd.insert(1, "--dry-run")
+    with open(patch_file, "r") as f:
+        result = subprocess.run(
+            cmd,
+            cwd=str(site_dir),
+            stdin=f,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True
+        )
+    if not dry_run:
+        print(f"[INFO] Applied patch: {patch_file}")
+    if result.returncode != 0:
+        print(f"[ERROR] Patch failed: {patch_file}")
+        print(result.stdout)
+        print(result.stderr)
+    return result.returncode == 0, result.stdout, result.stderr
+
+
+def extract_patch_targets(patch_file):
+    targets = []
+    with open(patch_file, "r") as f:
+        for line in f:
+            if line.startswith("--- a/") or line.startswith("+++ b/"):
+                path = line.split("\t")[0].split(" ", 1)[-1]
+                if path not in ("a/dev/null", "b/dev/null"):
+                    idx = len("a/vllm/")
+                    targets.append(path[idx:-1])
+    return list(set(targets))
+
+
+def backup_files(targets, site_dir, backup_root):
+    for rel in targets:
+        src = site_dir / rel
+        if src.exists():
+            dst = backup_root / rel
+            dst.parent.mkdir(parents=True, exist_ok=True)
+            print(f"[INFO] Backing up {src} to {dst}")
+            shutil.copy2(src, dst)
+
+
+def restore_backup(backup_root, site_dir):
+    if not backup_root.exists():
+        print("[WARN] No backup directory found.")
+        return
+    for root, _, files in os.walk(backup_root):
+        for f in files:
+            bfile = Path(root) / f
+            rel = bfile.relative_to(backup_root)
+            orig = site_dir / rel
+            orig.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy2(bfile, orig)
+    print("[INFO] Restore completed.")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--patch-dir", type=str, default="./third_party/vllm/",
+                        help="Directory containing .patch files")
+    args = parser.parse_args()
+    patch_dir = Path(args.patch_dir)
+
+    if not patch_dir.exists() or not patch_dir.is_dir():
+        print(f"[ERROR] patch-dir does not exist: {patch_dir}")
+        sys.exit(1)
+
+    site_dir = get_vllm_site()
+    print(f"[INFO] vLLM site-packages: {site_dir}")
+
+    patch_files = sorted(p for p in patch_dir.rglob("*.patch"))
+    if not patch_files:
+        print("[ERROR] No patch files found.")
+        sys.exit(1)
+
+    print(f"[INFO] Found {len(patch_files)} patch(es).")
+
+    # Backup root folder
+    backup_root = site_dir.parent / "vllm_patch_backup"
+    if backup_root.exists():
+        print("[WARN] Removing previous backup...")
+        shutil.rmtree(backup_root)
+    backup_root.mkdir(parents=True)
+
+    print("[INFO] Running dry-run...")
+    for p in patch_files:
+        ok, out, err = run_patch(p, site_dir, dry_run=True)
+        if not ok:
+            print(f"[FATAL] Dry-run failed for patch: {p}\n{err}")
+            sys.exit(2)
+    print("[INFO] Dry-run passed.")
+
+    print("[INFO] Backing up modified files...")
+    for p in patch_files:
+        targets = extract_patch_targets(p)
+        backup_files(targets, site_dir, backup_root)
+
+    print("[INFO] Applying patches...")
+    for p in patch_files:
+        ok, out, err = run_patch(p, site_dir, dry_run=False)
+        if not ok:
+            print(f"[ERROR] Failed to apply patch: {p}\n{err}")
+            print("[INFO] Restoring from backup...")
+            restore_backup(backup_root, site_dir)
+            sys.exit(3)
+    print("[SUCCESS] All patches applied successfully.")
+    print(f"[INFO] Backup stored at: {backup_root}")
+
+if __name__ == "__main__":
+    main()