Bläddra i källkod

Add support for diffusion_models directory - treat .gguf files as checkpoint models

Fszontagh 3 månader sedan
förälder
incheckning
4de1542abd
2 ändrade filer med 13 tillägg och 2 borttagningar
  1. 9 2
      include/model_manager.h
  2. 4 0
      src/model_manager.cpp

+ 9 - 2
include/model_manager.h

@@ -31,7 +31,8 @@ enum class ModelType : uint32_t {
     ESRGAN = 128,
     CONTROLNET = 256,
     UPSCALER = 512,
-    EMBEDDING = 1024
+    EMBEDDING = 1024,
+    DIFFUSION_MODELS = 2048
 };
 
 // Enable bitwise operations for ModelType
@@ -96,6 +97,12 @@ public:
     /**
      * @brief Scan the models directory to discover available models
      *
+     * Recursively scans all subdirectories within the models directory to find
+     * 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.
+     *
      * @return true if scanning was successful, false otherwise
      */
     bool scanModelsDirectory();
@@ -303,4 +310,4 @@ private:
     std::unique_ptr<Impl> pImpl; // Pimpl idiom
 };
 
-#endif // MODEL_MANAGER_H
+#endif // MODEL_MANAGER_H

+ 4 - 0
src/model_manager.cpp

@@ -231,6 +231,10 @@ public:
                         return ModelType::ESRGAN;
                     }
                 } else if (dirName == "embeddings" || dirName == "textual-inversion") {
+                } else if (dirName == "diffusion_models" || dirName == "diffusion") {
+                    if (isExtensionMatch(extension, ModelType::CHECKPOINT)) {
+                        return ModelType::CHECKPOINT;
+                    }
                     if (isExtensionMatch(extension, ModelType::EMBEDDING)) {
                         return ModelType::EMBEDDING;
                     }