382 lines
11 KiB
Go
382 lines
11 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"mime"
|
|
"mime/multipart"
|
|
"net"
|
|
"net/http"
|
|
"net/textproto"
|
|
"net/url"
|
|
"path/filepath"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
const maxReferenceBytes int64 = 64 << 20
|
|
|
|
type outputItem struct {
|
|
Filename string
|
|
Subfolder string
|
|
Type string
|
|
Kind string
|
|
}
|
|
|
|
func uploadReferences(comfy string, payload jsonMap) (map[string]string, error) {
|
|
files := map[string]string{}
|
|
for _, group := range []any{payload["referenceImages"], payload["referenceVideos"], payload["referenceAudios"], payload["mask"]} {
|
|
values := sliceValue(group)
|
|
if len(values) == 0 && group != nil {
|
|
values = []any{group}
|
|
}
|
|
for _, raw := range values {
|
|
media, ok := mapValue(raw)
|
|
if !ok {
|
|
continue
|
|
}
|
|
data, mimeType, err := mediaBytes(media)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(data) == 0 {
|
|
continue
|
|
}
|
|
name := safeName(firstNonEmpty(stringValue(media["id"]), randomID())) + extensionForMime(mimeType)
|
|
uploadedName, err := uploadComfyMedia(comfy, name, mimeType, data)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if id := stringValue(media["id"]); id != "" {
|
|
files[id] = uploadedName
|
|
}
|
|
}
|
|
}
|
|
return files, nil
|
|
}
|
|
|
|
func mediaBytes(media jsonMap) ([]byte, string, error) {
|
|
dataURL := stringValue(media["dataUrl"])
|
|
mimeType := firstNonEmpty(stringValue(media["mimeType"]), stringValue(media["type"]))
|
|
if strings.HasPrefix(dataURL, "data:") {
|
|
comma := strings.IndexByte(dataURL, ',')
|
|
if comma < 0 {
|
|
return nil, "", errors.New("参考素材 data URL 格式无效")
|
|
}
|
|
header := dataURL[5:comma]
|
|
if mimeType == "" {
|
|
mimeType = strings.Split(header, ";")[0]
|
|
}
|
|
data, err := base64.StdEncoding.DecodeString(dataURL[comma+1:])
|
|
if err != nil {
|
|
return nil, "", errors.New("参考素材 data URL 编码无效")
|
|
}
|
|
if int64(len(data)) > maxReferenceBytes {
|
|
return nil, "", errors.New("参考素材超过 64MB")
|
|
}
|
|
return data, firstNonEmpty(mimeType, "application/octet-stream"), nil
|
|
}
|
|
remoteURL := stringValue(media["url"])
|
|
if remoteURL == "" {
|
|
return nil, "", nil
|
|
}
|
|
if err := assertPublicRemoteURL(remoteURL); err != nil {
|
|
return nil, "", err
|
|
}
|
|
client := publicMediaClient()
|
|
response, err := client.Get(remoteURL)
|
|
if err != nil {
|
|
return nil, "", fmt.Errorf("下载远程参考素材失败：%w", err)
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
|
return nil, "", fmt.Errorf("下载远程参考素材失败 HTTP %d", response.StatusCode)
|
|
}
|
|
if response.ContentLength > maxReferenceBytes {
|
|
return nil, "", errors.New("远程参考素材超过 64MB")
|
|
}
|
|
data, err := io.ReadAll(io.LimitReader(response.Body, maxReferenceBytes+1))
|
|
if err != nil {
|
|
return nil, "", err
|
|
}
|
|
if int64(len(data)) > maxReferenceBytes {
|
|
return nil, "", errors.New("远程参考素材超过 64MB")
|
|
}
|
|
if mimeType == "" {
|
|
mimeType = strings.TrimSpace(strings.Split(response.Header.Get("Content-Type"), ";")[0])
|
|
}
|
|
return data, firstNonEmpty(mimeType, "application/octet-stream"), nil
|
|
}
|
|
|
|
func uploadComfyMedia(comfy, name, mimeType string, data []byte) (string, error) {
|
|
body := &bytes.Buffer{}
|
|
writer := multipart.NewWriter(body)
|
|
header := make(textproto.MIMEHeader)
|
|
header.Set("Content-Disposition", mime.FormatMediaType("form-data", map[string]string{"name": "image", "filename": name}))
|
|
header.Set("Content-Type", mimeType)
|
|
part, err := writer.CreatePart(header)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if _, err := part.Write(data); err != nil {
|
|
return "", err
|
|
}
|
|
if err := writer.Close(); err != nil {
|
|
return "", err
|
|
}
|
|
request, err := http.NewRequest(http.MethodPost, comfy+"/upload/image", body)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
request.Header.Set("Content-Type", writer.FormDataContentType())
|
|
response, err := comfyHTTP.Do(request)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer response.Body.Close()
|
|
dataResponse, _ := io.ReadAll(io.LimitReader(response.Body, 1<<20))
|
|
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
|
return "", fmt.Errorf("ComfyUI 上传参考素材失败 HTTP %d: %.500s", response.StatusCode, dataResponse)
|
|
}
|
|
var payload jsonMap
|
|
if json.Unmarshal(dataResponse, &payload) != nil {
|
|
return "", errors.New("ComfyUI 上传响应不是有效 JSON")
|
|
}
|
|
uploadedName := stringValue(payload["name"])
|
|
if uploadedName == "" {
|
|
return "", errors.New("ComfyUI 上传响应缺少文件名")
|
|
}
|
|
return uploadedName, nil
|
|
}
|
|
|
|
func assertPublicRemoteURL(raw string) error {
|
|
parsed, err := url.Parse(raw)
|
|
if err != nil || parsed.Hostname() == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.User != nil {
|
|
return errors.New("远程参考素材只允许无凭证的 HTTP/HTTPS 公网地址")
|
|
}
|
|
hostname := strings.ToLower(strings.Trim(parsed.Hostname(), "[]"))
|
|
if hostname == "localhost" || strings.HasSuffix(hostname, ".localhost") || strings.HasSuffix(hostname, ".local") {
|
|
return errors.New("拒绝访问本机或私网远程参考素材地址")
|
|
}
|
|
addresses, err := net.DefaultResolver.LookupIPAddr(context.Background(), hostname)
|
|
if err != nil || len(addresses) == 0 {
|
|
return errors.New("远程参考素材域名无法解析")
|
|
}
|
|
for _, address := range addresses {
|
|
if isPrivateIP(address.IP) {
|
|
return errors.New("拒绝访问解析到本机或私网的远程参考素材地址")
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func publicMediaClient() *http.Client {
|
|
transport := &http.Transport{DialContext: func(ctx context.Context, network, address string) (net.Conn, error) {
|
|
host, port, err := net.SplitHostPort(address)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
addresses, err := net.DefaultResolver.LookupIPAddr(ctx, host)
|
|
if err != nil || len(addresses) == 0 {
|
|
return nil, errors.New("远程参考素材域名无法解析")
|
|
}
|
|
for _, candidate := range addresses {
|
|
if isPrivateIP(candidate.IP) {
|
|
return nil, errors.New("拒绝访问解析到本机或私网的远程参考素材地址")
|
|
}
|
|
}
|
|
dialer := &net.Dialer{Timeout: 15 * time.Second}
|
|
return dialer.DialContext(ctx, network, net.JoinHostPort(addresses[0].IP.String(), port))
|
|
}}
|
|
return &http.Client{
|
|
Timeout: 2 * time.Minute,
|
|
Transport: transport,
|
|
CheckRedirect: func(request *http.Request, via []*http.Request) error {
|
|
if len(via) >= 5 {
|
|
return errors.New("远程参考素材重定向次数过多")
|
|
}
|
|
return assertPublicRemoteURL(request.URL.String())
|
|
},
|
|
}
|
|
}
|
|
|
|
func isPrivateIP(ip net.IP) bool {
|
|
if ip == nil || ip.IsUnspecified() || ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsMulticast() {
|
|
return true
|
|
}
|
|
if ipv4 := ip.To4(); ipv4 != nil {
|
|
return ipv4[0] == 0 || (ipv4[0] == 100 && ipv4[1] >= 64 && ipv4[1] <= 127) || (ipv4[0] == 192 && ipv4[1] == 0) || (ipv4[0] == 198 && (ipv4[1] == 18 || ipv4[1] == 19)) || ipv4[0] >= 224
|
|
}
|
|
return false
|
|
}
|
|
|
|
func collectResult(comfy, mode string, history jsonMap) (jsonMap, error) {
|
|
outputs := make([]outputItem, 0)
|
|
root, _ := mapValue(history["outputs"])
|
|
for _, nodeOutput := range root {
|
|
collectOutputItems(nodeOutput, "", &outputs)
|
|
}
|
|
if len(outputs) == 0 {
|
|
return nil, errors.New("ComfyUI 没有返回可下载产物")
|
|
}
|
|
images := make([]any, 0)
|
|
var video, audio jsonMap
|
|
for _, item := range outputs {
|
|
query := url.Values{"filename": {item.Filename}, "subfolder": {item.Subfolder}, "type": {firstNonEmpty(item.Type, "output")}}
|
|
response, err := comfyHTTP.Get(comfy + "/view?" + query.Encode())
|
|
if err != nil {
|
|
continue
|
|
}
|
|
data, readErr := io.ReadAll(io.LimitReader(response.Body, 64<<20))
|
|
response.Body.Close()
|
|
if readErr != nil || response.StatusCode < 200 || response.StatusCode >= 300 {
|
|
continue
|
|
}
|
|
mimeType := strings.TrimSpace(strings.Split(response.Header.Get("Content-Type"), ";")[0])
|
|
if mimeType == "" || mimeType == "application/octet-stream" {
|
|
mimeType = mimeForName(item.Filename)
|
|
}
|
|
value := jsonMap{"dataUrl": "data:" + mimeType + ";base64," + base64.StdEncoding.EncodeToString(data), "mimeType": mimeType, "bytes": len(data), "name": item.Filename}
|
|
switch item.Kind {
|
|
case "video":
|
|
video = value
|
|
case "audio":
|
|
audio = value
|
|
case "image":
|
|
images = append(images, value)
|
|
}
|
|
}
|
|
if mode == "video" && video != nil {
|
|
return jsonMap{"mode": "video", "video": video}, nil
|
|
}
|
|
if mode == "audio" && audio != nil {
|
|
return jsonMap{"mode": "audio", "audio": audio}, nil
|
|
}
|
|
if len(images) > 0 {
|
|
return jsonMap{"mode": "image", "images": images}, nil
|
|
}
|
|
if video != nil {
|
|
return jsonMap{"mode": "video", "video": video}, nil
|
|
}
|
|
if audio != nil {
|
|
return jsonMap{"mode": "audio", "audio": audio}, nil
|
|
}
|
|
return nil, errors.New("ComfyUI 产物下载失败")
|
|
}
|
|
|
|
func collectOutputItems(value any, forcedKind string, output *[]outputItem) {
|
|
if items := sliceValue(value); len(items) > 0 {
|
|
for _, child := range items {
|
|
collectOutputItems(child, forcedKind, output)
|
|
}
|
|
return
|
|
}
|
|
item, ok := mapValue(value)
|
|
if !ok {
|
|
return
|
|
}
|
|
if filename := stringValue(item["filename"]); filename != "" {
|
|
kind := forcedKind
|
|
if kind == "" {
|
|
kind = outputKind(filename)
|
|
}
|
|
if kind != "other" {
|
|
*output = append(*output, outputItem{Filename: filename, Subfolder: stringValue(item["subfolder"]), Type: firstNonEmpty(stringValue(item["type"]), "output"), Kind: kind})
|
|
}
|
|
return
|
|
}
|
|
for key, child := range item {
|
|
kind := ""
|
|
switch key {
|
|
case "videos":
|
|
kind = "video"
|
|
case "audio", "audios":
|
|
kind = "audio"
|
|
case "images", "gifs", "files":
|
|
kind = forcedKind
|
|
}
|
|
collectOutputItems(child, kind, output)
|
|
}
|
|
}
|
|
|
|
func safeName(value string) string {
|
|
var result strings.Builder
|
|
for _, character := range value {
|
|
if (character >= 'a' && character <= 'z') || (character >= 'A' && character <= 'Z') || (character >= '0' && character <= '9') || character == '_' || character == '-' {
|
|
result.WriteRune(character)
|
|
}
|
|
if result.Len() >= 64 {
|
|
break
|
|
}
|
|
}
|
|
if result.Len() == 0 {
|
|
return "reference"
|
|
}
|
|
return result.String()
|
|
}
|
|
|
|
func extensionForMime(value string) string {
|
|
lower := strings.ToLower(value)
|
|
switch {
|
|
case strings.Contains(lower, "jpeg"):
|
|
return ".jpg"
|
|
case strings.Contains(lower, "webp"):
|
|
return ".webp"
|
|
case strings.Contains(lower, "video"):
|
|
return ".mp4"
|
|
case strings.Contains(lower, "audio"):
|
|
return ".wav"
|
|
default:
|
|
return ".png"
|
|
}
|
|
}
|
|
|
|
func mimeForName(value string) string {
|
|
extension := strings.ToLower(filepath.Ext(value))
|
|
switch extension {
|
|
case ".webm":
|
|
return "video/webm"
|
|
case ".mov":
|
|
return "video/quicktime"
|
|
case ".mp4", ".m4v", ".mkv":
|
|
return "video/mp4"
|
|
case ".wav":
|
|
return "audio/wav"
|
|
case ".ogg":
|
|
return "audio/ogg"
|
|
case ".flac":
|
|
return "audio/flac"
|
|
case ".m4a", ".aac":
|
|
return "audio/mp4"
|
|
case ".mp3":
|
|
return "audio/mpeg"
|
|
case ".jpg", ".jpeg":
|
|
return "image/jpeg"
|
|
case ".webp":
|
|
return "image/webp"
|
|
case ".gif":
|
|
return "image/gif"
|
|
default:
|
|
return "image/png"
|
|
}
|
|
}
|
|
|
|
func outputKind(value string) string {
|
|
extension := strings.ToLower(filepath.Ext(value))
|
|
if stringIn(extension, ".mp4", ".webm", ".mov", ".m4v", ".mkv") {
|
|
return "video"
|
|
}
|
|
if stringIn(extension, ".mp3", ".wav", ".ogg", ".m4a", ".flac", ".aac") {
|
|
return "audio"
|
|
}
|
|
if stringIn(extension, ".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp") {
|
|
return "image"
|
|
}
|
|
return "other"
|
|
}
|