ai/models/Model.go

package models

import (
	"log"
	"main/configs"
)

type Model struct {
	ID           int    `json:"id"`
	Name         string `json:"name"`
	Type         string `json:"type"`          // (lora|ckp|hyper|ti)
	TriggerWords string `json:"trigger_words"` // 觸發詞
	BaseModel    string `json:"base_model"`    // (SD1.5|SD2)
	ModelPath    string `json:"model_path"`    // 模型路徑
	Status       string `json:"status"`        // (initial|ready|waiting|running|success|error)
	Progress     int    `json:"progress"`      // (0-100)
	Image        string `json:"image"`         // 封面圖片實際地址
	Tags         string `json:"tags"`
	CreatedAt    string `json:"created_at"`
	UpdatedAt    string `json:"updated_at"`
	UserID       int    `json:"user_id"`
}

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, 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.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 * FROM models WHERE id = ?", model.ID).Scan(&model.ID, &model.Name, &model.Type, &model.CreatedAt, &model.UpdatedAt)
	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, 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.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, 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.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
}