diff --git a/.gitignore b/.gitignore
index 6dea479..3b1f4ca 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,3 +9,4 @@ output/
prompts_log.jsonl
publish.sh
test.py
+.vscode/launch.json
diff --git a/ai_frame_image_server.py b/ai_frame_image_server.py
index 7209681..9ff113b 100644
--- a/ai_frame_image_server.py
+++ b/ai_frame_image_server.py
@@ -5,23 +5,34 @@ from flask import (
request,
jsonify,
redirect,
- url_for
+ url_for,
+ session,
+ render_template_string,
)
import os
import time
import threading
from apscheduler.schedulers.background import BackgroundScheduler
-from libs.generic import load_config, load_recent_prompts, get_details_from_png, get_current_version, load_models_from_config
+from libs.generic import (
+ load_config,
+ load_recent_prompts,
+ get_details_from_png,
+ get_current_version,
+ load_models_from_config,
+)
from libs.comfyui import cancel_current_job, create_image, select_model
from libs.ollama import create_prompt_on_openwebui
-#workflow test commit
+# workflow test commit
user_config = load_config()
+
app = Flask(__name__)
+app.secret_key = os.environ.get('SECRET_KEY')
image_folder = "./output"
+
@app.route("/", methods=["GET"])
def index() -> str:
"""
@@ -39,18 +50,34 @@ def index() -> str:
image=image_filename,
prompt=prompt,
reload_interval=user_config["frame"]["reload_interval"],
- version=version
+ version=version,
)
+
+@app.route('/login', methods=['GET', 'POST'])
+def login():
+ if request.method == 'POST':
+ if request.form['password'] == user_config["frame"]["password_for_auth"]:
+ session['authenticated'] = True
+ return render_template("create_image.html", models=load_models_from_config())
+ else:
+ return redirect(url_for('login'))
+ return render_template('login.html')
+
+
@app.route("/images", methods=["GET"])
def gallery() -> str:
images = []
for f in os.listdir(image_folder):
- if f.lower().endswith(('png', 'jpg', 'jpeg', 'gif')):
- images.append({'filename': f})
- images = sorted(images, key=lambda x: os.path.getmtime(os.path.join(image_folder, x['filename'])), reverse=True)
+ if f.lower().endswith(("png", "jpg", "jpeg", "gif")):
+ images.append({"filename": f})
+ images = sorted(
+ images,
+ key=lambda x: os.path.getmtime(os.path.join(image_folder, x["filename"])),
+ reverse=True,
+ )
return render_template("gallery.html", images=images)
-
+
@app.route("/image-details/<filename>", methods=["GET"])
def image_details(filename):
@@ -58,16 +85,12 @@ def image_details(filename):
if not os.path.exists(path):
return {"error": "File not found"}, 404
details = get_details_from_png(path)
- return {
- "prompt": details["p"],
- "model": details["m"],
- "date": details["d"]
- }
-
+ return {"prompt": details["p"], "model": details["m"], "date": details["d"]}
-@app.route('/images/thumbnails/<path:filename>')
+
+@app.route("/images/thumbnails/<path:filename>")
def serve_thumbnail(filename):
- return send_from_directory('output/thumbnails', filename)
+ return send_from_directory("output/thumbnails", filename)
@app.route("/images/<filename>", methods=["GET"])
@@ -105,8 +128,10 @@ def create():
# Start generation in background
threading.Thread(target=lambda: create_image(prompt, model)).start()
-
- return redirect(url_for("image_queued", prompt=prompt, model=model.split(".")[0]))
+
+ return redirect(
+ url_for("image_queued", prompt=prompt, model=model.split(".")[0])
+ )
# For GET requests, just show the form to enter prompt
return render_template("create_image.html", models=load_models_from_config())
@@ -118,23 +143,23 @@ def image_queued():
model = request.args.get("model", "No model selected.").split(".")[0]
return render_template("image_queued.html", prompt=prompt, model=model)
+
def scheduled_task() -> None:
"""Executes the scheduled image generation task."""
print(f"Executing scheduled task at {time.strftime('%Y-%m-%d %H:%M:%S')}")
create_image(None)
+
@app.route("/create_image", methods=["GET"])
def create_image_endpoint() -> str:
"""
Renders the create image template with image and prompt.
"""
-
+ if user_config["frame"]["create_requires_auth"] == "True" and not session.get('authenticated'):
+ return redirect(url_for("login"))
models = load_models_from_config()
- return render_template(
- "create_image.html", models=models
- )
-
+ return render_template("create_image.html", models=models)
if user_config["frame"]["auto_regen"] == "True":
@@ -148,10 +173,9 @@ if user_config["frame"]["auto_regen"] == "True":
minute=regen_time[1],
id="scheduled_task",
max_instances=1, # prevent overlapping
- replace_existing=True # don't double-schedule
+ replace_existing=True, # don't double-schedule
)
scheduler.start()
os.makedirs(image_folder, exist_ok=True)
app.run(host="0.0.0.0", port=user_config["frame"]["port"], debug=True)
-
diff --git a/templates/login.html b/templates/login.html
new file mode 100644
index 0000000..cee2868
--- /dev/null
+++ b/templates/login.html
@@ -0,0 +1,72 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+ <meta charset="UTF-8">
+ <meta name="viewport" content="width=device-width, initial-scale=1.0">
+ <title>Login</title>
+ <style>
+ * {
+ margin: 0;
+ padding: 0;
+ box-sizing: border-box;
+ }
+ body {
+ display: flex;
+ flex-direction: column;
+ align-items: center;
+ justify-content: center;
+ height: 100vh;
+ background: black;
+ color: white;
+ font-family: Arial, sans-serif;
+ padding: 20px;
+ text-align: center;
+ }
+ .message {
+ font-size: 22px;
+ margin-bottom: 20px;
+ }
+ .prompt-text {
+ font-size: 20px;
+ background: #111;
+ padding: 20px;
+ border-radius: 10px;
+ border: 1px solid #333;
+ max-width: 80vw;
+ margin-bottom: 30px;
+ }
+ input[type="password"] {
+ padding: 10px;
+ border-radius: 8px;
+ border: 1px solid #555;
+ background: #222;
+ color: white;
+ font-size: 16px;
+ margin-bottom: 20px;
+ width: 250px;
+ }
+ button {
+ background: #333;
+ color: white;
+ border: none;
+ padding: 10px 20px;
+ border-radius: 8px;
+ font-size: 16px;
+ cursor: pointer;
+ transition: background 0.3s;
+ }
+ button:hover {
+ background: #555;
+ }
+ </style>
+</head>
+<body>
+ <div class="message">Please enter the password to continue:</div>
+ <form method="post">
+ <div class="prompt-text">
+ <input type="password" name="password" placeholder="Password" required>
+ </div>
+ <button type="submit">Login</button>
+ </form>
+</body>
+</html>
diff --git a/user_config.cfg.sample b/user_config.cfg.sample
index c5bd5f2..58ec719 100644
--- a/user_config.cfg.sample
+++ b/user_config.cfg.sample
@@ -3,6 +3,8 @@ reload_interval = 30000
auto_regen = True
regen_time = 03:00
port = 5000
+create_requires_auth = False
+password_for_auth = create
[comfyui]
comfyui_url = http://comfyui