libnovel/backend/internal/cfai/image.go

// Image generation via Cloudflare Workers AI text-to-image models.
//
// API reference:
//
//	POST https://api.cloudflare.com/client/v4/accounts/{accountID}/ai/run/{model}
//	Authorization: Bearer {apiToken}
//	Content-Type:  application/json
//
// Text-only request (all models):
//
//	{ "prompt": "...", "num_steps": 20 }
//
// Reference-image request:
//   - FLUX models: { "prompt": "...", "image_b64": "<base64>" }
//   - SD img2img:  { "prompt": "...", "image": [r,g,b,a,...], "strength": 0.75 }
//
// All models return raw PNG bytes on success (Content-Type: image/png).
//
// Recommended models for LibNovel:
//   - Book covers (no reference): flux-2-dev, flux-2-klein-9b, lucid-origin
//   - Chapter images (speed):     flux-2-klein-4b, flux-1-schnell
//   - With reference image:       flux-2-dev, flux-2-klein-9b, sd-v1-5-img2img
package cfai

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"image"
	_ "image/jpeg" // register JPEG decoder
	_ "image/png"  // register PNG decoder
	"io"
	"net/http"
	"time"
)

// ImageModel identifies a Cloudflare Workers AI text-to-image model.
type ImageModel string

const (
	// ImageModelFlux2Dev — best quality, multi-reference. Recommended for covers.
	ImageModelFlux2Dev ImageModel = "@cf/black-forest-labs/flux-2-dev"
	// ImageModelFlux2Klein9B — 9B params, multi-reference. Good for covers.
	ImageModelFlux2Klein9B ImageModel = "@cf/black-forest-labs/flux-2-klein-9b"
	// ImageModelFlux2Klein4B — ultra-fast, unified gen+edit. Recommended for chapters.
	ImageModelFlux2Klein4B ImageModel = "@cf/black-forest-labs/flux-2-klein-4b"
	// ImageModelFlux1Schnell — fastest, text-only. Good for quick illustrations.
	ImageModelFlux1Schnell ImageModel = "@cf/black-forest-labs/flux-1-schnell"
	// ImageModelSDXLLightning — fast 1024px generation.
	ImageModelSDXLLightning ImageModel = "@cf/bytedance/stable-diffusion-xl-lightning"
	// ImageModelSD15Img2Img — explicit img2img with flat RGBA reference.
	ImageModelSD15Img2Img ImageModel = "@cf/runwayml/stable-diffusion-v1-5-img2img"
	// ImageModelSDXLBase — Stability AI SDXL base.
	ImageModelSDXLBase ImageModel = "@cf/stabilityai/stable-diffusion-xl-base-1.0"
	// ImageModelLucidOrigin — Leonardo AI; strong prompt adherence.
	ImageModelLucidOrigin ImageModel = "@cf/leonardo/lucid-origin"
	// ImageModelPhoenix10 — Leonardo AI; accurate text rendering.
	ImageModelPhoenix10 ImageModel = "@cf/leonardo/phoenix-1.0"

	// DefaultImageModel is the default model for book-cover generation.
	DefaultImageModel = ImageModelFlux2Dev
)

// ImageModelInfo describes a single image generation model.
type ImageModelInfo struct {
	ID             string   `json:"id"`
	Label          string   `json:"label"`
	Provider       string   `json:"provider"`
	SupportsRef    bool     `json:"supports_ref"`
	RecommendedFor []string `json:"recommended_for"` // "cover" and/or "chapter"
	Description    string   `json:"description"`
}

// AllImageModels returns metadata about every supported image model.
func AllImageModels() []ImageModelInfo {
	return []ImageModelInfo{
		{
			ID: string(ImageModelFlux2Dev), Label: "FLUX.2 Dev", Provider: "Black Forest Labs",
			SupportsRef: true, RecommendedFor: []string{"cover"},
			Description: "Best quality; multi-reference editing. Recommended for book covers.",
		},
		{
			ID: string(ImageModelFlux2Klein9B), Label: "FLUX.2 Klein 9B", Provider: "Black Forest Labs",
			SupportsRef: true, RecommendedFor: []string{"cover"},
			Description: "9B parameters with multi-reference support.",
		},
		{
			ID: string(ImageModelFlux2Klein4B), Label: "FLUX.2 Klein 4B", Provider: "Black Forest Labs",
			SupportsRef: true, RecommendedFor: []string{"chapter"},
			Description: "Ultra-fast unified gen+edit. Recommended for chapter images.",
		},
		{
			ID: string(ImageModelFlux1Schnell), Label: "FLUX.1 Schnell", Provider: "Black Forest Labs",
			SupportsRef: false, RecommendedFor: []string{"chapter"},
			Description: "Fastest inference. Good for quick chapter illustrations.",
		},
		{
			ID: string(ImageModelSDXLLightning), Label: "SDXL Lightning", Provider: "ByteDance",
			SupportsRef: false, RecommendedFor: []string{"chapter"},
			Description: "Lightning-fast 1024px images in a few steps.",
		},
		{
			ID: string(ImageModelSD15Img2Img), Label: "SD 1.5 img2img", Provider: "RunwayML",
			SupportsRef: true, RecommendedFor: []string{"cover", "chapter"},
			Description: "Explicit img2img: generates from a reference image + prompt.",
		},
		{
			ID: string(ImageModelSDXLBase), Label: "SDXL Base 1.0", Provider: "Stability AI",
			SupportsRef: false, RecommendedFor: []string{"cover"},
			Description: "Stable Diffusion XL base model.",
		},
		{
			ID: string(ImageModelLucidOrigin), Label: "Lucid Origin", Provider: "Leonardo AI",
			SupportsRef: false, RecommendedFor: []string{"cover"},
			Description: "Highly prompt-responsive; strong graphic design and HD renders.",
		},
		{
			ID: string(ImageModelPhoenix10), Label: "Phoenix 1.0", Provider: "Leonardo AI",
			SupportsRef: false, RecommendedFor: []string{"cover"},
			Description: "Exceptional prompt adherence; accurate text rendering.",
		},
	}
}

// ImageRequest is the input to GenerateImage / GenerateImageFromReference.
type ImageRequest struct {
	// Prompt is the text description of the desired image.
	Prompt string
	// Model is the CF Workers AI model. Defaults to DefaultImageModel when empty.
	Model ImageModel
	// NumSteps controls inference quality (default 20). Range: 1–20.
	NumSteps int
	// Width and Height in pixels. 0 = model default (typically 1024x1024).
	Width, Height int
	// Guidance controls prompt adherence (default 7.5).
	Guidance float64
	// Strength for img2img: 0.0 = copy reference, 1.0 = ignore reference (default 0.75).
	Strength float64
}

// ImageGenClient generates images via Cloudflare Workers AI.
type ImageGenClient interface {
	// GenerateImage creates an image from a text prompt only.
	// Returns raw PNG bytes.
	GenerateImage(ctx context.Context, req ImageRequest) ([]byte, error)

	// GenerateImageFromReference creates an image from a text prompt + reference image.
	// refImage should be PNG or JPEG bytes. Returns raw PNG bytes.
	GenerateImageFromReference(ctx context.Context, req ImageRequest, refImage []byte) ([]byte, error)

	// Models returns metadata about all supported image models.
	Models() []ImageModelInfo
}

// imageGenHTTPClient is the concrete CF AI image generation client.
type imageGenHTTPClient struct {
	accountID string
	apiToken  string
	http      *http.Client
}

// NewImageGen returns an ImageGenClient for the given Cloudflare account.
func NewImageGen(accountID, apiToken string) ImageGenClient {
	return &imageGenHTTPClient{
		accountID: accountID,
		apiToken:  apiToken,
		http:      &http.Client{Timeout: 5 * time.Minute},
	}
}

// GenerateImage generates an image from text only.
func (c *imageGenHTTPClient) GenerateImage(ctx context.Context, req ImageRequest) ([]byte, error) {
	req = applyImageDefaults(req)
	body := map[string]any{
		"prompt":    req.Prompt,
		"num_steps": req.NumSteps,
	}
	if req.Width > 0 {
		body["width"] = req.Width
	}
	if req.Height > 0 {
		body["height"] = req.Height
	}
	if req.Guidance > 0 {
		body["guidance"] = req.Guidance
	}
	return c.callImageAPI(ctx, req.Model, body)
}

// GenerateImageFromReference generates an image from a text prompt + reference image.
func (c *imageGenHTTPClient) GenerateImageFromReference(ctx context.Context, req ImageRequest, refImage []byte) ([]byte, error) {
	if len(refImage) == 0 {
		return c.GenerateImage(ctx, req)
	}
	req = applyImageDefaults(req)

	var body map[string]any
	if req.Model == ImageModelSD15Img2Img {
		pixels, err := decodeImageToRGBA(refImage)
		if err != nil {
			return nil, fmt.Errorf("cfai/image: decode reference: %w", err)
		}
		strength := req.Strength
		if strength <= 0 {
			strength = 0.75
		}
		body = map[string]any{
			"prompt":    req.Prompt,
			"image":     pixels,
			"strength":  strength,
			"num_steps": req.NumSteps,
		}
	} else {
		b64 := base64.StdEncoding.EncodeToString(refImage)
		body = map[string]any{
			"prompt":    req.Prompt,
			"image_b64": b64,
			"num_steps": req.NumSteps,
		}
		if req.Strength > 0 {
			body["strength"] = req.Strength
		}
	}
	if req.Width > 0 {
		body["width"] = req.Width
	}
	if req.Height > 0 {
		body["height"] = req.Height
	}
	if req.Guidance > 0 {
		body["guidance"] = req.Guidance
	}
	return c.callImageAPI(ctx, req.Model, body)
}

// Models returns all supported image model metadata.
func (c *imageGenHTTPClient) Models() []ImageModelInfo {
	return AllImageModels()
}

func (c *imageGenHTTPClient) callImageAPI(ctx context.Context, model ImageModel, body map[string]any) ([]byte, error) {
	encoded, err := json.Marshal(body)
	if err != nil {
		return nil, fmt.Errorf("cfai/image: marshal: %w", err)
	}
	url := fmt.Sprintf("https://api.cloudflare.com/client/v4/accounts/%s/ai/run/%s",
		c.accountID, string(model))
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(encoded))
	if err != nil {
		return nil, fmt.Errorf("cfai/image: build request: %w", err)
	}
	req.Header.Set("Authorization", "Bearer "+c.apiToken)
	req.Header.Set("Content-Type", "application/json")

	resp, err := c.http.Do(req)
	if err != nil {
		return nil, fmt.Errorf("cfai/image: http: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		errBody, _ := io.ReadAll(resp.Body)
		msg := string(errBody)
		if len(msg) > 300 {
			msg = msg[:300]
		}
		return nil, fmt.Errorf("cfai/image: model %s returned %d: %s", model, resp.StatusCode, msg)
	}

	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("cfai/image: read response: %w", err)
	}
	return data, nil
}

func applyImageDefaults(req ImageRequest) ImageRequest {
	if req.Model == "" {
		req.Model = DefaultImageModel
	}
	if req.NumSteps <= 0 {
		req.NumSteps = 20
	}
	return req
}

// decodeImageToRGBA decodes PNG/JPEG bytes to a flat []uint8 RGBA pixel array
// required by the stable-diffusion-v1-5-img2img model.
func decodeImageToRGBA(data []byte) ([]uint8, error) {
	img, _, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("decode image: %w", err)
	}
	bounds := img.Bounds()
	w := bounds.Max.X - bounds.Min.X
	h := bounds.Max.Y - bounds.Min.Y
	pixels := make([]uint8, w*h*4)
	idx := 0
	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			r, g, b, a := img.At(x, y).RGBA()
			pixels[idx] = uint8(r >> 8)
			pixels[idx+1] = uint8(g >> 8)
			pixels[idx+2] = uint8(b >> 8)
			pixels[idx+3] = uint8(a >> 8)
			idx += 4
		}
	}
	return pixels, nil
}