[WIP] Added Support for importing model from lm studio

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

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()

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"
+}

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)

transformerlab/plugins/lmstudio_server/setup.sh

@@ -0,0 +1 @@
+uv pip install lmstudio