atridad/himbot
1
0
Fork 0
Code Projects Releases Packages 1 Activity Actions

Fixed S3 upload names

Browse source
This commit is contained in:
Atridad Lahiji 2024-01-19 14:33:37 -07:00
parent 96e8318f06
commit f31739a59b
No known key found for this signature in database
2 changed files with 20 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
lib/replicate.go
View file

@ -54,7 +54,7 @@ func ReplicateTextGeneration(prompt string) (string, error) {
return result, nil return result, nil
} }
func ReplicateImageGeneration(prompt string) (*bytes.Buffer, error) { func ReplicateImageGeneration(prompt string, filename string) (*bytes.Buffer, error) {
client, clientError := replicate.NewClient(replicate.WithTokenFromEnv()) client, clientError := replicate.NewClient(replicate.WithTokenFromEnv())
if clientError != nil { if clientError != nil {
return nil, clientError return nil, clientError
@ -102,10 +102,19 @@ func ReplicateImageGeneration(prompt string) (*bytes.Buffer, error) {
} }
// Save image to a temporary file // Save image to a temporary file
tmpfile, err := os.CreateTemp("", "image.*.jpg") var tmpfile *os.File
var err error
if filename != "" {
tmpfile, err = os.CreateTemp("", filename)
} else {
tmpfile, err = os.CreateTemp("", "image.*.jpg")
}
if err != nil { if err != nil {
log.Fatal(err) log.Fatal(err)
} }
defer os.Remove(tmpfile.Name()) defer os.Remove(tmpfile.Name())
if _, err := tmpfile.Write(imageBytes); err != nil { if _, err := tmpfile.Write(imageBytes); err != nil {

11
main.go
View file

@ -10,6 +10,7 @@ import (
"net" "net"
"os" "os"
"os/signal" "os/signal"
"strconv"
"time" "time"
"github.com/diamondburned/arikawa/v3/api" "github.com/diamondburned/arikawa/v3/api"
@ -183,7 +184,13 @@ func (h *handler) cmdPic(ctx context.Context, data cmdroute.CommandData) *api.In
return errorResponse(err) return errorResponse(err)
} }
imageFile, err := lib.ReplicateImageGeneration(options.Prompt) // Get current epoch timestamp
timestamp := strconv.FormatInt(time.Now().Unix(), 10)
// Concatenate clean username and timestamp to form filename
filename := data.Event.Sender().Username + "_" + timestamp + ".jpg"
imageFile, err := lib.ReplicateImageGeneration(options.Prompt, filename)
if err != nil { if err != nil {
lib.CancelCooldown(data.Event.User.ID.String(), "pic") lib.CancelCooldown(data.Event.User.ID.String(), "pic")
@ -191,7 +198,7 @@ func (h *handler) cmdPic(ctx context.Context, data cmdroute.CommandData) *api.In
} }
file := sendpart.File{ file := sendpart.File{
Name: "himbot_response.png", Name: filename,
Reader: imageFile, Reader: imageFile,
} }