|
- package models
-
- import (
- "code.gitea.io/gitea/modules/log"
- "code.gitea.io/gitea/modules/timeutil"
- )
-
- const (
- HfTransferStatusWaiting = -1
- HfTransferStatusSuccess = 1
- HfTransferStatusOngoing = 0
- HfTransferStatusFailure = 2
- )
-
- type HfModelFile struct {
- ID int64 `xorm:"pk autoincr" json:"id"`
- Filename string `xorm:"varchar(400)" json:"filename"`
- Status int `xorm:"NOT NULL DEFAULT 0" json:"status"`
- Path string `xorm:"varchar(400) NULL" json:"path"`
- HfRepoId string `xorm:"varchar(400)" json:"hf_repo_id"`
- OpeniRepoId string `xorm:"varchar(400)" json:"openi_repo_id"`
- ModelId string `xorm:"varchar(400)" json:"model_id"`
- ModelName string `xorm:"varchar(400)" json:"model_name"`
- RepoId int64 `xorm:"DEFAULT 0 " json:"repo_id"`
- OwnerId int64 `xorm:"DEFAULT 0" json:"owner_id"`
- OwnerName string `xorm:"varchar(400)" json:"owner_name"`
- UserId int64 `xorm:"DEFAULT 0" json:"user_id"`
- Size int64 `xorm:"DEFAULT 0 NOT NULL" json:"size"`
- HfBlobId string `xorm:"varchar(400)" json:"hf_blob_id"`
- HfSha string `xorm:"varchar(400)" json:"hf_sha"`
- HfLicense string `xorm:"varchar(400)" json:"hf_license"`
- CreatedUnix timeutil.TimeStamp `xorm:"created" json:"created_unix"`
- UpdatedUnix timeutil.TimeStamp `xorm:"updated" json:"updated_unix"`
- }
-
- func SaveHfFile(record *HfModelFile) error {
- sess := x.NewSession()
- defer sess.Close()
-
- _, err := sess.Insert(record)
- if err != nil {
- return err
- }
- return nil
- }
-
- func QueryHfFilesByOwnerId(hfRepoId string, modelOwnerId int64) ([]*HfModelFile, error) {
- sess := x.NewSession()
- defer sess.Close()
- var files []*HfModelFile
- err := sess.Where("hf_repo_id = ? AND owner_id = ?", hfRepoId, modelOwnerId).OrderBy("filename asc").Find(&files)
- if err != nil {
- return nil, err
- }
- return files, nil
- }
-
- func UpdateHfFileStatus(fileRecordId int64, statusNew int) error {
- sess := x.NewSession()
- defer sess.Close()
- _, err := sess.Where("id = ?", fileRecordId).Cols("status").Update(&HfModelFile{Status: statusNew})
- if err != nil {
- return err
- }
- return nil
- }
-
- func IsTransferComplete(modelId string) (bool, error) {
- sess := x.NewSession()
- defer sess.Close()
-
- var files []HfModelFile
- err := sess.Where("model_id = ?", modelId).Find(&files)
- if err != nil {
- return false, err
- }
-
- for _, file := range files {
- if file.Status == HfTransferStatusOngoing || file.Status == HfTransferStatusWaiting {
- return false, nil
- }
- }
-
- return true, nil
- }
-
- func IsTransferSuccess(modelId string) (bool, error) {
- sess := x.NewSession()
- defer sess.Close()
-
- var files []HfModelFile
- err := sess.Where("model_id = ?", modelId).Find(&files)
- if err != nil {
- return false, err
- }
-
- var totalCount, totalStatus int
- for _, file := range files {
- totalCount++
- totalStatus += file.Status
- }
-
- return totalCount == totalStatus, nil
- }
-
- func DeleteHfFilesByModelId(id string) error {
- sess := x.NewSession()
- defer sess.Close()
- _, err := sess.Delete(&HfModelFile{
- ModelId: id,
- })
- if err != nil {
- return err
- }
- return nil
- }
-
- func GetUserIdByModelId(modelId string) (int64, error) {
- sess := x.NewSession()
- defer sess.Close()
- var file HfModelFile
- has, err := sess.Where("model_id = ?", modelId).Get(&file)
- if err != nil {
- return 0, err
- }
- if !has {
- return 0, nil
- }
- return file.UserId, nil
- }
-
- func GetHfRepoIdByModelId(modelId string) (string, error) {
- sess := x.NewSession()
- defer sess.Close()
- var file HfModelFile
- has, err := sess.Where("model_id = ?", modelId).Get(&file)
- if err != nil {
- return "", err
- }
- if !has {
- return "", nil
- }
- return file.HfRepoId, nil
- }
-
- func GetHfModelSizeByModelId(modelId string) (int64, error) {
- sess := x.NewSession()
- defer sess.Close()
- var files []HfModelFile
- err := sess.Where("model_id = ?", modelId).Find(&files)
- if err != nil {
- return 0, err
- }
- var totalSize int64
- for _, file := range files {
- totalSize += file.Size
- }
- return totalSize, nil
- }
-
- func GetOngoingFiles(ownerId int64) ([]*HfModelFile, error) {
- sess := x.NewSession()
- defer sess.Close()
- var files []*HfModelFile
- err := sess.Where("status = ? AND owner_id = ?", HfTransferStatusOngoing, ownerId).OrderBy("updated_unix asc").Find(&files)
- if err != nil {
- log.Error("[hf_model] GetOngoingFiles: %v", err)
- return nil, err
- }
- return files, nil
- }
-
- func GetWaitingFiles(ownerId int64, numFiles int) ([]*HfModelFile, error) {
- sess := x.NewSession()
- defer sess.Close()
- var files []*HfModelFile
- err := sess.Where("status = ? AND owner_id = ?", HfTransferStatusWaiting, ownerId).OrderBy("updated_unix, size").Limit(numFiles).Find(&files)
- if err != nil {
- return nil, err
- }
- return files, nil
- }
|