138 lines
4.2 KiB
Go
138 lines
4.2 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"os"
|
|
"reflect"
|
|
"regexp"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"waifu_gallery/openrouter"
|
|
"waifu_gallery/s3"
|
|
)
|
|
|
|
// STRONGLY recommended to use an uncensored LLM for funnier outputs
|
|
type Chatbots struct {
|
|
Endpoint string `json:"endpoint"` // required, API endpoint
|
|
Authorization string `json:"authorization"` // required, authorization tken in format `BEARER XXXXXXX`
|
|
|
|
// generate the enhanced prompt, 1:1 copy of chatGPT/openrouter's chat API
|
|
// the user input will be appended at the end with the "user" role
|
|
PromptCompletion openrouter.ChatCompletions `json:"promptCompletion"`
|
|
// generate the description for the artpiece
|
|
// Title: "%s" with the user input and Content: "%s" will be appened at the end with the "user" role
|
|
DescriptionCompletion openrouter.ChatCompletions `json:"descriptionCompletion"`
|
|
}
|
|
|
|
type Config struct {
|
|
// Fooocus serverless api endpoint
|
|
RunpodApi struct {
|
|
Endpoint string `json:"endpoint"`
|
|
Authorization string `json:"authorization"`
|
|
} `json:"runpodApi"`
|
|
|
|
// Image storage
|
|
S3 s3.Config `json:"s3"`
|
|
|
|
// Optional,LLMs
|
|
Chatbot Chatbots `json:"chatbot"`
|
|
|
|
// Where our web server listens
|
|
ListenAddress string `json:"listenAddress"`
|
|
MaxConcurrentJobs int `json:"maxConcurrentJobs"` // set this to limit the number of generation jobs that can be set in parallel, 0 to disable
|
|
|
|
// Content generation properties
|
|
Loras []Lora `json:"loras"`
|
|
AspectRatio string `json:"aspectRatio"` // necessary! used to compute display properties. e.g: 512*512
|
|
ImageSeed int `json:"imageSeed"` // set it to make image generation deterministic
|
|
RandomPrompt []string `json:"randomPrompt"` // prepended to base prompt, used to spice things up a bit (e.g multiple artist styles)
|
|
BasePrompt string `json:"basePrompt"` // user prompt will be appended to this
|
|
NegativePrompt string `json:"negativePrompt"` // self explanatory
|
|
UserPromptBias float64 `json:"userPromptBias"` // how much do we favor the user's input
|
|
GuidanceScale float64 `json:"GuidanceScale"` // The guidance scale, default 7.0
|
|
Sharpness float64 `json:"sharpness"` // Sharpness, defaults to 2.0
|
|
Steps int `json:"steps"` // Number of generation steps, default 25
|
|
Sampler string `json:"sampler"` // Which sampler to use, default euler_ancestral
|
|
BaseModel string `json:"baseModel"` // Base model, leave it empty to use fooocus' default
|
|
BannedWords []string `json:"bannedWords"` // user input will be blocked if it contains one of these
|
|
|
|
// Internal app state
|
|
state struct {
|
|
width int
|
|
height int
|
|
|
|
bannedWordsRegex *regexp.Regexp
|
|
|
|
canarySet bool
|
|
chatbotPromptEnabled bool
|
|
descriptionEnabled bool
|
|
}
|
|
}
|
|
|
|
func parseConfig(name string, cfg *Config) (err error) {
|
|
// read file
|
|
file, err := os.ReadFile(name)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// unmarshal
|
|
err = json.Unmarshal(file, &cfg)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// set defaults
|
|
if len(cfg.ListenAddress) == 0 {
|
|
cfg.ListenAddress = "127.0.0.1:8080"
|
|
}
|
|
if len(cfg.AspectRatio) == 0 {
|
|
cfg.AspectRatio = "1024*1024"
|
|
}
|
|
if len(cfg.Sampler) == 0 {
|
|
cfg.Sampler = "euler_ancestral"
|
|
}
|
|
if cfg.UserPromptBias == 0 {
|
|
cfg.UserPromptBias = 0.5
|
|
}
|
|
if cfg.GuidanceScale == 0 {
|
|
cfg.GuidanceScale = 7.0
|
|
}
|
|
if cfg.Steps == 0 {
|
|
cfg.Steps = 25
|
|
}
|
|
if cfg.Sharpness == 0 {
|
|
cfg.Sharpness = 2.0
|
|
}
|
|
|
|
// initialize aspect ratio
|
|
parts := strings.Split(cfg.AspectRatio, "*")
|
|
if len(parts) != 2 {
|
|
return errors.New("invalid aspect ratio format (must be something like 1024*1024)")
|
|
}
|
|
cfg.state.width, err = strconv.Atoi(parts[0])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
cfg.state.height, err = strconv.Atoi(parts[1])
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if cfg.state.width <= 0 || cfg.state.height <= 0 {
|
|
return errors.New("invalid aspect ratio (must be something like 1024*1024)")
|
|
}
|
|
// initialize badword regex
|
|
config.state.bannedWordsRegex = regexp.MustCompile(strings.Join(cfg.BannedWords, "|"))
|
|
|
|
// check for a bunch of chatbot related states
|
|
if !reflect.ValueOf(cfg.Chatbot.PromptCompletion).IsZero() {
|
|
cfg.state.chatbotPromptEnabled = true
|
|
}
|
|
if !reflect.ValueOf(cfg.Chatbot.DescriptionCompletion).IsZero() {
|
|
cfg.state.descriptionEnabled = true
|
|
}
|
|
|
|
return nil
|
|
}
|