gorm

models/Model.go

@@ -6,7 +6,7 @@ import (
 )
 
 type Model struct {
-	ID           int    `json:"id"`
+	ID           int    `json:"id" gorm:"primary_key"`
 	Name         string `json:"name"`
 	Type         string `json:"type"`                     // (lora|ckp|hyper|ti)
 	TriggerWords string `json:"trigger_words"`            // 觸發詞
@@ -21,6 +21,10 @@ type Model struct {
 	UserID       int    `json:"user_id"`
 }
 
+func init() {
+	configs.ORMDB().AutoMigrate(&Model{})
+}
+
 func (model *Model) SendToTrain() error {
 	db, err := configs.GetDB()
 	if err != nil {
@@ -43,142 +47,3 @@ func (model *Model) SendToTrain() error {
 	// 將任務放入隊列中, 將自動回調更新任務狀態
 	return nil
 }
-
-func (model *Model) Create() error {
-	db, err := configs.GetDB()
-	if err != nil {
-		log.Println(err)
-		return err
-	}
-	defer db.Close()
-	stmt, err := db.Prepare("INSERT INTO models(name, type, trigger_words, base_model, model_path, status, progress, image, tags, created_at, updated_at, user_id) values(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)")
-	if err != nil {
-		log.Println(err)
-		return err
-	}
-	defer stmt.Close()
-	result, err := stmt.Exec(model.Name, model.Type, model.TriggerWords, model.BaseModel, model.ModelPath, model.Status, model.Progress, model.Image, model.Tags, model.CreatedAt, model.UpdatedAt, model.UserID)
-	if err != nil {
-		log.Println(err)
-		return err
-	}
-	id, err := result.LastInsertId()
-	if err != nil {
-		return err
-	}
-	model.ID = int(id)
-	return nil
-}
-
-func (model *Model) Delete() error {
-	db, err := configs.GetDB()
-	if err != nil {
-		log.Println(err)
-		return err
-	}
-	defer db.Close()
-	stmt, err := db.Prepare("DELETE FROM models WHERE id = ?")
-	if err != nil {
-		log.Println(err)
-		return err
-	}
-	defer stmt.Close()
-	_, err = stmt.Exec(model.ID)
-	if err != nil {
-		log.Println(err)
-		return err
-	}
-	return nil
-}
-
-func (model *Model) Update() error {
-	db, err := configs.GetDB()
-	if err != nil {
-		log.Println(err)
-		return err
-	}
-	defer db.Close()
-	stmt, err := db.Prepare("UPDATE models SET name = ?, type = ?, updated_at = ? WHERE id = ?")
-	if err != nil {
-		log.Println(err)
-		return err
-	}
-	defer stmt.Close()
-	_, err = stmt.Exec(model.Name, model.Type, model.UpdatedAt, model.ID)
-	if err != nil {
-		log.Println(err)
-		return err
-	}
-	return nil
-}
-
-func (model *Model) Get() error {
-	db, err := configs.GetDB()
-	if err != nil {
-		log.Println(err)
-		return err
-	}
-	defer db.Close()
-	err = db.QueryRow("SELECT id,name,type,trigger_words,base_model,model_path,status,progress,image,tags,creaated_at,updated_at,user_id FROM models WHERE id = ?", model.ID).Scan(&model.ID, &model.Name, &model.Type, &model.TriggerWords, &model.BaseModel, &model.ModelPath, &model.Status, &model.Progress, &model.Image, &model.Tags, &model.CreatedAt, &model.UpdatedAt, &model.UserID)
-	if err != nil {
-		log.Println(err)
-		return err
-	}
-
-	return nil
-}
-
-func QueryModel(id int) (model Model) {
-	db, err := configs.GetDB()
-	if err != nil {
-		log.Println(err)
-		return
-	}
-	defer db.Close()
-	err = db.QueryRow("SELECT id,name,type,trigger_words,base_model,model_path,status,progress,image,tags,created_at,updated_at,user_id FROM models WHERE id = ?", id).Scan(&model.ID, &model.Name, &model.Type, &model.TriggerWords, &model.BaseModel, &model.ModelPath, &model.Status, &model.Progress, &model.Image, &model.Tags, &model.CreatedAt, &model.UpdatedAt, &model.UserID)
-	if err != nil {
-		log.Println(err)
-		return
-	}
-	return
-}
-
-func QueryModels(page int, pagesize int) (models []interface{}) {
-	db, err := configs.GetDB()
-	if err != nil {
-		log.Println(err)
-		return
-	}
-	defer db.Close()
-	rows, err := db.Query("SELECT id, name, type, trigger_words, base_model, model_path, status, progress, image, tags, created_at, updated_at, user_id FROM models LIMIT ?, ?", (page-1)*pagesize, pagesize)
-	if err != nil {
-		log.Println(err)
-		return
-	}
-	defer rows.Close()
-	for rows.Next() {
-		var model Model
-		err = rows.Scan(&model.ID, &model.Name, &model.Type, &model.TriggerWords, &model.BaseModel, &model.ModelPath, &model.Status, &model.Progress, &model.Image, &model.Tags, &model.CreatedAt, &model.UpdatedAt, &model.UserID)
-		if err != nil {
-			log.Println(err)
-			return
-		}
-		models = append(models, model)
-	}
-	return
-}
-
-func CountModels() (count int) {
-	db, err := configs.GetDB()
-	if err != nil {
-		log.Println(err)
-		return
-	}
-	defer db.Close()
-	err = db.QueryRow("SELECT COUNT(*) FROM models").Scan(&count)
-	if err != nil {
-		log.Println(err)
-		return
-	}
-	return
-}