llmgw/internal/router/images.go

package router

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/comfy"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
	"gitea.mynodes.xyz/ilker/llmgw/internal/version"
	"gitea.mynodes.xyz/ilker/llmgw/internal/web"
)

// /v1/images/generations, OpenAI's shape, made by ComfyUI: the request goes
// through the same checks as any other — key, allowlist, limits, admission,
// guardrails on its prompt — and only at the replica is it turned into a
// workflow, queued and waited on, and its pictures sent back as base64.

const imagesPath = "/v1/images/generations"

// imageRequest is what a request for pictures asks. Beyond OpenAI's fields,
// those a diffusion model has and DALL·E does not: seed, steps,
// negative_prompt and guidance.
type imageRequest struct {
	Prompt         string   `json:"prompt"`
	N              int      `json:"n"`
	Size           string   `json:"size"`
	ResponseFormat string   `json:"response_format"`
	Seed           *int64   `json:"seed"`
	Steps          int      `json:"steps"`
	NegativePrompt string   `json:"negative_prompt"`
	Guidance       *float64 `json:"guidance"`

	width, height int
}

// The bounds of a request. A FLUX latent is 1/16 of the picture, so a side
// that is not a multiple of 16 is not one it can make; the rest keep one
// request from holding a card for the length of a film.
const (
	maxImages   = 4
	minSide     = 256
	maxSide     = 2048
	maxSteps    = 100
	defaultSide = 1024
)

func parseImageRequest(body []byte) (imageRequest, error) {
	var in imageRequest
	if err := json.Unmarshal(body, &in); err != nil {
		return in, err
	}
	if strings.TrimSpace(in.Prompt) == "" {
		return in, errors.New("prompt is required")
	}
	switch {
	case in.N == 0:
		in.N = 1
	case in.N < 0 || in.N > maxImages:
		return in, fmt.Errorf("n is between 1 and %d", maxImages)
	}
	switch in.ResponseFormat {
	case "", "b64_json":
	case "url":
		// A URL would have to be served by somebody and kept for some time,
		// and a picture made from a prompt is the prompt's owner's: it is
		// not kept here.
		return in, errors.New(`response_format "url" is not offered: the picture comes back as b64_json and is kept nowhere`)
	default:
		return in, fmt.Errorf("response_format %q is not b64_json", in.ResponseFormat)
	}
	in.width, in.height = defaultSide, defaultSide
	if s := strings.TrimSpace(in.Size); s != "" && s != "auto" {
		w, h, ok := strings.Cut(strings.ToLower(s), "x")
		wi, err1 := strconv.Atoi(w)
		hi, err2 := strconv.Atoi(h)
		if !ok || err1 != nil || err2 != nil {
			return in, fmt.Errorf("size %q is not WIDTHxHEIGHT", s)
		}
		in.width, in.height = wi, hi
	}
	for _, side := range []int{in.width, in.height} {
		if side < minSide || side > maxSide || side%16 != 0 {
			return in, fmt.Errorf("size %dx%d: each side is a multiple of 16 from %d to %d", in.width, in.height, minSide, maxSide)
		}
	}
	if in.Steps < 0 || in.Steps > maxSteps {
		return in, fmt.Errorf("steps is between 1 and %d", maxSteps)
	}
	if in.Guidance != nil && (*in.Guidance <= 0 || *in.Guidance > 30) {
		return in, errors.New("guidance is above 0 and at most 30")
	}
	return in, nil
}

// imageRouteRefusal says why a request cannot be answered at the path it was
// sent to, before it takes a place in a queue: a picture asked of a model
// that writes text, or text of one that makes pictures. Empty when it can,
// and for a model the catalog does not know, which a hand-configured
// endpoint may serve as it likes.
func (r *Router) imageRouteRefusal(path, model string, body []byte) string {
	m, known := r.modelMap()[model]
	if path == imagesPath {
		if known && !m.MakesImages() {
			return fmt.Sprintf("%s does not make pictures", m.Name)
		}
		if _, err := parseImageRequest(body); err != nil {
			return err.Error()
		}
		return ""
	}
	if known && m.ServesComfyUI() {
		return fmt.Sprintf("%s makes pictures: ask for them at %s", m.Name, imagesPath)
	}
	return ""
}

// generateImages answers a request at a ComfyUI replica. It returns what the
// proxy would: a failure of this replica that another may not have is
// retryable, and nothing is written until the pictures are all in hand.
func (r *Router) generateImages(ctx context.Context, w http.ResponseWriter, req *parsedRequest, ep endpoint.Info, body []byte) (proxy.Outcome, error) {
	if req.Path != imagesPath {
		return proxy.Outcome{}, fmt.Errorf("%s is a ComfyUI server: it answers %s only", ep.ID, imagesPath)
	}
	in, err := parseImageRequest(body)
	if err != nil {
		return proxy.Outcome{}, err
	}
	var upstream struct {
		Model string `json:"model"`
	}
	_ = json.Unmarshal(body, &upstream)
	tpl, err := comfy.Load(ep.Workflow)
	if err != nil {
		return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err)
	}
	seed := comfy.RandomSeed()
	if in.Seed != nil {
		seed = *in.Seed
	}
	p := comfy.Params{
		Checkpoint: upstream.Model, Prompt: in.Prompt, NegativePrompt: in.NegativePrompt,
		Width: in.width, Height: in.height, N: in.N, Seed: seed, Steps: in.Steps,
	}
	if in.Guidance != nil {
		p.Guidance = *in.Guidance
	}
	workflow, err := tpl.Fill(p)
	if err != nil {
		return proxy.Outcome{}, fmt.Errorf("%s: %w", ep.ID, err)
	}

	if t := r.cfg.Timeouts.Total; t > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, t)
		defer cancel()
	}
	start := time.Now()
	images, err := r.comfy.Generate(ctx, ep.URL, ep.APIKey, workflow)
	if err != nil {
		var ce *comfy.Error
		switch {
		case errors.As(err, &ce):
			return proxy.Outcome{Retryable: ce.Retryable}, err
		case errors.Is(err, context.Canceled) && ctxCauseIsClient(ctx):
			return proxy.Outcome{}, proxy.ErrClientGone
		case errors.Is(err, context.DeadlineExceeded):
			return proxy.Outcome{}, fmt.Errorf("%s made no picture within %s", ep.ID, r.cfg.Timeouts.Total)
		}
		return proxy.Outcome{}, err
	}

	data := make([]map[string]any, 0, len(images))
	for _, im := range images {
		data = append(data, map[string]any{"b64_json": base64.StdEncoding.EncodeToString(im.Data)})
	}
	out, err := json.Marshal(map[string]any{"created": time.Now().Unix(), "data": data})
	if err != nil {
		return proxy.Outcome{}, err
	}
	// The seed that made them, so a picture liked can be made again.
	w.Header().Set("x-llmgw-seed", strconv.FormatInt(seed, 10))
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write(out)
	took := time.Since(start)
	return proxy.Outcome{Wrote: true, Status: http.StatusOK, Bytes: int64(len(out)), TTFT: took, Decode: took,
		Images: len(images)}, nil
}

// ctxCauseIsClient reports whether a context ended because the client went
// away, rather than because a deadline set here passed.
func ctxCauseIsClient(ctx context.Context) bool {
	return !errors.Is(context.Cause(ctx), context.DeadlineExceeded)
}

// registerImagesPage mounts the page that makes pictures, beside the chat
// and for the same reasons: a model tried without a client, with the
// browser holding a session and never a key, and the request sent down the
// one path every API request takes.
func (r *Router) registerImagesPage(mux *http.ServeMux, pages *web.AccountPages) {
	pages.SetImages(func() bool {
		for _, m := range r.modelMap() {
			if m.MakesImages() {
				return true
			}
		}
		return false
	})
	mux.HandleFunc("GET /images", pages.Guard(func(w http.ResponseWriter, req *http.Request, user store.User) {
		contexts, _ := r.chatContexts(req.Context(), user)
		current, err := r.pickContext(req.Context(), user, req.URL.Query().Get("context"))
		if err != nil {
			current = ""
		}
		if current == "" && !r.personalOff {
			current = personalContext
		}
		web.Images(user, r.modelsOf(req, user, true), contexts, current, version.Version)(w, req)
	}))
	mux.HandleFunc("POST /images/generate", pages.Guard(r.imagesSend))
}

// imagesSend makes pictures on behalf of a signed-in person.
func (r *Router) imagesSend(w http.ResponseWriter, req *http.Request, user store.User) {
	body, err := io.ReadAll(io.LimitReader(req.Body, r.cfg.MaxBodyBytes))
	if err != nil {
		http.Error(w, "could not read the request", http.StatusBadRequest)
		return
	}
	var in struct {
		Model          string `json:"model"`
		Prompt         string `json:"prompt"`
		Size           string `json:"size"`
		N              int    `json:"n"`
		Seed           *int64 `json:"seed"`
		Steps          int    `json:"steps"`
		NegativePrompt string `json:"negative_prompt"`
		// Context is what the pictures are charged to, checked as the
		// chat's is.
		Context string `json:"context"`
	}
	if err := json.Unmarshal(body, &in); err != nil || in.Model == "" {
		http.Error(w, "a model and a prompt are required", http.StatusBadRequest)
		return
	}
	principal, err := r.principalFor(req, user, in.Context)
	if err != nil {
		http.Error(w, err.Error(), http.StatusForbidden)
		return
	}
	// Rebuilt from the fields the page has, so it can set nothing else.
	out := map[string]any{"model": in.Model, "prompt": in.Prompt, "size": in.Size, "n": in.N}
	if in.Seed != nil {
		out["seed"] = *in.Seed
	}
	if in.Steps > 0 {
		out["steps"] = in.Steps
	}
	if s := strings.TrimSpace(in.NegativePrompt); s != "" {
		out["negative_prompt"] = s
	}
	upstream, err := json.Marshal(out)
	if err != nil {
		http.Error(w, "could not build the request", http.StatusInternalServerError)
		return
	}
	next := req.Clone(auth.WithPrincipal(req.Context(), principal))
	next.Body = io.NopCloser(bytes.NewReader(upstream))
	next.ContentLength = int64(len(upstream))
	next.Header.Set("Content-Type", "application/json")
	next.Header.Del(knowledgeHeader)
	r.handleInference(imagesPath)(w, next)
}