324 lines
12 KiB
Go
324 lines
12 KiB
Go
// 图片生成生图业务模块，负责提交用户级图片任务、查询历史并清理生成资源。
|
||
package productimage
|
||
|
||
import (
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"strings"
|
||
"time"
|
||
"unicode/utf8"
|
||
|
||
"juhe-factory/api/internal/billing"
|
||
"juhe-factory/api/internal/model"
|
||
queuepkg "juhe-factory/api/internal/queue"
|
||
|
||
"github.com/google/uuid"
|
||
"github.com/hibiken/asynq"
|
||
"gorm.io/gorm"
|
||
"gorm.io/gorm/clause"
|
||
)
|
||
|
||
var activeTaskStatuses = []string{"pending_submission", "submitting", "submitted", "processing", "result_ready", "downloading", "cancel_requested"}
|
||
|
||
// ReferenceUpload 表示已经上传到对象存储、等待绑定到生图任务的参考图。
|
||
type ReferenceUpload struct {
|
||
Asset *model.MediaAsset
|
||
Name string
|
||
}
|
||
|
||
// GenerateInput 是图片生成生图接口的配置参数和临时参考图。
|
||
type GenerateInput struct {
|
||
Prompt string
|
||
ModelID uuid.UUID
|
||
AspectRatio string
|
||
Resolution string
|
||
References []ReferenceUpload
|
||
}
|
||
|
||
// GenerationView 汇总任务状态和已经落库的图片结果。
|
||
type GenerationView struct {
|
||
TaskID uuid.UUID `json:"task_id"`
|
||
Status string `json:"status"`
|
||
Prompt string `json:"prompt"`
|
||
ModelName string `json:"model_name"`
|
||
ErrorMessage string `json:"error_message"`
|
||
MediaAssetID *uuid.UUID `json:"media_asset_id"`
|
||
PublicURL string `json:"public_url"`
|
||
MimeType string `json:"mime_type"`
|
||
EstimatedPoints string `json:"estimated_points"`
|
||
ActualPoints *string `json:"actual_points"`
|
||
CreatedAt time.Time `json:"created_at"`
|
||
FinishedAt *time.Time `json:"finished_at"`
|
||
}
|
||
|
||
// Service 组合图片生成模块需要的数据库和任务队列依赖。
|
||
type Service struct {
|
||
DB *gorm.DB
|
||
Queue *asynq.Client
|
||
}
|
||
|
||
// NewService 创建图片生成生图服务。
|
||
func NewService(db *gorm.DB, queue *asynq.Client) *Service { return &Service{DB: db, Queue: queue} }
|
||
|
||
// QueueGeneration 校验配置、保存临时参考图并提交一个独立图片任务。
|
||
func (s *Service) QueueGeneration(userID uuid.UUID, input GenerateInput) (*model.GenerationTask, error) {
|
||
if s.Queue == nil {
|
||
return nil, errors.New("生成任务队列不可用")
|
||
}
|
||
prompt := strings.TrimSpace(input.Prompt)
|
||
if prompt == "" {
|
||
return nil, errors.New("提示词不能为空")
|
||
}
|
||
if utf8.RuneCountInString(prompt) > 5000 {
|
||
return nil, errors.New("提示词不能超过 5000 个字符")
|
||
}
|
||
if len(input.References) > 4 {
|
||
return nil, errors.New("最多上传 4 张参考图")
|
||
}
|
||
if input.ModelID == uuid.Nil {
|
||
return nil, errors.New("请选择图片模型")
|
||
}
|
||
if !validRatio(input.AspectRatio) || !validResolution(input.Resolution) {
|
||
return nil, errors.New("图片生成参数无效")
|
||
}
|
||
var selected struct {
|
||
ModelID uuid.UUID
|
||
ChannelID uuid.UUID
|
||
ModelName string
|
||
Price string
|
||
PriceExists bool
|
||
}
|
||
if err := s.DB.Raw(`SELECT model.id AS model_id,model.channel_id,model.name AS model_name,
|
||
coalesce(price.price,0)::text AS price,(price.id IS NOT NULL) AS price_exists
|
||
FROM models model JOIN channels channel ON channel.id=model.channel_id
|
||
LEFT JOIN model_prices price ON price.model_id=model.id AND lower(price.price_key)=lower(?)
|
||
WHERE model.id=? AND model.model_type='image' AND model.enabled=true AND model.deleted_at IS NULL
|
||
AND channel.enabled=true AND channel.deleted_at IS NULL`, input.Resolution, input.ModelID).Scan(&selected).Error; err != nil {
|
||
return nil, err
|
||
}
|
||
if selected.ModelID == uuid.Nil {
|
||
return nil, errors.New("图片模型不可用")
|
||
}
|
||
if !selected.PriceExists {
|
||
return nil, fmt.Errorf("当前图片模型未配置 %s 价格", input.Resolution)
|
||
}
|
||
imageURLs := make([]string, 0, len(input.References))
|
||
referenceIDs := make([]string, 0, len(input.References))
|
||
referenceNames := make([]string, 0, len(input.References))
|
||
seenNames := make(map[string]struct{}, len(input.References))
|
||
for _, reference := range input.References {
|
||
if reference.Asset == nil || strings.TrimSpace(reference.Asset.PublicURL) == "" {
|
||
return nil, errors.New("参考图信息无效")
|
||
}
|
||
name := strings.TrimSpace(reference.Name)
|
||
if name == "" {
|
||
return nil, errors.New("参考图名称不能为空")
|
||
}
|
||
if _, exists := seenNames[name]; exists {
|
||
return nil, errors.New("参考图名称不能重复")
|
||
}
|
||
seenNames[name] = struct{}{}
|
||
imageURLs = append(imageURLs, reference.Asset.PublicURL)
|
||
referenceIDs = append(referenceIDs, reference.Asset.ID.String())
|
||
referenceNames = append(referenceNames, name)
|
||
}
|
||
providerPrompt := buildProviderPrompt(prompt, referenceNames)
|
||
payload, err := json.Marshal(map[string]any{
|
||
"product_image": true, "prompt": providerPrompt, "display_prompt": prompt,
|
||
"model": selected.ModelName, "size": input.AspectRatio, "resolution": input.Resolution,
|
||
"image_urls": imageURLs, "reference_media_asset_ids": referenceIDs, "reference_names": referenceNames, "n": 1,
|
||
})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
task := &model.GenerationTask{ID: uuid.New(), RequestID: "product_image_" + uuid.NewString(), UserID: userID, ChannelID: &selected.ChannelID, ModelID: &selected.ModelID, TaskType: "image_generation", Status: "pending_submission", InputData: payload}
|
||
err = s.DB.Transaction(func(tx *gorm.DB) error {
|
||
for _, reference := range input.References {
|
||
owner := userID
|
||
reference.Asset.OwnerUserID = &owner
|
||
if err := tx.Create(reference.Asset).Error; err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return billing.PrechargeGenerationTask(tx, task, selected.Price, 1, "图片生成")
|
||
})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if err := queuepkg.EnqueueID(s.Queue, queuepkg.TypeDispatchChannel, selected.ChannelID, 0); err != nil {
|
||
if refundErr := s.failQueuedTask(task.ID, err.Error()); refundErr != nil {
|
||
return nil, fmt.Errorf("生成任务入队失败且预扣返还失败: %w", refundErr)
|
||
}
|
||
return nil, errors.New("生成任务队列暂时不可用，请稍后重试")
|
||
}
|
||
return task, nil
|
||
}
|
||
|
||
// Generations 返回当前用户最近五十条图片生成记录。
|
||
func (s *Service) Generations(userID uuid.UUID) ([]GenerationView, error) {
|
||
items := make([]GenerationView, 0)
|
||
err := s.DB.Raw(`SELECT task.id AS task_id,task.status,
|
||
coalesce(task.input_data->>'display_prompt',task.input_data->>'prompt','') AS prompt,
|
||
coalesce(model.name,task.input_data->>'model','') AS model_name,task.error_message,
|
||
output.media_asset_id,coalesce(media.public_url,'') AS public_url,coalesce(media.mime_type,'') AS mime_type,
|
||
task.estimated_points,task.actual_points,task.created_at,task.finished_at
|
||
FROM generation_tasks task
|
||
LEFT JOIN models model ON model.id=task.model_id
|
||
LEFT JOIN generation_outputs output ON output.task_id=task.id AND output.sequence_no=1
|
||
LEFT JOIN media_assets media ON media.id=output.media_asset_id AND media.deleted_at IS NULL
|
||
WHERE task.user_id=? AND task.task_type='image_generation' AND task.input_data->>'product_image'='true'
|
||
ORDER BY task.created_at DESC LIMIT 50`, userID).Scan(&items).Error
|
||
return items, err
|
||
}
|
||
|
||
// DeleteGeneration 清理当前用户的图片任务、生成媒体和参考媒体；活动任务也允许主动删除。
|
||
func (s *Service) DeleteGeneration(userID, taskID uuid.UUID, deleteObjects func([]string) error) error {
|
||
objectKeys := make([]string, 0)
|
||
var channelID *uuid.UUID
|
||
err := s.DB.Transaction(func(tx *gorm.DB) error {
|
||
var task model.GenerationTask
|
||
if err := tx.Where("id=? AND user_id=? AND task_type='image_generation' AND input_data->>'product_image'='true'", taskID, userID).Take(&task).Error; err != nil {
|
||
return err
|
||
}
|
||
if !canDeleteGeneration(task.Status, task.ErrorMessage) {
|
||
return errors.New("图片仍在正常生成中，暂时不能删除")
|
||
}
|
||
channelID = task.ChannelID
|
||
var input struct {
|
||
ReferenceMediaAssetIDs []uuid.UUID `json:"reference_media_asset_ids"`
|
||
}
|
||
_ = json.Unmarshal(task.InputData, &input)
|
||
mediaIDs := append([]uuid.UUID(nil), input.ReferenceMediaAssetIDs...)
|
||
var outputs []struct {
|
||
MediaAssetID uuid.UUID
|
||
ObjectKey string
|
||
}
|
||
if err := tx.Table("generation_outputs output").Select("output.media_asset_id,media.object_key").Joins("JOIN media_assets media ON media.id=output.media_asset_id").Where("output.task_id=?", taskID).Scan(&outputs).Error; err != nil {
|
||
return err
|
||
}
|
||
for _, output := range outputs {
|
||
mediaIDs = append(mediaIDs, output.MediaAssetID)
|
||
if output.ObjectKey != "" {
|
||
objectKeys = append(objectKeys, output.ObjectKey)
|
||
}
|
||
}
|
||
if len(mediaIDs) > 0 {
|
||
var refs []struct{ ObjectKey string }
|
||
if err := tx.Table("media_assets").Select("object_key").Where("id IN ?", mediaIDs).Scan(&refs).Error; err != nil {
|
||
return err
|
||
}
|
||
for _, ref := range refs {
|
||
if ref.ObjectKey != "" {
|
||
objectKeys = append(objectKeys, ref.ObjectKey)
|
||
}
|
||
}
|
||
}
|
||
if deleteObjects != nil {
|
||
if err := deleteObjects(uniqueStrings(objectKeys)); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
if err := tx.Exec("DELETE FROM generation_outputs WHERE task_id=?", taskID).Error; err != nil {
|
||
return err
|
||
}
|
||
if len(mediaIDs) > 0 {
|
||
if err := tx.Exec("DELETE FROM channel_asset_cache WHERE media_asset_id IN ?", mediaIDs).Error; err != nil {
|
||
return err
|
||
}
|
||
if err := tx.Exec("DELETE FROM media_assets WHERE id IN ?", mediaIDs).Error; err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return tx.Exec("DELETE FROM generation_tasks WHERE id=?", taskID).Error
|
||
})
|
||
if err == nil && channelID != nil && s.Queue != nil {
|
||
_ = queuepkg.EnqueueID(s.Queue, queuepkg.TypeDispatchChannel, *channelID, 0)
|
||
}
|
||
return err
|
||
}
|
||
|
||
// canDeleteGeneration 判断图片任务是否已结束，或因错误停留在活动状态而允许用户清理。
|
||
func canDeleteGeneration(status, errorMessage string) bool {
|
||
for _, activeStatus := range activeTaskStatuses {
|
||
if status == activeStatus {
|
||
return strings.TrimSpace(errorMessage) != ""
|
||
}
|
||
}
|
||
return true
|
||
}
|
||
|
||
// failQueuedTask 在调度队列不可用时结束任务、返还积分并删除参考媒体记录。
|
||
func (s *Service) failQueuedTask(taskID uuid.UUID, message string) error {
|
||
return s.DB.Transaction(func(tx *gorm.DB) error {
|
||
var task model.GenerationTask
|
||
if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", taskID).Take(&task).Error; err != nil {
|
||
return err
|
||
}
|
||
refunded, err := billing.RefundGenerationTask(tx, &task, "图片生成入队失败返还")
|
||
if err != nil {
|
||
return err
|
||
}
|
||
var input struct {
|
||
ReferenceMediaAssetIDs []uuid.UUID `json:"reference_media_asset_ids"`
|
||
}
|
||
_ = json.Unmarshal(task.InputData, &input)
|
||
if len(input.ReferenceMediaAssetIDs) > 0 {
|
||
if err := tx.Exec("DELETE FROM media_assets WHERE id IN ?", input.ReferenceMediaAssetIDs).Error; err != nil {
|
||
return err
|
||
}
|
||
}
|
||
updates := map[string]any{"status": "failed", "error_code": "queue_unavailable", "error_message": message, "finished_at": time.Now(), "actual_points": "0.00"}
|
||
if refunded {
|
||
updates["cost_refunded"] = true
|
||
}
|
||
return tx.Model(&model.GenerationTask{}).Where("id=?", taskID).Updates(updates).Error
|
||
})
|
||
}
|
||
|
||
// validRatio 判断画幅比例是否属于图片生成工具公开支持的取值。
|
||
func validRatio(value string) bool {
|
||
for _, allowed := range []string{"9:21", "1:1", "3:4", "4:3", "9:16", "16:9"} {
|
||
if value == allowed {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
// validResolution 判断分辨率是否属于图片生成工具公开支持的档位。
|
||
func validResolution(value string) bool {
|
||
for _, allowed := range []string{"1k", "2k", "4k"} {
|
||
if value == allowed {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
// buildProviderPrompt 将文件名提及与参考图数组顺序建立明确对应关系。
|
||
func buildProviderPrompt(prompt string, names []string) string {
|
||
if len(names) == 0 {
|
||
return prompt
|
||
}
|
||
mappings := make([]string, 0, len(names))
|
||
for index, name := range names {
|
||
mappings = append(mappings, fmt.Sprintf("第%d张参考图名称为“%s”", index+1, name))
|
||
}
|
||
return "参考图对应关系：" + strings.Join(mappings, "；") + "。\n" + prompt
|
||
}
|
||
|
||
// uniqueStrings 按首次出现顺序去除空对象键和重复对象键。
|
||
func uniqueStrings(values []string) []string {
|
||
seen := make(map[string]struct{}, len(values))
|
||
result := make([]string, 0, len(values))
|
||
for _, value := range values {
|
||
if _, ok := seen[value]; ok || value == "" {
|
||
continue
|
||
}
|
||
seen[value] = struct{}{}
|
||
result = append(result, value)
|
||
}
|
||
return result
|
||
}
|