108 lines
3.0 KiB
Go
108 lines
3.0 KiB
Go
package api
|
|
|
|
import (
|
|
"io/fs"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"strconv"
|
|
|
|
"trle5.xyz/upscayl-server/configs"
|
|
"trle5.xyz/upscayl-server/task"
|
|
"trle5.xyz/upscayl-server/upscayl"
|
|
"trle5.xyz/upscayl-server/utils"
|
|
)
|
|
|
|
// GETAPIHandler handler `/` URI.
|
|
func GETAPIHandler(w http.ResponseWriter, r *http.Request) {
|
|
utils.Json(w, map[string]any{
|
|
"message": "upscayl-server running",
|
|
})
|
|
}
|
|
|
|
// GETTasksHandler handler GET `/tasks` URI.
|
|
func GETTasksHandler(w http.ResponseWriter, r *http.Request) {
|
|
utils.Json(w, map[string]any{
|
|
"message": "ok",
|
|
"pending_tasks": task.TaskCount(),
|
|
})
|
|
}
|
|
|
|
// GETTasksIDHandler handler GET `/tasks/{taskID}` URI.
|
|
func GETTasksIDHandler(w http.ResponseWriter, r *http.Request) {
|
|
taskID := r.PathValue("taskID")
|
|
utils.CloseIf(taskID == "", "taskID is required", "invalid request", http.StatusBadRequest)
|
|
|
|
task, err := task.Get(taskID)
|
|
utils.CloseIfErr(err, "task not found", http.StatusNotFound)
|
|
|
|
utils.Json(w, task)
|
|
}
|
|
|
|
// POSTTasksHandler handler POST `/tasks` URI.
|
|
func POSTTasksHandler(w http.ResponseWriter, r *http.Request) {
|
|
theToken := r.Header.Get("Authorization")
|
|
utils.CloseIf(theToken != configs.AccessToken, "require access", "request blocked", http.StatusUnauthorized)
|
|
|
|
taskID, hash, err := utils.SaveUpload(r)
|
|
utils.CloseIfErr(err, "failed to save file", http.StatusBadRequest)
|
|
|
|
format := r.FormValue("format")
|
|
if format == "" {
|
|
format = "png"
|
|
}
|
|
|
|
var scaleInt int
|
|
scale := r.FormValue("scale")
|
|
if scale != "" {
|
|
scaleInt, err = strconv.Atoi(scale)
|
|
utils.CloseIfErr(err, "invalid scale", http.StatusBadRequest)
|
|
}
|
|
|
|
task, err := task.Add(taskID, hash, upscayl.Params{
|
|
Input: filepath.Join(configs.InputDir, taskID + ".temp"),
|
|
Output: filepath.Join(configs.OuputDir, taskID + "." + format),
|
|
Format: format,
|
|
Model: r.FormValue("model"),
|
|
Scale: scaleInt,
|
|
})
|
|
utils.CloseIfErr(err, "failed to add task", http.StatusBadRequest)
|
|
|
|
w.WriteHeader(http.StatusCreated)
|
|
utils.Json(w, task)
|
|
}
|
|
|
|
// OPTIONSTasksHandler handler OPTIONS `/tasks` URI.
|
|
func OPTIONSTasksHandler(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Add("Access-Control-Allow-Origin", "*")
|
|
w.Header().Add("Access-Control-Allow-Methods", "POST, GET")
|
|
w.Header().Add("Access-Control-Allow-Headers", "Content-Type, Access-Control-Allow-Headers, Authorization, X-Requested-With")
|
|
w.WriteHeader(http.StatusOK)
|
|
}
|
|
|
|
// GETModelsHandler handler GET `/models` URI.
|
|
func GETModelsHandler(w http.ResponseWriter, r *http.Request) {
|
|
var models []string
|
|
|
|
err := filepath.WalkDir(configs.ModelsDir, func(path string, d fs.DirEntry, err error) error {
|
|
if err != nil { return err }
|
|
if d.IsDir() { return nil }
|
|
|
|
ext := filepath.Ext(path)
|
|
name := filepath.Base(path[:len(path)-len(ext)])
|
|
|
|
switch ext {
|
|
case ".bin":
|
|
_, err = os.Stat(path[:len(path)-len(ext)] + ".param")
|
|
if err == nil { models = append(models, name) }
|
|
}
|
|
|
|
return nil
|
|
})
|
|
utils.CloseIfErr(err, "failed to read models directory", http.StatusInternalServerError)
|
|
|
|
utils.Json(w, map[string]any{
|
|
"models": models,
|
|
})
|
|
}
|