334 lines
13 KiB
Python
334 lines
13 KiB
Python
import sys
|
|
import argparse
|
|
import logging
|
|
import time
|
|
import os
|
|
import json
|
|
import requests
|
|
import concurrent.futures
|
|
from core.epub_reader import read_epub, read_epub_metadata
|
|
from core.scene_splitter import split_scenes
|
|
from manga.panel_layout import group_into_pages
|
|
from manga.prompt_builder import build_page_prompt, get_cfg
|
|
from manga.speech_bubbles import add_speech_bubbles
|
|
from ai.scene_parser import parse_scene
|
|
from ai.character_memory import update as update_character, load as load_character_db, dump as dump_character_db, add_hint
|
|
from image.backend import generate_image
|
|
from image.sd_launcher import ensure_running, shutdown
|
|
import atexit
|
|
from export.epub_builder import build_epub
|
|
from utils.naming import make_output_name
|
|
import config
|
|
|
|
logging.basicConfig(level=logging.WARNING, format="%(levelname)s %(name)s: %(message)s")
|
|
log = logging.getLogger(__name__)
|
|
|
|
def fmt_time(seconds):
|
|
seconds = int(seconds)
|
|
if seconds < 60:
|
|
return f"{seconds}s"
|
|
m, s = divmod(seconds, 60)
|
|
if m < 60:
|
|
return f"{m}m {s:02d}s"
|
|
h, m = divmod(m, 60)
|
|
return f"{h}h {m:02d}m"
|
|
|
|
def progress_bar(current, total, label="", eta_str="", width=35):
|
|
pct = current / total if total else 0
|
|
filled = int(width * pct)
|
|
bar = "█" * filled + "░" * (width - filled)
|
|
eta = f" ETA {eta_str}" if eta_str else ""
|
|
print(f"\r [{bar}] {current}/{total} {label}{eta} ", end="", flush=True)
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(
|
|
prog="main.py",
|
|
description="Convert an EPUB novel into a manga-style EPUB using LLM scene parsing and Stable Diffusion.",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
examples:
|
|
python3 main.py book.epub
|
|
python3 main.py book.epub output.epub
|
|
python3 main.py book.epub --layout tiny --model mistral
|
|
python3 main.py book.epub --sd-url http://192.168.1.10:7860 --steps 30 --size 512x768
|
|
python3 main.py book.epub --dry-run
|
|
|
|
layout modes:
|
|
normal 3 scenes per page (default)
|
|
tiny 1 scene per page (more pages, more detail per image)
|
|
|
|
requirements:
|
|
Ollama running at OLLAMA_URL (default: http://localhost:11434)
|
|
Stable Diffusion WebUI running at SD_API_URL (default: http://127.0.0.1:7860)
|
|
Both URLs can be overridden via flags or by editing config.py
|
|
"""
|
|
)
|
|
|
|
parser.add_argument("input", help="Path to input .epub file")
|
|
parser.add_argument("output", nargs="?", help="Path for output .epub (default: manga-<input>.epub)")
|
|
parser.add_argument("--layout", choices=["normal", "tiny"], default="normal",
|
|
help="Page layout mode: normal=3 scenes/page, tiny=1 scene/page (default: normal)")
|
|
parser.add_argument("--model", default=config.OLLAMA_MODEL, metavar="MODEL",
|
|
help=f"Ollama model to use for scene parsing (default: {config.OLLAMA_MODEL})")
|
|
parser.add_argument("--ollama-url", default=config.OLLAMA_URL, metavar="URL",
|
|
help=f"Ollama API base URL (default: {config.OLLAMA_URL})")
|
|
parser.add_argument("--sd-url", default=config.SD_API_URL, metavar="URL",
|
|
help=f"Stable Diffusion WebUI API URL (default: {config.SD_API_URL})")
|
|
parser.add_argument("--sd-path", default=None, metavar="PATH",
|
|
help="Path to stable-diffusion-webui folder; if SD is not running, auto-starts it")
|
|
parser.add_argument("--steps", type=int, default=20, metavar="N",
|
|
help="Diffusion steps per image (default: 20, higher=better quality but slower)")
|
|
parser.add_argument("--style", choices=["lineart", "manga"], default="lineart",
|
|
help="Image style: lineart=simple clean outlines (default), manga=detailed screentone")
|
|
parser.add_argument("--size", default="768x1024", metavar="WxH",
|
|
help="Image dimensions in pixels (default: 768x1024)")
|
|
parser.add_argument("--workers", type=int, default=2, metavar="N",
|
|
help="Parallel image generation workers (default: 2)")
|
|
parser.add_argument("--hint", action="append", metavar="NAME:DESC",
|
|
help="Character description hint e.g. 'Rocky:alien who resembles a rock spider'. Can be used multiple times.")
|
|
parser.add_argument("--jpeg-quality", type=int, default=85, metavar="N",
|
|
help="JPEG quality for epub images 1-95 (default: 85, lower=smaller file)")
|
|
parser.add_argument("--dry-run", nargs="?", const=3, type=int, metavar="N",
|
|
help="Show N prompts without generating images (default: 3)")
|
|
parser.add_argument("--verbose", "-v", action="store_true",
|
|
help="Show debug logging")
|
|
|
|
return parser.parse_args()
|
|
|
|
def main():
|
|
args = parse_args()
|
|
|
|
if args.verbose:
|
|
logging.getLogger().setLevel(logging.DEBUG)
|
|
|
|
config.OLLAMA_MODEL = args.model
|
|
config.OLLAMA_URL = args.ollama_url
|
|
from image.sd_launcher import write_config as _wc
|
|
_wc("llm-model", args.model)
|
|
config.SD_API_URL = args.sd_url
|
|
|
|
atexit.register(shutdown)
|
|
ensure_running(args.sd_url, args.sd_path, style=args.style)
|
|
|
|
try:
|
|
width, height = (int(x) for x in args.size.split("x"))
|
|
except ValueError:
|
|
print(f"Error: --size must be in WxH format, e.g. 768x1024")
|
|
sys.exit(1)
|
|
|
|
if not os.path.isfile(args.input):
|
|
print(f"Error: input file not found: {args.input}")
|
|
sys.exit(1)
|
|
|
|
output = args.output or make_output_name(args.input)
|
|
output_dir = os.path.splitext(output)[0] + "_pages"
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
total_start = time.time()
|
|
|
|
print(f"\nepub-to-manga")
|
|
print(f" input: {args.input}")
|
|
print(f" output: {output}")
|
|
print(f" layout: {args.layout} | model: {args.model} | steps: {args.steps} | size: {width}x{height} | style: {args.style} | workers: {args.workers}")
|
|
|
|
cache_file = make_output_name(args.input).replace(".epub", ".cache.json")
|
|
|
|
print(f"\nReading & splitting epub...")
|
|
source_title, source_author = read_epub_metadata(args.input)
|
|
chapters = read_epub(args.input)
|
|
scenes = split_scenes(chapters)
|
|
print(f" {len(scenes)} scenes found across {len(chapters)} chapters")
|
|
|
|
dry_run_n = args.dry_run if args.dry_run is not None else None
|
|
needed = dry_run_n if dry_run_n is not None else len(scenes)
|
|
|
|
cached_scenes = []
|
|
cached_chars = {}
|
|
if os.path.exists(cache_file):
|
|
with open(cache_file) as f:
|
|
cache_data = json.load(f)
|
|
cached_scenes = cache_data.get("parsed_scenes", [])
|
|
cached_chars = cache_data.get("character_db", {})
|
|
|
|
load_character_db(cached_chars)
|
|
|
|
if args.hint:
|
|
for hint in args.hint:
|
|
if ":" in hint:
|
|
name, _, desc = hint.partition(":")
|
|
add_hint(name.strip(), desc.strip())
|
|
print(f" Character hint: {name.strip()} = {desc.strip()}")
|
|
|
|
if len(cached_scenes) >= needed:
|
|
print(f" Loaded {len(cached_scenes)} scenes from cache ({cache_file})")
|
|
parsed_scenes = cached_scenes
|
|
else:
|
|
if cached_scenes:
|
|
print(f" Resuming parse — {len(cached_scenes)}/{needed} scenes cached")
|
|
|
|
remaining_scenes = scenes[len(cached_scenes):needed]
|
|
|
|
print(f"\nParsing scenes with {args.model} ({needed} total)")
|
|
progress_bar(len(cached_scenes), needed)
|
|
run_start = time.time()
|
|
|
|
new_parsed = []
|
|
scene_times = []
|
|
|
|
def on_scene_done(batch_idx, batch_len):
|
|
pass
|
|
|
|
for idx, scene in enumerate(remaining_scenes):
|
|
t0 = time.time()
|
|
result = parse_scene(scene, timeout=90)
|
|
elapsed_scene = time.time() - t0
|
|
new_parsed.append(result)
|
|
scene_times.append(elapsed_scene)
|
|
if len(scene_times) > 10:
|
|
scene_times.pop(0)
|
|
|
|
for char in result.get("characters", []):
|
|
if char and char.strip():
|
|
update_character(char.strip(), "")
|
|
|
|
combined = cached_scenes + new_parsed
|
|
with open(cache_file, "w") as f:
|
|
json.dump({"parsed_scenes": combined, "character_db": dump_character_db()}, f)
|
|
|
|
done = len(cached_scenes) + len(new_parsed)
|
|
avg = sum(scene_times) / len(scene_times)
|
|
remaining_count = needed - done
|
|
eta = fmt_time(avg * remaining_count) if remaining_count > 0 else fmt_time(time.time() - run_start)
|
|
progress_bar(done, needed, eta_str=eta)
|
|
|
|
print()
|
|
parsed_scenes = cached_scenes + new_parsed
|
|
print(f" Saved to {cache_file}")
|
|
|
|
for ps in parsed_scenes:
|
|
for char in ps.get("characters", []):
|
|
if char and char.strip():
|
|
update_character(char.strip(), "")
|
|
|
|
pages = group_into_pages(parsed_scenes, mode=args.layout)
|
|
print(f" {len(pages)} pages ({args.layout} layout)")
|
|
|
|
if args.dry_run is not None:
|
|
n = args.dry_run
|
|
print(f"\nDry run — first {n} prompts:")
|
|
for i, page in enumerate(pages[:n]):
|
|
positive, negative = build_page_prompt(page, style=args.style)
|
|
print(f"\n--- Page {i+1} ---\nPositive: {positive}\nNegative: {negative}")
|
|
return
|
|
|
|
|
|
from image.sd_launcher import read_config as _rc
|
|
_sd_cfg = _rc()
|
|
_sd_path = _sd_cfg.get("sd-path", "")
|
|
_backend = _sd_cfg.get("sd-backend", "comfyui")
|
|
_port = "8188" if _backend == "comfyui" else "7860"
|
|
if _backend == "comfyui":
|
|
from image import comfy_api
|
|
_model = comfy_api.ensure_model_for_style(_sd_path, style=args.style)
|
|
from image.sd_launcher import write_config as _wc2
|
|
_wc2("comfy-model", _model)
|
|
|
|
total = len(pages)
|
|
already_done = [
|
|
os.path.join(output_dir, f"page_{i}.png")
|
|
for i in range(total)
|
|
if os.path.exists(os.path.join(output_dir, f"page_{i}.png"))
|
|
]
|
|
resumed = len(already_done)
|
|
|
|
print(f"\nGenerating images... (http://127.0.0.1:{_port})")
|
|
print(f" {total} pages | ~{args.steps} steps each | {width}x{height} | {args.workers} workers")
|
|
if resumed:
|
|
print(f" Resuming — {resumed}/{total} pages already done, skipping...")
|
|
|
|
images = [None] * total
|
|
for i in range(total):
|
|
img_path = os.path.join(output_dir, f"page_{i}.png")
|
|
if os.path.exists(img_path):
|
|
images[i] = img_path
|
|
|
|
progress_bar(resumed, total)
|
|
img_start = time.time()
|
|
times = []
|
|
completed = resumed
|
|
lock = __import__("threading").Lock()
|
|
|
|
def gen_page(args_tuple):
|
|
i, page = args_tuple
|
|
img_path = os.path.join(output_dir, f"page_{i}.png")
|
|
if os.path.exists(img_path):
|
|
return i, img_path, None
|
|
|
|
positive, negative = build_page_prompt(page, style=args.style)
|
|
t0 = time.time()
|
|
try:
|
|
img_bytes = generate_image(
|
|
positive, negative,
|
|
steps=args.steps, width=width, height=height,
|
|
cfg=get_cfg(args.style),
|
|
force_bw=(args.style == "lineart"),
|
|
style=args.style
|
|
)
|
|
elapsed = time.time() - t0
|
|
|
|
with open(img_path, "wb") as f:
|
|
f.write(img_bytes)
|
|
|
|
dialogue = []
|
|
if isinstance(page, list):
|
|
for scene in page:
|
|
if isinstance(scene, dict):
|
|
dialogue.extend(scene.get("dialogue", []))
|
|
if dialogue:
|
|
add_speech_bubbles(img_path, dialogue)
|
|
|
|
return i, img_path, elapsed
|
|
except (requests.RequestException, KeyError, IndexError, RuntimeError) as e:
|
|
return i, None, str(e)
|
|
|
|
pending = [(i, page) for i, page in enumerate(pages) if images[i] is None]
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as executor:
|
|
futures = {executor.submit(gen_page, item): item[0] for item in pending}
|
|
for future in concurrent.futures.as_completed(futures):
|
|
i, img_path, result = future.result()
|
|
with lock:
|
|
if img_path:
|
|
images[i] = img_path
|
|
if isinstance(result, float):
|
|
times.append(result)
|
|
if len(times) > 8:
|
|
times.pop(0)
|
|
else:
|
|
print(f"\n [!] Page {i} failed: {result}")
|
|
completed += 1
|
|
avg = sum(times) / len(times) if times else 0
|
|
remaining_count = total - completed
|
|
eta = fmt_time(avg * remaining_count / args.workers) if avg and remaining_count else ""
|
|
progress_bar(completed, total, eta_str=eta)
|
|
|
|
print()
|
|
|
|
final_images = [img for img in images if img]
|
|
|
|
if not final_images:
|
|
print(f"\nError: no images generated. Is the SD API running at {args.sd_url}?")
|
|
sys.exit(1)
|
|
|
|
print(f"\nBuilding epub...")
|
|
result = build_epub(final_images, output, title=source_title or "Manga Book", author="Manga", jpeg_quality=args.jpeg_quality)
|
|
|
|
total_elapsed = time.time() - total_start
|
|
print(f" Done in {fmt_time(total_elapsed)}: {result}\n")
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except KeyboardInterrupt:
|
|
print("\n\nStopped. Progress saved - re-run to resume.")
|