Gmarker689
96285014b1
- 后端 GetRecentAssets 增加 projectName/width/height/projectId 字段 - 前端 AssetGallery 点击图片弹出详情弹窗：提示词/工程/分辨率/格式 - 详情弹窗提供下载按钮，走 /api/v1/assets/download?key= 接口
397 lines
11 KiB
Go
Executable File
397 lines
11 KiB
Go
Executable File
package handler
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"gen2d/internal/db"
|
|
"gen2d/internal/logger"
|
|
"gen2d/internal/model"
|
|
"gen2d/internal/service"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
// generateQueue 全局生成任务队列，由 main 通过 SetGenerateQueue 注入。
|
|
var generateQueue *service.TaskQueue
|
|
|
|
// SetGenerateQueue 设置生成任务队列。
|
|
func SetGenerateQueue(q *service.TaskQueue) {
|
|
generateQueue = q
|
|
}
|
|
|
|
// GenerateRequest 素材生成请求。
|
|
type GenerateRequest struct {
|
|
ProjectID string `json:"projectId"`
|
|
Prompt string `json:"prompt"`
|
|
AssetType string `json:"assetType" binding:"required"`
|
|
Tags []string `json:"tags"`
|
|
UserNote string `json:"userNote"`
|
|
ProjectStyle map[string]string `json:"projectStyle"`
|
|
TaskStyle map[string]string `json:"taskStyle"`
|
|
Resolution int `json:"resolution"`
|
|
Directions int `json:"directions"`
|
|
FramesPerDir int `json:"framesPerDir"`
|
|
Format string `json:"format"`
|
|
}
|
|
|
|
// GenerateResponse 素材生成响应体。
|
|
type GenerateResponse struct {
|
|
TaskID string `json:"taskId"`
|
|
}
|
|
|
|
// AssetsResponse 素材列表响应。
|
|
type AssetsResponse struct {
|
|
Assets []model.AssetResponse `json:"assets"`
|
|
Metadata service.AssetMetadata `json:"metadata"`
|
|
}
|
|
|
|
// Generate 素材生成接口（异步）。
|
|
// 立即返回 taskId，任务进入 FIFO 队列串行执行，前端通过 GET /tasks/:taskId 轮询进度。
|
|
func Generate(c *gin.Context) {
|
|
var req GenerateRequest
|
|
if err := c.ShouldBindJSON(&req); err != nil {
|
|
c.JSON(http.StatusBadRequest, model.Fail(http.StatusBadRequest, "参数错误: "+err.Error()))
|
|
return
|
|
}
|
|
|
|
projectID := req.ProjectID
|
|
if projectID == "" {
|
|
projectID = "default"
|
|
}
|
|
taskID := fmt.Sprintf("task-%d", time.Now().UnixMilli())
|
|
|
|
// 保存任务到数据库，初始状态为 pending
|
|
if err := saveTaskToDB(c.Request.Context(), projectID, taskID, req); err != nil {
|
|
logger.FromCtx(c.Request.Context()).Error("failed to save task", "error", err)
|
|
c.JSON(http.StatusInternalServerError, model.Fail(http.StatusInternalServerError, "创建任务失败"))
|
|
return
|
|
}
|
|
|
|
// 加入 FIFO 队列
|
|
queuePos := 1
|
|
if generateQueue != nil {
|
|
queuePos = generateQueue.Enqueue(&service.TaskJob{
|
|
Ctx: context.Background(),
|
|
ProjectID: projectID,
|
|
TaskID: taskID,
|
|
Execute: func(ctx context.Context) error {
|
|
runPipelineBg(ctx, projectID, taskID, req)
|
|
return nil
|
|
},
|
|
})
|
|
}
|
|
|
|
c.JSON(http.StatusOK, model.OK(GenerateResponse{TaskID: taskID}))
|
|
_ = queuePos
|
|
}
|
|
|
|
// runPipelineBg 后台执行生成管线，更新任务状态。
|
|
func runPipelineBg(ctx context.Context, projectID, taskID string, req GenerateRequest) {
|
|
l := logger.With("task_id", taskID, "project_id", projectID)
|
|
l.Info("task started", "asset_type", req.AssetType, "prompt", req.Prompt)
|
|
|
|
// 注入进度上报回调
|
|
ctx = service.WithProgressReporter(ctx, func(stage string, progress int) {
|
|
l.Info("progress update", "stage", stage, "progress", progress)
|
|
updateTaskInDB(ctx, taskID, "running", stage, "", progress)
|
|
})
|
|
|
|
// 队列调度后才标记为 running，初始写入时是 pending
|
|
updateTaskInDB(ctx, taskID, "running", "prompt_builder", "", 5)
|
|
|
|
in := service.PipelineInput{
|
|
ProjectID: projectID,
|
|
TaskID: taskID,
|
|
Prompt: req.Prompt,
|
|
AssetType: req.AssetType,
|
|
Tags: req.Tags,
|
|
UserNote: req.UserNote,
|
|
ProjectStyle: req.ProjectStyle,
|
|
TaskStyle: req.TaskStyle,
|
|
Params: service.AssetParams{
|
|
Resolution: req.Resolution,
|
|
Frames: service.FrameParams{
|
|
Directions: req.Directions,
|
|
FramesPerDirection: req.FramesPerDir,
|
|
},
|
|
Format: req.Format,
|
|
},
|
|
}
|
|
|
|
l.Info("calling pipeline", "input_tags", in.Tags, "user_note", in.UserNote)
|
|
output, err := service.RunPipeline(ctx, in)
|
|
if err != nil {
|
|
l.Error("task pipeline failed", "error", err)
|
|
updateTaskInDB(ctx, taskID, "failed", "", err.Error(), 0)
|
|
return
|
|
}
|
|
l.Info("pipeline completed", "asset_count", len(output.Assets))
|
|
|
|
updateTaskInDB(ctx, taskID, "saving", "format_adapter", "", 90)
|
|
|
|
// 上传素材并保存到数据库
|
|
var lastGIFCDNURL string
|
|
for i, a := range output.Assets {
|
|
key := fmt.Sprintf("generation/%s/%s/%d.%s", projectID, taskID, i, a.Format)
|
|
cdnURL, err := storageSvc.Upload(ctx, key, a.Data)
|
|
if err != nil {
|
|
l.Error("upload asset failed", "index", i, "error", err)
|
|
updateTaskInDB(ctx, taskID, "failed", "", "上传素材失败: "+err.Error(), 0)
|
|
return
|
|
}
|
|
|
|
var assetMeta map[string]interface{}
|
|
if a.Format == "gif" {
|
|
assetMeta = map[string]interface{}{"index": i, "type": "preview"}
|
|
lastGIFCDNURL = cdnURL
|
|
} else if strings.Contains(a.URL, "spritesheet") {
|
|
assetMeta = map[string]interface{}{"index": i, "type": "spritesheet"}
|
|
} else {
|
|
assetMeta = map[string]interface{}{"index": i, "type": "frame"}
|
|
}
|
|
metadataJSON, _ := json.Marshal(assetMeta)
|
|
|
|
asset := &model.Asset{
|
|
TaskID: getTaskDBID(ctx, taskID),
|
|
Key: key,
|
|
URL: cdnURL,
|
|
Format: a.Format,
|
|
Metadata: string(metadataJSON),
|
|
}
|
|
|
|
if err := db.GetDB().WithContext(ctx).Create(asset).Error; err != nil {
|
|
l.Error("save asset to db failed", "index", i, "error", err)
|
|
}
|
|
}
|
|
|
|
// GIF URL 替换为实际上传后的 CDN 地址
|
|
if output.Metadata.GIFURL != "" && lastGIFCDNURL != "" {
|
|
output.Metadata.GIFURL = lastGIFCDNURL
|
|
}
|
|
|
|
var fullMetadata string
|
|
if metadataJSON, err := json.Marshal(output.Metadata); err == nil {
|
|
fullMetadata = string(metadataJSON)
|
|
}
|
|
|
|
db.GetDB().WithContext(ctx).Model(&model.Task{}).
|
|
Where("external_id = ?", taskID).
|
|
Updates(map[string]interface{}{
|
|
"status": "completed",
|
|
"progress": 100,
|
|
"metadata": fullMetadata,
|
|
})
|
|
|
|
l.Info("task completed")
|
|
}
|
|
|
|
// GetTask 查询任务信息。
|
|
func GetTask(c *gin.Context) {
|
|
taskID := c.Param("taskId")
|
|
|
|
var task model.Task
|
|
if err := db.GetDB().WithContext(c.Request.Context()).
|
|
Where("external_id = ?", taskID).
|
|
First(&task).Error; err != nil {
|
|
c.JSON(http.StatusNotFound, model.Fail(http.StatusNotFound, "任务不存在"))
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, model.OK(toTaskResponse(&task)))
|
|
}
|
|
|
|
// GetAssets 查询任务素材列表。
|
|
func GetAssets(c *gin.Context) {
|
|
taskID := c.Param("taskId")
|
|
|
|
var task model.Task
|
|
if err := db.GetDB().WithContext(c.Request.Context()).
|
|
Where("external_id = ?", taskID).
|
|
First(&task).Error; err != nil {
|
|
c.JSON(http.StatusNotFound, model.Fail(http.StatusNotFound, "任务不存在"))
|
|
return
|
|
}
|
|
|
|
if task.Status != "completed" {
|
|
c.JSON(http.StatusBadRequest, model.Fail(http.StatusBadRequest, "任务尚未完成，当前状态: "+task.Status))
|
|
return
|
|
}
|
|
|
|
var assets []model.Asset
|
|
if err := db.GetDB().WithContext(c.Request.Context()).
|
|
Where("task_id = ?", task.ID).
|
|
Find(&assets).Error; err != nil {
|
|
c.JSON(http.StatusInternalServerError, model.Fail(http.StatusInternalServerError, "查询素材失败"))
|
|
return
|
|
}
|
|
|
|
response := []model.AssetResponse{}
|
|
for _, a := range assets {
|
|
response = append(response, model.AssetResponse{
|
|
Key: a.Key,
|
|
URL: storageSvc.GetSignedURL(a.Key),
|
|
Format: a.Format,
|
|
Metadata: a.Metadata,
|
|
})
|
|
}
|
|
|
|
var metadata service.AssetMetadata
|
|
if task.Metadata != "" {
|
|
json.Unmarshal([]byte(task.Metadata), &metadata)
|
|
}
|
|
|
|
c.JSON(http.StatusOK, model.OK(AssetsResponse{
|
|
Assets: response,
|
|
Metadata: metadata,
|
|
}))
|
|
}
|
|
|
|
// saveTaskToDB 将任务记录保存到数据库。
|
|
func saveTaskToDB(ctx context.Context, projectID, taskID string, req GenerateRequest) error {
|
|
// "default" 工程没有数据库记录，跳过
|
|
if projectID == "default" {
|
|
return nil
|
|
}
|
|
|
|
projectIDUint, err := strconv.ParseUint(projectID, 10, 32)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to parse projectID: %w", err)
|
|
}
|
|
|
|
task := &model.Task{
|
|
ExternalID: taskID,
|
|
ProjectID: uint(projectIDUint),
|
|
Prompt: req.Prompt,
|
|
AssetType: req.AssetType,
|
|
Status: "pending",
|
|
Progress: 0,
|
|
Stage: "",
|
|
RetryCount: 0,
|
|
CreatedAt: time.Now(),
|
|
UpdatedAt: time.Now(),
|
|
}
|
|
|
|
return db.GetDB().WithContext(ctx).Create(task).Error
|
|
}
|
|
|
|
// updateTaskInDB 更新数据库中的任务状态。
|
|
func updateTaskInDB(ctx context.Context, taskID, status, stage, error string, progress int) {
|
|
updates := map[string]interface{}{
|
|
"status": status,
|
|
"progress": progress,
|
|
"updated_at": time.Now(),
|
|
}
|
|
if stage != "" {
|
|
updates["stage"] = stage
|
|
}
|
|
if error != "" {
|
|
updates["error"] = error
|
|
}
|
|
|
|
db.GetDB().WithContext(ctx).Model(&model.Task{}).
|
|
Where("external_id = ?", taskID).
|
|
Updates(updates)
|
|
}
|
|
|
|
// getTaskDBID 根据外部taskID获取数据库中的任务ID。
|
|
func getTaskDBID(ctx context.Context, taskID string) uint {
|
|
var task model.Task
|
|
db.GetDB().WithContext(ctx).Select("id").Where("external_id = ?", taskID).First(&task)
|
|
return task.ID
|
|
}
|
|
|
|
// RecentAssetItem 首页素材展示项。
|
|
type RecentAssetItem struct {
|
|
Key string `json:"key"`
|
|
URL string `json:"url"`
|
|
Format string `json:"format"`
|
|
Prompt string `json:"prompt"`
|
|
AssetType string `json:"assetType"`
|
|
MetaType string `json:"metaType"`
|
|
TaskID string `json:"taskId"`
|
|
ProjectID string `json:"projectId"`
|
|
ProjectName string `json:"projectName"`
|
|
Width int `json:"width"`
|
|
Height int `json:"height"`
|
|
CreatedAt string `json:"createdAt"`
|
|
}
|
|
|
|
// GetRecentAssets 获取最近生成的素材列表（已完成任务的全部素材）。
|
|
func GetRecentAssets(c *gin.Context) {
|
|
limit := 60
|
|
|
|
var dbRows []struct {
|
|
Key string
|
|
URL string
|
|
Format string
|
|
Prompt string
|
|
AssetType string
|
|
MetaType string
|
|
TaskID string
|
|
ProjectID uint
|
|
ProjectName string
|
|
Width int
|
|
Height int
|
|
CreatedAt string
|
|
}
|
|
|
|
db.GetDB().WithContext(c.Request.Context()).
|
|
Raw(`SELECT a.key, a.url, a.format, t.prompt, t.asset_type as asset_type,
|
|
COALESCE(json_extract(a.metadata, '$.type'), 'frame') as meta_type,
|
|
t.external_id as task_id, t.project_id,
|
|
p.name as project_name,
|
|
CAST(COALESCE(json_extract(t.metadata, '$.frameWidth'), '0') AS INTEGER) as width,
|
|
CAST(COALESCE(json_extract(t.metadata, '$.frameHeight'), '0') AS INTEGER) as height,
|
|
t.created_at as created_at
|
|
FROM assets a
|
|
INNER JOIN tasks t ON a.task_id = t.id
|
|
INNER JOIN projects p ON t.project_id = p.id
|
|
WHERE t.status = 'completed'
|
|
ORDER BY a.created_at DESC
|
|
LIMIT ?`, limit).
|
|
Scan(&dbRows)
|
|
|
|
result := make([]RecentAssetItem, len(dbRows))
|
|
for i, r := range dbRows {
|
|
result[i] = RecentAssetItem{
|
|
Key: r.Key,
|
|
URL: storageSvc.GetSignedURL(r.Key),
|
|
Format: r.Format,
|
|
Prompt: r.Prompt,
|
|
AssetType: r.AssetType,
|
|
MetaType: r.MetaType,
|
|
TaskID: r.TaskID,
|
|
ProjectID: fmt.Sprintf("%d", r.ProjectID),
|
|
ProjectName: r.ProjectName,
|
|
Width: r.Width,
|
|
Height: r.Height,
|
|
CreatedAt: r.CreatedAt,
|
|
}
|
|
}
|
|
|
|
c.JSON(http.StatusOK, model.OK(result))
|
|
}
|
|
|
|
// toTaskResponse 转换任务响应格式。
|
|
func toTaskResponse(task *model.Task) model.TaskResponse {
|
|
return model.TaskResponse{
|
|
ID: task.ExternalID,
|
|
ProjectID: fmt.Sprintf("%d", task.ProjectID),
|
|
Prompt: task.Prompt,
|
|
AssetType: task.AssetType,
|
|
Status: task.Status,
|
|
Stage: task.Stage,
|
|
Progress: task.Progress,
|
|
RetryCount: task.RetryCount,
|
|
Error: task.Error,
|
|
CreatedAt: task.CreatedAt.Format("2006-01-02T15:04:05Z07:00"),
|
|
UpdatedAt: task.UpdatedAt.Format("2006-01-02T15:04:05Z07:00"),
|
|
}
|
|
}
|