Luminary/internal/prediction/params.go

package prediction

import (
	"fmt"

	"github.com/rose_cat707/luminary/internal/model"
)

type ParamType string

const (
	ParamString ParamType = "string"
	ParamInt    ParamType = "int"
	ParamFloat  ParamType = "float"
)

type ParamDef struct {
	Name     string    `json:"name"`
	Label    string    `json:"label"`
	Type     ParamType `json:"type"`
	Required bool      `json:"required"`
	Default  any       `json:"default,omitempty"`
}

var Txt2ImgParams = []ParamDef{
	{Name: "prompt", Label: "正向提示词", Type: ParamString, Required: true},
	{Name: "negative_prompt", Label: "负向提示词", Type: ParamString, Default: ""},
	{Name: "seed", Label: "随机种子", Type: ParamInt, Default: -1},
	{Name: "width", Label: "宽度", Type: ParamInt, Default: 1024},
	{Name: "height", Label: "高度", Type: ParamInt, Default: 1024},
	{Name: "steps", Label: "采样步数", Type: ParamInt, Default: 20},
	{Name: "cfg_scale", Label: "CFG", Type: ParamFloat, Default: 7.0},
}

var Img2ImgParams = []ParamDef{
	{Name: "prompt", Label: "正向提示词", Type: ParamString, Required: true},
	{Name: "negative_prompt", Label: "负向提示词", Type: ParamString, Default: ""},
	{Name: "image", Label: "输入图 URL", Type: ParamString, Required: true},
	{Name: "seed", Label: "随机种子", Type: ParamInt, Default: -1},
	{Name: "steps", Label: "采样步数", Type: ParamInt, Default: 20},
	{Name: "cfg_scale", Label: "CFG", Type: ParamFloat, Default: 7.0},
	{Name: "denoise", Label: "重绘幅度", Type: ParamFloat, Default: 0.75},
}

func ParamsForWorkflowType(wt model.WorkflowType) []ParamDef {
	switch wt {
	case model.WorkflowImg2Img:
		return Img2ImgParams
	default:
		return Txt2ImgParams
	}
}

func RequiredBindings(wt model.WorkflowType) []string {
	switch wt {
	case model.WorkflowImg2Img:
		return []string{"prompt", "image"}
	default:
		return []string{"prompt"}
	}
}

func ValidateInput(wt model.WorkflowType, input map[string]any) (map[string]any, error) {
	defs := ParamsForWorkflowType(wt)
	out := make(map[string]any, len(defs))
	for _, d := range defs {
		v, ok := input[d.Name]
		if !ok || v == nil {
			if d.Required {
				return nil, fmt.Errorf("missing required field: %s", d.Name)
			}
			if d.Default != nil {
				out[d.Name] = d.Default
			}
			continue
		}
		coerced, err := coerceValue(d, v)
		if err != nil {
			return nil, fmt.Errorf("%s: %w", d.Name, err)
		}
		out[d.Name] = coerced
	}
	return out, nil
}

func coerceValue(d ParamDef, v any) (any, error) {
	switch d.Type {
	case ParamString:
		s, ok := v.(string)
		if !ok {
			return nil, fmt.Errorf("expected string")
		}
		if d.Required && s == "" {
			return nil, fmt.Errorf("cannot be empty")
		}
		return s, nil
	case ParamInt:
		switch n := v.(type) {
		case float64:
			return int(n), nil
		case int:
			return n, nil
		case int64:
			return int(n), nil
		default:
			return nil, fmt.Errorf("expected integer")
		}
	case ParamFloat:
		switch n := v.(type) {
		case float64:
			return n, nil
		case int:
			return float64(n), nil
		case int64:
			return float64(n), nil
		default:
			return nil, fmt.Errorf("expected number")
		}
	default:
		return v, nil
	}
}