aiforge/routers/api/v1/repo/images.go

package repo

import (
	"encoding/json"
	"net/http"
	"strconv"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/grampus"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/modelarts"
	"code.gitea.io/gitea/modules/setting"
)

type NPUImageINFO struct {
	ID    string `json:"id"`
	Value string `json:"value"`
}

func GetPublicImages(ctx *context.APIContext) {
	uid := getUID(ctx)
	opts := models.SearchImageOptions{
		IncludePublicOnly:   true,
		UID:                 uid,
		Keyword:             ctx.Query("q"),
		Topics:              ctx.Query("topic"),
		IncludeOfficialOnly: ctx.QueryBool("recommend"),
		SearchOrderBy:       "type desc, num_stars desc,id desc",
		Status:              models.IMAGE_STATUS_SUCCESS,
		CloudbrainType:      ctx.QueryInt("cloudbrainType"),
	}

	getImages(ctx, &opts)

}

func GetCustomImages(ctx *context.APIContext) {
	uid := getUID(ctx)
	opts := models.SearchImageOptions{
		UID:              uid,
		IncludeOwnerOnly: true,
		Keyword:          ctx.Query("q"),
		Topics:           ctx.Query("topic"),
		Status:           -1,
		SearchOrderBy:    "id desc",
	}
	getImages(ctx, &opts)

}
func GetStarImages(ctx *context.APIContext) {

	uid := getUID(ctx)
	opts := models.SearchImageOptions{
		UID:             uid,
		IncludeStarByMe: true,
		Keyword:         ctx.Query("q"),
		Topics:          ctx.Query("topic"),
		Status:          models.IMAGE_STATUS_SUCCESS,
		SearchOrderBy:   "id desc",
	}
	getImages(ctx, &opts)

}

func GetNpuImages(ctx *context.APIContext) {
	cloudbrainType := ctx.QueryInt("type")
	if cloudbrainType == 0 { //modelarts
		getModelArtsImages(ctx)
	} else { //c2net
		getC2netNpuImages(ctx)
	}
}

func getModelArtsImages(ctx *context.APIContext) {

	var versionInfos modelarts.VersionInfo
	_ = json.Unmarshal([]byte(setting.EngineVersions), &versionInfos)
	var npuImageInfos []NPUImageINFO
	for _, info := range versionInfos.Version {
		npuImageInfos = append(npuImageInfos, NPUImageINFO{
			ID:    strconv.Itoa(info.ID),
			Value: info.Value,
		})
	}
	ctx.JSON(http.StatusOK, npuImageInfos)

}

func getC2netNpuImages(ctx *context.APIContext) {
	images, err := grampus.GetImages(grampus.ProcessorTypeNPU)
	var npuImageInfos []NPUImageINFO
	if err != nil {
		log.Error("GetImages failed:", err.Error())
		ctx.JSON(http.StatusOK, []NPUImageINFO{})
	} else {
		for _, info := range images.Infos {
			npuImageInfos = append(npuImageInfos, NPUImageINFO{
				ID:    info.ID,
				Value: info.Name,
			})
		}
		ctx.JSON(http.StatusOK, npuImageInfos)
	}
}
func getImages(ctx *context.APIContext, opts *models.SearchImageOptions) {
	page := ctx.QueryInt("page")
	if page <= 0 {
		page = 1
	}

	pageSize := ctx.QueryInt("pageSize")
	if pageSize <= 0 {
		pageSize = 15
	}
	opts.ListOptions = models.ListOptions{
		Page:     page,
		PageSize: pageSize,
	}
	imageList, total, err := models.SearchImage(opts)
	if err != nil {
		log.Error("Can not get images:%v", err)
		ctx.JSON(http.StatusOK, models.ImagesPageResult{
			Count:  0,
			Images: []*models.Image{},
		})
	} else {
		ctx.JSON(http.StatusOK, models.ImagesPageResult{
			Count:  total,
			Images: imageList,
		})
	}
}

func getUID(ctx *context.APIContext) int64 {
	var uid int64 = -1
	if ctx.IsSigned {
		uid = ctx.User.ID
	}
	return uid
}