diff --git a/models/server.go b/models/server.go index 61ffe53..f078978 100644 --- a/models/server.go +++ b/models/server.go @@ -93,7 +93,51 @@ func InitDefaultServer() (err error) { // 初始化服务器中的模型列表 func (server *Server) InitModels() (err error) { - // 获取服务器中的模型列表 + // 获取服务器中的模型列表(Lora) + resp_lora, err := http.Get(fmt.Sprintf("http://%s:%d/sdapi/v1/loras", server.IP, server.Port)) + if err != nil { + return fmt.Errorf("获取服务器中的模型列表失败: %v", err) + } + defer resp_lora.Body.Close() + + // 解码JSON (数组) + var data_lora []map[string]interface{} + if err := json.NewDecoder(resp_lora.Body).Decode(&data_lora); err != nil { + return err + } + + // 从数据库检查此模型hash是否存在 + for _, item := range data_lora { + // 打印为格式化的JSON + b, _ := json.MarshalIndent(item, "", " ") + fmt.Println(string(b)) + + var model Model + if err := configs.ORMDB().Where("hash = ?", item["path"].(string)).First(&model).Error; err != nil { + // 不存在则添加 + model = Model{ + Name: item["name"].(string), + Hash: item["path"].(string), + ModelCheckpoint: item["alias"].(string), + ModelPath: item["path"].(string), + ServerID: server.ID, + Type: "lora", + } + // 添加到数据库 + if err := configs.ORMDB().Create(&model).Error; err != nil { + return fmt.Errorf("添加模型到数据库失败: %v", err) + } + // 添加到模型列表 + server.Models = append(server.Models, model.ID) + } + } + + // 更新数据库 + if err := configs.ORMDB().Save(&server).Error; err != nil { + return fmt.Errorf("更新数据库失败: %v", err) + } + + // 获取服务器中的模型列表(ckp) resp, err := http.Get(fmt.Sprintf("http://%s:%d/sdapi/v1/sd-models", server.IP, server.Port)) if err != nil { return fmt.Errorf("获取服务器中的模型列表失败: %v", err) @@ -130,7 +174,6 @@ func (server *Server) InitModels() (err error) { } // 添加到模型列表 server.Models = append(server.Models, model.ID) - } }