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

Added S3 backup of images generated

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

2
go.mod
View file

@ -5,11 +5,13 @@ go 1.21.6
require github.com/diamondburned/arikawa/v3 v3.3.4
require (
github.com/jmespath/go-jmespath v0.4.0 // indirect
golang.org/x/net v0.20.0 // indirect
golang.org/x/sync v0.6.0 // indirect
)
require (
github.com/aws/aws-sdk-go v1.50.0
github.com/gorilla/schema v1.2.1 // indirect
github.com/gorilla/websocket v1.5.1 // indirect
github.com/joho/godotenv v1.5.1

9
go.sum
View file

@ -1,3 +1,6 @@
github.com/aws/aws-sdk-go v1.50.0 h1:HBtrLeO+QyDKnc3t1+5DR1RxodOHCGr8ZcrHudpv7jI=
github.com/aws/aws-sdk-go v1.50.0/go.mod h1:LF8svs817+Nz+DmiMQKTO3ubZ/6IaTpq3TjupRn3Eqk=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/diamondburned/arikawa/v3 v3.3.4 h1:UXOjM7PRlWLJ8kVAydX/VetqV7W4/d4xU92JRy3SpU4=
@ -8,12 +11,16 @@ github.com/gorilla/schema v1.2.1/go.mod h1:Dg5SSm5PV60mhF2NFaTV1xuYYj8tV8NOPRo4F
github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
github.com/gorilla/websocket v1.5.1 h1:gmztn0JnHVt9JZquRuzLw3g4wouNVzKL15iLr/zn/QY=
github.com/gorilla/websocket v1.5.1/go.mod h1:x3kM2JMyaluk02fnUJpQuwD2dCS5NDG2ZHL0uE0tcaY=
github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=
github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U=
github.com/joho/godotenv v1.5.1 h1:7eLL/+HRGLY0ldzfGMeQkb7vMd0as4CfYvUVzLqw0N0=
github.com/joho/godotenv v1.5.1/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/replicate/replicate-go v0.14.2 h1:XgK+REvYrWs7qDeyugxHA93h31qBhEFk/3p1/p2w3W8=
github.com/replicate/replicate-go v0.14.2/go.mod h1:otIrl1vDmyjNhTzmVmp/mQU3Wt1+3387gFNEsAZq0ig=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
@ -51,5 +58,7 @@ golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGm
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

22
lib/replicate.go
View file

@ -5,7 +5,9 @@ import (
"context"
"errors"
"io"
"log"
"net/http"
"os"
"strings"
"github.com/replicate/replicate-go"
@ -99,6 +101,26 @@ func ReplicateImageGeneration(prompt string) (*bytes.Buffer, error) {
return nil, imgReadErr
}
// Save image to a temporary file
tmpfile, err := os.CreateTemp("", "image.*.jpg")
if err != nil {
log.Fatal(err)
}
defer os.Remove(tmpfile.Name())
if _, err := tmpfile.Write(imageBytes); err != nil {
log.Fatal(err)
}
if err := tmpfile.Close(); err != nil {
log.Fatal(err)
}
// Upload the image to S3
_, uploadErr := UploadToS3(tmpfile.Name())
if uploadErr != nil {
log.Printf("Failed to upload image to S3: %v", uploadErr)
}
imageFile := bytes.NewBuffer(imageBytes)
return imageFile, nil
}

53
lib/s3.go Normal file
View file

@ -0,0 +1,53 @@
package lib
import (
"fmt"
"os"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/credentials"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-sdk-go/service/s3/s3manager"
)
func UploadToS3(filePath string) (*s3manager.UploadOutput, error) {
bucket := os.Getenv("BUCKET_NAME")
if bucket == "" {
fmt.Println("No S3 bucket specified, skipping upload.")
return nil, nil
}
endpoint := os.Getenv("AWS_ENDPOINT_URL_S3")
accessKeyID := os.Getenv("AWS_ACCESS_KEY_ID")
secretAccessKey := os.Getenv("AWS_SECRET_ACCESS_KEY")
sess, err := session.NewSession(&aws.Config{
Credentials: credentials.NewStaticCredentials(
accessKeyID,
secretAccessKey,
"",
),
Endpoint: aws.String(endpoint),
})
if err != nil {
return nil, fmt.Errorf("failed to create session, %v", err)
}
file, err := os.Open(filePath)
if err != nil {
return nil, fmt.Errorf("failed to open file, %v", err)
}
defer file.Close()
uploader := s3manager.NewUploader(sess)
result, err := uploader.Upload(&s3manager.UploadInput{
Bucket: aws.String(bucket),
Key: aws.String(filePath),
Body: file,
})
if err != nil {
return nil, fmt.Errorf("failed to upload file, %v", err)
}
return result, nil
}