diff --git a/adapters/multimodal/client.go b/adapters/multimodal/client.go index 18dce42..d187a74 100644 --- a/adapters/multimodal/client.go +++ b/adapters/multimodal/client.go @@ -11,7 +11,9 @@ import ( "mime/multipart" "net/http" "net/textproto" + "net/url" "strings" + "time" ) var ErrEmptyResult = errors.New("multimodal provider returned an empty result") @@ -42,6 +44,20 @@ type GeneratedImage struct { 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") @@ -216,6 +232,155 @@ func (c ImageGenerationClient) Generate(ctx context.Context, model, prompt, size 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 + } + + 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) 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. @@ -286,6 +451,13 @@ func (c ImageGenerationClient) client() *http.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 diff --git a/adapters/multimodal/client_test.go b/adapters/multimodal/client_test.go index 4068320..6f1d981 100644 --- a/adapters/multimodal/client_test.go +++ b/adapters/multimodal/client_test.go @@ -8,6 +8,7 @@ import ( "net/http/httptest" "strings" "testing" + "time" ) func TestSpeechClientSendsMultipartAndReturnsTranscript(t *testing.T) { @@ -130,3 +131,39 @@ func TestVideoUnderstandingClientSendsMultipartAndReturnsDescription(t *testing. t.Fatalf("description=%q err=%v", text, err) } } + +func TestComfyImageGenerationClientRunsWorkflowAndReturnsImage(t *testing.T) { + polls := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/prompt": + var payload struct { + Prompt map[string]any `json:"prompt"` + } + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatal(err) + } + if payload.Prompt["6"].(map[string]any)["inputs"].(map[string]any)["text"] != "red bird" { + t.Fatalf("prompt was not forwarded: %+v", payload.Prompt["6"]) + } + _ = json.NewEncoder(w).Encode(map[string]string{"prompt_id": "test-prompt"}) + case "/history/test-prompt": + polls++ + if polls == 1 { + _, _ = w.Write([]byte(`{}`)) + return + } + _ = json.NewEncoder(w).Encode(map[string]any{"test-prompt": map[string]any{"outputs": map[string]any{"9": map[string]any{"images": []any{map[string]string{"filename": "out.png", "subfolder": "", "type": "output"}}}}}}) + case "/view": + w.Header().Set("Content-Type", "image/png") + _, _ = w.Write([]byte("png")) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + image, err := (ComfyImageGenerationClient{BaseURL: server.URL, Client: server.Client(), PollInterval: time.Millisecond, PollTimeout: time.Second}).Generate(context.Background(), "", "red bird", "256x256") + if err != nil || image.DataBase64 != "cG5n" || image.ContentType != "image/png" { + t.Fatalf("image=%+v err=%v", image, err) + } +} diff --git a/cmd/cognitive-worker/main.go b/cmd/cognitive-worker/main.go index 74f0557..e6f76e4 100644 --- a/cmd/cognitive-worker/main.go +++ b/cmd/cognitive-worker/main.go @@ -171,15 +171,20 @@ func main() { mux.HandleFunc("POST /v1/advisories", gateway.advisoryHandler) mux.HandleFunc("POST /v1/interpretations", gateway.interpretationHandler) mux.HandleFunc("POST /v1/task-evaluations/execute", gateway.taskEvaluationHandler) + imageGenerator := multimodal.ImageGenerator(multimodal.ImageGenerationClient{BaseURL: os.Getenv("LOCAL_IMAGE_PROVIDER_URL"), Client: &http.Client{Timeout: 300 * time.Second}}) + if strings.EqualFold(strings.TrimSpace(os.Getenv("LOCAL_IMAGE_PROVIDER_KIND")), "comfyui") { + imageGenerator = multimodal.ComfyImageGenerationClient{BaseURL: os.Getenv("LOCAL_IMAGE_PROVIDER_URL"), Client: &http.Client{Timeout: 300 * time.Second}, PollTimeout: 15 * time.Minute} + } multimodalGateway := multimodalGateway{ - speech: multimodal.SpeechClient{BaseURL: envOr("LOCAL_SPEECH_PROVIDER_URL", "http://llm-local-speech:8080"), Client: &http.Client{Timeout: 120 * time.Second}}, - vision: multimodal.VisionClient{BaseURL: envOr("LOCAL_VISION_PROVIDER_URL", "http://llm-local-vision:8080/v1"), Client: &http.Client{Timeout: 180 * time.Second}}, - imageGenerator: multimodal.ImageGenerationClient{BaseURL: os.Getenv("LOCAL_IMAGE_PROVIDER_URL"), Client: &http.Client{Timeout: 300 * time.Second}}, - video: multimodal.VideoUnderstandingClient{BaseURL: os.Getenv("LOCAL_VIDEO_PROVIDER_URL"), Client: &http.Client{Timeout: 300 * time.Second}}, - model: envOr("LOCAL_VISION_MODEL", "/models/gemma-4-26B-A4B-it-UD-Q4_K_M.gguf"), - imageModel: os.Getenv("LOCAL_IMAGE_MODEL"), - videoModel: os.Getenv("LOCAL_VIDEO_MODEL"), - now: time.Now, + speech: multimodal.SpeechClient{BaseURL: envOr("LOCAL_SPEECH_PROVIDER_URL", "http://llm-local-speech:8080"), Client: &http.Client{Timeout: 120 * time.Second}}, + vision: multimodal.VisionClient{BaseURL: envOr("LOCAL_VISION_PROVIDER_URL", "http://llm-local-vision:8080/v1"), Client: &http.Client{Timeout: 180 * time.Second}}, + imageGenerator: imageGenerator, + imageProviderURL: os.Getenv("LOCAL_IMAGE_PROVIDER_URL"), + video: multimodal.VideoUnderstandingClient{BaseURL: os.Getenv("LOCAL_VIDEO_PROVIDER_URL"), Client: &http.Client{Timeout: 300 * time.Second}}, + model: envOr("LOCAL_VISION_MODEL", "/models/gemma-4-26B-A4B-it-UD-Q4_K_M.gguf"), + imageModel: os.Getenv("LOCAL_IMAGE_MODEL"), + videoModel: os.Getenv("LOCAL_VIDEO_MODEL"), + now: time.Now, } mux.HandleFunc("POST /v1/multimodal/transcribe", multimodalGateway.handler(capability.MultimodalTranscribe)) mux.HandleFunc("POST /v1/multimodal/understand", multimodalGateway.handler(capability.MultimodalUnderstand)) diff --git a/cmd/cognitive-worker/multimodal.go b/cmd/cognitive-worker/multimodal.go index aea8baf..7d20de4 100644 --- a/cmd/cognitive-worker/multimodal.go +++ b/cmd/cognitive-worker/multimodal.go @@ -15,14 +15,15 @@ import ( const maxMultimodalUploadBytes = 12 << 20 type multimodalGateway struct { - speech multimodal.SpeechClient - vision multimodal.VisionClient - imageGenerator multimodal.ImageGenerationClient - video multimodal.VideoUnderstandingClient - model string - imageModel string - videoModel string - now func() time.Time + speech multimodal.SpeechClient + vision multimodal.VisionClient + imageGenerator multimodal.ImageGenerator + imageProviderURL string + video multimodal.VideoUnderstandingClient + model string + imageModel string + videoModel string + now func() time.Time } func (g multimodalGateway) handler(operation capability.MultimodalOperation) http.HandlerFunc { @@ -33,7 +34,7 @@ func (g multimodalGateway) handler(operation capability.MultimodalOperation) htt return } if operation == capability.MultimodalGenerateImage { - if strings.TrimSpace(g.imageGenerator.BaseURL) == "" { + if strings.TrimSpace(g.imageProviderURL) == "" { http.Error(w, "local image generation capability unavailable", http.StatusBadGateway) return }