338 lines
11 KiB
Python
338 lines
11 KiB
Python
|
|
import os
|
|
import subprocess
|
|
import time
|
|
import threading
|
|
import sys
|
|
import requests
|
|
|
|
_sd_process = None
|
|
|
|
CONFIG_FILE = os.path.expanduser("~/.epub-to-manga")
|
|
|
|
BACKENDS = [
|
|
{
|
|
"name": "AUTOMATIC1111 Stable Diffusion WebUI",
|
|
"repo": "https://github.com/AUTOMATIC1111/stable-diffusion-webui",
|
|
"dir": "stable-diffusion-webui",
|
|
"type": "a1111",
|
|
"launch": ["bash", "webui.sh", "--api", "--nowebui"],
|
|
"ready_url": "http://127.0.0.1:7860/sdapi/v1/options",
|
|
},
|
|
{
|
|
"name": "ComfyUI (lighter, faster on Apple Silicon)",
|
|
"repo": "https://github.com/comfyanonymous/ComfyUI",
|
|
"dir": "ComfyUI",
|
|
"type": "comfyui",
|
|
"launch": [sys.executable, "main.py", "--listen"],
|
|
"ready_url": "http://127.0.0.1:8188/system_stats",
|
|
},
|
|
{
|
|
"name": "InvokeAI",
|
|
"repo": "https://github.com/invoke-ai/InvokeAI",
|
|
"dir": "InvokeAI",
|
|
"type": "invokeai",
|
|
"launch": ["invokeai-web"],
|
|
"ready_url": "http://127.0.0.1:9090/api/v1/app/version",
|
|
},
|
|
]
|
|
|
|
KNOWN_ERRORS = [
|
|
{
|
|
"markers": ["No module named 'pkg_resources'", "Couldn't install clip"],
|
|
"message": "AUTOMATIC1111 is not compatible with Python {python_version}. A1111 requires Python 3.10.",
|
|
"solutions": [
|
|
{
|
|
"label": "Install Python 3.10 via pyenv and relaunch A1111",
|
|
"action": "fix_a1111_python",
|
|
},
|
|
{
|
|
"label": "Switch to ComfyUI (better Apple Silicon support)",
|
|
"action": "switch_backend",
|
|
"backend_index": 1,
|
|
},
|
|
],
|
|
},
|
|
{
|
|
"markers": ["CUDA out of memory", "out of memory"],
|
|
"message": "GPU ran out of memory.",
|
|
"solutions": [
|
|
{
|
|
"label": "Re-run with a smaller image size (--size 256x384)",
|
|
"action": "suggest_flag",
|
|
"flag": "--size 256x384",
|
|
},
|
|
],
|
|
},
|
|
]
|
|
|
|
def read_config():
|
|
config = {}
|
|
if os.path.exists(CONFIG_FILE):
|
|
with open(CONFIG_FILE) as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if "=" in line and not line.startswith("#"):
|
|
key, _, val = line.partition("=")
|
|
config[key.strip()] = val.strip()
|
|
return config
|
|
|
|
def write_config(key, value):
|
|
config = read_config()
|
|
config[key] = value
|
|
with open(CONFIG_FILE, "w") as f:
|
|
for k, v in config.items():
|
|
f.write(f"{k}={v}\n")
|
|
|
|
def get_backend_for_path(sd_path):
|
|
if not sd_path:
|
|
return BACKENDS[0]
|
|
for b in BACKENDS:
|
|
if b["dir"] in sd_path or os.path.exists(os.path.join(sd_path, b["dir"])):
|
|
return b
|
|
if b["type"] == "a1111" and os.path.exists(os.path.join(sd_path, "webui.sh")):
|
|
return b
|
|
if b["type"] == "comfyui" and os.path.exists(os.path.join(sd_path, "comfy_extras")):
|
|
return b
|
|
return BACKENDS[0]
|
|
|
|
def is_running(ready_url, timeout=3):
|
|
try:
|
|
requests.get(ready_url, timeout=timeout)
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
def find_installed():
|
|
home = os.path.expanduser("~")
|
|
search_dirs = [home, os.path.join(home, "Downloads"), os.path.join(home, "Documents"), os.getcwd()]
|
|
for backend in BACKENDS:
|
|
for base in search_dirs:
|
|
candidate = os.path.join(base, backend["dir"])
|
|
if os.path.isdir(candidate):
|
|
return candidate, backend
|
|
return None, None
|
|
|
|
def get_python_version():
|
|
return f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
|
|
|
|
def handle_known_error(error_def, sd_path, backend):
|
|
python_version = get_python_version()
|
|
msg = error_def["message"].format(python_version=python_version)
|
|
solutions = error_def["solutions"]
|
|
|
|
print(f"\n\n ⚠️ {msg}\n")
|
|
print(" Solutions:\n")
|
|
for i, s in enumerate(solutions, 1):
|
|
print(f" {i}) {s['label']}")
|
|
print()
|
|
|
|
while True:
|
|
choice = input(" Choose a solution (or q to quit): ").strip().lower()
|
|
if choice == "q":
|
|
sys.exit(0)
|
|
if choice.isdigit() and 1 <= int(choice) <= len(solutions):
|
|
solution = solutions[int(choice) - 1]
|
|
break
|
|
print(f" Enter a number between 1 and {len(solutions)}")
|
|
|
|
action = solution["action"]
|
|
|
|
if action == "fix_a1111_python":
|
|
print("\n Installing pyenv and Python 3.10.14...")
|
|
subprocess.run(["brew", "install", "pyenv"], check=False)
|
|
subprocess.run(["pyenv", "install", "3.10.14"], check=False)
|
|
pyenv_python = os.path.expanduser("~/.pyenv/versions/3.10.14/bin/python3")
|
|
if not os.path.exists(pyenv_python):
|
|
print("\n Error: pyenv install failed. See https://github.com/pyenv/pyenv")
|
|
sys.exit(1)
|
|
env = os.environ.copy()
|
|
env["PYTHON"] = pyenv_python
|
|
print(f"\n Relaunching A1111 with Python 3.10.14...")
|
|
global _sd_process
|
|
_sd_process = subprocess.Popen(
|
|
["bash", "webui.sh", "--api", "--nowebui"],
|
|
cwd=sd_path,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
env=env,
|
|
)
|
|
stream_and_wait(backend, sd_path)
|
|
|
|
elif action == "switch_backend":
|
|
new_backend = BACKENDS[solution["backend_index"]]
|
|
dest = os.path.join(os.path.expanduser("~"), new_backend["dir"])
|
|
if os.path.isdir(dest):
|
|
print(f"\n Found existing {new_backend['name']} at {dest} — using it.")
|
|
else:
|
|
confirm = input(f"\n Install {new_backend['name']} to {dest}? [Y/n] ").strip().lower()
|
|
if confirm in ("n", "no"):
|
|
print(" Aborted.")
|
|
sys.exit(0)
|
|
print(f"\n Cloning {new_backend['repo']}...")
|
|
result = subprocess.run(["git", "clone", "--recursive", new_backend["repo"], dest], check=False)
|
|
if result.returncode != 0:
|
|
print("\n Error: git clone failed.")
|
|
sys.exit(1)
|
|
print(f"\n Installed to {dest}")
|
|
write_config("sd-path", dest)
|
|
write_config("sd-backend", new_backend["type"])
|
|
print(f" Re-run your original command — switching will take effect automatically.")
|
|
sys.exit(0)
|
|
|
|
elif action == "suggest_flag":
|
|
print(f"\n Re-run with: {solution['flag']}")
|
|
sys.exit(0)
|
|
|
|
def stream_logs(proc, ready_event, error_detected, sd_path, backend):
|
|
interesting = [
|
|
"Loading", "Downloading", "Installing", "Running", "Creating",
|
|
"Model loaded", "Starting", "Applying", "torch", "CUDA", "MPS",
|
|
"checkpoint", "venv", "pip", "Startup", "listen",
|
|
]
|
|
|
|
for raw in proc.stdout:
|
|
if ready_event.is_set():
|
|
break
|
|
line = raw.decode("utf-8", errors="replace").rstrip()
|
|
|
|
for err_def in KNOWN_ERRORS:
|
|
if any(m in line for m in err_def["markers"]):
|
|
ready_event.set()
|
|
error_detected["def"] = err_def
|
|
return
|
|
|
|
if any(kw in line for kw in interesting):
|
|
print(f"\r > {line[:100]:<100}")
|
|
print(" ", end="", flush=True)
|
|
|
|
def stream_and_wait(backend, sd_path):
|
|
global _sd_process
|
|
ready_event = threading.Event()
|
|
error_detected = {}
|
|
|
|
log_thread = threading.Thread(
|
|
target=stream_logs,
|
|
args=(_sd_process, ready_event, error_detected, sd_path, backend),
|
|
daemon=True,
|
|
)
|
|
log_thread.start()
|
|
|
|
start = time.time()
|
|
while not ready_event.is_set():
|
|
elapsed = int(time.time() - start)
|
|
mins, secs = divmod(elapsed, 60)
|
|
elapsed_str = f"{mins}m {secs:02d}s" if mins else f"{secs}s"
|
|
print(f"\r Waiting for API... {elapsed_str} elapsed ", end="", flush=True)
|
|
|
|
if is_running(backend["ready_url"]):
|
|
ready_event.set()
|
|
elapsed = int(time.time() - start)
|
|
mins, secs = divmod(elapsed, 60)
|
|
elapsed_str = f"{mins}m {secs:02d}s" if mins else f"{secs}s"
|
|
print(f"\r Ready in {elapsed_str}! \n")
|
|
return
|
|
|
|
if _sd_process.poll() is not None and not ready_event.is_set():
|
|
ready_event.set()
|
|
break
|
|
|
|
time.sleep(2)
|
|
|
|
log_thread.join(timeout=2)
|
|
|
|
if "def" in error_detected:
|
|
handle_known_error(error_detected["def"], sd_path, backend)
|
|
elif not is_running(backend["ready_url"]):
|
|
print("\n\n Error: SD process exited unexpectedly.")
|
|
sys.exit(1)
|
|
|
|
def launch(sd_path, backend):
|
|
global _sd_process
|
|
print(f"\nStarting {backend['name']} from {sd_path}...")
|
|
|
|
_sd_process = subprocess.Popen(
|
|
backend["launch"],
|
|
cwd=sd_path,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
)
|
|
stream_and_wait(backend, sd_path)
|
|
|
|
def prompt_install():
|
|
print("\nNo Stable Diffusion backend found on your system.")
|
|
print("\nAvailable backends to install:\n")
|
|
for i, b in enumerate(BACKENDS, 1):
|
|
print(f" {i}) {b['name']}")
|
|
print(f" {b['repo']}")
|
|
print()
|
|
|
|
while True:
|
|
choice = input("Enter number to install (or q to quit): ").strip().lower()
|
|
if choice == "q":
|
|
sys.exit(0)
|
|
if choice.isdigit() and 1 <= int(choice) <= len(BACKENDS):
|
|
backend = BACKENDS[int(choice) - 1]
|
|
break
|
|
print(f" Please enter a number between 1 and {len(BACKENDS)}")
|
|
|
|
dest = os.path.join(os.path.expanduser("~"), backend["dir"])
|
|
confirm = input(f"\nInstall {backend['name']} to {dest}? [Y/n] ").strip().lower()
|
|
if confirm in ("n", "no"):
|
|
print("Aborted.")
|
|
sys.exit(0)
|
|
|
|
if os.path.isdir(dest):
|
|
print(f"\nFound existing {backend['name']} at {dest} — using it.")
|
|
else:
|
|
print(f"\nCloning {backend['repo']}...")
|
|
result = subprocess.run(["git", "clone", "--recursive", backend["repo"], dest], check=False)
|
|
if result.returncode != 0:
|
|
print("\nError: git clone failed. Is git installed?")
|
|
sys.exit(1)
|
|
print(f"\nInstalled to {dest}")
|
|
|
|
write_config("sd-path", dest)
|
|
write_config("sd-backend", backend["type"])
|
|
print(f"Saved to {CONFIG_FILE}")
|
|
print(f"\nRe-run your original command — no --sd-path needed.")
|
|
sys.exit(0)
|
|
|
|
def ensure_running(sd_url=None, sd_path=None, style="manga"):
|
|
config = read_config()
|
|
|
|
if not sd_path:
|
|
sd_path = config.get("sd-path")
|
|
|
|
backend = get_backend_for_path(sd_path)
|
|
|
|
if is_running(backend["ready_url"]):
|
|
return
|
|
|
|
if sd_path and os.path.isdir(sd_path):
|
|
from image import backend as backend_mod
|
|
backend_mod.setup(sd_path)
|
|
backend_mod.ensure_model(backend["type"], sd_path, style=style)
|
|
launch(sd_path, backend)
|
|
return
|
|
|
|
installed_path, found_backend = find_installed()
|
|
if installed_path:
|
|
print(f"\nFound {found_backend['name']} at {installed_path}")
|
|
write_config("sd-path", installed_path)
|
|
write_config("sd-backend", found_backend["type"])
|
|
print(f"Saved to {CONFIG_FILE}")
|
|
from image import backend as backend_mod
|
|
backend_mod.setup(installed_path)
|
|
backend_mod.ensure_model(found_backend["type"], installed_path, style=style)
|
|
launch(installed_path, found_backend)
|
|
return
|
|
|
|
prompt_install()
|
|
|
|
def shutdown():
|
|
global _sd_process
|
|
if _sd_process:
|
|
_sd_process.terminate()
|
|
_sd_process = None
|