Switch REST API to flask-restx and added APIv2 (#9)

* Refactored code to be able to handle multiple Hordes
* Switched to flask-restx for API
* Refactored API to be able to handle multiple API versions
* Added API documentation
* Added limiter per endpoint

.gitignore

@@ -131,4 +131,5 @@ dmypy.json
 
 db/*
 cliRequestsData.py
-*horde_generation*.png
+*horde_generation*.png
+test_commands.txt

README.md

@@ -22,7 +22,6 @@ The point of registering is to track your usage and your contributions. The more
 
 # Generating Prompts
 
-
 ### GUI
 
 We provide [a client interface](https://dbzer0.itch.io/stable-horde-client) requiring no installation and no technical expertise
@@ -47,28 +46,20 @@ You can make a copy of `cliRequestData_template.py` into `cliRequestData.py` and
 
 ## REST API
 
+[Full Documentation](https://stablehorde.net/api/v1)
+
+![](api_screenshot.png)
+
 You can also use the REST API directly. Be aware that this will return a base64 encoded image, so it will flood your output. This is not recommended unless you know what you're doing!
 
 ```
-curl -H "Content-Type: application/json" -d '{"prompt":"A horde of stable robots", "params":{"n":1, "width": 256, "height": 256}, "api_key":"0000000000"}' https://stablehorde.net/api/latest/generate/sync
+curl -H "Content-Type: application/json" -d '{"prompt":"A horde of stable robots", "params":{"n":1, "width": 256, "height": 256}, "api_key":"0000000000"}' https://stablehorde.net/api/v1/generate/sync
 ```
 
 The "params" dictionary is the same as use by the Stable API Webui. Documentation will be forthcoming.
 
 Pass an API Keyin order to track your usage.
 
-### Other API endpoints
-
-Full configuration on the API page coming soon
-
-API root: `/api/latest`
-
-* `/status/performance`: Performance of the Stable Horde as a whole
-* `/servers`: All active servers and their stats
-* `/servers/<id>`: Specific Server.
-* `/users`: All active users and their stats
-* `/users/<id>`: Specific User.
-
 ## Specifying servers
 
 You can optionally specify only specific servers to generate for you. Grab one or more server IDs from `/servers` and then send it with your payload as a list in the "servers" arg. Your generation will only be fulfiled by servers with the specified IDs

horde/__init__.py

@@ -0,0 +1,67 @@
+from .logger import logger, set_logger_verbosity, quiesce_logger
+from .argparser import args
+
+set_logger_verbosity(args.verbosity)
+quiesce_logger(args.quiet)
+
+from .classes.maintenance import Maintenance
+maintenance = Maintenance()
+
+from .limiter import limiter
+from flask import Flask, render_template, redirect, url_for, request, Blueprint
+from .flask import HORDE
+from . import routes
+from .apis import apiv1, apiv2
+from flask_dance.contrib.google import make_google_blueprint, google
+from flask_dance.contrib.discord import make_discord_blueprint, discord
+from flask_dance.contrib.github import make_github_blueprint, github
+import os
+
+HORDE.register_blueprint(apiv1)
+HORDE.register_blueprint(apiv2)
+
+@HORDE.before_request
+def limit_remote_addr():
+    logger.debug(request.remote_addr)
+    # if not allow_direct_connections and request.remote_addr != '127.0.0.1':
+    #     error_msg = get_error(ServerErrors.NO_PROXY)
+    #     abort(403, error_msg)
+
+
+@HORDE.after_request
+def after_request(response):
+    response.headers["Access-Control-Allow-Origin"] = "*"
+    response.headers["Access-Control-Allow-Methods"] = "POST, GET, OPTIONS, PUT, DELETE"
+    response.headers["Access-Control-Allow-Headers"] = "Accept, Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, apikey"
+    return response
+
+google_client_id = os.getenv("GOOGLE_CLIENT_ID")
+google_client_secret = os.getenv("GLOOGLE_CLIENT_SECRET")
+discord_client_id = os.getenv("DISCORD_CLIENT_ID")
+discord_client_secret = os.getenv("DISCORD_CLIENT_SECRET")
+github_client_id = os.getenv("GITHUB_CLIENT_ID")
+github_client_secret = os.getenv("GITHUB_CLIENT_SECRET")
+HORDE.secret_key = os.getenv("secret_key")
+os.environ['OAUTHLIB_RELAX_TOKEN_SCOPE'] = '1'
+google_blueprint = make_google_blueprint(
+    client_id = google_client_id,
+    client_secret = google_client_secret,
+    reprompt_consent = True,
+    redirect_url='/register',
+    scope = ["email"],
+)
+HORDE.register_blueprint(google_blueprint,url_prefix="/google")
+discord_blueprint = make_discord_blueprint(
+    client_id = discord_client_id,
+    client_secret = discord_client_secret,
+    scope = ["identify"],
+    redirect_url='/finish_dance',
+)
+HORDE.register_blueprint(discord_blueprint,url_prefix="/discord")
+github_blueprint = make_github_blueprint(
+    client_id = github_client_id,
+    client_secret = github_client_secret,
+    scope = ["identify"],
+    redirect_url='/finish_dance',
+)
+HORDE.register_blueprint(github_blueprint,url_prefix="/github")

horde/apis/__init__.py

@@ -0,0 +1,9 @@
+from .. import args
+from importlib import import_module
+
+ModelsV2 = import_module(name=f'horde.apis.models.{args.horde}_v2').Models
+ParsersV2 = import_module(name=f'horde.apis.models.{args.horde}_v2').Parsers
+
+from .apiv2 import blueprint as apiv2
+from .apiv1 import blueprint as apiv1
+

horde/apis/apiv1.py

@@ -0,0 +1,22 @@
+from flask import Blueprint
+from flask_restx import Api
+from .. import args
+from importlib import import_module
+
+from .v1.v1 import api as v1
+ModelsV2 = import_module(name=f'horde.apis.models.{args.horde}_v2').Models
+v2 = import_module(name=f'.{args.horde}', package=f'horde.apis.v2').api
+
+blueprint = Blueprint('apiv1', __name__, url_prefix='/api')
+api = Api(blueprint,
+    version='1.0', 
+    title='Stable Horde',
+    description='The API documentation for the Stable Horde',
+    contact_email="mail@dbzer0.com",
+    default="v1",
+    default_label="Obsolete Version",
+    ordered=True,
+)
+
+api.add_namespace(v2)
+api.add_namespace(v1)

horde/apis/apiv2.py

@@ -0,0 +1,21 @@
+from flask import Blueprint
+from flask_restx import Api
+from .. import args
+from importlib import import_module
+
+from .v1.v1 import api as v1
+v2 = import_module(name=f'.{args.horde}', package=f'horde.apis.v2').api
+
+blueprint = Blueprint('apiv2', __name__, url_prefix='/api')
+api = Api(blueprint,
+    version='2.0', 
+    title='Stable Horde',
+    description='The API documentation for the Stable Horde',
+    contact_email="mail@dbzer0.com",
+    default="v2",
+    default_label="Latest Version",
+    ordered=True,
+)
+
+api.add_namespace(v2)
+api.add_namespace(v1)

horde/apis/exceptions.py

@@ -0,0 +1,106 @@
+from werkzeug import exceptions as wze
+from ..logger import logger
+
+class MissingPrompt(wze.BadRequest):
+    def __init__(self, username):
+        self.specific = "You cannot specify an empty prompt."
+        self.log = f"User '{username}' sent an empty prompt. Aborting!"
+
+class KudosValidationError(wze.BadRequest):
+    def __init__(self, username, error_message):
+        self.specific = error_message
+        self.log = f"User '{username}' Failed to transfer Kudos."
+
+class NoValidActions(wze.BadRequest):
+    def __init__(self, error_message):
+        self.specific = error_message
+        self.log = None
+
+class InvalidSize(wze.BadRequest):
+    def __init__(self, username):
+        self.specific = "Invalid size. The image dimentions have to be multiples of 64."
+        self.log = f"User '{username}' sent an invalid size. Aborting!"
+
+class TooManySteps(wze.BadRequest):
+    def __init__(self, username, steps):
+        self.specific = "Too many sampling steps. To allow resources for everyone, we allow only up to 100 steps."
+        self.log = f"User '{username}' sent too many steps ({steps}). Aborting!"
+
+class InvalidAPIKey(wze.Unauthorized):
+    def __init__(self, subject):
+        self.specific = "No user matching sent API Key. Have you remembered to register at https://stablehorde.net/register ?"
+        self.log = f"Invalid API Key sent for {subject}"
+
+class WrongCredentials(wze.Forbidden):
+    def __init__(self, username, worker):
+        self.specific = "You cannot specify an empty prompt."
+        self.log = f"User '{username}' sent wrong credentials for utilizing worker {worker}"
+
+class NotAdmin(wze.Forbidden):
+    def __init__(self, username, endpoint):
+        self.specific = "You're not an admin. Sod off!"
+        self.log = f"Non-admin user '{username}' tried to use admin endpoint: '{endpoint}. Aborting!"
+
+class NotOwner(wze.Forbidden):
+    def __init__(self, username, worker_name):
+        self.specific = "You're not an admin. Sod off!"
+        self.log = f"User '{username}'' tried to modify worker they do not own '{worker_name}'. Aborting!"
+
+class WorkerMaintenance(wze.Forbidden):
+    def __init__(self):
+        self.specific = "This worker has been put into maintenance by its owner"
+        self.log = None
+
+class InvalidProcGen(wze.NotFound):
+    def __init__(self, gen_id):
+        self.specific = f"Processing Generation with ID {gen_id} does not exist."
+        self.log = f"Worker attempted to provide generation for {gen_id} but it did not exist"
+
+class RequestNotFound(wze.NotFound):
+    def __init__(self, wp_id):
+        self.specific = f"Request with ID '{wp_id}' not found."
+        self.log = f"Status of WP with ID '{wp_id}' does not exist"
+
+class WorkerNotFound(wze.NotFound):
+    def __init__(self, worker_id):
+        self.specific = f"Worker with ID '{worker_id}' not found."
+        self.log = f"Attempted to retrieve worker with non-existent ID '{worker_id}'"
+
+class UserNotFound(wze.NotFound):
+    def __init__(self, user_id):
+        self.specific = f"User with ID '{user_id}' not found."
+        self.log = f"Attempted to retrieve user with non-existent ID '{user_id}'"
+
+class DuplicateGen(wze.NotFound):
+    def __init__(self, worker, gen_id):
+        self.specific = f"Processing Generation with ID {gen_id} already submitted."
+        self.log = f"Worker '{worker}' attempted to provide duplicate generation for {gen_id}"
+
+class RequestExpired(wze.Gone):
+    def __init__(self, username):
+        self.specific = f"Prompt Request Expired"
+        self.log = f"Request from '{username}' took too long to complete and has been cancelled."
+
+class TooManyPrompts(wze.TooManyRequests):
+    def __init__(self, username, count):
+        self.specific = f"Parallel requests exceeded user limit ({count}). Please try again later or request to increase your concurrency."
+        self.log = f"User '{username}' has already requested too many parallel requests ({count}). Aborting!"
+
+class NoValidWorkers(wze.ServiceUnavailable):
+    retry_after = 600
+    def __init__(self, username):
+        self.specific = f"No active worker found to fulfill this request. Please Try again later..."
+        self.log = f"No active worker found to match the request from '{username}'. Aborting!"
+
+class MaintenanceMode(wze.ServiceUnavailable):
+    retry_after = 60
+    def __init__(self, endpoint):
+        self.specific = f"Horde has enterred maintenance mode. Please try again later."
+        self.log = f"Rejecting endpoint '{endpoint}' because horde in maintenance mode."
+
+
+def handle_bad_requests(error):
+    '''Namespace error handler'''
+    if error.log:
+        logger.warning(error.log)
+    return({'message': error.specific}, error.code)

horde/apis/models/__init__.py

@@ -0,0 +1 @@
+from .stable_v2 import Models

horde/apis/models/stable_v2.py

@@ -0,0 +1,71 @@
+from flask_restx import fields, reqparse
+from . import v2
+
+
+class Parsers(v2.Parsers):
+    def __init__(self):
+        self.job_pop_parser.add_argument("max_pixels", type=int, required=False, default=512, help="The maximum amount of pixels this worker can generate", location="json")
+        self.job_submit_parser.add_argument("seed", type=str, required=True, default=[], help="The seed of the generation", location="json")
+
+class Models(v2.Models):
+    def __init__(self,api):
+
+        super().__init__(api)
+
+        self.response_model_generation_result = api.inherit('GenerationStable', self.response_model_generation_result, {
+            'img': fields.String(title="Generated Image", description="The generated image as a Base64-encoded .webp file"),
+            'seed': fields.String(title="Generation Seed", description="The seed which generated this image"),
+        })
+        self.response_model_wp_status_full = api.inherit('RequestStatusStable', self.response_model_wp_status_lite, {
+            'generations': fields.List(fields.Nested(self.response_model_generation_result)),
+        })
+        self.response_model_generation_payload = api.model('ModelPayload', {
+            'prompt': fields.String(description="The prompt which will be sent to Stable Diffusion to generate an image"),
+            'ddim_steps': fields.Integer(example=50), 
+            'sampler_name': fields.String(enum=["k_lms", "k_heun", "k_euler", "k_euler_a", "k_dpm_2", "k_dpm_2_a", "DDIM", "PLMS"]), 
+            'toggles': fields.List(fields.Integer,example=[1,4], description="Special Toggles used in the SD Webui. To be documented."), 
+            'realesrgan_model_name': fields.String,
+            'ddim_eta': fields.Float, 
+            'n_iter': fields.Integer(example=1, description="The amount of images to generate"), 
+            'batch_size': fields.Integer(example=1), 
+            'cfg_scale': fields.Float(example=5.0), 
+            'seed': fields.String(description="The seed to use to generete this request"),
+            'height': fields.Integer(example=512,description="The height of the image to generate"), 
+            'width': fields.Integer(example=512,description="The width of the image to generate"), 
+            'fp': fields.Integer(example=512), 
+            'variant_amount': fields.Float, 
+            'variant_seed': fields.Integer
+        })
+        self.response_model_generations_skipped = api.inherit('NoValidRequestFoundStable', self.response_model_generations_skipped, {
+            'max_pixels': fields.Integer(example=0,description="How many waiting requests were skipped because they demanded a higher size than this worker provides"),
+        })
+        self.response_model_job_pop = api.model('GenerationPayload', {
+            'payload': fields.Nested(self.response_model_generation_payload,skip_none=True),
+            'id': fields.String(description="The UUID for this image generation"),
+            'skipped': fields.Nested(self.response_model_generations_skipped,skip_none=True)
+        })
+        self.response_model_worker_details = api.inherit('WorkerDetailsStable', self.response_model_worker_details, {
+            "max_pixels": fields.Integer(example=262144,description="The maximum pixels in resolution this workr can generate"),
+            "megapixelsteps_generated": fields.Float(description="How many megapixelsteps this worker has generated until now"),
+        })
+        self.response_model_contrib_details = api.inherit('ContributionsDetailsStable', self.response_model_contrib_details, {
+            "megapixelsteps": fields.Float(description="How many megapixelsteps this user has generated"),
+        })
+        self.response_model_use_details = api.inherit('UsageDetailsStable', self.response_model_use_details, {
+            "megapixelsteps": fields.Float(description="How many megapixelsteps this user has requested"),
+        })
+        self.response_model_user_details = api.model('UserDetails', {
+            "username": fields.String(description="The user's unique Username. It is a combination of their chosen alias plus their ID."),
+            "id": fields.Integer(description="The user unique ID. It is always an integer."),
+            "kudos": fields.Float(description="The amount of Kudos this user has. Can be negative. The amount of Kudos determines the priority when requesting image generations."),
+            "kudos_details": fields.Nested(self.response_model_user_kudos_details),
+            "usage": fields.Nested(self.response_model_use_details),
+            "contributions": fields.Nested(self.response_model_contrib_details),
+            "concurrency": fields.Integer(description="How many concurrent image generations this user may request."),    
+        })
+        self.response_model_horde_performance = api.inherit('HordePerformanceStable', self.response_model_horde_performance, {
+            "queued_requests": fields.Integer(description="The amount of waiting and processing requests currently in this Horde"),
+            "queued_megapixelsteps": fields.Float(description="The amount of megapixelsteps in waiting and processing requests currently in this Horde"),
+            "past_minute_megapixelsteps": fields.Float(description="How many megapixelsteps this Horde generated in the last minute"),
+            "worker_count": fields.Integer(description="How many workers are actively processing image generations in this Horde in the past 5 minutes"),
+        })

horde/apis/models/v2.py

@@ -0,0 +1,137 @@
+from flask_restx import fields, reqparse
+
+
+class Parsers:
+
+    generate_parser = reqparse.RequestParser()
+    generate_parser.add_argument("apikey", type=str, required=True, help="The API Key corresponding to a registered user", location='headers')
+    generate_parser.add_argument("prompt", type=str, required=True, help="The prompt to generate from", location="json")
+    generate_parser.add_argument("params", type=dict, required=False, default={}, help="Extra generate params to send to the worker", location="json")
+    generate_parser.add_argument("workers", type=str, action='append', required=False, default=[], help="If specified, only the worker with this ID will be able to generate this prompt", location="json")
+    # generate_parser.add_argument("models", type=str, action='append', required=False, default=[], help="Models", location="json")
+
+    # The parser for RequestPop
+    job_pop_parser = reqparse.RequestParser()
+    job_pop_parser.add_argument("apikey", type=str, required=True, help="The API Key corresponding to a registered user", location='headers')
+    job_pop_parser.add_argument("name", type=str, required=True, help="The worker's unique name, to track contributions", location="json")
+    job_pop_parser.add_argument("priority_usernames", type=str, action='append', required=False, default=[], help="The usernames which get priority use on this worker", location="json")
+
+    job_submit_parser = reqparse.RequestParser()
+    job_submit_parser.add_argument("apikey", type=str, required=True, help="The worker's owner API key", location='headers')
+    job_submit_parser.add_argument("id", type=str, required=True, help="The processing generation uuid", location="json")
+    job_submit_parser.add_argument("generation", type=str, required=False, default=[], help="The generated output", location="json")
+
+
+class Models:
+    def __init__(self,api):
+
+        self.response_model_generation_result = api.model('Generation', {
+            'worker_id': fields.String(title="Worker ID", description="The UUID of the worker which generated this image"),
+            'worker_name': fields.String(title="Worker Name", description="The name of the worker which generated this image"),
+        })
+        self.response_model_wp_status_lite = api.model('RequestStatusCheck', {
+            'finished': fields.Integer(description="The amount of finished images in this request"),
+            'processing': fields.Integer(description="The amount of still processing images in this request"),
+            'waiting': fields.Integer(description="The amount of images waiting to be picked up by a worker"),
+            'done': fields.Boolean(description="True when all images in this request are done. Else False."),
+            'wait_time': fields.Integer(description="The expected amount to wait (in seconds) to generate all images in this request"),
+            'queue_position': fields.Integer(description="The position in the requests queue. This position is determined by relative Kudos amounts."),
+        })
+        self.response_model_wp_status_full = api.inherit('RequestStatus', self.response_model_wp_status_lite, {
+            'generations': fields.List(fields.Nested(self.response_model_generation_result)),
+        })
+        self.response_model_async = api.model('RequestAsync', {
+            'id': fields.String(description="The UUID of the request. Use this to retrieve the request status in the future"),
+            'message': fields.String(default=None,description="Any extra information from the horde about this request"),
+        })
+        self.response_model_generation_payload = api.model('ModelPayload', {
+            'prompt': fields.String(description="The prompt which will be sent to Stable Diffusion to generate an image"),
+            'n': fields.Integer(example=1, description="The amount of images to generate"), 
+            'seed': fields.String(description="The seed to use to generete this request"),
+        })
+        self.response_model_generations_skipped = api.model('NoValidRequestFound', {
+            'worker_id': fields.Integer(description="How many waiting requests were skipped because they demanded a specific worker"),
+        })
+
+        self.response_model_job_pop = api.model('GenerationPayload', {
+            'payload': fields.Nested(self.response_model_generation_payload, skip_none=True),
+            'id': fields.String(description="The UUID for this image generation"),
+            'skipped': fields.Nested(self.response_model_generations_skipped, skip_none=True)
+        })
+
+        self.response_model_job_submit = api.model('GenerationSubmitted', {
+            'reward': fields.Float(example=10.0,description="The amount of kudos gained for submitting this request"),
+        })
+
+        self.response_model_kudos_transfer = api.model('KudosTransferred', {
+            'transferred': fields.Integer(example=100,description="The amount of Kudos tranferred"),
+        })
+
+        self.response_model_admin_maintenance = api.model('MaintenanceModeSet', {
+            'maintenance_mode': fields.Boolean(example=True,description="The current state of maintenance_mode"),
+        })
+
+        self.response_model_worker_kudos_details = api.model('WorkerKudosDetails', {
+            'generated': fields.Float(description="How much Kudos this worker has received for generating images"),
+            'uptime': fields.Integer(description="How much Kudos this worker has received for staying online longer"),
+        })
+
+        self.response_model_worker_details = api.model('WorkerDetails', {
+            "name": fields.String(description="The Name given to this worker."),
+            "id": fields.String(description="The UUID of this worker."),
+            "requests_fulfilled": fields.Integer(description="How many images this worker has generated."),
+            "kudos_rewards": fields.Float(description="How many Kudos this worker has been rewarded in total."),
+            "kudos_details": fields.Nested(self.response_model_worker_kudos_details),
+            "performance": fields.String(description="The average performance of this worker in human readable form."),
+            "uptime": fields.Integer(description="The amount of seconds this worker has been online for this Horde."),
+            "maintenance_mode": fields.Boolean(example=False,description="When True, this worker will not pick up any new requests"),
+            "paused": fields.Boolean(example=False,description="When True, this worker not be given any new requests."),
+        })
+
+        self.response_model_worker_modify = api.model('ModifyWorker', {
+            "maintenance": fields.Boolean(description="The new state of the 'maintenance' var for this worker. When True, this worker will not pick up any new requests"),
+            "paused": fields.Boolean(description="The new state of the 'paused' var for this worker. When True, this worker will not be given any new requests"),
+        })
+
+        self.response_model_user_kudos_details = api.model('UserKudosDetails', {
+            "accumulated": fields.Float(default=0,description="The ammount of Kudos accumulated or used for generating images."),
+            "gifted": fields.Float(default=0,description="The amount of Kudos this user has given to other users"),
+            "admin": fields.Float(default=0,description="The amount of Kudos this user has been given by the Horde admins"),
+            "received": fields.Float(default=0,description="The amount of Kudos this user has been given by other users"),
+        })
+
+        self.response_model_contrib_details = api.model('ContributionsDetails', {
+            "fulfillments": fields.Integer(description="How many images this user has generated")
+        })
+        self.response_model_use_details = api.model('UsageDetails', {
+            "requests": fields.Integer(description="How many images this user has requested")
+        })
+
+        self.response_model_user_details = api.model('UserDetails', {
+            "username": fields.String(description="The user's unique Username. It is a combination of their chosen alias plus their ID."),
+            "id": fields.Integer(description="The user unique ID. It is always an integer."),
+            "kudos": fields.Float(description="The amount of Kudos this user has. Can be negative. The amount of Kudos determines the priority when requesting image generations."),
+            "kudos_details": fields.Nested(self.response_model_user_kudos_details),
+            "usage": fields.Nested(self.response_model_use_details),
+            "contributions": fields.Nested(self.response_model_contrib_details),
+            "concurrency": fields.Integer(description="How many concurrent image generations this user may request."),    
+        })
+
+        self.response_model_user_modify = api.model('ModifyUser', {
+            "new_kudos": fields.Float(description="The new total Kudos this user has after this request"),
+            "concurrency": fields.Integer(example=30,description="The request concurrency this user has after this request"),
+            "usage_multiplier": fields.Float(example=1.0,description="Multiplies the amount of kudos lost when generating images."),
+        })
+
+        self.response_model_horde_performance = api.model('HordePerformance', {
+            "queued_requests": fields.Integer(description="The amount of waiting and processing requests currently in this Horde"),
+            "worker_count": fields.Integer(description="How many workers are actively processing image generations in this Horde in the past 5 minutes"),
+        })
+
+        self.response_model_horde_maintenance_mode = api.model('HordeMaintenanceMode', {
+            "maintenance_mode": fields.Boolean(description="When True, this Horde will not accept new requests for image generation, but will finish processing the ones currently in the queue."),
+        })
+
+        self.response_model_error = api.model('RequestError', {
+            'message': fields.String(description="The error message for this status code."),
+        })

horde/apis/v1/__init__.py

@@ -0,0 +1 @@
+api = None

horde/apis/v1/v1.py

@@ -0,0 +1,539 @@
+from flask_restx import Namespace, Resource, reqparse, fields, Api, abort
+from flask import request
+from ... import limiter
+from ...logger import logger
+from ...classes import db as _db
+from ...classes import processing_generations,waiting_prompts,Worker,User,WaitingPrompt
+from ... import maintenance
+from enum import Enum
+import os, time
+
+
+api = Namespace('v1', 'API Version 1' )
+
+response_model_generation = api.model('GenerationV1', {
+    'img': fields.String,
+    'seed': fields.String,
+    'server_id': fields.String(attribute='worker_id'),
+    'server_name': fields.String(attribute='worker_name'),
+    'queue_position': fields.Integer(description="The position in the requests queue. This position is determined by relative Kudos amounts."),
+})
+response_model_wp_status_lite = api.model('RequestStatusCheckV1', {
+    'finished': fields.Integer,
+    'processing': fields.Integer,
+    'waiting': fields.Integer,
+    'done': fields.Boolean,
+    'wait_time': fields.Integer,
+    'queue_position': fields.Integer(description="The position in the requests queue. This position is determined by relative Kudos amounts."),
+})
+response_model_wp_status_full = api.inherit('RequestStatusV1', response_model_wp_status_lite, {
+    'generations': fields.List(fields.Nested(response_model_generation)),
+})
+# Used to for the flas limiter, to limit requests per url paths
+def get_request_path():
+    return(request.path)
+
+class ServerErrors(Enum):
+    WRONG_CREDENTIALS = 0
+    INVALID_PROCGEN = 1
+    DUPLICATE_GEN = 2
+    TOO_MANY_PROMPTS = 3
+    EMPTY_PROMPT = 4
+    INVALID_API_KEY = 5
+    INVALID_SIZE = 6
+    NO_PROXY = 7
+    TOO_MANY_STEPS = 8
+    NOT_ADMIN = 9
+    MAINTENANCE_MODE = 10
+    NOT_OWNER = 11
+
+@logger.catch
+def get_error(error, **kwargs):
+    if error == ServerErrors.INVALID_API_KEY:
+        logger.warning(f'Invalid API Key sent for {kwargs["subject"]}.')
+        return(f'No user matching sent API Key. Have you remembered to register at https://stablehorde.net/register ?')
+    if error == ServerErrors.WRONG_CREDENTIALS:
+        logger.warning(f'User "{kwargs["username"]}" sent wrong credentials for utilizing instance {kwargs["kai_instance"]}')
+        return(f'wrong credentials for utilizing instance {kwargs["kai_instance"]}')
+    if error == ServerErrors.INVALID_PROCGEN:
+        logger.warning(f'Server attempted to provide generation for {kwargs["id"]} but it did not exist')
+        return(f'Processing Generation with ID {kwargs["id"]} does not exist')
+    if error == ServerErrors.DUPLICATE_GEN:
+        logger.warning(f'Server attempted to provide duplicate generation for {kwargs["id"]} ')
+        return(f'Processing Generation with ID {kwargs["id"]} already submitted')
+    if error == ServerErrors.TOO_MANY_PROMPTS:
+        logger.warning(f'User "{kwargs["username"]}" has already requested too many parallel requests ({kwargs["wp_count"]}). Aborting!')
+        return(f"Parallel requests exceeded user limit ({kwargs['wp_count']}). Please try again later or request to increase your concurrency.")
+    if error == ServerErrors.EMPTY_PROMPT:
+        logger.warning(f'User "{kwargs["username"]}" sent an empty prompt. Aborting!')
+        return("You cannot specify an empty prompt.")
+    if error == ServerErrors.INVALID_SIZE:
+        logger.warning(f'User "{kwargs["username"]}" sent an invalid size. Aborting!')
+        return("Invalid size. The image dimentions have to be multiples of 64.")
+    if error == ServerErrors.TOO_MANY_STEPS:
+        logger.warning(f'User "{kwargs["username"]}" sent too many steps ({kwargs["steps"]}). Aborting!')
+        return("Too many sampling steps. To allow resources for everyone, we allow only up to 100 steps.")
+    if error == ServerErrors.NO_PROXY:
+        logger.warning(f'Attempt to access outside reverse proxy')
+        return(f'Access allowed only through https')
+    if error == ServerErrors.NOT_ADMIN:
+        logger.warning(f'Non-admin user "{kwargs["username"]}" tried to use admin endpoint: "{kwargs["endpoint"]}". Aborting!')
+        return("You're not an admin. Sod off!")
+    if error == ServerErrors.MAINTENANCE_MODE:
+        logger.info(f'Rejecting endpoint "{kwargs["endpoint"]}" because server in maintenance mode.')
+        return("Server has enterred maintenance mode. Please try again later.")
+    if error == ServerErrors.NOT_OWNER:
+        logger.warning(f'User "{kwargs["username"]}" tried to modify server they do not own: "{kwargs["server_name"]}". Aborting!')
+        return("You're not the owner of this server!")
+
+
+
+class SyncGenerate(Resource):
+    # model_generation = api.model('Successful Sync Generation', {
+    # 'generations': fields.List(fields.String),
+    # })
+    parser = reqparse.RequestParser()
+    parser.add_argument("prompt", type=str, required=True, help="The prompt to generate from")
+    parser.add_argument("api_key", type=str, required=True, help="The API Key corresponding to a registered user")
+    parser.add_argument("params", type=dict, required=False, default={}, help="Extra generate params to send to the SD server")
+    parser.add_argument("servers", type=str, action='append', required=False, default=[], help="If specified, only the server with this ID will be able to generate this prompt")
+    @api.expect(parser)
+    @api.response(200, 'Success', response_model_generation)
+    @api.response(400, 'Validation Error')
+    def post(self):
+        args = self.parser.parse_args()
+        username = 'Anonymous'
+        user = None
+        if maintenance.active:
+            return(f"{get_error(ServerErrors.MAINTENANCE_MODE, endpoint = 'SyncGenerate')}",503)
+        if args.api_key:
+            user = _db.find_user_by_api_key(args['api_key'])
+        if not user:
+            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'prompt generation')}",401)
+        username = user.get_unique_alias()
+        if args['prompt'] == '':
+            return(f"{get_error(ServerErrors.EMPTY_PROMPT, username = username)}",400)
+        wp_count = waiting_prompts.count_waiting_requests(user)
+        if wp_count >= user.concurrency:
+            return(f"{get_error(ServerErrors.TOO_MANY_PROMPTS, username = username, wp_count = wp_count)}",503)
+        if args["params"].get("length",512)%64:
+            return(f"{get_error(ServerErrors.INVALID_SIZE, username = username)}",400)
+        if args["params"].get("width",512)%64:
+            return(f"{get_error(ServerErrors.INVALID_SIZE, username = username)}",400)
+        if args["params"].get("steps",50) > 100:
+            return(f"{get_error(ServerErrors.TOO_MANY_STEPS, username = username, steps = args['params']['steps'])}",400)
+        wp = WaitingPrompt(
+            _db,
+            waiting_prompts,
+            processing_generations,
+            args["prompt"],
+            user,
+            args["params"],
+            servers=args["servers"],
+        )
+        server_found = False
+        for server in _db.workers.values():
+            if len(args.servers) and server.id not in args.servers:
+                continue
+            if server.can_generate(wp)[0]:
+                server_found = True
+                break
+        if not server_found:
+            del wp # Normally garbage collection will handle it, but doesn't hurt to be thorough
+            return("No active server found to fulfill this request. Please Try again later...", 503)
+        # if a server is available to fulfil this prompt, we activate it and add it to the queue to be generated
+        wp.activate()
+        while True:
+            time.sleep(1)
+            if wp.is_stale():
+                return("Prompt Request Expired", 500)
+            if wp.is_completed():
+                break
+        ret_dict = wp.get_status()['generations']
+        # We delete it from memory immediately to ensure we don't run out
+        wp.delete()
+        return(ret_dict, 200)
+
+
+class AsyncStatus(Resource):
+    decorators = [limiter.limit("2/minute", key_func = get_request_path)]
+    @logger.catch
+    @api.marshal_with(response_model_wp_status_full, code=200, description='Images Generated')
+    def get(self, id = ''):
+        wp = waiting_prompts.get_item(id)
+        if not wp:
+            return("ID not found", 404)
+        wp_status = wp.get_status()
+        # If the status is retrieved after the wp is done we clear it to free the ram
+        if wp_status["done"]:
+            wp.delete()
+        return(wp_status, 200)
+
+
+class AsyncCheck(Resource):
+    # Increasing this until I can figure out how to pass original IP from reverse proxy
+    decorators = [limiter.limit("10/second")]
+    @logger.catch
+    def get(self, id = ''):
+        wp = waiting_prompts.get_item(id)
+        if not wp:
+            return("ID not found", 404)
+        return(wp.get_lite_status(), 200)
+
+
+class AsyncGenerate(Resource):
+    parser = reqparse.RequestParser()
+    parser.add_argument("prompt", type=str, required=True, help="The prompt to generate from")
+    parser.add_argument("api_key", type=str, required=True, help="The API Key corresponding to a registered user")
+    parser.add_argument("params", type=dict, required=False, default={}, help="Extra generate params to send to the SD server")
+    parser.add_argument("servers", type=str, action='append', required=False, default=[], help="If specified, only the server with this ID will be able to generate this prompt")
+
+    @api.expect(parser)
+    def post(self):
+        args = self.parser.parse_args()
+        username = 'Anonymous'
+        user = None
+        if maintenance.active:
+            return(f"{get_error(ServerErrors.MAINTENANCE_MODE, endpoint = 'AsyncGenerate')}",503)
+        if args.api_key:
+            user = _db.find_user_by_api_key(args['api_key'])
+        if not user:
+            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'prompt generation')}",401)
+        username = user.get_unique_alias()
+        if args['prompt'] == '':
+            return(f"{get_error(ServerErrors.EMPTY_PROMPT, username = username)}",400)
+        wp_count = waiting_prompts.count_waiting_requests(user)
+        if wp_count >= user.concurrency:
+            return(f"{get_error(ServerErrors.TOO_MANY_PROMPTS, username = username, wp_count = wp_count)}",503)
+        if args["params"].get("length",512)%64:
+            return(f"{get_error(ServerErrors.INVALID_SIZE, username = username)}",400)
+        if args["params"].get("width",512)%64:
+            return(f"{get_error(ServerErrors.INVALID_SIZE, username = username)}",400)
+        if args["params"].get("steps",50) > 100:
+            return(f"{get_error(ServerErrors.TOO_MANY_STEPS, username = username, steps = args['params']['steps'])}",400)
+        wp = WaitingPrompt(
+            _db,
+            waiting_prompts,
+            processing_generations,
+            args["prompt"],
+            user,
+            args["params"],
+            servers=args["servers"],
+        )
+        server_found = False
+        for server in _db.workers.values():
+            if len(args.servers) and server.id not in args.servers:
+                continue
+            if server.can_generate(wp)[0]:
+                server_found = True
+                break
+        if not server_found:
+            del wp # Normally garbage collection will handle it, but doesn't hurt to be thorough
+            return("No active server found to fulfill this request. Please Try again later...", 503)
+        # if a server is available to fulfil this prompt, we activate it and add it to the queue to be generated
+        wp.activate()
+        return({"id":wp.id}, 200)
+
+
+class PromptPop(Resource):
+    parser = reqparse.RequestParser()
+    parser.add_argument("api_key", type=str, required=True, help="The API Key corresponding to a registered user")
+    parser.add_argument("name", type=str, required=True, help="The server's unique name, to track contributions")
+    parser.add_argument("max_pixels", type=int, required=False, default=512, help="The maximum amount of pixels this server can generate")
+    parser.add_argument("priority_usernames", type=str, action='append', required=False, default=[], help="The usernames which get priority use on this server")
+
+    decorators = [limiter.limit("45/second")]
+    @api.expect(parser)
+    def post(self):
+        args = self.parser.parse_args()
+        skipped = {}
+        user = _db.find_user_by_api_key(args['api_key'])
+        if not user:
+            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'server promptpop: ' + args['name'])}",401)
+        server = _db.find_worker_by_name(args['name'])
+        if not server:
+            server = Worker(_db)
+            server.create(user, args['name'])
+        if user != server.user:
+            return(f"{get_error(ServerErrors.WRONG_CREDENTIALS,kai_instance = args['name'], username = user.get_unique_alias())}",401)
+        server.check_in(args['max_pixels'])
+        if server.maintenance:
+            return(f"Server has been put into maintenance mode by the owner",403)
+        if server.paused:
+            return({"id": None, "skipped": {}},200)
+        # This ensures that the priority requested by the bridge is respected
+        prioritized_wp = []
+        priority_users = [user]
+        ## Start prioritize by bridge request ##
+        for priority_username in args.priority_usernames:
+            priority_user = _db.find_user_by_username(priority_username)
+            if priority_user:
+                priority_users.append(priority_user)
+        for priority_user in priority_users:
+            for wp in waiting_prompts.get_all():
+                if wp.user == priority_user and wp.needs_gen():
+                    prioritized_wp.append(wp)
+        ## End prioritize by bridge request ##
+        for wp in waiting_prompts.get_waiting_wp_by_kudos():
+            if wp not in prioritized_wp:
+                prioritized_wp.append(wp)
+        for wp in prioritized_wp:
+            check_gen = server.can_generate(wp)
+            if not check_gen[0]:
+                skipped_reason = check_gen[1]
+                skipped[skipped_reason] = skipped.get(skipped_reason,0) + 1
+                continue
+            ret = wp.start_generation(server)
+            return(ret, 200)
+        return({"id": None, "skipped": skipped}, 200)
+
+
+class SubmitGeneration(Resource):
+    parser = reqparse.RequestParser()
+    parser.add_argument("id", type=str, required=True, help="The processing generation uuid")
+    parser.add_argument("api_key", type=str, required=True, help="The server's owner API key")
+    parser.add_argument("generation", type=str, required=False, default=[], help="The download location of the image")
+    parser.add_argument("seed", type=str, required=True, default=[], help="The seed of the generated image")
+
+    @api.expect(parser)
+    def post(self):
+        args = self.parser.parse_args()
+        procgen = processing_generations.get_item(args['id'])
+        if not procgen:
+            return(f"{get_error(ServerErrors.INVALID_PROCGEN,id = args['id'])}",404)
+        user = _db.find_user_by_api_key(args['api_key'])
+        if not user:
+            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'server submit: ' + args['name'])}",401)
+        if user != procgen.worker.user:
+            return(f"{get_error(ServerErrors.WRONG_CREDENTIALS,kai_instance = args['name'], username = user.get_unique_alias())}",401)
+        kudos = procgen.set_generation(args['generation'], args['seed'])
+        if kudos == 0:
+            return(f"{get_error(ServerErrors.DUPLICATE_GEN,id = args['id'])}",400)
+        return({"reward": kudos}, 200)
+
+class TransferKudos(Resource):
+    parser = reqparse.RequestParser()
+    parser.add_argument("username", type=str, required=True, help="The user ID which will receive the kudos")
+    parser.add_argument("api_key", type=str, required=True, help="The sending user's API key")
+    parser.add_argument("amount", type=int, required=False, default=100, help="The amount of kudos to transfer")
+
+    @api.expect(parser)
+    def post(self):
+        args = self.parser.parse_args()
+        user = _db.find_user_by_api_key(args['api_key'])
+        if not user:
+            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'kudos transfer to: ' + args['username'])}",401)
+        ret = _db.transfer_kudos_from_apikey_to_username(args['api_key'],args['username'],args['amount'])
+        kudos = ret[0]
+        error = ret[1]
+        if error != 'OK':
+            return(f"{error}",400)
+        return({"transfered": kudos}, 200)
+
+class AdminMaintenanceMode(Resource):
+    parser = reqparse.RequestParser()
+    parser.add_argument("api_key", type=str, required=True, help="The Admin API key")
+    parser.add_argument("active", type=bool, required=True, help="Star or stop maintenance mode")
+
+    decorators = [limiter.limit("30/minute")]
+    @api.expect(parser)
+    def put(self):
+        args = self.parser.parse_args()
+        admin = _db.find_user_by_api_key(args['api_key'])
+        if not admin:
+            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'Admin action: ' + 'AdminMaintenanceMode')}",401)
+        if not os.getenv("ADMINS") or admin.get_unique_alias() not in os.getenv("ADMINS"):
+            return(f"{get_error(ServerErrors.NOT_ADMIN, username = admin.get_unique_alias(), endpoint = 'AdminMaintenanceMode')}",401)
+        logger.debug(maintenance)
+        maintenance.toggle(args['active'])
+        return({"maintenance_mode": maintenance.active}, 200)
+
+class Servers(Resource):
+    @logger.catch
+    def get(self):
+        servers_ret = []
+        for server in _db.workers.values():
+            if server.is_stale():
+                continue
+            sdict = {
+                "name": server.name,
+                "id": server.id,
+                "max_pixels": server.max_pixels,
+                "megapixelsteps_generated": server.contributions,
+                "requests_fulfilled": server.fulfilments,
+                "kudos_rewards": server.kudos,
+                "kudos_details": server.kudos_details,
+                "performance": server.get_performance(),
+                "uptime": server.uptime,
+                "maintenance_mode": server.maintenance,
+            }
+            servers_ret.append(sdict)
+        return(servers_ret,200)
+
+class ServerSingle(Resource):
+    @logger.catch
+    def get(self, server_id = ''):
+        server = _db.find_worker_by_id(server_id)
+        if server:
+            sdict = {
+                "name": server.name,
+                "id": server.id,
+                "max_pixels": server.max_pixels,
+                "megapixelsteps_generated": server.contributions,
+                "requests_fulfilled": server.fulfilments,
+                "latest_performance": server.get_performance(),
+                "maintenance_mode": server.maintenance,
+            }
+            return(sdict,200)
+        else:
+            return("Not found", 404)
+
+    parser = reqparse.RequestParser()
+    parser.add_argument("api_key", type=str, required=True, help="The Admin or server owner API key")
+    parser.add_argument("maintenance", type=bool, required=False, help="Set to true to put this server into maintenance.")
+    parser.add_argument("paused", type=bool, required=False, help="Set to true to pause this server.")
+
+    decorators = [limiter.limit("30/minute")]
+    @api.expect(parser)
+    def put(self, server_id = ''):
+        server = _db.find_worker_by_id(server_id)
+        if not server:
+            return("Invalid Server ID", 404)
+        args = self.parser.parse_args()
+        admin = _db.find_user_by_api_key(args['api_key'])
+        if not admin:
+            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'User action: ' + 'PUT ServerSingle')}",401)
+        ret_dict = {}
+        # Both admins and owners can set the server to maintenance
+        if args.maintenance != None:
+            if not os.getenv("ADMINS") or admin.get_unique_alias() not in os.getenv("ADMINS"):
+                if admin != server.user:
+                    return(f"{get_error(ServerErrors.NOT_OWNER, username = admin.get_unique_alias(), server_name = server.name)}",401)
+            server.maintenance = args.maintenance
+            ret_dict["maintenance"] = server.maintenance
+        # Only admins can set a server as paused
+        if args.paused != None:
+            if not os.getenv("ADMINS") or admin.get_unique_alias() not in os.getenv("ADMINS"):
+                return(f"{get_error(ServerErrors.NOT_ADMIN, username = admin.get_unique_alias(), endpoint = 'AdminModifyServer')}",401)
+            server.paused = args.paused
+            ret_dict["paused"] = server.paused
+        if not len(ret_dict):
+            return("No server modification selected!", 400)
+        return(ret_dict, 200)
+
+    # parser = reqparse.RequestParser()
+    # parser.add_argument("api_key", type=str, required=True, help="The Admin or server owner API key")
+
+    # # post shows also hidden server info
+    # decorators = [limiter.limit("30/minute")]
+    # def post(self, server_id = ''):
+    #     server = _db.find_worker_by_id(server_id)
+    #     if not server:
+    #         return("Invalid Server ID", 404)
+    #     args = self.parser.parse_args()
+    #     admin = _db.find_user_by_api_key(args['api_key'])
+    #     if not admin:
+    #         return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'User action: ' + 'PUT ServerSingle')}",401)
+    #     sdict = {
+    #         "name": server.name,
+    #         "id": server.id,
+    #         "max_pixels": server.max_pixels,
+    #         "megapixelsteps_generated": server.contributions,
+    #         "requests_fulfilled": server.fulfilments,
+    #         "latest_performance": server.get_performance(),
+    #         "maintenance": server.maintenance,
+    #         "paused": server.paused,
+    #         "owner": server.user.get_unique_alias(),
+    #     }
+    #     return(sdict,200)
+
+
+class Users(Resource):
+    @logger.catch
+    def get(self):
+        user_dict = {}
+        for user in _db.users.values():
+            user_dict[user.get_unique_alias()] = {
+                "id": user.id,
+                "kudos": user.kudos,
+                "kudos_details": user.kudos_details,
+                "usage": user.usage,
+                "contributions": user.contributions,
+                "concurrency": user.concurrency,
+            }
+        return(user_dict,200)
+
+
+class UserSingle(Resource):
+    @logger.catch
+    def get(self, user_id = ''):
+        logger.debug(user_id)
+        user = _db.find_user_by_id(user_id)
+        if user:
+            udict = {
+                "username": user.get_unique_alias(),
+                "kudos": user.kudos,
+                "usage": user.usage,
+                "contributions": user.contributions,
+                "concurrency": user.concurrency,
+            }
+            return(udict,200)
+        else:
+            return("Not found", 404)
+
+    parser = reqparse.RequestParser()
+    parser.add_argument("api_key", type=str, required=True, help="The Admin API key")
+    parser.add_argument("kudos", type=int, required=False, help="The amount of kudos to modify (can be negative)")
+    parser.add_argument("concurrency", type=int, required=False, help="The amount of concurrent request this user can have")
+    parser.add_argument("usage_multiplier", type=float, required=False, help="The amount by which to multiply the users kudos consumption")
+
+    decorators = [limiter.limit("30/minute")]
+    @api.expect(parser)
+    def put(self, user_id = ''):
+        user = user = _db.find_user_by_id(user_id)
+        if not user:
+            return(f"Invalid user_id: {user_id}",400)        
+        args = self.parser.parse_args()
+        admin = _db.find_user_by_api_key(args['api_key'])
+        if not admin:
+            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'Admin action: ' + 'PUT UserSingle')}",401)
+        if not os.getenv("ADMINS") or admin.get_unique_alias() not in os.getenv("ADMINS"):
+            return(f"{get_error(ServerErrors.NOT_ADMIN, username = admin.get_unique_alias(), endpoint = 'AdminModifyUser')}",401)
+        ret_dict = {}
+        if args.kudos:
+            user.modify_kudos(args.kudos, 'admin')
+            ret_dict["new_kudos"] = user.kudos
+        if args.concurrency:
+            user.concurrency = args.concurrency
+            ret_dict["concurrency"] = user.concurrency
+        if args.usage_multiplier:
+            user.usage_multiplier = args.usage_multiplier
+            ret_dict["usage_multiplier"] = user.usage_multiplier
+        if not len(ret_dict):
+            return("No usermod operations selected!", 400)
+        return(ret_dict, 200)
+
+
+class HordeLoad(Resource):
+    @logger.catch
+    def get(self):
+        load_dict = waiting_prompts.count_totals()
+        load_dict["megapixelsteps_per_min"] = _db.stats.get_things_per_min()
+        load_dict["server_count"] = _db.count_active_workers()
+        load_dict["maintenance_mode"] = maintenance.active
+        return(load_dict,200)
+
+api.add_resource(SyncGenerate, "/generate/sync")
+# Async is disabled due to the memory requirements of keeping images in running memory
+api.add_resource(AsyncGenerate, "/generate/async")
+api.add_resource(AsyncStatus, "/generate/prompt/<string:id>")
+api.add_resource(AsyncCheck, "/generate/check/<string:id>")
+api.add_resource(PromptPop, "/generate/pop")
+api.add_resource(SubmitGeneration, "/generate/submit")
+api.add_resource(Users, "/users")
+api.add_resource(UserSingle, "/users/<string:user_id>")
+api.add_resource(Servers, "/servers")
+api.add_resource(ServerSingle, "/servers/<string:server_id>")
+api.add_resource(TransferKudos, "/kudos/transfer")
+api.add_resource(HordeLoad, "/status/performance")
+api.add_resource(AdminMaintenanceMode, "/admin/maintenance")

horde/apis/v2/__init__.py

@@ -0,0 +1 @@
+api = None

horde/apis/v2/stable.py

@@ -0,0 +1,56 @@
+from .v2 import *
+
+class AsyncGenerate(AsyncGenerate):
+    
+    def validate(self):
+        super().validate()
+        if self.args["params"].get("length",512)%64:
+            raise e.InvalidSize(self.username)
+        if self.args["params"].get("width",512)%64:
+            raise e.InvalidSize(self.username)
+        if self.args["params"].get("steps",50) > 100:
+            raise e.TooManySteps(self.username, self.args['params']['steps'])
+
+
+class SyncGenerate(SyncGenerate):
+
+    def validate(self):
+        super().validate()
+        if self.args["params"].get("length",512)%64:
+            raise e.InvalidSize(self.username)
+        if self.args["params"].get("width",512)%64:
+            raise e.InvalidSize(self.username)
+        if self.args["params"].get("steps",50) > 100:
+            raise e.TooManySteps(self.username, self.args['params']['steps'])
+
+class JobPop(JobPop):
+
+    def check_in(self):
+        self.worker.check_in(self.args['max_pixels'])
+  
+class HordeLoad(HordeLoad):
+    decorators = [limiter.limit("2/second")]
+    # When we extend the actual method, we need to re-apply the decorators
+    @logger.catch
+    @api.marshal_with(models.response_model_horde_maintenance_mode, code=200, description='Horde Maintenance')
+    def get(self):
+        '''Details about the current performance of this Horde
+        '''
+        load_dict = super().get()
+        load_dict["past_minute_megapixelsteps"] = db.stats.get_things_per_min()
+        return(load_dict,200)
+
+
+api.add_resource(SyncGenerate, "/generate/sync")
+api.add_resource(AsyncGenerate, "/generate/async")
+api.add_resource(AsyncStatus, "/generate/status/<string:id>")
+api.add_resource(AsyncCheck, "/generate/check/<string:id>")
+api.add_resource(JobPop, "/generate/pop")
+api.add_resource(JobSubmit, "/generate/submit")
+api.add_resource(Users, "/users")
+api.add_resource(UserSingle, "/users/<string:user_id>")
+api.add_resource(Workers, "/workers")
+api.add_resource(WorkerSingle, "/workers/<string:worker_id>")
+api.add_resource(TransferKudos, "/kudos/transfer")
+api.add_resource(HordeLoad, "/status/performance")
+api.add_resource(HordeMaintenance, "/status/maintenance")

horde/apis/v2/v2.py

@@ -0,0 +1,518 @@
+from flask_restx import Namespace, Resource, reqparse, fields, Api, abort
+from flask import request
+from ... import limiter
+from ...logger import logger
+from ...classes import db
+from ...classes import processing_generations,waiting_prompts,Worker,User,WaitingPrompt
+from ... import maintenance
+from enum import Enum
+from .. import exceptions as e
+import os, time
+from .. import ModelsV2, ParsersV2
+
+api = Namespace('v2', 'API Version 2' )
+models = ModelsV2(api)
+parsers = ParsersV2()
+
+handle_missing_prompts = api.errorhandler(e.MissingPrompt)(e.handle_bad_requests)
+handle_kudos_validation_error = api.errorhandler(e.KudosValidationError)(e.handle_bad_requests)
+handle_invalid_size = api.errorhandler(e.InvalidSize)(e.handle_bad_requests)
+handle_too_many_steps = api.errorhandler(e.TooManySteps)(e.handle_bad_requests)
+handle_invalid_api = api.errorhandler(e.InvalidAPIKey)(e.handle_bad_requests)
+handle_wrong_credentials = api.errorhandler(e.WrongCredentials)(e.handle_bad_requests)
+handle_not_admin = api.errorhandler(e.NotAdmin)(e.handle_bad_requests)
+handle_not_owner = api.errorhandler(e.NotOwner)(e.handle_bad_requests)
+handle_worker_maintenance = api.errorhandler(e.WorkerMaintenance)(e.handle_bad_requests)
+handle_invalid_procgen = api.errorhandler(e.InvalidProcGen)(e.handle_bad_requests)
+handle_request_not_found = api.errorhandler(e.RequestNotFound)(e.handle_bad_requests)
+handle_worker_not_found = api.errorhandler(e.WorkerNotFound)(e.handle_bad_requests)
+handle_user_not_found = api.errorhandler(e.UserNotFound)(e.handle_bad_requests)
+handle_duplicate_gen = api.errorhandler(e.DuplicateGen)(e.handle_bad_requests)
+handle_too_many_prompts = api.errorhandler(e.TooManyPrompts)(e.handle_bad_requests)
+handle_no_valid_workers = api.errorhandler(e.NoValidWorkers)(e.handle_bad_requests)
+handle_maintenance_mode = api.errorhandler(e.MaintenanceMode)(e.handle_bad_requests)
+
+# Used to for the flas limiter, to limit requests per url paths
+def get_request_path():
+    return(request.path)
+
+# I have to put it outside the class as I can't figure out how to extend the argparser and also pass it to the @api.expect decorator inside the class
+class GenerateTemplate(Resource):
+
+    def post(self):
+        self.args = parsers.generate_parser.parse_args()
+        self.username = 'Anonymous'
+        self.user = None
+        self.validate()
+        self.initiate_waiting_prompt()
+        worker_found = False
+        for worker in db.workers.values():
+            if len(self.args.workers) and worker.id not in self.args.workers:
+                continue
+            if worker.can_generate(self.wp)[0]:
+                worker_found = True
+                break
+        self.activate_waiting_prompt()
+
+    # We split this into its own function, so that it may be overriden and extended
+    def validate(self):
+        if maintenance.active:
+            raise e.MaintenanceMode('SyncGenerate')
+        if self.args.apikey:
+            self.user = db.find_user_by_api_key(self.args['apikey'])
+        if not self.user:
+            raise e.InvalidAPIKey('async generation')
+        self.username = self.user.get_unique_alias()
+        if self.args['prompt'] == '':
+            raise e.MissingPrompt(self.username)
+        wp_count = waiting_prompts.count_waiting_requests(self.user)
+        if wp_count >= self.user.concurrency:
+            raise e.TooManyPrompts(self.username, wp_count)
+    
+    # We split this into its own function, so that it may be overriden
+    def initiate_waiting_prompt(self):
+        self.wp = WaitingPrompt(
+            db,
+            waiting_prompts,
+            processing_generations,
+            self.args["prompt"],
+            self.user,
+            self.args["params"],
+            workers=self.args["workers"],
+        )
+    
+    # We split this into its own function, so that it may be overriden and extended
+    def activate_waiting_prompt(self):
+        self.wp.activate()
+
+    def has_valid_workers(self):
+        worker_found = False
+        for worker in db.workers.values():
+            if len(self.args.workers) and worker.id not in self.args.workers:
+                continue
+            if worker.can_generate(self.wp)[0]:
+                worker_found = True
+                break
+        return(worker_found)
+
+class AsyncGenerate(GenerateTemplate):
+
+    @api.expect(parsers.generate_parser)
+    @api.marshal_with(models.response_model_async, code=202, description='Generation Queued', skip_none=True)
+    @api.response(400, 'Validation Error', models.response_model_error)
+    @api.response(401, 'Invalid API Key', models.response_model_error)
+    @api.response(503, 'Maintenance Mode', models.response_model_error)
+    @api.response(429, 'Too Many Prompts', models.response_model_error)
+    def post(self):
+        '''Initiate an Asynchronous request to generate images.
+        This endpoint will immediately return with the UUID of the request for generation.
+        This endpoint will always be accepted, even if there are no workers available currently to fulfill this request. 
+        Perhaps some will appear in the next 10 minutes.
+        Asynchronous requests live for 10 minutes before being considered stale and being deleted.
+        '''
+        super().post()
+        ret_dict = {"id":self.wp.id}
+        if not self.has_valid_workers():
+            ret_dict['message'] = "Warning: No available workers can fulfill this request. It will expire in 10 minutes. Consider reducing the size to 512x512"
+        return(ret_dict, 202)
+
+class SyncGenerate(GenerateTemplate):
+
+    @api.expect(parsers.generate_parser)
+     # If I marshal it here, it overrides the marshalling of the child class unfortunately
+    @api.marshal_with(models.response_model_wp_status_full, code=200, description='Images Generated')
+    @api.response(400, 'Validation Error', models.response_model_error)
+    @api.response(401, 'Invalid API Key', models.response_model_error)
+    @api.response(503, 'Maintenance Mode', models.response_model_error)
+    @api.response(429, 'Too Many Prompts', models.response_model_error)
+    def post(self):
+        '''Initiate a Synchronous request to generate images.
+        This connection will only terminate when the images have been generated, or an error occured.
+        If you connection is interrupted, you will not have the request UUID, so you cannot retrieve the images asynchronously.
+        '''
+        super().post()
+        while True:
+            time.sleep(1)
+            if self.wp.is_stale():
+                raise e.RequestExpired(self.username)
+            if self.wp.is_completed():
+                break
+        ret_dict = self.wp.get_status()
+        # We delete it from memory immediately to ensure we don't run out
+        self.wp.delete()
+        return(ret_dict, 200)
+
+    # We extend this function so we can check if any workers can fulfil the request, before adding it to the queue
+    def activate_waiting_prompt(self):
+        # We don't want to keep synchronous requests up unless there's someone who can fulfill them
+        if not self.has_valid_workers():
+            # We don't need to call .delete() on the wp because it's not activated yet
+            # And therefore not added to the waiting_prompt dict.
+            raise e.NoValidWorkers(self.username)
+        # if a worker is available to fulfil this prompt, we activate it and add it to the queue to be generated
+        super().activate_waiting_prompt()
+
+class AsyncStatus(Resource):
+    decorators = [limiter.limit("2/minute", key_func = get_request_path)]
+     # If I marshal it here, it overrides the marshalling of the child class unfortunately
+    @api.marshal_with(models.response_model_wp_status_full, code=200, description='Async Request Full Status')
+    @api.response(404, 'Request Not found', models.response_model_error)
+    def get(self, id = ''):
+        '''Retrieve the full status of an Asynchronous generation request.
+        This request will include all already generated images in base64 encoded .webp files.
+        As such, you are requested to not retrieve this endpoint often. Instead use the /check/ endpoint first
+        This endpoint is limited to 1 request per minute
+        '''
+        wp = waiting_prompts.get_item(id)
+        if not wp:
+            raise e.RequestNotFound(id)
+        wp_status = wp.get_status()
+        # If the status is retrieved after the wp is done we clear it to free the ram
+        if wp_status["done"]:
+            wp.delete()
+        return(wp_status, 200)
+
+    @api.marshal_with(models.response_model_wp_status_full, code=200, description='Async Request Full Status')
+    @api.response(404, 'Request Not found', models.response_model_error)
+    def delete(self, id = ''):
+        '''Cancel an unfinished request.
+        This request will include all already generated images in base64 encoded .webp files.
+        '''
+        wp = waiting_prompts.get_item(id)
+        if not wp:
+            raise e.RequestNotFound(id)
+        wp_status = wp.get_status()
+        logger.info(f"Request with ID {wp.id} has been cancelled.")
+        wp.delete()
+        return(wp_status, 200)
+
+
+class AsyncCheck(Resource):
+    # Increasing this until I can figure out how to pass original IP from reverse proxy
+    decorators = [limiter.limit("10/second")]
+    @api.marshal_with(models.response_model_wp_status_lite, code=200, description='Async Request Status Check')
+    @api.response(404, 'Request Not found', models.response_model_error)
+    def get(self, id = ''):
+        '''Retrieve the status of an Asynchronous generation request without images.
+        Use this request to check the status of a currently running asynchronous request without consuming bandwidth.
+        '''
+        wp = waiting_prompts.get_item(id)
+        if not wp:
+            raise e.RequestNotFound(id)
+        return(wp.get_lite_status(), 200)
+
+
+class JobPop(Resource):
+
+    decorators = [limiter.limit("2/second")]
+    @api.expect(parsers.job_pop_parser)
+    @api.marshal_with(models.response_model_job_pop, code=200, description='Generation Popped')
+    @api.response(401, 'Invalid API Key', models.response_model_error)
+    @api.response(403, 'Access Denied', models.response_model_error)
+    def post(self):
+        '''Check if there are generation requests queued for fulfillment.
+        This endpoint is used by registered workers only
+        '''
+        self.args = parsers.job_pop_parser.parse_args()
+        self.validate()
+        self.check_in()
+        # Paused worker return silently
+        if self.worker.paused:
+            return({"id": None, "skipped": {}},200)
+        # This ensures that the priority requested by the bridge is respected
+        self.prioritized_wp = []
+        self.priority_users = [self.user]
+        ## Start prioritize by bridge request ##
+        for priority_username in self.args.priority_usernames:
+            priority_user = db.find_user_by_username(priority_username)
+            if priority_user:
+               self.priority_users.append(priority_user)
+        for priority_user in self.priority_users:
+            for wp in waiting_prompts.get_all():
+                if wp.user == priority_user and wp.needs_gen():
+                    self.prioritized_wp.append(wp)
+        ## End prioritize by bridge request ##
+        for wp in waiting_prompts.get_waiting_wp_by_kudos():
+            if wp not in self.prioritized_wp:
+                self.prioritized_wp.append(wp)
+        for wp in self.prioritized_wp:
+            check_gen = self.worker.can_generate(wp)
+            if not check_gen[0]:
+                skipped_reason = check_gen[1]
+                self.skipped[skipped_reason] = self.skipped.get(skipped_reason,0) + 1
+                continue
+            ret = wp.start_generation(self.worker)
+            return(ret, 200)
+        return({"id": None, "skipped": self.skipped}, 200)
+
+    # We split this into its own function, so that it may be overriden and extended
+    def validate(self):
+        self.skipped = {}
+        self.user = db.find_user_by_api_key(self.args['apikey'])
+        if not self.user:
+            raise e.InvalidAPIKey('prompt pop')
+        self.worker = db.find_worker_by_name(self.args['name'])
+        if not self.worker:
+            self.worker = Worker(db)
+            self.worker.create(self.user, self.args['name'])
+        if self.user != self.worker.user:
+            raise e.WrongCredentials(self.user.get_unique_alias(), self.args['name'])
+        if self.worker.maintenance:
+            raise e.WorkerMaintenance()
+    
+    # We split this to its own function so that it can be extended with the specific vars needed to check in
+    # You typically never want to use this template's function without extending it
+    def check_in(self):
+        self.worker.check_in()
+
+
+class JobSubmit(Resource):
+    @api.expect(parsers.job_submit_parser)
+    @api.marshal_with(models.response_model_job_submit, code=200, description='Generation Submitted')
+    @api.response(400, 'Generation Already Submitted', models.response_model_error)
+    @api.response(401, 'Invalid API Key', models.response_model_error)
+    @api.response(402, 'Access Denied', models.response_model_error)
+    @api.response(404, 'Request Not Found', models.response_model_error)
+    def post(self):
+        '''Submit a generated image.
+        This endpoint is used by registered workers only
+        '''
+        self.args = parsers.job_submit_parser.parse_args()
+        self.validate()
+        return({"reward": self.kudos}, 200)
+
+    def validate(self):
+        self.procgen = processing_generations.get_item(self.args['id'])
+        if not self.procgen:
+            raise e.InvalidProcGen(self.args['id'])
+        self.user = db.find_user_by_api_key(self.args['apikey'])
+        if not self.user:
+            raise e.InvalidAPIKey('worker submit:' + self.args['name'])
+        if self.user != self.procgen.worker.user:
+            raise e.WrongCredentials(user.get_unique_alias(), self.args['name'])
+        self.kudos = self.procgen.set_generation(self.args['generation'], self.args['seed'])
+        if self.kudos == 0:
+            raise e.DuplicateGen(self.procgen.worker.name, self.args['id'])
+
+
+class TransferKudos(Resource):
+    parser = reqparse.RequestParser()
+    parser.add_argument("apikey", type=str, required=True, help="The sending user's API key", location='headers')
+    parser.add_argument("username", type=str, required=True, help="The user ID which will receive the kudos", location="json")
+    parser.add_argument("amount", type=int, required=False, default=100, help="The amount of kudos to transfer", location="json")
+
+    @api.expect(parser)
+    @api.marshal_with(models.response_model_kudos_transfer, code=200, description='Generation Submitted')
+    @api.response(400, 'Validation Error', models.response_model_error)
+    @api.response(401, 'Invalid API Key', models.response_model_error)
+    def post(self):
+        '''Transfer Kudos to another registed user
+        '''
+        self.args = self.parser.parse_args()
+        user = db.find_user_by_api_key(self.args['apikey'])
+        if not user:
+            raise e.InvalidAPIKey('kudos transfer to: ' + self.args['username'])
+        ret = db.transfer_kudos_from_apikey_to_username(self.args['apikey'],self.args['username'],self.args['amount'])
+        kudos = ret[0]
+        error = ret[1]
+        if error != 'OK':
+            raise e.KudosValidationError(user.get_unique_alias(), error)
+        return({"transfered": kudos}, 200)
+
+class Workers(Resource):
+    @logger.catch
+    @api.marshal_with(models.response_model_worker_details, code=200, description='Workers List', as_list=True, skip_none=True)
+    def get(self):
+        '''A List with the details of all registered and active workers
+        '''
+        workers_ret = []
+        # I could do this with a comprehension, but this is clearer to understand
+        for worker in db.workers.values():
+            if worker.is_stale():
+                continue
+            workers_ret.append(worker.get_details())
+        return(workers_ret,200)
+
+class WorkerSingle(Resource):
+
+    get_parser = reqparse.RequestParser()
+    get_parser.add_argument("apikey", type=str, required=False, help="The Admin or Owner API key", location='headers')
+
+    @api.marshal_with(models.response_model_worker_details, code=200, description='Worker Details', skip_none=True)
+    @api.response(404, 'Worker Not Found', models.response_model_error)
+    def get(self, worker_id = ''):
+        '''Details of a registered worker
+        Can retrieve the details of a worker even if inactive
+        (A worker is considered inactive if it has not checked in for 5 minutes)
+        '''
+        worker = db.find_worker_by_id(worker_id)
+        if not worker:
+            raise e.WorkerNotFound(worker_id)
+        is_privileged = False
+        self.args = self.get_parser.parse_args()
+        if self.args.apikey:
+            admin = db.find_user_by_api_key(self.args['apikey'])
+            if not admin:
+                raise e.InvalidAPIKey('admin worker details')
+            if not os.getenv("ADMINS") or admin.get_unique_alias() not in os.getenv("ADMINS"):
+                raise e.NotAdmin(admin.get_unique_alias(), 'AdminWorkerDetails')
+            is_privileged = True
+        return(worker.get_details(is_privileged),200)
+
+    put_parser = reqparse.RequestParser()
+    put_parser.add_argument("apikey", type=str, required=True, help="The Admin or Owner API key", location='headers')
+    put_parser.add_argument("maintenance", type=bool, required=False, help="Set to true to put this worker into maintenance.", location="json")
+    put_parser.add_argument("paused", type=bool, required=False, help="Set to true to pause this worker.", location="json")
+
+
+    decorators = [limiter.limit("30/minute")]
+    # @api.expect(parser)
+    @api.marshal_with(models.response_model_worker_modify, code=200, description='Modify Worker', skip_none=True)
+    @api.response(400, 'Validation Error', models.response_model_error)
+    @api.response(401, 'Invalid API Key', models.response_model_error)
+    @api.response(402, 'Access Denied', models.response_model_error)
+    @api.response(404, 'Worker Not Found', models.response_model_error)
+    def put(self, worker_id = ''):
+        '''Put the worker into maintenance or pause mode
+        Maintenance can be set by the owner of the serve or an admin. 
+        When in maintenance, the worker will receive a 503 request when trying to retrieve new requests. Use this to avoid disconnecting your worker in the middle of a generation
+        Paused can be set only by the admins of this Horde.
+        When in paused mode, the worker will not be given any requests to generate.
+        '''
+        worker = db.find_worker_by_id(worker_id)
+        if not worker:
+            raise e.WorkerNotFound(worker_id)
+        self.args = self.put_parser.parse_args()
+        admin = db.find_user_by_api_key(self.args['apikey'])
+        if not admin:
+            raise e.InvalidAPIKey('User action: ' + 'PUT WorkerSingle')
+        ret_dict = {}
+        # Both admins and owners can set the worker to maintenance
+        if self.args.maintenance != None:
+            if not os.getenv("ADMINS") or admin.get_unique_alias() not in os.getenv("ADMINS"):
+                if admin != worker.user:
+                    raise e.NotOwner(admin.get_unique_alias(), worker.name)
+            worker.maintenance = self.args.maintenance
+            ret_dict["maintenance"] = worker.maintenance
+        # Only admins can set a worker as paused
+        if self.args.paused != None:
+            if not os.getenv("ADMINS") or admin.get_unique_alias() not in os.getenv("ADMINS"):
+                raise e.NotAdmin(admin.get_unique_alias(), 'AdminModifyWorker')
+            worker.paused = self.args.paused
+            ret_dict["paused"] = worker.paused
+        if not len(ret_dict):
+            raise e.NoValidActions("No worker modification selected!")
+        return(ret_dict, 200)
+
+class Users(Resource):
+    decorators = [limiter.limit("2/minute")]
+    @logger.catch
+    @api.marshal_with(models.response_model_user_details, code=200, description='Users List')
+    def get(self):
+        '''A List with the details and statistic of all registered users
+        '''
+        users_list = [user.get_details() for user in db.users.values()]
+        return(users_list,200)
+
+
+class UserSingle(Resource):
+    decorators = [limiter.limit("30/minute")]
+    @api.marshal_with(models.response_model_user_details, code=200, description='User Details')
+    @api.response(404, 'User Not Found', models.response_model_error)
+    def get(self, user_id = ''):
+        '''Details and statistics about a specific user
+        '''
+        user = db.find_user_by_id(user_id)
+        if not user:
+            raise e.UserNotFound(user_id)
+        return(user.get_details(),200)
+
+    parser = reqparse.RequestParser()
+    parser.add_argument("apikey", type=str, required=True, help="The Admin API key", location='headers')
+    parser.add_argument("kudos", type=int, required=False, help="The amount of kudos to modify (can be negative)", location="json")
+    parser.add_argument("concurrency", type=int, required=False, help="The amount of concurrent request this user can have", location="json")
+    parser.add_argument("usage_multiplier", type=float, required=False, help="The amount by which to multiply the users kudos consumption", location="json")
+
+    decorators = [limiter.limit("30/minute")]
+    @api.expect(parser)
+    @api.marshal_with(models.response_model_user_modify, code=200, description='Modify User', skip_none=True)
+    @api.response(400, 'Validation Error', models.response_model_error)
+    @api.response(401, 'Invalid API Key', models.response_model_error)
+    @api.response(402, 'Access Denied', models.response_model_error)
+    @api.response(404, 'Worker Not Found', models.response_model_error)
+    def put(self, user_id = ''):
+        '''Endpoint for horde admins to perform operations on users
+        '''
+        user = user = db.find_user_by_id(user_id)
+        if not user:
+            raise e.UserNotFound(user_id)
+        self.args = self.parser.parse_args()
+        admin = db.find_user_by_api_key(self.args['apikey'])
+        if not admin:
+            raise e.InvalidAPIKey('Admin action: ' + 'PUT UserSingle')
+        if not os.getenv("ADMINS") or admin.get_unique_alias() not in os.getenv("ADMINS"):
+            raise e.NotAdmin(admin.get_unique_alias(), 'AdminModifyUser')
+        ret_dict = {}
+        if self.args.kudos:
+            user.modify_kudos(self.args.kudos, 'admin')
+            ret_dict["new_kudos"] = user.kudos
+        if self.args.concurrency:
+            user.concurrency = self.args.concurrency
+            ret_dict["concurrency"] = user.concurrency
+        if self.args.usage_multiplier:
+            user.usage_multiplier = self.args.usage_multiplier
+            ret_dict["usage_multiplier"] = user.usage_multiplier
+        if not len(ret_dict):
+            raise e.NoValidActions("No usermod operations selected!")
+        return(ret_dict, 200)
+
+
+class HordeLoad(Resource):
+    decorators = [limiter.limit("20/minute")]
+    @logger.catch
+    @api.marshal_with(models.response_model_horde_performance, code=200, description='Horde Performance')
+    def get(self):
+        '''Details about the current performance of this Horde
+        '''
+        load_dict = waiting_prompts.count_totals()
+        load_dict["worker_count"] = db.count_active_workers()
+        return(load_dict,200)
+
+class HordeMaintenance(Resource):
+    decorators = [limiter.limit("2/second")]
+    @logger.catch
+    @api.marshal_with(models.response_model_horde_maintenance_mode, code=200, description='Horde Maintenance')
+    def get(self):
+        '''Horde Maintenance Mode Status
+        Use this endpoint to quicky determine if this horde is in maintenance.
+        '''
+        ret_dict = {
+            "maintenance_mode": maintenance.active
+        }
+        return(ret_dict,200)
+
+    parser = reqparse.RequestParser()
+    parser.add_argument("apikey", type=str, required=True, help="The Admin API key", location="headers")
+    parser.add_argument("active", type=bool, required=True, help="Star or stop maintenance mode", location="json")
+
+    decorators = [limiter.limit("30/minute")]
+    @api.expect(parser)
+    @api.marshal_with(models.response_model_admin_maintenance, code=200, description='Maintenance Mode Set')
+    @api.response(401, 'Invalid API Key', models.response_model_error)
+    @api.response(402, 'Access Denied', models.response_model_error)
+    def put(self):
+        '''Change Horde Maintenance Mode 
+        Endpoint for admins to (un)set the horde into maintenance.
+        When in maintenance no new requests for generation will be accepted
+        but requests currently in the queue will be completed.
+        '''
+        self.args = self.parser.parse_args()
+        admin = db.find_user_by_api_key(self.args['apikey'])
+        if not admin:
+            raise e.InvalidAPIKey('Admin action: ' + 'AdminMaintenanceMode')
+        if not os.getenv("ADMINS") or admin.get_unique_alias() not in os.getenv("ADMINS"):
+            raise e.NotAdmin(admin.get_unique_alias(), 'AdminMaintenanceMode')
+        maintenance.toggle(self.args['active'])
+        return({"maintenance_mode": maintenance.active}, 200)
+
+

horde/argparser.py

@@ -0,0 +1,10 @@
+import  argparse
+
+arg_parser = argparse.ArgumentParser()
+arg_parser.add_argument('-i', '--insecure', action="store_true", help="If set, will use http instead of https (useful for testing)")
+arg_parser.add_argument('-v', '--verbosity', action='count', default=0, help="The default logging level is ERROR or higher. This value increases the amount of logging seen in your screen")
+arg_parser.add_argument('-q', '--quiet', action='count', default=0, help="The default logging level is ERROR or higher. This value decreases the amount of logging seen in your screen")
+arg_parser.add_argument('-c', '--convert_flag', action='store', default=None, required=False, type=str, help="A special flag to convert from previous DB entries to newer and exit")
+arg_parser.add_argument('-p', '--port', action='store', default=7001, required=False, type=int, help="Provide a different port to start with")
+arg_parser.add_argument('--horde', action='store', default='stable', required=True, type=str, help="Which Horde to Start. This customizes endpoints and methods")
+args = arg_parser.parse_args()

horde/classes/__init__.py

@@ -0,0 +1,29 @@
+from .. import logger, args
+from importlib import import_module
+
+thing_names = {
+    "stable": "megapixelsteps",
+    "kobold": "kilotokens",
+}
+raw_thing_names = {
+    "stable": "pixelsteps",
+    "kobold": "tokens",
+}
+
+thing_name = thing_names[args.horde]
+raw_thing_name = raw_thing_names[args.horde]
+
+# Should figure out an elegant way to do this with a for loop
+WaitingPrompt = import_module(name=f'horde.classes.{args.horde}').WaitingPrompt
+ProcessingGeneration = import_module(name=f'horde.classes.{args.horde}').ProcessingGeneration
+Worker = import_module(name=f'horde.classes.{args.horde}').Worker
+PromptsIndex = import_module(name=f'horde.classes.{args.horde}').PromptsIndex
+GenerationsIndex = import_module(name=f'horde.classes.{args.horde}').GenerationsIndex
+User = import_module(name=f'horde.classes.{args.horde}').User
+Database = import_module(name=f'horde.classes.{args.horde}').Database
+
+# from .base import WaitingPrompt,ProcessingGeneration,Worker,PromptsIndex,GenerationsIndex,User,Database
+
+db = Database(convert_flag=args.convert_flag)
+waiting_prompts = PromptsIndex()
+processing_generations = GenerationsIndex()

horde/classes/base/__init__.py

@@ -2,68 +2,56 @@ import json, os, sys
 from uuid import uuid4
 from datetime import datetime
 import threading, time
-from logger import logger
-
+from .. import logger
+from .. import thing_name,raw_thing_name
 
 class WaitingPrompt:
     def __init__(self, db, wps, pgs, prompt, user, params, **kwargs):
-        self._db = db
+        self.db = db
         self._waiting_prompts = wps
         self._processing_generations = pgs
         self.prompt = prompt
         self.user = user
         self.params = params
-        self.n = params.pop('n', 1)
-        self.steps = params.pop('steps', 50)
-        # We assume more than 20 is not needed. But I'll re-evalute if anyone asks.
-        if self.n > 20:
-            logger.warning(f"User {self.user.get_unique_alias()} requested {self.n} gens per action. Reducing to 20...")
-            self.n = 20
-        self.width = params.get("width", 512)
-        self.height = params.get("height", 512)
-        # To avoid unnecessary calculations, we do it once here.
-        self.pixelsteps = self.width * self.height * self.steps
-        self.total_usage = round(self.pixelsteps * self.n / 1000000,2)
-        # The total amount of to pixelsteps requested.
+        self.total_usage = 0
+        self.extract_params(params)
         self.id = str(uuid4())
-        # This is what we send to KoboldAI to the /generate/ API
-        self.gen_payload = params
-        self.gen_payload["prompt"] = prompt
-        # We always send only 1 iteration to KoboldAI
-        self.gen_payload["batch_size"] = 1
-        self.gen_payload["ddim_steps"] = self.steps
         # The generations that have been created already
         self.processing_gens = []
         self.last_process_time = datetime.now()
-        self.servers = kwargs.get("servers", [])
+        self.workers = kwargs.get("workers", [])
         # Prompt requests are removed after 1 mins of inactivity per n, to a max of 5 minutes
         self.stale_time = 180 * self.n
         if self.stale_time > 600:
             self.stale_time = 600
 
+    # These are typically worker-specific so they will be defined in the specific class for this horde type
+    def extract_params(self, params):
+        self.n = params.pop('n', 1)
+        self.prepare_job_payload(params)
+
+    def prepare_job_payload(self, initial_dict = {}):
+        # This is what we send to the worker
+        self.gen_payload = initial_dict
 
     def activate(self):
-        # We separate the activation from __init__ as often we want to check if there's a valid server for it
-        # Before we add it to the queue
+        '''We separate the activation from __init__ as often we want to check if there's a valid worker for it
+        Before we add it to the queue
+        '''
         self._waiting_prompts.add_item(self)
-        logger.info(f"New prompt by {self.user.get_unique_alias()}: w:{self.width} * h:{self.height} * s:{self.steps} * n:{self.n} == {self.total_usage} Total MPs")
         thread = threading.Thread(target=self.check_for_stale, args=())
         thread.daemon = True
         thread.start()
 
-    # The mps still queued to be generated for this WP
-    def get_queued_megapixelsteps(self):
-        return(round(self.pixelsteps * self.n/1000000,2))
-
     def needs_gen(self):
         if self.n > 0:
             return(True)
         return(False)
 
-    def start_generation(self, server):
+    def start_generation(self, worker):
         if self.n <= 0:
             return
-        new_gen = ProcessingGeneration(self, self._processing_generations, server)
+        new_gen = self.new_procgen(worker)
         self.processing_gens.append(new_gen)
         self.n -= 1
         self.refresh()
@@ -73,6 +61,10 @@ class WaitingPrompt:
         }
         return(prompt_payload)
 
+    # Using this function so that I can extend it to have it grab the correct extended class
+    def new_procgen(self, worker):
+        return(ProcessingGeneration(self, self._processing_generations, worker))
+
     def is_completed(self):
         if self.needs_gen():
             return(False)
@@ -93,59 +85,41 @@ class WaitingPrompt:
                 ret_dict["processing"] += 1
         return(ret_dict)
 
+    # Should be overwritten
+    def get_queued_things(self):
+        return(0)
+
     def get_status(self, lite = False):
         ret_dict = self.count_processing_gens()
         ret_dict["waiting"] = self.n
         ret_dict["done"] = self.is_completed()
-        queue_pos, queued_mps, queued_n = self.get_own_queue_stats()
-        # We increment the priority by 1, because it starts at 0
-        # This means when all our requests are currently processing or done, with nothing else in the queue, we'll show queue position 0 which is appropriate.
-        ret_dict["queue_position"] = queue_pos + 1
-        active_servers = self._db.count_active_servers()
-        # If there's less requests than the number of active servers
-        # Then we need to adjust the parallelization accordingly
-        if queued_n < active_servers:
-            active_servers = queued_n
-        mpss = (self._db.stats.get_request_avg() / 1000000) * active_servers
-        # Is this is 0, it means one of two things:
-        # 1. This horde hasn't had any requests yet. So we'll initiate it to 1mpss
-        # 2. All gens for this WP are being currently processed, so we'll just set it to 1 to avoid a div by zero, but it's not used anyway as it will just divide 0/1
-        if mpss == 0:
-            mpss = 1
-        wait_time = queued_mps / mpss
-        # We add the expected running time of our processing gens
-        for procgen in self.processing_gens:
-            wait_time += procgen.get_expected_time_left()
-        ret_dict["wait_time"] = round(wait_time)
         # Lite mode does not include the generations, to spare me download size
         if not lite:
             ret_dict["generations"] = []
             for procgen in self.processing_gens:
                 if procgen.is_completed():
-                    gen_dict = {
-                        "img": procgen.generation,
-                        "seed": procgen.seed,
-                        "server_id": procgen.server.id,
-                        "server_name": procgen.server.name,
-                    }
-                    ret_dict["generations"].append(gen_dict)
+                    ret_dict["generations"].append(procgen.get_details())
         return(ret_dict)
 
-    # Same as status, but without the images to avoid unnecessary size
     def get_lite_status(self):
+        '''Same as get_status(), but without the images to avoid unnecessary size'''
         ret_dict = self.get_status(True)
         return(ret_dict)
 
-    # Get out position in the working prompts queue sorted by kudos
-    # If this gen is completed, we return (-1,-1) which represents this, to avoid doing operations.
     def get_own_queue_stats(self):
+        '''Get out position in the working prompts queue sorted by kudos
+        If this gen is completed, we return (-1,-1) which represents this, to avoid doing operations.
+        '''
         if self.needs_gen():
             return(self._waiting_prompts.get_wp_queue_stats(self))
         return(-1,0,0)
 
-    # Record that we received a requested generation and how much kudos it costs us
-    def record_usage(self, pixelsteps, kudos):
-        self.user.record_usage(pixelsteps, kudos)
+    def record_usage(self, thing, kudos):
+        '''Record that we received a requested generation and how much kudos it costs us
+        We use 'thing' here as we do not care what type of thing we're recording at this point
+        This avoids me having to extend this just to change a var name
+        '''
+        self.user.record_usage(thing, kudos)
         self.refresh()
 
     def check_for_stale(self):
@@ -173,29 +147,21 @@ class WaitingPrompt:
 
 
 class ProcessingGeneration:
-    def __init__(self, owner, pgs, server):
+    def __init__(self, owner, pgs, worker):
         self._processing_generations = pgs
         self.id = str(uuid4())
         self.owner = owner
-        self.server = server
-        self.generation = None
-        self.seed = None
-        self.kudos = 0
+        self.worker = worker
         self.start_time = datetime.now()
         self._processing_generations.add_item(self)
+        self.generation = None
 
-    def set_generation(self, generation, seed):
+    # This should be extended by every horde type
+    def set_generation(self, generation):
         if self.is_completed():
             return(0)
         self.generation = generation
-        self.seed = seed
-        pixelsteps_per_sec = self.owner._db.stats.record_fulfilment(self.owner.pixelsteps, self.start_time)
-        self.kudos = self.owner._db.convert_pixelsteps_to_kudos(self.owner.pixelsteps)
-        self.server.record_contribution(self.owner.pixelsteps, self.kudos, pixelsteps_per_sec)
-        self.owner.record_usage(self.owner.pixelsteps, self.kudos)
-        
-        logger.info(f"New Generation worth {self.kudos} kudos, delivered by server: {self.server.name}")
-        return(self.kudos)
+        return(0)
 
     def is_completed(self):
         if self.generation:
@@ -206,10 +172,14 @@ class ProcessingGeneration:
         self._processing_generations.del_item(self)
         del self
 
+    # This should be extended by every horde type
+    def get_seconds_needed(self):
+        return(0)
+
     def get_expected_time_left(self):
         if self.is_completed():
             return(0)
-        seconds_needed = self.owner.pixelsteps / self.server.get_performance_average()
+        seconds_needed = self.get_seconds_needed()
         seconds_elapsed = (datetime.now() - self.start_time).seconds
         expected_time = seconds_needed - seconds_elapsed
         # In case we run into a slow request
@@ -217,22 +187,30 @@ class ProcessingGeneration:
             expected_time = 0
         return(expected_time)
 
+    # This should be extended by every horde type
+    def get_details(self):
+        '''Returns a dictionary with details about this processing generation'''
+        ret_dict = {
+            "gen": procgen.generation,
+            "worker_id": procgen.worker.id,
+            "worker_name": procgen.worker.name,
+        }
+        return(ret_dict)
 
-
-class KAIServer:
+class Worker:
     def __init__(self, db):
-        self._db = db
+        self.db = db
         self.kudos_details = {
             "generated": 0,
             "uptime": 0,
         }
         self.last_reward_uptime = 0
-        # Every how many seconds does this server get a kudos reward
+        # Every how many seconds does this worker get a kudos reward
         self.uptime_reward_threshold = 600
-        # Maintenance can be requested by the owner of the server (to allow them to not pick up more requests)
+        # Maintenance can be requested by the owner of the worker (to allow them to not pick up more requests)
         self.maintenance = False
-        # Paused is set by the admins to prevent that server from seeing any more requests
-        # This can be used for stopping servers who misbhevave for example, without informing their owners
+        # Paused is set by the admins to prevent that worker from seeing any more requests
+        # This can be used for stopping workers who misbhevave for example, without informing their owners
         self.paused = False
 
     def create(self, user, name):
@@ -244,9 +222,10 @@ class KAIServer:
         self.kudos = 0
         self.performances = []
         self.uptime = 0
-        self._db.register_new_server(self)
+        self.db.register_new_worker(self)
 
-    def check_in(self, max_pixels):
+    # This should be overwriten by each specific horde
+    def check_in(self):
         if not self.is_stale():
             self.uptime += (datetime.now() - self.last_check_in).seconds
             # Every 10 minutes of uptime gets 100 kudos rewarded
@@ -254,15 +233,14 @@ class KAIServer:
                 kudos = 100
                 self.modify_kudos(kudos,'uptime')
                 self.user.record_uptime(kudos)
-                logger.debug(f"server '{self.name}' received {kudos} kudos for uptime of {self.uptime_reward_threshold} seconds.")
+                logger.debug(f"worker '{self.name}' received {kudos} kudos for uptime of {self.uptime_reward_threshold} seconds.")
                 self.last_reward_uptime = self.uptime
         else:
-            # If the server comes back from being stale, we just reset their last_reward_uptime
+            # If the worker comes back from being stale, we just reset their last_reward_uptime
             # So that they have to stay up at least 10 mins to get uptime kudos
             self.last_reward_uptime = self.uptime
         self.last_check_in = datetime.now()
-        self.max_pixels = max_pixels
-        logger.debug(f"Server {self.name} checked-in")
+        logger.debug(f"Worker {self.name} checked-in")
 
     def get_human_readable_uptime(self):
         if self.uptime < 60:
@@ -275,25 +253,32 @@ class KAIServer:
             return(f"{round(self.uptime/60/60/24,2)} days")
 
     def can_generate(self, waiting_prompt):
-        # takes as an argument a WaitingPrompt class and checks if this server is valid for generating it
+        # takes as an argument a WaitingPrompt class and checks if this worker is valid for generating it
         is_matching = True
         skipped_reason = None
-        # if thes server is paused, we return OK, but skip everything
-        if len(waiting_prompt.servers) >= 1 and self.id not in waiting_prompt.servers:
+        if self.is_stale():
+            # We don't consider stale workers in the request, so we don't need to report a reason
             is_matching = False
-            skipped_reason = 'server_id'
-        if self.max_pixels < waiting_prompt.width * waiting_prompt.height:
+        # if thes worker is paused, we return OK, but skip everything
+        if len(waiting_prompt.workers) >= 1 and self.id not in waiting_prompt.workers:
             is_matching = False
-            skipped_reason = 'max_pixels'
+            skipped_reason = 'worker_id'
         return([is_matching,skipped_reason])
 
+    # We split it to its own function to make it extendable for specialized calculations
+    def convert_contribution(self, thing):
+        self.contributions = round(self.contributions + thing,2)
+
     @logger.catch
-    def record_contribution(self, pixelsteps, kudos, pixelsteps_per_sec):
-        self.user.record_contributions(pixelsteps, kudos)
+    def record_contribution(self, thing, kudos, thing_per_sec):
+        '''We record the servers newest contribution
+        We do not need to know what type the contribution is, to avoid unnecessarily extending this method
+        '''
+        self.user.record_contributions(thing, kudos)
         self.modify_kudos(kudos,'generated')
-        self.contributions = round(self.contributions + pixelsteps/1000000,2) # We store them as Megapixelsteps
+        self.convert_contribution(thing)
         self.fulfilments += 1
-        self.performances.append(pixelsteps_per_sec)
+        self.performances.append(thing_per_sec)
         if len(self.performances) > 20:
             del self.performances[0]
 
@@ -305,13 +290,13 @@ class KAIServer:
         if len(self.performances):
             ret_num = sum(self.performances) / len(self.performances)
         else:
-            # Always sending at least 1 pixelstep per second, to avoid divisions by zero
+            # Always sending at least 1 thing per second, to avoid divisions by zero
             ret_num = 1
         return(ret_num)
 
     def get_performance(self):
         if len(self.performances):
-            ret_str = f'{round(sum(self.performances) / len(self.performances),1)} pixelsteps per second'
+            ret_str = f'{round(sum(self.performances) / len(self.performances),1)} {raw_thing_name} per second'
         else:
             ret_str = f'No requests fulfilled yet'
         return(ret_str)
@@ -320,17 +305,34 @@ class KAIServer:
         try:
             if (datetime.now() - self.last_check_in).seconds > 300:
                 return(True)
-        # If the last_check_in isn't set, it's a new server, so it's stale by default
+        # If the last_check_in isn't set, it's a new worker, so it's stale by default
         except AttributeError:
             return(True)
         return(False)
 
+    # Should be extended by each specific horde
+    def get_details(self, is_privileged = False):
+        '''We display these in the workers list json'''
+        ret_dict = {
+            "name": self.name,
+            "id": self.id,
+            "requests_fulfilled": self.fulfilments,
+            "kudos_rewards": self.kudos,
+            "kudos_details": self.kudos_details,
+            "performance": self.get_performance(),
+            "uptime": self.uptime,
+            "maintenance_mode": self.maintenance,
+        }
+        if is_privileged:
+            ret_dict['paused'] = self.paused
+        return(ret_dict)
+
+    # Should be extended by each specific horde
     @logger.catch
     def serialize(self):
         ret_dict = {
             "oauth_id": self.user.oauth_id,
             "name": self.name,
-            "max_pixels": self.max_pixels,
             "contributions": self.contributions,
             "fulfilments": self.fulfilments,
             "kudos": self.kudos,
@@ -346,12 +348,9 @@ class KAIServer:
 
     @logger.catch
     def deserialize(self, saved_dict, convert_flag = None):
-        self.user = self._db.find_user_by_oauth_id(saved_dict["oauth_id"])
+        self.user = self.db.find_user_by_oauth_id(saved_dict["oauth_id"])
         self.name = saved_dict["name"]
-        self.max_pixels = saved_dict["max_pixels"]
         self.contributions = saved_dict["contributions"]
-        if convert_flag == 'pixelsteps':
-            self.contributions = round(self.contributions / 50,2)
         self.fulfilments = saved_dict["fulfilments"]
         self.kudos = saved_dict.get("kudos",0)
         self.kudos_details = saved_dict.get("kudos_details",self.kudos_details)
@@ -361,7 +360,7 @@ class KAIServer:
         self.uptime = saved_dict.get("uptime",0)
         self.maintenance = saved_dict.get("maintenance",False)
         self.paused = saved_dict.get("paused",False)
-        self._db.servers[self.name] = self
+        self.db.workers[self.name] = self
 
 
 class Index:
@@ -369,14 +368,12 @@ class Index:
         self._index = {}
 
     def add_item(self, item):
-        logger.debug(item)
         self._index[item.id] = item
 
     def get_item(self, uuid):
         return(self._index.get(uuid))
 
     def del_item(self, item):
-        logger.debug(item)
         del self._index[item.id]
 
     def get_all(self):
@@ -405,15 +402,9 @@ class PromptsIndex(Index):
     def count_totals(self):
         ret_dict = {
             "queued_requests": 0,
-            # mps == Megapixelsteps
-            "queued_megapixelsteps": 0,
         }
         for wp in self._index.values():
-            ret_dict["queued_requests"] += wp.n
-            if wp.n > 0:
-                ret_dict["queued_megapixelsteps"] += wp.pixelsteps / 1000000
-        # We round the end result to avoid to many decimals
-        ret_dict["queued_megapixelsteps"] = round(ret_dict["queued_megapixelsteps"],2)
+            ret_dict["queued_requests"] += wp.n + wp.count_processing_gens()["processing"]
         return(ret_dict)
 
     def get_waiting_wp_by_kudos(self):
@@ -425,18 +416,18 @@ class PromptsIndex(Index):
         return(final_wp_list)
 
     # Returns the queue position of the provided WP based on kudos
-    # Also returns the amount of mps until the wp is generated
+    # Also returns the amount of things until the wp is generated
     # Also returns the amount of different gens queued
     def get_wp_queue_stats(self, wp):
-        mps_ahead_in_queue = 0
+        things_ahead_in_queue = 0
         n_ahead_in_queue = 0
         priority_sorted_list = self.get_waiting_wp_by_kudos()
         for iter in range(len(priority_sorted_list)):
-            mps_ahead_in_queue += priority_sorted_list[iter].get_queued_megapixelsteps()
+            things_ahead_in_queue += priority_sorted_list[iter].get_queued_things()
             n_ahead_in_queue += priority_sorted_list[iter].n
             if priority_sorted_list[iter] == wp:
-                mps_ahead_in_queue = round(mps_ahead_in_queue,2)
-                return(iter, mps_ahead_in_queue, n_ahead_in_queue)
+                things_ahead_in_queue = round(things_ahead_in_queue,2)
+                return(iter, things_ahead_in_queue, n_ahead_in_queue)
         # -1 means the WP is done and not in the queue
         return(-1,0,0)
                 
@@ -447,7 +438,7 @@ class GenerationsIndex(Index):
 
 class User:
     def __init__(self, db):
-        self._db = db
+        self.db = db
         self.kudos = 0
         self.kudos_details = {
             "accumulated": 0,
@@ -467,11 +458,11 @@ class User:
         self.last_active = datetime.now()
         self.id = 0
         self.contributions = {
-            "megapixelsteps": 0,
+            thing_name: 0,
             "fulfillments": 0
         }
         self.usage = {
-            "megapixelsteps": 0,
+            thing_name: 0,
             "requests": 0
         }
         # We allow anonymous users more leeway for the max amount of concurrent requests
@@ -485,13 +476,13 @@ class User:
         self.invite_id = invite_id
         self.creation_date = datetime.now()
         self.last_active = datetime.now()
-        self.id = self._db.register_new_user(self)
+        self.id = self.db.register_new_user(self)
         self.contributions = {
-            "megapixelsteps": 0,
+            thing_name: 0,
             "fulfillments": 0
         }
         self.usage = {
-            "megapixelsteps": 0,
+            thing_name: 0,
             "requests": 0
         }
 
@@ -504,13 +495,11 @@ class User:
     def get_unique_alias(self):
         return(f"{self.username}#{self.id}")
 
-    def record_usage(self, pixelsteps, kudos):
-        self.usage["megapixelsteps"] = round(self.usage["megapixelsteps"] + (pixelsteps * self.usage_multiplier / 1000000),2)
+    def record_usage(self, kudos):
         self.usage["requests"] += 1
         self.modify_kudos(-kudos,"accumulated")
 
-    def record_contributions(self, pixelsteps, kudos):
-        self.contributions["megapixelsteps"] = round(self.contributions["megapixelsteps"] + pixelsteps/1000000,2)
+    def record_contributions(self, kudos):
         self.contributions["fulfillments"] += 1
         self.modify_kudos(kudos,"accumulated")
 
@@ -522,6 +511,18 @@ class User:
         self.kudos = round(self.kudos + kudos, 2)
         self.kudos_details[action] = round(self.kudos_details.get(action,0) + kudos, 2)
 
+    def get_details(self):
+        ret_dict = {
+            "username": self.get_unique_alias(),
+            "id": self.id,
+            "kudos": self.kudos,
+            "kudos_details": self.kudos_details,
+            "usage": self.usage,
+            "contributions": self.contributions,
+            "concurrency": self.concurrency,
+        }
+        return(ret_dict)
+
 
     @logger.catch
     def serialize(self):
@@ -557,12 +558,6 @@ class User:
         self.usage_multiplier = saved_dict.get("usage_multiplier", 1.0)
         if self.api_key == '0000000000':
             self.concurrency = 200
-        if convert_flag == 'pixelsteps':
-            # I average to 25 steps, to convert pixels to pixelsteps, since I wasn't tracking it until now
-            self.contributions['megapixelsteps'] = round(self.contributions['pixels'] / 50,2)
-            del self.contributions['pixels']
-            self.usage['megapixelsteps'] = round(self.usage['pixels'] / 50,2)
-            del self.usage['pixels']
         self.creation_date = datetime.strptime(saved_dict["creation_date"],"%Y-%m-%d %H:%M:%S")
         self.last_active = datetime.strptime(saved_dict["last_active"],"%Y-%m-%d %H:%M:%S")
 
@@ -570,46 +565,45 @@ class User:
 class Stats:
     def __init__(self, db, convert_flag = None, interval = 60):
         self.db = db
-        self.server_performances = []
+        self.worker_performances = []
         self.fulfillments = []
         self.interval = interval
         self.last_pruning = datetime.now()
 
-    def record_fulfilment(self, pixelsteps, starting_time):
+    def record_fulfilment(self, things, starting_time):
         seconds_taken = (datetime.now() - starting_time).seconds
         if seconds_taken == 0:
-            pixelsteps_per_sec = 1
+            things_per_sec = 1
         else:
-            pixelsteps_per_sec = round(pixelsteps / seconds_taken,1)
-        if len(self.server_performances) >= 10:
-            del self.server_performances[0]
-        self.server_performances.append(pixelsteps_per_sec)
+            things_per_sec = round(things / seconds_taken,1)
+        if len(self.worker_performances) >= 10:
+            del self.worker_performances[0]
+        self.worker_performances.append(things_per_sec)
         fulfillment_dict = {
-            "pixelsteps": pixelsteps,
+            raw_thing_name: things,
             "start_time": starting_time,
             "deliver_time": datetime.now(),
         }
         self.fulfillments.append(fulfillment_dict)
-        return(pixelsteps_per_sec)
+        return(things_per_sec)
 
-    def get_megapixelsteps_per_min(self):
-        total_pixelsteps = 0
+    def get_things_per_min(self):
+        total_things = 0
         pruned_array = []
         for fulfillment in self.fulfillments:
             if (datetime.now() - fulfillment["deliver_time"]).seconds <= 60:
                 pruned_array.append(fulfillment)
-                total_pixelsteps += fulfillment["pixelsteps"]
+                total_things += fulfillment[raw_thing_name]
         if (datetime.now() - self.last_pruning).seconds > self.interval:
             self.last_pruning = datetime.now()
             self.fulfillments = pruned_array
             logger.debug("Pruned fulfillments")
-        megapixelsteps_per_min = round(total_pixelsteps / 1000000,2)
-        return(megapixelsteps_per_min)
+        return(total_things)
 
     def get_request_avg(self):
-        if len(self.server_performances) == 0:
+        if len(self.worker_performances) == 0:
             return(0)
-        avg = sum(self.server_performances) / len(self.server_performances)
+        avg = sum(self.worker_performances) / len(self.worker_performances)
         return(round(avg,1))
 
     @logger.catch
@@ -617,14 +611,13 @@ class Stats:
         serialized_fulfillments = []
         for fulfillment in self.fulfillments.copy():
             json_fulfillment = {
-                "pixelsteps": fulfillment["pixelsteps"],
+                raw_thing_name: fulfillment[raw_thing_name],
                 "start_time": fulfillment["start_time"].strftime("%Y-%m-%d %H:%M:%S"),
                 "deliver_time": fulfillment["deliver_time"].strftime("%Y-%m-%d %H:%M:%S"),
             }
             serialized_fulfillments.append(json_fulfillment)
         ret_dict = {
-            "server_performances": self.server_performances,
-            "model_mulitpliers": self.model_mulitpliers,
+            "worker_performances": self.worker_performances,
             "fulfillments": serialized_fulfillments,
         }
         return(ret_dict)
@@ -632,19 +625,18 @@ class Stats:
     @logger.catch
     def deserialize(self, saved_dict, convert_flag = None):
         # Convert old key
-        if "fulfilment_times" in saved_dict:
-            self.server_performances = saved_dict["fulfilment_times"]
+        if "server_performances" in saved_dict:
+            self.worker_performances = saved_dict["server_performances"]
         else:
-            self.server_performances = saved_dict["server_performances"]
+            self.worker_performances = saved_dict["worker_performances"]
         deserialized_fulfillments = []
         for fulfillment in saved_dict.get("fulfillments", []):
             class_fulfillment = {
-                "pixelsteps": fulfillment["pixelsteps"],
+                raw_thing_name: fulfillment[raw_thing_name],
                 "start_time": datetime.strptime(fulfillment["start_time"],"%Y-%m-%d %H:%M:%S"),
                 "deliver_time":datetime.strptime(fulfillment["deliver_time"],"%Y-%m-%d %H:%M:%S"),
             }
             deserialized_fulfillments.append(class_fulfillment)
-        self.model_mulitpliers = saved_dict["model_mulitpliers"]
         self.fulfillments = deserialized_fulfillments
        
 class Database:
@@ -652,11 +644,11 @@ class Database:
         self.interval = interval
         self.ALLOW_ANONYMOUS = True
         # This is used for synchronous generations
-        self.SERVERS_FILE = "db/servers.json"
-        self.servers = {}
+        self.WORKERS_FILE = "db/workers.json"
+        self.workers = {}
         # Other miscellaneous statistics
         self.STATS_FILE = "db/stats.json"
-        self.stats = Stats(self)
+        self.stats = self.new_stats()
         self.USERS_FILE = "db/users.json"
         self.users = {}
         # Increments any time a new user is added
@@ -672,7 +664,7 @@ class Database:
                     if not user_dict:
                         logger.error("Found null user on db load. Bypassing")
                         continue
-                    new_user = User(self)
+                    new_user = self.new_user()
                     new_user.deserialize(user_dict,convert_flag)
                     self.users[new_user.oauth_id] = new_user
                     if new_user.id > self.last_user_id:
@@ -682,16 +674,16 @@ class Database:
             self.anon = User(self)
             self.anon.create_anon()
             self.users[self.anon.oauth_id] = self.anon
-        if os.path.isfile(self.SERVERS_FILE):
-            with open(self.SERVERS_FILE) as db:
-                serialized_servers = json.load(db)
-                for server_dict in serialized_servers:
-                    if not server_dict:
-                        logger.error("Found null server on db load. Bypassing")
+        if os.path.isfile(self.WORKERS_FILE):
+            with open(self.WORKERS_FILE) as db:
+                serialized_workers = json.load(db)
+                for worker_dict in serialized_workers:
+                    if not worker_dict:
+                        logger.error("Found null worker on db load. Bypassing")
                         continue
-                    new_server = KAIServer(self)
-                    new_server.deserialize(server_dict,convert_flag)
-                    self.servers[new_server.name] = new_server
+                    new_worker = self.new_worker()
+                    new_worker.deserialize(worker_dict,convert_flag)
+                    self.workers[new_worker.name] = new_worker
         if os.path.isfile(self.STATS_FILE):
             with open(self.STATS_FILE) as stats_db:
                 self.stats.deserialize(json.load(stats_db),convert_flag)
@@ -705,6 +697,15 @@ class Database:
         thread.start()
         logger.init_ok(f"Database Load", status="Completed")
 
+    # I don't know if I'm doing this right,  but I'm using these so that I can extend them from the extended DB
+    # So that it will grab the extended classes from each horde type, and not the internal classes in this package
+    def new_worker(self):
+        return(Worker(self))
+    def new_user(self):
+        return(User(self))
+    def new_stats(self):
+        return(Stats(self))
+
     def write_files(self):
         logger.init_ok("Database Store Thread", status="Started")
         while True:
@@ -714,14 +715,14 @@ class Database:
     def write_files_to_disk(self):
         if not os.path.exists('db'):
             os.mkdir('db')
-        server_serialized_list = []
+        worker_serialized_list = []
         logger.debug("Saving DB")
-        for server in self.servers.copy().values():
-            # We don't store data for anon servers
-            if server.user == self.anon: continue
-            server_serialized_list.append(server.serialize())
-        with open(self.SERVERS_FILE, 'w') as db:
-            json.dump(server_serialized_list,db)
+        for worker in self.workers.copy().values():
+            # We don't store data for anon workers
+            if worker.user == self.anon: continue
+            worker_serialized_list.append(worker.serialize())
+        with open(self.WORKERS_FILE, 'w') as db:
+            json.dump(worker_serialized_list,db)
         with open(self.STATS_FILE, 'w') as db:
             json.dump(self.stats.serialize(),db)
         user_serialized_list = []
@@ -735,35 +736,35 @@ class Database:
         top_contributor = None
         user = None
         for user in self.users.values():
-            if user.contributions['megapixelsteps'] > top_contribution and user != self.anon:
+            if user.contributions[thing_name] > top_contribution and user != self.anon:
                 top_contributor = user
-                top_contribution = user.contributions['megapixelsteps']
+                top_contribution = user.contributions[thing_name]
         return(top_contributor)
 
-    def get_top_server(self):
-        top_server = None
-        top_server_contribution = 0
-        for server in self.servers:
-            if self.servers[server].contributions > top_server_contribution:
-                top_server = self.servers[server]
-                top_server_contribution = self.servers[server].contributions
-        return(top_server)
+    def get_top_worker(self):
+        top_worker = None
+        top_worker_contribution = 0
+        for worker in self.workers:
+            if self.workers[worker].contributions > top_worker_contribution:
+                top_worker = self.workers[worker]
+                top_worker_contribution = self.workers[worker].contributions
+        return(top_worker)
 
-    def count_active_servers(self):
+    def count_active_workers(self):
         count = 0
-        for server in self.servers.values():
-            if not server.is_stale():
+        for worker in self.workers.values():
+            if not worker.is_stale():
                 count += 1
         return(count)
 
     def get_total_usage(self):
         totals = {
-            "megapixelsteps": 0,
+            thing_name: 0,
             "fulfilments": 0,
         }
-        for server in self.servers.values():
-            totals["megapixelsteps"] += server.contributions
-            totals["fulfilments"] += server.fulfilments
+        for worker in self.workers.values():
+            totals[thing_name] += worker.contributions
+            totals["fulfilments"] += worker.fulfilments
         return(totals)
 
 
@@ -773,16 +774,16 @@ class Database:
         logger.info(f'New user created: {user.username}#{self.last_user_id}')
         return(self.last_user_id)
 
-    def register_new_server(self, server):
-        self.servers[server.name] = server
-        logger.info(f'New server checked-in: {server.name} by {server.user.get_unique_alias()}')
+    def register_new_worker(self, worker):
+        self.workers[worker.name] = worker
+        logger.info(f'New worker checked-in: {worker.name} by {worker.user.get_unique_alias()}')
 
     def find_user_by_oauth_id(self,oauth_id):
         if oauth_id == 'anon' and not self.ALLOW_ANONYMOUS:
             return(None)
         return(self.users.get(oauth_id))
 
-    def find_user_by_username(self, id):
+    def find_user_by_username(self, username):
         for user in self.users.values():
             ulist = username.split('#')
             # This approach handles someone cheekily putting # in their username
@@ -809,13 +810,13 @@ class Database:
                 return(user)
         return(None)
 
-    def find_server_by_name(self,server_name):
-        return(self.servers.get(server_name))
+    def find_worker_by_name(self,worker_name):
+        return(self.workers.get(worker_name))
 
-    def find_server_by_id(self,server_id):
-        for server in self.servers.values():
-            if server.id == server_id:
-                return(server)
+    def find_worker_by_id(self,worker_id):
+        for worker in self.workers.values():
+            if worker.id == worker_id:
+                return(worker)
         return(None)
 
     def transfer_kudos(self, source_user, dest_user, amount):
@@ -845,8 +846,10 @@ class Database:
         kudos = self.transfer_kudos_to_username(source_user, dest_username, amount)
         return(kudos)
 
-    def convert_pixelsteps_to_kudos(self, pixelsteps):
+    # Should be overriden
+    def convert_things_to_kudos(self, things):
         # The baseline for a standard generation of 512x512, 50 steps is 10 kudos
-        kudos = round(pixelsteps / (512*512*5),2)
+        kudos = round(things,2)
         # logger.info([pixels,multiplier,kudos])
         return(kudos)
+

horde/classes/maintenance.py

@@ -0,0 +1,13 @@
+class Maintenance:
+    active = False
+
+    def activate(self):
+        self.active = True
+    
+    def disable(self):
+        self.active = False
+
+    def toggle(self,value):
+        self.active = value
+
+

horde/classes/stable/__init__.py

@@ -0,0 +1,205 @@
+from ..base import *
+
+class WaitingPrompt(WaitingPrompt):
+
+    def extract_params(self, params):
+        self.n = params.pop('n', 1)
+        self.steps = params.pop('steps', 50)
+        # We assume more than 20 is not needed. But I'll re-evalute if anyone asks.
+        if self.n > 20:
+            logger.warning(f"User {self.user.get_unique_alias()} requested {self.n} gens per action. Reducing to 20...")
+            self.n = 20
+        self.width = params.get("width", 512)
+        self.height = params.get("height", 512)
+        # To avoid unnecessary calculations, we do it once here.
+        self.pixelsteps = self.width * self.height * self.steps
+        # The total amount of to pixelsteps requested.
+        self.total_usage = round(self.pixelsteps * self.n / 1000000,2)
+        self.prepare_job_payload(params)
+
+    def prepare_job_payload(self, initial_dict = {}):
+        # This is what we send to KoboldAI to the /generate/ API
+        self.gen_payload = initial_dict
+        self.gen_payload["prompt"] = self.prompt
+        # We always send only 1 iteration to KoboldAI
+        self.gen_payload["batch_size"] = 1
+        self.gen_payload["ddim_steps"] = self.steps
+
+    def activate(self):
+        # We separate the activation from __init__ as often we want to check if there's a valid worker for it
+        # Before we add it to the queue
+        super().activate()
+        logger.info(f"New prompt by {self.user.get_unique_alias()}: w:{self.width} * h:{self.height} * s:{self.steps} * n:{self.n} == {self.total_usage} Total MPs")
+
+    def new_procgen(self, worker):
+        return(ProcessingGeneration(self, self._processing_generations, worker))
+
+    def get_queued_things(self):
+        '''The mps still queued to be generated for this WP
+        For the Stable Horde, things == megapixelsteps
+        '''
+        return(round(self.pixelsteps * self.n/1000000,2))
+
+    def get_status(self, lite = False):
+        ret_dict = super().get_status(lite)
+        queue_pos, queued_mps, queued_n = self.get_own_queue_stats()
+        # We increment the priority by 1, because it starts at 0
+        # This means when all our requests are currently processing or done, with nothing else in the queue, we'll show queue position 0 which is appropriate.
+        ret_dict["queue_position"] = queue_pos + 1
+        active_workers = self.db.count_active_workers()
+        # If there's less requests than the number of active workers
+        # Then we need to adjust the parallelization accordingly
+        if queued_n < active_workers:
+            active_workers = queued_n
+        mpss = (self.db.stats.get_request_avg() / 1000000) * active_workers
+        # Is this is 0, it means one of two things:
+        # 1. This horde hasn't had any requests yet. So we'll initiate it to 1mpss
+        # 2. All gens for this WP are being currently processed, so we'll just set it to 1 to avoid a div by zero, but it's not used anyway as it will just divide 0/1
+        if mpss == 0:
+            mpss = 1
+        wait_time = queued_mps / mpss
+        # We add the expected running time of our processing gens
+        for procgen in self.processing_gens:
+            wait_time += procgen.get_expected_time_left()
+        ret_dict["wait_time"] = round(wait_time)
+        return(ret_dict)
+
+
+class ProcessingGeneration(ProcessingGeneration):
+
+    def set_generation(self, generation, seed):
+        if self.is_completed():
+            return(0)
+        self.generation = generation
+        self.seed = seed
+        pixelsteps_per_sec = self.owner.db.stats.record_fulfilment(self.owner.pixelsteps, self.start_time)
+        self.kudos = self.owner.db.convert_things_to_kudos(self.owner.pixelsteps)
+        self.worker.record_contribution(self.owner.pixelsteps, self.kudos, pixelsteps_per_sec)
+        self.owner.record_usage(self.owner.pixelsteps, self.kudos)
+        logger.info(f"New Generation worth {self.kudos} kudos, delivered by worker: {self.worker.name}")
+        return(self.kudos)
+
+    def get_seconds_needed(self):
+        return(self.owner.pixelsteps / self.worker.get_performance_average())
+
+    def get_details(self):
+        '''Returns a dictionary with details about this processing generation'''
+        ret_dict = {
+            "img": self.generation,
+            "seed": self.seed,
+            "worker_id": self.worker.id,
+            "worker_name": self.worker.name,
+        }
+        return(ret_dict)
+
+
+class Worker(Worker):
+
+    def check_in(self, max_pixels):
+        if not self.is_stale():
+            self.uptime += (datetime.now() - self.last_check_in).seconds
+            # Every 10 minutes of uptime gets 100 kudos rewarded
+            if self.uptime - self.last_reward_uptime > self.uptime_reward_threshold:
+                kudos = 100
+                self.modify_kudos(kudos,'uptime')
+                self.user.record_uptime(kudos)
+                logger.debug(f"worker '{self.name}' received {kudos} kudos for uptime of {self.uptime_reward_threshold} seconds.")
+                self.last_reward_uptime = self.uptime
+        else:
+            # If the worker comes back from being stale, we just reset their last_reward_uptime
+            # So that they have to stay up at least 10 mins to get uptime kudos
+            self.last_reward_uptime = self.uptime
+        self.last_check_in = datetime.now()
+        self.max_pixels = max_pixels
+        logger.debug(f"Worker {self.name} checked-in, offering {self.max_pixels} max pixels")
+
+    def can_generate(self, waiting_prompt):
+        can_generate = super().can_generate(waiting_prompt)
+        is_matching = can_generate[0]
+        skipped_reason = can_generate[1]
+        if self.max_pixels < waiting_prompt.width * waiting_prompt.height:
+            is_matching = False
+            skipped_reason = 'max_pixels'
+        return([is_matching,skipped_reason])
+
+    # We split it to its own function to make it extendable
+    def convert_contribution(self,pixelsteps):
+        self.contributions = round(self.contributions + pixelsteps/1000000,2)
+
+    def get_details(self, is_privileged = False):
+        ret_dict = super().get_details(is_privileged)
+        ret_dict["max_pixels"] = self.max_pixels
+        ret_dict["megapixelsteps_generated"] = self.contributions
+        return(ret_dict)
+
+    def get_performance(self):
+        if len(self.performances):
+            ret_str = f'{round(sum(self.performances) / len(self.performances) / 1000000,1)} {thing_name} per second'
+        else:
+            ret_str = f'No requests fulfilled yet'
+        return(ret_str)
+
+    @logger.catch
+    def serialize(self):
+        ret_dict = super().serialize()
+        ret_dict["max_pixels"] = self.max_pixels
+        return(ret_dict)
+
+    @logger.catch
+    def deserialize(self, saved_dict, convert_flag = None):
+        super().deserialize(saved_dict, convert_flag)
+        self.max_pixels = saved_dict["max_pixels"]
+        if convert_flag == 'pixelsteps':
+            self.contributions = round(self.contributions / 50,2)
+
+class PromptsIndex(PromptsIndex):
+
+    def count_totals(self):
+        ret_dict = {
+            "queued_requests": 0,
+            # mps == Megapixelsteps
+            "queued_megapixelsteps": 0,
+        }
+        for wp in self._index.values():
+            current_wp_queue = wp.n + wp.count_processing_gens()["processing"]
+            ret_dict["queued_requests"] += current_wp_queue
+            if current_wp_queue > 0:
+                ret_dict["queued_megapixelsteps"] += wp.pixelsteps * current_wp_queue / 1000000
+        # We round the end result to avoid to many decimals
+        ret_dict["queued_megapixelsteps"] = round(ret_dict["queued_megapixelsteps"],2)
+        return(ret_dict)
+
+                              
+class User(User):
+
+    def record_usage(self, pixelsteps, kudos):
+        super().record_usage(kudos)
+        self.usage[thing_name] = round(self.usage[thing_name] + (pixelsteps * self.usage_multiplier / 1000000),2)
+
+    def record_contributions(self, pixelsteps, kudos):
+        super().record_contributions(kudos)
+        self.contributions[thing_name] = round(self.contributions[thing_name] + pixelsteps/1000000,2)
+
+
+class Stats(Stats):
+
+    def get_things_per_min(self):
+        pixelsteps_per_min = super().get_things_per_min()
+        megapixelsteps_per_min = round(pixelsteps_per_min / 1000000,2)
+        return(megapixelsteps_per_min)
+
+class Database(Database):
+
+    def convert_things_to_kudos(self, pixelsteps):
+        # The baseline for a standard generation of 512x512, 50 steps is 10 kudos
+        kudos = round(pixelsteps / (512*512*5),2)
+        # logger.info([pixels,multiplier,kudos])
+        return(kudos)
+
+    def new_worker(self):
+        return(Worker(self))
+    def new_user(self):
+        return(User(self))
+    def new_stats(self):
+        return(Stats(self))
+

horde/flask.py

@@ -0,0 +1,6 @@
+from flask import Flask
+from werkzeug.middleware.proxy_fix import ProxyFix
+
+HORDE = Flask(__name__)
+HORDE.wsgi_app = ProxyFix(HORDE.wsgi_app, x_for=1)
+

horde/limiter.py

@@ -0,0 +1,21 @@
+from flask_limiter import Limiter
+from flask_limiter.util import get_remote_address
+from .flask import HORDE
+
+# Very basic DOS prevention
+try:
+    limiter = Limiter(
+        HORDE,
+        key_func=get_remote_address,
+        storage_uri="redis://localhost:6379/1",
+        # storage_options={"connect_timeout": 30},
+        strategy="fixed-window", # or "moving-window"
+        default_limits=["90 per minute"]
+    )
+# Allow local workatation run
+except:
+    limiter = Limiter(
+        HORDE,
+        key_func=get_remote_address,
+        default_limits=["90 per minute"]
+    )

horde/logger.py

@@ -97,3 +97,8 @@ config = {
     ],
 }
 logger.configure(**config)
+logger.disable("__main__")
+logger.warning("disabled")
+logger.enable("")
+logger.enable(None)
+

horde/routes.py

@@ -0,0 +1,221 @@
+from flask import render_template, redirect, url_for, request
+import random, time, os, oauthlib, secrets, logging
+from flask_dance.contrib.google import google
+from flask_dance.contrib.discord import discord
+from flask_dance.contrib.github import github
+from markdown import markdown
+from uuid import uuid4
+from . import logger, maintenance, HORDE
+from .classes import db
+from .classes import waiting_prompts,User
+import bleach
+
+dance_return_to = '/'
+
+@logger.catch
+@HORDE.route('/')
+def index():
+    with open('index.md') as index_file:
+        index = index_file.read()
+    top_contributor = db.get_top_contributor()
+    top_worker = db.get_top_worker()
+    align_image = 0
+    big_image = align_image
+    while big_image == align_image:
+        big_image = random.randint(1, 5)
+    if not top_contributor or not top_worker:
+        top_contributors = f'\n<img src="https://github.com/db0/Stable-Horde/blob/master/img/{big_image}.png?raw=true" width="800" />'
+    else:
+        top_contributors = f"""\n## Top Contributors
+These are the people and workers who have contributed most to this horde.
+### Users
+This is the person whose worker(s) have generated the most pixels for the horde.
+#### {top_contributor.get_unique_alias()}
+* {round(top_contributor.contributions['megapixelsteps'] / 1000,2)} Gigapixelsteps generated.
+* {top_contributor.contributions['fulfillments']} requests fulfilled.
+### Workers
+This is the worker which has generated the most pixels for the horde.
+#### {top_worker.name}
+* {round(top_worker.contributions/1000,2)} Gigapixelsteps generated.
+* {top_worker.fulfilments} request fulfillments.
+* {top_worker.get_human_readable_uptime()} uptime.
+"""
+    policies = """
+## Policies
+
+[Privacy Policy](/privacy)
+
+[Terms of Service](/terms)"""
+    totals = db.get_total_usage()
+    wp_totals = waiting_prompts.count_totals()
+    findex = index.format(
+        stable_image = align_image,
+        avg_performance= round(db.stats.get_request_avg() / 1000000,2),
+        total_pixels = round(totals["megapixelsteps"] / 1000,2),
+        total_fulfillments = totals["fulfilments"],
+        active_workers = db.count_active_workers(),
+        total_queue = wp_totals["queued_requests"],
+        total_mps = wp_totals["queued_megapixelsteps"],
+        maintenance_mode = maintenance.active,
+    )
+    head = """<head>
+    <title>Stable Horde</title>
+    <meta name="google-site-verification" content="pmLKyCEPKM5csKT9mW1ZbGLu2TX_wD0S5FCxWlmg_iI" />
+    </head>
+    """
+    return(head + markdown(findex + top_contributors + policies))
+
+
+@logger.catch
+def get_oauth_id():
+    google_data = None
+    discord_data = None
+    github_data = None
+    authorized = False
+    if google.authorized:
+        google_user_info_endpoint = '/oauth2/v2/userinfo'
+        try:
+            google_data = google.get(google_user_info_endpoint).json()
+            authorized = True
+        except oauthlib.oauth2.rfc6749.errors.TokenExpiredError:
+            pass
+    if not authorized and discord.authorized:
+        discord_info_endpoint = '/api/users/@me'
+        try:
+            discord_data = discord.get(discord_info_endpoint).json()
+            authorized = True
+        except oauthlib.oauth2.rfc6749.errors.TokenExpiredError:
+            pass
+    if not authorized and github.authorized:
+        github_info_endpoint = '/user'
+        try:
+            github_data = github.get(github_info_endpoint).json()
+            authorized = True
+        except oauthlib.oauth2.rfc6749.errors.TokenExpiredError:
+            pass
+    oauth_id = None
+    if google_data:
+        oauth_id = f'g_{google_data["id"]}'
+    elif discord_data:
+        oauth_id = f'd_{discord_data["id"]}'
+    elif github_data:
+        oauth_id = f'gh_{github_data["id"]}'
+    return(oauth_id)
+
+
+@logger.catch
+@HORDE.route('/register', methods=['GET', 'POST'])
+def register():
+    api_key = None
+    user = None
+    welcome = 'Welcome'
+    username = ''
+    pseudonymous = False
+    oauth_id = get_oauth_id()
+    if oauth_id:
+        user = db.find_user_by_oauth_id(oauth_id)
+        if user:
+            username = user.username
+    if request.method == 'POST':
+        api_key = secrets.token_urlsafe(16)
+        if user:
+            username = bleach.clean(request.form['username'])
+            user.username = username
+            user.api_key = api_key
+        else:
+            # Triggered when the user created a username without logging in
+            if not oauth_id:
+                oauth_id = str(uuid4())
+                pseudonymous = True
+            user = User(db)
+            user.create(request.form['username'], oauth_id, api_key, None)
+            username = bleach.clean(request.form['username'])
+    if user:
+        welcome = f"Welcome back {user.get_unique_alias()}"
+    return render_template('register.html',
+                           page_title="Join the Stable Horde!",
+                           welcome=welcome,
+                           user=user,
+                           api_key=api_key,
+                           username=username,
+                           pseudonymous=pseudonymous,
+                           oauth_id=oauth_id)
+
+
+@logger.catch
+@HORDE.route('/transfer', methods=['GET', 'POST'])
+def transfer():
+    src_api_key = None
+    src_user = None
+    dest_username = None
+    kudos = None
+    error = None
+    welcome = 'Welcome'
+    oauth_id = get_oauth_id()
+    if oauth_id:
+        src_user = db.find_user_by_oauth_id(oauth_id)
+        if not src_user:
+            # This probably means the user was deleted
+            oauth_id = None
+    if request.method == 'POST':
+        dest_username = request.form['username']
+        amount = request.form['amount']
+        if not amount.isnumeric():
+            kudos = 0
+            error = "Please enter a number in the kudos field"
+        # Triggered when the user submited without logging in
+        elif src_user:
+            ret = db.transfer_kudos_to_username(src_user,dest_username,int(amount))
+            kudos = ret[0]
+            error = ret[1]
+        else:
+            ret = db.transfer_kudos_from_apikey_to_username(request.form['src_api_key'],dest_username,int(amount))
+            kudos = ret[0]
+            error = ret[1]
+    if src_user:
+        welcome = f"Welcome back {src_user.get_unique_alias()}. You have {src_user.kudos} kudos remaining"
+    return render_template('transfer_kudos.html',
+                           page_title="Kudos Transfer",
+                           welcome=welcome,
+                           kudos=kudos,
+                           error=error,
+                           dest_username=dest_username,
+                           oauth_id=oauth_id)
+
+
+@HORDE.route('/google/<return_to>')
+def google_login(return_to):
+    global dance_return_to
+    dance_return_to = '/' + return_to
+    return redirect(url_for('google.login'))
+
+
+@HORDE.route('/discord/<return_to>')
+def discord_login(return_to):
+    global dance_return_to
+    dance_return_to = '/' + return_to
+    return redirect(url_for('discord.login'))
+
+
+@HORDE.route('/github/<return_to>')
+def github_login(return_to):
+    global dance_return_to
+    dance_return_to = '/' + return_to
+    return redirect(url_for('github.login'))
+
+
+@HORDE.route('/finish_dance')
+def finish_dance():
+    global dance_return_to
+    redirect_url = dance_return_to
+    dance_return_to = '/'
+    return redirect(redirect_url)
+
+
+@HORDE.route('/privacy')
+def privacy():
+    return render_template('privacy_policy.html')
+
+@HORDE.route('/terms')
+def terms():
+    return render_template('terms_of_service.html')

index.md

@@ -7,10 +7,10 @@ Also check out our sister project: [KoboldAI Horde](https://koboldai.net)
 ## Stats 
 
 * Average Recent Performance: {avg_performance} Megapixelsteps per second
-* Total GP⋅s generated: {total_pixels} Gigapixelsteps
+* Total GPS generated: {total_pixels} Gigapixelsteps
 * Total requests fulfilled: {total_fulfillments}
-* Active [Servers](/api/v1/servers): {active_servers}
-* Current Request Queue: {total_queue}
+* Active [Workers](/api/v2/workers): {active_workers}
+* Queue: {total_queue} requests for a total of {total_mps} Megapixelsteps
 
 ## Usage
 
@@ -40,6 +40,10 @@ You can make a copy of `cliRequestData_template.py` into `cliRequestData.py` and
 
 We provide a [Godot Engine plugin](https://github.com/db0/Stable-Horde-Client-Addon) to integrate Stable Horde image generation into your games.
 
+## REST API
+
+[Full Documentation](/api)
+
 ## Services
 
 * [Register New Account](/register)
@@ -47,7 +51,8 @@ We provide a [Godot Engine plugin](https://github.com/db0/Stable-Horde-Client-Ad
 
 ## Community
 
-Join us on [Discord](https://discord.gg/3DxrhksKzn)
+* Join us on [Discord](https://discord.gg/3DxrhksKzn)
+* Support the development of the Stable Horde on [Patreon](https://www.patreon.com/db0) or [Github](https://github.com/db0)
 
 ## Credits
 

requirements.txt

@@ -1,5 +1,6 @@
+werkzeug==2.1.2
 Flask
-flask_restful
+flask-restx
 flask_limiter
 waitress
 requests >= 2.27
@@ -7,4 +8,5 @@ Markdown
 flask-dance[sqla]
 blinker
 python-dotenv
-loguru
+loguru
+bleach

server.py

@@ -1,805 +1,20 @@
-from flask import Flask, render_template, redirect, url_for, request, abort
-from flask_restful import Resource, reqparse, Api
-from flask_limiter import Limiter
-from flask_limiter.util import get_remote_address
-from werkzeug.middleware.proxy_fix import ProxyFix
-from flask_dance.contrib.google import make_google_blueprint, google
-from flask_dance.contrib.discord import make_discord_blueprint, discord
-from flask_dance.contrib.github import make_github_blueprint, github
-import requests, random, time, os, oauthlib, secrets, argparse, logging
-from enum import Enum
-from markdown import markdown
 from dotenv import load_dotenv
-from uuid import uuid4
-from werkzeug.middleware.proxy_fix import ProxyFix
-from server_classes import WaitingPrompt,ProcessingGeneration,KAIServer,PromptsIndex,GenerationsIndex,User,Database
-from logger import logger, set_logger_verbosity, quiesce_logger
-
-class ServerErrors(Enum):
-    WRONG_CREDENTIALS = 0
-    INVALID_PROCGEN = 1
-    DUPLICATE_GEN = 2
-    TOO_MANY_PROMPTS = 3
-    EMPTY_PROMPT = 4
-    INVALID_API_KEY = 5
-    INVALID_SIZE = 6
-    NO_PROXY = 7
-    TOO_MANY_STEPS = 8
-    NOT_ADMIN = 9
-    MAINTENANCE_MODE = 10
-    NOT_OWNER = 11
-
-REST_API = Flask(__name__)
-REST_API.wsgi_app = ProxyFix(REST_API.wsgi_app, x_for=1)
-# Very basic DOS prevention
-try:
-    limiter = Limiter(
-        REST_API,
-        key_func=get_remote_address,
-        storage_uri="redis://localhost:6379/1",
-        # storage_options={"connect_timeout": 30},
-        strategy="fixed-window", # or "moving-window"
-        default_limits=["90 per minute"]
-    )
-# Allow local workatation run
-except:
-    limiter = Limiter(
-        REST_API,
-        key_func=get_remote_address,
-        default_limits=["90 per minute"]
-    )
-
-api = Api(REST_API)
-dance_return_to = '/'
-maintenance_mode = False
-allow_direct_connections = False
 load_dotenv()
 
-@logger.catch
-def get_error(error, **kwargs):
-    if error == ServerErrors.INVALID_API_KEY:
-        logger.warning(f'Invalid API Key sent for {kwargs["subject"]}.')
-        return(f'No user matching sent API Key. Have you remembered to register at https://stablehorde.net/register ?')
-    if error == ServerErrors.WRONG_CREDENTIALS:
-        logger.warning(f'User "{kwargs["username"]}" sent wrong credentials for utilizing instance {kwargs["kai_instance"]}')
-        return(f'wrong credentials for utilizing instance {kwargs["kai_instance"]}')
-    if error == ServerErrors.INVALID_PROCGEN:
-        logger.warning(f'Server attempted to provide generation for {kwargs["id"]} but it did not exist')
-        return(f'Processing Generation with ID {kwargs["id"]} does not exist')
-    if error == ServerErrors.DUPLICATE_GEN:
-        logger.warning(f'Server attempted to provide duplicate generation for {kwargs["id"]} ')
-        return(f'Processing Generation with ID {kwargs["id"]} already submitted')
-    if error == ServerErrors.TOO_MANY_PROMPTS:
-        logger.warning(f'User "{kwargs["username"]}" has already requested too many parallel requests ({kwargs["wp_count"]}). Aborting!')
-        return(f"Parallel requests exceeded user limit ({kwargs['wp_count']}). Please try again later or request to increase your concurrency.")
-    if error == ServerErrors.EMPTY_PROMPT:
-        logger.warning(f'User "{kwargs["username"]}" sent an empty prompt. Aborting!')
-        return("You cannot specify an empty prompt.")
-    if error == ServerErrors.INVALID_SIZE:
-        logger.warning(f'User "{kwargs["username"]}" sent an invalid size. Aborting!')
-        return("Invalid size. The image dimentions have to be multiples of 64.")
-    if error == ServerErrors.TOO_MANY_STEPS:
-        logger.warning(f'User "{kwargs["username"]}" sent too many steps ({kwargs["steps"]}). Aborting!')
-        return("Too many sampling steps. To allow resources for everyone, we allow only up to 100 steps.")
-    if error == ServerErrors.NO_PROXY:
-        logger.warning(f'Attempt to access outside reverse proxy')
-        return(f'Access allowed only through https')
-    if error == ServerErrors.NOT_ADMIN:
-        logger.warning(f'Non-admin user "{kwargs["username"]}" tried to use admin endpoint: "{kwargs["endpoint"]}". Aborting!')
-        return("You're not an admin. Sod off!")
-    if error == ServerErrors.MAINTENANCE_MODE:
-        logger.info(f'Rejecting endpoint "{kwargs["endpoint"]}" because server in maintenance mode.')
-        return("Server has enterred maintenance mode. Please try again later.")
-    if error == ServerErrors.NOT_OWNER:
-        logger.warning(f'User "{kwargs["username"]}" tried to modify server they do not own: "{kwargs["server_name"]}". Aborting!')
-        return("You're not the owner of this server!")
-
-@REST_API.before_request
-def limit_remote_addr():
-    logger.debug(request.remote_addr)
-    # if not allow_direct_connections and request.remote_addr != '127.0.0.1':
-    #     error_msg = get_error(ServerErrors.NO_PROXY)
-    #     abort(403, error_msg)
+import os, logging
+from horde import logger, args, HORDE
 
 
-@REST_API.after_request
-def after_request(response):
-    response.headers["Access-Control-Allow-Origin"] = "*"
-    response.headers["Access-Control-Allow-Methods"] = "POST, GET, OPTIONS, PUT, DELETE"
-    response.headers["Access-Control-Allow-Headers"] = "Accept, Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization"
-    return response
-
-class SyncGenerate(Resource):
-    def post(self, api_version = None):
-        parser = reqparse.RequestParser()
-        parser.add_argument("prompt", type=str, required=True, help="The prompt to generate from")
-        parser.add_argument("api_key", type=str, required=True, help="The API Key corresponding to a registered user")
-        parser.add_argument("params", type=dict, required=False, default={}, help="Extra generate params to send to the SD server")
-        parser.add_argument("servers", type=str, action='append', required=False, default=[], help="If specified, only the server with this ID will be able to generate this prompt")
-        args = parser.parse_args()
-        username = 'Anonymous'
-        user = None
-        if maintenance_mode:
-            return(f"{get_error(ServerErrors.MAINTENANCE_MODE, endpoint = 'SyncGenerate')}",503)
-        if args.api_key:
-            user = _db.find_user_by_api_key(args['api_key'])
-        if not user:
-            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'prompt generation')}",401)
-        username = user.get_unique_alias()
-        if args['prompt'] == '':
-            return(f"{get_error(ServerErrors.EMPTY_PROMPT, username = username)}",400)
-        wp_count = _waiting_prompts.count_waiting_requests(user)
-        if wp_count >= user.concurrency:
-            return(f"{get_error(ServerErrors.TOO_MANY_PROMPTS, username = username, wp_count = wp_count)}",503)
-        if args["params"].get("length",512)%64:
-            return(f"{get_error(ServerErrors.INVALID_SIZE, username = username)}",400)
-        if args["params"].get("width",512)%64:
-            return(f"{get_error(ServerErrors.INVALID_SIZE, username = username)}",400)
-        if args["params"].get("steps",50) > 100:
-            return(f"{get_error(ServerErrors.TOO_MANY_STEPS, username = username, steps = args['params']['steps'])}",400)
-        wp = WaitingPrompt(
-            _db,
-            _waiting_prompts,
-            _processing_generations,
-            args["prompt"],
-            user,
-            args["params"],
-            servers=args["servers"],
-        )
-        server_found = False
-        for server in _db.servers.values():
-            if len(args.servers) and server.id not in args.servers:
-                continue
-            if server.can_generate(wp)[0]:
-                server_found = True
-                break
-        if not server_found:
-            del wp # Normally garbage collection will handle it, but doesn't hurt to be thorough
-            return("No active server found to fulfill this request. Please Try again later...", 503)
-        # if a server is available to fulfil this prompt, we activate it and add it to the queue to be generated
-        wp.activate()
-        while True:
-            time.sleep(1)
-            if wp.is_stale():
-                return("Prompt Request Expired", 500)
-            if wp.is_completed():
-                break
-        ret_dict = wp.get_status()['generations']
-        # We delete it from memory immediately to ensure we don't run out
-        wp.delete()
-        return(ret_dict, 200)
-
-
-class AsyncGeneratePrompt(Resource):
-    decorators = [limiter.limit("3/minute")]
-    @logger.catch
-    def get(self, api_version = None, id = ''):
-        wp = _waiting_prompts.get_item(id)
-        if not wp:
-            return("ID not found", 404)
-        wp_status = wp.get_status()
-        # If the status is retrieved after the wp is done we clear it to free the ram
-        if wp_status["done"]:
-            wp.delete()
-        return(wp_status, 200)
-
-
-class AsyncCheck(Resource):
-    # Increasing this until I can figure out how to pass original IP from reverse proxy
-    decorators = [limiter.limit("10/second")]
-    @logger.catch
-    def get(self, api_version = None, id = ''):
-        wp = _waiting_prompts.get_item(id)
-        if not wp:
-            return("ID not found", 404)
-        return(wp.get_lite_status(), 200)
-
-
-class AsyncGenerate(Resource):
-    def post(self, api_version = None):
-        parser = reqparse.RequestParser()
-        parser.add_argument("prompt", type=str, required=True, help="The prompt to generate from")
-        parser.add_argument("api_key", type=str, required=True, help="The API Key corresponding to a registered user")
-        parser.add_argument("params", type=dict, required=False, default={}, help="Extra generate params to send to the SD server")
-        parser.add_argument("servers", type=str, action='append', required=False, default=[], help="If specified, only the server with this ID will be able to generate this prompt")
-        args = parser.parse_args()
-        username = 'Anonymous'
-        user = None
-        if maintenance_mode:
-            return(f"{get_error(ServerErrors.MAINTENANCE_MODE, endpoint = 'AsyncGenerate')}",503)
-        if args.api_key:
-            user = _db.find_user_by_api_key(args['api_key'])
-        if not user:
-            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'prompt generation')}",401)
-        username = user.get_unique_alias()
-        if args['prompt'] == '':
-            return(f"{get_error(ServerErrors.EMPTY_PROMPT, username = username)}",400)
-        wp_count = _waiting_prompts.count_waiting_requests(user)
-        if wp_count >= user.concurrency:
-            return(f"{get_error(ServerErrors.TOO_MANY_PROMPTS, username = username, wp_count = wp_count)}",503)
-        if args["params"].get("length",512)%64:
-            return(f"{get_error(ServerErrors.INVALID_SIZE, username = username)}",400)
-        if args["params"].get("width",512)%64:
-            return(f"{get_error(ServerErrors.INVALID_SIZE, username = username)}",400)
-        if args["params"].get("steps",50) > 100:
-            return(f"{get_error(ServerErrors.TOO_MANY_STEPS, username = username, steps = args['params']['steps'])}",400)
-        wp = WaitingPrompt(
-            _db,
-            _waiting_prompts,
-            _processing_generations,
-            args["prompt"],
-            user,
-            args["params"],
-            servers=args["servers"],
-        )
-        server_found = False
-        for server in _db.servers.values():
-            if len(args.servers) and server.id not in args.servers:
-                continue
-            if server.can_generate(wp)[0]:
-                server_found = True
-                break
-        if not server_found:
-            del wp # Normally garbage collection will handle it, but doesn't hurt to be thorough
-            return("No active server found to fulfill this request. Please Try again later...", 503)
-        # if a server is available to fulfil this prompt, we activate it and add it to the queue to be generated
-        wp.activate()
-        return({"id":wp.id}, 200)
-
-
-class PromptPop(Resource):
-    decorators = [limiter.limit("45/second")]
-    def post(self, api_version = None):
-        parser = reqparse.RequestParser()
-        parser.add_argument("api_key", type=str, required=True, help="The API Key corresponding to a registered user")
-        parser.add_argument("name", type=str, required=True, help="The server's unique name, to track contributions")
-        parser.add_argument("max_pixels", type=int, required=False, default=512, help="The maximum amount of pixels this server can generate")
-        parser.add_argument("priority_usernames", type=str, action='append', required=False, default=[], help="The usernames which get priority use on this server")
-        args = parser.parse_args()
-        skipped = {}
-        user = _db.find_user_by_api_key(args['api_key'])
-        if not user:
-            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'server promptpop: ' + args['name'])}",401)
-        server = _db.find_server_by_name(args['name'])
-        if not server:
-            server = KAIServer(_db)
-            server.create(user, args['name'])
-        if user != server.user:
-            return(f"{get_error(ServerErrors.WRONG_CREDENTIALS,kai_instance = args['name'], username = user.get_unique_alias())}",401)
-        server.check_in(args['max_pixels'])
-        if server.maintenance:
-            return(f"Server has been put into maintenance mode by the owner",403)
-        if server.paused:
-            return({"id": None, "skipped": {}},200)
-        # This ensures that the priority requested by the bridge is respected
-        prioritized_wp = []
-        priority_users = [user]
-        ## Start prioritize by bridge request ##
-        for priority_username in args.priority_usernames:
-            priority_user = _db.find_user_by_username(priority_username)
-            if priority_user:
-                priority_users.append(priority_user)
-        for priority_user in priority_users:
-            for wp in _waiting_prompts.get_all():
-                if wp.user == priority_user and wp.needs_gen():
-                    prioritized_wp.append(wp)
-        ## End prioritize by bridge request ##
-        for wp in _waiting_prompts.get_waiting_wp_by_kudos():
-            if wp not in prioritized_wp:
-                prioritized_wp.append(wp)
-        for wp in prioritized_wp:
-            check_gen = server.can_generate(wp)
-            if not check_gen[0]:
-                skipped_reason = check_gen[1]
-                skipped[skipped_reason] = skipped.get(skipped_reason,0) + 1
-                continue
-            ret = wp.start_generation(server)
-            return(ret, 200)
-        return({"id": None, "skipped": skipped}, 200)
-
-
-class SubmitGeneration(Resource):
-    def post(self, api_version = None):
-        parser = reqparse.RequestParser()
-        parser.add_argument("id", type=str, required=True, help="The processing generation uuid")
-        parser.add_argument("api_key", type=str, required=True, help="The server's owner API key")
-        parser.add_argument("generation", type=str, required=False, default=[], help="The download location of the image")
-        parser.add_argument("seed", type=str, required=True, default=[], help="The seed of the generated image")
-        args = parser.parse_args()
-        procgen = _processing_generations.get_item(args['id'])
-        if not procgen:
-            return(f"{get_error(ServerErrors.INVALID_PROCGEN,id = args['id'])}",404)
-        user = _db.find_user_by_api_key(args['api_key'])
-        if not user:
-            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'server submit: ' + args['name'])}",401)
-        if user != procgen.server.user:
-            return(f"{get_error(ServerErrors.WRONG_CREDENTIALS,kai_instance = args['name'], username = user.get_unique_alias())}",401)
-        kudos = procgen.set_generation(args['generation'], args['seed'])
-        if kudos == 0:
-            return(f"{get_error(ServerErrors.DUPLICATE_GEN,id = args['id'])}",400)
-        return({"reward": kudos}, 200)
-
-class TransferKudos(Resource):
-    def post(self, api_version = None):
-        parser = reqparse.RequestParser()
-        parser.add_argument("username", type=str, required=True, help="The user ID which will receive the kudos")
-        parser.add_argument("api_key", type=str, required=True, help="The sending user's API key")
-        parser.add_argument("amount", type=int, required=False, default=100, help="The amount of kudos to transfer")
-        args = parser.parse_args()
-        user = _db.find_user_by_api_key(args['api_key'])
-        if not user:
-            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'kudos transfer to: ' + args['username'])}",401)
-        ret = _db.transfer_kudos_from_apikey_to_username(args['api_key'],args['username'],args['amount'])
-        kudos = ret[0]
-        error = ret[1]
-        if error != 'OK':
-            return(f"{error}",400)
-        return({"transfered": kudos}, 200)
-
-class AdminMaintenanceMode(Resource):
-    decorators = [limiter.limit("30/minute")]
-    def put(self, api_version = None):
-        global maintenance_mode
-        parser = reqparse.RequestParser()
-        parser.add_argument("api_key", type=str, required=True, help="The Admin API key")
-        parser.add_argument("active", type=bool, required=True, help="Star or stop maintenance mode")
-        args = parser.parse_args()
-        admin = _db.find_user_by_api_key(args['api_key'])
-        if not admin:
-            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'Admin action: ' + 'AdminMaintenanceMode')}",401)
-        if not os.getenv("ADMINS") or admin.get_unique_alias() not in os.getenv("ADMINS"):
-            return(f"{get_error(ServerErrors.NOT_ADMIN, username = admin.get_unique_alias(), endpoint = 'AdminMaintenanceMode')}",401)
-        maintenance_mode = args['active']
-        return({"maintenance_mode": maintenance_mode}, 200)
-
-class Servers(Resource):
-    @logger.catch
-    def get(self, api_version = None):
-        servers_ret = []
-        for server in _db.servers.values():
-            if server.is_stale():
-                continue
-            sdict = {
-                "name": server.name,
-                "id": server.id,
-                "max_pixels": server.max_pixels,
-                "megapixelsteps_generated": server.contributions,
-                "requests_fulfilled": server.fulfilments,
-                "kudos_rewards": server.kudos,
-                "kudos_details": server.kudos_details,
-                "performance": server.get_performance(),
-                "uptime": server.uptime,
-                "maintenance_mode": server.maintenance,
-            }
-            servers_ret.append(sdict)
-        return(servers_ret,200)
-
-class ServerSingle(Resource):
-    @logger.catch
-    def get(self, api_version = None, server_id = ''):
-        server = _db.find_server_by_id(server_id)
-        if server:
-            sdict = {
-                "name": server.name,
-                "id": server.id,
-                "max_pixels": server.max_pixels,
-                "megapixelsteps_generated": server.contributions,
-                "requests_fulfilled": server.fulfilments,
-                "latest_performance": server.get_performance(),
-                "maintenance_mode": server.maintenance,
-            }
-            return(sdict,200)
-        else:
-            return("Not found", 404)
-
-    decorators = [limiter.limit("30/minute")]
-    def put(self, api_version = None, server_id = ''):
-        server = _db.find_server_by_id(server_id)
-        if not server:
-            return("Invalid Server ID", 404)
-        parser = reqparse.RequestParser()
-        parser.add_argument("api_key", type=str, required=True, help="The Admin or server owner API key")
-        parser.add_argument("maintenance", type=bool, required=False, help="Set to true to put this server into maintenance.")
-        parser.add_argument("paused", type=bool, required=False, help="Set to true to pause this server.")
-        args = parser.parse_args()
-        admin = _db.find_user_by_api_key(args['api_key'])
-        if not admin:
-            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'User action: ' + 'PUT ServerSingle')}",401)
-        ret_dict = {}
-        # Both admins and owners can set the server to maintenance
-        if args.maintenance != None:
-            if not os.getenv("ADMINS") or admin.get_unique_alias() not in os.getenv("ADMINS"):
-                if admin != server.user:
-                    return(f"{get_error(ServerErrors.NOT_OWNER, username = admin.get_unique_alias(), server_name = server.name)}",401)
-            server.maintenance = args.maintenance
-            ret_dict["maintenance"] = server.maintenance
-        # Only admins can set a server as paused
-        if args.paused != None:
-            if not os.getenv("ADMINS") or admin.get_unique_alias() not in os.getenv("ADMINS"):
-                return(f"{get_error(ServerErrors.NOT_ADMIN, username = admin.get_unique_alias(), endpoint = 'AdminModifyServer')}",401)
-            server.paused = args.paused
-            ret_dict["paused"] = server.paused
-        if not len(ret_dict):
-            return("No server modification selected!", 400)
-        return(ret_dict, 200)
-
-    # post shows also hidden server info
-    decorators = [limiter.limit("30/minute")]
-    def post(self, api_version = None, server_id = ''):
-        server = _db.find_server_by_id(server_id)
-        if not server:
-            return("Invalid Server ID", 404)
-        parser = reqparse.RequestParser()
-        parser.add_argument("api_key", type=str, required=True, help="The Admin or server owner API key")
-        args = parser.parse_args()
-        admin = _db.find_user_by_api_key(args['api_key'])
-        if not admin:
-            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'User action: ' + 'PUT ServerSingle')}",401)
-        sdict = {
-            "name": server.name,
-            "id": server.id,
-            "max_pixels": server.max_pixels,
-            "megapixelsteps_generated": server.contributions,
-            "requests_fulfilled": server.fulfilments,
-            "latest_performance": server.get_performance(),
-            "maintenance": server.maintenance,
-            "paused": server.paused,
-            "owner": server.user.get_unique_alias(),
-        }
-        return(sdict,200)
-
-
-class Users(Resource):
-    @logger.catch
-    def get(self, api_version = None):
-        user_dict = {}
-        for user in _db.users.values():
-            user_dict[user.get_unique_alias()] = {
-                "id": user.id,
-                "kudos": user.kudos,
-                "kudos_details": user.kudos_details,
-                "usage": user.usage,
-                "contributions": user.contributions,
-                "concurrency": user.concurrency,
-            }
-        return(user_dict,200)
-
-
-class UserSingle(Resource):
-    @logger.catch
-    def get(self, api_version = None, user_id = ''):
-        logger.debug(user_id)
-        user = _db.find_user_by_id(user_id)
-        if user:
-            udict = {
-                "username": user.get_unique_alias(),
-                "kudos": user.kudos,
-                "usage": user.usage,
-                "contributions": user.contributions,
-                "concurrency": user.concurrency,
-            }
-            return(udict,200)
-        else:
-            return("Not found", 404)
-
-    decorators = [limiter.limit("30/minute")]
-    def put(self, api_version = None, user_id = ''):
-        user = user = _db.find_user_by_id(user_id)
-        if not user:
-            return(f"Invalid user_id: {user_id}",400)        
-        parser = reqparse.RequestParser()
-        parser.add_argument("api_key", type=str, required=True, help="The Admin API key")
-        parser.add_argument("kudos", type=int, required=False, help="The amount of kudos to modify (can be negative)")
-        parser.add_argument("concurrency", type=int, required=False, help="The amount of concurrent request this user can have")
-        parser.add_argument("usage_multiplier", type=float, required=False, help="The amount by which to multiply the users kudos consumption")
-        args = parser.parse_args()
-        admin = _db.find_user_by_api_key(args['api_key'])
-        if not admin:
-            return(f"{get_error(ServerErrors.INVALID_API_KEY, subject = 'Admin action: ' + 'PUT UserSingle')}",401)
-        if not os.getenv("ADMINS") or admin.get_unique_alias() not in os.getenv("ADMINS"):
-            return(f"{get_error(ServerErrors.NOT_ADMIN, username = admin.get_unique_alias(), endpoint = 'AdminModifyUser')}",401)
-        ret_dict = {}
-        if args.kudos:
-            user.modify_kudos(args.kudos, 'admin')
-            ret_dict["new_kudos"] = user.kudos
-        if args.concurrency:
-            user.concurrency = args.concurrency
-            ret_dict["concurrency"] = user.concurrency
-        if args.usage_multiplier:
-            user.usage_multiplier = args.usage_multiplier
-            ret_dict["usage_multiplier"] = user.usage_multiplier
-        if not len(ret_dict):
-            return("No usermod operations selected!", 400)
-        return(ret_dict, 200)
-
-
-class HordeLoad(Resource):
-    @logger.catch
-    def get(self, api_version = None):
-        load_dict = _waiting_prompts.count_totals()
-        load_dict["megapixelsteps_per_min"] = _db.stats.get_megapixelsteps_per_min()
-        load_dict["server_count"] = _db.count_active_servers()
-        load_dict["maintenance_mode"] = maintenance_mode
-        return(load_dict,200)
-
-
-@logger.catch
-@REST_API.route('/')
-def index():
-    with open('index.md') as index_file:
-        index = index_file.read()
-    top_contributor = _db.get_top_contributor()
-    top_server = _db.get_top_server()
-    align_image = 0
-    big_image = align_image
-    while big_image == align_image:
-        big_image = random.randint(1, 5)
-    if not top_contributor or not top_server:
-        top_contributors = f'\n<img src="https://github.com/db0/Stable-Horde/blob/master/img/{big_image}.png?raw=true" width="800" />'
-    else:
-        top_contributors = f"""\n## Top Contributors
-These are the people and servers who have contributed most to this horde.
-### Users
-This is the person whose server(s) have generated the most pixels for the horde.
-#### {top_contributor.get_unique_alias()}
-* {round(top_contributor.contributions['megapixelsteps'] / 1000,2)} Gigapixelsteps generated.
-* {top_contributor.contributions['fulfillments']} requests fulfilled.
-### Servers
-This is the server which has generated the most pixels for the horde.
-#### {top_server.name}
-* {round(top_server.contributions/1000,2)} Gigapixelsteps generated.
-* {top_server.fulfilments} request fulfillments.
-* {top_server.get_human_readable_uptime()} uptime.
-"""
-    policies = """
-## Policies
-
-[Privacy Policy](/privacy)
-
-[Terms of Service](/terms)"""
-    totals = _db.get_total_usage()
-    findex = index.format(
-        stable_image = align_image,
-        avg_performance= round(_db.stats.get_request_avg() / 1000000,2),
-        total_pixels = round(totals["megapixelsteps"] / 1000,2),
-        total_fulfillments = totals["fulfilments"],
-        active_servers = _db.count_active_servers(),
-        total_queue = _waiting_prompts.count_total_waiting_generations(),
-        maintenance_mode = maintenance_mode,
-    )
-    head = """<head>
-    <title>Stable Horde</title>
-    <meta name="google-site-verification" content="pmLKyCEPKM5csKT9mW1ZbGLu2TX_wD0S5FCxWlmg_iI" />
-    </head>
-    """
-    return(head + markdown(findex + top_contributors + policies))
-
-@logger.catch
-def get_oauth_id():
-    google_data = None
-    discord_data = None
-    github_data = None
-    authorized = False
-    if google.authorized:
-        google_user_info_endpoint = '/oauth2/v2/userinfo'
-        try:
-            google_data = google.get(google_user_info_endpoint).json()
-            authorized = True
-        except oauthlib.oauth2.rfc6749.errors.TokenExpiredError:
-            pass
-    if not authorized and discord.authorized:
-        discord_info_endpoint = '/api/users/@me'
-        try:
-            discord_data = discord.get(discord_info_endpoint).json()
-            authorized = True
-        except oauthlib.oauth2.rfc6749.errors.TokenExpiredError:
-            pass
-    if not authorized and github.authorized:
-        github_info_endpoint = '/user'
-        try:
-            github_data = github.get(github_info_endpoint).json()
-            authorized = True
-        except oauthlib.oauth2.rfc6749.errors.TokenExpiredError:
-            pass
-    oauth_id = None
-    if google_data:
-        oauth_id = f'g_{google_data["id"]}'
-    elif discord_data:
-        oauth_id = f'd_{discord_data["id"]}'
-    elif github_data:
-        oauth_id = f'gh_{github_data["id"]}'
-    return(oauth_id)
-
-
-@logger.catch
-@REST_API.route('/register', methods=['GET', 'POST'])
-def register():
-    api_key = None
-    user = None
-    welcome = 'Welcome'
-    username = ''
-    pseudonymous = False
-    oauth_id = get_oauth_id()
-    if oauth_id:
-        user = _db.find_user_by_oauth_id(oauth_id)
-        if user:
-            username = user.username
-    if request.method == 'POST':
-        api_key = secrets.token_urlsafe(16)
-        if user:
-            username = request.form['username']
-            user.username = request.form['username']
-            user.api_key = api_key
-        else:
-            # Triggered when the user created a username without logging in
-            if not oauth_id:
-                oauth_id = str(uuid4())
-                pseudonymous = True
-            user = User(_db)
-            user.create(request.form['username'], oauth_id, api_key, None)
-            username = request.form['username']
-    if user:
-        welcome = f"Welcome back {user.get_unique_alias()}"
-    return render_template('register.html',
-                           page_title="Join the Stable Horde!",
-                           welcome=welcome,
-                           user=user,
-                           api_key=api_key,
-                           username=username,
-                           pseudonymous=pseudonymous,
-                           oauth_id=oauth_id)
-
-
-@logger.catch
-@REST_API.route('/transfer', methods=['GET', 'POST'])
-def transfer():
-    src_api_key = None
-    src_user = None
-    dest_username = None
-    kudos = None
-    error = None
-    welcome = 'Welcome'
-    oauth_id = get_oauth_id()
-    if oauth_id:
-        src_user = _db.find_user_by_oauth_id(oauth_id)
-        if not src_user:
-            # This probably means the user was deleted
-            oauth_id = None
-    if request.method == 'POST':
-        dest_username = request.form['username']
-        amount = request.form['amount']
-        if not amount.isnumeric():
-            kudos = 0
-            error = "Please enter a number in the kudos field"
-        # Triggered when the user submited without logging in
-        elif src_user:
-            ret = _db.transfer_kudos_to_username(src_user,dest_username,int(amount))
-            kudos = ret[0]
-            error = ret[1]
-        else:
-            ret = _db.transfer_kudos_from_apikey_to_username(request.form['src_api_key'],dest_username,int(amount))
-            kudos = ret[0]
-            error = ret[1]
-    if src_user:
-        welcome = f"Welcome back {src_user.get_unique_alias()}. You have {src_user.kudos} kudos remaining"
-    return render_template('transfer_kudos.html',
-                           page_title="Kudos Transfer",
-                           welcome=welcome,
-                           kudos=kudos,
-                           error=error,
-                           dest_username=dest_username,
-                           oauth_id=oauth_id)
-
-
-@REST_API.route('/google/<return_to>')
-def google_login(return_to):
-    global dance_return_to
-    dance_return_to = '/' + return_to
-    return redirect(url_for('google.login'))
-
-
-@REST_API.route('/discord/<return_to>')
-def discord_login(return_to):
-    global dance_return_to
-    dance_return_to = '/' + return_to
-    return redirect(url_for('discord.login'))
-
-
-@REST_API.route('/github/<return_to>')
-def github_login(return_to):
-    global dance_return_to
-    dance_return_to = '/' + return_to
-    return redirect(url_for('github.login'))
-
-
-@REST_API.route('/finish_dance')
-def finish_dance():
-    global dance_return_to
-    redirect_url = dance_return_to
-    dance_return_to = '/'
-    return redirect(redirect_url)
-
-
-@REST_API.route('/privacy')
-def privacy():
-    return render_template('privacy_policy.html')
-
-@REST_API.route('/terms')
-def terms():
-    return render_template('terms_of_service.html')
-
-
-arg_parser = argparse.ArgumentParser()
-arg_parser.add_argument('-i', '--insecure', action="store_true", help="If set, will use http instead of https (useful for testing)")
-arg_parser.add_argument('-v', '--verbosity', action='count', default=0, help="The default logging level is ERROR or higher. This value increases the amount of logging seen in your screen")
-arg_parser.add_argument('-q', '--quiet', action='count', default=0, help="The default logging level is ERROR or higher. This value decreases the amount of logging seen in your screen")
-arg_parser.add_argument('-c', '--convert_flag', action='store', default=None, required=False, type=str, help="A special flag to convert from previous DB entries to newer and exit")
-arg_parser.add_argument('-p', '--port', action='store', default=7001, required=False, type=int, help="Provide a different port to start with")
-arg_parser.add_argument('--allow_direct_connections', action="store_true", required=False, default=False, help="If set, will allow connections outside the reverse proxy (useful for testing)")
 
 if __name__ == "__main__":
-    global _db
-    global _waiting_prompts
-    global _processing_generations
-    args = arg_parser.parse_args()
-    allow_direct_connections = args.allow_direct_connections
-    set_logger_verbosity(args.verbosity)
-    quiesce_logger(args.quiet)    
     # Only setting this for the WSGI logs
     logging.basicConfig(format='%(asctime)s - %(levelname)s - %(module)s:%(lineno)d - %(message)s',level=logging.WARNING)
-    _db = Database(convert_flag=args.convert_flag)
-    _waiting_prompts = PromptsIndex()
-    _processing_generations = GenerationsIndex()
-    google_client_id = os.getenv("GOOGLE_CLIENT_ID")
-    google_client_secret = os.getenv("GLOOGLE_CLIENT_SECRET")
-    discord_client_id = os.getenv("DISCORD_CLIENT_ID")
-    discord_client_secret = os.getenv("DISCORD_CLIENT_SECRET")
-    github_client_id = os.getenv("GITHUB_CLIENT_ID")
-    github_client_secret = os.getenv("GITHUB_CLIENT_SECRET")
-    REST_API.secret_key = os.getenv("secret_key")
-    os.environ['OAUTHLIB_RELAX_TOKEN_SCOPE'] = '1'
+    from waitress import serve
+    logger.init("WSGI Server", status="Starting")
     url_scheme = 'https'
     if args.insecure:
         os.environ['OAUTHLIB_INSECURE_TRANSPORT'] = '1' # Disable this on prod
         url_scheme = 'http'
-    google_blueprint = make_google_blueprint(
-        client_id = google_client_id,
-        client_secret = google_client_secret,
-        reprompt_consent = True,
-        redirect_url='/register',
-        scope = ["email"],
-    )
-    REST_API.register_blueprint(google_blueprint,url_prefix="/google")
-    discord_blueprint = make_discord_blueprint(
-        client_id = discord_client_id,
-        client_secret = discord_client_secret,
-        scope = ["identify"],
-        redirect_url='/finish_dance',
-    )
-    REST_API.register_blueprint(discord_blueprint,url_prefix="/discord")
-    github_blueprint = make_github_blueprint(
-        client_id = github_client_id,
-        client_secret = github_client_secret,
-        scope = ["identify"],
-        redirect_url='/finish_dance',
-    )
-    REST_API.register_blueprint(github_blueprint,url_prefix="/github")
-    api.add_resource(SyncGenerate, "/generate/sync","/api/<string:api_version>/generate/sync")
-    # Async is disabled due to the memory requirements of keeping images in running memory
-    api.add_resource(AsyncGenerate, "/generate/async","/api/<string:api_version>/generate/async")
-    api.add_resource(AsyncGeneratePrompt, "/generate/prompt/<string:id>","/api/<string:api_version>/generate/prompt/<string:id>")
-    api.add_resource(AsyncCheck, "/generate/prompt/<string:id>","/api/<string:api_version>/generate/check/<string:id>")
-    api.add_resource(PromptPop, "/generate/pop","/api/<string:api_version>/generate/pop")
-    api.add_resource(SubmitGeneration, "/generate/submit","/api/<string:api_version>/generate/submit")
-    api.add_resource(Users, "/users","/api/<string:api_version>/users")
-    api.add_resource(UserSingle, "/users/<string:user_id>","/api/<string:api_version>/users/<string:user_id>")
-    api.add_resource(Servers, "/servers","/api/<string:api_version>/servers")
-    api.add_resource(ServerSingle, "/servers/<string:server_id>","/api/<string:api_version>/servers/<string:server_id>")
-    api.add_resource(TransferKudos, "/api/<string:api_version>/kudos/transfer")
-    api.add_resource(HordeLoad, "/api/<string:api_version>/status/performance")
-    api.add_resource(AdminMaintenanceMode, "/api/<string:api_version>/admin/maintenance")
-    from waitress import serve
-    logger.init("WSGI Server", status="Starting")
-    serve(REST_API, host="0.0.0.0", port=args.port, url_scheme=url_scheme, threads=100, connection_limit=4096)
-    # REST_API.run(debug=True,host="0.0.0.0",port="5001")
+    serve(HORDE, host="127.0.0.1", port=args.port, url_scheme=url_scheme, threads=300, connection_limit=4096)
+    # HORDE.run(debug=True,host="0.0.0.0",port="5001")
     logger.init("WSGI Server", status="Stopped")