mememan/waifu_gallery
mememan/waifu_gallery
1
0
Fork 0
Code Issues Pull requests Releases Activity

Separate generation endpoint from gallery, interactive loading

Browse source
This commit is contained in:
Mememan 2024-04-21 03:08:10 +02:00
parent f177dc2a31
commit 07e9b6ddb0
8 changed files with 214 additions and 95 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
README.md
View file

@ -24,5 +24,8 @@ Example:
- [x] Poll request
- [x] Guard from redundant API calls
- [x] Guard from excessive number of img generation jobs
- [ ] Make user wait for result interactively and reload
- [x] Separate img generation endpoint and call it from HTML
- [x] Make user wait for result interactively and reload
- [ ] Convert results to WebP
- [ ] Better querying via LLMs
- [ ] Better front-end

1
display.html
View file

@ -4,6 +4,7 @@
<head>
<title>{{ .Label.Title }}</title>
<meta charset="UTF-8" />
<meta name="darkreader-lock">
<meta name="viewport" content="width=device-width, initial-scale=1" />
<link rel="shortcut icon" type="image/png" href="{{ .Img.URL }}" />
<meta property="og:image" content="https://{{ .Img.FullURL }}" />

52
html.go → gallery.go
View file

@ -6,6 +6,7 @@ import (
"net/http"
"strconv"
"time"
"waifu_gallery/s3"
"golang.org/x/text/cases"
"golang.org/x/text/language"
@ -42,11 +43,54 @@ type Label struct {
Technique string
}
type PageContent struct {
type GalleryPageContent struct {
Img Img
Label Label
}
func serveGallery(w http.ResponseWriter, r *http.Request) {
// Check that the root got called
if r.RequestURI != "/" {
return
}
// Validate user input
input, err := handleSubdomain(w, r)
if err != nil {
InfoLogger.Println(err)
return
}
metadata, err := s3.GetMetadata(input)
if err != nil {
switch s3.ToErrorResponse(err).Code {
case "NoSuchKey":
// Display loading page and finish
w.Header().Set("Cache-Control", "no-cache")
w.Header().Set("Refresh", "20")
writeWaitingPage(w, input)
return
default:
serveInternalError(w)
ErrorLogger.Println(err)
return
}
}
// Picture exists, let's show it
err = writeDisplayPage(w, r, FileData{
URL: "img.png",
Name: metadata.Key,
Size: metadata.Size,
Date: metadata.LastModified,
Metadata: metadata.UserMetadata,
})
if err != nil {
serveBadSubdomainError(w)
ErrorLogger.Println(err)
}
}
func writeDisplayPage(w http.ResponseWriter, r *http.Request, m FileData) (err error) {
tmpl, err := template.New("display").Parse(displayTmpl)
if err != nil {
@ -63,7 +107,7 @@ func writeDisplayPage(w http.ResponseWriter, r *http.Request, m FileData) (err e
heightCm := convToCm(height, PPI)
// Template building
tmplData := PageContent{
tmplData := GalleryPageContent{
Img: Img{
URL: "img.png",
FullURL: r.Host + "/img.png",
@ -79,10 +123,6 @@ func writeDisplayPage(w http.ResponseWriter, r *http.Request, m FileData) (err e
},
}
err = tmpl.Execute(w, tmplData)
if err != nil {
ErrorLogger.Println(err)
return
}
return
}

96
http.go
View file

@ -14,93 +14,10 @@ import (
var generationQueue sync.Map
func serveGallery(w http.ResponseWriter, r *http.Request) {
// Check that the root got called
if r.RequestURI != "/" {
return
}
// Validate user input
input, err := handleSubdomain(w, r)
if err != nil {
InfoLogger.Println(err)
return
}
metadata, err := s3.GetMetaData(input)
if err != nil {
switch s3.ToErrorResponse(err).Code {
case "NoSuchKey":
// Generate picture or wait on existing generation job
// Probably not the best way but... if it works?
if ch, ok := generationQueue.Load(input); !ok {
// Check number of job requests
// TODO make this less horrible inefficient (lol)
var jobNbr int
generationQueue.Range(func(k, v interface{}) bool {
jobNbr++
return true
})
if config.MaxConcurrentJobs > 0 && jobNbr >= config.MaxConcurrentJobs {
w.WriteHeader(http.StatusTooManyRequests)
w.Write([]byte("Too many images currently being generated, please try again in a few minutes."))
return
}
// Create channel
generationQueue.Store(input, make(chan string))
// Generate new picture
InfoLogger.Printf("generating with input '%s'\n", input)
err := createPicture(input)
if err != nil {
w.WriteHeader(http.StatusInternalServerError)
ErrorLogger.Println(err)
return
}
val, _ := generationQueue.Load(input)
generationQueue.Delete(input)
if c, ok := val.(chan string); ok {
c <- input
}
} else {
// Wait for job to finish if there's already a generation going on
if c, ok := ch.(chan string); ok {
<-c
}
}
// get metadata again
metadata, err = s3.GetMetaData(input)
if err != nil {
serverError(w)
ErrorLogger.Println(err)
return
}
default:
serverError(w)
ErrorLogger.Println(err)
return
}
}
// Picture exists, let's show it
err = writeDisplayPage(w, r, FileData{
URL: "img.png",
Name: metadata.Key,
Size: metadata.Size,
Date: metadata.LastModified,
Metadata: metadata.UserMetadata,
})
if err != nil {
serverError(w)
ErrorLogger.Println(err)
}
}
// Display a picture from S3, greatly benefits from being cached
func servePicture(w http.ResponseWriter, r *http.Request) {
// TODO no need to check for bad words here
// Validate user input
input, err := handleSubdomain(w, r)
if err != nil {
@ -119,7 +36,12 @@ func servePicture(w http.ResponseWriter, r *http.Request) {
http.ServeContent(w, r, input+".png", time.Now(), obj)
}
func serverError(w http.ResponseWriter) {
func serveBadSubdomainError(w http.ResponseWriter) {
w.WriteHeader(http.StatusBadRequest)
w.Write([]byte("Invalid subdomain provided"))
}
func serveInternalError(w http.ResponseWriter) {
w.WriteHeader(http.StatusInternalServerError)
w.Write([]byte("Server encountered an error processing your request, sorry."))
}

40
loading.html Normal file
View file

File diff suppressed because one or more lines are too long

1
main.go
View file

@ -35,6 +35,7 @@ func main() {
// http handler
http.HandleFunc("/img.png", servePicture)
http.HandleFunc("/generate", serveGenerationEndpoint)
http.HandleFunc("/", serveGallery)
err = http.ListenAndServe(config.ListenAddress, nil)

2
s3/s3.go
View file

@ -39,7 +39,7 @@ func InitMinio(cfg Config) (err error) {
func ToErrorResponse(e error) minio.ErrorResponse {
return minio.ToErrorResponse(e)
}
func GetMetaData(name string) (info minio.ObjectInfo, err error) {
func GetMetadata(name string) (info minio.ObjectInfo, err error) {
return minioClient.StatObject(ctx, bucketName, name, minio.GetObjectOptions{})
}

112
submit.go Normal file
View file

@ -0,0 +1,112 @@
package main
import (
"html/template"
"net/http"
"waifu_gallery/s3"
_ "embed"
)
type LoadingPageContent struct {
Query string
GenerationEndpoint string
}
//go:embed loading.html
var loadingTmpl string
// You really wanna serve this endpoint behind WAF / Captcha or whatever and rate limit it to protect your precious tokens
func serveGenerationEndpoint(w http.ResponseWriter, r *http.Request) {
// HACK: Making this a /GET route to serve a CF challenge
// Validate user input
input, err := handleSubdomain(w, r)
if err != nil {
serveBadSubdomainError(w)
return
}
_, err = s3.GetMetadata(input)
// We want an error to happen here actually
if err == nil {
// Image already generated, bai bai
http.Redirect(w, r, "/", http.StatusMovedPermanently)
return
} else {
switch s3.ToErrorResponse(err).Code {
// Check that we error'd because the image doesn't exist
case "NoSuchKey":
// Generate picture or wait on existing generation job
// Probably not the best way but... if it works?
if ch, ok := generationQueue.Load(input); !ok {
// Check number of job requests
// TODO make this less horribly inefficient (lol)
var jobNbr int
generationQueue.Range(func(k, v interface{}) bool {
jobNbr++
return true
})
if config.MaxConcurrentJobs > 0 && jobNbr >= config.MaxConcurrentJobs {
w.WriteHeader(http.StatusTooManyRequests)
w.Write([]byte("Too many images currently being generated, please try again in a few minutes."))
return
}
// Create channel
generationQueue.Store(input, make(chan string))
// Generate new picture
InfoLogger.Printf("generating with input '%s'\n", input)
err := createPicture(input) // where the magic happens
if err != nil {
w.WriteHeader(http.StatusInternalServerError)
ErrorLogger.Println(err)
return
}
val, _ := generationQueue.Load(input)
generationQueue.Delete(input)
if c, ok := val.(chan string); ok {
// Non-blocking send to queue
select {
case c <- input:
default:
}
}
} else {
// Wait for job to finish if there's already a generation with the same input going on
if c, ok := ch.(chan string); ok {
<-c
}
}
default:
ErrorLogger.Println(err)
serveInternalError(w)
}
}
// Generation succeeded
w.Write([]byte("picture generated!"))
}
func writeWaitingPage(w http.ResponseWriter, input string) (err error) {
tmpl, err := template.New("loading").Parse(loadingTmpl)
if err != nil {
ErrorLogger.Println(err)
return
}
tmppData := LoadingPageContent{
Query: input,
GenerationEndpoint: "/generate",
}
err = tmpl.Execute(w, tmppData)
if err != nil {
ErrorLogger.Println(err)
return
}
return
}