qwen_image_outpaint.go

package gocomfyui

import (
	"bytes"
	_ "embed"
	"encoding/json"
	"fmt"
	"log/slog"
	"strings"
	"text/template"
)

var (
	qwenImageOutpaintTemplate *template.Template
	//go:embed qwen_image_outpaint.tmpl
	qwenImageOutpaintTemplateData string
)

func init() {
	var err error
	qwenImageOutpaintTemplate, err = template.New("qwen_image_outpaint").Parse(qwenImageOutpaintTemplateData)
	if err != nil {
		panic(fmt.Errorf("qwen_image_outpaint: %v", err))
	}
}

func QwenImageOutpaint(params QwenImageOutpaintParams) ([]byte, error) {
	err := escapeStrings(&params)
	if err != nil {
		return nil, fmt.Errorf("escape: %w", err)
	}
	var b bytes.Buffer
	err = qwenImageOutpaintTemplate.Execute(&b, params)
	if err != nil {
		return nil, fmt.Errorf("execute template: %w", err)
	}
	data := b.Bytes()
	if !json.Valid(data) {
		return nil, fmt.Errorf("invalid json: %s", data)
	}
	return data, nil
}

type QwenImageOutpaintParams struct {
	FilenamePrefix string

	Image string

	PositivePrompt string
	NegativePrompt string

	Left   int
	Right  int
	Top    int
	Bottom int

	Feathering int

	Steps int

	CFG float64

	Seed int64
}

func (p QwenImageOutpaintParams) LogValue() slog.Value {
	as := make([]slog.Attr, 0, 9)
	as = append(as, slog.String("filename_prefix", p.FilenamePrefix))
	as = append(as, slog.String("image", p.Image))
	as = append(as, slog.String("positive_prompt", p.PositivePrompt))
	if p.NegativePrompt != "" {
		as = append(as, slog.String("negative_prompt", p.NegativePrompt))
	}
	as = append(as, slog.Group("boundaries",
		slog.Int("left", p.Left),
		slog.Int("right", p.Right),
		slog.Int("top", p.Top),
		slog.Int("bottom", p.Bottom),
	))
	as = append(as, slog.Int("feathering", p.Feathering))
	as = append(as, slog.Int("steps", p.Steps))
	as = append(as, slog.Float64("cfg", p.CFG))
	as = append(as, slog.Int64("seed", p.Seed))
	return slog.GroupValue(as...)
}

func (p QwenImageOutpaintParams) String() string {
	var b strings.Builder
	b.WriteString(fmt.Sprintf(`filename_prefix="%s"`, p.FilenamePrefix))
	b.WriteRune('\n')
	b.WriteString(fmt.Sprintf(`image="%s"`, p.Image))
	if p.PositivePrompt != "" {
		b.WriteRune('\n')
		b.WriteString(fmt.Sprintf(`positive_prompt="%s"`, p.PositivePrompt))
	}
	if p.NegativePrompt != "" {
		b.WriteRune('\n')
		b.WriteString(fmt.Sprintf(`negative_prompt="%s"`, p.NegativePrompt))
	}
	b.WriteRune('\n')
	b.WriteString(fmt.Sprintf(`left=%d, right=%d, top=%d, bottom=%d`,
		p.Left, p.Right, p.Top, p.Bottom))
	b.WriteRune('\n')
	b.WriteString(fmt.Sprintf(`feathering=%d`, p.Feathering))
	b.WriteRune('\n')
	b.WriteString(fmt.Sprintf(`steps=%d`, p.Steps))
	b.WriteRune('\n')
	b.WriteString(fmt.Sprintf(`cfg=%.1f`, p.CFG))
	b.WriteRune('\n')
	b.WriteString(fmt.Sprintf(`seed=%d`, p.Seed))
	return b.String()
}

func NewQwenImageOutpaintParams(opts ...QwenImageOutpaintParamsOption) QwenImageOutpaintParams {
	def := QwenImageOutpaintParams{
		FilenamePrefix: FilenamePrefix,

		Right:  250,
		Left:   250,
		Top:    130,
		Bottom: 130,

		Feathering: 32,

		Steps: 20,

		CFG: 4.0,
	}
	for _, opt := range opts {
		opt(&def)
	}
	return def
}

type QwenImageOutpaintParamsOption func(*QwenImageOutpaintParams)

func QwenImageOutpaintFilenamePrefix(filenamePrefix string) QwenImageOutpaintParamsOption {
	return func(params *QwenImageOutpaintParams) {
		params.FilenamePrefix = filenamePrefix
	}
}

func QwenImageOutpaintImage(img string) QwenImageOutpaintParamsOption {
	return func(params *QwenImageOutpaintParams) {
		params.Image = img
	}
}

func QwenImageOutpaintPositivePrompt(prompt string) QwenImageOutpaintParamsOption {
	return func(params *QwenImageOutpaintParams) {
		params.PositivePrompt = prompt
	}
}

func QwenImageOutpaintNegativePrompt(prompt string) QwenImageOutpaintParamsOption {
	return func(params *QwenImageOutpaintParams) {
		params.NegativePrompt = prompt
	}
}

func QwenImageOutpaintRight(right int) QwenImageOutpaintParamsOption {
	return func(params *QwenImageOutpaintParams) {
		params.Right = right
	}
}

func QwenImageOutpaintLeft(left int) QwenImageOutpaintParamsOption {
	return func(params *QwenImageOutpaintParams) {
		params.Left = left
	}
}

func QwenImageOutpaintTop(top int) QwenImageOutpaintParamsOption {
	return func(params *QwenImageOutpaintParams) {
		params.Top = top
	}
}

func QwenImageOutpaintBottom(bottom int) QwenImageOutpaintParamsOption {
	return func(params *QwenImageOutpaintParams) {
		params.Bottom = bottom
	}
}

func QwenImageOutpaintBoundaries(right, left, top, bottom int) QwenImageOutpaintParamsOption {
	return func(params *QwenImageOutpaintParams) {
		params.Right = right
		params.Left = left
		params.Top = top
		params.Bottom = bottom
	}
}

func QwenImageOutpaintFeathering(feathering int) QwenImageOutpaintParamsOption {
	return func(params *QwenImageOutpaintParams) {
		params.Feathering = feathering
	}
}

func QwenImageOutpaintSteps(steps int) QwenImageOutpaintParamsOption {
	return func(params *QwenImageOutpaintParams) {
		params.Steps = steps
	}
}

func QwenImageOutpaintCFG(cfg float64) QwenImageOutpaintParamsOption {
	return func(params *QwenImageOutpaintParams) {
		params.CFG = cfg
	}
}

func QwenImageOutpaintSeed(seed int64) QwenImageOutpaintParamsOption {
	return func(params *QwenImageOutpaintParams) {
		params.Seed = seed
	}
}