臨

models/Model.go

@@ -40,7 +40,7 @@ func (model *Model) Train() (err error) {
 
 	// 獲取一臺空閒的訓練機
 	var server Server
-	if err = configs.ORMDB().Where("status = ?", "閒置").First(&server).Error; err != nil {
+	if err = configs.ORMDB().Where("status = ?", "正常").First(&server).Error; err != nil {
 		fmt.Println(err)
 		// TOOD: 沒有空閒的訓練機, 訓練排隊, 等待訓練機空閒
 		// TODO: 如果訓練機數量低於10臺, 則創建新的訓練機
@@ -73,6 +73,14 @@ func (model *Model) Train() (err error) {
 	for index, url := range dataset.Images {
 		fmt.Println("下載數據到本地:", index, url)
 
+		// 檢查文件是否已經存在
+		filename := fmt.Sprintf("%x", md5.Sum([]byte(url)))
+		filePath := filepath.Join(dirPath, filename)
+		if _, err := os.Stat(filePath); err == nil {
+			fmt.Println("文件已經存在:", filePath)
+			continue
+		}
+
 		// 下載到臨時目錄
 		resp, err := http.Get(url)
 		if err != nil {
@@ -87,14 +95,7 @@ func (model *Model) Train() (err error) {
 			continue
 		}
 
-		// 保存文件到本地目錄下(自動創建目錄,文件名為url的md5值)
-		filename := fmt.Sprintf("%x", md5.Sum([]byte(url)))
-		filePath := filepath.Join(dirPath, filename)
-		if err := os.MkdirAll(dirPath, 0755); err != nil {
-			fmt.Println(err)
-			continue
-		}
-
+		// 保存文件到本地目錄下
 		if err := ioutil.WriteFile(filePath, data, 0644); err != nil {
 			fmt.Println(err)
 			continue
@@ -106,16 +107,13 @@ func (model *Model) Train() (err error) {
 
 	// 檢查目錄下是否有文件, 如果沒有文件則返回錯誤
 	files, err := ioutil.ReadDir(dirPath)
-	if err != nil {
-		fmt.Println(err)
-		return err
-	}
-
-	if len(files) == 0 {
+	if err != nil || len(files) == 0 {
 		fmt.Println("目錄下沒有文件")
 		return fmt.Errorf("目錄下沒有文件")
 	}
 
+	// 將數據上傳到訓練機
+
 	// 按類型執行訓練任務
 	if model.Type == "dreambooth" {
 		// 創建數據庫模型