Skip to content

Commit 9af293f

Browse files
inciner8rinciner8r
committed
fix upscale, variation commands
1 parent 359cf05 commit 9af293f

File tree

4 files changed

+60
-64
lines changed

4 files changed

+60
-64
lines changed

‎app/app.go

+9-3
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,14 @@ import (
44
"fmt"
55
"log"
66
"os"
7+
"os/signal"
8+
"syscall"
79

810
"github.com/MyriadFlow/airbot/app/handlers"
911
"github.com/bwmarrin/discordgo"
1012
)
1113

12-
func Init() *discordgo.Session {
14+
func Init() {
1315
bot_token := os.Getenv("BOT_TOKEN")
1416

1517
sess, err := discordgo.New("Bot " + bot_token)
@@ -25,8 +27,12 @@ func Init() *discordgo.Session {
2527
log.Fatal(err)
2628
}
2729
defer sess.Close()
28-
30+
// Wait here until CTRL-C or other term signal is received.
2931
fmt.Println("Bot is now running. Press CTRL-C to exit.")
30-
return sess
32+
sc := make(chan os.Signal, 1)
33+
signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt, os.Kill)
34+
<-sc
3135

36+
// Cleanly close down the Discord session.
37+
sess.Close()
3238
}

‎app/handlers/functions.go

+35-27
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,6 @@ func Generate(prompt string) {
5454
"attachments": []
5555
}
5656
}`
57-
fmt.Println(jsonStr)
5857
req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(jsonStr))
5958
if err != nil {
6059
log.Fatal(err)
@@ -92,10 +91,9 @@ func Upscale(number int, messageid string, imageid string) {
9291
"session_id": "937a1c8132cd7ce3940aa8f59dedf961",
9392
"data": {
9493
"component_type": 2,
95-
"custom_id": "MJ::JOB::upsample::"` + numberString + `"::"` + imageid + `"
94+
"custom_id": "MJ::JOB::upsample::` + numberString + `::` + imageid + `"
9695
}
9796
}`
98-
fmt.Println(jsonStr)
9997
req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(jsonStr))
10098
if err != nil {
10199
log.Fatal(err)
@@ -109,14 +107,9 @@ func Upscale(number int, messageid string, imageid string) {
109107
panic(err)
110108
}
111109
defer resp.Body.Close()
112-
113-
fmt.Println("response Status:", resp.Status)
114-
fmt.Println("response Headers:", resp.Header)
115-
body, _ := io.ReadAll(resp.Body)
116-
fmt.Println("response Body:", string(body))
117110
}
118111

119-
func UpscaleMax(number int, messageid string, imageid string) {
112+
func UpscaleMax(messageid string, imageid string) {
120113
url := "https://discord.com/api/v9/interactions"
121114
server_id := os.Getenv("SERVER_ID")
122115
user_token := os.Getenv("USER_TOKEN")
@@ -131,10 +124,9 @@ func UpscaleMax(number int, messageid string, imageid string) {
131124
"session_id": "1f3dbdf09efdf93d81a3a6420882c92c",
132125
"data": {
133126
"component_type": 2,
134-
"custom_id": "MJ::JOB::upsample_max::1::"` + imageid + `"::SOLO"
127+
"custom_id": "MJ::JOB::upsample_max::1::` + imageid + `::SOLO"
135128
}
136129
}`
137-
fmt.Println(jsonStr)
138130
req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(jsonStr))
139131
if err != nil {
140132
log.Fatal(err)
@@ -148,11 +140,6 @@ func UpscaleMax(number int, messageid string, imageid string) {
148140
panic(err)
149141
}
150142
defer resp.Body.Close()
151-
152-
fmt.Println("response Status:", resp.Status)
153-
fmt.Println("response Headers:", resp.Header)
154-
body, _ := io.ReadAll(resp.Body)
155-
fmt.Println("response Body:", string(body))
156143
}
157144

158145
func Variation(number int, messageid string, imageid string) {
@@ -172,10 +159,9 @@ func Variation(number int, messageid string, imageid string) {
172159
"session_id": "937a1c8132cd7ce3940aa8f59dedf961",
173160
"data": {a
174161
"component_type": 2,
175-
"custom_id": "MJ::JOB::variation::"` + numberString + `"::"` + imageid + `"
162+
"custom_id": "MJ::JOB::variation::` + numberString + `::` + imageid + `"
176163
}
177164
}`
178-
fmt.Println(jsonStr)
179165
req, err := http.NewRequest(http.MethodPost, url, strings.NewReader(jsonStr))
180166
if err != nil {
181167
log.Fatal(err)
@@ -196,21 +182,43 @@ func Variation(number int, messageid string, imageid string) {
196182
fmt.Println("response Body:", string(body))
197183
}
198184

199-
func getImageURLByMessageID(s *discordgo.Session, channelID, messageID string) (string, error) {
200-
message, err := s.ChannelMessage(channelID, messageID)
185+
func getImageFromMessageID(session *discordgo.Session, channelID, messageID string) (string, string, error) {
186+
message, err := session.ChannelMessage(channelID, messageID)
201187
if err != nil {
202-
return "", err
188+
return "", "", fmt.Errorf("error retrieving message: %s", err)
203189
}
204190

205-
imageURL := ""
191+
var imageURL string
192+
var imageID string
206193

194+
// Check if there are any attachments in the message
207195
if len(message.Attachments) > 0 {
208-
imageURL = message.Attachments[0].URL
196+
attachment := message.Attachments[0]
197+
imageURL = attachment.URL
198+
imageID = attachment.ID
199+
} else {
200+
// Check if there are any embeds in the message
201+
if len(message.Embeds) > 0 {
202+
embed := message.Embeds[0]
203+
204+
// Check if the embed contains an image
205+
if len(embed.Image.URL) > 0 {
206+
imageURL = embed.Image.URL
207+
}
208+
209+
// Check if the embed contains an image ID
210+
if len(embed.Image.ProxyURL) > 0 {
211+
imageID = embed.Image.ProxyURL
212+
}
213+
}
209214
}
210215

211-
if len(message.Embeds) > 0 && len(message.Embeds[0].Image.URL) > 0 {
212-
imageURL = message.Embeds[0].Image.URL
213-
}
216+
return imageURL, imageID, nil
217+
}
214218

215-
return imageURL, nil
219+
func getImageId(url string) string {
220+
arr := strings.Split(url, "_")
221+
png := arr[len(arr)-1]
222+
imageId := strings.Split(png, ".")[0]
223+
return imageId
216224
}

‎app/handlers/handlers.go

+15-22
Original file line numberDiff line numberDiff line change
@@ -63,48 +63,41 @@ func AddHandlers(sess *discordgo.Session) {
6363
return
6464
}
6565
if args[1] == "upscale" {
66-
if m.Message.Reference() != nil {
67-
repliedMessageID := m.Message.Reference().MessageID
68-
fmt.Println(repliedMessageID)
69-
imageID, err := getImageURLByMessageID(s, os.Getenv("CHANNEL_ID"), repliedMessageID)
66+
if m.MessageReference.MessageID != "" {
67+
repliedMessageID := m.MessageReference.MessageID
68+
imageURL, _, err := getImageFromMessageID(s, os.Getenv("CHANNEL_ID"), repliedMessageID)
69+
imageID := getImageId(imageURL)
7070
if err != nil {
7171
fmt.Println("error", err)
7272
return
7373
}
74-
fmt.Println("error", err)
75-
fmt.Println("imageid", imageID)
7674
number, _ := strconv.Atoi(args[2])
77-
Upscale(number, repliedMessageID, args[3])
75+
Upscale(number, repliedMessageID, imageID)
7876
}
7977
}
8078
if args[1] == "variation" {
81-
if m.Message.Reference() != nil {
82-
repliedMessageID := m.Message.Reference().MessageID
83-
fmt.Println(repliedMessageID)
84-
imageID, err := getImageURLByMessageID(s, os.Getenv("CHANNEL_ID"), repliedMessageID)
79+
if m.MessageReference.MessageID != "" {
80+
repliedMessageID := m.MessageReference.MessageID
81+
imageURL, _, err := getImageFromMessageID(s, os.Getenv("CHANNEL_ID"), repliedMessageID)
82+
imageID := getImageId(imageURL)
8583
if err != nil {
8684
fmt.Println("error", err)
8785
return
8886
}
89-
fmt.Println("error", err)
90-
fmt.Println("imageid", imageID)
9187
number, _ := strconv.Atoi(args[2])
92-
Variation(number, repliedMessageID, args[3])
88+
Variation(number, repliedMessageID, imageID)
9389
}
9490
}
9591
if args[1] == "maxupscale" {
96-
if m.Message.Reference() != nil {
97-
repliedMessageID := m.Message.Reference().MessageID
98-
fmt.Println(repliedMessageID)
99-
imageID, err := getImageURLByMessageID(s, os.Getenv("CHANNEL_ID"), repliedMessageID)
92+
if m.MessageReference.MessageID != "" {
93+
repliedMessageID := m.MessageReference.MessageID
94+
imageURL, _, err := getImageFromMessageID(s, os.Getenv("CHANNEL_ID"), repliedMessageID)
95+
imageID := getImageId(imageURL)
10096
if err != nil {
10197
fmt.Println("error", err)
10298
return
10399
}
104-
fmt.Println("error", err)
105-
fmt.Println("imageid", imageID)
106-
number, _ := strconv.Atoi(args[2])
107-
Upscale(number, repliedMessageID, args[3])
100+
UpscaleMax(repliedMessageID, imageID)
108101
}
109102
}
110103

‎main.go

+1-12
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,11 @@
11
package main
22

33
import (
4-
"os"
5-
"os/signal"
6-
"sync"
7-
"syscall"
8-
94
"github.com/MyriadFlow/airbot/app"
105
"github.com/joho/godotenv"
116
)
127

13-
var wg sync.WaitGroup
14-
158
func main() {
169
godotenv.Load()
17-
sess := app.Init()
18-
sc := make(chan os.Signal, 1)
19-
signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt, os.Kill)
20-
<-sc
21-
sess.Close()
10+
app.Init()
2211
}

0 commit comments

Comments
 (0)