Gmarker689
9ca1e65221
- PipelineInput 新增 Tags/UserNote 字段 - 管线重构为: START → PromptOptimizer → AssetGenerator → QualitySupervisor → FormatAdapter - promptOptimizerNode: 有标签调用 PromptAgent，无标签补技术参数，合并风格与重试信息 - PromptBuilder 移除，提示词构建逻辑并入 PromptOptimizer - 新增 POST /api/v1/prompt/optimize 接口 - main.go 注入 LLM/ImageGen 配置，注册 prompt 路由 - inference.go 新增 InitImageGenConfig 注入
113 lines
3.2 KiB
Go
113 lines
3.2 KiB
Go
package service
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/rand"
|
|
"fmt"
|
|
"image"
|
|
"image/color"
|
|
"image/png"
|
|
|
|
"gen2d/internal/config"
|
|
)
|
|
|
|
// imgCfg 保存文生图配置，由 main 通过 InitImageGenConfig 注入。
|
|
var imgCfg config.ImageGenConfig
|
|
|
|
// InitImageGenConfig 注入文生图配置。
|
|
func InitImageGenConfig(cfg config.ImageGenConfig) {
|
|
imgCfg = cfg
|
|
}
|
|
|
|
// GenerateImages 调用 AI 推理 API 生成图片。
|
|
// MVP 阶段返回 mock 占位图。
|
|
func GenerateImages(ctx context.Context, prompt string, params AssetParams) ([]GeneratedImage, error) {
|
|
size := params.Resolution
|
|
if size <= 0 {
|
|
size = 64
|
|
}
|
|
|
|
count := 1
|
|
if params.Frames.Directions > 0 && params.Frames.FramesPerDirection > 0 {
|
|
count = params.Frames.Directions * params.Frames.FramesPerDirection
|
|
}
|
|
|
|
images := make([]GeneratedImage, count)
|
|
for i := 0; i < count; i++ {
|
|
data, err := generateMockImage(size, i)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("generate mock image %d: %w", i, err)
|
|
}
|
|
images[i] = GeneratedImage{
|
|
Data: data,
|
|
Width: size,
|
|
Height: size,
|
|
Format: "png",
|
|
}
|
|
}
|
|
return images, nil
|
|
}
|
|
|
|
// QualityChecker 质检函数，可替换用于测试。
|
|
// 签名：(ctx, images, style) → (pass, reason, error)
|
|
var QualityChecker = defaultCheckQuality
|
|
|
|
// CheckQuality 调用当前 QualityChecker。
|
|
func CheckQuality(ctx context.Context, images []GeneratedImage, style map[string]string) (bool, string, error) {
|
|
return QualityChecker(ctx, images, style)
|
|
}
|
|
|
|
// defaultCheckQuality 默认 mock 质检，始终返回 pass。
|
|
func defaultCheckQuality(ctx context.Context, images []GeneratedImage, style map[string]string) (bool, string, error) {
|
|
return true, "", nil
|
|
}
|
|
|
|
// NewCountedQualityChecker 创建一个在第 passOnRetry 次调用时返回 pass 的质检函数。
|
|
func NewCountedQualityChecker(passOnRetry int) func(context.Context, []GeneratedImage, map[string]string) (bool, string, error) {
|
|
var callCount int
|
|
return func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) {
|
|
callCount++
|
|
if callCount >= passOnRetry {
|
|
return true, "", nil
|
|
}
|
|
return false, fmt.Sprintf("风格不一致（第 %d 次质检）", callCount), nil
|
|
}
|
|
}
|
|
|
|
// AlwaysFailQualityChecker 始终返回 fail 的质检函数。
|
|
func AlwaysFailQualityChecker() func(context.Context, []GeneratedImage, map[string]string) (bool, string, error) {
|
|
return func(_ context.Context, _ []GeneratedImage, _ map[string]string) (bool, string, error) {
|
|
return false, "风格不一致", nil
|
|
}
|
|
}
|
|
|
|
// generateMockImage 生成一张带随机色块的 PNG 占位图
|
|
func generateMockImage(size int, seed int) ([]byte, error) {
|
|
img := image.NewRGBA(image.Rect(0, 0, size, size))
|
|
|
|
// 用 seed 生成不同颜色
|
|
r := uint8((seed*47 + 13) % 256)
|
|
g := uint8((seed*83 + 37) % 256)
|
|
b := uint8((seed*61 + 71) % 256)
|
|
|
|
for y := 0; y < size; y++ {
|
|
for x := 0; x < size; x++ {
|
|
img.Set(x, y, color.RGBA{R: r, G: g, B: b, A: 255})
|
|
}
|
|
}
|
|
|
|
var buf bytes.Buffer
|
|
if err := png.Encode(&buf, img); err != nil {
|
|
return nil, err
|
|
}
|
|
return buf.Bytes(), nil
|
|
}
|
|
|
|
// generateRandomBytes 用于生成随机数据（备用）
|
|
func generateRandomBytes(n int) ([]byte, error) {
|
|
b := make([]byte, n)
|
|
_, err := rand.Read(b)
|
|
return b, err
|
|
}
|