majordomo/provider/llamaswap/image.go

package llamaswap

import (
	"context"
	"encoding/base64"
	"fmt"
	"net/http"
	"strconv"
	"strings"

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

// ImageModel implements imagegen.Provider, binding an image-generation model
// served by llama-swap (routed to a stable-diffusion.cpp upstream). The id is
// passed through verbatim and selects which upstream llama-swap loads.
func (p *Provider) ImageModel(id string, opts ...imagegen.ModelOption) (imagegen.Model, error) {
	if err := p.requireBaseURL(); err != nil {
		return nil, err
	}
	_ = imagegen.ApplyModelOptions(opts)
	return &imageModel{p: p, id: id}, nil
}

type imageModel struct {
	p  *Provider
	id string
}

// txt2imgRequest is the stable-diffusion.cpp sd-server A1111 request shape
// (POST /sdapi/v1/txt2img). We use this endpoint rather than the OpenAI
// /v1/images/generations one because that endpoint IGNORES `seed` on this
// sd-server build — every render of a given prompt comes back byte-identical,
// so a batch of N collapses to one image. /sdapi/v1/txt2img honours `seed`,
// giving real variety. llama-swap still routes by the `model` field in the
// body. Optional fields are pointers/omitempty so an unset value falls back to
// the model's baked default (the per-model --steps/--cfg-scale/etc. flags).
type txt2imgRequest struct {
	Model          string   `json:"model"`
	Prompt         string   `json:"prompt"`
	NegativePrompt string   `json:"negative_prompt,omitempty"`
	Seed           *int64   `json:"seed,omitempty"`
	Steps          *int     `json:"steps,omitempty"`
	CFGScale       *float64 `json:"cfg_scale,omitempty"`
	Width          *int     `json:"width,omitempty"`
	Height         *int     `json:"height,omitempty"`
	SampleMethod   string   `json:"sample_method,omitempty"`
	BatchCount     int      `json:"batch_count,omitempty"`
}

type txt2imgResponse struct {
	Images []string `json:"images"`
}

// sdWire validates the generation knobs shared by Generate and Edit and
// builds the common txt2img wire fields. verb labels validation errors.
func (m *imageModel) sdWire(verb, prompt, negativePrompt, sampler, size string, seed *int64, steps *int, cfgScale *float64, n int) (txt2imgRequest, error) {
	if strings.TrimSpace(prompt) == "" {
		return txt2imgRequest{}, fmt.Errorf("%w: image %s requires a prompt", llm.ErrUnsupported, verb)
	}
	if n < 0 {
		return txt2imgRequest{}, fmt.Errorf("%w: image count N must be >= 0, got %d", llm.ErrUnsupported, n)
	}
	width, height, err := parseSize(size)
	if err != nil {
		return txt2imgRequest{}, fmt.Errorf("%w: %v", llm.ErrUnsupported, err)
	}
	return txt2imgRequest{
		Model:          m.id,
		Prompt:         prompt,
		NegativePrompt: negativePrompt,
		Seed:           seed,
		Steps:          steps,
		CFGScale:       cfgScale,
		Width:          width,
		Height:         height,
		SampleMethod:   sampler,
		BatchCount:     n,
	}, nil
}

// Generate implements imagegen.Model via POST {base}/sdapi/v1/txt2img.
func (m *imageModel) Generate(ctx context.Context, req imagegen.Request, opts ...imagegen.Option) (*imagegen.Result, error) {
	req = req.Apply(opts...)
	wire, err := m.sdWire("generation", req.Prompt, req.NegativePrompt, req.Sampler, req.Size, req.Seed, req.Steps, req.CFGScale, req.N)
	if err != nil {
		return nil, err
	}

	var resp txt2imgResponse
	if err := m.p.doJSON(ctx, http.MethodPost, "/sdapi/v1/txt2img", m.id, &wire, &resp); err != nil {
		return nil, err
	}
	return decodeImages(m.p.name, m.id, &resp)
}

// decodeImages converts an SDAPI response's base64 images into an
// imagegen.Result, erroring when nothing decodable came back.
func decodeImages(provider, model string, resp *txt2imgResponse) (*imagegen.Result, error) {
	out := &imagegen.Result{Raw: resp}
	for i, b64 := range resp.Images {
		if b64 == "" {
			continue
		}
		raw, err := base64.StdEncoding.DecodeString(b64)
		if err != nil {
			return nil, fmt.Errorf("llama-swap: decode image %d: %w", i, err)
		}
		out.Images = append(out.Images, llm.ImagePart{MIME: sniffImageMIME(raw), Data: raw})
	}
	if len(out.Images) == 0 {
		return nil, &llm.APIError{
			Provider: provider,
			Model:    model,
			Message:  "image response contained no images",
		}
	}
	return out, nil
}

// img2imgRequest is the stable-diffusion.cpp sd-server A1111 request shape
// (POST /sdapi/v1/img2img): txt2img's fields plus the init image(s) and
// denoising strength. Same endpoint-family choice as txt2img — the OpenAI
// /v1/images/edits route is multipart and drops `seed` on this sd-server
// build, while the SDAPI shape reuses doJSON and keeps seed parity.
type img2imgRequest struct {
	txt2imgRequest
	InitImages        []string `json:"init_images"`
	DenoisingStrength *float64 `json:"denoising_strength,omitempty"`
	// Mask enables inpainting: base64 image, white = repaint, black = keep
	// (sd-server also accepts a data URL; plain base64 keeps symmetry with
	// init_images). sd-server has no mask_blur/inpaint_full_res — callers
	// wanting soft edges pre-feather the mask.
	Mask string `json:"mask,omitempty"`
}

// Edit implements imagegen.Editor via POST {base}/sdapi/v1/img2img.
func (m *imageModel) Edit(ctx context.Context, req imagegen.EditRequest, opts ...imagegen.EditOption) (*imagegen.Result, error) {
	req = req.Apply(opts...)
	if len(req.Init.Data) == 0 {
		return nil, fmt.Errorf("%w: image edit requires an init image", llm.ErrUnsupported)
	}
	if req.Strength != nil && (*req.Strength < 0 || *req.Strength > 1) {
		return nil, fmt.Errorf("%w: edit strength must be in [0,1], got %g", llm.ErrUnsupported, *req.Strength)
	}
	base, err := m.sdWire("edit", req.Prompt, req.NegativePrompt, req.Sampler, req.Size, req.Seed, req.Steps, req.CFGScale, req.N)
	if err != nil {
		return nil, err
	}
	wire := img2imgRequest{
		txt2imgRequest:    base,
		InitImages:        []string{base64.StdEncoding.EncodeToString(req.Init.Data)},
		DenoisingStrength: req.Strength,
	}
	if len(req.Mask.Data) > 0 {
		wire.Mask = base64.StdEncoding.EncodeToString(req.Mask.Data)
	}

	var resp txt2imgResponse
	if err := m.p.doJSON(ctx, http.MethodPost, "/sdapi/v1/img2img", m.id, &wire, &resp); err != nil {
		return nil, err
	}
	return decodeImages(m.p.name, m.id, &resp)
}

// parseSize splits a "WxH" string into width/height pointers. "" yields
// (nil, nil) so the model's own default resolution applies.
func parseSize(size string) (*int, *int, error) {
	size = strings.TrimSpace(size)
	if size == "" {
		return nil, nil, nil
	}
	parts := strings.SplitN(strings.ToLower(size), "x", 2)
	if len(parts) != 2 {
		return nil, nil, fmt.Errorf("invalid size %q (want WxH)", size)
	}
	w, err1 := strconv.Atoi(strings.TrimSpace(parts[0]))
	h, err2 := strconv.Atoi(strings.TrimSpace(parts[1]))
	if err1 != nil || err2 != nil || w <= 0 || h <= 0 {
		return nil, nil, fmt.Errorf("invalid size %q (want WxH)", size)
	}
	return &w, &h, nil
}

// sniffImageMIME identifies the image format from its leading bytes, defaulting
// to image/png (stable-diffusion.cpp emits PNG) when detection is inconclusive.
func sniffImageMIME(data []byte) string {
	mime := http.DetectContentType(data)
	if !strings.HasPrefix(mime, "image/") {
		return "image/png"
	}
	return mime
}