|
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() |
|
|
|
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() |
|
|
|
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) |
|
|
