feat(workflow): 工作流提供方 - 集成 RunningHub 与 ComfyUI Bridge
This commit is contained in:
yepwait authored and ddCat committed
1 parent
6141bd688c
commit
86fc3fda2f
82 files changed
+11404
-247
No files matched your search
@@ -0,0 +1,381 @@
|
||||
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"
|
||||
}
|
||||
Reference in new issue
Block a user