Sync from GitHub a877e8c

.env.example

@@ -11,6 +11,7 @@ GROQ_API_KEY=
 # DEBUG=false
 # MAX_UPLOAD_MB=50
 # PRELOAD_MODELS=true          # load models at startup instead of on the first request
+# KEEP_WARM_INTERVAL_S=300     # after this many idle seconds, run a dummy inference so the OS keeps models in RAM (0 = off)
 # APP_BASE_DIR=                # defaults to the project folder
 
 # --- LLM ---
@@ -29,6 +30,7 @@ GROQ_API_KEY=
 # USE_IMAGE_SIGNATURE=true
 # MATCH_WEIGHTS=0.5,0.3,0.2    # embedding, phash, image-signature
 # MIN_MATCH_COSINE=0.4         # reject figures whose best raw DINOv2 cosine is below this (0 = off); read "cos=" in the logs to calibrate
+# MATCH_MAX_SIDE=512           # downscale figures to this longest side before enhancement (0 = off); ~2.4x faster matching
 # REJECT_DIR=_rejects          # folder in the asset library with non-asset examples (code-block screenshots); figures closer to these are dropped
 # DINOV2_MODEL=facebook/dinov2-small   # must be the model the embeddings were generated with
 # TORCH_CPU_THREADS=4
@@ -36,5 +38,7 @@ GROQ_API_KEY=
 
 # --- Pipeline ---
 # PDF_DPI=300
+# FIGURE_EXTRACTOR=layout      # layout = YOLOX on the rendered pages (~1 s/page); hi_res = unstructured partition_pdf with OCR (~10 s/page)
+# LAYOUT_MODEL=yolox           # unstructured_inference model for the layout extractor (yolox, yolox_tiny, yolox_quantized)
 # GRAPH_RECURSION_LIMIT=200    # raised automatically for long PDFs
 # ZOOM_CODE_REGION=true        # also send a full-resolution crop of the code blocks (fixes misread values/dropped repeats)

Dockerfile

@@ -6,7 +6,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
     PYTHONUNBUFFERED=1 \
     PYTHONDONTWRITEBYTECODE=1
 
-# poppler: pdf2image page rendering; tesseract: OCR inside unstructured's hi_res strategy;
+# poppler: pdf2image page rendering; tesseract: OCR for the FIGURE_EXTRACTOR=hi_res fallback only;
 # libgl1/libglib2.0-0: OpenCV runtime.
 RUN apt-get update && apt-get install -y --no-install-recommends \
         poppler-utils \

scratch_vision/__init__.py

@@ -26,7 +26,7 @@ def create_app(settings: Optional[Settings] = None) -> Flask:
     """Application factory: configure logging, validate settings, register routes, load models."""
     from .api.errors import register_error_handlers
     from .api.routes import bp
-    from .core.models import warmup
+    from .core.models import start_keep_warm, warmup
 
     settings = settings or get_settings()
     setup_logging(settings)
@@ -44,6 +44,7 @@ def create_app(settings: Optional[Settings] = None) -> Flask:
 
     if settings.preload_models:
         warmup()
+        start_keep_warm(settings.keep_warm_interval_s)
     return app
 
 

scratch_vision/config.py

@@ -46,6 +46,7 @@ class Settings:
     public_base_url: str  # e.g. https://<space>.hf.space ; empty -> derived from the request
     max_upload_mb: int
     preload_models: bool
+    keep_warm_interval_s: float  # re-run a dummy inference after this much idle time so models stay in RAM (0 = off)
 
     # --- LLM ---
     groq_api_key: str
@@ -64,8 +65,11 @@ class Settings:
     use_image_signature: bool  # include the Goldberg image-signature signal in matching
     min_match_cosine: float  # reject a match whose raw DINOv2 cosine is below this (0 = off); real matches ~0.65+, code blocks ~0.1
     reject_dir_name: str  # folder in the asset library with images that are NOT assets (e.g. code blocks)
+    match_max_side: int  # downscale figures to this longest side before enhancement (0 = off); enhancing is O(pixels)
 
     # --- Pipeline ---
+    figure_extractor: str  # "layout": layout model on the rendered pages (fast); "hi_res": unstructured partition_pdf (OCR, slow)
+    layout_model: str  # unstructured_inference layout model used by the "layout" extractor
     pdf_dpi: int
     graph_recursion_limit: int
     zoom_code_region: bool  # also send a full-resolution crop of the code blocks with each page
@@ -121,6 +125,8 @@ class Settings:
             raise RuntimeError("GROQ_API_KEY is not set. Add it to .env or the Space secrets.")
         if not self.assets_dir.is_dir():
             raise RuntimeError(f"Asset library not found at {self.assets_dir}")
+        if self.figure_extractor not in ("layout", "hi_res"):
+            raise RuntimeError(f"FIGURE_EXTRACTOR must be 'layout' or 'hi_res', got '{self.figure_extractor}'")
 
 
 def _resolve_base_dir() -> Path:
@@ -161,6 +167,7 @@ def get_settings() -> Settings:
         public_base_url=_env_str("PUBLIC_BASE_URL").rstrip("/"),
         max_upload_mb=_env_int("MAX_UPLOAD_MB", 50),
         preload_models=_env_bool("PRELOAD_MODELS", True),
+        keep_warm_interval_s=_env_float("KEEP_WARM_INTERVAL_S", 300.0),
         groq_api_key=_env_str("GROQ_API_KEY"),
         llm_model=_env_str("LLM_MODEL", "qwen/qwen3.8-27b"),
         llm_temperature=_env_float("LLM_TEMPERATURE", 0.0),
@@ -175,6 +182,9 @@ def get_settings() -> Settings:
         use_image_signature=_env_bool("USE_IMAGE_SIGNATURE", True),
         min_match_cosine=_env_float("MIN_MATCH_COSINE", 0.4),
         reject_dir_name=_env_str("REJECT_DIR", "_rejects"),
+        match_max_side=_env_int("MATCH_MAX_SIDE", 512),
+        figure_extractor=_env_str("FIGURE_EXTRACTOR", "layout").lower(),
+        layout_model=_env_str("LAYOUT_MODEL", "yolox"),
         pdf_dpi=_env_int("PDF_DPI", 300),
         graph_recursion_limit=_env_int("GRAPH_RECURSION_LIMIT", 200),
         zoom_code_region=_env_bool("ZOOM_CODE_REGION", True),

scratch_vision/core/models.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 
 import logging
 import threading
+import time
 from dataclasses import dataclass
 from typing import Any
 
@@ -26,6 +27,7 @@ class Dinov2Bundle:
 _dinov2 = None
 _phash = None
 _gis = None
+_layout = None
 _llm = None
 
 
@@ -70,6 +72,20 @@ def get_image_signature():
     return _gis
 
 
+def get_layout_model():
+    """Layout detection model (YOLOX by default) that finds the figures on each PDF page."""
+    global _layout
+    if _layout is None:
+        with _lock:
+            if _layout is None:
+                from unstructured_inference.models.base import get_model
+
+                name = get_settings().layout_model
+                logger.info("Loading layout model %s", name)
+                _layout = get_model(name)
+    return _layout
+
+
 def get_llm():
     global _llm
     if _llm is None:
@@ -97,6 +113,69 @@ def warmup() -> None:
     get_dinov2()
     get_phash()
     get_image_signature()
+    if get_settings().figure_extractor == "layout":
+        get_layout_model()
     get_llm()
     get_reference_index()
     logger.info("All models and the reference index are loaded")
+    run_dummy_inference()
+
+
+def run_dummy_inference() -> None:
+    """One tiny pass through every local model and the PDF renderer.
+
+    Loading a model is not enough: the first real inference also pays for lazy kernel/session setup,
+    and after a long idle the OS pages the weights out, which made the first request take ~1 minute
+    instead of ~10s. Running this at boot and periodically while idle keeps that cost off requests.
+    """
+    import tempfile
+    from pathlib import Path
+
+    import numpy as np
+    from PIL import Image
+
+    from ..pdf.extraction import render_pages
+    from ..vision.embeddings import get_dinov2_embedding
+
+    start = time.perf_counter()
+    try:
+        blank = Image.new("RGB", (620, 877), "white")  # A4 at 75 dpi
+        with tempfile.TemporaryDirectory() as tmp:
+            pdf_path = Path(tmp) / "warmup.pdf"
+            blank.save(pdf_path, format="PDF")
+            render_pages(pdf_path, 72)
+        if get_settings().figure_extractor == "layout":
+            get_layout_model().image_processing(blank)
+        get_dinov2_embedding(blank.resize((224, 224)))
+        get_phash().encode_image(image_array=np.asarray(blank.convert("L").resize((64, 64))))
+    except Exception as e:  # warming is best effort; never take the app down over it
+        logger.warning("Dummy inference failed: %s", e)
+        return
+    logger.info("Dummy inference finished in %.2fs", time.perf_counter() - start)
+
+
+_last_activity = time.monotonic()
+
+
+def mark_active() -> None:
+    """Called by each request so the keep-warm thread only runs while the app is idle."""
+    global _last_activity
+    _last_activity = time.monotonic()
+
+
+def start_keep_warm(interval_s: float) -> None:
+    """Daemon thread that re-runs the dummy inference after every `interval_s` seconds of idle time."""
+    if interval_s <= 0:
+        return
+
+    def loop():
+        while True:
+            idle = time.monotonic() - _last_activity
+            if idle < interval_s:
+                time.sleep(interval_s - idle)
+                continue
+            run_dummy_inference()
+            mark_active()
+
+    threading.Thread(target=loop, name="keep-warm", daemon=True).start()
+    logger.info("Keep-warm thread started (every %ss of idle time)", interval_s)

scratch_vision/pdf/extraction.py

@@ -1,7 +1,8 @@
-"""PDF input handling: persist the upload, extract embedded figures, render pages."""
+"""PDF input handling: persist the upload, render pages, extract the figures on them."""
 from __future__ import annotations
 
 import base64
+import io
 import logging
 import uuid
 from pathlib import Path
@@ -60,6 +61,53 @@ def extract_figures(pdf_path: Path, output_dir: Path) -> Dict[str, dict]:
     return figures
 
 
+MIN_FIGURE_SIDE = 32  # px at PDF_DPI; smaller boxes are icons or detection noise
+ROW_TOLERANCE = 0.05  # boxes whose tops are within 5% of the page height count as one row
+
+
+def extract_figures_from_pages(pages: List[Image.Image]) -> Dict[str, dict]:
+    """Detect figures on already-rendered pages with the layout model and crop them out.
+
+    Same result shape as `extract_figures`, but skips partition_pdf's second render, full-page OCR and
+    text merging, which the app never uses. Figures are numbered page by page, top to bottom, left to right.
+    """
+    from unstructured_inference.constants import ElementType
+
+    from ..core.models import get_layout_model
+
+    model = get_layout_model()
+    pdf_id = uuid.uuid4().hex
+    figures: Dict[str, dict] = {}
+    for page_no, page in enumerate(pages, start=1):
+        rgb = page.convert("RGB")
+        try:
+            layout = model.image_processing(rgb)
+        except Exception as e:
+            raise PdfExtractionError(f"Layout detection failed on page {page_no}: {e}") from e
+
+        boxes = []
+        for coords, class_id in zip(layout.element_coords, layout.element_class_ids):
+            if layout.element_class_id_map.get(int(class_id)) != ElementType.PICTURE:
+                continue
+            x1, y1 = max(0, int(coords[0])), max(0, int(coords[1]))
+            x2, y2 = min(rgb.width, int(round(coords[2]))), min(rgb.height, int(round(coords[3])))
+            if x2 - x1 >= MIN_FIGURE_SIDE and y2 - y1 >= MIN_FIGURE_SIDE:
+                boxes.append((x1, y1, x2, y2))
+        boxes.sort(key=lambda b: (round(b[1] / (rgb.height * ROW_TOLERANCE)), b[0]))
+
+        for box in boxes:
+            buffer = io.BytesIO()
+            rgb.crop(box).save(buffer, format="PNG")
+            figures[f"Sprite {len(figures) + 1}"] = {
+                "base64": base64.b64encode(buffer.getvalue()).decode("utf-8"),
+                "file-path": None,  # cropped in memory; nothing is written to disk
+                "pdf-id": pdf_id,
+                "image-uuid": str(uuid.uuid4()),
+            }
+    logger.info("Extracted %d figure(s) from %d page(s) with the layout model", len(figures), len(pages))
+    return figures
+
+
 def render_pages(pdf_path: Path, dpi: int) -> List[Image.Image]:
     from pdf2image import convert_from_path
 

scratch_vision/services/project_service.py

@@ -10,8 +10,9 @@ from pathlib import Path
 from typing import Dict, List, Optional
 
 from ..config import get_settings
+from ..core.models import mark_active
 from ..exceptions import InvalidUploadError
-from ..pdf.extraction import PDF_MAGIC, extract_figures, render_pages, save_pdf
+from ..pdf.extraction import PDF_MAGIC, extract_figures, extract_figures_from_pages, render_pages, save_pdf
 from ..pipeline.graph import get_graph
 from ..pipeline.nodes import NO_CODE_RESULT, find_target
 from ..pipeline.state import initial_state
@@ -103,6 +104,7 @@ def _attach_code(described: List[Dict], pages: List[Dict], targets: list) -> Lis
 def process_pdf(pdf_bytes: bytes, project_id: str) -> ProcessResult:
     settings = get_settings()
     validate_pdf(pdf_bytes)
+    mark_active()
 
     project_folder = settings.output_dir / project_id  # becomes the .sb3 contents
     work_dir = settings.work_dir / project_id          # private scratch space, deleted below
@@ -112,17 +114,22 @@ def process_pdf(pdf_bytes: bytes, project_id: str) -> ProcessResult:
     try:
         pdf_path = save_pdf(pdf_bytes, settings.generated_dir / project_id / f"{project_id}.pdf")
 
+        # Rendered once: the layout model finds figures on these pages and the LLM reads them.
+        with timed(f"[{project_id}] page rendering", logger):
+            pages = render_pages(pdf_path, settings.pdf_dpi)
+
         # 1. Figures in the PDF -> matching sprites/backdrops from the library -> project skeleton
-        with timed(f"[{project_id}] figure extraction", logger):
-            figures = extract_figures(pdf_path, work_dir / "figures")
+        with timed(f"[{project_id}] figure extraction ({settings.figure_extractor})", logger):
+            if settings.figure_extractor == "hi_res":
+                figures = extract_figures(pdf_path, work_dir / "figures")
+            else:
+                figures = extract_figures_from_pages(pages)
         with timed(f"[{project_id}] asset matching", logger):
             matches = match_images({sid: f["base64"] for sid, f in figures.items()}, work_dir)
             project_json_path = build_project_skeleton([m.path for m in matches if m.path], project_folder)
         project_json = json.loads(project_json_path.read_text(encoding="utf-8"))
 
         # 2. Each page image -> LLM pseudocode -> Scratch blocks
-        with timed(f"[{project_id}] page rendering", logger):
-            pages = render_pages(pdf_path, settings.pdf_dpi)
         recursion_limit = max(settings.graph_recursion_limit, STEPS_PER_PAGE * len(pages) + 10)
         with timed(f"[{project_id}] LangGraph pipeline ({len(pages)} page(s))", logger):
             final_state = get_graph().invoke(

scratch_vision/vision/image_utils.py

@@ -56,6 +56,14 @@ def preprocess_for_model(pil_img: Image.Image) -> Optional[Image.Image]:
         return None
 
 
+def cap_size(pil_img: Image.Image, max_side: int) -> Image.Image:
+    """Downscale so the longest side is at most max_side (0 = no cap). Never upscales."""
+    if max_side <= 0 or max(pil_img.size) <= max_side:
+        return pil_img
+    ratio = max_side / max(pil_img.size)
+    return pil_img.resize((max(1, round(pil_img.width * ratio)), max(1, round(pil_img.height * ratio))), Image.LANCZOS)
+
+
 # ---------- OpenCV enhancement ----------
 def _pil_to_bgr(pil_img: Image.Image) -> np.ndarray:
     return cv2.cvtColor(np.array(pil_img.convert("RGB")), cv2.COLOR_RGB2BGR)

scratch_vision/vision/reference_index.py

@@ -89,7 +89,7 @@ def _embed_reject_examples(folder: Path) -> np.ndarray:
     from PIL import Image
 
     from .embeddings import get_dinov2_embedding
-    from .image_utils import enhance_image, preprocess_for_model
+    from .image_utils import cap_size, enhance_image, preprocess_for_model
 
     if not folder.is_dir():
         return np.empty((0, 0))
@@ -98,7 +98,8 @@ def _embed_reject_examples(folder: Path) -> np.ndarray:
         try:
             with Image.open(img_path) as img:
                 img.load()
-                enhanced = enhance_image(img, scale=2) or img
+                capped = cap_size(img, get_settings().match_max_side)
+                enhanced = enhance_image(capped, scale=2) or capped
                 emb = get_dinov2_embedding(preprocess_for_model(enhanced))
         except Exception as e:
             logger.warning("Could not embed reject example %s: %s", img_path, e)

scratch_vision/vision/search.py

@@ -14,7 +14,7 @@ from PIL import Image
 from ..config import get_settings
 from ..core.models import get_image_signature, get_phash
 from .embeddings import get_dinov2_embedding
-from .image_utils import base64_to_pil, enhance_image, preprocess_for_hash, preprocess_for_model
+from .image_utils import base64_to_pil, cap_size, enhance_image, preprocess_for_hash, preprocess_for_model
 from .reference_index import ReferenceIndex, get_reference_index
 
 logger = logging.getLogger(__name__)
@@ -36,6 +36,8 @@ class MatchResult:
 
 # ---------- 1. features of the query image ----------
 def compute_query_features(pil_img: Image.Image, work_dir: Path, name: str) -> Optional[QueryFeatures]:
+    # Enhancement cost grows with pixel count, and DINOv2 resizes to 224 px anyway: cap the size first.
+    pil_img = cap_size(pil_img, get_settings().match_max_side)
     enhanced = enhance_image(pil_img, scale=2) or pil_img
 
     embedding = get_dinov2_embedding(preprocess_for_model(enhanced))