1
This commit is contained in:
@@ -0,0 +1,223 @@
|
||||
package main
|
||||
|
||||
// 为 wenwu901 真实生成一套 AI 穿搭方案（真实调用 imagegen 与 LLM，非 mock）：
|
||||
// go run scripts/gen_outfit_plan/main.go
|
||||
// 步骤：补衣橱（8 件单品，imagegen 生成服装图）→ 调 OutfitService.Generate →
|
||||
// 轮询任务到 done → 选主方案（触发效果图异步生成）→ 等 3 张效果图完成。
|
||||
// 前置：config.yml 已配置 geo.amap_key + weather.qweather_key（天气硬依赖）。
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
_ "github.com/gogf/gf/contrib/drivers/sqlite/v2"
|
||||
"github.com/gogf/gf/v2/frame/g"
|
||||
|
||||
"slogan-agent/styleagent/agent"
|
||||
"slogan-agent/styleagent/consts"
|
||||
"slogan-agent/styleagent/dao"
|
||||
"slogan-agent/styleagent/model/dto"
|
||||
"slogan-agent/styleagent/model/entity"
|
||||
"slogan-agent/styleagent/service"
|
||||
)
|
||||
|
||||
const username = "wenwu901"
|
||||
|
||||
type garment struct {
|
||||
name string
|
||||
category string
|
||||
style string
|
||||
color string
|
||||
prompt string
|
||||
}
|
||||
|
||||
var garments = []garment{
|
||||
{"白色长袖衬衫", "上衣", "休闲", "白色", "纯白背景的白色长袖衬衫商品图，正面展示，高清，电商风格"},
|
||||
{"灰色圆领T恤", "上衣", "休闲", "灰色", "纯白背景的灰色圆领T恤商品图，正面展示，高清，电商风格"},
|
||||
{"深蓝夹克外套", "上衣", "外套", "深蓝", "纯白背景的深蓝色夹克外套商品图，正面展示，高清，电商风格"},
|
||||
{"深灰休闲长裤", "下装", "休闲", "深灰", "纯白背景的深灰色休闲长裤商品图，正面展示，高清，电商风格"},
|
||||
{"蓝色牛仔裤", "下装", "休闲", "蓝色", "纯白背景的蓝色牛仔裤商品图，正面展示，高清，电商风格"},
|
||||
{"白色运动鞋", "鞋", "休闲", "白色", "纯白背景的白色运动鞋商品图，侧面展示，高清，电商风格"},
|
||||
{"棕色皮鞋", "鞋", "商务", "棕色", "纯白背景的棕色皮鞋商品图，侧面展示，高清，电商风格"},
|
||||
{"黑色双肩背包", "配饰", "休闲", "黑色", "纯白背景的黑色双肩背包商品图，正面展示，高清，电商风格"},
|
||||
}
|
||||
|
||||
func main() {
|
||||
ctx := context.Background()
|
||||
|
||||
var user entity.User
|
||||
if err := g.DB().Model(consts.TableNameUser).Ctx(ctx).
|
||||
Where("username", username).Scan(&user); err != nil || user.Id == 0 {
|
||||
panic(fmt.Sprintf("用户 %s 不存在: %v", username, err))
|
||||
}
|
||||
fmt.Printf("用户: %s (id=%d)\n", username, user.Id)
|
||||
|
||||
if err := ensureWardrobe(ctx, user.Id); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
// 幂等：已有方案则只重试效果图（主方案 → select-main → 等 3 张 done）
|
||||
var mainPlan *entity.OutfitPlan
|
||||
existing, _ := dao.OutfitPlan.ListByUser(ctx, user.Id)
|
||||
for _, p := range existing {
|
||||
if p.MainFlag == 1 {
|
||||
mainPlan = p
|
||||
}
|
||||
}
|
||||
if len(existing) == 0 {
|
||||
taskId, err := service.OutfitService.Generate(ctx, user.Id, &dto.OutfitGenerateReq{
|
||||
StartDate: "2026-08-01", EndDate: "2026-08-07", Location: "上海", Occasion: "通勤",
|
||||
})
|
||||
if err != nil {
|
||||
if strings.Contains(err.Error(), "未配置") {
|
||||
panic(fmt.Sprintf("%v\n请先在 config.yml 配置 geo.amap_key / weather.qweather_key 后重跑", err))
|
||||
}
|
||||
panic(fmt.Sprintf("发起方案生成失败: %v", err))
|
||||
}
|
||||
fmt.Printf("生成任务已提交: task_id=%d，轮询中...\n", taskId)
|
||||
|
||||
waitTaskDone(ctx, taskId, user.Id)
|
||||
plans, err := dao.OutfitPlan.ListByTask(ctx, taskId)
|
||||
if err != nil || len(plans) == 0 {
|
||||
panic(fmt.Sprintf("任务完成但无方案: %v", err))
|
||||
}
|
||||
mainPlan = plans[0]
|
||||
for _, p := range plans {
|
||||
if p.MainFlag == 1 {
|
||||
mainPlan = p
|
||||
}
|
||||
}
|
||||
} else if mainPlan == nil {
|
||||
panic("已有方案但无主方案，请先选主方案")
|
||||
}
|
||||
fmt.Printf("主方案: id=%d %s 评分=%d\n", mainPlan.Id, mainPlan.Title, mainPlan.Score)
|
||||
if err := service.OutfitPlanService.SelectMain(ctx, user.Id, mainPlan.Id); err != nil {
|
||||
panic(fmt.Sprintf("选主方案失败: %v", err))
|
||||
}
|
||||
|
||||
waitEffects(ctx, mainPlan.Id)
|
||||
|
||||
all, _ := dao.OutfitPlan.ListByUser(ctx, user.Id)
|
||||
fmt.Printf("完成！wenwu901 现有 %d 套方案：\n", len(all))
|
||||
for _, p := range all {
|
||||
fmt.Printf(" - plan %d: %s（评分 %d，%s）\n", p.Id, p.Title, p.Score, p.Source)
|
||||
}
|
||||
}
|
||||
|
||||
// ensureWardrobe 为指定用户补齐 8 件衣橱单品（同 Category 已有则跳过该分类），服装图用 imagegen 生成
|
||||
func ensureWardrobe(ctx context.Context, userId int64) error {
|
||||
existing, err := dao.WardrobeItem.ListAllByUser(ctx, userId)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
have := map[string]bool{}
|
||||
for _, it := range existing {
|
||||
have[it.Category] = true
|
||||
}
|
||||
need := make([]garment, 0, len(garments))
|
||||
for _, ga := range garments {
|
||||
if !have[ga.category] {
|
||||
need = append(need, ga)
|
||||
}
|
||||
}
|
||||
if len(need) == 0 {
|
||||
fmt.Println("衣橱 4 类已齐，跳过补衣橱")
|
||||
return nil
|
||||
}
|
||||
|
||||
client, err := agent.NewClient(g.Cfg().MustGet(ctx, "imagegen.supplier", "wanx").String())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
dir := filepath.Join("workspace", fmt.Sprintf("user_%d", userId), "wardrobe")
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, ga := range need {
|
||||
fmt.Printf("生成服装图: %s...\n", ga.name)
|
||||
url, err := client.Generate(ctx, &agent.GenerateReq{
|
||||
Prompt: ga.prompt, Seed: time.Now().UnixNano() % 1_000_000,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("生成 %s 服装图失败: %w", ga.name, err)
|
||||
}
|
||||
path := filepath.Join(dir, fmt.Sprintf("%d_%s.png", time.Now().UnixNano(), ga.name))
|
||||
if err := download(url, path); err != nil {
|
||||
return fmt.Errorf("保存 %s 失败: %w", ga.name, err)
|
||||
}
|
||||
if _, err := dao.WardrobeItem.Insert(ctx, &entity.WardrobeItem{
|
||||
UserId: userId, PhotoUrl: "/" + filepath.ToSlash(path),
|
||||
Name: ga.name, Category: ga.category, Season: "四季", StyleTags: ga.style,
|
||||
ColorInfo: ga.color, Status: 1,
|
||||
}); err != nil {
|
||||
return fmt.Errorf("入库 %s 失败: %w", ga.name, err)
|
||||
}
|
||||
fmt.Printf("%s 完成: %s\n", ga.name, path)
|
||||
}
|
||||
fmt.Println("衣橱补齐完毕")
|
||||
return nil
|
||||
}
|
||||
|
||||
func waitTaskDone(ctx context.Context, taskId, userId int64) {
|
||||
for i := 0; i < 30; i++ {
|
||||
task, err := dao.OutfitGenTask.GetOne(ctx, taskId, userId)
|
||||
if err != nil || task == nil {
|
||||
panic(fmt.Sprintf("读取任务失败: %v", err))
|
||||
}
|
||||
switch task.Status {
|
||||
case consts.TaskStatusDone:
|
||||
fmt.Println("方案生成完成")
|
||||
return
|
||||
case consts.TaskStatusFailed:
|
||||
panic(fmt.Sprintf("方案生成失败: %s", task.Error))
|
||||
}
|
||||
time.Sleep(10 * time.Second)
|
||||
}
|
||||
panic("方案生成超时（5 分钟）")
|
||||
}
|
||||
|
||||
// waitEffects 等主方案的 3 张效果图（正面/侧面/背面）全部 done
|
||||
func waitEffects(ctx context.Context, planId int64) {
|
||||
for i := 0; i < 20; i++ {
|
||||
images, _ := dao.PlanEffectImage.ListByPlan(ctx, planId)
|
||||
done := 0
|
||||
for _, im := range images {
|
||||
if im.Status == consts.EffectStatusDone {
|
||||
done++
|
||||
}
|
||||
}
|
||||
if done >= 3 {
|
||||
fmt.Printf("效果图 3 张完成\n")
|
||||
return
|
||||
}
|
||||
if i == 19 {
|
||||
fmt.Printf("警告: 效果图超时（完成 %d/3），可稍后查看\n", done)
|
||||
return
|
||||
}
|
||||
time.Sleep(15 * time.Second)
|
||||
}
|
||||
}
|
||||
|
||||
func download(url, dest string) error {
|
||||
resp, err := http.Get(url)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("下载失败: http %d", resp.StatusCode)
|
||||
}
|
||||
out, err := os.Create(dest)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer out.Close()
|
||||
_, err = io.Copy(out, resp.Body)
|
||||
return err
|
||||
}
|
||||
Reference in New Issue
Block a user