File size: 4,667 Bytes
3a464db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """Fast-dVLM sglang chatbot with selectable MDM or speculative decoding."""
import argparse
import os
import sys
# `import sglang` must resolve to the pip-installed fork (third_party/sglang),
# so drop this script's dir from sys.path to avoid shadowing by a stray
# local `sglang/` next to this file.
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path = [p for p in sys.path if os.path.abspath(p) != _HERE]
ALGO_MAP = {
"mdm": "HierarchyBlock",
"spec": "SpeculativeBlock",
}
def build_inputs(processor, image, prompt):
from qwen_vl_utils import process_vision_info
content = []
if image:
content.append({"type": "image", "image": image})
content.append({"type": "text", "text": prompt})
messages = [{"role": "user", "content": content}]
text = processor.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True,
)
image_inputs, video_inputs = process_vision_info(messages)
inputs = processor(
text=[text], images=image_inputs, videos=video_inputs,
padding=True, return_tensors="pt",
)
return inputs.input_ids[0].tolist()
def main():
parser = argparse.ArgumentParser(description="Fast-dVLM sglang chatbot")
parser.add_argument(
"--model-path", default="Efficient-Large-Model/Fast_dVLM_3B",
help="HuggingFace model id or local path to Fast_dVLM checkpoint.",
)
parser.add_argument(
"--processor-path", default="Qwen/Qwen2.5-VL-3B-Instruct",
help="HF processor for chat template + image preprocessing.",
)
parser.add_argument(
"--algorithm", choices=list(ALGO_MAP.keys()), default="mdm",
help="mdm = HierarchyBlock (block diffusion); spec = SpeculativeBlock.",
)
parser.add_argument("--image", default=None, help="Image URL or local path. Empty for text-only.")
parser.add_argument("--prompt", default=None, help="Text prompt. If omitted, enters interactive mode.")
parser.add_argument("--max-tokens", type=int, default=512)
parser.add_argument("--mem-fraction-static", type=float, default=0.75)
parser.add_argument(
"--quantization", default=None,
choices=["w8a8_fp8"],
help="Quantization format of the checkpoint. Leave unset for BF16. "
"w8a8_fp8 requires SM89+ (4090 / L40 / H100 / H200).",
)
args = parser.parse_args()
os.environ.setdefault("SGLANG_DISABLE_CUDNN_CHECK", "1")
import sglang as sgl
from transformers import AutoProcessor, AutoTokenizer
processor = AutoProcessor.from_pretrained(args.processor_path, use_fast=False)
tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
processor.tokenizer = tokenizer
dllm_algo = ALGO_MAP[args.algorithm]
engine_kwargs = dict(
model_path=args.model_path,
trust_remote_code=True,
dtype="bfloat16",
mem_fraction_static=args.mem_fraction_static,
max_running_requests=1,
chunked_prefill_size=16384,
dllm_algorithm=dllm_algo,
disable_cuda_graph=False,
log_level="warning",
enable_metrics=True,
mm_attention_backend="triton_attn",
)
if args.quantization:
engine_kwargs["quantization"] = args.quantization
print(
f"Launching sglang Engine with dllm_algorithm={dllm_algo}"
f"{f', quantization={args.quantization}' if args.quantization else ''} ..."
)
engine = sgl.Engine(**engine_kwargs)
sampling = {"max_new_tokens": args.max_tokens, "temperature": 0.0}
def run_once(image, prompt):
input_ids = build_inputs(processor, image, prompt)
out = engine.generate(
input_ids=input_ids,
image_data=[image] if image else None,
sampling_params=sampling,
)
if isinstance(out, list):
out = out[0]
return out["text"]
try:
if args.prompt:
image = args.image or "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
print(f"\n{run_once(image, args.prompt)}")
else:
print("Interactive mode. Type 'exit' to quit, 'clear' to reset.")
while True:
prompt = input("\nYou: ").strip()
if prompt.lower() == "exit":
break
if prompt.lower() == "clear":
print("History cleared.")
continue
if not prompt:
continue
print(f"\nAssistant: {run_once(args.image, prompt)}")
finally:
engine.shutdown()
if __name__ == "__main__":
main()
|