stable-diffusion.cpp-rest/include/model_manager.h

#ifndef MODEL_MANAGER_H
#define MODEL_MANAGER_H

#include <string>
#include <memory>
#include <map>
#include <shared_mutex>
#include <vector>
#include <filesystem>
#include <cstdint>

// Forward declarations
class StableDiffusionWrapper;

#include "server_config.h"

/**
 * @brief Model type enumeration
 *
 * These values are bit flags that can be combined to filter model types.
 */
enum class ModelType : uint32_t {
    NONE = 0,
    LORA = 1,
    CHECKPOINT = 2,
    VAE = 4,
    PRESETS = 8,
    PROMPTS = 16,
    NEG_PROMPTS = 32,
    TAESD = 64,
    ESRGAN = 128,
    CONTROLNET = 256,
    UPSCALER = 512,
    EMBEDDING = 1024
};

// Enable bitwise operations for ModelType
inline ModelType operator|(ModelType a, ModelType b) {
    return static_cast<ModelType>(static_cast<uint32_t>(a) | static_cast<uint32_t>(b));
}

inline ModelType operator&(ModelType a, ModelType b) {
    return static_cast<ModelType>(static_cast<uint32_t>(a) & static_cast<uint32_t>(b));
}

inline ModelType& operator|=(ModelType& a, ModelType b) {
    a = a | b;
    return a;
}

/**
 * @brief Model manager class for loading and managing stable-diffusion models
 *
 * This class handles loading, unloading, and managing multiple stable-diffusion models.
 * It provides thread-safe access to models and manages model resources efficiently.
 */
class ModelManager {
public:
    /**
     * @brief Model information structure
     */
    struct ModelInfo {
        std::string name;                ///< Model name
        std::string path;                ///< Model file path
        std::string fullPath;            ///< Absolute path to the model file
        ModelType type;                  ///< Model type
        bool isLoaded;                   ///< Whether the model is currently loaded
        size_t fileSize;                 ///< File size in bytes
        std::string sha256;              ///< SHA256 hash of the file
        std::filesystem::file_time_type createdAt;       ///< File creation time
        std::filesystem::file_time_type modifiedAt;     ///< Last modification time
        std::string description;         ///< Model description
        std::map<std::string, std::string> metadata;   ///< Additional metadata

        // Architecture detection fields
        std::string architecture;        ///< Detected architecture (e.g., "Stable Diffusion XL Base", "Flux Dev")
        std::string recommendedVAE;      ///< Recommended VAE for this model
        int recommendedWidth = 0;        ///< Recommended image width
        int recommendedHeight = 0;       ///< Recommended image height
        int recommendedSteps = 0;        ///< Recommended number of steps
        std::string recommendedSampler;  ///< Recommended sampler
        std::vector<std::string> requiredModels; ///< List of required auxiliary models (VAE, CLIP, etc.)
        std::vector<std::string> missingModels;  ///< List of missing required models
    };

    /**
     * @brief Construct a new Model Manager object
     */
    ModelManager();

    /**
     * @brief Destroy the Model Manager object
     */
    virtual ~ModelManager();

    /**
     * @brief Scan the models directory to discover available models
     *
     * @return true if scanning was successful, false otherwise
     */
    bool scanModelsDirectory();

    /**
     * @brief Cancel any ongoing model directory scanning
     */
    void cancelScan();

    /**
     * @brief Load a model from the specified path
     *
     * @param name The name to assign to the model
     * @param path The file path to the model
     * @param type The type of model
     * @return true if the model was loaded successfully, false otherwise
     */
    bool loadModel(const std::string& name, const std::string& path, ModelType type);

    /**
     * @brief Load a model by name (must be discovered first)
     *
     * @param name The name of the model to load
     * @return true if the model was loaded successfully, false otherwise
     */
    bool loadModel(const std::string& name);

    /**
     * @brief Unload a model
     *
     * @param name The name of the model to unload
     * @return true if the model was unloaded successfully, false otherwise
     */
    bool unloadModel(const std::string& name);

    /**
     * @brief Get a pointer to a loaded model
     *
     * @param name The name of the model
     * @return StableDiffusionWrapper* Pointer to the model wrapper, or nullptr if not found
     */
    StableDiffusionWrapper* getModel(const std::string& name);

    /**
     * @brief Get information about all models
     *
     * @return std::map<std::string, ModelInfo> Map of model names to their information
     */
    std::map<std::string, ModelInfo> getAllModels() const;

    /**
     * @brief Get information about models of a specific type
     *
     * @param type The model type to filter by
     * @return std::vector<ModelInfo> List of model information
     */
    std::vector<ModelInfo> getModelsByType(ModelType type) const;

    /**
     * @brief Get information about a specific model
     *
     * @param name The name of the model
     * @return ModelInfo Model information, or empty if not found
     */
    ModelInfo getModelInfo(const std::string& name) const;

    /**
     * @brief Check if a model is loaded
     *
     * @param name The name of the model
     * @return true if the model is loaded, false otherwise
     */
    bool isModelLoaded(const std::string& name) const;

    /**
     * @brief Get the number of loaded models
     *
     * @return size_t Number of loaded models
     */
    size_t getLoadedModelsCount() const;

    /**
     * @brief Get the number of available models
     *
     * @return size_t Number of available models
     */
    size_t getAvailableModelsCount() const;

    /**
     * @brief Set the models directory path
     *
     * @param path The path to the models directory
     */
    void setModelsDirectory(const std::string& path);

    /**
     * @brief Get the models directory path
     *
     * @return std::string The models directory path
     */
    std::string getModelsDirectory() const;

    /**
     * @brief Set directory for a specific model type
     *
     * @param type The model type
     * @param path The directory path
     * @return true if the directory was set successfully, false otherwise
     */
    bool setModelTypeDirectory(ModelType type, const std::string& path);

    /**
     * @brief Get directory for a specific model type
     *
     * @param type The model type
     * @return std::string The directory path, empty if not set
     */
    std::string getModelTypeDirectory(ModelType type) const;

    /**
     * @brief Set all model type directories at once
     *
     * @param directories Map of model types to directory paths
     * @return true if all directories were set successfully, false otherwise
     */
    bool setAllModelTypeDirectories(const std::map<ModelType, std::string>& directories);

    /**
     * @brief Get all model type directories
     *
     * @return std::map<ModelType, std::string> Map of model types to directory paths
     */
    std::map<ModelType, std::string> getAllModelTypeDirectories() const;

    /**
     * @brief Reset to legacy directory mode (single models directory)
     */
    void resetToLegacyDirectories();

    /**
     * @brief Configure ModelManager with ServerConfig
     *
     * @param config The server configuration
     * @return true if configuration was successful, false otherwise
     */
    bool configureFromServerConfig(const struct ServerConfig& config);

    /**
     * @brief Convert ModelType to string
     *
     * @param type The model type
     * @return std::string String representation of the model type
     */
    static std::string modelTypeToString(ModelType type);

    /**
     * @brief Convert string to ModelType
     *
     * @param typeStr String representation of the model type
     * @return ModelType The model type
     */
    static ModelType stringToModelType(const std::string& typeStr);

    /**
     * @brief Compute SHA256 hash of a model file
     *
     * @param modelName The name of the model
     * @return std::string The SHA256 hash, or empty string on error
     */
    std::string computeModelHash(const std::string& modelName);

    /**
     * @brief Load hash from JSON file for a model
     *
     * @param modelName The name of the model
     * @return std::string The loaded hash, or empty string if not found
     */
    std::string loadModelHashFromFile(const std::string& modelName);

    /**
     * @brief Save hash to JSON file for a model
     *
     * @param modelName The name of the model
     * @param hash The SHA256 hash to save
     * @return true if saved successfully, false otherwise
     */
    bool saveModelHashToFile(const std::string& modelName, const std::string& hash);

    /**
     * @brief Find model by hash (full or partial - minimum 10 chars)
     *
     * @param hash Full or partial SHA256 hash (minimum 10 characters)
     * @return std::string Model name, or empty string if not found
     */
    std::string findModelByHash(const std::string& hash);

    /**
     * @brief Load hash for a model (from file or compute if missing)
     *
     * @param modelName The name of the model
     * @param forceCompute Force recomputation even if hash file exists
     * @return std::string The SHA256 hash, or empty string on error
     */
    std::string ensureModelHash(const std::string& modelName, bool forceCompute = false);

private:
    class Impl;
    std::unique_ptr<Impl> pImpl; // Pimpl idiom
};

#endif // MODEL_MANAGER_H