majordomo/provider/llamaswap/edit_test.go

package llamaswap

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"

	"gitea.stevedudenhoeffer.com/steve/majordomo/imagegen"
	"gitea.stevedudenhoeffer.com/steve/majordomo/llm"
)

func editInit(t *testing.T) imagegen.Image {
	t.Helper()
	raw, err := base64.StdEncoding.DecodeString(onePixelPNG)
	if err != nil {
		t.Fatalf("decode fixture: %v", err)
	}
	return imagegen.Image{MIME: "image/png", Data: raw}
}

func TestImageEdit(t *testing.T) {
	var gotBody map[string]any
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/sdapi/v1/img2img" {
			t.Errorf("path = %q", r.URL.Path)
		}
		_ = json.NewDecoder(r.Body).Decode(&gotBody)
		_, _ = w.Write([]byte(`{"images":["` + onePixelPNG + `"]}`))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	im, _ := p.ImageModel("sd")
	ed, ok := im.(imagegen.Editor)
	if !ok {
		t.Fatal("imageModel does not implement imagegen.Editor")
	}

	res, err := ed.Edit(context.Background(),
		imagegen.EditRequest{Prompt: "make it night", Init: editInit(t)},
		imagegen.WithEditStrength(0.6),
	)
	if err != nil {
		t.Fatalf("Edit: %v", err)
	}
	if len(res.Images) != 1 || res.Images[0].MIME != "image/png" {
		t.Fatalf("images = %+v", res.Images)
	}

	if gotBody["model"] != "sd" || gotBody["prompt"] != "make it night" {
		t.Errorf("model/prompt = %v/%v", gotBody["model"], gotBody["prompt"])
	}
	inits, ok := gotBody["init_images"].([]any)
	if !ok || len(inits) != 1 || inits[0] != onePixelPNG {
		t.Errorf("init_images = %v, want the b64 fixture", gotBody["init_images"])
	}
	if gotBody["denoising_strength"] != 0.6 {
		t.Errorf("denoising_strength = %v, want 0.6", gotBody["denoising_strength"])
	}
}

func TestImageEditOmitsUnsetOverrides(t *testing.T) {
	var gotBody map[string]any
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewDecoder(r.Body).Decode(&gotBody)
		_, _ = w.Write([]byte(`{"images":["` + onePixelPNG + `"]}`))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	im, _ := p.ImageModel("sd")
	ed := im.(imagegen.Editor)
	if _, err := ed.Edit(context.Background(), imagegen.EditRequest{Prompt: "x", Init: editInit(t)}); err != nil {
		t.Fatalf("Edit: %v", err)
	}
	for _, k := range []string{"denoising_strength", "steps", "cfg_scale", "negative_prompt", "sample_method", "seed", "width", "height"} {
		if v, ok := gotBody[k]; ok {
			t.Errorf("unset request sent %q = %v, want omitted", k, v)
		}
	}
}

func TestImageEditValidation(t *testing.T) {
	p := New(WithBaseURL("http://example.invalid"))
	im, _ := p.ImageModel("sd")
	ed := im.(imagegen.Editor)

	cases := []struct {
		name string
		req  imagegen.EditRequest
	}{
		{"empty prompt", imagegen.EditRequest{Prompt: " ", Init: imagegen.Image{Data: []byte{1}}}},
		{"missing init", imagegen.EditRequest{Prompt: "x"}},
		{"negative N", imagegen.EditRequest{Prompt: "x", Init: imagegen.Image{Data: []byte{1}}, N: -1}},
	}
	for _, tc := range cases {
		if _, err := ed.Edit(context.Background(), tc.req); !errors.Is(err, llm.ErrUnsupported) {
			t.Errorf("%s: err = %v, want ErrUnsupported", tc.name, err)
		}
	}

	bad := 1.5
	if _, err := ed.Edit(context.Background(), imagegen.EditRequest{Prompt: "x", Init: imagegen.Image{Data: []byte{1}}, Strength: &bad}); !errors.Is(err, llm.ErrUnsupported) {
		t.Errorf("out-of-range strength: err = %v, want ErrUnsupported", err)
	}
}

func TestImageEditWithMask(t *testing.T) {
	var gotBody map[string]any
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewDecoder(r.Body).Decode(&gotBody)
		_, _ = w.Write([]byte(`{"images":["` + onePixelPNG + `"]}`))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	im, _ := p.ImageModel("sd")
	ed := im.(imagegen.Editor)

	mask := imagegen.Image{MIME: "image/png", Data: []byte{0xDE, 0xAD}}
	if _, err := ed.Edit(context.Background(),
		imagegen.EditRequest{Prompt: "replace the sky", Init: editInit(t)},
		imagegen.WithEditMask(mask),
	); err != nil {
		t.Fatalf("Edit: %v", err)
	}
	if got := gotBody["mask"]; got != base64.StdEncoding.EncodeToString(mask.Data) {
		t.Errorf("mask = %v, want the b64 mask", got)
	}
}

func TestImageEditWithoutMaskOmitsField(t *testing.T) {
	var gotBody map[string]any
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewDecoder(r.Body).Decode(&gotBody)
		_, _ = w.Write([]byte(`{"images":["` + onePixelPNG + `"]}`))
	}))
	defer srv.Close()

	p := New(WithBaseURL(srv.URL), WithHTTPClient(srv.Client()))
	im, _ := p.ImageModel("sd")
	ed := im.(imagegen.Editor)
	if _, err := ed.Edit(context.Background(),
		imagegen.EditRequest{Prompt: "p", Init: editInit(t)}); err != nil {
		t.Fatalf("Edit: %v", err)
	}
	if _, ok := gotBody["mask"]; ok {
		t.Error("mask field sent for unmasked edit; want omitted")
	}
}