Add Flask app for managing AI stacks

This commit is contained in:
Outis
2026-06-01 09:32:30 -04:00
committed by GitHub
parent bb38aefb9d
commit 112be30317
+117
View File
@@ -0,0 +1,117 @@
import os
import subprocess
import threading
from flask import Flask, render_template, redirect, jsonify
app = Flask(__name__)
IMAGE_STACK = os.environ.get("IMAGE_STACK", "/docker/ai-image-gen")
LLM_STACK = os.environ.get("LLM_STACK", "/docker/ai-llm")
# Track current active stack and swap status
state = {
"active": None, # "images" | "chat" | None
"swapping": False,
"message": "No stack active"
}
def run_compose(stack_dir, command):
"""Run a docker compose command in a stack directory."""
subprocess.run(
["docker", "compose", "-f", f"{stack_dir}/docker-compose.yml"] + command,
capture_output=True
)
def is_stack_running(stack_dir):
"""Check if the GPU container for this stack is running."""
container = "ai-image-gen-invokeai-1" if "image" in stack_dir else "ai-llm-ollama-1"
result = subprocess.run(
["docker", "inspect", "--format", "{{.State.Running}}", container],
capture_output=True,
text=True
)
return result.stdout.strip() == "true"
def do_swap(target):
"""Stop the GPU-heavy container of inactive stack, start target."""
state["swapping"] = True
if target == "images":
state["message"] = "Stopping Ollama to free VRAM..."
subprocess.run(["docker", "stop", "ai-llm-ollama-1"], capture_output=True)
state["message"] = "Starting InvokeAI..."
subprocess.run(
["docker", "compose", "-f", f"{IMAGE_STACK}/docker-compose.yml", "up", "-d"],
capture_output=True
)
else:
state["message"] = "Stopping InvokeAI to free VRAM..."
subprocess.run(["docker", "stop", "ai-image-gen-invokeai-1"], capture_output=True)
state["message"] = "Starting Ollama..."
subprocess.run(["docker", "start", "ai-llm-ollama-1"], capture_output=True)
state["active"] = target
state["swapping"] = False
state["message"] = f"{target.capitalize()} is ready."
@app.route("/")
@app.route("/portal")
def portal():
image_running = is_stack_running(IMAGE_STACK)
llm_running = is_stack_running(LLM_STACK)
return render_template(
"portal.html",
active=state["active"],
swapping=state["swapping"],
message=state["message"],
image_running=image_running,
llm_running=llm_running
)
@app.route("/switch/<target>")
def switch(target):
if target not in ("images", "chat"):
return "Invalid target", 400
if state["swapping"]:
return redirect("/portal")
# Warn if same stack requested
if state["active"] == target:
urls = {"chat": "https://chat.mydomain.com", "images": "https://images.mydomain.com"}
return redirect(urls[target])
# Run swap in background so browser gets immediate response
thread = threading.Thread(target=do_swap, args=(target,))
thread.daemon = True
thread.start()
return render_template(
"switching.html",
target=target,
message=f"Switching to {target}... this takes about 10 seconds."
)
@app.route("/chat")
def chat():
return redirect(URLS["chat"])
@app.route("/images")
def images():
return redirect(URLS["images"])
@app.route("/status")
def status():
"""Polled by the switching page to know when swap is done."""
return jsonify(state)
if __name__ == "__main__":
# Detect what's already running on startup
if is_stack_running(IMAGE_STACK):
state["active"] = "images"
state["message"] = "Images stack is active."
elif is_stack_running(LLM_STACK):
state["active"] = "chat"
state["message"] = "Chat stack is active."
app.run(host="0.0.0.0", port=8080)