Selaa lähdekoodia

Fix model manager path calculation bug

- Fixed relative path calculation to use model type directory as base
- Model names now exclude default directory prefixes (checkpoints/, esrgan/, etc.)
- Fixed ESRGAN model duplication issue
- Checkpoints are no longer misidentified when in wrong directories
- ESRGAN models now correctly identified as type 'esrgan' instead of 'checkpoint'
Fszontagh 3 kuukautta sitten
vanhempi
sitoutus
ea197f1e75
2 muutettua tiedostoa jossa 137 lisäystä ja 39 poistoa
  1. 68 39
      src/model_manager.cpp
  2. 69 0
      test_subfolder_scanning.sh

+ 68 - 39
src/model_manager.cpp

@@ -9,6 +9,7 @@
 #include <chrono>
 #include <future>
 #include <atomic>
+#include <set>
 #include <openssl/evp.h>
 #include <sstream>
 #include <iomanip>
@@ -148,7 +149,10 @@ public:
     }
 
     /**
-     * @brief Determine model type based on file path and extension
+     * @brief FIXED: Determine model type based on file path and extension
+     *
+     * THIS IS THE FIXED VERSION - no extension-based fallback!
+     * Only returns a model type if the file is actually in the right directory.
      *
      * @param filePath The file path
      * @return ModelType The determined model type
@@ -237,24 +241,8 @@ public:
             }
         }
 
-        // Fall back to extension-based detection
-        // Only return a model type if the extension matches expected extensions for that type
-        if (isExtensionMatch(extension, ModelType::CHECKPOINT)) {
-            return ModelType::CHECKPOINT;
-        } else if (isExtensionMatch(extension, ModelType::LORA)) {
-            return ModelType::LORA;
-        } else if (isExtensionMatch(extension, ModelType::VAE)) {
-            return ModelType::VAE;
-        } else if (isExtensionMatch(extension, ModelType::TAESD)) {
-            return ModelType::TAESD;
-        } else if (isExtensionMatch(extension, ModelType::ESRGAN)) {
-            return ModelType::ESRGAN;
-        } else if (isExtensionMatch(extension, ModelType::CONTROLNET)) {
-            return ModelType::CONTROLNET;
-        } else if (isExtensionMatch(extension, ModelType::EMBEDDING)) {
-            return ModelType::EMBEDDING;
-        }
-
+        // NO EXTENSION-BASED FALLBACK - this was the bug!
+        // Files must be in the correct directory to be recognized
         return ModelType::NONE;
     }
 
@@ -287,14 +275,19 @@ public:
     }
 
     /**
-     * @brief Scan a directory for models of a specific type (without holding mutex)
+     * @brief Scan a directory for models recursively (all types, without holding mutex)
+     *
+     * Recursively walks the directory tree to find all model files. For each model
+     * found, constructs the display name as 'relative_path/model_name' where
+     * relative_path is the path from the models root directory to the file's
+     * containing folder (using forward slashes). Models in the root directory
+     * appear without a prefix.
      *
      * @param directory The directory to scan
-     * @param type The model type to look for
      * @param modelsMap Reference to the map to store results
      * @return bool True if scanning completed without cancellation
      */
-    bool scanDirectory(const fs::path& directory, ModelType type, std::map<std::string, ModelInfo>& modelsMap) {
+    bool scanDirectory(const fs::path& directory, std::map<std::string, ModelInfo>& modelsMap) {
         if (scanCancelled.load()) {
             return false;
         }
@@ -314,11 +307,33 @@ public:
                     ModelType detectedType = determineModelType(filePath);
 
                     // Only add files that have a valid model type (not NONE)
-                    if (detectedType != ModelType::NONE && (type == ModelType::NONE || detectedType == type)) {
+                    if (detectedType != ModelType::NONE) {
                         ModelInfo info;
 
-                        // Calculate relative path from the scanned directory (not base models directory)
-                        fs::path relativePath = fs::relative(filePath, directory);
+                        // Calculate relative path from the model type directory to exclude model type folder names
+                        fs::path relativePath;
+                        try {
+                            // Get the specific model type directory for this detected type
+                            std::string modelTypeDir = getModelTypeDirectory(detectedType);
+
+                            if (!modelTypeDir.empty()) {
+                                fs::path typeBaseDir(modelTypeDir);
+                                // Get relative path from the model type directory
+                                relativePath = fs::relative(filePath, typeBaseDir);
+                            } else {
+                                // Fallback: use the base models directory
+                                if (!modelsDirectory.empty()) {
+                                    fs::path baseDir(modelsDirectory);
+                                    relativePath = fs::relative(filePath, baseDir);
+                                } else {
+                                    relativePath = fs::relative(filePath, directory);
+                                }
+                            }
+                        } catch (const fs::filesystem_error&) {
+                            // If relative path calculation fails, use filename only
+                            relativePath = filePath.filename();
+                        }
+
                         std::string modelName = relativePath.string();
 
                         // Normalize path separators for consistency
@@ -460,22 +475,36 @@ bool ModelManager::scanModelsDirectory() {
     // Create temporary map to store scan results (outside of lock)
     std::map<std::string, ModelInfo> tempModels;
 
-    // Explicit mode: scan configured directories for each model type
-    std::vector<std::pair<ModelType, std::string>> directoriesToScan = {
-        {ModelType::CHECKPOINT, pImpl->getModelTypeDirectory(ModelType::CHECKPOINT)},
-        {ModelType::CONTROLNET, pImpl->getModelTypeDirectory(ModelType::CONTROLNET)},
-        {ModelType::LORA, pImpl->getModelTypeDirectory(ModelType::LORA)},
-        {ModelType::VAE, pImpl->getModelTypeDirectory(ModelType::VAE)},
-        {ModelType::TAESD, pImpl->getModelTypeDirectory(ModelType::TAESD)},
-        {ModelType::ESRGAN, pImpl->getModelTypeDirectory(ModelType::ESRGAN)},
-        {ModelType::EMBEDDING, pImpl->getModelTypeDirectory(ModelType::EMBEDDING)}
+    // Scan all configured directories for all model types recursively
+    // We scan each directory once and detect all model types within it
+    std::set<std::string> scannedDirectories; // Avoid scanning the same directory multiple times
+    std::vector<std::string> directoriesToScan;
+
+    // Collect unique directories to scan
+    std::vector<ModelType> allTypes = {
+        ModelType::CHECKPOINT, ModelType::CONTROLNET, ModelType::LORA,
+        ModelType::VAE, ModelType::TAESD, ModelType::ESRGAN, ModelType::EMBEDDING
     };
 
-    for (const auto& [type, dirPath] : directoriesToScan) {
-        if (!dirPath.empty()) {
-            if (!pImpl->scanDirectory(dirPath, type, tempModels)) {
-                return false;
-            }
+    for (const auto& type : allTypes) {
+        std::string dirPath = pImpl->getModelTypeDirectory(type);
+        if (!dirPath.empty() && scannedDirectories.find(dirPath) == scannedDirectories.end()) {
+            directoriesToScan.push_back(dirPath);
+            scannedDirectories.insert(dirPath);
+        }
+    }
+
+    // Also scan the base models directory if it exists and isn't already covered
+    if (!pImpl->modelsDirectory.empty() &&
+        fs::exists(pImpl->modelsDirectory) &&
+        scannedDirectories.find(pImpl->modelsDirectory) == scannedDirectories.end()) {
+        directoriesToScan.push_back(pImpl->modelsDirectory);
+    }
+
+    // Scan each unique directory recursively for all model types
+    for (const auto& dirPath : directoriesToScan) {
+        if (!pImpl->scanDirectory(dirPath, tempModels)) {
+            return false;
         }
     }
 

+ 69 - 0
test_subfolder_scanning.sh

@@ -0,0 +1,69 @@
+#!/bin/bash
+
+# Test script for subfolder model scanning
+# This script creates a test directory structure and verifies that subfolder models are detected
+
+set -e
+
+TEST_BASE="/tmp/sd_test_models"
+TEST_CHECKPOINTS="$TEST_BASE/checkpoints"
+
+echo "Setting up test environment..."
+
+# Clean up previous test
+rm -rf "$TEST_BASE"
+
+# Create test directory structure
+mkdir -p "$TEST_CHECKPOINTS/subfolder1"
+mkdir -p "$TEST_CHECKPOINTS/subfolder2/deep"
+mkdir -p "$TEST_CHECKPOINTS/subfolder2/nested"
+
+# Create dummy model files
+touch "$TEST_CHECKPOINTS/model1.safetensors"
+touch "$TEST_CHECKPOINTS/subfolder1/model2.safetensors"
+touch "$TEST_CHECKPOINTS/subfolder2/model3.safetensors"
+touch "$TEST_CHECKPOINTS/subfolder2/deep/model4.safetensors"
+touch "$TEST_CHECKPOINTS/subfolder2/nested/model5.safetensors"
+
+# Create other model types
+mkdir -p "$TEST_BASE/loras"
+touch "$TEST_BASE/loras/lora1.safetensors"
+mkdir -p "$TEST_BASE/loras/subfolder"
+touch "$TEST_BASE/loras/subfolder/lora2.safetensors"
+
+echo "Test directory structure created:"
+echo "- $TEST_CHECKPOINTS/model1.safetensors"
+echo "- $TEST_CHECKPOINTS/subfolder1/model2.safetensors"
+echo "- $TEST_CHECKPOINTS/subfolder2/model3.safetensors"
+echo "- $TEST_CHECKPOINTS/subfolder2/deep/model4.safetensors"
+echo "- $TEST_CHECKPOINTS/subfolder2/nested/model5.safetensors"
+echo "- $TEST_BASE/loras/lora1.safetensors"
+echo "- $TEST_BASE/loras/subfolder/lora2.safetensors"
+
+echo ""
+echo "Running model scanning test..."
+
+# Run the server with test directory
+cd /data/stable-diffusion.cpp-rest/build
+timeout 10s ./stable-diffusion-rest-server \
+    --models-dir "$TEST_BASE" \
+    --checkpoints "$TEST_CHECKPOINTS" \
+    --lora-dir "$TEST_BASE/loras" \
+    --host 127.0.0.1 \
+    --port 0 \
+    --verbose || true
+
+echo ""
+echo "Test completed. Check the output above to verify that subfolder models were detected."
+echo ""
+echo "Expected models to be found:"
+echo "- model1.safetensors"
+echo "- subfolder1/model2.safetensors"
+echo "- subfolder2/model3.safetensors"
+echo "- subfolder2/deep/model4.safetensors"
+echo "- subfolder2/nested/model5.safetensors"
+echo "- loras/lora1.safetensors"
+echo "- loras/subfolder/lora2.safetensors"
+
+# Cleanup
+rm -rf "$TEST_BASE"