admin
a6de9ebd12
git-subtree-dir: server git-subtree-mainline:c4e617ada7git-subtree-split:e64421295f
115 lines
3.4 KiB
Go
115 lines
3.4 KiB
Go
package main
|
|
|
|
// 为指定用户生成一套三视角全身照（真实调用 imagegen，非 mock）：
|
|
// go run scripts/gen_user_photos/main.go [username]
|
|
// 默认用户 wenwu901。已存在同视角照片时跳过；图片存 workspace/user_{id}/photos/，记录写入 slogan_user_photo。
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"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/entity"
|
|
)
|
|
|
|
const personDesc = "一位穿浅蓝色衬衫与深灰色西裤的亚洲年轻女性，干净利落的黑色短发，身材匀称"
|
|
|
|
var views = []struct {
|
|
angle string
|
|
photoT int
|
|
prompt string
|
|
}{
|
|
{angle: "front", photoT: consts.PhotoTypeFullFront, prompt: personDesc + "，全身正面照，站直面对镜头，双手自然下垂，纯白背景，高清写实，全身入镜"},
|
|
{angle: "side", photoT: consts.PhotoTypeFullSide, prompt: personDesc + "，全身侧面照，侧身站立目视前方，纯白背景，高清写实，全身入镜"},
|
|
{angle: "back", photoT: consts.PhotoTypeFullBack, prompt: personDesc + "，全身背面照，背对镜头站立，纯白背景，高清写实，全身入镜"},
|
|
}
|
|
|
|
func main() {
|
|
username := "wenwu901"
|
|
if len(os.Args) > 1 {
|
|
username = os.Args[1]
|
|
}
|
|
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)
|
|
|
|
existing, err := dao.UserPhoto.ListByUser(ctx, user.Id, 0)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
have := map[int]bool{}
|
|
for _, p := range existing {
|
|
have[p.Type] = true
|
|
}
|
|
|
|
client, err := agent.NewClient(g.Cfg().MustGet(ctx, "imagegen.supplier", "wanx").String())
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
dir := filepath.Join("workspace", fmt.Sprintf("user_%d", user.Id), "photos")
|
|
if err := os.MkdirAll(dir, 0o755); err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
// 三视角用同一 seed，保证人物一致
|
|
seed := time.Now().UnixNano() % 1_000_000
|
|
for _, v := range views {
|
|
if have[v.photoT] {
|
|
fmt.Printf("视角 %s 已有照片，跳过\n", v.angle)
|
|
continue
|
|
}
|
|
fmt.Printf("生成 %s 视角...\n", v.angle)
|
|
url, err := client.Generate(ctx, &agent.GenerateReq{
|
|
Prompt: v.prompt, Angle: v.angle, Seed: seed,
|
|
})
|
|
if err != nil {
|
|
panic(fmt.Sprintf("生成 %s 失败: %v", v.angle, err))
|
|
}
|
|
path := filepath.Join(dir, fmt.Sprintf("%d_%s.png", time.Now().UnixNano(), v.angle))
|
|
if err := download(url, path); err != nil {
|
|
panic(fmt.Sprintf("保存 %s 失败: %v", v.angle, err))
|
|
}
|
|
if _, err := dao.UserPhoto.Insert(ctx, &entity.UserPhoto{
|
|
UserId: user.Id, Type: v.photoT, Url: "/" + filepath.ToSlash(path), Status: 1,
|
|
}); err != nil {
|
|
panic(fmt.Sprintf("入库 %s 失败: %v", v.angle, err))
|
|
}
|
|
fmt.Printf("%s 完成: %s\n", v.angle, path)
|
|
}
|
|
fmt.Println("照片套生成完毕")
|
|
}
|
|
|
|
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
|
|
}
|