import os
from lmdeploy import pipeline, TurbomindEngineConfig, GenerationConfig

IMAGE_TOKEN = "[IMAGE_TOKEN]"

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

# Model initialization
model = "OpenGVLab/InternVL2-26B-AWQ"
pipe = pipeline(
    model,
    backend_config=TurbomindEngineConfig(
        model_format="awq",
        tp=2, 
        # tp=4,
        session_len=2048,  # 4096, 8192, 16384, 32768
        max_batch_size=1,
        cache_max_entry_count=0.1, # 0.05
        cache_block_seq_len=4096,  # 8192, 16384, 32768
        quant_policy=4,
        # precision="fp16",
    ), 
    # log_level='DEBUG'
)