imagine.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
)

const defaultBaseURL = "https://api.vyro.ai/v2"

var validStyles = map[string]bool{
	"realistic":     true,
	"anime":         true,
	"flux-schnell":  true,
	"flux-dev":      true,
	"flux-dev-fast": true,
	"sdxl-1.0":      true,
	"imagine-turbo": true,
}

var validAspectRatios = map[string]bool{
	"1:1":  true,
	"3:2":  true,
	"4:3":  true,
	"3:4":  true,
	"16:9": true,
	"9:16": true,
}

type Client struct {
	apiKey string
	base   string
	http   *http.Client
}

func NewClient(apiKey string) *Client {
	return &Client{apiKey: apiKey, base: defaultBaseURL, http: &http.Client{}}
}

type apiError struct {
	Status  int    `json:"status"`
	Code    string `json:"code"`
	Error   string `json:"error"`
	Message string `json:"message"`
}

type TextToImageParams struct {
	Prompt      string
	Style       string
	AspectRatio string
	Seed        *int
}

func (c *Client) TextToImage(ctx context.Context, p TextToImageParams) ([]byte, error) {
	if p.Prompt == "" {
		return nil, fmt.Errorf("prompt is required")
	}
	if p.Style == "" {
		p.Style = "imagine-turbo"
	}
	if !validStyles[p.Style] {
		return nil, fmt.Errorf("invalid style %q (must be one of: realistic, anime, flux-schnell, flux-dev, flux-dev-fast, sdxl-1.0, imagine-turbo)", p.Style)
	}
	if p.AspectRatio == "" {
		p.AspectRatio = "1:1"
	}
	if !validAspectRatios[p.AspectRatio] {
		return nil, fmt.Errorf("invalid aspect_ratio %q (must be one of: 1:1, 3:2, 4:3, 3:4, 16:9, 9:16)", p.AspectRatio)
	}

	var body bytes.Buffer
	w := multipart.NewWriter(&body)
	if err := w.WriteField("prompt", p.Prompt); err != nil {
		return nil, err
	}
	if err := w.WriteField("style", p.Style); err != nil {
		return nil, err
	}
	if err := w.WriteField("aspect_ratio", p.AspectRatio); err != nil {
		return nil, err
	}
	if p.Seed != nil {
		if err := w.WriteField("seed", fmt.Sprintf("%d", *p.Seed)); err != nil {
			return nil, err
		}
	}
	if err := w.Close(); err != nil {
		return nil, err
	}

	return c.do(ctx, "/image/generations", w.FormDataContentType(), &body)
}

func (c *Client) RemoveBackground(ctx context.Context, imagePath string) ([]byte, error) {
	f, err := os.Open(imagePath)
	if err != nil {
		return nil, fmt.Errorf("open image: %w", err)
	}
	defer f.Close()

	var body bytes.Buffer
	w := multipart.NewWriter(&body)
	part, err := w.CreateFormFile("image", filepath.Base(imagePath))
	if err != nil {
		return nil, err
	}
	if _, err := io.Copy(part, f); err != nil {
		return nil, err
	}
	if err := w.Close(); err != nil {
		return nil, err
	}

	return c.do(ctx, "/image/background/remover", w.FormDataContentType(), &body)
}

func (c *Client) do(ctx context.Context, path, contentType string, body io.Reader) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.base+path, body)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+c.apiKey)
	req.Header.Set("Content-Type", contentType)

	resp, err := c.http.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode >= 200 && resp.StatusCode < 300 {
		return respBody, nil
	}

	var apiErr apiError
	if err := json.Unmarshal(respBody, &apiErr); err == nil && apiErr.Message != "" {
		return nil, fmt.Errorf("imagine: %d %s — %s", resp.StatusCode, apiErr.Error, apiErr.Message)
	}
	return nil, fmt.Errorf("imagine: %d %s", resp.StatusCode, string(respBody))
}