package handler import ( "fmt" "net/http" "os" "path/filepath" "sync" "time" "gen2d/internal/model" "gen2d/internal/service" "github.com/gin-gonic/gin" ) // 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"` Assets []AssetResponse `json:"assets"` Metadata service.AssetMetadata `json:"metadata"` } // AssetResponse 单个素材响应。 type AssetResponse struct { URL string `json:"url"` Format string `json:"format"` } // TaskResponse 任务查询响应。 type TaskResponse struct { ID string `json:"id"` ProjectID string `json:"projectId"` Prompt string `json:"prompt"` AssetType string `json:"assetType"` Status string `json:"status"` Progress int `json:"progress"` RetryCount int `json:"retryCount"` Error string `json:"error,omitempty"` CreatedAt string `json:"createdAt"` } // AssetsResponse 素材列表响应。 type AssetsResponse struct { Assets []AssetResponse `json:"assets"` Metadata service.AssetMetadata `json:"metadata"` } // taskRecord 内存中的任务记录。 type taskRecord struct { task TaskResponse assets []AssetResponse metadata service.AssetMetadata } var ( taskStore = sync.Map{} // taskID → *taskRecord ) // Generate 素材生成接口。 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()) 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, }, } output, err := service.RunPipeline(c.Request.Context(), in) if err != nil { c.JSON(http.StatusInternalServerError, model.Fail(http.StatusInternalServerError, "素材生成失败: "+err.Error())) return } // 保存图片到 ../generation/{projectId}/{taskId}/ genDir := filepath.Join("..", "generation", projectID, taskID) if err := os.MkdirAll(genDir, 0755); err != nil { c.JSON(http.StatusInternalServerError, model.Fail(http.StatusInternalServerError, "创建输出目录失败: "+err.Error())) return } assets := make([]AssetResponse, len(output.Assets)) for i, a := range output.Assets { filename := fmt.Sprintf("%d.%s", i, a.Format) filePath := filepath.Join(genDir, filename) if err := os.WriteFile(filePath, a.Data, 0644); err != nil { c.JSON(http.StatusInternalServerError, model.Fail(http.StatusInternalServerError, "保存图片失败: "+err.Error())) return } assets[i] = AssetResponse{ URL: fmt.Sprintf("/generation/%s/%s/%s", projectID, taskID, filename), Format: a.Format, } } // 存储任务记录到内存 taskStore.Store(taskID, &taskRecord{ task: TaskResponse{ ID: taskID, ProjectID: projectID, Prompt: req.Prompt, AssetType: req.AssetType, Status: "completed", Progress: 100, CreatedAt: time.Now().Format(time.RFC3339), }, assets: assets, metadata: output.Metadata, }) c.JSON(http.StatusOK, model.OK(GenerateResponse{ TaskID: taskID, Assets: assets, Metadata: output.Metadata, })) } // GetTask 查询任务信息。 func GetTask(c *gin.Context) { taskID := c.Param("taskId") rec, ok := taskStore.Load(taskID) if !ok { c.JSON(http.StatusNotFound, model.Fail(http.StatusNotFound, "任务不存在")) return } r := rec.(*taskRecord) c.JSON(http.StatusOK, model.OK(r.task)) } // GetAssets 查询任务素材列表。 func GetAssets(c *gin.Context) { taskID := c.Param("taskId") rec, ok := taskStore.Load(taskID) if !ok { c.JSON(http.StatusNotFound, model.Fail(http.StatusNotFound, "任务不存在")) return } r := rec.(*taskRecord) c.JSON(http.StatusOK, model.OK(AssetsResponse{ Assets: r.assets, Metadata: r.metadata, })) }