92ee465a41
Bundle the current backend and web changes into one publishable snapshot from the server worktree so the remote mirror reflects the deployed source. Constraint: Push the existing mixed worktree from /docker/new-api/src without reshaping files Rejected: Split into feature-specific commits | current worktree is already a single undifferentiated snapshot Confidence: medium Scope-risk: moderate Directive: Split future feature work before deployment to preserve clearer history Tested: git status inspection; git diff --stat review Not-tested: build, lint, unit tests, integration tests
1088 lines
30 KiB
Go
1088 lines
30 KiB
Go
package controller
|
||
|
||
import (
|
||
"bytes"
|
||
"crypto/hmac"
|
||
"crypto/sha256"
|
||
"encoding/base64"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"mime/multipart"
|
||
"net/http"
|
||
"net/textproto"
|
||
"os"
|
||
"path/filepath"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/QuantumNous/new-api/common"
|
||
"github.com/QuantumNous/new-api/constant"
|
||
"github.com/QuantumNous/new-api/model"
|
||
"github.com/QuantumNous/new-api/service"
|
||
"github.com/QuantumNous/new-api/setting"
|
||
"github.com/gin-gonic/gin"
|
||
)
|
||
|
||
const (
|
||
imageStudioGroup = "image"
|
||
imageStudioTaskStatusPending = "pending"
|
||
imageStudioTaskStatusRunning = "running"
|
||
imageStudioTaskStatusSuccess = "success"
|
||
imageStudioTaskStatusFailure = "failure"
|
||
imageStudioTaskTTL = 30 * time.Minute
|
||
imageStudioTaskCleanupEvery = 5 * time.Minute
|
||
imageStudioHTTPTimeout = 10 * time.Minute
|
||
imageStudioGenerateAction = "generate"
|
||
imageStudioEditAction = "edit"
|
||
imageStudioResultBaseDir = "/data/image-studio-results"
|
||
imageStudioResultURLPrefix = "/image-studio-results"
|
||
)
|
||
|
||
type imageStudioGenerateRequest struct {
|
||
TokenID int `json:"token_id"`
|
||
Model string `json:"model"`
|
||
Prompt string `json:"prompt"`
|
||
Size string `json:"size"`
|
||
Quality string `json:"quality"`
|
||
}
|
||
|
||
type imageStudioUploadedImage struct {
|
||
Filename string
|
||
ContentType string
|
||
Data []byte
|
||
}
|
||
|
||
type imageStudioTask struct {
|
||
ID string `json:"task_id"`
|
||
UserID int `json:"-"`
|
||
Mode string `json:"mode"`
|
||
Model string `json:"model"`
|
||
Size string `json:"size"`
|
||
Quality string `json:"quality"`
|
||
Status string `json:"status"`
|
||
Error string `json:"error,omitempty"`
|
||
Result json.RawMessage `json:"result,omitempty"`
|
||
ResultDir string `json:"-"`
|
||
CreatedAt int64 `json:"created_at"`
|
||
UpdatedAt int64 `json:"updated_at"`
|
||
}
|
||
|
||
var (
|
||
imageStudioTasks = struct {
|
||
sync.RWMutex
|
||
items map[string]*imageStudioTask
|
||
}{items: make(map[string]*imageStudioTask)}
|
||
imageStudioCleanupOnce sync.Once
|
||
)
|
||
|
||
func ensureImageStudioTaskCleanup() {
|
||
imageStudioCleanupOnce.Do(func() {
|
||
go func() {
|
||
ticker := time.NewTicker(imageStudioTaskCleanupEvery)
|
||
defer ticker.Stop()
|
||
for range ticker.C {
|
||
cutoff := time.Now().Add(-imageStudioTaskTTL).Unix()
|
||
imageStudioTasks.Lock()
|
||
for id, task := range imageStudioTasks.items {
|
||
if task.UpdatedAt < cutoff {
|
||
cleanupImageStudioTaskArtifacts(task)
|
||
delete(imageStudioTasks.items, id)
|
||
}
|
||
}
|
||
imageStudioTasks.Unlock()
|
||
}
|
||
}()
|
||
})
|
||
}
|
||
|
||
func newImageStudioTask(userID int, mode string, req imageStudioGenerateRequest) (*imageStudioTask, error) {
|
||
ensureImageStudioTaskCleanup()
|
||
key, err := common.GenerateRandomCharsKey(24)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
now := time.Now().Unix()
|
||
task := &imageStudioTask{
|
||
ID: "imgtask_" + key,
|
||
UserID: userID,
|
||
Mode: mode,
|
||
Model: strings.TrimSpace(req.Model),
|
||
Size: resolveImageStudioSize(req.Model, req.Size),
|
||
Quality: normalizeImageStudioQuality(req.Quality),
|
||
Status: imageStudioTaskStatusPending,
|
||
CreatedAt: now,
|
||
UpdatedAt: now,
|
||
}
|
||
imageStudioTasks.Lock()
|
||
imageStudioTasks.items[task.ID] = task
|
||
imageStudioTasks.Unlock()
|
||
return task, nil
|
||
}
|
||
|
||
func getImageStudioTask(taskID string) (*imageStudioTask, bool) {
|
||
imageStudioTasks.RLock()
|
||
defer imageStudioTasks.RUnlock()
|
||
task, ok := imageStudioTasks.items[taskID]
|
||
return task, ok
|
||
}
|
||
|
||
func updateImageStudioTask(taskID string, updateFn func(task *imageStudioTask)) {
|
||
imageStudioTasks.Lock()
|
||
defer imageStudioTasks.Unlock()
|
||
if task, ok := imageStudioTasks.items[taskID]; ok {
|
||
updateFn(task)
|
||
task.UpdatedAt = time.Now().Unix()
|
||
}
|
||
}
|
||
|
||
func cleanupImageStudioTaskArtifacts(task *imageStudioTask) {
|
||
if task == nil || strings.TrimSpace(task.ResultDir) == "" {
|
||
return
|
||
}
|
||
baseDir := filepath.Clean(imageStudioResultBaseDir)
|
||
resultDir := filepath.Clean(task.ResultDir)
|
||
prefix := baseDir + string(os.PathSeparator)
|
||
if resultDir == baseDir || !strings.HasPrefix(resultDir, prefix) {
|
||
return
|
||
}
|
||
_ = os.RemoveAll(resultDir)
|
||
task.ResultDir = ""
|
||
}
|
||
|
||
func normalizeImageStudioSize(v string) string {
|
||
value := strings.TrimSpace(strings.ToLower(v))
|
||
switch value {
|
||
case "", "auto":
|
||
return "auto"
|
||
case "1024x1024", "1536x1024", "1024x1536", "2048x2048", "2048x1152", "1152x2048", "3840x2160", "2160x3840":
|
||
return value
|
||
default:
|
||
return "auto"
|
||
}
|
||
}
|
||
|
||
func normalizeImageStudioModel(v string) string {
|
||
return strings.TrimSpace(strings.ToLower(v))
|
||
}
|
||
|
||
func isGPTImage2Model(modelName string) bool {
|
||
return normalizeImageStudioModel(modelName) == "gpt-image-2"
|
||
}
|
||
|
||
func isValidGPTImage2Size(size string) bool {
|
||
if size == "" || size == "auto" {
|
||
return true
|
||
}
|
||
parts := strings.Split(size, "x")
|
||
if len(parts) != 2 {
|
||
return false
|
||
}
|
||
width, err := strconv.Atoi(parts[0])
|
||
if err != nil {
|
||
return false
|
||
}
|
||
height, err := strconv.Atoi(parts[1])
|
||
if err != nil {
|
||
return false
|
||
}
|
||
if width <= 0 || height <= 0 {
|
||
return false
|
||
}
|
||
if width > 3840 || height > 3840 {
|
||
return false
|
||
}
|
||
if width%16 != 0 || height%16 != 0 {
|
||
return false
|
||
}
|
||
longEdge := width
|
||
shortEdge := height
|
||
if height > width {
|
||
longEdge = height
|
||
shortEdge = width
|
||
}
|
||
if longEdge > shortEdge*3 {
|
||
return false
|
||
}
|
||
pixels := width * height
|
||
if pixels < 655360 || pixels > 8294400 {
|
||
return false
|
||
}
|
||
return true
|
||
}
|
||
|
||
func resolveImageStudioSize(modelName string, size string) string {
|
||
normalizedSize := normalizeImageStudioSize(size)
|
||
if isGPTImage2Model(modelName) {
|
||
candidate := strings.TrimSpace(strings.ToLower(size))
|
||
if isValidGPTImage2Size(candidate) {
|
||
return candidate
|
||
}
|
||
if normalizedSize != "auto" {
|
||
return normalizedSize
|
||
}
|
||
}
|
||
return normalizedSize
|
||
}
|
||
|
||
func normalizeImageStudioQuality(v string) string {
|
||
value := strings.TrimSpace(strings.ToLower(v))
|
||
switch value {
|
||
case "", "auto":
|
||
return "auto"
|
||
case "low", "medium", "high":
|
||
return value
|
||
default:
|
||
return "auto"
|
||
}
|
||
}
|
||
|
||
func normalizeImageStudioBearerToken(key string) string {
|
||
trimmed := strings.TrimSpace(key)
|
||
if strings.HasPrefix(trimmed, "sk-") {
|
||
return trimmed
|
||
}
|
||
return "sk-" + trimmed
|
||
}
|
||
|
||
var imageStudioDisallowedKeywords = []string{
|
||
"暴力",
|
||
"血腥",
|
||
"伤害",
|
||
"虐待",
|
||
"自残",
|
||
"自杀",
|
||
"恐怖",
|
||
"战争",
|
||
"武器",
|
||
"尸体",
|
||
"残肢",
|
||
"恶心",
|
||
"色情",
|
||
"裸露",
|
||
"性暗示",
|
||
"卖淫",
|
||
"约炮",
|
||
"成人视频",
|
||
"强奷",
|
||
"sm",
|
||
"性暴力",
|
||
"性服务",
|
||
"赌博",
|
||
"博彩",
|
||
"赌场",
|
||
"下注",
|
||
"洗钱",
|
||
"毒品",
|
||
"制毒",
|
||
"吸毒",
|
||
"贩毒",
|
||
"违禁药物",
|
||
"政治敏感",
|
||
"反动",
|
||
"分裂",
|
||
"颠覆",
|
||
"煽动",
|
||
"国家领导人",
|
||
"国旗",
|
||
"国徽",
|
||
"国歌",
|
||
"恐怖主义",
|
||
"暴恐",
|
||
}
|
||
|
||
func buildImageStudioSafetyPrompt() string {
|
||
return setting.ImageStudioSafetyPrompt
|
||
}
|
||
|
||
func normalizeImageStudioPromptForSafety(prompt string) string {
|
||
replacer := strings.NewReplacer(
|
||
" ", "",
|
||
"\n", "",
|
||
"\r", "",
|
||
"\t", "",
|
||
"-", "",
|
||
"_", "",
|
||
"/", "",
|
||
"\\", "",
|
||
"|", "",
|
||
"*", "",
|
||
".", "",
|
||
",", "",
|
||
"，", "",
|
||
"。", "",
|
||
":", "",
|
||
"：", "",
|
||
";", "",
|
||
"；", "",
|
||
"(", "",
|
||
")", "",
|
||
"（", "",
|
||
"）", "",
|
||
"[", "",
|
||
"]", "",
|
||
"【", "",
|
||
"】", "",
|
||
"{", "",
|
||
"}", "",
|
||
"<", "",
|
||
">", "",
|
||
"《", "",
|
||
"》", "",
|
||
"'", "",
|
||
"\"", "",
|
||
"`", "",
|
||
"!", "",
|
||
"！", "",
|
||
"?", "",
|
||
"？", "",
|
||
)
|
||
return strings.ToLower(replacer.Replace(strings.TrimSpace(prompt)))
|
||
}
|
||
|
||
func containsImageStudioDisallowedKeyword(prompt string) (bool, string) {
|
||
checkText := normalizeImageStudioPromptForSafety(prompt)
|
||
for _, keyword := range setting.ImageStudioBannedWords {
|
||
k := normalizeImageStudioPromptForSafety(keyword)
|
||
if k == "" {
|
||
continue
|
||
}
|
||
if strings.Contains(checkText, k) {
|
||
return true, keyword
|
||
}
|
||
}
|
||
return false, ""
|
||
}
|
||
|
||
func validateImageStudioPromptSafety(prompt string) error {
|
||
trimmed := strings.TrimSpace(prompt)
|
||
if trimmed == "" {
|
||
return fmt.Errorf("请先填写提示词")
|
||
}
|
||
if ok, keyword := containsImageStudioDisallowedKeyword(trimmed); ok {
|
||
return fmt.Errorf("提示词包含禁止内容：%s", keyword)
|
||
}
|
||
if ok, words := service.SensitiveWordContains(trimmed); ok {
|
||
if len(words) > 0 {
|
||
return fmt.Errorf("提示词包含敏感内容：%s", strings.Join(words, ", "))
|
||
}
|
||
return fmt.Errorf("提示词包含敏感内容，请修改后重试")
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func imageStudioModeEndpointType(mode string) constant.EndpointType {
|
||
if mode == imageStudioEditAction {
|
||
return constant.EndpointTypeImageEdit
|
||
}
|
||
return constant.EndpointTypeImageGeneration
|
||
}
|
||
|
||
func imageStudioModePath(mode string) string {
|
||
if mode == imageStudioEditAction {
|
||
return "/v1/images/edits"
|
||
}
|
||
return "/v1/images/generations"
|
||
}
|
||
|
||
func resolveImageStudioGenerateRelayAction(images []imageStudioUploadedImage) string {
|
||
if len(images) > 0 {
|
||
return imageStudioEditAction
|
||
}
|
||
return imageStudioGenerateAction
|
||
}
|
||
|
||
func validateImageStudioTokenAndModel(userID int, req imageStudioGenerateRequest, mode string) (*model.Token, error) {
|
||
token, err := model.GetTokenByIds(req.TokenID, userID)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("未找到所选令牌，请刷新后重试")
|
||
}
|
||
if token.Status != common.TokenStatusEnabled {
|
||
return nil, fmt.Errorf("当前令牌状态不可用，请去控制台-令牌管理检查该令牌")
|
||
}
|
||
if strings.TrimSpace(token.Group) != imageStudioGroup {
|
||
return nil, fmt.Errorf("所选令牌不是 image 分组，请去控制台-令牌管理创建或选择 image 分组令牌")
|
||
}
|
||
modelName := strings.TrimSpace(req.Model)
|
||
if modelName == "" {
|
||
return nil, fmt.Errorf("请选择图片模型")
|
||
}
|
||
if isGPTImage2Model(modelName) && !isValidGPTImage2Size(strings.TrimSpace(strings.ToLower(req.Size))) {
|
||
return nil, fmt.Errorf("gpt-image-2 尺寸不合法：最长边不能超过 3840，边长需为 16 的倍数，长宽比不能超过 3:1，总像素需在 655360 到 8294400 之间")
|
||
}
|
||
if token.ModelLimitsEnabled {
|
||
limits := token.GetModelLimitsMap()
|
||
if len(limits) == 0 || !limits[modelName] {
|
||
return nil, fmt.Errorf("当前令牌没有放行该图片模型，请去控制台-令牌管理调整模型限制")
|
||
}
|
||
}
|
||
|
||
pricingItems := model.GetPricing()
|
||
var matched *model.Pricing
|
||
for index := range pricingItems {
|
||
item := pricingItems[index]
|
||
if item.ModelName != modelName {
|
||
continue
|
||
}
|
||
if !containsString(item.EnableGroup, imageStudioGroup) {
|
||
continue
|
||
}
|
||
matched = &item
|
||
break
|
||
}
|
||
if matched == nil {
|
||
return nil, fmt.Errorf("当前 image 分组未配置该图片模型，请联系管理员检查 image 分组和模型配置")
|
||
}
|
||
if !containsEndpointType(matched.SupportedEndpointTypes, imageStudioModeEndpointType(mode)) {
|
||
if mode == imageStudioEditAction {
|
||
return nil, fmt.Errorf("当前模型不支持编辑图片，请联系管理员检查模型能力配置")
|
||
}
|
||
return nil, fmt.Errorf("当前模型不支持生成图片，请联系管理员检查模型能力配置")
|
||
}
|
||
return token, nil
|
||
}
|
||
|
||
func containsString(items []string, target string) bool {
|
||
for _, item := range items {
|
||
if strings.TrimSpace(item) == target {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
func containsEndpointType(items []constant.EndpointType, target constant.EndpointType) bool {
|
||
for _, item := range items {
|
||
if item == target {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
func readImageStudioUploads(c *gin.Context, required bool) ([]imageStudioUploadedImage, error) {
|
||
form, err := c.MultipartForm()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
fileHeaders := form.File["image[]"]
|
||
if len(fileHeaders) == 0 {
|
||
fileHeaders = form.File["images[]"]
|
||
}
|
||
if len(fileHeaders) == 0 {
|
||
fileHeaders = form.File["image"]
|
||
}
|
||
if len(fileHeaders) == 0 {
|
||
if required {
|
||
return nil, fmt.Errorf("编辑图片至少需要上传一张参考图")
|
||
}
|
||
return nil, nil
|
||
}
|
||
images := make([]imageStudioUploadedImage, 0, len(fileHeaders))
|
||
for _, header := range fileHeaders {
|
||
image, err := readImageStudioUpload(header)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
images = append(images, image)
|
||
}
|
||
return images, nil
|
||
}
|
||
|
||
func readImageStudioUpload(header *multipart.FileHeader) (imageStudioUploadedImage, error) {
|
||
file, err := header.Open()
|
||
if err != nil {
|
||
return imageStudioUploadedImage{}, err
|
||
}
|
||
defer file.Close()
|
||
data, err := io.ReadAll(file)
|
||
if err != nil {
|
||
return imageStudioUploadedImage{}, err
|
||
}
|
||
contentType := header.Header.Get("Content-Type")
|
||
if contentType == "" {
|
||
contentType = http.DetectContentType(data)
|
||
}
|
||
return imageStudioUploadedImage{
|
||
Filename: header.Filename,
|
||
ContentType: contentType,
|
||
Data: data,
|
||
}, nil
|
||
}
|
||
|
||
func buildImageStudioInternalURL(mode string) string {
|
||
return fmt.Sprintf("http://127.0.0.1:%d%s", *common.Port, imageStudioModePath(mode))
|
||
}
|
||
|
||
func decodeImageStudioError(body []byte, fallback string) string {
|
||
if len(body) == 0 {
|
||
return fallback
|
||
}
|
||
var payload map[string]any
|
||
if err := json.Unmarshal(body, &payload); err == nil {
|
||
if errorPart, ok := payload["error"].(map[string]any); ok {
|
||
if message, ok := errorPart["message"].(string); ok && strings.TrimSpace(message) != "" {
|
||
return message
|
||
}
|
||
}
|
||
if message, ok := payload["message"].(string); ok && strings.TrimSpace(message) != "" {
|
||
return message
|
||
}
|
||
}
|
||
return fallback
|
||
}
|
||
|
||
func getImageStudioResultExtension(outputFormat string) string {
|
||
switch strings.ToLower(strings.TrimSpace(outputFormat)) {
|
||
case "jpg", "jpeg":
|
||
return "jpg"
|
||
case "webp":
|
||
return "webp"
|
||
default:
|
||
return "png"
|
||
}
|
||
}
|
||
|
||
func signImageStudioResult(relativePath string, expiresAt int64) string {
|
||
mac := hmac.New(sha256.New, []byte(common.CryptoSecret))
|
||
_, _ = mac.Write([]byte(relativePath))
|
||
_, _ = mac.Write([]byte("\n"))
|
||
_, _ = mac.Write([]byte(strconv.FormatInt(expiresAt, 10)))
|
||
return hex.EncodeToString(mac.Sum(nil))
|
||
}
|
||
|
||
func buildImageStudioResultFileURL(taskID string, index int, outputFormat string) string {
|
||
relativePath := fmt.Sprintf("%s/%d.%s", taskID, index, getImageStudioResultExtension(outputFormat))
|
||
expiresAt := time.Now().Add(imageStudioTaskTTL).Unix()
|
||
signature := signImageStudioResult(relativePath, expiresAt)
|
||
return fmt.Sprintf("%s/%s?expires=%d&sig=%s", imageStudioResultURLPrefix, relativePath, expiresAt, signature)
|
||
}
|
||
|
||
func persistImageStudioTaskResult(taskID string, respBody []byte) (json.RawMessage, string, error) {
|
||
var payload map[string]any
|
||
if err := json.Unmarshal(respBody, &payload); err != nil {
|
||
return nil, "", err
|
||
}
|
||
data, ok := payload["data"].([]any)
|
||
if !ok || len(data) == 0 {
|
||
return json.RawMessage(respBody), "", nil
|
||
}
|
||
resultDir := filepath.Join(imageStudioResultBaseDir, taskID)
|
||
wroteFiles := false
|
||
for index, rawItem := range data {
|
||
item, ok := rawItem.(map[string]any)
|
||
if !ok {
|
||
continue
|
||
}
|
||
b64, _ := item["b64_json"].(string)
|
||
b64 = strings.TrimSpace(b64)
|
||
if b64 == "" {
|
||
continue
|
||
}
|
||
decoded, err := base64.StdEncoding.DecodeString(b64)
|
||
if err != nil {
|
||
_ = os.RemoveAll(resultDir)
|
||
return nil, "", fmt.Errorf("图片结果解码失败")
|
||
}
|
||
if !wroteFiles {
|
||
_ = os.RemoveAll(resultDir)
|
||
if err := os.MkdirAll(resultDir, 0o755); err != nil {
|
||
return nil, "", fmt.Errorf("图片结果目录创建失败")
|
||
}
|
||
}
|
||
outputFormat, _ := item["output_format"].(string)
|
||
ext := getImageStudioResultExtension(outputFormat)
|
||
filename := fmt.Sprintf("%d.%s", index, ext)
|
||
fullPath := filepath.Join(resultDir, filename)
|
||
if err := os.WriteFile(fullPath, decoded, 0o644); err != nil {
|
||
_ = os.RemoveAll(resultDir)
|
||
return nil, "", fmt.Errorf("图片结果写入失败")
|
||
}
|
||
delete(item, "b64_json")
|
||
item["url"] = buildImageStudioResultFileURL(taskID, index, outputFormat)
|
||
wroteFiles = true
|
||
}
|
||
if !wroteFiles {
|
||
return json.RawMessage(respBody), "", nil
|
||
}
|
||
encoded, err := json.Marshal(payload)
|
||
if err != nil {
|
||
_ = os.RemoveAll(resultDir)
|
||
return nil, "", fmt.Errorf("图片结果序列化失败")
|
||
}
|
||
return json.RawMessage(encoded), resultDir, nil
|
||
}
|
||
|
||
func executeImageStudioGenerateTask(taskID string, tokenKey string, req imageStudioGenerateRequest, images []imageStudioUploadedImage) {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusRunning
|
||
task.Error = ""
|
||
})
|
||
|
||
relayAction := resolveImageStudioGenerateRelayAction(images)
|
||
var body bytes.Buffer
|
||
contentType := "application/json"
|
||
if len(images) > 0 {
|
||
writer := multipart.NewWriter(&body)
|
||
_ = writer.WriteField("model", strings.TrimSpace(req.Model))
|
||
_ = writer.WriteField("prompt", buildImageStudioSafetyPrompt()+"\n\n"+strings.TrimSpace(req.Prompt))
|
||
_ = writer.WriteField("size", resolveImageStudioSize(req.Model, req.Size))
|
||
_ = writer.WriteField("quality", normalizeImageStudioQuality(req.Quality))
|
||
_ = writer.WriteField("response_format", "b64_json")
|
||
_ = writer.WriteField("output_format", "png")
|
||
|
||
for _, image := range images {
|
||
headers := make(textproto.MIMEHeader)
|
||
headers.Set("Content-Disposition", fmt.Sprintf(`form-data; name="image[]"; filename="%s"`, image.Filename))
|
||
headers.Set("Content-Type", image.ContentType)
|
||
|
||
part, err := writer.CreatePart(headers)
|
||
if err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = "图片生成参考图封装失败"
|
||
})
|
||
_ = writer.Close()
|
||
return
|
||
}
|
||
if _, err = part.Write(image.Data); err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = "图片生成参考图写入失败"
|
||
})
|
||
_ = writer.Close()
|
||
return
|
||
}
|
||
}
|
||
if err := writer.Close(); err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = "图片生成请求构建失败"
|
||
})
|
||
return
|
||
}
|
||
contentType = writer.FormDataContentType()
|
||
} else {
|
||
payload, err := json.Marshal(gin.H{
|
||
"model": strings.TrimSpace(req.Model),
|
||
"prompt": buildImageStudioSafetyPrompt() + "\n\n" + strings.TrimSpace(req.Prompt),
|
||
"size": resolveImageStudioSize(req.Model, req.Size),
|
||
"quality": normalizeImageStudioQuality(req.Quality),
|
||
"response_format": "b64_json",
|
||
"output_format": "png",
|
||
})
|
||
if err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = "图片生成请求构建失败"
|
||
})
|
||
return
|
||
}
|
||
body.Write(payload)
|
||
}
|
||
|
||
httpReq, err := http.NewRequest(http.MethodPost, buildImageStudioInternalURL(relayAction), bytes.NewReader(body.Bytes()))
|
||
if err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = "图片生成请求创建失败"
|
||
})
|
||
return
|
||
}
|
||
httpReq.Header.Set("Authorization", "Bearer "+normalizeImageStudioBearerToken(tokenKey))
|
||
httpReq.Header.Set("Content-Type", contentType)
|
||
|
||
client := &http.Client{Timeout: imageStudioHTTPTimeout}
|
||
resp, err := client.Do(httpReq)
|
||
if err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = "图片生成请求失败：" + err.Error()
|
||
})
|
||
return
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
respBody, err := io.ReadAll(resp.Body)
|
||
if err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = "读取图片生成结果失败"
|
||
})
|
||
return
|
||
}
|
||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = decodeImageStudioError(respBody, "图片生成失败")
|
||
})
|
||
return
|
||
}
|
||
|
||
sanitizedResult, resultDir, err := persistImageStudioTaskResult(taskID, respBody)
|
||
if err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = err.Error()
|
||
})
|
||
return
|
||
}
|
||
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusSuccess
|
||
cleanupImageStudioTaskArtifacts(task)
|
||
task.Result = sanitizedResult
|
||
task.ResultDir = resultDir
|
||
task.Error = ""
|
||
})
|
||
}
|
||
|
||
func executeImageStudioEditTask(taskID string, tokenKey string, req imageStudioGenerateRequest, images []imageStudioUploadedImage) {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusRunning
|
||
task.Error = ""
|
||
})
|
||
|
||
var body bytes.Buffer
|
||
writer := multipart.NewWriter(&body)
|
||
_ = writer.WriteField("model", strings.TrimSpace(req.Model))
|
||
_ = writer.WriteField("prompt", buildImageStudioSafetyPrompt()+"\n\n"+strings.TrimSpace(req.Prompt))
|
||
_ = writer.WriteField("size", resolveImageStudioSize(req.Model, req.Size))
|
||
_ = writer.WriteField("quality", normalizeImageStudioQuality(req.Quality))
|
||
_ = writer.WriteField("response_format", "b64_json")
|
||
_ = writer.WriteField("output_format", "png")
|
||
|
||
for _, image := range images {
|
||
headers := make(textproto.MIMEHeader)
|
||
headers.Set("Content-Disposition", fmt.Sprintf(`form-data; name="image[]"; filename="%s"`, image.Filename))
|
||
headers.Set("Content-Type", image.ContentType)
|
||
|
||
part, err := writer.CreatePart(headers)
|
||
if err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = "编辑图片参考图封装失败"
|
||
})
|
||
_ = writer.Close()
|
||
return
|
||
}
|
||
if _, err = part.Write(image.Data); err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = "编辑图片参考图写入失败"
|
||
})
|
||
_ = writer.Close()
|
||
return
|
||
}
|
||
}
|
||
if err := writer.Close(); err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = "编辑图片请求构建失败"
|
||
})
|
||
return
|
||
}
|
||
|
||
httpReq, err := http.NewRequest(http.MethodPost, buildImageStudioInternalURL(imageStudioEditAction), bytes.NewReader(body.Bytes()))
|
||
if err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = "编辑图片请求创建失败"
|
||
})
|
||
return
|
||
}
|
||
httpReq.Header.Set("Authorization", "Bearer "+normalizeImageStudioBearerToken(tokenKey))
|
||
httpReq.Header.Set("Content-Type", writer.FormDataContentType())
|
||
|
||
client := &http.Client{Timeout: imageStudioHTTPTimeout}
|
||
resp, err := client.Do(httpReq)
|
||
if err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = "编辑图片请求失败：" + err.Error()
|
||
})
|
||
return
|
||
}
|
||
defer resp.Body.Close()
|
||
|
||
respBody, err := io.ReadAll(resp.Body)
|
||
if err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = "读取编辑图片结果失败"
|
||
})
|
||
return
|
||
}
|
||
if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = decodeImageStudioError(respBody, "图片编辑失败")
|
||
})
|
||
return
|
||
}
|
||
|
||
sanitizedResult, resultDir, err := persistImageStudioTaskResult(taskID, respBody)
|
||
if err != nil {
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusFailure
|
||
task.Error = err.Error()
|
||
})
|
||
return
|
||
}
|
||
|
||
updateImageStudioTask(taskID, func(task *imageStudioTask) {
|
||
task.Status = imageStudioTaskStatusSuccess
|
||
cleanupImageStudioTaskArtifacts(task)
|
||
task.Result = sanitizedResult
|
||
task.ResultDir = resultDir
|
||
task.Error = ""
|
||
})
|
||
}
|
||
|
||
func SubmitImageStudioGenerate(c *gin.Context) {
|
||
userID := c.GetInt("id")
|
||
var req imageStudioGenerateRequest
|
||
var images []imageStudioUploadedImage
|
||
var err error
|
||
|
||
if strings.Contains(strings.ToLower(c.GetHeader("Content-Type")), "multipart/form-data") {
|
||
req = imageStudioGenerateRequest{
|
||
Model: c.PostForm("model"),
|
||
Prompt: strings.TrimSpace(c.PostForm("prompt")),
|
||
Size: c.PostForm("size"),
|
||
Quality: c.PostForm("quality"),
|
||
}
|
||
if tokenID := strings.TrimSpace(c.PostForm("token_id")); tokenID != "" {
|
||
_, _ = fmt.Sscanf(tokenID, "%d", &req.TokenID)
|
||
}
|
||
images, err = readImageStudioUploads(c, false)
|
||
if err != nil {
|
||
common.ApiErrorMsg(c, "图片生成参考图读取失败，请重新上传后重试")
|
||
return
|
||
}
|
||
} else if err = c.ShouldBindJSON(&req); err != nil {
|
||
common.ApiErrorMsg(c, "图片生成参数格式不正确，请刷新后重试")
|
||
return
|
||
}
|
||
|
||
req.Prompt = strings.TrimSpace(req.Prompt)
|
||
if err := validateImageStudioPromptSafety(req.Prompt); err != nil {
|
||
common.ApiErrorMsg(c, err.Error())
|
||
return
|
||
}
|
||
validationMode := imageStudioGenerateAction
|
||
if len(images) > 0 {
|
||
validationMode = imageStudioEditAction
|
||
}
|
||
token, err := validateImageStudioTokenAndModel(userID, req, validationMode)
|
||
if err != nil {
|
||
common.ApiErrorMsg(c, err.Error())
|
||
return
|
||
}
|
||
task, err := newImageStudioTask(userID, imageStudioGenerateAction, req)
|
||
if err != nil {
|
||
common.ApiErrorMsg(c, "创建图片任务失败，请稍后重试")
|
||
return
|
||
}
|
||
go executeImageStudioGenerateTask(task.ID, token.GetFullKey(), req, images)
|
||
common.ApiSuccess(c, task)
|
||
}
|
||
|
||
func SubmitImageStudioEdit(c *gin.Context) {
|
||
userID := c.GetInt("id")
|
||
req := imageStudioGenerateRequest{
|
||
Model: c.PostForm("model"),
|
||
Prompt: strings.TrimSpace(c.PostForm("prompt")),
|
||
Size: c.PostForm("size"),
|
||
Quality: c.PostForm("quality"),
|
||
}
|
||
if tokenID := strings.TrimSpace(c.PostForm("token_id")); tokenID != "" {
|
||
_, _ = fmt.Sscanf(tokenID, "%d", &req.TokenID)
|
||
}
|
||
if err := validateImageStudioPromptSafety(req.Prompt); err != nil {
|
||
common.ApiErrorMsg(c, err.Error())
|
||
return
|
||
}
|
||
images, err := readImageStudioUploads(c, true)
|
||
if err != nil {
|
||
common.ApiErrorMsg(c, err.Error())
|
||
return
|
||
}
|
||
token, err := validateImageStudioTokenAndModel(userID, req, imageStudioEditAction)
|
||
if err != nil {
|
||
common.ApiErrorMsg(c, err.Error())
|
||
return
|
||
}
|
||
task, err := newImageStudioTask(userID, imageStudioEditAction, req)
|
||
if err != nil {
|
||
common.ApiErrorMsg(c, "创建图片任务失败，请稍后重试")
|
||
return
|
||
}
|
||
go executeImageStudioEditTask(task.ID, token.GetFullKey(), req, images)
|
||
common.ApiSuccess(c, task)
|
||
}
|
||
|
||
func buildImageStudioTaskResultURL(taskID string, index int) string {
|
||
return fmt.Sprintf("/api/image-studio/task/%s/result/%d", taskID, index)
|
||
}
|
||
|
||
func getImageStudioResultContentType(outputFormat string) string {
|
||
switch strings.ToLower(strings.TrimSpace(outputFormat)) {
|
||
case "jpg", "jpeg":
|
||
return "image/jpeg"
|
||
case "webp":
|
||
return "image/webp"
|
||
default:
|
||
return "image/png"
|
||
}
|
||
}
|
||
|
||
func sanitizeImageStudioTaskForResponse(task *imageStudioTask) *imageStudioTask {
|
||
if task == nil {
|
||
return nil
|
||
}
|
||
clone := *task
|
||
if clone.Status != imageStudioTaskStatusSuccess || len(clone.Result) == 0 {
|
||
return &clone
|
||
}
|
||
var payload map[string]any
|
||
if err := json.Unmarshal(clone.Result, &payload); err != nil {
|
||
return &clone
|
||
}
|
||
data, ok := payload["data"].([]any)
|
||
if !ok || len(data) == 0 {
|
||
return &clone
|
||
}
|
||
changed := false
|
||
for index, rawItem := range data {
|
||
item, ok := rawItem.(map[string]any)
|
||
if !ok {
|
||
continue
|
||
}
|
||
if b64, _ := item["b64_json"].(string); strings.TrimSpace(b64) != "" {
|
||
delete(item, "b64_json")
|
||
item["url"] = buildImageStudioTaskResultURL(task.ID, index)
|
||
changed = true
|
||
}
|
||
}
|
||
if !changed {
|
||
return &clone
|
||
}
|
||
if encoded, err := json.Marshal(payload); err == nil {
|
||
clone.Result = json.RawMessage(encoded)
|
||
}
|
||
return &clone
|
||
}
|
||
|
||
func GetImageStudioTask(c *gin.Context) {
|
||
userID := c.GetInt("id")
|
||
taskID := strings.TrimSpace(c.Param("id"))
|
||
if taskID == "" {
|
||
common.ApiErrorMsg(c, "任务不存在")
|
||
return
|
||
}
|
||
task, ok := getImageStudioTask(taskID)
|
||
if !ok || task == nil || task.UserID != userID {
|
||
common.ApiErrorMsg(c, "任务不存在或已过期，请重新提交")
|
||
return
|
||
}
|
||
common.ApiSuccess(c, sanitizeImageStudioTaskForResponse(task))
|
||
}
|
||
|
||
func GetImageStudioTaskResult(c *gin.Context) {
|
||
userID := c.GetInt("id")
|
||
taskID := strings.TrimSpace(c.Param("id"))
|
||
if taskID == "" {
|
||
common.ApiErrorMsg(c, "任务不存在")
|
||
return
|
||
}
|
||
index, err := strconv.Atoi(strings.TrimSpace(c.Param("index")))
|
||
if err != nil || index < 0 {
|
||
common.ApiErrorMsg(c, "图片结果不存在")
|
||
return
|
||
}
|
||
task, ok := getImageStudioTask(taskID)
|
||
if !ok || task == nil || task.UserID != userID {
|
||
common.ApiErrorMsg(c, "任务不存在或已过期，请重新提交")
|
||
return
|
||
}
|
||
if task.Status != imageStudioTaskStatusSuccess || len(task.Result) == 0 {
|
||
common.ApiErrorMsg(c, "图片结果尚未生成完成")
|
||
return
|
||
}
|
||
var payload map[string]any
|
||
if err := json.Unmarshal(task.Result, &payload); err != nil {
|
||
common.ApiErrorMsg(c, "图片结果解析失败")
|
||
return
|
||
}
|
||
data, ok := payload["data"].([]any)
|
||
if !ok || index >= len(data) {
|
||
common.ApiErrorMsg(c, "图片结果不存在")
|
||
return
|
||
}
|
||
item, ok := data[index].(map[string]any)
|
||
if !ok {
|
||
common.ApiErrorMsg(c, "图片结果不存在")
|
||
return
|
||
}
|
||
if url, _ := item["url"].(string); strings.TrimSpace(url) != "" {
|
||
c.Redirect(http.StatusTemporaryRedirect, url)
|
||
return
|
||
}
|
||
b64, _ := item["b64_json"].(string)
|
||
if strings.TrimSpace(b64) == "" {
|
||
common.ApiErrorMsg(c, "图片结果不存在")
|
||
return
|
||
}
|
||
decoded, err := base64.StdEncoding.DecodeString(b64)
|
||
if err != nil {
|
||
common.ApiErrorMsg(c, "图片结果解码失败")
|
||
return
|
||
}
|
||
contentType := getImageStudioResultContentType(fmt.Sprintf("%v", item["output_format"]))
|
||
c.Header("Cache-Control", "private, no-store, max-age=0")
|
||
c.Data(http.StatusOK, contentType, decoded)
|
||
}
|
||
|
||
func GetImageStudioResultFile(c *gin.Context) {
|
||
relativePath := strings.TrimPrefix(c.Param("filepath"), "/")
|
||
relativePath = strings.TrimSpace(relativePath)
|
||
if relativePath == "" {
|
||
c.AbortWithStatus(http.StatusNotFound)
|
||
return
|
||
}
|
||
expiresAt, err := strconv.ParseInt(strings.TrimSpace(c.Query("expires")), 10, 64)
|
||
if err != nil || expiresAt <= 0 || time.Now().Unix() > expiresAt {
|
||
c.AbortWithStatus(http.StatusForbidden)
|
||
return
|
||
}
|
||
sig := strings.TrimSpace(c.Query("sig"))
|
||
if sig == "" {
|
||
c.AbortWithStatus(http.StatusForbidden)
|
||
return
|
||
}
|
||
expected := signImageStudioResult(relativePath, expiresAt)
|
||
if !hmac.Equal([]byte(strings.ToLower(sig)), []byte(expected)) {
|
||
c.AbortWithStatus(http.StatusForbidden)
|
||
return
|
||
}
|
||
baseDir := filepath.Clean(imageStudioResultBaseDir)
|
||
fullPath := filepath.Clean(filepath.Join(baseDir, filepath.Clean(relativePath)))
|
||
prefix := baseDir + string(os.PathSeparator)
|
||
if fullPath == baseDir || !strings.HasPrefix(fullPath, prefix) {
|
||
c.AbortWithStatus(http.StatusNotFound)
|
||
return
|
||
}
|
||
info, err := os.Stat(fullPath)
|
||
if err != nil || info.IsDir() {
|
||
c.AbortWithStatus(http.StatusNotFound)
|
||
return
|
||
}
|
||
c.Header("Cache-Control", "private, no-store, no-cache, max-age=0")
|
||
c.Header("Pragma", "no-cache")
|
||
c.File(fullPath)
|
||
}
|