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 openAPIType(pt ParamType) string {
	switch pt {
	case ParamInt:
		return "integer"
	case ParamFloat:
		return "number"
	default:
		return "string"
	}
}

// OpenAPISchemaForWorkflow builds a Replicate-compatible schema with components.schemas.Input.
func OpenAPISchemaForWorkflow(wt model.WorkflowType, title, description string) map[string]any {
	defs := ParamsForWorkflowType(wt)
	properties := make(map[string]any, len(defs))
	required := make([]string, 0, len(defs))
	for i, d := range defs {
		prop := map[string]any{
			"type":        openAPIType(d.Type),
			"title":       d.Label,
			"description": d.Label,
			"x-order":     i,
		}
		if d.Default != nil {
			prop["default"] = d.Default
		}
		properties[d.Name] = prop
		if d.Required {
			required = append(required, d.Name)
		}
	}
	inputSchema := map[string]any{
		"type":       "object",
		"properties": properties,
	}
	if len(required) > 0 {
		inputSchema["required"] = required
	}
	return map[string]any{
		"openapi": "3.0.0",
		"info": map[string]any{
			"title":       title,
			"version":     "1.0.0",
			"description": description,
		},
		"components": map[string]any{
			"schemas": map[string]any{
				"Input": inputSchema,
			},
		},
	}
}

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
	}
}