100 lines
2.5 KiB
Go
100 lines
2.5 KiB
Go
package lib
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/replicate/replicate-go"
|
|
)
|
|
|
|
func ReplicateTextGeneration(prompt string) (string, error) {
|
|
client, clientError := replicate.NewClient(replicate.WithTokenFromEnv())
|
|
if clientError != nil {
|
|
return "", clientError
|
|
}
|
|
|
|
input := replicate.PredictionInput{
|
|
"prompt": prompt,
|
|
}
|
|
webhook := replicate.Webhook{
|
|
URL: "https://example.com/webhook",
|
|
Events: []replicate.WebhookEventType{"start", "completed"},
|
|
}
|
|
|
|
prediction, predictionError := client.Run(context.Background(), "mistralai/mixtral-8x7b-instruct-v0.1:7b3212fbaf88310cfef07a061ce94224e82efc8403c26fc67e8f6c065de51f21", input, &webhook)
|
|
|
|
if predictionError != nil {
|
|
return "", predictionError
|
|
}
|
|
|
|
test, ok := prediction.([]interface{})
|
|
|
|
if !ok {
|
|
return "", errors.New("there was an error generating the image based on this prompt... this usually happens when the generated image violates safety requirements")
|
|
}
|
|
|
|
strs := make([]string, len(test))
|
|
for i, v := range test {
|
|
str, ok := v.(string)
|
|
if !ok {
|
|
return "", errors.New("element is not a string")
|
|
}
|
|
strs[i] = str
|
|
}
|
|
|
|
result := strings.Join(strs, "")
|
|
|
|
return result, nil
|
|
}
|
|
|
|
func ReplicateImageGeneration(prompt string) (*bytes.Buffer, error) {
|
|
client, clientError := replicate.NewClient(replicate.WithTokenFromEnv())
|
|
if clientError != nil {
|
|
return nil, clientError
|
|
}
|
|
|
|
input := replicate.PredictionInput{
|
|
"prompt": prompt,
|
|
}
|
|
webhook := replicate.Webhook{
|
|
URL: "https://example.com/webhook",
|
|
Events: []replicate.WebhookEventType{"start", "completed"},
|
|
}
|
|
|
|
prediction, predictionError := client.Run(context.Background(), "stability-ai/sdxl:c221b2b8ef527988fb59bf24a8b97c4561f1c671f73bd389f866bfb27c061316", input, &webhook)
|
|
|
|
if predictionError != nil {
|
|
return nil, predictionError
|
|
}
|
|
|
|
test, ok := prediction.([]interface{})
|
|
|
|
if !ok {
|
|
return nil, errors.New("there was an error generating the image based on this prompt... this usually happens when the generated image violates safety requirements")
|
|
}
|
|
|
|
imgUrl, ok := test[0].(string)
|
|
|
|
if !ok {
|
|
return nil, errors.New("there was an error generating the image based on this prompt... this usually happens when the generated image violates safety requirements")
|
|
}
|
|
|
|
imageRes, imageGetErr := http.Get(imgUrl)
|
|
if imageGetErr != nil {
|
|
return nil, imageGetErr
|
|
}
|
|
|
|
defer imageRes.Body.Close()
|
|
|
|
imageBytes, imgReadErr := io.ReadAll(imageRes.Body)
|
|
if imgReadErr != nil {
|
|
return nil, imgReadErr
|
|
}
|
|
|
|
imageFile := bytes.NewBuffer(imageBytes)
|
|
return imageFile, nil
|
|
}
|