fix sd_checkpoint mode setting

SpenserCai · Sep 4, 2023 · 2b3e2f3 · 2b3e2f3
1 parent 27e90a8
commit 2b3e2f3
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/dbot/slash_handler/img2img.go b/dbot/slash_handler/img2img.go
@@ -262,6 +262,8 @@ func (shdl SlashHandler) Img2imgSetOptions(cmd discordgo.ApplicationCommandInter
 	opt.InpaintFullResPadding = 32
 	resizeByScale := 1.0
 	isSetSize := false
+	isSetCheckpoints := false
+	defaultCheckpoints := shdl.GetDefaultSettingFromUser("sd_model_checkpoint", "", i).(string)
 	for _, v := range cmd.Options {
 		switch v.Name {
 		case "prompt":
@@ -302,6 +304,7 @@ func (shdl SlashHandler) Img2imgSetOptions(cmd discordgo.ApplicationCommandInter
 			tmpOverrideSettings := opt.OverrideSettings.(map[string]interface{})
 			tmpOverrideSettings["sd_model_checkpoint"] = v.StringValue()
 			opt.OverrideSettings = tmpOverrideSettings
+			isSetCheckpoints = true
 		case "refiner_checkpoint":
 			opt.RefinerCheckpoint = v.StringValue()
 		case "refiner_switch_at":
@@ -336,6 +339,11 @@ func (shdl SlashHandler) Img2imgSetOptions(cmd discordgo.ApplicationCommandInter
 			opt.Height = func() *int64 { v := int64(float64(height) * resizeByScale); return &v }()
 		}
 	}
+	if !isSetCheckpoints && defaultCheckpoints != "" {
+		tmpOverrideSettings := opt.OverrideSettings.(map[string]interface{})
+		tmpOverrideSettings["sd_model_checkpoint"] = defaultCheckpoints
+		opt.OverrideSettings = tmpOverrideSettings
+	}
 }
 
 func (shdl SlashHandler) Img2imgAction(s *discordgo.Session, i *discordgo.InteractionCreate, opt *intersvc.SdapiV1Img2imgRequest, node *cluster.ClusterNode) {

diff --git a/dbot/slash_handler/txt2img.go b/dbot/slash_handler/txt2img.go
@@ -3,7 +3,7 @@
  * @Date: 2023-08-22 17:13:19
  * @version:
  * @LastEditors: SpenserCai
- * @LastEditTime: 2023-09-03 23:06:53
+ * @LastEditTime: 2023-09-04 17:30:43
  * @Description: file content
  */
 package slash_handler
@@ -172,6 +172,8 @@ func (shdl SlashHandler) Txt2imgSetOptions(dsOpt []*discordgo.ApplicationCommand
 	opt.ScriptArgs = []interface{}{}
 	opt.AlwaysonScripts = map[string]interface{}{}
 	opt.OverrideSettings = map[string]interface{}{}
+	isSetCheckpoints := false
+	defaultCheckpoints := shdl.GetDefaultSettingFromUser("sd_model_checkpoint", "", i).(string)
 
 	for _, v := range dsOpt {
 		switch v.Name {
@@ -209,12 +211,18 @@ func (shdl SlashHandler) Txt2imgSetOptions(dsOpt []*discordgo.ApplicationCommand
 			tmpOverrideSettings := opt.OverrideSettings.(map[string]interface{})
 			tmpOverrideSettings["sd_model_checkpoint"] = v.StringValue()
 			opt.OverrideSettings = tmpOverrideSettings
+			isSetCheckpoints = true
 		case "refiner_checkpoint":
 			opt.RefinerCheckpoint = v.StringValue()
 		case "refiner_switch_at":
 			opt.RefinerSwitchAt = v.FloatValue()
 		}
 	}
+	if !isSetCheckpoints && defaultCheckpoints != "" {
+		tmpOverrideSettings := opt.OverrideSettings.(map[string]interface{})
+		tmpOverrideSettings["sd_model_checkpoint"] = defaultCheckpoints
+		opt.OverrideSettings = tmpOverrideSettings
+	}
 
 	// optJson, _ := json.Marshal(opt)
 	// log.Println(string(optJson))