329 lines
10 KiB
Go
329 lines
10 KiB
Go
package artifact
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/base64"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"mime"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"net/textproto"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
|
|
"neuralhunt/internal/settings"
|
|
)
|
|
|
|
type referenceImage struct {
|
|
Name string
|
|
ContentType string
|
|
Bytes []byte
|
|
}
|
|
|
|
var collectionAnchorMu sync.Mutex
|
|
|
|
var ErrCharacterAnchorExists = errors.New("character anchor already exists")
|
|
|
|
func validateOpenAIImageSize(model string, width, height int) error {
|
|
if !strings.EqualFold(strings.TrimSpace(model), "gpt-image-2") {
|
|
return nil
|
|
}
|
|
pixels := width * height
|
|
longEdge, shortEdge := width, height
|
|
if height > width {
|
|
longEdge, shortEdge = height, width
|
|
}
|
|
if width%16 != 0 || height%16 != 0 || longEdge > 3840 || shortEdge <= 0 || longEdge > 3*shortEdge || pixels < 655360 || pixels > 8294400 {
|
|
return fmt.Errorf("gpt-image-2 size %dx%d is invalid: edges must be multiples of 16, max edge 3840, ratio <=3:1, pixels 655360..8294400", width, height)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (w *Worker) characterAnchorPath() string {
|
|
return filepath.Join(w.dir, "_collection", "character_anchor.png")
|
|
}
|
|
|
|
func (w *Worker) createCharacterAnchorLocked(ctx context.Context, cfg settings.Runtime, failIfExists bool) ([]byte, bool, error) {
|
|
anchorPath := w.characterAnchorPath()
|
|
if b, err := os.ReadFile(anchorPath); err == nil && len(b) > 1024 {
|
|
if failIfExists {
|
|
return nil, false, ErrCharacterAnchorExists
|
|
}
|
|
return b, false, nil
|
|
}
|
|
if !strings.HasPrefix(strings.ToLower(strings.TrimSpace(cfg.ArtifactModel)), "gpt-image-") {
|
|
return nil, false, fmt.Errorf("collection character consistency requires a GPT Image model; got %q", cfg.ArtifactModel)
|
|
}
|
|
if err := validateOpenAIImageSize(cfg.ArtifactModel, cfg.ArtifactWidth, cfg.ArtifactHeight); err != nil {
|
|
return nil, false, err
|
|
}
|
|
if err := os.MkdirAll(filepath.Dir(anchorPath), 0o750); err != nil {
|
|
return nil, false, err
|
|
}
|
|
// The global anchor is intentionally generated without a style image. It is
|
|
// an identity reference only; visual style is supplied independently per
|
|
// task as Image 2 during actual card generation.
|
|
res, err := w.openAITrackedRequest(ctx, cfg, "", "character_anchor", characterAnchorPrompt, nil)
|
|
if err != nil {
|
|
return nil, false, fmt.Errorf("create canonical RIFT anchor: %w", err)
|
|
}
|
|
if len(res.Bytes) == 0 {
|
|
return nil, false, errors.New("create canonical RIFT anchor: OpenAI returned empty image")
|
|
}
|
|
if err := atomicWrite(anchorPath, res.Bytes, 0o640); err != nil {
|
|
return nil, false, fmt.Errorf("store canonical RIFT anchor: %w", err)
|
|
}
|
|
return res.Bytes, true, nil
|
|
}
|
|
|
|
// CreateCharacterAnchor creates the collection-wide RIFT identity reference on
|
|
// demand from the admin UI. It deliberately refuses to overwrite an existing
|
|
// anchor so the identity cannot be changed accidentally after a collection has
|
|
// started.
|
|
func (w *Worker) CreateCharacterAnchor(ctx context.Context) error {
|
|
collectionAnchorMu.Lock()
|
|
defer collectionAnchorMu.Unlock()
|
|
_, _, err := w.createCharacterAnchorLocked(ctx, w.settings.Get(), true)
|
|
return err
|
|
}
|
|
|
|
func (w *Worker) ensureCharacterAnchor(ctx context.Context, cfg settings.Runtime) ([]byte, bool, error) {
|
|
collectionAnchorMu.Lock()
|
|
defer collectionAnchorMu.Unlock()
|
|
return w.createCharacterAnchorLocked(ctx, cfg, false)
|
|
}
|
|
|
|
func (w *Worker) openAI(ctx context.Context, cfg settings.Runtime, x win, prompt string) (imageResult, error) {
|
|
if err := validateOpenAIImageSize(cfg.ArtifactModel, cfg.ArtifactWidth, cfg.ArtifactHeight); err != nil {
|
|
return imageResult{}, err
|
|
}
|
|
if !strings.EqualFold(strings.TrimSpace(cfg.ArtifactPreset), collectionPresetRaccoon) {
|
|
return w.openAITrackedRequest(ctx, cfg, x.ID, "artifact", prompt, nil)
|
|
}
|
|
anchor, created, err := w.ensureCharacterAnchor(ctx, cfg)
|
|
if err != nil {
|
|
return imageResult{}, err
|
|
}
|
|
styleRef, err := w.loadStyleReference(x.StyleReference)
|
|
if err != nil && strings.EqualFold(strings.TrimSpace(x.Origin), "admin_drop") {
|
|
log.Printf("artifact worker admin drop %s: template style %q unavailable (%v); falling back to bundled default style", x.ID, x.StyleReference, err)
|
|
styleRef, err = w.loadStyleReference("")
|
|
}
|
|
if err != nil {
|
|
return imageResult{}, err
|
|
}
|
|
res, err := w.openAITrackedRequest(ctx, cfg, x.ID, "artifact", prompt, []referenceImage{
|
|
{Name: "character_anchor.png", ContentType: "image/png", Bytes: anchor},
|
|
{Name: styleRef.Name, ContentType: styleRef.ContentType, Bytes: styleRef.Bytes},
|
|
})
|
|
if err != nil {
|
|
return imageResult{}, err
|
|
}
|
|
h := sha256.Sum256(anchor)
|
|
if res.Meta == nil {
|
|
res.Meta = map[string]any{}
|
|
}
|
|
res.Meta["character_anchor_sha256"] = hex.EncodeToString(h[:])
|
|
res.Meta["character_anchor_created"] = created
|
|
res.Meta["style_reference_sha256"] = styleRef.SHA256
|
|
res.Meta["style_reference_name"] = styleRef.Name
|
|
res.Meta["style_reference_custom"] = styleRef.Custom
|
|
res.Meta["reference_mode"] = "character-plus-task-style"
|
|
return res, nil
|
|
}
|
|
|
|
func (w *Worker) openAIRequest(ctx context.Context, cfg settings.Runtime, prompt string, refs []referenceImage) (imageResult, error) {
|
|
key := strings.TrimSpace(os.Getenv("OPENAI_API_KEY"))
|
|
if key == "" {
|
|
return imageResult{}, errors.New("OPENAI_API_KEY is not configured")
|
|
}
|
|
base := strings.TrimRight(strings.TrimSpace(os.Getenv("OPENAI_BASE_URL")), "/")
|
|
if base == "" {
|
|
base = "https://api.openai.com"
|
|
}
|
|
size := fmt.Sprintf("%dx%d", cfg.ArtifactWidth, cfg.ArtifactHeight)
|
|
quality := strings.ToLower(strings.TrimSpace(cfg.ArtifactQuality))
|
|
if quality == "" {
|
|
quality = "medium"
|
|
}
|
|
|
|
var req *http.Request
|
|
var err error
|
|
endpoint := ""
|
|
if len(refs) > 0 {
|
|
endpoint = base + "/v1/images/edits"
|
|
if strings.HasSuffix(base, "/v1") {
|
|
endpoint = base + "/images/edits"
|
|
}
|
|
var body bytes.Buffer
|
|
mw := multipart.NewWriter(&body)
|
|
fields := map[string]string{
|
|
"model": cfg.ArtifactModel,
|
|
"prompt": prompt,
|
|
"n": "1",
|
|
"size": size,
|
|
"quality": quality,
|
|
"output_format": "png",
|
|
}
|
|
for k, v := range fields {
|
|
if err := mw.WriteField(k, v); err != nil {
|
|
return imageResult{}, err
|
|
}
|
|
}
|
|
for i, ref := range refs {
|
|
if len(ref.Bytes) == 0 {
|
|
continue
|
|
}
|
|
name := strings.TrimSpace(ref.Name)
|
|
if name == "" {
|
|
name = fmt.Sprintf("reference_%d.png", i+1)
|
|
}
|
|
contentType := strings.TrimSpace(ref.ContentType)
|
|
if contentType == "" {
|
|
contentType = styleContentType(name, ref.Bytes)
|
|
}
|
|
switch contentType {
|
|
case "image/jpeg", "image/png", "image/webp":
|
|
default:
|
|
return imageResult{}, fmt.Errorf("unsupported reference image content type %q for %s", contentType, name)
|
|
}
|
|
|
|
header := make(textproto.MIMEHeader)
|
|
header.Set("Content-Disposition", mime.FormatMediaType("form-data", map[string]string{
|
|
"name": "image[]",
|
|
"filename": name,
|
|
}))
|
|
header.Set("Content-Type", contentType)
|
|
part, err := mw.CreatePart(header)
|
|
if err != nil {
|
|
return imageResult{}, err
|
|
}
|
|
if _, err := part.Write(ref.Bytes); err != nil {
|
|
return imageResult{}, err
|
|
}
|
|
}
|
|
if err := mw.Close(); err != nil {
|
|
return imageResult{}, err
|
|
}
|
|
req, err = http.NewRequestWithContext(ctx, http.MethodPost, endpoint, &body)
|
|
if err == nil {
|
|
req.Header.Set("Content-Type", mw.FormDataContentType())
|
|
}
|
|
} else {
|
|
endpoint = base + "/v1/images/generations"
|
|
if strings.HasSuffix(base, "/v1") {
|
|
endpoint = base + "/images/generations"
|
|
}
|
|
payload := map[string]any{
|
|
"model": cfg.ArtifactModel,
|
|
"prompt": prompt,
|
|
"n": 1,
|
|
"size": size,
|
|
"quality": quality,
|
|
"output_format": "png",
|
|
}
|
|
body, marshalErr := json.Marshal(payload)
|
|
if marshalErr != nil {
|
|
return imageResult{}, marshalErr
|
|
}
|
|
req, err = http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
|
|
if err == nil {
|
|
req.Header.Set("Content-Type", "application/json")
|
|
}
|
|
}
|
|
if err != nil {
|
|
return imageResult{}, err
|
|
}
|
|
req.Header.Set("Authorization", "Bearer "+key)
|
|
resp, err := w.http.Do(req)
|
|
if err != nil {
|
|
return imageResult{}, err
|
|
}
|
|
defer resp.Body.Close()
|
|
raw, err := io.ReadAll(io.LimitReader(resp.Body, 128<<20))
|
|
if err != nil {
|
|
return imageResult{}, err
|
|
}
|
|
requestID := strings.TrimSpace(resp.Header.Get("x-request-id"))
|
|
if resp.StatusCode/100 != 2 {
|
|
if requestID != "" {
|
|
return imageResult{}, fmt.Errorf("OpenAI images API HTTP %d (request %s): %s", resp.StatusCode, requestID, truncate(string(raw), 1400))
|
|
}
|
|
return imageResult{}, fmt.Errorf("OpenAI images API HTTP %d: %s", resp.StatusCode, truncate(string(raw), 1400))
|
|
}
|
|
var out struct {
|
|
Data []struct {
|
|
B64 string `json:"b64_json"`
|
|
} `json:"data"`
|
|
Usage *struct {
|
|
TotalTokens int64 `json:"total_tokens"`
|
|
InputTokens int64 `json:"input_tokens"`
|
|
OutputTokens int64 `json:"output_tokens"`
|
|
InputDetails struct {
|
|
TextTokens int64 `json:"text_tokens"`
|
|
ImageTokens int64 `json:"image_tokens"`
|
|
} `json:"input_tokens_details"`
|
|
} `json:"usage"`
|
|
}
|
|
if err := json.Unmarshal(raw, &out); err != nil {
|
|
return imageResult{}, fmt.Errorf("decode OpenAI image response: %w", err)
|
|
}
|
|
if len(out.Data) == 0 || out.Data[0].B64 == "" {
|
|
return imageResult{}, errors.New("OpenAI images API returned no b64_json")
|
|
}
|
|
img, err := base64.StdEncoding.DecodeString(out.Data[0].B64)
|
|
if err != nil {
|
|
return imageResult{}, err
|
|
}
|
|
meta := map[string]any{
|
|
"model": cfg.ArtifactModel,
|
|
"size": size,
|
|
"quality": quality,
|
|
"endpoint": map[bool]string{true: "images/edits", false: "images/generations"}[len(refs) > 0],
|
|
"references": len(refs),
|
|
}
|
|
if requestID != "" {
|
|
meta["request_id"] = requestID
|
|
}
|
|
result := imageResult{Bytes: img, Ext: "png", Provider: "openai", Meta: meta}
|
|
if out.Usage != nil {
|
|
u := imageUsage{
|
|
TotalTokens: out.Usage.TotalTokens,
|
|
InputTokens: out.Usage.InputTokens,
|
|
OutputTokens: out.Usage.OutputTokens,
|
|
TextInputTokens: out.Usage.InputDetails.TextTokens,
|
|
ImageInputTokens: out.Usage.InputDetails.ImageTokens,
|
|
}
|
|
if u.available() {
|
|
result.Usage = &u
|
|
meta["usage"] = map[string]any{
|
|
"total_tokens": u.TotalTokens,
|
|
"input_tokens": u.InputTokens,
|
|
"input_text_tokens": u.TextInputTokens,
|
|
"input_image_tokens": u.ImageInputTokens,
|
|
"output_tokens": u.OutputTokens,
|
|
}
|
|
if cost, basis, ok := estimateOpenAIImageCost(cfg.ArtifactModel, u); ok {
|
|
result.EstimatedCostUSD = &cost
|
|
result.PricingBasis = basis
|
|
meta["estimated_cost_usd"] = cost
|
|
meta["pricing_basis"] = basis
|
|
} else if basis != "" {
|
|
result.PricingBasis = basis
|
|
meta["pricing_basis"] = basis
|
|
meta["estimated_cost_available"] = false
|
|
}
|
|
}
|
|
}
|
|
return result, nil
|
|
}
|