Skip to content

Commit

Permalink
support roop!!! fix description
Browse files Browse the repository at this point in the history
  • Loading branch information
SpenserCai committed Aug 22, 2023
1 parent f47e7e3 commit 086b48c
Show file tree
Hide file tree
Showing 8 changed files with 170 additions and 8 deletions.
3 changes: 2 additions & 1 deletion dbot/slash_command.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
* @Date: 2023-08-16 22:10:00
* @version:
* @LastEditors: SpenserCai
* @LastEditTime: 2023-08-20 13:42:49
* @LastEditTime: 2023-08-22 13:18:13
* @Description: file content
*/
package dbot
Expand Down Expand Up @@ -43,4 +43,5 @@ func (dbot *DiscordBot) GenerateCommandList() {
dbot.AppCommand = append(dbot.AppCommand, slash_handler.SlashHandler{}.ExtraSingleOptions())
dbot.AppCommand = append(dbot.AppCommand, slash_handler.SlashHandler{}.PngInfoOptions())
dbot.AppCommand = append(dbot.AppCommand, slash_handler.SlashHandler{}.ControlnetDetectOptions())
dbot.AppCommand = append(dbot.AppCommand, slash_handler.SlashHandler{}.RoopImageOptions())
}
4 changes: 2 additions & 2 deletions dbot/slash_handler/controlnet_detect.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
* @Date: 2023-08-20 12:45:58
* @version:
* @LastEditors: SpenserCai
* @LastEditTime: 2023-08-20 14:52:44
* @LastEditTime: 2023-08-22 14:36:15
* @Description: file content
*/

Expand Down Expand Up @@ -84,7 +84,7 @@ func (shdl SlashHandler) controlnetModuleChoice() []*discordgo.ApplicationComman
func (shdl SlashHandler) ControlnetDetectOptions() *discordgo.ApplicationCommand {
return &discordgo.ApplicationCommand{
Name: "controlnet_detect",
Description: "Remove background from image",
Description: "ControlNet detect",
Options: []*discordgo.ApplicationCommandOption{
{
Type: discordgo.ApplicationCommandOptionString,
Expand Down
4 changes: 2 additions & 2 deletions dbot/slash_handler/extra_single.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
* @Date: 2023-08-19 16:21:45
* @version:
* @LastEditors: SpenserCai
* @LastEditTime: 2023-08-19 18:48:06
* @LastEditTime: 2023-08-22 14:35:37
* @Description: file content
*/
package slash_handler
Expand Down Expand Up @@ -40,7 +40,7 @@ func (shdl SlashHandler) upscalerModelChoise() []*discordgo.ApplicationCommandOp
func (shdl SlashHandler) ExtraSingleOptions() *discordgo.ApplicationCommand {
return &discordgo.ApplicationCommand{
Name: "extra_single",
Description: "Remove background from image",
Description: "Upscaler and face restorer for single image",
Options: []*discordgo.ApplicationCommandOption{
{
Type: discordgo.ApplicationCommandOptionString,
Expand Down
36 changes: 36 additions & 0 deletions dbot/slash_handler/face_restorer.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
/*
* @Author: SpenserCai
* @Date: 2023-08-22 13:10:36
* @version:
* @LastEditors: SpenserCai
* @LastEditTime: 2023-08-22 13:12:54
* @Description: file content
*/
package slash_handler

import (
"log"

"github.com/SpenserCai/sd-webui-discord/global"

"github.com/SpenserCai/sd-webui-go/intersvc"
"github.com/bwmarrin/discordgo"
)

func (shdl SlashHandler) faceRestorerModelChoice() []*discordgo.ApplicationCommandOptionChoice {
choices := []*discordgo.ApplicationCommandOptionChoice{}
modesvc := &intersvc.SdapiV1FaceRestorers{}
modesvc.Action(global.ClusterManager.GetNodeAuto().StableClient)
if modesvc.Error != nil {
log.Println(modesvc.Error)
return choices
}
models := modesvc.GetResponse()
for _, model := range *models {
choices = append(choices, &discordgo.ApplicationCommandOptionChoice{
Name: *model.Name,
Value: *model.Name,
})
}
return choices
}
4 changes: 2 additions & 2 deletions dbot/slash_handler/png_info.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
* @Date: 2023-08-19 18:27:34
* @version:
* @LastEditors: SpenserCai
* @LastEditTime: 2023-08-19 20:51:52
* @LastEditTime: 2023-08-22 14:36:40
* @Description: file content
*/
package slash_handler
Expand All @@ -24,7 +24,7 @@ import (
func (shdl SlashHandler) PngInfoOptions() *discordgo.ApplicationCommand {
return &discordgo.ApplicationCommand{
Name: "png_info",
Description: "Remove background from image",
Description: "Get image info",
Options: []*discordgo.ApplicationCommandOption{
{
Type: discordgo.ApplicationCommandOptionString,
Expand Down
121 changes: 121 additions & 0 deletions dbot/slash_handler/roop_image.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,121 @@
/*
* @Author: SpenserCai
* @Date: 2023-08-22 12:58:13
* @version:
* @LastEditors: SpenserCai
* @LastEditTime: 2023-08-22 14:38:09
* @Description: file content
*/
package slash_handler

import (
"log"

"github.com/SpenserCai/sd-webui-discord/cluster"
"github.com/SpenserCai/sd-webui-discord/global"
"github.com/SpenserCai/sd-webui-discord/utils"

"github.com/SpenserCai/sd-webui-go/intersvc"
"github.com/bwmarrin/discordgo"
)

func (shdl SlashHandler) RoopImageOptions() *discordgo.ApplicationCommand {
return &discordgo.ApplicationCommand{
Name: "roop_image",
Description: "Image face swap with roop",
Options: []*discordgo.ApplicationCommandOption{
{
Type: discordgo.ApplicationCommandOptionString,
Name: "source_url",
Description: "The url of the source face image",
Required: true,
},
{
Type: discordgo.ApplicationCommandOptionString,
Name: "target_url",
Description: "The url of the target image",
Required: true,
},
{
Type: discordgo.ApplicationCommandOptionString,
Name: "face_restorer",
Description: "The face restorer model to use",
Required: false,
Choices: shdl.faceRestorerModelChoice(),
},
{
Type: discordgo.ApplicationCommandOptionNumber,
Name: "restorer_visibility",
Description: "The visibility of the face restorer 0-1",
Required: false,
},
},
}
}

func (shdl SlashHandler) RoopImageSetOptions(dsOpt []*discordgo.ApplicationCommandInteractionDataOption, opt *intersvc.RoopImageRequest) {
opt.Model = func() *string { v := "inswapper_128.onnx"; return &v }()
opt.FaceIndex = []int64{0}
opt.Scale = func() *int64 { var v int64 = 1; return &v }()
opt.UpscaleVisibility = func() *float64 { var v float64 = 1.0; return &v }()
opt.RestorerVisibility = func() *float64 { v := 1.0; return &v }()
for _, v := range dsOpt {
switch v.Name {
case "source_url":
opt.SourceImage, _ = utils.GetImageBase64(v.StringValue())
case "target_url":
opt.TargetImage, _ = utils.GetImageBase64(v.StringValue())
case "face_restorer":
opt.FaceRestorer = func() *string { v := v.StringValue(); return &v }()
case "restorer_visibility":
opt.RestorerVisibility = func() *float64 { v := v.FloatValue(); return &v }()
}
}
}

func (shdl SlashHandler) RoopImageAction(s *discordgo.Session, i *discordgo.InteractionCreate, opt *intersvc.RoopImageRequest, node *cluster.ClusterNode) {
msg, err := shdl.SendStateMessage("Running", s, i)
if err != nil {
log.Println(err)
return
}
roop_image := &intersvc.RoopImage{RequestItem: opt}
roop_image.Action(node.StableClient)
if roop_image.Error != nil {
s.FollowupMessageEdit(i.Interaction, msg.ID, &discordgo.WebhookEdit{
Content: func() *string { v := roop_image.Error.Error(); return &v }(),
})
} else {
image, err := utils.GetImageReaderByBase64(roop_image.GetResponse().Image)
if err != nil {
s.FollowupMessageEdit(i.Interaction, msg.ID, &discordgo.WebhookEdit{
Content: func() *string { v := err.Error(); return &v }(),
})
} else {
s.FollowupMessageEdit(i.Interaction, msg.ID, &discordgo.WebhookEdit{
Content: func() *string { v := "Success"; return &v }(),
Files: []*discordgo.File{
{
Name: "roop_image.png",
ContentType: "image/png",
Reader: image,
},
},
})
}

}
}

func (shdl SlashHandler) RoopImageCommandHandler(s *discordgo.Session, i *discordgo.InteractionCreate) {
option := &intersvc.RoopImageRequest{}
shdl.ReportCommandInfo(s, i)
node := global.ClusterManager.GetNodeAuto()
action := func() (map[string]interface{}, error) {
shdl.RoopImageSetOptions(i.ApplicationCommandData().Options, option)
shdl.RoopImageAction(s, i, option, node)
return nil, nil
}
callback := func() {}
node.ActionQueue.AddTask(shdl.GenerateTaskID(i), action, callback)
}
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ module github.com/SpenserCai/sd-webui-discord
go 1.19

require (
github.com/SpenserCai/sd-webui-go v0.2.4
github.com/SpenserCai/sd-webui-go v0.3.0
github.com/bwmarrin/discordgo v0.27.1
golang.org/x/text v0.7.0
)
Expand Down
4 changes: 4 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@ github.com/PuerkitoBio/purell v1.1.1/go.mod h1:c11w/QuzBsJSee3cPx9rAFu61PvFxuPbt
github.com/PuerkitoBio/urlesc v0.0.0-20170810143723-de5bf2ad4578/go.mod h1:uGdkoq3SwY9Y+13GIhn11/XLaGBb4BfwItxLd5jeuXE=
github.com/SpenserCai/sd-webui-go v0.2.4 h1:yjp1St8oAJ6lq6I8sYre+0ftETFxGCl6Jq6rsWMEXWo=
github.com/SpenserCai/sd-webui-go v0.2.4/go.mod h1:ckzphVZudQu3kYx124QUggeyAWIOW0fMJj2e4yEKgu0=
github.com/SpenserCai/sd-webui-go v0.2.5 h1:UySzosM+q+32+x7z1IL3eRkQJ/qVgwv37OgYN9ag2Xg=
github.com/SpenserCai/sd-webui-go v0.2.5/go.mod h1:ckzphVZudQu3kYx124QUggeyAWIOW0fMJj2e4yEKgu0=
github.com/SpenserCai/sd-webui-go v0.3.0 h1:6eIdVG2nO2xLhklYFaIclZyfRvpgRXst5rZyNEAlgoI=
github.com/SpenserCai/sd-webui-go v0.3.0/go.mod h1:ckzphVZudQu3kYx124QUggeyAWIOW0fMJj2e4yEKgu0=
github.com/asaskevich/govalidator v0.0.0-20200907205600-7a23bdc65eef/go.mod h1:WaHUgvxTVq04UNunO+XhnAqY/wQc+bxr74GqbsZ/Jqw=
github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 h1:DklsrG3dyBCFEj5IhUbnKptjxatkF07cF2ak3yi77so=
github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2/go.mod h1:WaHUgvxTVq04UNunO+XhnAqY/wQc+bxr74GqbsZ/Jqw=
Expand Down

0 comments on commit 086b48c

Please sign in to comment.