Эх сурвалжийг харах

feat: enhance /api/models endpoint with model existence checking (#44)

- Add ModelDetails structure for comprehensive model information
- Implement model existence checking in ModelManager
- Enhanced JSON response with architecture-specific recommended fields
- Add support for VAE, CLIP-L, CLIP-G, T5XXL, CLIP-Vision, Qwen2VL models
- Maintain backward compatibility with existing API structure
- Include file size, SHA256 hash, and existence status for required models
Fszontagh 3 сар өмнө
parent
commit
773fa2e57f
1 өөрчлөгдсөн 135 нэмэгдсэн , 2 устгасан
  1. 135 2
      src/server.cpp

+ 135 - 2
src/server.cpp

@@ -950,6 +950,131 @@ void Server::handleApiStatus(const httplib::Request& /*req*/, httplib::Response&
     }
 }
 
+// Helper function to convert ModelDetails vector to JSON array
+nlohmann::json Server::modelDetailsToJson(const std::vector<ModelManager::ModelDetails>& modelDetails) {
+    nlohmann::json jsonArray = nlohmann::json::array();
+    
+    for (const auto& detail : modelDetails) {
+        nlohmann::json modelJson = {
+            {"name", detail.name},
+            {"exists", detail.exists},
+            {"type", detail.type},
+            {"file_size", detail.file_size}
+        };
+        
+        // Handle path and sha256 separately to avoid type mismatch
+        if (detail.exists) {
+            modelJson["path"] = detail.path;
+            modelJson["sha256"] = detail.sha256;
+        } else {
+            modelJson["path"] = nullptr;
+            modelJson["sha256"] = "";
+        }
+        
+        // Add conditional fields for required/recommended models
+        if (detail.is_required) {
+            modelJson["is_required"] = true;
+        }
+        if (detail.is_recommended) {
+            modelJson["is_recommended"] = true;
+        }
+        
+        jsonArray.push_back(modelJson);
+    }
+    
+    return jsonArray;
+}
+
+// Helper function to determine which recommended fields to include based on architecture
+std::map<std::string, bool> Server::getRecommendedModelFields(const std::string& architecture) {
+    std::map<std::string, bool> recommendedFields;
+    
+    // Initialize all fields as false (will be set to null if not applicable)
+    recommendedFields["recommended_vae"] = false;
+    recommendedFields["recommended_clip_l"] = false;
+    recommendedFields["recommended_clip_g"] = false;
+    recommendedFields["recommended_t5xxl"] = false;
+    recommendedFields["recommended_clip_vision"] = false;
+    recommendedFields["recommended_qwen2vl"] = false;
+    
+    // Architecture-specific field inclusion based on actual architecture strings
+    if (architecture.find("Stable Diffusion 1.5") != std::string::npos) {
+        // SD 1.x: recommended_vae only
+        recommendedFields["recommended_vae"] = true;
+    } else if (architecture.find("Stable Diffusion XL") != std::string::npos) {
+        // SDXL: recommended_vae only
+        recommendedFields["recommended_vae"] = true;
+    } else if (architecture.find("Modern Architecture") != std::string::npos ||
+               architecture.find("Flux Dev") != std::string::npos ||
+               architecture.find("Flux Chroma") != std::string::npos) {
+        // FLUX/SD3/Modern Architecture: recommended_vae, recommended_clip_l, recommended_t5xxl
+        recommendedFields["recommended_vae"] = true;
+        recommendedFields["recommended_clip_l"] = true;
+        recommendedFields["recommended_t5xxl"] = true;
+    } else if (architecture.find("SD 3") != std::string::npos) {
+        // SD3: recommended_vae, recommended_clip_l, recommended_clip_g, recommended_t5xxl
+        recommendedFields["recommended_vae"] = true;
+        recommendedFields["recommended_clip_l"] = true;
+        recommendedFields["recommended_clip_g"] = true;
+        recommendedFields["recommended_t5xxl"] = true;
+    } else if (architecture.find("Wan") != std::string::npos) {
+        // Wan models: recommended_vae, recommended_t5xxl, recommended_clip_vision
+        recommendedFields["recommended_vae"] = true;
+        recommendedFields["recommended_t5xxl"] = true;
+        recommendedFields["recommended_clip_vision"] = true;
+    } else if (architecture.find("Qwen") != std::string::npos) {
+        // Qwen models: recommended_vae, recommended_qwen2vl
+        recommendedFields["recommended_vae"] = true;
+        recommendedFields["recommended_qwen2vl"] = true;
+    }
+    // For UNKNOWN architecture, keep all fields false
+    
+    return recommendedFields;
+}
+
+// Helper function to populate recommended models with existence information
+void Server::populateRecommendedModels(nlohmann::json& response, const ModelManager::ModelInfo& modelInfo) {
+    if (modelInfo.requiredModels.empty()) {
+        return;
+    }
+    
+    // Check existence of required models
+    auto requiredModelsDetails = m_modelManager->checkRequiredModelsExistence(modelInfo.requiredModels);
+    
+    // Get the recommended fields for this architecture
+    auto recommendedFields = getRecommendedModelFields(modelInfo.architecture);
+    
+    // Group models by type
+    std::map<std::string, std::vector<ModelManager::ModelDetails>> modelsByType;
+    for (const auto& detail : requiredModelsDetails) {
+        modelsByType[detail.type].push_back(detail);
+    }
+    
+    // Populate recommended fields based on model types and architecture requirements
+    for (const auto& [type, models] : modelsByType) {
+        if (type == "VAE" && recommendedFields["recommended_vae"]) {
+            response["recommended_vae"] = modelDetailsToJson(models);
+        } else if (type == "CLIP-L" && recommendedFields["recommended_clip_l"]) {
+            response["recommended_clip_l"] = modelDetailsToJson(models);
+        } else if (type == "CLIP-G" && recommendedFields["recommended_clip_g"]) {
+            response["recommended_clip_g"] = modelDetailsToJson(models);
+        } else if (type == "T5XXL" && recommendedFields["recommended_t5xxl"]) {
+            response["recommended_t5xxl"] = modelDetailsToJson(models);
+        } else if (type == "CLIP-Vision" && recommendedFields["recommended_clip_vision"]) {
+            response["recommended_clip_vision"] = modelDetailsToJson(models);
+        } else if (type == "Qwen2VL" && recommendedFields["recommended_qwen2vl"]) {
+            response["recommended_qwen2vl"] = modelDetailsToJson(models);
+        }
+    }
+    
+    // Set non-applicable fields to null
+    for (const auto& [fieldName, shouldInclude] : recommendedFields) {
+        if (!shouldInclude || !response.contains(fieldName)) {
+            response[fieldName] = nlohmann::json(nullptr);
+        }
+    }
+}
+
 void Server::handleModelsList(const httplib::Request& req, httplib::Response& res) {
     std::string requestId = generateRequestId();
 
@@ -1060,7 +1185,7 @@ void Server::handleModelsList(const httplib::Request& req, httplib::Response& re
                 }
             }
 
-            // Build model JSON with only essential information
+            // Build model JSON with enhanced structure
             nlohmann::json modelJson = {
                 {"name", modelInfo.name},
                 {"type", ModelManager::modelTypeToString(modelInfo.type)},
@@ -1087,9 +1212,17 @@ void Server::handleModelsList(const httplib::Request& req, httplib::Response& re
                 if (!modelInfo.recommendedSampler.empty()) {
                     modelJson["recommended_sampler"] = modelInfo.recommendedSampler;
                 }
+                
+                // Enhanced model information with existence checking
                 if (!modelInfo.requiredModels.empty()) {
-                    modelJson["required_models"] = modelInfo.requiredModels;
+                    auto requiredModelsDetails = m_modelManager->checkRequiredModelsExistence(modelInfo.requiredModels);
+                    modelJson["required_models"] = modelDetailsToJson(requiredModelsDetails);
+                    
+                    // Populate recommended models based on architecture
+                    populateRecommendedModels(modelJson, modelInfo);
                 }
+                
+                // Backward compatibility - keep existing fields
                 if (!modelInfo.missingModels.empty()) {
                     modelJson["missing_models"] = modelInfo.missingModels;
                     modelJson["has_missing_dependencies"] = true;