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")
	}
}

// TestImageEditByReferenceUsesTxt2ImgExtraImages pins the instruction-edit
// wire shape. It is a DIFFERENT endpoint and a DIFFERENT field from img2img,
// and the difference is not cosmetic: measured against FLUX.1-Kontext on
// 2026-07-30, the same prompt sent as init_images left the thing it was told
// to change untouched and drifted everything else, while extra_images changed
// exactly what was asked and left the rest of the frame numerically
// unchanged. Routing a reference edit down the img2img path would look like
// a working call and silently produce the wrong picture.
func TestImageEditByReferenceUsesTxt2ImgExtraImages(t *testing.T) {
	var gotPath string
	var gotBody map[string]any
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = 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("imagegen-flux-kontext")
	ed := im.(imagegen.Editor)

	ref := editInit(t)
	if _, err := ed.Edit(context.Background(), imagegen.EditRequest{
		Prompt: "make the sign read OPEN",
		// Init/Mask/Strength are set and must be IGNORED — they describe a
		// pipeline this model does not run.
		Init:     ref,
		Mask:     ref,
		Strength: func() *float64 { s := 0.75; return &s }(),
	}, imagegen.WithEditRefImages(ref)); err != nil {
		t.Fatalf("reference edit: %v", err)
	}

	if gotPath != "/sdapi/v1/txt2img" {
		t.Errorf("path = %q, want /sdapi/v1/txt2img (there is no init latent to denoise)", gotPath)
	}
	extra, ok := gotBody["extra_images"].([]any)
	if !ok || len(extra) != 1 {
		t.Fatalf("extra_images = %v, want the one reference image", gotBody["extra_images"])
	}
	if _, present := gotBody["init_images"]; present {
		t.Error("init_images must NOT be sent on the reference path — it re-noises the picture")
	}
	if _, present := gotBody["denoising_strength"]; present {
		t.Error("denoising_strength must NOT be sent on the reference path")
	}
	if _, present := gotBody["mask"]; present {
		t.Error("mask must NOT be sent on the reference path")
	}
}

// TestImageEditByReferenceRejectsEmptyRefs guards the case that would
// otherwise silently become a plain txt2img: a reference edit whose only
// reference carries no bytes has nothing to condition on, and rendering the
// prompt from scratch is not what the caller asked for.
func TestImageEditByReferenceRejectsEmptyRefs(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_, _ = w.Write([]byte(`{"images":["` + onePixelPNG + `"]}`))
	}))
	defer srv.Close()

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

	_, err := ed.Edit(context.Background(), imagegen.EditRequest{Prompt: "anything"},
		imagegen.WithEditRefImages(imagegen.Image{MIME: "image/png"}))
	if !errors.Is(err, llm.ErrUnsupported) {
		t.Fatalf("err = %v, want ErrUnsupported for an all-empty reference set", err)
	}
}