|
|
@@ -132,19 +132,22 @@ class ModelSingleton:
|
|
|
if "cache_max_entry_count" not in kwargs:
|
|
|
kwargs["cache_max_entry_count"] = 0.5
|
|
|
|
|
|
- if "lmdeploy_device" in kwargs:
|
|
|
- device_type = kwargs.pop("lmdeploy_device")
|
|
|
- if device_type not in ["cuda", "ascend", "maca", "camb"]:
|
|
|
- raise ValueError(f"Unsupported lmdeploy device type: {device_type}")
|
|
|
- else:
|
|
|
- device_type = "cuda"
|
|
|
-
|
|
|
- if "lmdeploy_backend" in kwargs:
|
|
|
- lm_backend = kwargs.pop("lmdeploy_backend")
|
|
|
- if lm_backend not in ["pytorch", "turbomind"]:
|
|
|
- raise ValueError(f"Unsupported lmdeploy backend: {lm_backend}")
|
|
|
- else:
|
|
|
- lm_backend = set_lmdeploy_backend(device_type)
|
|
|
+ device_type = os.getenv("MINERU_LMDEPLOY_DEVICE", "")
|
|
|
+ if device_type == "":
|
|
|
+ if "lmdeploy_device" in kwargs:
|
|
|
+ device_type = kwargs.pop("lmdeploy_device")
|
|
|
+ if device_type not in ["cuda", "ascend", "maca", "camb"]:
|
|
|
+ raise ValueError(f"Unsupported lmdeploy device type: {device_type}")
|
|
|
+ else:
|
|
|
+ device_type = "cuda"
|
|
|
+ lm_backend = os.getenv("MINERU_LMDEPLOY_BACKEND", "")
|
|
|
+ if lm_backend == "":
|
|
|
+ if "lmdeploy_backend" in kwargs:
|
|
|
+ lm_backend = kwargs.pop("lmdeploy_backend")
|
|
|
+ if lm_backend not in ["pytorch", "turbomind"]:
|
|
|
+ raise ValueError(f"Unsupported lmdeploy backend: {lm_backend}")
|
|
|
+ else:
|
|
|
+ lm_backend = set_lmdeploy_backend(device_type)
|
|
|
logger.info(f"lmdeploy device is: {device_type}, lmdeploy backend is: {lm_backend}")
|
|
|
|
|
|
if lm_backend == "pytorch":
|