diff --git a/api/handler/admin/image_handler.go b/api/handler/admin/image_handler.go index 966192cd..f1737435 100644 --- a/api/handler/admin/image_handler.go +++ b/api/handler/admin/image_handler.go @@ -194,7 +194,6 @@ func (h *ImageHandler) Remove(c *gin.Context) { remark = fmt.Sprintf("任务失败,退回算力。任务ID:%d,Err: %s", job.Id, job.ErrMsg) progress = job.Progress imgURL = job.ImgURL - break case "sd": var job model.SdJob if res := h.DB.Where("id", id).First(&job); res.Error != nil { @@ -210,7 +209,6 @@ func (h *ImageHandler) Remove(c *gin.Context) { remark = fmt.Sprintf("任务失败,退回算力。任务ID:%d,Err: %s", job.Id, job.ErrMsg) progress = job.Progress imgURL = job.ImgURL - break case "dall": var job model.DallJob if res := h.DB.Where("id", id).First(&job); res.Error != nil { @@ -226,7 +224,6 @@ func (h *ImageHandler) Remove(c *gin.Context) { remark = fmt.Sprintf("任务失败,退回算力。任务ID:%d,Err: %s", job.Id, job.ErrMsg) progress = job.Progress imgURL = job.ImgURL - break default: resp.ERROR(c, types.InvalidArgs) return diff --git a/api/handler/admin/jimeng_handler.go b/api/handler/admin/jimeng_handler.go index 417e65bc..0637f69a 100644 --- a/api/handler/admin/jimeng_handler.go +++ b/api/handler/admin/jimeng_handler.go @@ -1,12 +1,15 @@ package admin import ( + "fmt" "strconv" "geekai/core" "geekai/core/types" "geekai/handler" + "geekai/service" "geekai/service/jimeng" + "geekai/service/oss" "geekai/store/model" "geekai/utils" "geekai/utils/resp" @@ -19,13 +22,17 @@ import ( type AdminJimengHandler struct { handler.BaseHandler jimengService *jimeng.Service + userService *service.UserService + uploader *oss.UploaderManager } // NewAdminJimengHandler 创建管理后台即梦AI处理器 -func NewAdminJimengHandler(app *core.AppServer, db *gorm.DB, jimengService *jimeng.Service) *AdminJimengHandler { +func NewAdminJimengHandler(app *core.AppServer, db *gorm.DB, jimengService *jimeng.Service, userService *service.UserService, uploader *oss.UploaderManager) *AdminJimengHandler { return &AdminJimengHandler{ BaseHandler: handler.BaseHandler{App: app, DB: db}, jimengService: jimengService, + userService: userService, + uploader: uploader, } } @@ -34,8 +41,7 @@ func (h *AdminJimengHandler) RegisterRoutes() { rg := h.App.Engine.Group("/api/admin/jimeng/") rg.GET("/jobs", h.Jobs) rg.GET("/jobs/:id", h.JobDetail) - rg.DELETE("/jobs/:id", h.Remove) - rg.POST("/jobs/batch-remove", h.BatchRemove) + rg.POST("/jobs/remove", h.BatchRemove) rg.GET("/stats", h.Stats) rg.GET("/config", h.GetConfig) rg.POST("/config/update", h.UpdateConfig) @@ -107,24 +113,6 @@ func (h *AdminJimengHandler) JobDetail(c *gin.Context) { resp.SUCCESS(c, job) } -// Remove 删除任务 -func (h *AdminJimengHandler) Remove(c *gin.Context) { - idStr := c.Param("id") - jobId, err := strconv.ParseUint(idStr, 10, 32) - if err != nil { - resp.ERROR(c, "参数错误") - return - } - - err = h.DB.Where("id = ?", jobId).Delete(&model.JimengJob{}).Error - if err != nil { - resp.ERROR(c, "删除任务失败") - return - } - - resp.SUCCESS(c, gin.H{}) -} - // BatchRemove 批量删除任务 func (h *AdminJimengHandler) BatchRemove(c *gin.Context) { var req struct { @@ -136,23 +124,57 @@ func (h *AdminJimengHandler) BatchRemove(c *gin.Context) { return } - result := h.DB.Where("id IN ?", req.JobIds).Delete(&model.JimengJob{}) - if result.Error != nil { - resp.ERROR(c, "批量删除失败") - return + var deletedCount int64 = 0 + for _, jobId := range req.JobIds { + var job model.JimengJob + err := h.DB.Where("id = ?", jobId).First(&job).Error + if err != nil { + continue // 跳过不存在的 + } + tx := h.DB.Begin() + if job.Status != model.JMTaskStatusSuccess && job.Power > 0 { + remark := fmt.Sprintf("任务未成功,退回算力。任务ID:%d,Err: %s", job.Id, job.ErrMsg) + err = h.userService.IncreasePower(job.UserId, job.Power, model.PowerLog{ + Type: types.PowerRefund, + Model: "jimeng", + Remark: remark, + }) + if err != nil { + tx.Rollback() + continue + } + } + err = tx.Where("id = ?", jobId).Delete(&model.JimengJob{}).Error + if err != nil { + tx.Rollback() + continue + } + tx.Commit() + deletedCount++ + if job.ImgURL != "" { + err = h.uploader.GetUploadHandler().Delete(job.ImgURL) + if err != nil { + logger.Error("remove image failed: ", err) + } + } + if job.VideoURL != "" { + err = h.uploader.GetUploadHandler().Delete(job.VideoURL) + if err != nil { + logger.Error("remove video failed: ", err) + } + } } - resp.SUCCESS(c, gin.H{ "message": "批量删除成功", - "deleted_count": result.RowsAffected, + "deleted_count": deletedCount, }) } // Stats 获取统计信息 func (h *AdminJimengHandler) Stats(c *gin.Context) { type StatResult struct { - Status string `json:"status"` - Count int64 `json:"count"` + Status model.JMTaskStatus `json:"status"` + Count int64 `json:"count"` } var stats []StatResult @@ -177,14 +199,14 @@ func (h *AdminJimengHandler) Stats(c *gin.Context) { for _, stat := range stats { result["totalTasks"] = result["totalTasks"].(int64) + stat.Count switch stat.Status { - case "completed": - result["completedTasks"] = stat.Count - case "processing": - result["processingTasks"] = stat.Count - case "failed": - result["failedTasks"] = stat.Count - case "pending": + case model.JMTaskStatusInQueue: result["pendingTasks"] = stat.Count + case model.JMTaskStatusSuccess: + result["completedTasks"] = stat.Count + case model.JMTaskStatusGenerating: + result["processingTasks"] = stat.Count + case model.JMTaskStatusFailed: + result["failedTasks"] = stat.Count } } @@ -193,33 +215,7 @@ func (h *AdminJimengHandler) Stats(c *gin.Context) { // GetConfig 获取即梦AI配置 func (h *AdminJimengHandler) GetConfig(c *gin.Context) { - var config model.Config - err := h.DB.Debug().Where("name", "jimeng").First(&config).Error - if err != nil { - // 如果配置不存在,返回默认配置 - defaultConfig := types.JimengConfig{ - AccessKey: "", - SecretKey: "", - Power: types.JimengPower{ - TextToImage: 10, - ImageToImage: 15, - ImageEdit: 20, - ImageEffects: 25, - TextToVideo: 30, - ImageToVideo: 35, - }, - } - resp.SUCCESS(c, defaultConfig) - return - } - - var jimengConfig types.JimengConfig - err = utils.JsonDecode(config.Value, &jimengConfig) - if err != nil { - resp.ERROR(c, "解析配置失败: "+err.Error()) - return - } - + jimengConfig := h.jimengService.GetConfig() resp.SUCCESS(c, jimengConfig) } diff --git a/api/handler/jimeng_handler.go b/api/handler/jimeng_handler.go index a00f54d1..cd8a5ef7 100644 --- a/api/handler/jimeng_handler.go +++ b/api/handler/jimeng_handler.go @@ -8,6 +8,8 @@ import ( "geekai/core/types" "geekai/service/jimeng" "geekai/store/model" + "geekai/store/vo" + "geekai/utils" "geekai/utils/resp" "github.com/gin-gonic/gin" @@ -33,7 +35,7 @@ func (h *JimengHandler) RegisterRoutes() { rg := h.App.Engine.Group("/api/jimeng") rg.POST("task", h.CreateTask) // 只保留统一任务接口 rg.GET("power-config", h.GetPowerConfig) // 新增算力配置接口 - rg.GET("jobs", h.Jobs) + rg.POST("jobs", h.Jobs) rg.GET("remove", h.Remove) rg.GET("retry", h.Retry) } @@ -253,28 +255,66 @@ func (h *JimengHandler) CreateTask(c *gin.Context) { // Jobs 获取任务列表 func (h *JimengHandler) Jobs(c *gin.Context) { - user, err := h.GetLoginUser(c) - if err != nil { - resp.NotAuth(c) + userId := h.GetLoginUserId(c) + + var req struct { + Page int `json:"page"` + PageSize int `json:"page_size"` + Filter string `json:"filter"` + Ids []uint `json:"ids"` + } + + if err := c.ShouldBindJSON(&req); err != nil { + resp.ERROR(c, types.InvalidArgs) return } - page := h.GetInt(c, "page", 1) - pageSize := h.GetInt(c, "page_size", 20) + var jobs []model.JimengJob + var total int64 + query := h.DB.Model(&model.JimengJob{}).Where("user_id = ?", userId) + if req.Filter == "image" { + query = query.Where("type IN (?)", []model.JMTaskType{ + model.JMTaskTypeTextToImage, + model.JMTaskTypeImageToImage, + model.JMTaskTypeImageEdit, + model.JMTaskTypeImageEffects, + }) + } else if req.Filter == "video" { + query = query.Where("type IN (?)", []model.JMTaskType{ + model.JMTaskTypeTextToVideo, + model.JMTaskTypeImageToVideo, + }) + } - jobs, total, err := h.jimengService.GetUserJobs(user.Id, page, pageSize) - if err != nil { - logger.Errorf("get user jimeng jobs failed: %v", err) - resp.ERROR(c, "获取任务列表失败") + if len(req.Ids) > 0 { + query = query.Where("id IN (?)", req.Ids) + } + + // 统计总数 + if err := query.Count(&total).Error; err != nil { + resp.ERROR(c, err.Error()) return } - resp.SUCCESS(c, gin.H{ - "jobs": jobs, - "total": total, - "page": page, - "page_size": pageSize, - }) + // 分页查询 + offset := (req.Page - 1) * req.PageSize + if err := query.Order("created_at DESC").Offset(offset).Limit(req.PageSize).Find(&jobs).Error; err != nil { + resp.ERROR(c, err.Error()) + return + } + + // 填充 VO + var jobVos []vo.JimengJob + for _, job := range jobs { + var jobVo vo.JimengJob + err := utils.CopyObject(job, &jobVo) + if err != nil { + continue + } + jobVo.CreatedAt = job.CreatedAt.Unix() + jobVos = append(jobVos, jobVo) + } + resp.SUCCESS(c, vo.NewPage(total, req.Page, req.PageSize, jobVos)) } // Remove 删除任务 @@ -355,7 +395,7 @@ func (h *JimengHandler) Retry(c *gin.Context) { } // 重新推送到队列 - task := map[string]interface{}{ + task := map[string]any{ "job_id": jobId, "type": job.Type, } @@ -393,27 +433,7 @@ func (h *JimengHandler) subUserPower(userId uint, power int, powerLog model.Powe // getPowerFromConfig 从配置中获取指定类型的算力消耗 func (h *JimengHandler) getPowerFromConfig(taskType model.JMTaskType) int { - config, err := h.jimengService.GetConfig() - if err != nil { - logger.Errorf("获取即梦AI配置失败: %v", err) - // 返回默认值 - switch taskType { - case model.JMTaskTypeTextToImage: - return 10 - case model.JMTaskTypeImageToImage: - return 15 - case model.JMTaskTypeImageEdit: - return 20 - case model.JMTaskTypeImageEffects: - return 25 - case model.JMTaskTypeTextToVideo: - return 30 - case model.JMTaskTypeImageToVideo: - return 35 - default: - return 10 - } - } + config := h.jimengService.GetConfig() switch taskType { case model.JMTaskTypeTextToImage: @@ -435,11 +455,7 @@ func (h *JimengHandler) getPowerFromConfig(taskType model.JMTaskType) int { // GetPowerConfig 获取即梦各任务类型算力消耗配置 func (h *JimengHandler) GetPowerConfig(c *gin.Context) { - config, err := h.jimengService.GetConfig() - if err != nil || config == nil { - resp.ERROR(c, "获取算力配置失败") - return - } + config := h.jimengService.GetConfig() resp.SUCCESS(c, gin.H{ "text_to_image": config.Power.TextToImage, "image_to_image": config.Power.ImageToImage, diff --git a/api/service/jimeng/service.go b/api/service/jimeng/service.go index 62e2c098..69ef0a2c 100644 --- a/api/service/jimeng/service.go +++ b/api/service/jimeng/service.go @@ -609,12 +609,15 @@ func (s *Service) GetJob(jobId uint) (*model.JimengJob, error) { return &job, nil } -// GetUserJobs 获取用户任务列表 -func (s *Service) GetUserJobs(userId uint, page, pageSize int) ([]*model.JimengJob, int64, error) { +// GetJobByPage 分页获取任务列表 +func (s *Service) GetJobByPage(userId uint, page, pageSize int) ([]*model.JimengJob, int64, error) { var jobs []*model.JimengJob var total int64 - query := s.db.Model(&model.JimengJob{}).Where("user_id = ?", userId) + query := s.db.Model(&model.JimengJob{}) + if userId > 0 { + query = query.Where("user_id = ?", userId) + } // 统计总数 if err := query.Count(&total).Error; err != nil { @@ -688,8 +691,17 @@ func (s *Service) UpdateClientConfig(accessKey, secretKey string) error { return nil } +var defaultPower = types.JimengPower{ + TextToImage: 20, + ImageToImage: 20, + ImageEdit: 20, + ImageEffects: 20, + TextToVideo: 300, + ImageToVideo: 300, +} + // GetConfig 获取即梦AI配置 -func (s *Service) GetConfig() (*types.JimengConfig, error) { +func (s *Service) GetConfig() *types.JimengConfig { var config model.Config err := s.db.Where("name", "jimeng").First(&config).Error if err != nil { @@ -697,22 +709,19 @@ func (s *Service) GetConfig() (*types.JimengConfig, error) { return &types.JimengConfig{ AccessKey: "", SecretKey: "", - Power: types.JimengPower{ - TextToImage: 10, - ImageToImage: 15, - ImageEdit: 20, - ImageEffects: 25, - TextToVideo: 30, - ImageToVideo: 35, - }, - }, nil + Power: defaultPower, + } } var jimengConfig types.JimengConfig err = utils.JsonDecode(config.Value, &jimengConfig) if err != nil { - return nil, fmt.Errorf("解析配置失败: %w", err) + return &types.JimengConfig{ + AccessKey: "", + SecretKey: "", + Power: defaultPower, + } } - return &jimengConfig, nil + return &jimengConfig } diff --git a/api/store/vo/jimeng_job.go b/api/store/vo/jimeng_job.go index 14b76817..2f1d869d 100644 --- a/api/store/vo/jimeng_job.go +++ b/api/store/vo/jimeng_job.go @@ -1,21 +1,23 @@ package vo +import "geekai/store/model" + // JimengJob 即梦AI任务VO type JimengJob struct { - Id uint `json:"id"` - UserId uint `json:"user_id"` - TaskId string `json:"task_id"` - Type string `json:"type"` - ReqKey string `json:"req_key"` - Prompt string `json:"prompt"` - TaskParams string `json:"task_params"` - ImgURL string `json:"img_url"` - VideoURL string `json:"video_url"` - RawData string `json:"raw_data"` - Progress int `json:"progress"` - Status string `json:"status"` - ErrMsg string `json:"err_msg"` - Power int `json:"power"` - CreatedAt int64 `json:"created_at"` // 时间戳 - UpdatedAt int64 `json:"updated_at"` // 时间戳 + Id uint `json:"id"` + UserId uint `json:"user_id"` + TaskId string `json:"task_id"` + Type model.JMTaskType `json:"type"` + ReqKey string `json:"req_key"` + Prompt string `json:"prompt"` + TaskParams string `json:"task_params"` + ImgURL string `json:"img_url"` + VideoURL string `json:"video_url"` + RawData string `json:"raw_data"` + Progress int `json:"progress"` + Status model.JMTaskStatus `json:"status"` + ErrMsg string `json:"err_msg"` + Power int `json:"power"` + CreatedAt int64 `json:"created_at"` // 时间戳 + UpdatedAt int64 `json:"updated_at"` // 时间戳 } diff --git a/api/utils/common.go b/api/utils/common.go index cb76c4c1..37726a12 100644 --- a/api/utils/common.go +++ b/api/utils/common.go @@ -11,9 +11,6 @@ import ( "bytes" "encoding/json" "fmt" - "github.com/lionsoul2014/ip2region/binding/golang/xdb" - "github.com/nfnt/resize" - "github.com/skip2/go-qrcode" "image" "image/color" "image/draw" @@ -22,11 +19,22 @@ import ( "reflect" "strconv" "strings" + + "github.com/lionsoul2014/ip2region/binding/golang/xdb" + "github.com/nfnt/resize" + "github.com/skip2/go-qrcode" ) // CopyObject 拷贝对象 func CopyObject(src interface{}, dst interface{}) error { + // 这里做异常处理 + defer func() { + if r := recover(); r != nil { + logger.Errorf("copy object failed: %v", r) + } + }() + srcType := reflect.TypeOf(src) srcValue := reflect.ValueOf(src) dstValue := reflect.ValueOf(dst).Elem() diff --git a/web/src/assets/css/jimeng.styl b/web/src/assets/css/jimeng.styl index 8d6b043c..45a4823b 100644 --- a/web/src/assets/css/jimeng.styl +++ b/web/src/assets/css/jimeng.styl @@ -235,8 +235,6 @@ border-radius: 12px; box-shadow: var(--card-shadow, 0 2px 8px rgba(0,0,0,0.1)); overflow: hidden; - min-height: 420px; - height: 100%; transition: box-shadow 0.2s; &:hover { box-shadow: 0 4px 24px rgba(88,101,242,0.12); diff --git a/web/src/store/jimeng.js b/web/src/store/jimeng.js index 43aaeb97..27752596 100644 --- a/web/src/store/jimeng.js +++ b/web/src/store/jimeng.js @@ -5,7 +5,6 @@ // * @Author yangjian102621@163.com // * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ -import nodata from '@/assets/img/no-data.png' import { checkSession } from '@/store/cache' import { showMessageError, showMessageOK } from '@/utils/dialog' import { httpGet, httpPost } from '@/utils/http' @@ -26,8 +25,6 @@ export const useJimengStore = defineStore('jimeng', () => { // 共同状态 const loading = ref(false) const submitting = ref(false) - const list = ref([]) - const noData = ref(true) const page = ref(1) const pageSize = ref(10) const total = ref(0) @@ -186,6 +183,8 @@ export const useJimengStore = defineStore('jimeng', () => { userPower.value = user.power // 获取任务列表 await fetchData(1) + // 开始轮询 + startPolling() } catch (error) { console.error('初始化失败:', error) } @@ -257,58 +256,40 @@ export const useJimengStore = defineStore('jimeng', () => { // 切换任务筛选 const switchTaskFilter = (filter) => { taskFilter.value = filter - updateCurrentList() - } - - // 更新当前列表 - const updateCurrentList = () => { - if (taskFilter.value === 'all') { - currentList.value = list.value - } else if (taskFilter.value === 'image') { - currentList.value = list.value.filter((item) => - ['text_to_image', 'image_to_image_portrait', 'image_edit', 'image_effects'].includes( - item.type - ) - ) - } else if (taskFilter.value === 'video') { - currentList.value = list.value.filter((item) => - ['text_to_video', 'image_to_video'].includes(item.type) - ) - } + fetchData(1) } // 轮询定时器 let pollHandler = null - // 获取任务列表 const fetchData = async (pageNum = 1) => { try { loading.value = true page.value = pageNum - const response = await httpGet('/api/jimeng/jobs', { + const response = await httpPost('/api/jimeng/jobs', { page: pageNum, page_size: pageSize.value, + filter: taskFilter.value, }) + const data = response.data + if (data.total === 0) { + isOver.value = true + currentList.value = [] + return + } - if (response.data) { - list.value = response.data.jobs || [] - total.value = response.data.total || 0 - noData.value = list.value.length === 0 - updateCurrentList() - // 判断是否有未完成任务 - const hasPending = list.value.some( - (item) => item.status === 'in_queue' || item.status === 'processing' - ) - if (hasPending) { - startPolling() - } else { - stopPolling() - } + total.value = data.total || 0 + if (data.items.length < pageSize.value) { + isOver.value = true + } + if (pageNum === 1) { + currentList.value = data.items + } else { + currentList.value = currentList.value.concat(data.items) } } catch (error) { - console.error('获取任务列表失败:', error) - showMessageError('获取任务列表失败') + showMessageError('获取任务列表失败:' + error.message) } finally { loading.value = false } @@ -317,8 +298,30 @@ export const useJimengStore = defineStore('jimeng', () => { // 简单轮询逻辑 const startPolling = () => { if (pollHandler) return - pollHandler = setInterval(() => { - fetchData(page.value) + pollHandler = setInterval(async () => { + const response = await httpPost('/api/jimeng/jobs', { + page: 1, + page_size: 20, + }) + const data = response.data + if (data.items.length === 0) { + stopPolling() + return + } + + const todoList = data.items.filter( + (item) => item.status === 'in_queue' || item.status === 'generating' + ) + // 更新当前列表 + currentList.value.forEach((item) => { + const index = data.items.findIndex((i) => i.id === item.id) + if (index !== -1) { + Object.assign(item, data.items[index]) + } + }) + if (todoList.length === 0) { + stopPolling() + } }, 5000) } @@ -533,18 +536,16 @@ export const useJimengStore = defineStore('jimeng', () => { useImageInput, loading, submitting, - list, - noData, page, pageSize, total, taskFilter, currentList, + isOver, isLogin, userPower, showDialog, currentVideoUrl, - nodata, // 配置 categories, @@ -577,7 +578,6 @@ export const useJimengStore = defineStore('jimeng', () => { getTaskStatusText, getStatusType, switchTaskFilter, - updateCurrentList, fetchData, submitTask, retryTask, diff --git a/web/src/views/Jimeng.vue b/web/src/views/Jimeng.vue index 2fc3074f..fb5668f5 100644 --- a/web/src/views/Jimeng.vue +++ b/web/src/views/Jimeng.vue @@ -276,7 +276,7 @@ -
+

你的作品

@@ -306,119 +306,134 @@
-
- - - - +
+
+
@@ -467,9 +482,8 @@ onUnmounted(() => { store.cleanup() }) -// 自动加载下一页逻辑 function onWaterfallAfterRender() { - if (!store.loading && store.currentList.length < store.total) { + if (!store.loading && !store.isOver) { store.fetchData(store.page + 1) } } diff --git a/web/src/views/admin/jimeng/JimengJobs.vue b/web/src/views/admin/jimeng/JimengJobs.vue index cf87ceec..f85840fa 100644 --- a/web/src/views/admin/jimeng/JimengJobs.vue +++ b/web/src/views/admin/jimeng/JimengJobs.vue @@ -1,10 +1,48 @@ - \ No newline at end of file +