diff --git a/transformerlab/models/lmstudiomodel.py b/transformerlab/models/lmstudiomodel.py new file mode 100644 index 000000000..414bb5d24 --- /dev/null +++ b/transformerlab/models/lmstudiomodel.py @@ -0,0 +1,110 @@ +from transformerlab.models import basemodel + +import os +import json +import errno +import shutil + +_LM_MODEL_EXTS = (".gguf", ".safetensors", ".pt", ".bin") + + +async def list_models(): + try: + models_dir = lmstudio_models_dir() + except Exception as e: + print("Failed to locate LM Studio models directory:") + print(str(e)) + return [] + + if not models_dir: + return [] + + models = [] + for root, _, files in os.walk(models_dir): + for fname in files: + if fname.lower().endswith(_LM_MODEL_EXTS): + model_path = os.path.join(root, fname) + models.append(LMStudioModel(model_path)) + + return models + + +class LMStudioModel(basemodel.BaseModel): + def __init__(self, model_path: str): + filename = os.path.basename(model_path) + super().__init__(model_id=filename) + + self.source = "lmstudio" + self.name = f"{os.path.splitext(filename)[0]} (LM Studio)" + self.source_id_or_path = os.path.abspath(model_path) + self.model_filename = filename + + async def get_json_data(self): + json_data = await super().get_json_data() + + ext = os.path.splitext(self.model_filename)[1].lower() + if ext == ".gguf": + json_data["architecture"] = "GGUF" + json_data["formats"] = ["GGUF"] + elif ext in (".safetensors", ".pt", ".bin"): + json_data["architecture"] = "PyTorch" + json_data["formats"] = ["safetensors" if ext == ".safetensors" else "pt"] + else: + json_data["architecture"] = "" + json_data["formats"] = [] + + json_data["source_id_or_path"] = self.source_id_or_path + return json_data + + async def install(self): + input_model_path = self.source_id_or_path + if not input_model_path or not os.path.isfile(input_model_path): + raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), input_model_path) + + from lab.dirs import get_models_dir + + output_filename = self.id + output_path = os.path.join(get_models_dir(), output_filename) + + if os.path.exists(output_path): + raise FileExistsError(errno.EEXIST, "Model already exists", output_path) + os.makedirs(output_path, exist_ok=True) + + link_name = os.path.join(output_path, output_filename) + os.symlink(input_model_path, link_name) + + json_data = await self.get_json_data() + + model_description = { + "model_id": self.id, + "model_filename": output_filename, + "name": self.name, + "source": self.source, + "json_data": { + "uniqueID": self.id, + "name": self.name, + "model_filename": output_filename, + "description": f"LM Studio model {self.source_id_or_path}", + "source": self.source, + "architecture": json_data["architecture"], + }, + } + + model_info_file = os.path.join(output_path, "index.json") + with open(model_info_file, "w") as f: + json.dump(model_description, f) + + +def lmstudio_models_dir(): + try: + lm_dir = os.environ["LMSTUDIO_MODELS"] + except KeyError: + lm_dir = os.path.join(os.path.expanduser("~"), ".lmstudio", "models") + + if os.path.isdir(lm_dir): + return lm_dir + + if shutil.which("lmstudio"): + return lm_dir + + return None diff --git a/transformerlab/models/model_helper.py b/transformerlab/models/model_helper.py index 0cae6d5d3..695704973 100644 --- a/transformerlab/models/model_helper.py +++ b/transformerlab/models/model_helper.py @@ -10,6 +10,7 @@ from transformerlab.models import ollamamodel from transformerlab.models import huggingfacemodel from transformerlab.models import localmodel +from transformerlab.models import lmstudiomodel import traceback @@ -49,7 +50,7 @@ def list_model_sources(): Supported strings that can be passsed as model_source to the functons that follow. """ - return ["huggingface", "ollama"] + return ["huggingface", "ollama", "lmstudio"] def get_model_by_source_id(model_source: str, model_source_id: str): @@ -65,6 +66,8 @@ def get_model_by_source_id(model_source: str, model_source_id: str): return ollamamodel.OllamaModel(model_source_id) case "huggingface": return huggingfacemodel.HuggingFaceModel(model_source_id) + case "lmstudio": + return lmstudiomodel.LMStudioModel(model_source_id) except Exception: print(f"Caught exception getting model {model_source_id} from {model_source}:") traceback.print_exc() diff --git a/transformerlab/plugins/lmstudio_server/index.json b/transformerlab/plugins/lmstudio_server/index.json new file mode 100644 index 000000000..1c6a8a0a8 --- /dev/null +++ b/transformerlab/plugins/lmstudio_server/index.json @@ -0,0 +1,18 @@ +{ + "name": "Google LM Studio Server", + "uniqueId": "lmstudio_server", + "description": "Google LM Studio loads models for inference using LM Studio for generation.", + "plugin-format": "python", + "type": "loader", + "version": "0.0.47", + "supports": ["chat", "completion", "embeddings"], + "files": ["main.py", "setup.sh"], + "parameters": { + "port": { + "title": "Server Port", + "type": "integer", + "default": 1234 + } + }, + "setup-script": "setup.sh" +} diff --git a/transformerlab/plugins/lmstudio_server/main.py b/transformerlab/plugins/lmstudio_server/main.py new file mode 100644 index 000000000..bc5199695 --- /dev/null +++ b/transformerlab/plugins/lmstudio_server/main.py @@ -0,0 +1,80 @@ +""" +LM Studio Server + +This plugin integrates LM Studio models into TransformerLab, allowing users to utilize models stored in the LM Studio format. +""" + +import argparse +import os +import subprocess +import json +import uuid +from hashlib import sha256 +from pathlib import Path +import sys +import lmstudio +import time +import requests + +worker_id = str(uuid.uuid4())[:8] + +LMSTUDIO_STARTUP_TIMEOUT = 180 # seconds + +try: + from transformerlab.plugin import register_process +except ImportError: + from transformerlab.plugin_sdk.transformerlab.plugin import register_process + +parser = argparse.ArgumentParser() +parser.add_argument("--model-path", type=str) +parser.add_argument("--parameters", type=str, default="{}") + +args, unknown = parser.parse_known_args() +# model_path can be a hugging face ID or a local file in Transformer Lab +# But LM Studio models are always stored as a local path because +# we are using a specific LM Studio model file +if os.path.exists(args.model_path): + model_path = args.model_path +else: + raise FileNotFoundError( + f"The specified LM Studio model '{args.model_path}' was not found. Please select a valid LM Studio model file to proceed." + ) + +llmlab_root_dir = os.getenv("LLM_LAB_ROOT_PATH") + +parameters = args.parameters +parameters = json.loads(parameters) +# Now go through the parameters object and remove the key that is equal to "inferenceEngine": +if "inferenceEngine" in parameters: + del parameters["inferenceEngine"] + +if "inferenceEngineFriendlyName" in parameters: + del parameters["inferenceEngineFriendlyName"] + +# Get plugin directory +real_plugin_dir = os.path.realpath(os.path.dirname(__file__)) + +port = int(parameters.get("port", 1234)) +env = os.environ.copy() +env["LMSTUDIO_HOST"] = f"127.0.0.1:{port}" +print("Starting LM Studio server...", file=sys.stderr) + +process = subprocess.Popen(["lms", "server", "start", f"--port {port}"], env=env) + +lmstudio_models_url = f"http://127.0.0.1:{port}/v1/models" +start_time = time.time() +while True: + try: + response = requests.get(lmstudio_models_url) + if response.status_code == 200: + print("LM Studio server is up and running.", file=sys.stderr) + break + except requests.ConnectionError: + pass + if time.time() - start_time > LMSTUDIO_STARTUP_TIMEOUT: + print("Timeout waiting for LM Studio server to start.", file=sys.stderr) + process.terminate() + sys.exit(1) + time.sleep(1) + +register_process(process.pid) diff --git a/transformerlab/plugins/lmstudio_server/setup.sh b/transformerlab/plugins/lmstudio_server/setup.sh new file mode 100644 index 000000000..eacda34ee --- /dev/null +++ b/transformerlab/plugins/lmstudio_server/setup.sh @@ -0,0 +1 @@ +uv pip install lmstudio \ No newline at end of file