Merge pull request #670 from transformerlab/fix/loading-local-models

Fix loading locally available models and also saving local models to have filenames in them

Author: deep1401

transformerlab/models/localmodel.py

@@ -142,33 +142,88 @@ async def list_models(self, embedding=False):
                 models.remove(model)
                 continue
             # Only set model["stored_in_filesystem"] to True if the model is a local model and not a Hugging Face model
-            if (
-                not model.get("json_data", {}).get("source", "") == "huggingface"
-                and not model.get("json_data", {}).get("model_filename", "") == ""
-            ):
+            # model_filename can be:
+            # - A filename (e.g., "model.gguf") for file-based models
+            # - "." for directory-based models (indicates the directory itself)
+            # - Empty string for legacy models (should be treated as directory-based)
+            model_filename = model.get("json_data", {}).get("model_filename", "")
+            is_huggingface = model.get("json_data", {}).get("source", "") == "huggingface"
+            has_model_filename = model_filename != ""
+            
+            # Determine the potential model directory path
+            # This applies to both HuggingFace models stored locally and local models
+            model_id = model.get("model_id", "")
+            potential_path = os.path.join(models_dir, secure_filename(model_id))
+            # Check if local path exists
+            if not os.path.exists(potential_path):
+                # Remove the Starting TransformerLab/ prefix to handle the save_transformerlab_model function
+                potential_path = os.path.join(models_dir, secure_filename("/".join(model_id.split("/")[1:])))
+            
+            # Check if model should be considered local:
+            # 1. If it has a model_filename set (and is not a HuggingFace model, OR is a HuggingFace model stored locally), OR
+            # 2. If the directory exists and has files other than index.json
+            is_local_model = False
+            if not is_huggingface:
+                # For non-HuggingFace models, check if it has model_filename or files in directory
+                if has_model_filename:
+                    is_local_model = True
+                elif os.path.exists(potential_path) and os.path.isdir(potential_path):
+                    # Check if directory has files other than index.json
+                    try:
+                        files = os.listdir(potential_path)
+                        # Filter out index.json and other metadata files
+                        model_files = [f for f in files if f not in ["index.json", "_tlab_provenance.json"]]
+                        if model_files:
+                            is_local_model = True
+                    except (OSError, PermissionError):
+                        # If we can't read the directory, skip it
+                        pass
+            elif is_huggingface and has_model_filename:
+                # For HuggingFace models, if they have a model_filename and the file/directory exists locally,
+                # treat them as stored locally (e.g., downloaded GGUF files)
+                if os.path.exists(potential_path):
+                    is_local_model = True
+            
+            if is_local_model:
                 # tells the app this model was loaded from workspace directory
                 model["stored_in_filesystem"] = True
-
-                # Set local_path to the filesystem location
-                # this will tell Hugging Face to not try downloading
-                model_id = model.get("model_id", "")
-                model_filename = model.get("json_data", {}).get("model_filename", "")
-                model["local_path"] = os.path.join(models_dir, secure_filename(model_id))
-                # Check if local path exists
-                if not os.path.exists(model["local_path"]):
-                    # Remove the Starting TransformerLab/ prefix to handle the save_transformerlab_model function
-                    model["local_path"] = os.path.join(models_dir, secure_filename("/".join(model_id.split("/")[1:])))
-
-                # Some models are a single file (possibly of many in a directory, e.g. GGUF)
-                # For models that have model_filename set we should link directly to that specific file
-                if "model_filename" in model.get("json_data", {}):
-                    model_filename = model["json_data"]["model_filename"]
-                    if model_filename.endswith(".gguf"):
-                        model["local_path"] = os.path.join(
-                            os.path.join(models_dir, secure_filename(model_id)), model_filename
-                        )
+                model["local_path"] = potential_path
+
+                # Handle different model_filename cases
+                if model_filename == ".":
+                    # Directory-based model - convert to absolute path so it can be used anywhere
+                    model["local_path"] = os.path.abspath(model["local_path"])
+                elif model_filename and model_filename.endswith(".gguf"):
+                    # GGUF file - append the filename to the model directory and convert to absolute path
+                    # This ensures we get the full path like: /path/to/models/dir/model.gguf
+                    base_path = model["local_path"]
+                    model_path = os.path.join(base_path, model_filename)
+                    if os.path.exists(model_path):
+                        if os.path.isdir(model_path):
+                            # List all files in the directory ending with .gguf
+                            gguf_files = [f for f in os.listdir(model_path) if f.endswith(".gguf")]
+                            if gguf_files:
+                                model_path = os.path.join(model_path, gguf_files[0])
                     else:
-                        model["local_path"] = os.path.join(model["local_path"], model["json_data"]["model_filename"])
+                        # Seearch for files ending with .gguf in the directory
+                        gguf_files = [f for f in os.listdir(model["local_path"]) if f.endswith(".gguf")]
+                        if gguf_files:
+                            gguf_file = gguf_files[0]
+                            model_path = os.path.join(base_path, gguf_file)
+                            if os.path.isdir(model_path):
+                                gguf_files = [f for f in os.listdir(model_path) if f.endswith(".gguf")]
+                                if gguf_files:
+                                    model_path = os.path.join(model_path, gguf_files[0])
+                                
+                                
+
+                    model["local_path"] = os.path.abspath(model_path)
+                elif model_filename:
+                    # Other file-based models - append the filename and convert to absolute path
+                    model["local_path"] = os.path.abspath(os.path.join(model["local_path"], model_filename))
+                else:
+                    # Legacy model without model_filename but with files - use directory path
+                    model["local_path"] = os.path.abspath(model["local_path"])
 
         # Filter out models based on whether they are embedding models or not
         models = await self.filter_embedding_models(models, embedding)

transformerlab/plugin_sdk/transformerlab/sdk/v1/train.py

@@ -389,14 +389,31 @@ def create_transformerlab_model(
             json_data = json_data.copy() if json_data else {}
             json_data["pipeline_tag"] = pipeline_tag
 
-        if generate_json:
-            generate_model_json(fused_model_name, model_architecture, json_data=json_data, output_directory=output_dir)
-
         if output_dir is None:
             fused_model_location = os.path.join(WORKSPACE_DIR, "models", fused_model_name)
         else:
             fused_model_location = os.path.join(output_dir, fused_model_name)
 
+        # Determine model_filename based on architecture
+        # Most models are directory-based, only GGUF models are file-based
+        # Default to directory-based (use "." to indicate the directory itself)
+        model_filename = "."
+        
+        # GGUF architecture indicates a file-based model
+        # The actual filename will be set by the export process, so we don't set it here
+        # For now, if it's GGUF and the file exists, use the filename
+        if "GGUF" in model_architecture.upper() or model_architecture.upper() == "GGUF":
+            if os.path.exists(fused_model_location):
+                if os.path.isfile(fused_model_location):
+                    # File-based model - use the filename
+                    model_filename = os.path.basename(fused_model_location)
+                # If it's a directory for GGUF, keep "." (directory-based)
+                # This shouldn't normally happen for GGUF, but handle it gracefully
+            # If GGUF file doesn't exist yet, the export process will set the filename
+
+        if generate_json:
+            generate_model_json(fused_model_name, model_architecture, model_filename=model_filename, json_data=json_data, output_directory=output_dir)
+
         # Create the hash files for the model
         md5_objects = self.create_md5_checksum_model_files(fused_model_location)
 

transformerlab/routers/experiment/export.py

@@ -82,6 +82,9 @@ async def run_exporter_script(
             output_model_id = f"{input_model_id_without_author}-{conversion_time}-{q_type}.gguf"
 
         output_filename = output_model_id
+    else:
+        # For directory-based models (non-GGUF), set model_filename to "." to indicate the directory itself
+        output_filename = "."
 
     # Figure out plugin and model output directories
     script_directory = lab_dirs.plugin_dir_by_name(plugin_name)