Luminary/internal/handler/admin/workflows.go

package admin

import (
	"encoding/json"
	"net/http"

	"github.com/gin-gonic/gin"
	"github.com/rose_cat707/luminary/internal/model"
	"github.com/rose_cat707/luminary/internal/prediction"
)

func (h *Handler) GetWorkflowParamMeta(c *gin.Context) {
	c.JSON(http.StatusOK, gin.H{
		"txt2img": prediction.Txt2ImgParams,
		"img2img": prediction.Img2ImgParams,
	})
}

func (h *Handler) AnalyzeWorkflow(c *gin.Context) {
	var req struct {
		WorkflowJSON string             `json:"workflow_json"`
		WorkflowType model.WorkflowType `json:"workflow_type"`
	}
	if err := c.ShouldBindJSON(&req); err != nil || req.WorkflowJSON == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "workflow_json required"})
		return
	}
	if req.WorkflowType == "" {
		req.WorkflowType = model.WorkflowTxt2Img
	}
	result, err := prediction.AnalyzeWorkflow(req.WorkflowJSON, req.WorkflowType)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, result)
}

func (h *Handler) GetPrediction(c *gin.Context) {
	if h.predictions == nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	p, err := h.predictions.GetDBPrediction(c.Param("id"))
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	images, _ := h.images.ListByPrediction(p.ID)
	c.JSON(http.StatusOK, gin.H{
		"prediction": h.predictions.APIView(p, h.publicBaseURL()),
		"images":     images,
	})
}

func (h *Handler) ServeAdminImage(c *gin.Context) {
	if h.images == nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	img, err := h.images.Get(c.Param("id"))
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
		return
	}
	c.Header("Content-Type", img.Mime)
	c.File(img.LocalPath)
}

func bindingFromSuggestions(suggestions []prediction.BindingSuggestion, overrides map[string]*prediction.NodeField) prediction.InputBinding {
	binding := prediction.InputBinding{}
	for _, s := range suggestions {
		if o, ok := overrides[s.Param]; ok && o != nil {
			if o.Node != "" && o.Field != "" {
				binding[s.Param] = o
			}
			continue
		}
		if s.Node != "" && s.Field != "" {
			binding[s.Param] = &prediction.NodeField{Node: s.Node, Field: s.Field}
		}
	}
	return binding
}

func validateImageModelCreate(req struct {
	ModelID          string
	DisplayName      string
	WorkflowType     model.WorkflowType
	WorkflowJSON     string
	InputBindingJSON string
	InputBinding     map[string]*prediction.NodeField
}) error {
	if req.WorkflowJSON == "" {
		return nil
	}
	var workflow map[string]any
	if err := json.Unmarshal([]byte(req.WorkflowJSON), &workflow); err != nil {
		return err
	}
	var binding prediction.InputBinding
	if req.InputBindingJSON != "" {
		b, err := prediction.ParseBinding(req.InputBindingJSON)
		if err != nil {
			return err
		}
		binding = b
	} else if req.InputBinding != nil {
		binding = req.InputBinding
	}
	wt := req.WorkflowType
	if wt == "" {
		wt = model.WorkflowTxt2Img
	}
	return binding.Validate(workflow, prediction.RequiredBindings(wt))
}