diff --git a/internal/server/artifact_upload_test.go b/internal/server/artifact_upload_test.go index e0ef795..3c05168 100644 --- a/internal/server/artifact_upload_test.go +++ b/internal/server/artifact_upload_test.go @@ -130,3 +130,42 @@ func TestArtifactUploadCreatesReadableArtifactForSourceArtifactID(t *testing.T) t.Fatal("expected resolved reference content to match uploaded bytes") } } + +func TestResolveImageToImageReferenceAcceptsCanonicalImageURL(t *testing.T) { + cfg := Config{ + SQLitePath: filepath.Join(t.TempDir(), "popiart.db"), + SkillhubDir: makeEmptySkillhub(t), + SessionSecret: "test-secret", + } + server, err := NewWithConfig(cfg) + if err != nil { + t.Fatalf("NewWithConfig: %v", err) + } + + imageBytes := tinyPNG(t) + refSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "image/png") + _, _ = w.Write(imageBytes) + })) + defer refSrv.Close() + + editRef, err := server.resolveImageToImageReference(context.Background(), &job{ + UserID: "user_test", + UpstreamKey: "sk-upstream", + SkillID: "popiskill-image-img2img-basic-v1", + }, map[string]any{ + "image": refSrv.URL + "/reference.png", + }) + if err != nil { + t.Fatalf("resolveImageToImageReference: %v", err) + } + if editRef.URL != refSrv.URL+"/reference.png" { + t.Fatalf("expected resolved URL to match input image, got %q", editRef.URL) + } + if editRef.ContentType != "image/png" { + t.Fatalf("expected image/png, got %q", editRef.ContentType) + } + if !bytes.Equal(editRef.Content, imageBytes) { + t.Fatal("expected resolved canonical image content to match downloaded bytes") + } +} diff --git a/internal/server/routes_test.go b/internal/server/routes_test.go index e9ec763..abe735b 100644 --- a/internal/server/routes_test.go +++ b/internal/server/routes_test.go @@ -31,8 +31,23 @@ func TestInferRouteKeyForModelRecognizesViduAsVideo(t *testing.T) { } for modelID, want := range cases { - if got := inferRouteKeyForModel(modelID); got != want { + if got := inferRouteKeyForModel(modelID, nil); got != want { t.Fatalf("inferRouteKeyForModel(%q) = %q, want %q", modelID, got, want) } } } + +func TestInferRouteKeyForModelRecognizesImageEditInputs(t *testing.T) { + cases := []map[string]any{ + {"image": "https://example.com/reference.jpg"}, + {"image_url": "https://example.com/reference.jpg"}, + {"reference_image_url": "https://example.com/reference.jpg"}, + {"source_artifact_id": "art_123"}, + } + + for _, input := range cases { + if got := inferRouteKeyForModel("seedream-4-5-251128", input); got != "image.img2img" { + t.Fatalf("inferRouteKeyForModel(seedream, %#v) = %q, want image.img2img", input, got) + } + } +} diff --git a/internal/server/server.go b/internal/server/server.go index 82173f1..bf3d420 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -1046,7 +1046,7 @@ func (s *Server) handleModelsInfer(w http.ResponseWriter, r *http.Request) { writeError(w, http.StatusBadRequest, "VALIDATION_ERROR", "model_id is required", nil) return } - routeKey := inferRouteKeyForModel(req.ModelID) + routeKey := inferRouteKeyForModel(req.ModelID, req.Input) record, statusCode, err := s.store.createJob( "", routeKey, @@ -1334,10 +1334,20 @@ func (s *Server) resolveModelID(routeKey, projectID string) string { } } -func inferRouteKeyForModel(modelID string) string { +func inferRouteKeyForModel(modelID string, input map[string]any) string { if modelType, _ := classifyModelIDFallback(modelID); modelType == "video" { return "video.image2video" } + if input != nil { + if ref := strings.TrimSpace(stringValue( + input["source_artifact_id"], + input["image"], + input["reference_image_url"], + input["image_url"], + )); ref != "" { + return "image.img2img" + } + } return "image.text2image" } @@ -1371,6 +1381,7 @@ func (s *Server) resolveImageToImageReference(ctx context.Context, record *job, } refURL := strings.TrimSpace(stringValue( + input["image"], input["reference_image_url"], input["image_url"], )) @@ -1540,7 +1551,7 @@ func buildGenericImageToImagePrompt(input map[string]any, prompt string) string if prompt != "" { parts = append(parts, prompt) } - if refURL := strings.TrimSpace(stringValue(input["reference_image_url"], input["image_url"])); refURL != "" { + if refURL := strings.TrimSpace(stringValue(input["image"], input["reference_image_url"], input["image_url"])); refURL != "" { parts = append(parts, "preserve the same main subject identity and key visual traits from the reference image") } if aspectRatio := strings.TrimSpace(stringValue(input["aspect_ratio"])); aspectRatio != "" {