134 lines
2.9 KiB
Go
134 lines
2.9 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
|
|
"danlegt.com/stablediffusion-friends/lib/rtypes"
|
|
goaway "github.com/TwiN/go-away"
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
var nodes [1]rtypes.SDNode = [...]rtypes.SDNode{
|
|
rtypes.SDNode{
|
|
URL: "127.0.0.1",
|
|
Port: 7860,
|
|
Status: false,
|
|
},
|
|
}
|
|
|
|
var imagesToDisplay []string
|
|
|
|
func main() {
|
|
r := gin.Default()
|
|
runLoaders(r)
|
|
|
|
r.GET("/", func(c *gin.Context) {
|
|
c.HTML(http.StatusOK, "index.html", nil)
|
|
})
|
|
|
|
r.POST("/api/image/generate", func(c *gin.Context) {
|
|
|
|
imageDescriptior := c.PostForm("image_description")
|
|
|
|
if goaway.IsProfane(imageDescriptior) {
|
|
imageDescriptior = goaway.Censor(imageDescriptior)
|
|
fmt.Printf("Found profanity, censoring to: %s\n", imageDescriptior)
|
|
}
|
|
|
|
// Build the request
|
|
request := rtypes.SDTextToImageRequest{
|
|
Prompt: imageDescriptior,
|
|
SamplerName: "DPM2",
|
|
BatchSize: 1,
|
|
Steps: 25,
|
|
CfgScale: 8,
|
|
Width: 512,
|
|
Height: 512,
|
|
NegativePrompt: "nsfw, porn, naked, nude, nipple, penis, dick, vagina, asshole, visible nipple, nsfl, not safe for work, nudity, artifact, deformed, multiple limbs, ugly, gore, blood, sex, pornography, penis, dick, genitalia, male genitalia, anus, penetration, double penetration, cock",
|
|
SamplerIndex: "Euler",
|
|
SendImages: true,
|
|
SaveImages: false,
|
|
}
|
|
|
|
body, err := json.Marshal(request)
|
|
if err != nil {
|
|
c.AbortWithError(500, err)
|
|
return
|
|
}
|
|
|
|
r, err := http.NewRequest("POST", nodes[0].GetTextToImageLink(), bytes.NewBuffer(body))
|
|
if err != nil {
|
|
c.AbortWithError(500, err)
|
|
return
|
|
}
|
|
|
|
r.Header.Add("Content-Type", "application/json")
|
|
client := &http.Client{}
|
|
res, err := client.Do(r)
|
|
if err != nil {
|
|
c.AbortWithError(500, err)
|
|
return
|
|
}
|
|
|
|
defer res.Body.Close()
|
|
|
|
responseJson := &rtypes.SDTextToImageResponse{}
|
|
err = json.NewDecoder(res.Body).Decode(responseJson)
|
|
if err != nil {
|
|
fmt.Println("Failed to decode the JSON")
|
|
c.AbortWithError(500, err)
|
|
return
|
|
}
|
|
|
|
c.JSON(http.StatusOK, responseJson)
|
|
})
|
|
|
|
r.POST("/api/image/submit", func(c *gin.Context) {
|
|
|
|
imageData := c.PostForm("image_data")
|
|
imagesToDisplay = append(imagesToDisplay[:], imageData)
|
|
|
|
fmt.Printf("Got a new image to display, currently %v in queue\n", len(imagesToDisplay))
|
|
|
|
c.JSON(200, gin.H{
|
|
"queue_position": len(imagesToDisplay),
|
|
})
|
|
})
|
|
|
|
r.GET("/api/image", func(c *gin.Context) {
|
|
|
|
if len(imagesToDisplay) <= 0 {
|
|
c.JSON(200, [0]string{})
|
|
return
|
|
}
|
|
|
|
newImage, newImagesToDisplay := imagesToDisplay[0], imagesToDisplay[1:]
|
|
|
|
// Actually pop it
|
|
if 3 > 2 {
|
|
imagesToDisplay = newImagesToDisplay
|
|
}
|
|
|
|
c.JSON(200, gin.H{
|
|
"image": newImage,
|
|
"queue_left": len(imagesToDisplay),
|
|
})
|
|
|
|
})
|
|
|
|
r.GET("/api/image/queue", func(c *gin.Context) {
|
|
c.JSON(200, imagesToDisplay)
|
|
})
|
|
|
|
r.Run()
|
|
|
|
}
|
|
|
|
func runLoaders(r *gin.Engine) {
|
|
r.Static("/public", "./public")
|
|
r.LoadHTMLGlob("templates/**/*.html")
|
|
}
|