510 lines
16 KiB
Go
510 lines
16 KiB
Go
package multimodal
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"net/textproto"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
var ErrEmptyResult = errors.New("multimodal provider returned an empty result")
|
|
|
|
type SpeechClient struct {
|
|
BaseURL string
|
|
Client *http.Client
|
|
}
|
|
|
|
type VisionClient struct {
|
|
BaseURL string
|
|
Client *http.Client
|
|
}
|
|
|
|
type ImageGenerationClient struct {
|
|
BaseURL string
|
|
Client *http.Client
|
|
}
|
|
|
|
type VideoUnderstandingClient struct {
|
|
BaseURL string
|
|
Client *http.Client
|
|
}
|
|
|
|
type GeneratedImage struct {
|
|
DataBase64 string
|
|
SourceURL string
|
|
ContentType string
|
|
}
|
|
|
|
type ImageGenerator interface {
|
|
Generate(ctx context.Context, model, prompt, size string) (GeneratedImage, error)
|
|
}
|
|
|
|
// ComfyImageGenerationClient is the replaceable transport for a ComfyUI
|
|
// image-generation skill. The workflow is deliberately kept in this adapter:
|
|
// the core only selects a capability and supplies the user's prompt.
|
|
type ComfyImageGenerationClient struct {
|
|
BaseURL string
|
|
Client *http.Client
|
|
PollInterval time.Duration
|
|
PollTimeout time.Duration
|
|
}
|
|
|
|
func (c SpeechClient) Transcribe(ctx context.Context, filename, contentType string, data []byte) (string, error) {
|
|
if len(data) == 0 {
|
|
return "", errors.New("audio payload is empty")
|
|
}
|
|
body := &bytes.Buffer{}
|
|
writer := multipart.NewWriter(body)
|
|
header := make(textproto.MIMEHeader)
|
|
header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, strings.ReplaceAll(filename, `"`, "")))
|
|
header.Set("Content-Type", contentType)
|
|
part, err := writer.CreatePart(header)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if _, err := part.Write(data); err != nil {
|
|
return "", err
|
|
}
|
|
if err := writer.WriteField("response_format", "json"); err != nil {
|
|
return "", err
|
|
}
|
|
if err := writer.Close(); err != nil {
|
|
return "", err
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/inference", body)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
request.Header.Set("Content-Type", writer.FormDataContentType())
|
|
response, err := c.client().Do(request)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode != http.StatusOK {
|
|
return "", fmt.Errorf("speech provider status %d", response.StatusCode)
|
|
}
|
|
var result struct {
|
|
Text string `json:"text"`
|
|
}
|
|
if err := json.NewDecoder(io.LimitReader(response.Body, 1<<20)).Decode(&result); err != nil {
|
|
return "", err
|
|
}
|
|
if strings.TrimSpace(result.Text) == "" {
|
|
return "", ErrEmptyResult
|
|
}
|
|
return strings.TrimSpace(result.Text), nil
|
|
}
|
|
|
|
func (c VisionClient) Describe(ctx context.Context, model, contentType string, data []byte, instruction string) (string, error) {
|
|
if len(data) == 0 {
|
|
return "", errors.New("image payload is empty")
|
|
}
|
|
payload := struct {
|
|
Model string `json:"model"`
|
|
Messages []struct {
|
|
Role string `json:"role"`
|
|
Content []struct {
|
|
Type string `json:"type"`
|
|
Text string `json:"text,omitempty"`
|
|
ImageURL *struct {
|
|
URL string `json:"url"`
|
|
} `json:"image_url,omitempty"`
|
|
} `json:"content"`
|
|
} `json:"messages"`
|
|
MaxTokens int `json:"max_tokens"`
|
|
}{}
|
|
payload.Model = model
|
|
payload.MaxTokens = 1024
|
|
payload.Messages = append(payload.Messages, struct {
|
|
Role string `json:"role"`
|
|
Content []struct {
|
|
Type string `json:"type"`
|
|
Text string `json:"text,omitempty"`
|
|
ImageURL *struct {
|
|
URL string `json:"url"`
|
|
} `json:"image_url,omitempty"`
|
|
} `json:"content"`
|
|
}{Role: "user", Content: []struct {
|
|
Type string `json:"type"`
|
|
Text string `json:"text,omitempty"`
|
|
ImageURL *struct {
|
|
URL string `json:"url"`
|
|
} `json:"image_url,omitempty"`
|
|
}{
|
|
{Type: "text", Text: instruction},
|
|
{Type: "image_url", ImageURL: &struct {
|
|
URL string `json:"url"`
|
|
}{URL: "data:" + contentType + ";base64," + encodeBase64(data)},
|
|
}}})
|
|
encoded, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/chat/completions", bytes.NewReader(encoded))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
request.Header.Set("Content-Type", "application/json")
|
|
response, err := c.client().Do(request)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode != http.StatusOK {
|
|
return "", fmt.Errorf("vision provider status %d", response.StatusCode)
|
|
}
|
|
var result struct {
|
|
Choices []struct {
|
|
Message struct {
|
|
Content string `json:"content"`
|
|
ReasoningContent string `json:"reasoning_content"`
|
|
} `json:"message"`
|
|
} `json:"choices"`
|
|
}
|
|
if err := json.NewDecoder(io.LimitReader(response.Body, 2<<20)).Decode(&result); err != nil {
|
|
return "", err
|
|
}
|
|
if len(result.Choices) == 0 {
|
|
return "", ErrEmptyResult
|
|
}
|
|
text := strings.TrimSpace(result.Choices[0].Message.Content)
|
|
if text == "" {
|
|
text = strings.TrimSpace(result.Choices[0].Message.ReasoningContent)
|
|
}
|
|
if text == "" {
|
|
return "", ErrEmptyResult
|
|
}
|
|
return text, nil
|
|
}
|
|
|
|
// Generate uses an OpenAI-compatible image-generation skill. The provider is
|
|
// selected outside the core by the capability binding; this adapter contains
|
|
// transport only and never decides what the user meant.
|
|
func (c ImageGenerationClient) Generate(ctx context.Context, model, prompt, size string) (GeneratedImage, error) {
|
|
if strings.TrimSpace(prompt) == "" {
|
|
return GeneratedImage{}, errors.New("image prompt is empty")
|
|
}
|
|
payload := struct {
|
|
Model string `json:"model"`
|
|
Prompt string `json:"prompt"`
|
|
Size string `json:"size,omitempty"`
|
|
ResponseFormat string `json:"response_format,omitempty"`
|
|
}{Model: model, Prompt: prompt, Size: size, ResponseFormat: "b64_json"}
|
|
encoded, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return GeneratedImage{}, err
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/images/generations", bytes.NewReader(encoded))
|
|
if err != nil {
|
|
return GeneratedImage{}, err
|
|
}
|
|
request.Header.Set("Content-Type", "application/json")
|
|
response, err := c.client().Do(request)
|
|
if err != nil {
|
|
return GeneratedImage{}, err
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode < 200 || response.StatusCode > 299 {
|
|
return GeneratedImage{}, fmt.Errorf("image provider status %d", response.StatusCode)
|
|
}
|
|
var result struct {
|
|
Data []struct {
|
|
Base64 string `json:"b64_json"`
|
|
URL string `json:"url"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.NewDecoder(io.LimitReader(response.Body, 32<<20)).Decode(&result); err != nil {
|
|
return GeneratedImage{}, err
|
|
}
|
|
if len(result.Data) == 0 || (strings.TrimSpace(result.Data[0].Base64) == "" && strings.TrimSpace(result.Data[0].URL) == "") {
|
|
return GeneratedImage{}, ErrEmptyResult
|
|
}
|
|
return GeneratedImage{DataBase64: result.Data[0].Base64, SourceURL: result.Data[0].URL, ContentType: "image/png"}, nil
|
|
}
|
|
|
|
func (c ComfyImageGenerationClient) Generate(ctx context.Context, model, prompt, size string) (GeneratedImage, error) {
|
|
if strings.TrimSpace(prompt) == "" {
|
|
return GeneratedImage{}, errors.New("image prompt is empty")
|
|
}
|
|
width, height := imageSize(size)
|
|
workflow := map[string]any{
|
|
"3": map[string]any{"inputs": map[string]any{"seed": time.Now().UnixNano(), "steps": 8, "cfg": 1.0, "sampler_name": "euler", "scheduler": "simple", "denoise": 1.0, "model": []any{"16", 0}, "positive": []any{"6", 0}, "negative": []any{"7", 0}, "latent_image": []any{"13", 0}}, "class_type": "KSampler"},
|
|
"6": map[string]any{"inputs": map[string]any{"text": prompt, "clip": []any{"18", 0}}, "class_type": "CLIPTextEncode"},
|
|
"7": map[string]any{"inputs": map[string]any{"text": "blurry, ugly, distorted", "clip": []any{"18", 0}}, "class_type": "CLIPTextEncode"},
|
|
"8": map[string]any{"inputs": map[string]any{"samples": []any{"3", 0}, "vae": []any{"17", 0}}, "class_type": "VAEDecode"},
|
|
"9": map[string]any{"inputs": map[string]any{"filename_prefix": "domovoy", "images": []any{"8", 0}}, "class_type": "SaveImage"},
|
|
"13": map[string]any{"inputs": map[string]any{"width": width, "height": height, "batch_size": 1}, "class_type": "EmptySD3LatentImage"},
|
|
"16": map[string]any{"inputs": map[string]any{"unet_name": "z_image_turbo_bf16.safetensors", "weight_dtype": "default"}, "class_type": "UNETLoader"},
|
|
"17": map[string]any{"inputs": map[string]any{"vae_name": "ae.safetensors"}, "class_type": "VAELoader"},
|
|
"18": map[string]any{"inputs": map[string]any{"clip_name": "qwen_3_4b.safetensors", "type": "lumina2", "device": "default"}, "class_type": "CLIPLoader"},
|
|
}
|
|
payload := map[string]any{"prompt": workflow, "client_id": fmt.Sprintf("domovoy-%d", time.Now().UnixNano())}
|
|
encoded, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return GeneratedImage{}, err
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/prompt", bytes.NewReader(encoded))
|
|
if err != nil {
|
|
return GeneratedImage{}, err
|
|
}
|
|
request.Header.Set("Content-Type", "application/json")
|
|
response, err := c.client().Do(request)
|
|
if err != nil {
|
|
return GeneratedImage{}, err
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode < 200 || response.StatusCode > 299 {
|
|
return GeneratedImage{}, fmt.Errorf("comfy provider status %d", response.StatusCode)
|
|
}
|
|
var accepted struct {
|
|
PromptID string `json:"prompt_id"`
|
|
}
|
|
if err := json.NewDecoder(io.LimitReader(response.Body, 1<<20)).Decode(&accepted); err != nil {
|
|
return GeneratedImage{}, err
|
|
}
|
|
if accepted.PromptID == "" {
|
|
return GeneratedImage{}, ErrEmptyResult
|
|
}
|
|
// A caller timeout must remove only this prompt. ComfyUI's global
|
|
// /interrupt endpoint can otherwise stop unrelated users' work and leave
|
|
// queued requests in an ambiguous state.
|
|
defer func() {
|
|
if ctx.Err() != nil {
|
|
cancelCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
if err := c.cancelPrompt(cancelCtx, accepted.PromptID); err != nil {
|
|
// The original context error is more useful to the caller; the
|
|
// provider cleanup failure remains observable in the server log.
|
|
log.Printf("comfy prompt cleanup failed for %s: %v", accepted.PromptID, err)
|
|
}
|
|
}
|
|
}()
|
|
|
|
interval := c.PollInterval
|
|
if interval <= 0 {
|
|
interval = time.Second
|
|
}
|
|
timeout := c.PollTimeout
|
|
if timeout <= 0 {
|
|
timeout = 10 * time.Minute
|
|
}
|
|
deadline := time.NewTimer(timeout)
|
|
defer deadline.Stop()
|
|
ticker := time.NewTicker(interval)
|
|
defer ticker.Stop()
|
|
for {
|
|
image, done, err := c.historyImage(ctx, accepted.PromptID)
|
|
if err != nil {
|
|
return GeneratedImage{}, err
|
|
}
|
|
if done {
|
|
return image, nil
|
|
}
|
|
select {
|
|
case <-ctx.Done():
|
|
return GeneratedImage{}, ctx.Err()
|
|
case <-deadline.C:
|
|
return GeneratedImage{}, context.DeadlineExceeded
|
|
case <-ticker.C:
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c ComfyImageGenerationClient) cancelPrompt(ctx context.Context, promptID string) error {
|
|
if strings.TrimSpace(promptID) == "" {
|
|
return errors.New("comfy prompt id is empty")
|
|
}
|
|
payload, err := json.Marshal(map[string][]string{"delete": {promptID}})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/queue", bytes.NewReader(payload))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
request.Header.Set("Content-Type", "application/json")
|
|
response, err := c.client().Do(request)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode < 200 || response.StatusCode > 299 {
|
|
return fmt.Errorf("comfy prompt cleanup status %d", response.StatusCode)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c ComfyImageGenerationClient) historyImage(ctx context.Context, promptID string) (GeneratedImage, bool, error) {
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimRight(c.BaseURL, "/")+"/history/"+url.PathEscape(promptID), nil)
|
|
if err != nil {
|
|
return GeneratedImage{}, false, err
|
|
}
|
|
response, err := c.client().Do(request)
|
|
if err != nil {
|
|
return GeneratedImage{}, false, err
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode == http.StatusNotFound {
|
|
return GeneratedImage{}, false, nil
|
|
}
|
|
if response.StatusCode < 200 || response.StatusCode > 299 {
|
|
return GeneratedImage{}, false, fmt.Errorf("comfy history status %d", response.StatusCode)
|
|
}
|
|
var history map[string]struct {
|
|
Outputs map[string]struct {
|
|
Images []struct {
|
|
Filename string `json:"filename"`
|
|
Subfolder string `json:"subfolder"`
|
|
Type string `json:"type"`
|
|
} `json:"images"`
|
|
} `json:"outputs"`
|
|
}
|
|
if err := json.NewDecoder(io.LimitReader(response.Body, 4<<20)).Decode(&history); err != nil {
|
|
return GeneratedImage{}, false, err
|
|
}
|
|
record, ok := history[promptID]
|
|
if !ok {
|
|
return GeneratedImage{}, false, nil
|
|
}
|
|
for _, output := range record.Outputs {
|
|
for _, image := range output.Images {
|
|
if image.Filename == "" {
|
|
continue
|
|
}
|
|
query := url.Values{"filename": {image.Filename}, "subfolder": {image.Subfolder}, "type": {image.Type}}
|
|
imageRequest, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimRight(c.BaseURL, "/")+"/view?"+query.Encode(), nil)
|
|
if err != nil {
|
|
return GeneratedImage{}, false, err
|
|
}
|
|
imageResponse, err := c.client().Do(imageRequest)
|
|
if err != nil {
|
|
return GeneratedImage{}, false, err
|
|
}
|
|
data, readErr := io.ReadAll(io.LimitReader(imageResponse.Body, 32<<20))
|
|
imageResponse.Body.Close()
|
|
if readErr != nil {
|
|
return GeneratedImage{}, false, readErr
|
|
}
|
|
if imageResponse.StatusCode < 200 || imageResponse.StatusCode > 299 || len(data) == 0 {
|
|
return GeneratedImage{}, false, ErrEmptyResult
|
|
}
|
|
contentType := imageResponse.Header.Get("Content-Type")
|
|
if contentType == "" {
|
|
contentType = "image/png"
|
|
}
|
|
return GeneratedImage{DataBase64: encodeBase64(data), ContentType: contentType}, true, nil
|
|
}
|
|
}
|
|
return GeneratedImage{}, true, ErrEmptyResult
|
|
}
|
|
|
|
func imageSize(size string) (int, int) {
|
|
parts := strings.Split(strings.TrimSpace(size), "x")
|
|
if len(parts) == 2 {
|
|
var width, height int
|
|
if _, err := fmt.Sscanf(size, "%dx%d", &width, &height); err == nil && width >= 256 && height >= 256 && width <= 2048 && height <= 2048 {
|
|
return width, height
|
|
}
|
|
}
|
|
return 512, 512
|
|
}
|
|
|
|
// Describe accepts a video skill endpoint that performs its own bounded frame
|
|
// sampling. Frame extraction is provider-owned so the core does not acquire a
|
|
// media-processing dependency or hide a policy in orchestration code.
|
|
func (c VideoUnderstandingClient) Describe(ctx context.Context, model, filename, contentType string, data []byte, instruction string) (string, error) {
|
|
if len(data) == 0 {
|
|
return "", errors.New("video payload is empty")
|
|
}
|
|
body := &bytes.Buffer{}
|
|
writer := multipart.NewWriter(body)
|
|
header := make(textproto.MIMEHeader)
|
|
header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="file"; filename="%s"`, strings.ReplaceAll(filename, `"`, "")))
|
|
header.Set("Content-Type", contentType)
|
|
part, err := writer.CreatePart(header)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if _, err := part.Write(data); err != nil {
|
|
return "", err
|
|
}
|
|
_ = writer.WriteField("model", model)
|
|
_ = writer.WriteField("instruction", instruction)
|
|
if err := writer.Close(); err != nil {
|
|
return "", err
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/understand", body)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
request.Header.Set("Content-Type", writer.FormDataContentType())
|
|
response, err := c.client().Do(request)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode < 200 || response.StatusCode > 299 {
|
|
return "", fmt.Errorf("video provider status %d", response.StatusCode)
|
|
}
|
|
var result struct {
|
|
Text string `json:"text"`
|
|
}
|
|
if err := json.NewDecoder(io.LimitReader(response.Body, 2<<20)).Decode(&result); err != nil {
|
|
return "", err
|
|
}
|
|
if strings.TrimSpace(result.Text) == "" {
|
|
return "", ErrEmptyResult
|
|
}
|
|
return strings.TrimSpace(result.Text), nil
|
|
}
|
|
|
|
func (c SpeechClient) client() *http.Client {
|
|
if c.Client != nil {
|
|
return c.Client
|
|
}
|
|
return http.DefaultClient
|
|
}
|
|
|
|
func (c VisionClient) client() *http.Client {
|
|
if c.Client != nil {
|
|
return c.Client
|
|
}
|
|
return http.DefaultClient
|
|
}
|
|
|
|
func (c ImageGenerationClient) client() *http.Client {
|
|
if c.Client != nil {
|
|
return c.Client
|
|
}
|
|
return http.DefaultClient
|
|
}
|
|
|
|
func (c ComfyImageGenerationClient) client() *http.Client {
|
|
if c.Client != nil {
|
|
return c.Client
|
|
}
|
|
return http.DefaultClient
|
|
}
|
|
|
|
func (c VideoUnderstandingClient) client() *http.Client {
|
|
if c.Client != nil {
|
|
return c.Client
|
|
}
|
|
return http.DefaultClient
|
|
}
|
|
|
|
func encodeBase64(data []byte) string {
|
|
return base64.StdEncoding.EncodeToString(data)
|
|
}
|