gen2d/backend/internal/service/inference.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
}