From 96e8318f06274f842564400842d067dc769e1efa Mon Sep 17 00:00:00 2001 From: atridadl Date: Fri, 19 Jan 2024 14:08:13 -0700 Subject: [PATCH] Added S3 backup of images generated --- go.mod | 2 ++ go.sum | 9 ++++++++ lib/replicate.go | 22 ++++++++++++++++++++ lib/s3.go | 53 ++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 86 insertions(+) create mode 100644 lib/s3.go diff --git a/go.mod b/go.mod index 3a2d36b..4a223b5 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index eab30d0..1ec78c4 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/lib/replicate.go b/lib/replicate.go index cf46095..b8670df 100644 --- a/lib/replicate.go +++ b/lib/replicate.go @@ -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 } diff --git a/lib/s3.go b/lib/s3.go new file mode 100644 index 0000000..11e8a2f --- /dev/null +++ b/lib/s3.go @@ -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 +}