70 lines
1.6 KiB
Go
70 lines
1.6 KiB
Go
package main
|
|
|
|
import (
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
|
|
"waifu_gallery/runpod"
|
|
"waifu_gallery/s3"
|
|
)
|
|
|
|
var config Config
|
|
var InfoLogger = log.New(os.Stdout, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile)
|
|
var ErrorLogger = log.New(os.Stderr, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile)
|
|
|
|
func main() {
|
|
// Read config file
|
|
err := parseConfig("config.json", &config)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
// Initialize Runpod
|
|
if len(config.RunpodApi.Endpoint) <= 0 || len(config.RunpodApi.Authorization) <= 0 {
|
|
log.Fatal("Runpod credentials not set")
|
|
}
|
|
runpod.Endpoint = config.RunpodApi.Endpoint
|
|
runpod.Authorization = config.RunpodApi.Authorization
|
|
// Check if Runpod serverless API is up
|
|
_, err = runpod.Health()
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
// Initialize S3
|
|
err = s3.InitMinio(s3.Config(config.S3))
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
// Check if we can access the bucket
|
|
_, err = s3.GetMetadata("could_be_any_object")
|
|
if err != nil {
|
|
if s3.ToErrorResponse(err).Code != "NoSuchKey" {
|
|
log.Fatal(err)
|
|
}
|
|
}
|
|
// Check chatbot credentials / setup
|
|
if config.state.chatbotPromptEnabled {
|
|
_, err = getCompletion("Write a short haiku (10 words max)", config.Chatbot.PromptCompletion)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
}
|
|
if config.state.descriptionEnabled {
|
|
_, err = getCompletion("Write a short haiku (10 words max)", config.Chatbot.DescriptionCompletion)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
}
|
|
|
|
// http handler
|
|
http.HandleFunc("/img.png", servePicture)
|
|
http.HandleFunc("/generate", serveGenerationEndpoint)
|
|
http.HandleFunc("/", serveGallery)
|
|
|
|
err = http.ListenAndServe(config.ListenAddress, nil)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
}
|