app.py

import os
from pathlib import Path
import json
from flask import Flask, request, jsonify, g
from flask_expects_json import expects_json
from flask_cors import CORS
from PIL import Image
from huggingface_hub import Repository
from flask_apscheduler import APScheduler
import shutil
import sqlite3
import subprocess
from jsonschema import ValidationError

MODE = os.environ.get("FLASK_ENV", "production")
IS_DEV = MODE == "development"
app = Flask(__name__, static_url_path="/static")
app.config["JSONIFY_PRETTYPRINT_REGULAR"] = False

schema = {
    "type": "object",
    "properties": {
        "prompt": {"type": "string"},
        "images": {
            "type": "array",
            "items": {
                "type": "object",
                "minProperties": 2,
                "maxProperties": 2,
                "properties": {
                    "colors": {
                        "type": "array",
                        "items": {"type": "string"},
                        "maxItems": 5,
                        "minItems": 5,
                    },
                    "imgURL": {"type": "string"},
                },
            },
        },
    },
    "minProperties": 2,
    "maxProperties": 2,
}

CORS(app)

DB_FILE = Path("./data.db")
TOKEN = os.environ.get("HUGGING_FACE_HUB_TOKEN")
repo = Repository(
    local_dir="data",
    repo_type="dataset",
    clone_from="huggingface-projects/color-palettes-sd",
    use_auth_token=TOKEN,
)
repo.git_pull()
# copy db on db to local path
shutil.copyfile("./data/data.db", DB_FILE)

db = sqlite3.connect(DB_FILE)
try:
    data = db.execute("SELECT * FROM palettes").fetchall()
    if IS_DEV:
        print(f"Loaded {len(data)} palettes from local db")
    db.close()
except sqlite3.OperationalError:
    db.execute(
        "CREATE TABLE palettes (id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, data json, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL)"
    )
    db.commit()


def get_db():
    db = getattr(g, "_database", None)
    if db is None:
        db = g._database = sqlite3.connect(DB_FILE)
        db.row_factory = sqlite3.Row
    return db


@app.teardown_appcontext
def close_connection(exception):
    db = getattr(g, "_database", None)
    if db is not None:
        db.close()


def update_repository():
    repo.git_pull()
    # copy db on db to local path
    shutil.copyfile(DB_FILE, "./data/data.db")

    with sqlite3.connect("./data/data.db") as db:
        db.row_factory = sqlite3.Row
        palettes = db.execute("SELECT * FROM palettes").fetchall()
        data = [
            {
                "id": row["id"],
                "data": json.loads(row["data"]),
                "created_at": row["created_at"],
            }
            for row in palettes
        ]

    with open("./data/data.json", "w") as f:
        json.dump(data, f, separators=(",", ":"))

    print("Updating repository")
    subprocess.Popen(
        "git add . && git commit --amend -m 'update' && git push --force",
        cwd="./data",
        shell=True,
    )
    repo.push_to_hub(blocking=False)


@app.route("/")
def index():
    return app.send_static_file("index.html")


@app.route("/force_push")
def push():
    if request.headers["token"] == TOKEN:
        update_repository()
        return jsonify({"success": True})
    else:
        return "Error", 401


def getAllData():
    palettes = get_db().execute("SELECT * FROM palettes").fetchall()
    data = [
        {
            "id": row["id"],
            "data": json.loads(row["data"]),
            "created_at": row["created_at"],
        }
        for row in palettes
    ]
    return data


@app.route("/data")
def getdata():
    return jsonify(getAllData())


@app.route("/new_palette", methods=["POST"])
@expects_json(schema)
def create():
    data = g.data
    db = get_db()
    cursor = db.cursor()
    cursor.execute("INSERT INTO palettes(data) VALUES (?)", [json.dumps(data)])
    db.commit()
    return jsonify(getAllData())


@app.errorhandler(400)
def bad_request(error):
    if isinstance(error.description, ValidationError):
        original_error = error.description
        return jsonify({"error": original_error.message}), 400
    return error


if __name__ == "__main__":
    if not IS_DEV:
        print("Starting scheduler -- Running Production")
        scheduler = APScheduler()
        scheduler.add_job(
            id="Update Dataset Repository",
            func=update_repository,
            trigger="interval",
            hours=1,
        )
        scheduler.start()
    else:
        print("Not Starting scheduler -- Running Development")
    app.run(
        host="0.0.0.0",
        port=int(os.environ.get("PORT", 7860)),
        debug=True,
        use_reloader=IS_DEV,
    )