diff --git a/internal/handler/admin/providers_keys.go b/internal/handler/admin/providers_keys.go index 52007c7..b2c107e 100644 --- a/internal/handler/admin/providers_keys.go +++ b/internal/handler/admin/providers_keys.go @@ -275,6 +275,9 @@ func (h *Handler) CreateModel(c *gin.Context) { VersionHash: prediction.HashWorkflowJSON(req.WorkflowJSON), Enabled: true, } + if prov.Category == model.CategoryImage { + m.Alias = "" + } if m.DisplayName == "" { m.DisplayName = req.ModelID } @@ -347,6 +350,9 @@ func (h *Handler) UpdateModel(c *gin.Context) { m.InputBindingJSON = string(b) } prov, _ := h.cache.GetProvider(m.ProviderID) + if prov.Category == model.CategoryImage { + m.Alias = "" + } if prov.Category == model.CategoryImage && m.WorkflowJSON != "" { if err := validateImageModelCreate(struct { ModelID, DisplayName string diff --git a/internal/prediction/inject.go b/internal/prediction/inject.go index 42e185a..baf9bd2 100644 --- a/internal/prediction/inject.go +++ b/internal/prediction/inject.go @@ -22,10 +22,7 @@ func BuildPrompt(workflowJSON string, binding InputBinding, input map[string]any continue } if name == "seed" { - if iv, ok := val.(int); ok && iv < 0 { - n, _ := rand.Int(rand.Reader, big.NewInt(1<<31-1)) - val = int(n.Int64()) - } + val = resolveSeed(val) } node, ok := prompt[target.Node].(map[string]any) if !ok { @@ -37,9 +34,61 @@ func BuildPrompt(workflowJSON string, binding InputBinding, input map[string]any } inputs[target.Field] = val } + sanitizeSamplerSeeds(prompt) return prompt, nil } +// resolveSeed maps seed < 0 (random) to a non-negative value ComfyUI accepts. +func resolveSeed(val any) any { + n, ok := seedAsInt64(val) + if !ok || n >= 0 { + return val + } + r, err := rand.Int(rand.Reader, big.NewInt(1<<31-1)) + if err != nil { + return int(0) + } + return int(r.Int64()) +} + +func seedAsInt64(v any) (int64, bool) { + switch n := v.(type) { + case int: + return int64(n), true + case int64: + return n, true + case float64: + return int64(n), true + case json.Number: + i, err := n.Int64() + return i, err == nil + default: + return 0, false + } +} + +func sanitizeSamplerSeeds(prompt map[string]any) { + seedFields := []string{"seed", "noise_seed"} + for _, raw := range prompt { + node, ok := raw.(map[string]any) + if !ok { + continue + } + if class, _ := node["class_type"].(string); class != "KSampler" { + continue + } + inputs, ok := node["inputs"].(map[string]any) + if !ok { + continue + } + for _, field := range seedFields { + if v, ok := inputs[field]; ok { + inputs[field] = resolveSeed(v) + } + } + } +} + func deepCopyMap(src map[string]any) map[string]any { b, _ := json.Marshal(src) var dst map[string]any diff --git a/internal/prediction/inject_test.go b/internal/prediction/inject_test.go new file mode 100644 index 0000000..2eaf384 --- /dev/null +++ b/internal/prediction/inject_test.go @@ -0,0 +1,75 @@ +package prediction + +import ( + "encoding/json" + "testing" +) + +func TestResolveSeedNegativeInt(t *testing.T) { + got := resolveSeed(-1) + n, ok := got.(int) + if !ok || n < 0 { + t.Fatalf("expected non-negative int, got %v", got) + } +} + +func TestResolveSeedNegativeFloat64(t *testing.T) { + got := resolveSeed(float64(-1)) + n, ok := got.(int) + if !ok || n < 0 { + t.Fatalf("expected non-negative int from float64 -1, got %v", got) + } +} + +func TestResolveSeedExplicit(t *testing.T) { + if got := resolveSeed(42); got != 42 { + t.Fatalf("expected 42, got %v", got) + } +} + +func TestBuildPromptSanitizesUnboundKSamplerSeed(t *testing.T) { + workflow := map[string]any{ + "17": map[string]any{ + "class_type": "KSampler", + "inputs": map[string]any{ + "seed": float64(-1), + "steps": 20, + }, + }, + } + raw, _ := json.Marshal(workflow) + prompt, err := BuildPrompt(string(raw), InputBinding{}, map[string]any{}) + if err != nil { + t.Fatal(err) + } + node := prompt["17"].(map[string]any) + inputs := node["inputs"].(map[string]any) + seed, ok := inputs["seed"].(int) + if !ok || seed < 0 { + t.Fatalf("expected sanitized seed, got %v", inputs["seed"]) + } +} + +func TestBuildPromptSeedBindingAfterCloneMap(t *testing.T) { + workflow := map[string]any{ + "17": map[string]any{ + "class_type": "KSampler", + "inputs": map[string]any{"seed": 0}, + }, + } + raw, _ := json.Marshal(workflow) + binding := InputBinding{ + "seed": {Node: "17", Field: "seed"}, + } + // cloneMap round-trip produces float64, as in execute(). + input := cloneMap(map[string]any{"seed": -1}) + prompt, err := BuildPrompt(string(raw), binding, input) + if err != nil { + t.Fatal(err) + } + inputs := prompt["17"].(map[string]any)["inputs"].(map[string]any) + seed, ok := inputs["seed"].(int) + if !ok || seed < 0 { + t.Fatalf("expected non-negative seed after binding, got %v", inputs["seed"]) + } +} diff --git a/web/src/components/WorkflowModelDialog.vue b/web/src/components/WorkflowModelDialog.vue index f59eac3..9a62e63 100644 --- a/web/src/components/WorkflowModelDialog.vue +++ b/web/src/components/WorkflowModelDialog.vue @@ -1,11 +1,9 @@