Enhancement: Improve MVP Stability by Adding Configurable Florence Batching with OOM Fallback and Finalizing Singleton Service Wiring

[Santosh69]

Description

This PR keeps only the latest 3 commits for this branch and focuses on:

1. Using one shared app-level instance (app.state) for OmniParser, RAG KB, and Heuristic Engine.
2. Running Florence captioning in batches with a configurable batch size.
3. Retrying with smaller batches when memory errors happen, instead of failing the full request.
4. Making startup logs clearly show singleton initialization.

Note: This is the Sub-Issue of MVP currently working on and there are followed commits

---

Included Commits (this PR only)

• 8dcbb6d - refactor(singleton): wire OmniParser, RAG KB, and engine via app.state
• 5227b47 - feat(omniparser): add configurable Florence batching with adaptive OOM fallback
• 1aa0033 - chore(logging): label RAG and heuristic engine startup logs as singleton

---

Review Focus (Suggested Order)

1. app/services/omniparser_client.py - batching + adaptive fallback logic
2. app/core/config.py and .env.example - batch-size config contract
3. main.py, app/services/heuristic_engine.py, route files - singleton wiring

Key File Changes

1. OmniParser batching + memory fallback (app/services/omniparser_client.py)

Commit: 5227b47

Adds configurable Florence batching and safe fallback when memory is low.
This is the main logic change maintainers should review.

+ from app.core.config import ALLOWED_IMAGE_TYPES, MAX_IMAGE_SIZE_BYTES, FLORENCE_BATCH_SIZE
...
+ def _is_probable_oom(self, err: Exception) -> bool:
+     ...
+
+ def florence_batch_caption_adaptive(self, jobs, elements, max_batch_size: int) -> None:
+     ...
+     chunk_size = max(1, chunk_size // 2)
...
+ batch_crops: list[tuple[int, Image.Image]] = []
+ if len(batch_crops) == FLORENCE_BATCH_SIZE:
+     self.florence_batch_caption_adaptive(batch_crops, elements, FLORENCE_BATCH_SIZE)
+     batch_crops.clear()
...
+ if batch_crops:
+     self.florence_batch_caption_adaptive(batch_crops, elements, FLORENCE_BATCH_SIZE)
+     batch_crops.clear()

---

2. Batch size config from env (app/core/config.py, .env.example)

Commit: 5227b47

Makes Florence batch size configurable and validated.

class Settings(BaseSettings):
+   FLORENCE_BATCH_SIZE: int = Field(default=3, env="FLORENCE_BATCH_SIZE", ge=1)

settings = Settings()
+ FLORENCE_BATCH_SIZE = settings.FLORENCE_BATCH_SIZE

+ FLORENCE_BATCH_SIZE=3

---

3. Singleton service wiring in startup and routes (main.py, app/services/heuristic_engine.py, app/api/routes/evaluation.py, app/api/routes/heuristic.py)

Commit: 8dcbb6d

Ensures startup creates shared services once, and routes reuse those instances.

+ app.state.omniparser_client = OmniParserClient()
+ await app.state.omniparser_client.initialize()

+ app.state.rag_knowledge_base = RAGKnowledgeBase()
+ await app.state.rag_knowledge_base.initialize()

+ app.state.heuristic_engine = HeuristicEvaluationEngine(
+     rag_kb=app.state.rag_knowledge_base
+ )
+ await app.state.heuristic_engine.initialize()

- evaluation_engine = HeuristicEvaluationEngine()
- await evaluation_engine.initialize()
+ evaluation_engine = request.app.state.heuristic_engine

- kb = RAGKnowledgeBase()
- await kb.initialize()
+ kb = request.app.state.rag_knowledge_base

- if self.rag_kb is None:
-     self.rag_kb = RAGKnowledgeBase()
+ if self.rag_kb is None:
+     raise RuntimeError("RAGKnowledgeBase dependency missing. Inject via startup.")

---

4. Clear singleton startup logs (main.py)

Commit: 1aa0033

Small log clarity update so startup clearly shows singleton usage.

- logger.info("RAG knowledge base initialized")
+ logger.info("RAG knowledge base initialized (singleton)")

- logger.info("Heuristic evaluation engine initialized")
+ logger.info("Heuristic evaluation engine initialized (singleton)")

---

Before / After Flow

Before

• OmniParser, RAG KB, and Heuristic Engine could be created in more than one place (startup + route-level paths), which could lead to inconsistent instance usage.
• Florence captioning did not have a configurable batch-size control from env.
• If a Florence batch hit memory pressure, there was no adaptive chunk-size fallback path.

After

• Startup creates one shared instance for each core service, and routes use those via request.app.state.*.
• Florence batch size is configurable through FLORENCE_BATCH_SIZE.
• On probable OOM during captioning, batch processing retries with smaller chunk sizes (N -> N/2 -> ... -> 1) instead of failing the whole request.

---

Proof of Work

Custom Florence-Batching (OmniParser)

[Copilot]

Pull request overview

This PR improves MVP stability by moving heavyweight services (OmniParser, RAG KB, Heuristic Engine) to app-level singletons, and adds configurable Florence caption batching with adaptive retry on probable OOM errors. It also adds dev tooling to download required weights and start the stack more consistently across platforms.

Changes:

• Initialize and reuse OmniParser, RAGKnowledgeBase, and HeuristicEvaluationEngine via app.state singletons.
• Add configurable Florence batching (FLORENCE_BATCH_SIZE) with adaptive batch halving on probable OOM.
• Add cross-platform dev launcher and OmniParser weight setup helper; update docs and ignores for local weights.

Reviewed changes

Copilot reviewed 11 out of 13 changed files in this pull request and generated 10 comments.

Show a summary per file

File	Description
start-dev.py	New cross-platform dev launcher that preps weights and starts backend (+ optional firebase-functions).
scripts/setup_omniparser.py	New helper to download/lay out OmniParser v2 weights into weights/.
requirements.txt	Pins/extends ML stack deps (torch/transformers/numpy/einops/huggingface-hub/accelerate).
main.py	Startup now wires OmniParser/RAG/Heuristic Engine as singletons with clearer logs.
app/utils/transformers_patches.py	Adds a transformers monkey-patch to avoid flash-attn import checks on non-CUDA.
app/services/omniparser_client.py	Adds Florence batching + adaptive OOM fallback; adjusts model initialization and validation.
app/services/heuristic_engine.py	Makes RAG KB a constructor dependency and avoids reinitialization when already initialized.
app/core/config.py	Adds FLORENCE_BATCH_SIZE setting sourced from env with validation.
app/api/routes/heuristic.py	Routes now reuse the singleton OmniParser client.
app/api/routes/evaluation.py	Routes now reuse singleton engine and pass upload content_type into detection.
README.md	Updates quick start and documents weight setup paths and options.
.gitignore	Stops ignoring docs, and starts ignoring weights + local firebase config files.
.env.example	Documents FLORENCE_BATCH_SIZE env var.

Comments suppressed due to low confidence (2)

app/services/omniparser_client.py:9

• Unused imports: asyncio and json are imported but not referenced anywhere in this module. Removing them avoids confusion and keeps dependencies clear.

import asyncio
from typing import List, Dict, Any, Optional, Tuple
from PIL import Image
import io
import json

app/services/omniparser_client.py:361

• detect_elements still exposes an image_url parameter, but the implementation only processes image_data bytes and never fetches/uses image_url. As a result, callers like detect_elements(image_url=...) will always fail validation (no image bytes). Either implement URL fetching inside detect_elements (and validate content-type/size), or remove image_url from the API and update the routes accordingly.

    async def detect_elements(
        self,
        image_data: bytes,
        image_url: Optional[str] = None,
        content_type: str = "image/jpeg"

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The caption model is loaded from the local snapshot (weights/icon_caption_florence), but the processor is loaded from Hugging Face ("microsoft/Florence-2-base-ft"). This introduces an unexpected runtime network dependency and risks processor/model mismatch. Consider loading the processor from the same local directory (or ensure setup_omniparser downloads the processor artifacts and point AutoProcessor to the local path).

Suggested change

	self.processor = AutoProcessor.from_pretrained("microsoft/Florence-2-base-ft", trust_remote_code=True)
	self.processor = AutoProcessor.from_pretrained(
	"weights/icon_caption_florence",
	trust_remote_code=True,
	)

[Copilot]

Unused import: re is imported but not used (caption cleaning is done with .replace(...)). Either remove it or use it for the intended cleaning logic.

Suggested change

import re # For cleaning Florence output

[Copilot]

element_content is assigned but never read (captions are written directly into elements[element_index].content inside the batch captioner). Removing this variable and related dead assignments will make the captioning flow easier to follow.

[Copilot]

analyze_interface calls detect_elements(contents) without passing the upload content_type. Since detect_elements now validates content_type, this endpoint will always use the default ("image/jpeg"), which can incorrectly accept unsupported types (or reject supported ones if defaults change). Pass content_type=image.content_type (with a safe fallback) to keep validation consistent with /evaluation/evaluate.

Suggested change

	result = await client.detect_elements(contents)
	content_type = image.content_type or "image/jpeg"
	result = await client.detect_elements(contents, image.filename, content_type=content_type)

[Copilot]

The launcher only waits on the Uvicorn process. If Uvicorn exits/crashes, the script will exit without terminating firebase_process, leaving it running in the background. Consider wrapping the wait loop in a try/finally (or monitoring both processes) to ensure all child processes are terminated and to propagate the backend exit code.

[Copilot]

initialize() now raises if rag_kb wasn’t injected. This is a breaking behavior change for any direct HeuristicEvaluationEngine() usage (e.g., the existing integration test instantiates it without dependencies). If standalone construction should remain supported, consider defaulting to a new RAGKnowledgeBase() internally (or provide a clear factory/helper used by tests) rather than raising at runtime.

Suggested change

	raise RuntimeError("RAGKnowledgeBase dependency missing. Inject via startup.")
	self.logger.warning(
	"RAGKnowledgeBase dependency was not injected; creating a default instance."
	)
	self.rag_kb = RAGKnowledgeBase()

[Copilot]

Comment typo: "Initilize Hueristic" should be "Initialize Heuristic" for clarity/professionalism in startup logs/comments.

Suggested change

	# Initilize Hueristic Evaluation Engine (Singleton)
	# Initialize Heuristic Evaluation Engine (Singleton)

[Copilot]

detect_elements(..., content_type=content_type) can raise InvalidInputError (and OmniParser errors), but this handler currently falls through to the generic except Exception at the bottom of the endpoint, returning HTTP 500 for client errors. Add explicit except InvalidInputError / service-exception branches in this endpoint so invalid content types and validation failures return 400/422/503 as intended.

[Copilot]

Unused import: shutil is imported but never used in this script. Removing it keeps the setup script minimal.

Suggested change

import shutil

[Copilot]

The early-exit check only verifies weights/icon_caption_florence/model.safetensors, so the script may skip even if weights/icon_detect/model.pt is missing (or the caption folder is incomplete). Consider validating both required weight sets (detector + caption) before skipping, so start-dev.py reliably produces a runnable environment.

Suggested change

	# if model exists, do nothing
	safetensors_path = os.path.join(target_caption_folder, "model.safetensors")
	if os.path.exists(safetensors_path):
	print(f"✅ Models found in {target_caption_folder}. Skipping download.")
	# if both caption and detector models exist, do nothing
	safetensors_path = os.path.join(target_caption_folder, "model.safetensors")
	detector_model_path = os.path.join(weights_dir, "icon_detect", "model.pt")
	if os.path.exists(safetensors_path) and os.path.exists(detector_model_path):
	print(f"✅ Models found in {weights_dir}. Skipping download.")