ai-art

  1"""
  2A simple flask server for serving a web interface to change the config and view
  3the contents of the "/data/outputs" folder.
  4"""
  5import datetime
  6import os
  7import json
  8import subprocess
  9
 10from flask import (
 11    Blueprint,
 12    Flask,
 13    render_template,
 14    send_from_directory,
 15    request,
 16    redirect,
 17    send_file,
 18)
 19from PIL import Image
 20
 21from .check import check_files_and_folders
 22
 23
 24web = Blueprint("web", __name__, template_folder="templates")
 25
 26
 27@web.route("/")
 28def index():
 29    current_config = open("/data/config.json", "r").read()
 30    current_config = json.loads(current_config)
 31    current_config["prompts"] = ", ".join(current_config["prompts"])
 32    return render_template("index.html", current_config=current_config)
 33
 34
 35@web.route("/outputs")
 36def outputs():
 37    """
 38    Returns the contents of the "/data/outputs" folder
 39    """
 40
 41    files = os.listdir("/data/outputs")
 42
 43    # sort files by date created
 44    files.sort(key=lambda x: -os.path.getmtime("/data/outputs/" + x))
 45
 46    # remove 'steps' from the list of directories
 47    if "steps" in files:
 48        files.remove("steps")
 49
 50    # create a list of links to the outputs
 51    links = []
 52    for file in files[:25]:
 53        date = os.path.getmtime("/data/outputs/" + file)
 54        date = datetime.datetime.fromtimestamp(date).strftime("%Y-%m-%d %H:%M:%S")
 55        id = file.split("--")[0]
 56        title = " ".join(file.split("--")[1].split(".")[0].split("-"))
 57        width, height = Image.open("/data/outputs/" + file).size
 58        links.append(
 59            {
 60                "file": "/outputs/" + file,
 61                "date": date,
 62                "id": id,
 63                "title": title,
 64                "width": width,
 65                "height": height,
 66            }
 67        )
 68
 69    return render_template("outputs.html", links=links)
 70
 71
 72@web.route("/outputs/<path:path>")
 73def outputs_path(path):
 74    """
 75    Returns the image at the given path.
 76    """
 77    return send_from_directory("/data/outputs", path)
 78
 79
 80@web.route("/generate")
 81def generate():
 82    """
 83    Runs the script "scripts/generate.py" and writes the output to /data/outputs
 84    as it is generate with threading
 85    """
 86
 87    command = [
 88        "python",
 89        "-m",
 90        "scripts.generate",
 91        "-c",
 92        "/data/config.json",
 93    ]
 94
 95    p = subprocess.Popen(command, stdout=subprocess.PIPE)
 96    output = p.communicate()[0]
 97
 98    return output.decode("utf-8")
 99
100
101@web.route("/update_config", methods=["POST"])
102def update_config():
103    """
104    Updates the config.json file with the given values. The values that can be
105    updated and are in the POST request are:
106
107    prompts: A string of comma separated prompts to use, turn into a list of strings
108    init_image: An image that we need to download and set this path to.
109    width: The width of the image. (is the first dimension of size list)
110    height: The height of the image. (is the second dimension of size list)
111    max_iterations: The number of iterations to run.
112    """
113
114    old_config = json.loads(open("/data/config.json", "r").read())
115
116    # get the values from the POST request
117    prompts = request.form.get("prompts")
118    width = request.form.get("width")
119    height = request.form.get("height")
120    max_iterations = request.form.get("max_iterations")
121
122    # get the init_image from the POST request if we have one
123    init_image = request.files.get("init_image")
124    if init_image.filename != "":
125        # make "/data/init_images" if it doesn't exist
126        if not os.path.exists("/data/init_images"):
127            os.mkdir("/data/init_images")
128        init_image.save("/data/init_images/" + init_image.filename)
129        init_image = "/data/init_images/" + init_image.filename
130    elif request.form.get("clear_init_image") == "true":
131        init_image = ""
132    else:
133        init_image = old_config["init_image"]
134
135    # create a new config object
136    config = {
137        "prompts": [x.strip() for x in prompts.split(",")],
138        "init_image": init_image,
139        "size": [int(width), int(height)],
140        "max_iterations": int(max_iterations),
141    }
142
143    # write the config to the config file
144    with open("/data/config.json", "w") as f:
145        json.dump(config, f)
146
147    return redirect("/")
148
149
150@web.route("/latest_output")
151def latest_output():
152    """
153    Returns the latest output file.
154    """
155    if os.path.exists("/data/outputs/steps"):
156        files = os.listdir("/data/outputs/steps")
157        files.sort(key=lambda x: -os.path.getmtime("/data/outputs/steps/" + x))
158        try:
159            return send_file("/data/outputs/steps/" + files[0], mimetype="image/png")
160        except IndexError:
161            return "No output found"
162    else:
163        return "No output found"
164
165
166app = Flask(__name__)
167app.register_blueprint(web)
168
169check_files_and_folders()
170
171print("-----------------------------------------------------")
172print("")
173print("AI-Art by Isaac Bythewood")
174print("https://github.com/overshard/ai-art")
175print("")
176print("Open the following in you browser to see the UI:")
177print("")
178print("    http://localhost:3000/")
179print("")
180print("-----------------------------------------------------")
181
182if __name__ == "__main__":
183    app.run(host="0.0.0.0", port=3000, debug=True)