Gmarker689
96ac56cb4e
三个修复： 1. runPipelineBg 改用 context.Background()，避免 HTTP 响应返回后 Gin 取消 request context 导致后台管线静默失败 2. 新增 TaskQueue FIFO 串行队列，任务提交后进入 pending 状态排队， 按提交顺序逐个执行，前端轮询显示排队中 3. promptOptimizerNode 移除 RunPromptAgent 调用，提示词优化仅由 前端在提交前通过 /api/v1/prompt/optimize 执行一次，管线内只做 风格合并和技术参数追加
270 lines
8.1 KiB
Go
Executable File
270 lines
8.1 KiB
Go
Executable File
package service
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"fmt"
|
||
"image/png"
|
||
"strings"
|
||
|
||
"gen2d/internal/logger"
|
||
"gen2d/pkg/gifmaker"
|
||
"gen2d/pkg/splitsprite"
|
||
|
||
"github.com/cloudwego/eino/compose"
|
||
)
|
||
|
||
// promptOptimizerNode 节点：合并风格描述、重试信息与技术参数，输出最终提示词。
|
||
// 提示词优化已由前端在提交前完成，管线内不再重复调用 PromptAgent。
|
||
var promptOptimizerNode = compose.InvokableLambda(func(ctx context.Context, in PipelineInput) (string, error) {
|
||
// 合并风格描述
|
||
styleDesc := buildStyleDescription(in.ProjectStyle, in.TaskStyle)
|
||
if styleDesc != "" {
|
||
if in.Prompt != "" {
|
||
in.Prompt = in.Prompt + "。" + styleDesc
|
||
} else {
|
||
in.Prompt = styleDesc
|
||
}
|
||
}
|
||
|
||
// 注入重试原因
|
||
if in.RejectReason != "" {
|
||
if in.Prompt != "" {
|
||
in.Prompt = in.Prompt + "。注意修正以下问题：" + in.RejectReason
|
||
} else {
|
||
in.Prompt = "修正以下问题：" + in.RejectReason
|
||
}
|
||
}
|
||
|
||
if in.Prompt == "" {
|
||
return "", fmt.Errorf("pipeline: prompt is empty")
|
||
}
|
||
|
||
// 追加技术参数段，不再调用 PromptAgent
|
||
return appendTechNotes(in.Prompt, in.AssetType, in.Params), nil
|
||
})
|
||
|
||
// promptOptimizerPreHandler 首次运行时保存输入到 state；重试时注入 RejectReason。
|
||
func promptOptimizerPreHandler(ctx context.Context, in PipelineInput, state *PipelineState) (PipelineInput, error) {
|
||
if state.RetryCount == 0 {
|
||
state.Input = in
|
||
reportProgress(ctx, "prompt_builder", 10)
|
||
} else if state.RejectReason != "" {
|
||
in.RejectReason = state.RejectReason
|
||
reportProgress(ctx, "prompt_builder", 30+state.RetryCount*10)
|
||
}
|
||
return in, nil
|
||
}
|
||
|
||
// promptOptimizerPostHandler 将最终提示词写入全局状态。
|
||
func promptOptimizerPostHandler(ctx context.Context, out string, state *PipelineState) (string, error) {
|
||
state.FinalPrompt = out
|
||
reportProgress(ctx, "asset_generator", 35)
|
||
return out, nil
|
||
}
|
||
|
||
// AssetGenerator 节点：调用 AI 推理 API 出图。
|
||
var assetGeneratorNode = compose.InvokableLambda(func(ctx context.Context, prompt string) ([]GeneratedImage, error) {
|
||
var params AssetParams
|
||
_ = compose.ProcessState[*PipelineState](ctx, func(_ context.Context, state *PipelineState) error {
|
||
params = state.Input.Params
|
||
return nil
|
||
})
|
||
return GenerateImages(ctx, prompt, params)
|
||
})
|
||
|
||
// assetGeneratorPostHandler 将原始图片写入全局状态。
|
||
func assetGeneratorPostHandler(ctx context.Context, out []GeneratedImage, state *PipelineState) ([]GeneratedImage, error) {
|
||
state.RawImages = out
|
||
reportProgress(ctx, "quality_supervisor", 60)
|
||
return out, nil
|
||
}
|
||
|
||
// QualitySupervisor 节点：质检，设置路由目标。
|
||
var qualitySupervisorNode = compose.InvokableLambda(func(ctx context.Context, images []GeneratedImage) (PipelineInput, error) {
|
||
var input PipelineInput
|
||
err := compose.ProcessState[*PipelineState](ctx, func(_ context.Context, state *PipelineState) error {
|
||
state.RawImages = images
|
||
|
||
style := mergeStyle(state.Input.ProjectStyle, state.Input.TaskStyle)
|
||
pass, reason, checkErr := CheckQuality(ctx, images, style)
|
||
if checkErr != nil {
|
||
return fmt.Errorf("quality check: %w", checkErr)
|
||
}
|
||
|
||
state.PassQuality = pass
|
||
if !pass {
|
||
state.RejectReason = reason
|
||
}
|
||
|
||
if pass {
|
||
state.NextNode = nodeFormatAdapter
|
||
reportProgress(ctx, "format_adapter", 85)
|
||
} else if state.RetryCount >= 3 {
|
||
state.NextNode = nodeFormatAdapter
|
||
reportProgress(ctx, "format_adapter", 85)
|
||
} else {
|
||
state.RetryCount++
|
||
state.NextNode = nodePromptOptimizer
|
||
reportProgress(ctx, "quality_supervisor", 50+state.RetryCount*10)
|
||
}
|
||
|
||
input = state.Input
|
||
return nil
|
||
})
|
||
if err != nil {
|
||
return PipelineInput{}, err
|
||
}
|
||
return input, nil
|
||
})
|
||
|
||
// formatAdapterNode 节点：精灵表格式时调用 splitsprite 拆分 + gifmaker 生成 GIF 预览。
|
||
var formatAdapterNode = compose.InvokableLambda(func(ctx context.Context, input PipelineInput) (PipelineOutput, error) {
|
||
var images []GeneratedImage
|
||
_ = compose.ProcessState[*PipelineState](ctx, func(_ context.Context, state *PipelineState) error {
|
||
images = state.RawImages
|
||
return nil
|
||
})
|
||
|
||
params := input.Params
|
||
resolution := params.Resolution
|
||
if resolution <= 0 {
|
||
resolution = 64
|
||
}
|
||
|
||
// 精灵表模式：单张图时拆分 + GIF 预览；多图时已是独立帧，透传
|
||
if params.Format == "spritesheet" && len(images) == 1 {
|
||
return processSpriteSheet(ctx, images[0], params, resolution)
|
||
}
|
||
|
||
// 普通模式：原样透传
|
||
assets := make([]Asset, len(images))
|
||
for i, img := range images {
|
||
assets[i] = Asset{
|
||
Data: img.Data,
|
||
Format: img.Format,
|
||
URL: fmt.Sprintf("output/%d.%s", i, img.Format),
|
||
}
|
||
}
|
||
|
||
return PipelineOutput{
|
||
Assets: assets,
|
||
Metadata: AssetMetadata{
|
||
FrameWidth: resolution,
|
||
FrameHeight: resolution,
|
||
FrameCount: len(images),
|
||
Directions: params.Frames.Directions,
|
||
},
|
||
}, nil
|
||
})
|
||
|
||
// processSpriteSheet 将单张精灵表拆分为独立帧并生成 GIF 预览。
|
||
func processSpriteSheet(ctx context.Context, img GeneratedImage, params AssetParams, resolution int) (PipelineOutput, error) {
|
||
l := logger.FromCtx(ctx)
|
||
src, err := png.Decode(bytes.NewReader(img.Data))
|
||
if err != nil {
|
||
l.Error("format_adapter decode sprite sheet failed", "error", err)
|
||
return PipelineOutput{}, fmt.Errorf("decode sprite sheet: %w", err)
|
||
}
|
||
|
||
opts := splitsprite.DefaultOptions()
|
||
if params.GridRows > 0 && params.GridCols > 0 {
|
||
opts.GridRows = params.GridRows
|
||
opts.GridCols = params.GridCols
|
||
}
|
||
|
||
frames, err := splitsprite.Process(src, opts)
|
||
if err != nil {
|
||
l.Error("format_adapter split sprite sheet failed", "error", err)
|
||
return PipelineOutput{}, fmt.Errorf("split sprite sheet: %w", err)
|
||
}
|
||
l.Info("format_adapter split sprite sheet", "frame_count", len(frames))
|
||
|
||
// 帧 → Asset
|
||
assets := make([]Asset, 0, len(frames))
|
||
for i, f := range frames {
|
||
var buf bytes.Buffer
|
||
if err := png.Encode(&buf, f); err != nil {
|
||
l.Error("format_adapter encode frame failed", "error", err)
|
||
return PipelineOutput{}, fmt.Errorf("encode frame %d: %w", i, err)
|
||
}
|
||
assets = append(assets, Asset{
|
||
Data: buf.Bytes(),
|
||
Format: "png",
|
||
URL: fmt.Sprintf("output/frame_%03d.png", i),
|
||
})
|
||
}
|
||
|
||
// GIF 预览
|
||
var gifBuf bytes.Buffer
|
||
if err := gifmaker.Encode(&gifBuf, frames, nil); err != nil {
|
||
l.Warn("format_adapter generate GIF preview failed", "error", err)
|
||
} else {
|
||
l.Info("format_adapter generated GIF preview", "size_bytes", gifBuf.Len())
|
||
}
|
||
|
||
fw, fh := 0, 0
|
||
if len(frames) > 0 {
|
||
b := frames[0].Bounds()
|
||
fw, fh = b.Dx(), b.Dy()
|
||
}
|
||
|
||
return PipelineOutput{
|
||
Assets: assets,
|
||
Metadata: AssetMetadata{
|
||
FrameWidth: fw,
|
||
FrameHeight: fh,
|
||
FrameCount: len(frames),
|
||
Directions: params.Frames.Directions,
|
||
GIFPreview: gifBuf.Bytes(),
|
||
},
|
||
}, nil
|
||
}
|
||
|
||
// buildStyleDescription 将风格键值对转为自然语言描述，供 PromptAgent 注入。
|
||
func buildStyleDescription(projectStyle, taskStyle map[string]string) string {
|
||
merged := mergeStyle(projectStyle, taskStyle)
|
||
if len(merged) == 0 {
|
||
return ""
|
||
}
|
||
var parts []string
|
||
for k, v := range merged {
|
||
parts = append(parts, fmt.Sprintf("%s: %s", k, v))
|
||
}
|
||
return "风格约束：" + strings.Join(parts, "；")
|
||
}
|
||
|
||
// appendTechNotes 在无标签（不走 PromptAgent）时补上技术参数段。
|
||
func appendTechNotes(prompt, assetType string, params AssetParams) string {
|
||
var parts []string
|
||
if prompt != "" {
|
||
parts = append(parts, prompt)
|
||
}
|
||
parts = append(parts, fmt.Sprintf("素材类型: %s", assetType))
|
||
if params.Resolution > 0 {
|
||
parts = append(parts, fmt.Sprintf("分辨率: %d", params.Resolution))
|
||
}
|
||
if params.Frames.Directions > 0 {
|
||
parts = append(parts, fmt.Sprintf("方向数: %d", params.Frames.Directions))
|
||
}
|
||
if params.Frames.FramesPerDirection > 0 {
|
||
parts = append(parts, fmt.Sprintf("每方向帧数: %d", params.Frames.FramesPerDirection))
|
||
}
|
||
if params.Format != "" {
|
||
parts = append(parts, fmt.Sprintf("输出格式: %s", params.Format))
|
||
}
|
||
return strings.Join(parts, "；")
|
||
}
|
||
|
||
// mergeStyle 合并工程风格与任务风格覆盖，任务同名键覆盖工程。
|
||
func mergeStyle(projectStyle, taskStyle map[string]string) map[string]string {
|
||
result := make(map[string]string)
|
||
for k, v := range projectStyle {
|
||
result[k] = v
|
||
}
|
||
for k, v := range taskStyle {
|
||
result[k] = v
|
||
}
|
||
return result
|
||
}
|