Upload engine.py with huggingface_hub
231dc45 verified 4 months ago
31.9 kB
"""
KDP PDF Generation Engine (V8 — AI Image Pipeline)
==================================================
A minimal pipeline: cover image + theme note -> OpenRouter image gen
-> Pillow cover stamp -> Pillow upscale -> ReportLab PDF.
Public API
----------
``generate_book(cover_bytes, filename, theme_note, page_count, page_size)``
returns ``(ok, output_path)``. Everything else is a private helper.
"""

from __future__ import annotations
import base64
import io
import json
import logging
import os
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
import requests
from PIL import Image, ImageEnhance, ImageFilter, ImageOps
from reportlab.lib.units import inch
from reportlab.lib.utils import ImageReader
from reportlab.pdfgen.canvas import Canvas
from renderer import render_layout_to_pdf, load_layout_dict

# Configuration
# ---------------------------------------------------------------------------

BASE_DIR = Path(__file__).resolve().parent
OUTPUTS_DIR = BASE_DIR / "outputs"
PROMPT_PATH = BASE_DIR / "OPENROUTER_PROMPT.txt"

OPENROUTER_IMAGES_URL = "https://openrouter.ai/api/v1/images"
OPENROUTER_CHAT_URL = "https://openrouter.ai/api/v1/chat/completions"
DEFAULT_MODEL = "google/gemini-3.1-flash-lite-image"
# V9.0: text model used to generate JSON layouts (true vector output).
LAYOUT_MODEL = "google/gemini-3.1-flash-lite"
# Recraft V4.1 returns native SVG (true vector) — best for sharp text & shapes.
RECRAFT_VECTOR_MODEL = "recraft/recraft-v4.1-vector"
VECTOR_MODELS = {RECRAFT_VECTOR_MODEL}
RECRAFT_PROMPT_PATH = BASE_DIR / "RECRAFT_PROMPT.txt"
LAYOUT_PROMPT_PATH = BASE_DIR / "LAYOUT_PROMPT.txt"
DEFAULT_PAGE_SIZE_INCHES = (6.0, 9.0)
TOP_BANNER_INCHES = 1.5   # cover band height
TITLE_BUFFER_INCHES = 0.4  # extra space between cover and design (V8.8: prevent title clipping)
TARGET_DPI = 1800         # V8.8: pushed to max (was 1200)
WORK_WIDTH_PX = 3600      # intermediate working resolution
REQUEST_TIMEOUT_SEC = 90

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
log = logging.getLogger("kdp-engine")

OUTPUTS_DIR.mkdir(parents=True, exist_ok=True)


# API key resolution
# ---------------------------------------------------------------------------

def _get_api_key() -> str:
    """Read OPENROUTER_API_KEY from Streamlit secrets or env. Empty if missing."""
    try:
        import streamlit as st  # type: ignore
        return str(st.secrets.get("OPENROUTER_API_KEY", "") or "")
    except Exception:
        return os.environ.get("OPENROUTER_API_KEY", "") or ""


# Prompt
# ---------------------------------------------------------------------------

def _load_prompt(theme_note: str, model: str = DEFAULT_MODEL) -> str:
    """Read the prompt file for the chosen model and inject the theme note.
    Gemini uses ``OPENROUTER_PROMPT.txt``; Recraft V4.1 vector uses
    ``RECRAFT_PROMPT.txt`` (vector-specific instructions).
    """
    if model in VECTOR_MODELS:
        path = RECRAFT_PROMPT_PATH
        fallback = (
            "A beautiful flat-design 6x9 inch journal page on a white background. "
            "Theme: {THEME_NOTE}. Vector style, sharp lines, no gradients."
        )
    else:
        path = PROMPT_PATH
        fallback = "Design a beautiful 6x9 inch journal interior. Theme: {THEME_NOTE}."
    try:
        template = path.read_text(encoding="utf-8")
    except FileNotFoundError:
        log.warning("Prompt file not found at %s; using fallback.", path)
        template = fallback
    return template.replace("{THEME_NOTE}", (theme_note or "").strip())


# OpenRouter image generation
# ---------------------------------------------------------------------------

def _call_openrouter_images(
    theme_note: str,
    model: str = DEFAULT_MODEL,
) -> Tuple[bytes, str]:
    """POST to OpenRouter /v1/images.
    Returns ``(raw_bytes, media_type)`` where ``media_type`` is one of
    ``"image/png"`` or ``"image/svg+xml"``. Vector models like
    Recraft V4.1 return SVG; raster models return PNG.
    """
    api_key = _get_api_key()
    if not api_key:
        raise RuntimeError("OPENROUTER_API_KEY is not set.")

    prompt = _load_prompt(theme_note, model=model)
    payload = {"model": model, "prompt": prompt}
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    log.info("Calling OpenRouter /v1/images model=%s (prompt %d chars)...", model, len(prompt))
    try:
        resp = requests.post(
            OPENROUTER_IMAGES_URL,
            headers=headers,
            data=json.dumps(payload),
            timeout=REQUEST_TIMEOUT_SEC,
        )
    except requests.Timeout as exc:
        raise RuntimeError(f"OpenRouter request timed out after {REQUEST_TIMEOUT_SEC}s.") from exc
    except requests.RequestException as exc:
        raise RuntimeError(f"OpenRouter request failed: {exc}") from exc

    if resp.status_code >= 400:
        # Try to surface a useful error message.
        try:
            err = resp.json()
        except ValueError:
            err = resp.text[:500]
        raise RuntimeError(
            f"OpenRouter returned HTTP {resp.status_code}: {err}"
        )

    try:
        result = resp.json()
    except ValueError as exc:
        raise RuntimeError(f"OpenRouter returned non-JSON: {resp.text[:500]}") from exc

    data = result.get("data") or []
    if not data or "b64_json" not in data[0]:
        raise RuntimeError(
            f"OpenRouter response missing data[0].b64_json. Keys: {list(result.keys())}"
        )

    raw = base64.b64decode(data[0]["b64_json"])
    media_type = data[0].get("media_type") or "image/png"
    # Some models omit media_type — sniff from the first bytes.
    if media_type == "image/png" and raw.lstrip()[:5] in (b"<?xml", b"<svg ", b"<svg"):
        media_type = "image/svg+xml"
    log.info("OpenRouter returned %d bytes of %s data.", len(raw), media_type)
    return raw, media_type


# V9.0: AI layout generation
# ---------------------------------------------------------------------------

def _call_openrouter_chat(theme_note: str,
                          model: str = LAYOUT_MODEL) -> str:
    """POST to OpenRouter /v1/chat/completions. Returns the assistant text."""
    api_key = _get_api_key()
    if not api_key:
        raise RuntimeError("OPENROUTER_API_KEY is not set.")

    try:
        template = LAYOUT_PROMPT_PATH.read_text(encoding="utf-8")
    except FileNotFoundError:
        raise RuntimeError(f"LAYOUT_PROMPT.txt not found at {LAYOUT_PROMPT_PATH}")
    user_msg = template.replace("{THEME_NOTE}", (theme_note or "").strip())

    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": "You are a JSON layout engine. Output only valid JSON. No markdown, no commentary."},
            {"role": "user", "content": user_msg},
        ],
        "temperature": 0.6,
    }
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    log.info("Calling OpenRouter /v1/chat/completions model=%s (prompt %d chars)...", model, len(user_msg))
    try:
        resp = requests.post(OPENROUTER_CHAT_URL, headers=headers, data=json.dumps(payload), timeout=REQUEST_TIMEOUT_SEC)
    except requests.RequestException as exc:
        raise RuntimeError(f"OpenRouter chat request failed: {exc}") from exc
    if resp.status_code >= 400:
        try:
            err = resp.json()
        except ValueError:
            err = resp.text[:500]
        raise RuntimeError(f"OpenRouter chat returned HTTP {resp.status_code}: {err}")
    result = resp.json()
    choices = result.get("choices") or []
    if not choices:
        raise RuntimeError(f"OpenRouter chat returned no choices: {result}")
    msg = choices[0].get("message") or {}
    text = (msg.get("content") or "").strip()
    if not text:
        raise RuntimeError(f"OpenRouter chat returned empty content: {result}")
    log.info("OpenRouter chat returned %d chars.", len(text))
    return text


def generate_book_layout(cover_bytes: bytes,
                         filename: str,
                         theme_note: str = "",
                         page_count: int = 1,
                         page_size: Tuple[float, float] = DEFAULT_PAGE_SIZE_INCHES,
                         layout_model: str = LAYOUT_MODEL
                         ) -> Tuple[bool, Optional[Path], List[str], Dict[str, Any]]:
    """V9.0 vector pipeline: AI text -> JSON layout -> vector PDF.
    The AI is used only to suggest a layout (titles, table headers,
    theme colors, which elements to use). The actual PDF is drawn
    with ReportLab primitives — true vector, perfect typography,
    ~200 KB output instead of 70 MB.
    Returns ``(ok, output_path, warnings, cover_info)``. ``cover_info``
    is a dict ``{"ok": bool, "reason": str, "size_in": (w, h),
    "crop_method": str}`` describing the actual cover stamp outcome
    so the UI can display a one-line status caption.
    """
    warnings: List[str] = []
    cover_info: Dict[str, Any] = {
        "ok": False, "reason": "not attempted",
        "size_in": (0.0, 0.0), "crop_method": "none",
    }
    if not filename.endswith(".pdf"):
        filename += ".pdf"
    out_path = OUTPUTS_DIR / filename

    # Pre-flight: verify the cover is openable. A corrupt cover
    # used to crash deep inside the renderer; now we surface a
    # clear error and still attempt to draw a placeholder.
    if not cover_bytes:
        warnings.append(
            "No cover image was provided — a placeholder will be stamped instead."
        )
        log.warning("[engine] no cover bytes; renderer will draw a placeholder")
    else:
        try:
            from PIL import Image as _PILImage  # local import
            _PILImage.open(io.BytesIO(cover_bytes))
        except Exception as exc:
            warnings.append(
                f"Cover image could not be opened: {exc}. "
                "A placeholder will be stamped instead."
            )
            log.warning("[engine] cover image invalid: %s", exc)

    try:
        log.info("Starting vector book generation for %s (model=%s)", filename, layout_model)
        raw = _call_openrouter_chat(theme_note, model=layout_model)
        layout = load_layout_dict(raw)
        log.info("Loaded layout: %d elements", len(layout.get("elements") or []))

        pdf_bytes, info = render_layout_to_pdf(
            layout=layout,
            cover_bytes=cover_bytes,
            page_count=page_count,
            page_size_in=page_size,
            output_path=out_path,
        )
        cover_info = info.get("cover_info") or cover_info
        for w in info.get("cover_warnings") or []:
            if w not in warnings:
                warnings.append(w)
        log.info("Wrote %d bytes to %s (cover_stamped=%s)",
                 len(pdf_bytes), out_path, info.get("cover_stamped"))
        return True, out_path, warnings, cover_info
    except Exception as exc:
        log.error("Vector book generation failed: %s", exc, exc_info=True)
        return False, None, warnings, cover_info


# Pillow helpers
# ---------------------------------------------------------------------------


def _render_svg_to_png(svg_bytes: bytes,
                       work_w: int = WORK_WIDTH_PX,
                       work_h: Optional[int] = None) -> bytes:
    """Render an SVG byte string to high-resolution PNG bytes.
    Uses ``resvg-py`` (a self-contained Rust SVG renderer with pre-built
    wheels — no system ``cairo`` or other native deps required, so it
    works on stock Hugging Face Spaces containers).
    The output is sized to fit the working canvas (default ``work_w``
    wide, ``work_h`` tall — derived from the 6:9 aspect ratio when not
    provided). Returns raw RGBA PNG bytes.
    """
    if work_h is None:
        work_h = int(work_w * 9 / 6)

    try:
        import resvg_py  # type: ignore
    except ImportError as exc:
        raise RuntimeError(
            "resvg-py is not installed. Add `resvg-py>=0.3` to requirements.txt."
        ) from exc

    log.info("[svg->png] start: target=%dx%d px, svg=%d bytes", work_w, work_h, len(svg_bytes))
    import time
    t0 = time.time()
    png_bytes = resvg_py.svg_to_bytes(
        svg_string=svg_bytes.decode("utf-8", errors="replace"),
        width=work_w,
        height=work_h,
    )
    log.info(
        "[svg->png] done: %d bytes in %.2fs (resvg-py)",
        len(png_bytes), time.time() - t0,
    )
    return png_bytes

def _stamp_cover(raw_bytes: bytes, cover_bytes: bytes,
                 top_in: float = TOP_BANNER_INCHES,
                 buffer_in: float = TITLE_BUFFER_INCHES,
                 work_w: int = WORK_WIDTH_PX) -> bytes:
    """Composite the cover on the top of the page with a title-safe buffer.
    The cover occupies the top ``top_in`` inches. An additional
    ``buffer_in`` inches of empty space separates the cover from the
    AI's design so the title is never clipped by the cover band
    (V8.8: fixed "ACTIVITY LOG & TRACKER" title-clipping bug).
    The cover is cropped to its visible (non-transparent) content
    bbox, then scaled to FIT inside the band (preserving aspect). The
    band itself is sized to ``work_w`` wide and ``stamp_h`` tall, and
    the cover is centered within it. Sticker-style covers with large
    transparent margins around the subject are handled correctly.
    """
    work_h = int(work_w * 9 / 6)
    stamp_h = int(work_w * top_in / 6)
    buffer_px = int(work_w * buffer_in / 6)
    safe_h = stamp_h + buffer_px  # total empty zone at top (cover + buffer)

    ai = Image.open(io.BytesIO(raw_bytes)).convert("RGB")
    ai = ai.resize((work_w, work_h), Image.LANCZOS)

    # SAFETY NET: paint a solid rectangle over the top ``safe_h``
    # pixels of the AI image. This wipes any AI content in both the
    # cover band AND the buffer zone, so the design starts cleanly
    # below the buffer. We sample the background color from the AI
    # image so we don't create a harsh white block if the image
    # isn't pure white.
    from PIL import ImageDraw
    bg_color = ai.getpixel((10, 10))
    ImageDraw.Draw(ai).rectangle(
        [(0, 0), (work_w, safe_h)], fill=bg_color
    )

    cover_rgba = Image.open(io.BytesIO(cover_bytes)).convert("RGBA")
    if cover_rgba.size[0] <= 0 or cover_rgba.size[1] <= 0:
        raise ValueError("Cover image has zero dimensions.")

    # bbox of the non-transparent region, using a threshold to ignore
    # faint anti-aliased edges far from the subject. Fall back to the
    # full image if everything is opaque or everything is transparent.
    alpha = cover_rgba.split()[-1]
    # Threshold: any pixel with alpha > 128 is "solid content". This
    # ignores the wispy anti-aliased halo around a sticker subject.
    solid = alpha.point(lambda a: 255 if a > 128 else 0)
    bbox = solid.getbbox()
    if bbox is None:
        cropped_rgba = cover_rgba
    else:
        x0, y0, x1, y1 = bbox
        x0 = max(0, x0 - 2)
        y0 = max(0, y0 - 2)
        x1 = min(cover_rgba.size[0], x1 + 2)
        y1 = min(cover_rgba.size[1], y1 + 2)
        cropped_rgba = cover_rgba.crop((x0, y0, x1, y1))

    # Flatten onto white so transparency becomes white.
    cw, ch = cropped_rgba.size
    flat = Image.new("RGB", (cw, ch), (255, 255, 255))
    flat.paste(cropped_rgba, (0, 0), cropped_rgba)

    # Scale to FIT inside the band, preserving aspect ratio. This way
    # the cover's subject is always fully visible inside the band.
    band_aspect = work_w / stamp_h            # ~4.0 for 1.5 in band
    cover_aspect = cw / ch
    if cover_aspect > band_aspect:
        # Cover is wider than band — limit by width.
        new_w = work_w
        new_h = max(1, int(work_w / cover_aspect))
    else:
        # Cover is taller than band — limit by height.
        new_h = stamp_h
        new_w = max(1, int(stamp_h * cover_aspect))
    scaled = flat.resize((new_w, new_h), Image.LANCZOS)

    # Center the scaled cover inside the band.
    band = Image.new("RGB", (work_w, stamp_h), bg_color)
    x0 = (work_w - new_w) // 2
    y0 = (stamp_h - new_h) // 2
    band.paste(scaled, (x0, y0))

    # The AI often ignores the "leave top blank" rule and draws its
    # title overlapping the banner. Instead of just painting over the
    # AI image (which cuts the title in half), we shift the entire AI
    # image down by ``safe_h`` so the design starts cleanly below the
    # cover band + buffer. The bottom of the AI is cropped to fit the
    # page (V8.8: extra buffer prevents title clipping).

    final_canvas = Image.new("RGB", (work_w, work_h), bg_color)
    final_canvas.paste(ai, (0, safe_h))
    final_canvas.paste(band, (0, 0))

    buf = io.BytesIO()
    final_canvas.save(buf, format="PNG")
    return buf.getvalue()

def _should_vectorize() -> bool:
    """Return True iff the VECTORIZE environment flag is enabled."""
    return os.environ.get("VECTORIZE") == "1"


def _quantize_palette(img: Image.Image, n_colors: int = 32) -> Image.Image:
    """Posterize + quantize an image down to ``n_colors`` flat colors.
    Pipeline: PIL.ImageOps.posterize(5 bits/channel) -> MEDIANCUT quantize
    to ``n_colors`` -> back to RGB. Produces a true flat-color Illustrator
    look with no anti-aliased halos.
    """
    log.info("[quantize] start: posterize=5 bits, n_colors=%d, input=%s", n_colors, img.size)
    img_p = ImageOps.posterize(img.convert("RGB"), 5)
    quantized = img_p.quantize(colors=n_colors, method=Image.Quantize.MEDIANCUT)
    result = quantized.convert("RGB")
    try:
        unique = len(set(result.getdata()))
    except Exception:
        unique = -1
    log.info("[quantize] done: %d unique colors", unique)
    return result


def _vectorize_lines(img: Image.Image, target_px: Tuple[int, int]) -> Image.Image:
    """Best-effort vectorization via ``potrace`` + ``svglib``/``reportlab``.
    Converts the image to a binary B&W mask, traces it to SVG with
    ``potrace``, then re-rasterizes the SVG at ``target_px`` and returns
    the resulting PIL image. This is opt-in (controlled by the
    ``VECTORIZE=1`` env var). On ANY failure (missing binary, subprocess
    error, library error) the original ``img`` is returned unchanged
    and a warning is logged — the default path is always safe to run
    on a stock HF Spaces free-tier container.
    """
    if not _should_vectorize():
        return img

    log.info("[vectorize] start: target_px=%s, input=%s", target_px, img.size)
    try:
        # Convert to grayscale, then threshold to pure B&W.
        gray = img.convert("L")
        bw = gray.point(lambda p: 0 if p < 128 else 255).convert("1").convert("RGB")

        with tempfile.TemporaryDirectory() as tmpdir:
            png_path = os.path.join(tmpdir, "input.png")
            svg_path = os.path.join(tmpdir, "output.svg")
            bw.save(png_path)

            # 1) Trace PNG -> SVG with potrace.
            proc = subprocess.run(
                ["potrace", "-s", "-o", svg_path, png_path],
                capture_output=True,
                text=True,
                timeout=60,
            )
            if proc.returncode != 0:
                log.warning(
                    "[vectorize] potrace failed (rc=%d): %s",
                    proc.returncode, (proc.stderr or "")[:200].strip(),
                )
                return img

            # 2) Re-rasterize SVG at target_px.
            from svglib.svglib import svg2rlg  # type: ignore
            from reportlab.graphics import renderPM  # type: ignore

            drawing = svg2rlg(svg_path)
            if drawing is None:
                log.warning("[vectorize] svglib returned None drawing.")
                return img

            # Scale drawing so its largest side matches target_px.
            try:
                dw = float(drawing.width) if drawing.width else 1.0
                dh = float(drawing.height) if drawing.height else 1.0
                scale = min(target_px[0] / dw, target_px[1] / dh)
                drawing.scale(scale, scale)
                drawing.width = target_px[0]
                drawing.height = target_px[1]
            except Exception:
                # If scaling fails, still try to render at native size.
                pass

            png_buf = io.BytesIO()
            renderPM.drawToFile(drawing, png_buf, fmt="PNG")
            png_buf.seek(0)
            traced = Image.open(png_buf).convert("RGB")
            log.info("[vectorize] done: %s -> %s", img.size, traced.size)
            return traced

    except FileNotFoundError as exc:
        log.warning("[vectorize] potrace binary not found (%s). Falling back.", exc)
        return img
    except subprocess.TimeoutExpired as exc:
        log.warning("[vectorize] potrace timed out (%s). Falling back.", exc)
        return img
    except Exception as exc:
        log.warning(
            "[vectorize] failed (%s: %s). Falling back to raster path.",
            type(exc).__name__, exc,
        )
        return img


def _bilateral_filter(img: Image.Image, d: int = 9,
                      sigma_color: float = 75,
                      sigma_space: float = 75) -> Image.Image:
    """Edge-preserving smoothing filter via OpenCV (cv2.bilateralFilter).
    Bilateral filter smooths flat-color regions while keeping hard
    edges sharp — the key to a "vector / illustrator" look on a
    raster image. Falls back to the input image if OpenCV is
    unavailable.
    """
    try:
        import cv2  # type: ignore
        import numpy as np
    except ImportError:
        log.warning("[bilateral] opencv-python not available; skipping.")
        return img

    arr = np.asarray(img)
    # cv2.bilateralFilter expects uint8 BGR; we work in RGB.
    filtered = cv2.bilateralFilter(arr, d, sigma_color, sigma_space)
    out = Image.fromarray(filtered, mode=img.mode)
    log.info(
        "[bilateral] d=%d, sigma_color=%.1f, sigma_space=%.1f, %s -> %s",
        d, sigma_color, sigma_space, img.size, out.size,
    )
    return out


def _posterize_safe(img: Image.Image, bits: int = 6) -> Image.Image:
    """Posterize to ``bits`` per channel, then recover edge sharpness.
    The posterize step produces the "flat color" illustrator look but
    flattens anti-aliased text edges. We follow it with a mild unsharp
    mask so the text regains crispness without re-introducing the
    noise the posterize removed.
    """
    log.info("[posterize-safe] bits=%d, input=%s", bits, img.size)
    p = ImageOps.posterize(img.convert("RGB"), bits)
    # Recover edge sharpness lost to quantization.
    p = p.filter(ImageFilter.UnsharpMask(radius=1, percent=120, threshold=2))
    log.info("[posterize-safe] done: %s", p.size)
    return p


def _vector_overlay(img: Image.Image,
                    page_in: Tuple[float, float] = DEFAULT_PAGE_SIZE_INCHES,
                    draw_border: bool = True,
                    draw_dividers: bool = True) -> Image.Image:
    """Overlay crisp vector primitives on top of the raster image.
    Adds:
        * A 1px inner border (just inside the trim).
        * A subtle top divider line under the cover band.
        * A subtle bottom divider above the bottom margin.
    These are drawn with PIL primitives (Draw.line / Draw.rectangle)
    so they're pixel-perfect, anti-aliased but mathematically straight.
    """
    from PIL import ImageDraw

    w, h = img.size
    draw = ImageDraw.Draw(img, "RGB")
    margin_px = int(w * 0.04)  # ~0.24" inside trim

    if draw_border:
        # Outer border (1px), inset by margin_px on all sides.
        draw.rectangle(
            [(margin_px, margin_px), (w - margin_px, h - margin_px)],
            outline=(0, 0, 0), width=2,
        )
        log.info("[vector-overlay] border: %d px inset, %dx%d", margin_px, w, h)

    if draw_dividers:
        # Horizontal divider just below the buffer zone (cover + 0.4 in
        # buffer = 1.9 in). Draws at 1.95 in to leave a small gap.
        div_top_y = int(h * (TOP_BANNER_INCHES + TITLE_BUFFER_INCHES + 0.05) / page_in[1])
        draw.line(
            [(margin_px, div_top_y), (w - margin_px, div_top_y)],
            fill=(0, 0, 0), width=2,
        )
        # Horizontal divider above bottom margin.
        div_bot_y = int(h * (page_in[1] - 0.4) / page_in[1])
        draw.line(
            [(margin_px, div_bot_y), (w - margin_px, div_bot_y)],
            fill=(0, 0, 0), width=2,
        )
        log.info("[vector-overlay] dividers: y=%d, y=%d", div_top_y, div_bot_y)

    return img


def _ultra_sharpen(raw_bytes: bytes,
                   target_dpi: int = TARGET_DPI,
                   page_in: Tuple[float, float] = DEFAULT_PAGE_SIZE_INCHES,
                   vectorize: bool = False) -> bytes:
    """V8.7 illustrator pipeline: smooth rescale + bilateral + posterize
    + edge enhance + vector overlay.
    Heavy filters (bilateral, posterize) run at WORK_WIDTH_PX for speed
    — O(N*d^2) bilateral at 4x fewer pixels is ~4x faster. The final
    LANCZOS upscale to ``target_dpi`` is essentially free and recovers
    every pixel of detail.
    Passes (in order):
        1. LANCZOS resize to WORK_WIDTH_PX (3600x5400)
        2. Bilateral filter (cv2) — smooth flat areas, keep edges crisp
        3. Posterize-safe (6 bits) — flat color palette, edges recovered
        4. EDGE_ENHANCE_MORE — push outlines harder
        5. ImageEnhance.Sharpness(1.5) — final crisp pass
        6. LANCZOS upscale to target_dpi
        7. Vector overlay — crisp border + section dividers (PIL primitives)
    Returns PNG bytes tagged with ``dpi=(target_dpi, target_dpi)``.
    """
    target_px = (int(page_in[0] * target_dpi), int(page_in[1] * target_dpi))
    work_px = (WORK_WIDTH_PX, int(WORK_WIDTH_PX * page_in[1] / page_in[0]))
    log.info(
        "[ultra-sharpen] start: target_dpi=%d, page_in=%s -> work=%s -> final=%s px",
        target_dpi, page_in, work_px, target_px,
    )

    # Pass 1: LANCZOS resize to working resolution.
    img = Image.open(io.BytesIO(raw_bytes)).convert("RGB")
    img = img.resize(work_px, Image.LANCZOS)
    log.info("[ultra-sharpen] pass 1/7: LANCZOS resize -> %s", img.size)

    # Pass 2: Bilateral filter (edge-preserving smoothing) at work res.
    img = _bilateral_filter(img, d=9, sigma_color=75, sigma_space=75)
    log.info("[ultra-sharpen] pass 2/7: bilateral filter")

    # Pass 3: Posterize-safe (flat color palette + edge recovery).
    img = _posterize_safe(img, bits=6)
    log.info("[ultra-sharpen] pass 3/7: posterize-safe (6 bits)")

    # Pass 4: Edge enhance.
    img = img.filter(ImageFilter.EDGE_ENHANCE_MORE)
    log.info("[ultra-sharpen] pass 4/7: EDGE_ENHANCE_MORE")

    # Pass 5: Final sharpness enhancement.
    img = ImageEnhance.Sharpness(img).enhance(1.5)
    log.info("[ultra-sharpen] pass 5/7: ImageEnhance.Sharpness(1.5)")

    # Pass 6: LANCZOS upscale to print resolution.
    img = img.resize(target_px, Image.LANCZOS)
    log.info("[ultra-sharpen] pass 6/7: LANCZOS upscale -> %s", img.size)

    # Pass 7: Vector overlay (border + dividers) at print resolution.
    img = _vector_overlay(img, page_in=page_in)
    log.info("[ultra-sharpen] pass 7/7: vector overlay")

    buf = io.BytesIO()
    img.save(buf, format="PNG", dpi=(target_dpi, target_dpi))
    log.info("[ultra-sharpen] done: %d bytes, %s", buf.tell(), img.size)
    return buf.getvalue()


def _embed_in_pdf(image_bytes: bytes,
                  page_count: int,
                  page_size: Tuple[float, float] = DEFAULT_PAGE_SIZE_INCHES) -> bytes:
    """Build a full-bleed PDF with ``page_count`` copies of the image."""
    w_pt = page_size[0] * inch
    h_pt = page_size[1] * inch
    buf = io.BytesIO()
    c = Canvas(buf, pagesize=(w_pt, h_pt))
    for _ in range(max(1, int(page_count))):
        c.drawImage(
            ImageReader(io.BytesIO(image_bytes)),
            0, 0,
            width=w_pt, height=h_pt,
            preserveAspectRatio=True,
            mask="auto",
        )
        c.showPage()
    c.save()
    return buf.getvalue()


# Public entry point
# ---------------------------------------------------------------------------

def generate_book(
    cover_bytes: bytes,
    filename: str,
    theme_note: str = "",
    page_count: int = 1,
    page_size: Tuple[float, float] = DEFAULT_PAGE_SIZE_INCHES,
    model: str = DEFAULT_MODEL,
) -> Tuple[bool, Optional[Path]]:
    """Orchestrate the V8 generation pipeline.
    ``model`` selects the image generator:
        * ``google/gemini-3.1-flash-lite-image`` (default) — fast PNG
        * ``recraft/recraft-v4.1-vector`` — true SVG vectors
          (rasterized to high-res PNG via ``resvg-py`` for the cover
          stamp overlay; no quality loss because the underlying art is
          vector-crisp at any zoom level).
    """
    if not filename.endswith(".pdf"):
        filename += ".pdf"
    out_path = OUTPUTS_DIR / filename

    try:
        # 1) OpenRouter image gen (PNG or SVG)
        log.info("Starting book generation for %s (model=%s)", filename, model)
        raw_ai, media_type = _call_openrouter_images(theme_note, model=model)

        # 1a) Vector models (Recraft) come back as SVG — rasterize to PNG
        #     at the working resolution. The cover stamp needs a raster
        #     canvas to composite onto, but because the source is vector,
        #     the resvg render is mathematically perfect at any DPI.
        if media_type == "image/svg+xml":
            log.info("Vector response detected (model=%s); rasterizing via resvg-py.", model)
            raw_ai = _render_svg_to_png(raw_ai, work_w=WORK_WIDTH_PX)

        # 2) Stamp cover
        stamped = _stamp_cover(raw_ai, cover_bytes, top_in=TOP_BANNER_INCHES, work_w=WORK_WIDTH_PX)

        # 3) Upscale to print
        vectorize_flag = _should_vectorize()
        final_bytes = _ultra_sharpen(stamped, target_dpi=TARGET_DPI, page_in=page_size, vectorize=vectorize_flag)

        # 4) PDF
        pdf_bytes = _embed_in_pdf(final_bytes, page_count, page_size)

        out_path.write_bytes(pdf_bytes)
        log.info("Successfully wrote %d bytes to %s", len(pdf_bytes), out_path)
        return True, out_path

    except Exception as exc:
        log.error("Book generation failed: %s", exc, exc_info=True)
        return False, None


# CLI entry point (optional, for ad-hoc testing)
# ---------------------------------------------------------------------------

def _cli() -> int:
    import argparse

    parser = argparse.ArgumentParser(description="V8 AI-image KDP PDF generator")
    parser.add_argument("--cover", required=True, help="Path to cover image (PNG/JPG/WebP)")
    parser.add_argument("--out", default="book", help="Output filename stem (no .pdf)")
    parser.add_argument("--theme", default="", help="Optional theme note")
    parser.add_argument("--pages", type=int, default=1, help="Page count")
    parser.add_argument("--width", type=float, default=6.0, help="Page width (in)")
    parser.add_argument("--height", type=float, default=9.0, help="Page height (in)")
    parser.add_argument(
        "--model",
        default=DEFAULT_MODEL,
        help=(
            "OpenRouter image model. Default: gemini-3.1-flash-lite-image. "
            "Use 'recraft/recraft-v4.1-vector' for true vector output."
        ),
    )
    args = parser.parse_args()

    cover_path = Path(args.cover)
    if not cover_path.is_file():
        log.error("Cover not found: %s", cover_path)
        return 2

    ok, out = generate_book(
        cover_path.read_bytes(),
        args.out,
        args.theme,
        args.pages,
        (args.width, args.height),
        model=args.model,
    )
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(_cli())