diff --git a/PATCHES.md b/PATCHES.md
new file mode 100644
index 000000000..3be63e315
--- /dev/null
+++ b/PATCHES.md
@@ -0,0 +1,533 @@
+# Patches — parseq-marathi
+
+This document records every infrastructure and compatibility change made to this
+fork of the official [PARSeq](https://github.com/baudm/parseq) repository.
+**No model architecture, encoder/decoder, tokenizer, or recognition-pipeline code
+was modified.**
+
+---
+
+## 1. Unicode Compatibility Fix
+
+**File:** `strhub/data/dataset.py`
+
+### Problem
+The original `_preprocess_labels` method unconditionally applied
+`unicodedata.normalize('NFKD', label).encode('ascii', 'ignore').decode()` to
+every label.  
+For Devanagari (and any other non-ASCII script) this silently discards all
+characters, producing empty labels and an empty dataset.
+
+### Fix
+Added a helper function `_is_ascii_charset(charset: str) -> bool` that checks
+whether every character in the configured charset is a plain ASCII codepoint
+(< 128).
+
+`_preprocess_labels` now branches on this flag:
+
+| Charset type | Normalisation applied |
+|---|---|
+| ASCII-only (English) | `NFKD` → strip non-ASCII bytes *(original behaviour, unchanged)* |
+| Non-ASCII (e.g. Devanagari) | `NFC` only — composed characters such as Devanagari matras are preserved |
+
+### Backward compatibility
+English datasets continue to use exactly the original normalisation path.
+The change is gated entirely on whether the charset itself contains non-ASCII
+characters, so existing configs and checkpoints are unaffected.
+
+---
+
+## 2. NumPy 2.x / imgaug Removal
+
+**Files:**  
+- `strhub/data/augment.py`  
+- `requirements/train.in`, `requirements/train.txt`  
+- `requirements/tune.in`, `requirements/tune.txt`  
+- `requirements/constraints.txt`
+
+### Problem
+`imgaug==0.4.0` uses internal NumPy APIs removed in NumPy 2.0
+(`np.bool`, `np.int`, `np.float`, etc.), causing an `AttributeError` at
+import time on any environment with NumPy ≥ 2.0.
+
+### Fix
+Replaced the three `imgaug` augmenters with equivalent implementations using
+only **NumPy** and **SciPy** (both already required by the project):
+
+| Original (imgaug) | Replacement |
+|---|---|
+| `iaa.MotionBlur(k)` | `scipy.ndimage.uniform_filter1d` along the horizontal axis with a box kernel of width `k` |
+| `iaa.AdditiveGaussianNoise(scale=s)` | `np.random.default_rng().normal(0, s, shape)` added to the image array |
+| `iaa.AdditivePoissonNoise(lam=l)` | `np.random.default_rng().poisson(l, shape)` added to the image array |
+
+All three functions retain the same signature (`(img: PIL.Image, param, **kwargs) -> PIL.Image`)
+so the `timm` RandAugment machinery that calls them is unaffected.
+
+The public API `rand_augment_transform(magnitude, num_layers)` is **unchanged**.
+
+`imgaug` and its imgaug-only transitive dependencies (`imageio` via imgaug,
+`opencv-python` via imgaug, `scikit-image` via imgaug, `shapely`, `tifffile` as
+a scikit-image dep, `lazy-loader`) have been removed from all requirement files.
+`scipy` was already a transitive dependency and remains.
+
+---
+
+## 3. Marathi Dataset Support
+
+**New files:**  
+- `configs/charset/marathi.yaml`  
+- `configs/dataset/marathi.yaml`
+
+### charset — `configs/charset/marathi.yaml`
+Defines `model.charset_train` with the full Devanagari character inventory used
+in written Marathi:
+
+- Independent vowels (स्वर): अ आ इ ई उ ऊ ऋ ए ऐ ओ औ
+- Consonants (व्यंजन): क–ह including ळ क्ष ज्ञ
+- Dependent vowel signs (मात्रा) and halant (्)
+- Diacritics: anusvara (ं), visarga (ः), chandrabindu (ँ), nukta (़)
+- Devanagari digits: ०–९
+- Special symbols: avagraha (ऽ), Om (ॐ)
+
+### dataset — `configs/dataset/marathi.yaml`
+Sets `data.train_dir: marathi` and opts into Unicode-preserving normalisation
+(`normalize_unicode: true`, which now applies NFC for this non-ASCII charset).
+
+### Usage
+```bash
+python train.py charset=marathi dataset=marathi
+```
+
+Expected LMDB layout under `data.root_dir`:
+```
+<root_dir>/
+  train/
+    marathi/          ← data.train_dir
+      data.mdb        ← single LMDB  OR  subdirectories, each with data.mdb
+      lock.mdb
+  val/
+    data.mdb
+    lock.mdb
+  test/
+    data.mdb
+    lock.mdb
+```
+
+---
+
+## 4. SceneTextDataModule — Flat LMDB Layout Support
+
+**File:** `strhub/data/dataset.py`  (`build_tree_dataset`)
+
+### Problem
+The recursive `glob('**/data.mdb')` pattern only matches LMDB stores that are
+**inside** a sub-directory of `root`.  
+When the LMDB lives directly at `root/data.mdb` (the flat layout used by the
+Marathi dataset and many single-split datasets), the glob returns nothing and
+the returned `ConcatDataset` is empty.
+
+### Fix
+After the recursive glob, if no datasets were found and `root/data.mdb` exists,
+the LMDB at `root` itself is opened directly.  
+Both the tree layout and the flat layout are supported transparently without any
+caller changes.
+
+---
+
+## 5. PyTorch Lightning 2.x / PyTorch 2.x Compatibility
+
+**Files modified:** `train.py`, `tune.py`, `strhub/models/base.py`,
+`strhub/models/parseq/system.py`, `strhub/models/abinet/system.py`,
+`strhub/models/crnn/system.py`, `strhub/models/vitstr/system.py`,
+`strhub/models/trba/system.py`
+
+### 5a. `STEP_OUTPUT` removed from all model imports
+
+**Problem:** Every model system file imported `STEP_OUTPUT` from
+`pytorch_lightning.utilities.types`. This internal type alias is not guaranteed
+stable across PL minor versions and triggers deprecation warnings in PL 2.x.
+
+**Fix:** Removed the import entirely. `training_step` return types are now
+annotated as `Tensor` (the real returned value). `validation_step`,
+`test_step`, and `_eval_step` in `base.py` are annotated as
+`Optional[dict[str, Any]]`. No runtime behavior changes.
+
+### 5b. `summarize()` → `ModelSummary` in `train.py`
+
+**Problem:** `pytorch_lightning.utilities.model_summary.summarize()` was
+deprecated in PL 1.8 and removed in PL 2.x.
+
+**Fix:** Replaced with `ModelSummary(model, max_depth=2)` — the canonical PL
+2.x API that produces identical console output.
+
+### 5c. `torch.get_autocast_gpu_dtype()` → `torch.get_autocast_dtype('cuda')`
+
+**Problem:** `torch.get_autocast_gpu_dtype()` was deprecated in PyTorch 2.1.
+The replacement is `torch.get_autocast_dtype('cuda')`.
+
+**Fix:** Added a private helper `_get_autocast_dtype(device_type)` in
+`train.py` that calls the new API first and falls back to the old one for
+PyTorch < 2.1, avoiding warnings on modern installs.
+
+### 5d. Integer precision `16` → string `'16-mixed'` in `tune.py`
+
+**Problem:** `config.trainer.precision = 16` (integer) is deprecated in PL 2.x.
+PL 2.x requires the explicit string `'16-mixed'` to distinguish mixed-precision
+from full `float16`.
+
+**Fix:** Set `config.trainer.precision = '16-mixed'`, consistent with `train.py`.
+
+### 5e. `config.trainer.get('gpus', 0)` → `accelerator == 'gpu'` in `tune.py`
+
+**Problem:** The `gpus` Trainer argument was removed in PL 2.0 in favour of
+`accelerator='gpu'` + `devices=N`. The old key was never populated, so GPU
+precision was silently never set.
+
+**Fix:** Changed the guard to `config.trainer.get('accelerator') == 'gpu'`,
+matching `train.py`.
+
+### 5f. `air.RunConfig(local_dir=…)` → `storage_path=` in `tune.py`
+
+**Problem:** Ray renamed `RunConfig.local_dir` to `storage_path` in Ray 2.7.
+`requirements/tune.txt` pins Ray 2.9.2, so `local_dir` raises a deprecation
+warning and may be removed in a later Ray release.
+
+**Fix:** Replaced `local_dir=str(out_dir.parent.absolute())` with
+`storage_path=str(out_dir.parent.absolute())`.
+
+---
+
+## 6. Safe Multilingual Pretrained Weight Loading
+
+**Files modified:** `strhub/models/utils.py`, `train.py`
+
+### Problem
+`nn.Module.load_state_dict()` in strict mode raises a `RuntimeError` when
+any tensor in the checkpoint has a different shape than the current model.
+This is always the case when fine-tuning on a different charset (e.g. English
+→ Marathi): the classifier/output head dimension equals
+`len(charset) + num_special_tokens`, so it changes with every distinct
+character set. The encoder, decoder, and all other layers are fully reusable.
+
+### New helpers (private)
+
+#### `_extract_state_dict(checkpoint)`
+Detects and unwraps three common checkpoint formats:
+
+| Format | Detection | Action |
+|---|---|---|
+| Plain state dict | No `state_dict` or `model` key | Use as-is |
+| PyTorch Lightning | `"state_dict"` key present | Extract `checkpoint["state_dict"]` |
+| Generic wrapper | `"model"` key present | Extract `checkpoint["model"]` |
+
+#### `_strip_prefix(state_dict, prefixes)`
+Strips leading `model.` or `module.` from checkpoint keys before matching,
+so checkpoints saved from a `LightningModule` wrapper or `DataParallel` model
+align correctly with a bare `nn.Module` state dict without manual key renaming.
+
+### `safe_load_pretrained(model, experiment)` — updated algorithm
+
+1. Download via `get_pretrained_weights(experiment)`.
+2. **Extract** flat state dict via `_extract_state_dict()`.
+3. **Normalise** key prefixes via `_strip_prefix()`.
+4. Compare every checkpoint key against the current model:
+   - shape matches → **loaded**
+   - shape differs → **skipped** (classifier head for cross-charset case)
+   - absent from model → **unexpected**
+5. Collect model keys absent from checkpoint → **missing**.
+6. Merge loaded tensors into a copy of the current state dict.
+7. `load_state_dict(strict=True)` — PyTorch validates the final result.
+8. Print detailed summary; return all four lists.
+
+**Why `strict=False` is not used**
+
+`strict=False` silently ignores *all* missing and unexpected keys, hiding
+typos, refactoring errors, and format changes. Explicit filtering surfaces
+every discrepancy in the printed summary while still producing a model that
+PyTorch considers fully valid (`strict=True` on the filtered dict).
+
+### Return value (updated)
+
+```python
+{
+    "loaded":     [...],   # shape-compatible, copied from checkpoint
+    "skipped":    [...],   # shape mismatch (classifier head)
+    "missing":    [...],   # in model, absent from checkpoint
+    "unexpected": [...],   # in checkpoint, absent from model
+}
+```
+
+### Behaviour by scenario
+
+| Scenario | Loaded | Skipped | Missing | Unexpected |
+|---|---|---|---|---|
+| English → English (same charset) | All | 0 | 0 | 0 |
+| English → Marathi (different charset) | Encoder + decoder | `head.weight`, `head.bias` | 0 | 0 |
+| Old checkpoint (missing new layers) | All compatible | 0 | New layer keys | 0 |
+| Lightning ckpt with `model.` prefix | All (after strip) | shape mismatches only | 0 | 0 |
+
+### Console output (multilingual fine-tuning)
+
+```
+================================================
+           Safe Pretrained Loading
+================================================
+  Checkpoint tensors : 289
+  Model tensors      : 289
+  Loaded             : 287
+  Skipped            : 2
+  Missing            : 0
+  Unexpected         : 0
+  Skipped layers:
+    head.weight
+    head.bias
+================================================
+```
+
+---
+
+## 7. Transfer Learning Support
+
+**Context:** `strhub/models/utils.py`, `train.py`, `configs/charset/marathi.yaml`,
+`configs/dataset/marathi.yaml`
+
+### Overview
+
+The English pretrained PARSeq checkpoint can serve as a strong initialisation
+point for Marathi scene-text recognition, following standard transfer learning
+practice for OCR.
+
+### What transfers
+
+| Component | Transfers? | Reason |
+|---|---|---|
+| ViT encoder (patch embedding + transformer blocks) | ✅ Yes | Learns script-agnostic stroke and spatial features |
+| Autoregressive decoder (attention layers) | ✅ Yes | Sequence modelling is character-system independent |
+| Position embeddings | ✅ Yes | Same image resolution and patch layout |
+| Layer norms, projection layers | ✅ Yes | Same shapes regardless of charset |
+| Classifier head (`head.weight`, `head.bias`) | ❌ No | Output dimension = `len(charset) + specials`; differs per script |
+
+### What is re-initialised
+
+Only the classifier head layers are skipped and re-initialised from scratch
+(random truncated normal, matching `init_weights()`). All other parameters
+start from the English pretrained values and are fine-tuned jointly.
+
+### Training command
+
+```bash
+python train.py \
+  charset=marathi \
+  dataset=marathi \
+  pretrained=parseq
+```
+
+`safe_load_pretrained()` handles the shape mismatch automatically and prints
+a loading summary confirming which layers were transferred and which were skipped.
+
+### Why this works
+
+Transfer learning from Latin-script OCR to Indic scripts is well-established.
+The visual encoder trained on English text images already learns robust
+low-level features (edges, curves, junctions) that are equally relevant for
+Devanagari. Starting from these weights rather than random initialisation
+typically accelerates convergence and improves final accuracy, especially when
+the Marathi training set is smaller than the English one.
+
+---
+
+## 8. Evaluation and Inference Tools
+
+**New files:** `tools/evaluate_marathi.py`, `tools/infer_marathi.py`
+
+No existing files were modified.
+
+### `tools/evaluate_marathi.py`
+
+Evaluates a checkpoint against any LMDB validation split and writes a CSV.
+
+```bash
+python tools/evaluate_marathi.py \
+    --checkpoint outputs/parseq/best.ckpt \
+    --data_root  data \
+    --lmdb_path  data/val \
+    --batch_size 64 \
+    --output     evaluation_predictions.csv
+```
+
+**What it reuses from the repository**
+- `load_from_checkpoint()` — model loading, hyperparameter recovery
+- `LmdbDataset` — LMDB reading, label filtering
+- `SceneTextDataModule.get_transform()` — image preprocessing pipeline
+- `model.tokenizer.decode()` — sequence decoding (no duplication)
+- `model.charset_adapter` — post-decode character filtering
+
+**Metrics** (all NFC-normalised, not NFKD)
+
+| Metric | Definition |
+|---|---|
+| Exact Match Accuracy | `correct / total × 100` |
+| Character Accuracy | `(1 - CER) × 100`, clamped to 0 |
+| CER | `edit_distance(pred, gt) / len(gt)`, averaged |
+| NED | `1 - mean(edit_distance / max(len(pred), len(gt)))` — ICDAR 2019 |
+
+**Output CSV columns:** `index`, `ground_truth`, `prediction`, `confidence`
+
+---
+
+### `tools/infer_marathi.py`
+
+Runs inference on a single image or every image in a folder.
+
+```bash
+# Single image
+python tools/infer_marathi.py --checkpoint best.ckpt --image word.jpg
+
+# Folder
+python tools/infer_marathi.py --checkpoint best.ckpt --folder images/
+```
+
+**Single-image stdout output**
+
+```
+Prediction : मराठी
+Confidence : 0.9231
+```
+
+**Output CSV columns:** `filename`, `prediction`, `confidence`
+
+Supported extensions: `.jpg`, `.jpeg`, `.png`, `.bmp`, `.tiff`, `.webp`.
+Images in a folder are processed in alphabetical order for reproducibility.
+Predictions are NFC-normalised before printing and saving.
+
+---
+
+## 9. Configurable Staged Fine-tuning
+
+**Files modified:** `strhub/models/base.py`, `strhub/models/parseq/system.py`
+**New file:** `configs/experiment/finetune_marathi.yaml`
+
+No changes to the training loop, loss computation, tokenizer, decoding,
+inference, checkpointing, scheduler, validation, or testing.
+
+### Design
+
+All logic lives in `base.py` as module-level helpers and a small extension to
+`BaseSystem.__init__` and `configure_optimizers`. `PARSeq.__init__` receives
+the new options via explicit parameters (before `**kwargs`) and applies them
+after building `self.model`. Every other model (ABINet, CRNN, ViTSTR, TRBA)
+is unaffected — they absorb unknown keys via `**kwargs`.
+
+### New module-level symbols in `strhub/models/base.py`
+
+| Symbol | Kind | Purpose |
+|---|---|---|
+| `_COMPONENT_MAP` | `dict` | Maps config key → `nn.Module` attribute name |
+| `_BACKBONE_COMPONENTS` | `tuple` | `('encoder', 'decoder')` — share `backbone_lr` |
+| `_HEAD_COMPONENTS` | `tuple` | `('head', 'text_embed')` — share `head_lr` |
+| `_apply_freeze(inner, cfg)` | function | Sets `requires_grad=False` on frozen submodules |
+| `_print_finetune_summary(...)` | function | Prints the one-time startup config table |
+
+### Changes to `BaseSystem`
+
+- `__init__` accepts `backbone_lr=None` and `head_lr=None`. When both are
+  `None` (the default) behaviour is **identical** to the original.
+- `configure_optimizers` branches on whether both differential LRs are set:
+  - **Not set (default):** original single-LR path via `create_optimizer_v2` — unchanged.
+  - **Set:** three `AdamW` parameter groups (backbone, head, other), each with
+    its own scaled LR. `OneCycleLR` receives a list of `max_lr` values so every
+    group warms up to its own peak independently.
+- `_lr_scale()` helper factored out to avoid repeating the DDP scaling formula.
+
+### Changes to `CrossEntropySystem` / `CTCSystem`
+
+Both now forward `backbone_lr` and `head_lr` to `BaseSystem`. Signatures
+remain fully backward-compatible (both kwargs default to `None`).
+
+### Changes to `PARSeq.__init__`
+
+- `freeze: Optional[dict] = None`, `backbone_lr`, `head_lr` added as explicit
+  parameters before `**kwargs`.
+- After `self.model` is built, `_apply_freeze(self.model, self._freeze_cfg)` is
+  called — no-op when `freeze` is absent or all-false.
+- `on_train_start` prints the fine-tuning summary on rank 0 only (no duplicate
+  output in DDP multi-GPU runs).
+
+### New config: `configs/experiment/finetune_marathi.yaml`
+
+Ready-to-use preset:
+
+```bash
+python train.py \
+  experiment=finetune_marathi \
+  charset=marathi \
+  dataset=marathi \
+  pretrained=parseq
+```
+
+Defaults: encoder frozen, decoder/head/text_embed trainable,
+`backbone_lr=1e-5`, `head_lr=5e-4`.
+
+### Startup summary example
+
+```
+================================================
+       Fine-tuning Configuration
+================================================
+  Encoder     : Frozen
+  Decoder     : Trainable
+  Head        : Trainable
+  Text Embed  : Trainable
+  Backbone LR : 1e-05
+  Head LR     : 0.0005
+================================================
+```
+
+### Backward compatibility
+
+All three experiment variants work via Hydra CLI overrides with zero code changes:
+
+```bash
+# Experiment 1 — freeze encoder + decoder, differential LRs
+python train.py experiment=finetune_marathi \
+  model.freeze.encoder=true model.freeze.decoder=true \
+  model.backbone_lr=1e-5 model.head_lr=5e-4
+
+# Experiment 2 — freeze encoder only
+python train.py experiment=finetune_marathi \
+  model.freeze.encoder=true model.freeze.decoder=false
+
+# Experiment 3 — train everything (original behaviour unchanged)
+python train.py charset=marathi dataset=marathi pretrained=parseq
+```
+
+---
+
+## Summary of Modified Files
+
+| File | Change |
+|---|---|
+| `strhub/data/dataset.py` | Unicode auto-detection; flat LMDB layout support |
+| `strhub/data/augment.py` | imgaug replaced with NumPy/SciPy; public API unchanged |
+| `train.py` | `summarize` → `ModelSummary`; autocast compat helper; `safe_load_pretrained` |
+| `tune.py` | `gpus` → `accelerator` check; integer precision → `'16-mixed'`; `local_dir` → `storage_path` |
+| `strhub/models/utils.py` | Added `_extract_state_dict()`, `_strip_prefix()`, `safe_load_pretrained()`; updated `create_model()` |
+| `strhub/models/base.py` | `STEP_OUTPUT` removed; staged fine-tuning helpers; `configure_optimizers` extension |
+| `strhub/models/parseq/system.py` | `STEP_OUTPUT` removed; `freeze`/`backbone_lr`/`head_lr` params; `on_train_start` |
+| `strhub/models/abinet/system.py` | Removed `STEP_OUTPUT`; `training_step` → `Tensor` |
+| `strhub/models/crnn/system.py` | Removed `STEP_OUTPUT`; `training_step` → `Tensor` |
+| `strhub/models/vitstr/system.py` | Removed `STEP_OUTPUT`; `training_step` → `Tensor` |
+| `strhub/models/trba/system.py` | Removed `STEP_OUTPUT`; both `training_step`s → `Tensor` |
+| `requirements/train.in` | Removed `imgaug` |
+| `requirements/train.txt` | Removed `imgaug` and imgaug-only transitive deps |
+| `requirements/tune.in` | Removed `imgaug` |
+| `requirements/tune.txt` | Removed `imgaug` and imgaug-only transitive deps |
+| `requirements/constraints.txt` | Removed `imgaug` entry; updated `# via` annotations |
+| `configs/charset/marathi.yaml` | **New** — Devanagari charset for Marathi |
+| `configs/dataset/marathi.yaml` | **New** — Marathi dataset config |
+| `configs/experiment/finetune_marathi.yaml` | **New** — staged fine-tuning preset |
+| `tools/evaluate_marathi.py` | **New** — LMDB evaluation with metrics + CSV export |
+| `tools/infer_marathi.py` | **New** — single-image and folder inference + CSV export |
+| `tests/__init__.py` | **New** — test package marker |
+| `tests/test_safe_pretrained.py` | **New** — 33 offline tests for `safe_load_pretrained` |
+| `PATCHES.md` | **New** — this file |
diff --git a/README.md b/README.md
index 925555436..a100aa699 100644
--- a/README.md
+++ b/README.md
@@ -10,29 +10,26 @@
 <div align="center">
 
 # Scene Text Recognition with<br/>Permuted Autoregressive Sequence Models
+### Extended for Marathi (Devanagari) Scene Text Recognition
+
 [![Apache License 2.0](https://img.shields.io/github/license/baudm/parseq)](https://github.com/baudm/parseq/blob/main/LICENSE)
 [![arXiv preprint](http://img.shields.io/badge/arXiv-2207.06966-b31b1b)](https://arxiv.org/abs/2207.06966)
 [![In Proc. ECCV 2022](http://img.shields.io/badge/ECCV-2022-6790ac)](https://www.ecva.net/papers/eccv_2022/papers_ECCV/html/556_ECCV_2022_paper.php)
 [![Gradio demo](https://img.shields.io/badge/%F0%9F%A4%97%20demo-Gradio-ff7c00)](https://huggingface.co/spaces/baudm/PARSeq-OCR)
 
-[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/scene-text-recognition-with-permuted/scene-text-recognition-on-coco-text)](https://paperswithcode.com/sota/scene-text-recognition-on-coco-text?p=scene-text-recognition-with-permuted)
-[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/scene-text-recognition-with-permuted/scene-text-recognition-on-ic19-art)](https://paperswithcode.com/sota/scene-text-recognition-on-ic19-art?p=scene-text-recognition-with-permuted)
-[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/scene-text-recognition-with-permuted/scene-text-recognition-on-icdar2013)](https://paperswithcode.com/sota/scene-text-recognition-on-icdar2013?p=scene-text-recognition-with-permuted)
-[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/scene-text-recognition-with-permuted/scene-text-recognition-on-iiit5k)](https://paperswithcode.com/sota/scene-text-recognition-on-iiit5k?p=scene-text-recognition-with-permuted)
-[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/scene-text-recognition-with-permuted/scene-text-recognition-on-cute80)](https://paperswithcode.com/sota/scene-text-recognition-on-cute80?p=scene-text-recognition-with-permuted)
-[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/scene-text-recognition-with-permuted/scene-text-recognition-on-icdar2015)](https://paperswithcode.com/sota/scene-text-recognition-on-icdar2015?p=scene-text-recognition-with-permuted)
-[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/scene-text-recognition-with-permuted/scene-text-recognition-on-svt)](https://paperswithcode.com/sota/scene-text-recognition-on-svt?p=scene-text-recognition-with-permuted)
-[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/scene-text-recognition-with-permuted/scene-text-recognition-on-svtp)](https://paperswithcode.com/sota/scene-text-recognition-on-svtp?p=scene-text-recognition-with-permuted)
-
 [**Darwin Bautista**](https://github.com/baudm) and [**Rowel Atienza**](https://github.com/roatienza)
 
 Electrical and Electronics Engineering Institute<br/>
 University of the Philippines, Diliman
 
-[Method](#method-tldr) | [Sample Results](#sample-results) | [Getting Started](#getting-started) | [FAQ](#frequently-asked-questions) | [Training](#training) | [Evaluation](#evaluation) | [Citation](#citation)
+[Method](#method-tldr) | [Marathi Extension](#marathi-extension) | [Getting Started](#getting-started) | [Training](#training) | [Marathi Training](#marathi-training) | [Evaluation](#evaluation) | [Inference](#inference) | [FAQ](#frequently-asked-questions) | [Citation](#citation)
 
 </div>
 
+> **This fork** extends the official PARSeq repository with full Marathi (Devanagari) scene-text recognition support, NumPy 2.x / PyTorch Lightning 2.x compatibility fixes, safe multilingual weight loading, staged fine-tuning, and evaluation/inference tools. The model architecture, tokenizer, and training loop are **unchanged**. See [`PATCHES.md`](PATCHES.md) for a complete change log.
+
+---
+
 Scene Text Recognition (STR) models use language context to be more robust against noisy or corrupted images. Recent approaches like ABINet use a standalone or external Language Model (LM) for prediction refinement. In this work, we show that the external LM&mdash;which requires upfront allocation of dedicated compute capacity&mdash;is inefficient for STR due to its poor performance vs cost characteristics. We propose a more efficient approach using **p**ermuted **a**uto**r**egressive **seq**uence (PARSeq) models. View our ECCV [poster](https://drive.google.com/file/d/19luOT_RMqmafLMhKQQHBnHNXV7fOCRfw/view) and [presentation](https://drive.google.com/file/d/11VoZW4QC5tbMwVIjKB44447uTiuCJAAD/view) for a brief overview.
 
 ![PARSeq](.github/gh-teaser.png)
@@ -47,8 +44,6 @@ Our main insight is that with an ensemble of autoregressive (AR) models, we coul
 A single Transformer can realize different models by merely varying its attention mask. With the correct decoder parameterization, it can be trained with Permutation Language Modeling to enable inference for arbitrary output positions given arbitrary subsets of the input context. This *arbitrary decoding* characteristic results in a _unified_ STR model&mdash;PARSeq&mdash;capable of context-free and context-aware inference, as well as iterative prediction refinement using bidirectional context **without** requiring a standalone language model. PARSeq can be considered an ensemble of AR models with shared architecture and weights:
 
 ![System](.github/system.png)
-**NOTE:** _LayerNorm and Dropout layers are omitted. `[B]`, `[E]`, and `[P]` stand for beginning-of-sequence (BOS), end-of-sequence (EOS), and padding tokens, respectively. `T` = 25 results in 26 distinct position tokens. The position tokens both serve as query vectors and position embeddings for the input context. For `[B]`, no position embedding is added. Attention
-masks are generated from the given permutations and are used only for the context-position attention. L<sub>ce</sub> pertains to the cross-entropy loss._
 
 ### Sample Results
 <div align="center">
@@ -65,205 +60,375 @@ masks are generated from the given permutations and are used only for the contex
 **NOTE:** _Bold letters and underscores indicate wrong and missing character predictions, respectively._
 </div>
 
+---
+
+## Marathi Extension
+
+This fork adds first-class Marathi (Devanagari) support to PARSeq. All changes
+are infrastructure-only — the model architecture, tokenizer, and training loop
+are identical to the upstream repository.
+
+### What was added
+
+| Area | Change |
+|---|---|
+| **Unicode** | NFC normalisation (not NFKD) for non-ASCII charsets; English datasets unchanged |
+| **imgaug removal** | Replaced with NumPy/SciPy equivalents compatible with NumPy 2.x |
+| **LMDB layout** | Flat `root/data.mdb` layout supported alongside the existing tree layout |
+| **Charset** | `configs/charset/marathi.yaml` — full Devanagari inventory |
+| **Dataset config** | `configs/dataset/marathi.yaml` |
+| **Safe weight loading** | `safe_load_pretrained()` skips mismatched classifier layers for cross-charset transfer |
+| **Staged fine-tuning** | Per-component freeze flags + differential learning rates via Hydra config |
+| **Evaluation tool** | `tools/evaluate_marathi.py` — Exact Match, CER, NED, CSV export |
+| **Inference tool** | `tools/infer_marathi.py` — single image or folder, CSV export |
+| **PL 2.x compat** | Removed deprecated `STEP_OUTPUT`, `summarize()`, integer precision, `gpus` flag |
+
+### Charset
+
+The Marathi charset covers the full Devanagari Unicode block used in written Marathi:
+vowels (अ–औ), consonants (क–ह, ळ, क्ष, ज्ञ), dependent vowel signs (मात्रा),
+halant (्), diacritics (anusvara ं, visarga ः, chandrabindu ँ), Devanagari digits (०–९),
+and common punctuation.
+
+---
+
 ## Getting Started
-This repository contains the reference implementation for PARSeq and reproduced models (collectively referred to as _Scene Text Recognition Model Hub_). See `NOTICE` for copyright information.
-Majority of the code is licensed under the Apache License v2.0 (see `LICENSE`) while ABINet and CRNN sources are
-released under the BSD and MIT licenses, respectively (see corresponding `LICENSE` files for details).
 
-### Demo
-An [interactive Gradio demo](https://huggingface.co/spaces/baudm/PARSeq-OCR) hosted at Hugging Face is available. The pretrained weights released here are used for the demo.
+Requires Python ≥ 3.9 and PyTorch ≥ 2.0.
 
-### Installation
-Requires Python >= 3.9 and PyTorch >= 2.0. The default requirements files will install the latest versions of the dependencies (as of February 22, 2024).
 ```bash
 # Use specific platform build. Other PyTorch 2.0 options: cu118, cu121, rocm5.7
 platform=cpu
-# Generate requirements files for specified PyTorch platform
 make torch-${platform}
-# Install the project and core + train + test dependencies. Subsets: [dev,train,test,bench,tune]
 pip install -r requirements/core.${platform}.txt -e .[train,test]
- ```
+```
+
 #### Updating dependency version pins
 ```bash
 pip install pip-tools
-make clean-reqs reqs  # Regenerate all the requirements files
- ```
+make clean-reqs reqs
+```
+
 ### Datasets
+
 Download the [datasets](Datasets.md) from the following links:
 1. [LMDB archives](https://drive.google.com/drive/folders/1NYuoi7dfJVgo-zUJogh8UQZgIMpLviOE) for MJSynth, SynthText, IIIT5k, SVT, SVTP, IC13, IC15, CUTE80, ArT, RCTW17, ReCTS, LSVT, MLT19, COCO-Text, and Uber-Text.
 2. [LMDB archives](https://drive.google.com/drive/folders/1D9z_YJVa6f-O0juni-yG5jcwnhvYw-qC) for TextOCR and OpenVINO.
 
+For Marathi datasets, the expected LMDB layout under `data.root_dir` is:
+
+```
+data/
+  train/
+    marathi/
+      data.mdb        ← single LMDB or subdirectories each with data.mdb
+      lock.mdb
+  val/
+    data.mdb
+    lock.mdb
+  test/
+    data.mdb
+    lock.mdb
+```
+
 ### Pretrained Models via Torch Hub
-Available models are: `abinet`, `crnn`, `trba`, `vitstr`, `parseq_tiny`, `parseq_patch16_224`, and `parseq`.
+
 ```python
 import torch
 from PIL import Image
 from strhub.data.module import SceneTextDataModule
 
-# Load model and image transforms
 parseq = torch.hub.load('baudm/parseq', 'parseq', pretrained=True).eval()
 img_transform = SceneTextDataModule.get_transform(parseq.hparams.img_size)
 
 img = Image.open('/path/to/image.png').convert('RGB')
-# Preprocess. Model expects a batch of images with shape: (B, C, H, W)
 img = img_transform(img).unsqueeze(0)
 
 logits = parseq(img)
-logits.shape  # torch.Size([1, 26, 95]), 94 characters + [EOS] symbol
-
-# Greedy decoding
 pred = logits.softmax(-1)
 label, confidence = parseq.tokenizer.decode(pred)
 print('Decoded label = {}'.format(label[0]))
 ```
 
-## Frequently Asked Questions
-- How do I train on a new language? See Issues [#5](https://github.com/baudm/parseq/issues/5) and [#9](https://github.com/baudm/parseq/issues/9).
-- Can you export to TorchScript or ONNX? Yes, see Issue [#12](https://github.com/baudm/parseq/issues/12#issuecomment-1267842315).
-- How do I test on my own dataset? See Issue [#27](https://github.com/baudm/parseq/issues/27).
-- How do I finetune and/or create a custom dataset? See Issue [#7](https://github.com/baudm/parseq/issues/7).
-- What is `val_NED`? See Issue [#10](https://github.com/baudm/parseq/issues/10).
+---
 
 ## Training
-The training script can train any supported model. You can override any configuration using the command line. Please refer to [Hydra](https://hydra.cc) docs for more info about the syntax. Use `./train.py --help` to see the default configuration.
 
-<details><summary>Sample commands for different training configurations</summary><p>
+The training script can train any supported model. Use `./train.py --help` to see the default configuration.
+
+<details><summary>Sample commands for standard training</summary><p>
 
 ### Finetune using pretrained weights
 ```bash
-./train.py +experiment=parseq-tiny pretrained=parseq-tiny  # Not all experiments have pretrained weights
+./train.py +experiment=parseq-tiny pretrained=parseq-tiny
 ```
 
-### Train a model variant/preconfigured experiment
-The base model configurations are in `configs/model/`, while variations are stored in `configs/experiment/`.
+### Train a model variant
 ```bash
-./train.py +experiment=parseq-tiny  # Some examples: abinet-sv, trbc
+./train.py +experiment=parseq-tiny
 ```
 
-### Specify the character set for training
+### Specify the character set
 ```bash
-./train.py charset=94_full  # Other options: 36_lowercase or 62_mixed-case. See configs/charset/
+./train.py charset=94_full  # Other options: 36_lowercase, 62_mixed-case
 ```
 
 ### Specify the training dataset
 ```bash
-./train.py dataset=real  # Other option: synth. See configs/dataset/
+./train.py dataset=real  # Other option: synth
 ```
 
-### Change general model training parameters
+### Change model parameters
 ```bash
 ./train.py model.img_size=[32, 128] model.max_label_length=25 model.batch_size=384
 ```
 
-### Change data-related training parameters
+### Change data parameters
 ```bash
 ./train.py data.root_dir=data data.num_workers=2 data.augment=true
 ```
 
-### Change `pytorch_lightning.Trainer` parameters
+### Change Trainer parameters
 ```bash
 ./train.py trainer.max_epochs=20 trainer.accelerator=gpu trainer.devices=2
 ```
-Note that you can pass any [Trainer parameter](https://pytorch-lightning.readthedocs.io/en/stable/common/trainer.html),
-you just need to prefix it with `+` if it is not originally specified in `configs/main.yaml`.
 
-### Resume training from checkpoint (experimental)
+### Resume from checkpoint
 ```bash
 ./train.py +experiment=<model_exp> ckpt_path=outputs/<model>/<timestamp>/checkpoints/<checkpoint>.ckpt
 ```
 
 </p></details>
 
+---
+
+## Marathi Training
+
+### Quick start — fine-tune from English pretrained weights
+
+```bash
+python train.py \
+  experiment=finetune_marathi \
+  charset=marathi \
+  dataset=marathi \
+  pretrained=parseq
+```
+
+`safe_load_pretrained()` automatically loads all compatible layers (encoder,
+decoder, position embeddings) and skips only the classifier head, which is
+re-initialised for the Marathi vocab size. A loading summary is printed:
+
+```
+================================================
+        Safe Pretrained Loading
+================================================
+  Checkpoint tensors : 289
+  Model tensors      : 289
+  Loaded             : 287
+  Skipped            : 2
+  Skipped layers:
+    head.weight
+    head.bias
+================================================
+```
+
+### Staged fine-tuning experiments
+
+The `finetune_marathi` experiment config supports three common transfer learning
+strategies via Hydra CLI overrides — no code changes needed.
+
+**Experiment 1 — freeze encoder + decoder, train head only with differential LRs**
+```bash
+python train.py experiment=finetune_marathi charset=marathi dataset=marathi \
+  pretrained=parseq \
+  model.freeze.encoder=true \
+  model.freeze.decoder=true \
+  model.backbone_lr=1e-5 \
+  model.head_lr=5e-4
+```
+
+**Experiment 2 — freeze encoder only**
+```bash
+python train.py experiment=finetune_marathi charset=marathi dataset=marathi \
+  pretrained=parseq \
+  model.freeze.encoder=true \
+  model.freeze.decoder=false
+```
+
+**Experiment 3 — train all layers (standard fine-tuning)**
+```bash
+python train.py charset=marathi dataset=marathi pretrained=parseq
+```
+
+A fine-tuning summary is printed at the start of training:
+
+```
+================================================
+       Fine-tuning Configuration
+================================================
+  Encoder     : Frozen
+  Decoder     : Trainable
+  Head        : Trainable
+  Text Embed  : Trainable
+  Backbone LR : 1e-05
+  Head LR     : 0.0005
+================================================
+```
+
+### Freeze configuration reference
+
+```yaml
+# In your experiment config or as CLI overrides
+model:
+  freeze:
+    encoder:    true   # freeze the ViT backbone
+    decoder:    false  # keep decoder trainable
+    head:       false  # always train (re-initialised for new charset)
+    text_embed: false  # always train (new vocab embeddings)
+
+  backbone_lr: 1.0e-5   # LR for encoder + decoder (when not null)
+  head_lr:     5.0e-4   # LR for head + text_embed (when not null)
+```
+
+If `backbone_lr` and `head_lr` are both `null` (the default), the standard
+single learning-rate optimizer is used unchanged.
+
+---
+
 ## Evaluation
-The test script, ```test.py```, can be used to evaluate any model trained with this project. For more info, see ```./test.py --help```.
 
-PARSeq runtime parameters can be passed using the format `param:type=value`. For example, PARSeq NAR decoding can be invoked via `./test.py parseq.ckpt refine_iters:int=2 decode_ar:bool=false`.
+### Standard English benchmark evaluation
+
+```bash
+./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt
+# or
+./test.py pretrained=parseq
+```
 
-<details><summary>Sample commands for reproducing results</summary><p>
+### Marathi validation set evaluation
 
-### Lowercase alphanumeric comparison on benchmark datasets (Table 6)
 ```bash
-./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt  # or use the released weights: ./test.py pretrained=parseq
+python tools/evaluate_marathi.py \
+  --checkpoint outputs/parseq/<timestamp>/checkpoints/best.ckpt \
+  --lmdb_path  data/val \
+  --batch_size 64 \
+  --output     evaluation_predictions.csv
 ```
+
 **Sample output:**
-| Dataset   | # samples | Accuracy | 1 - NED | Confidence | Label Length |
-|:---------:|----------:|---------:|--------:|-----------:|-------------:|
-| IIIT5k    |      3000 |    99.00 |   99.79 |      97.09 |         5.09 |
-| SVT       |       647 |    97.84 |   99.54 |      95.87 |         5.86 |
-| IC13_1015 |      1015 |    98.13 |   99.43 |      97.19 |         5.31 |
-| IC15_2077 |      2077 |    89.22 |   96.43 |      91.91 |         5.33 |
-| SVTP      |       645 |    96.90 |   99.36 |      94.37 |         5.86 |
-| CUTE80    |       288 |    98.61 |   99.80 |      96.43 |         5.53 |
-| **Combined** | **7672** | **95.95** | **98.78** | **95.34** | **5.33** |
---------------------------------------------------------------------------
-
-### Benchmark using different evaluation character sets (Table 4)
+```
+================================================
+         Evaluation Results
+================================================
+  Samples          : 5000
+  Exact Match Acc  : 82.34 %
+  Character Acc    : 96.12 %
+  CER              :  3.88 %
+  NED              : 94.57 %
+================================================
+Predictions saved to: evaluation_predictions.csv
+```
+
+Output CSV columns: `index`, `ground_truth`, `prediction`, `confidence`.
+
+All string comparisons use NFC normalisation (not NFKD) to correctly handle
+Devanagari matras and composed characters.
+
+<details><summary>Standard benchmark commands</summary><p>
+
+### Lowercase alphanumeric comparison (Table 6)
+```bash
+./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt
+```
+
+### Mixed-case and punctuation
 ```bash
-./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt  # lowercase alphanumeric (36-character set)
-./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased  # mixed-case alphanumeric (62-character set)
-./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased --punctuation  # mixed-case alphanumeric + punctuation (94-character set)
+./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased
+./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased --punctuation
 ```
 
-### Lowercase alphanumeric comparison on more challenging datasets (Table 5)
+### New benchmark datasets (Table 5)
 ```bash
 ./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --new
 ```
 
-### Benchmark Model Compute Requirements (Figure 5)
+### Benchmark compute requirements
 ```bash
 ./bench.py model=parseq model.decode_ar=false model.refine_iters=3
-<torch.utils.benchmark.utils.common.Measurement object at 0x7f8fcae67ee0>
-model(x)
-  Median: 14.87 ms
-  IQR:    0.33 ms (14.78 to 15.12)
-  7 measurements, 10 runs per measurement, 1 thread
-| module                | #parameters   | #flops   | #activations   |
-|:----------------------|:--------------|:---------|:---------------|
-| model                 | 23.833M       | 3.255G   | 8.214M         |
-|  encoder              |  21.381M      |  2.88G   |  7.127M        |
-|  decoder              |  2.368M       |  0.371G  |  1.078M        |
-|  head                 |  36.575K      |  3.794M  |  9.88K         |
-|  text_embed.embedding |  37.248K      |  0       |  0             |
-```
-
-### Latency Measurements vs Output Label Length (Appendix I)
-```bash
-./bench.py model=parseq model.decode_ar=false model.refine_iters=3 +range=true
 ```
 
-### Orientation robustness benchmark (Appendix J)
+### Orientation robustness
 ```bash
-./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased --punctuation  # no rotation
-./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased --punctuation --rotation 90
-./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased --punctuation --rotation 180
-./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --cased --punctuation --rotation 270
+./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --rotation 90
+./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --rotation 180
+./test.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --rotation 270
 ```
 
-### Using trained models to read text from images (Appendix L)
+</p></details>
+
+---
+
+## Inference
+
+### Original inference script (English)
+
 ```bash
-./read.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --images demo_images/*  # Or use ./read.py pretrained=parseq
-Additional keyword arguments: {}
-demo_images/art-01107.jpg: CHEWBACCA
-demo_images/coco-1166773.jpg: Chevrol
-demo_images/cute-184.jpg: SALMON
-demo_images/ic13_word_256.png: Verbandsteffe
-demo_images/ic15_word_26.png: Kaopa
-demo_images/uber-27491.jpg: 3rdAve
-
-# use NAR decoding + 2 refinement iterations for PARSeq
+./read.py outputs/<model>/<timestamp>/checkpoints/last.ckpt --images demo_images/*
 ./read.py pretrained=parseq refine_iters:int=2 decode_ar:bool=false --images demo_images/*
 ```
-</p></details>
+
+### Marathi inference tool
+
+**Single image**
+```bash
+python tools/infer_marathi.py \
+  --checkpoint best.ckpt \
+  --image      word.jpg
+```
+
+Output:
+```
+Prediction : मराठी
+Confidence : 0.9231
+```
+
+**Folder of images**
+```bash
+python tools/infer_marathi.py \
+  --checkpoint best.ckpt \
+  --folder     images/ \
+  --output     predictions.csv
+```
+
+Output CSV columns: `filename`, `prediction`, `confidence`.
+
+Supported image formats: `.jpg`, `.jpeg`, `.png`, `.bmp`, `.tiff`, `.webp`.
+Images in a folder are processed in alphabetical order.
+
+---
 
 ## Tuning
 
-We use [Ray Tune](https://www.ray.io/ray-tune) for automated parameter tuning of the learning rate. See `./tune.py --help`. Extend `tune.py` to support tuning of other hyperparameters.
 ```bash
-./tune.py tune.num_samples=20  # find optimum LR for PARSeq's default config using 20 trials
-./tune.py +experiment=tune_abinet-lm  # find the optimum learning rate for ABINet's language model
+./tune.py tune.num_samples=20
+./tune.py +experiment=tune_abinet-lm
 ```
 
+---
+
+## Frequently Asked Questions
+
+- How do I train on a new language? See Issues [#5](https://github.com/baudm/parseq/issues/5) and [#9](https://github.com/baudm/parseq/issues/9). For Marathi specifically, see the [Marathi Training](#marathi-training) section above.
+- Can you export to TorchScript or ONNX? Yes, see Issue [#12](https://github.com/baudm/parseq/issues/12#issuecomment-1267842315).
+- How do I test on my own dataset? See Issue [#27](https://github.com/baudm/parseq/issues/27).
+- How do I finetune a custom dataset? See Issue [#7](https://github.com/baudm/parseq/issues/7) and the [Staged Fine-tuning](#staged-fine-tuning-experiments) section above.
+- What is `val_NED`? See Issue [#10](https://github.com/baudm/parseq/issues/10).
+- Why does loading a pretrained checkpoint fail with a shape mismatch? The classifier head size depends on the charset. Use `safe_load_pretrained()` (called automatically via the `pretrained=` flag) which skips incompatible layers instead of failing.
+
+---
+
 ## Citation
+
 ```bibtex
 @InProceedings{bautista2022parseq,
   title={Scene Text Recognition with Permuted Autoregressive Sequence Models},
diff --git a/configs/charset/marathi.yaml b/configs/charset/marathi.yaml
new file mode 100644
index 000000000..8b7a26c4f
--- /dev/null
+++ b/configs/charset/marathi.yaml
@@ -0,0 +1,5 @@
+# @package _global_
+
+model:
+  charset_train: ",-./:|ँंःअआइईउऊएऐऑओऔकखगघचछजझञटठडढणतथदधनपफबभमयरऱलळवशषसह़ािीुूृॅेैॉोौ्०१२३४५६७८९ॲ"
+  charset_test: ",-./:|ँंःअआइईउऊएऐऑओऔकखगघचछजझञटठडढणतथदधनपफबभमयरऱलळवशषसह़ािीुूृॅेैॉोौ्०१२३४५६७८९ॲ"
diff --git a/configs/dataset/marathi.yaml b/configs/dataset/marathi.yaml
new file mode 100644
index 000000000..a58b48635
--- /dev/null
+++ b/configs/dataset/marathi.yaml
@@ -0,0 +1,24 @@
+# @package _global_
+# Marathi scene-text dataset configuration.
+#
+# Expected LMDB layout (relative to data.root_dir):
+#   train/
+#     data.mdb        <- single LMDB  OR
+#     <subset>/       <- one LMDB per subset (build_tree_dataset handles both)
+#       data.mdb
+#   val/
+#     data.mdb
+#   test/
+#     data.mdb
+#
+# Override normalize_unicode=false if your LMDB labels are already NFC-normalised
+# (the dataset loader applies NFC automatically for non-ASCII charsets, so this
+#  is safe to leave at true).
+model:
+  charset_train: ",-./:|ँंःअआइईउऊएऐऑओऔकखगघचछजझञटठडढणतथदधनपफबभमयरऱलळवशषसह़ािीुूृॅेैॉोौ्०१२३४५६७८९ॲ"
+  charset_test: ",-./:|ँंःअआइईउऊएऐऑओऔकखगघचछजझञटठडढणतथदधनपफबभमयरऱलळवशषसह़ािीुूृॅेैॉोौ्०१२३४५६७८९ॲ"
+
+data:
+  train_dir: "."
+  remove_whitespace: true
+  normalize_unicode: true
diff --git a/configs/experiment/finetune_marathi.yaml b/configs/experiment/finetune_marathi.yaml
new file mode 100644
index 000000000..0e4d89890
--- /dev/null
+++ b/configs/experiment/finetune_marathi.yaml
@@ -0,0 +1,40 @@
+# @package _global_
+# Staged fine-tuning configuration for Marathi (Devanagari) OCR.
+#
+# Experiment: freeze the visual encoder, train everything else with differential LRs.
+# To run:
+#   python train.py \
+#     experiment=finetune_marathi \
+#     charset=marathi \
+#     dataset=marathi \
+#     pretrained=parseq
+#
+# Adjust freeze flags and learning rates per experiment:
+#
+#   Experiment 1 — freeze encoder + decoder, train head only:
+#     model.freeze.encoder=true  model.freeze.decoder=true
+#     model.backbone_lr=1e-5     model.head_lr=5e-4
+#
+#   Experiment 2 — freeze encoder only:
+#     model.freeze.encoder=true
+#
+#   Experiment 3 — train everything (default behaviour, identical to no freeze):
+#     (just omit any freeze/lr overrides)
+defaults:
+  - override /model: parseq
+
+model:
+  # --- Layer freezing ---
+  # Set a component to true to freeze it (requires_grad=False).
+  # Omitting this section entirely preserves the original trainable behaviour.
+  freeze:
+    encoder:    true   # freeze the ViT backbone — reuse English visual features
+    decoder:    false  # keep decoder trainable for Marathi sequence modelling
+    head:       false  # always train the head (it was re-initialised for Marathi)
+    text_embed: false  # always train embeddings for the new charset
+
+  # --- Differential learning rates ---
+  # encoder + decoder use backbone_lr; head + text_embed use head_lr.
+  # Set both to null (or omit entirely) to use the single global lr instead.
+  backbone_lr: 1.0e-5
+  head_lr:     5.0e-4
diff --git a/configs/main.yaml b/configs/main.yaml
index fe890d345..5f63dd4db 100644
--- a/configs/main.yaml
+++ b/configs/main.yaml
@@ -3,6 +3,7 @@ defaults:
   - model: parseq
   - charset: 94_full
   - dataset: real
+  - experiment: null
 
 model:
   _convert_: all
diff --git a/requirements/constraints.txt b/requirements/constraints.txt
index b2b4887fe..30a05e720 100644
--- a/requirements/constraints.txt
+++ b/requirements/constraints.txt
@@ -84,13 +84,7 @@ idna==3.6
     #   requests
     #   yarl
 imageio==2.34.0
-    # via
-    #   imgaug
-    #   scikit-image
-imgaug==0.4.0
-    # via
-    #   -r requirements/train.in
-    #   -r requirements/tune.in
+    # via scikit-image
 importlib-resources==6.1.1
     # via matplotlib
 iopath==0.1.10
@@ -137,7 +131,7 @@ lmdb==1.4.1
 markupsafe==2.1.5
     # via jinja2
 matplotlib==3.8.3
-    # via imgaug
+    # via ax-platform
 matplotlib-inline==0.1.6
     # via ipython
 mpmath==1.3.0
@@ -153,9 +147,7 @@ multipledispatch==1.0.0
 mypy-extensions==1.0.0
     # via typing-inspect
 networkx==3.2.1
-    # via
-    #   scikit-image
-    #   torch
+    # via torch
 nltk==3.8.1
     # via -r requirements/core.in
 nodeenv==1.8.0
@@ -164,28 +156,22 @@ numpy==1.26.4
     # via
     #   contourpy
     #   fvcore
-    #   imageio
-    #   imgaug
     #   jaxtyping
     #   matplotlib
-    #   opencv-python
     #   opt-einsum
     #   pandas
     #   pyarrow
     #   pyro-ppl
     #   pytorch-lightning
-    #   scikit-image
     #   scikit-learn
     #   scipy
-    #   shapely
     #   tensorboardx
-    #   tifffile
     #   torchmetrics
     #   torchvision
 omegaconf==2.3.0
     # via hydra-core
 opencv-python==4.9.0.80
-    # via imgaug
+    # via -r requirements/bench.in
 opt-einsum==3.3.0
     # via pyro-ppl
 packaging==23.2
@@ -197,7 +183,6 @@ packaging==23.2
     #   plotly
     #   pytorch-lightning
     #   ray
-    #   scikit-image
     #   tensorboardx
     #   torchmetrics
 pandas==2.2.0
@@ -214,10 +199,7 @@ pillow==10.2.0
     #   -r requirements/train.in
     #   -r requirements/tune.in
     #   fvcore
-    #   imageio
-    #   imgaug
     #   matplotlib
-    #   scikit-image
     #   torchvision
 platformdirs==4.2.0
     # via virtualenv
@@ -286,8 +268,10 @@ rpds-py==0.18.0
     #   referencing
 safetensors==0.4.2
     # via timm
+lazy-loader==0.3
+    # via scikit-image
 scikit-image==0.22.0
-    # via imgaug
+    # via -r requirements/bench.in
 scikit-learn==1.4.1.post1
     # via
     #   ax-platform
@@ -296,16 +280,11 @@ scipy==1.12.0
     # via
     #   ax-platform
     #   botorch
-    #   imgaug
     #   linear-operator
-    #   scikit-image
     #   scikit-learn
-shapely==2.0.3
-    # via imgaug
 six==1.16.0
     # via
     #   asttokens
-    #   imgaug
     #   python-dateutil
 stack-data==0.6.3
     # via ipython
diff --git a/requirements/train.in b/requirements/train.in
index 1fa8d896b..c7fbdbefb 100644
--- a/requirements/train.in
+++ b/requirements/train.in
@@ -2,6 +2,5 @@
 
 lmdb
 Pillow
-imgaug
 hydra-core >=1.2.0
 tensorboardx
diff --git a/requirements/train.txt b/requirements/train.txt
index d30f8ed83..ee64150a0 100644
--- a/requirements/train.txt
+++ b/requirements/train.txt
@@ -3,27 +3,20 @@ contourpy==1.2.0
 cycler==0.12.1
 fonttools==4.49.0
 hydra-core==1.3.2
-imageio==2.34.0
-imgaug==0.4.0
 importlib-resources==6.1.1
 kiwisolver==1.4.5
-lazy-loader==0.3
 lmdb==1.4.1
 matplotlib==3.8.3
 networkx==3.2.1
 numpy==1.26.4
 omegaconf==2.3.0
-opencv-python==4.9.0.80
 packaging==23.2
 pillow==10.2.0
 protobuf==4.25.3
 pyparsing==3.1.1
 python-dateutil==2.8.2
 pyyaml==6.0.1
-scikit-image==0.22.0
 scipy==1.12.0
-shapely==2.0.3
 six==1.16.0
 tensorboardx==2.6.2.2
-tifffile==2024.2.12
 zipp==3.17.0
diff --git a/requirements/tune.in b/requirements/tune.in
index db1be2492..e13a1fd12 100644
--- a/requirements/tune.in
+++ b/requirements/tune.in
@@ -2,7 +2,6 @@
 
 lmdb
 Pillow
-imgaug
 hydra-core >=1.2.0
 ray[tune] >=2.0.0
 ax-platform
diff --git a/requirements/tune.txt b/requirements/tune.txt
index 86ed86985..46ee68f36 100644
--- a/requirements/tune.txt
+++ b/requirements/tune.txt
@@ -20,8 +20,6 @@ fsspec==2024.2.0
 gpytorch==1.11
 hydra-core==1.3.2
 idna==3.6
-imageio==2.34.0
-imgaug==0.4.0
 importlib-resources==6.1.1
 ipython==8.18.1
 ipywidgets==8.1.2
@@ -33,7 +31,6 @@ jsonschema==4.21.1
 jsonschema-specifications==2023.12.1
 jupyterlab-widgets==3.0.10
 kiwisolver==1.4.5
-lazy-loader==0.3
 linear-operator==0.5.1
 lmdb==1.4.1
 markupsafe==2.1.5
@@ -46,7 +43,6 @@ mypy-extensions==1.0.0
 networkx==3.2.1
 numpy==1.26.4
 omegaconf==2.3.0
-opencv-python==4.9.0.80
 opt-einsum==3.3.0
 packaging==23.2
 pandas==2.2.0
@@ -71,17 +67,14 @@ ray==2.9.2
 referencing==0.33.0
 requests==2.31.0
 rpds-py==0.18.0
-scikit-image==0.22.0
 scikit-learn==1.4.1.post1
 scipy==1.12.0
-shapely==2.0.3
 six==1.16.0
 stack-data==0.6.3
 sympy==1.12
 tenacity==8.2.3
 tensorboardx==2.6.2.2
 threadpoolctl==3.3.0
-tifffile==2024.2.12
 tqdm==4.66.2
 traitlets==5.14.1
 typeguard==2.13.3
diff --git a/strhub/data/augment.py b/strhub/data/augment.py
index ed8832503..e2c118cd4 100644
--- a/strhub/data/augment.py
+++ b/strhub/data/augment.py
@@ -12,10 +12,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+#
+# NOTE: imgaug has been replaced with NumPy / SciPy equivalents so that this
+# module is compatible with NumPy 2.x.  The public API (rand_augment_transform)
+# is unchanged.
 
 from functools import partial
 
-import imgaug.augmenters as iaa
 import numpy as np
 from PIL import Image, ImageFilter
 
@@ -42,6 +45,10 @@ def _get_param(level, img, max_dim_factor, min_level=1):
     return round(min(level, max_level))
 
 
+# ---------------------------------------------------------------------------
+# Blur ops
+# ---------------------------------------------------------------------------
+
 def gaussian_blur(img, radius, **__):
     radius = _get_param(radius, img, 0.02)
     key = 'gaussian_blur_' + str(radius)
@@ -50,26 +57,47 @@ def gaussian_blur(img, radius, **__):
 
 
 def motion_blur(img, k, **__):
+    """Horizontal motion blur implemented with scipy.ndimage (no imgaug)."""
     k = _get_param(k, img, 0.08, 3) | 1  # bin to odd values
-    key = 'motion_blur_' + str(k)
-    op = _get_op(key, lambda: iaa.MotionBlur(k))
-    return Image.fromarray(op(image=np.asarray(img)))
-
+    from scipy.ndimage import uniform_filter1d
+    arr = np.asarray(img, dtype=np.float32)
+    if arr.ndim == 2:
+        result = uniform_filter1d(arr, size=k, axis=1, mode='reflect')
+    else:
+        result = np.stack(
+            [uniform_filter1d(arr[..., c], size=k, axis=1, mode='reflect')
+             for c in range(arr.shape[2])],
+            axis=2,
+        )
+    return Image.fromarray(np.clip(result, 0, 255).astype(np.uint8))
+
+
+# ---------------------------------------------------------------------------
+# Noise ops
+# ---------------------------------------------------------------------------
 
 def gaussian_noise(img, scale, **_):
+    """Additive Gaussian noise implemented with NumPy (no imgaug)."""
     scale = _get_param(scale, img, 0.25) | 1  # bin to odd values
-    key = 'gaussian_noise_' + str(scale)
-    op = _get_op(key, lambda: iaa.AdditiveGaussianNoise(scale=scale))
-    return Image.fromarray(op(image=np.asarray(img)))
+    arr = np.asarray(img, dtype=np.float32)
+    rng = np.random.default_rng()
+    noise = rng.normal(0.0, float(scale), arr.shape).astype(np.float32)
+    return Image.fromarray(np.clip(arr + noise, 0, 255).astype(np.uint8))
 
 
 def poisson_noise(img, lam, **_):
+    """Additive Poisson noise implemented with NumPy (no imgaug)."""
     lam = _get_param(lam, img, 0.2) | 1  # bin to odd values
-    key = 'poisson_noise_' + str(lam)
-    op = _get_op(key, lambda: iaa.AdditivePoissonNoise(lam))
-    return Image.fromarray(op(image=np.asarray(img)))
+    arr = np.asarray(img, dtype=np.float32)
+    rng = np.random.default_rng()
+    noise = rng.poisson(float(lam), arr.shape).astype(np.float32)
+    return Image.fromarray(np.clip(arr + noise, 0, 255).astype(np.uint8))
 
 
+# ---------------------------------------------------------------------------
+# RandAugment integration
+# ---------------------------------------------------------------------------
+
 def _level_to_arg(level, _hparams, max):
     level = max * level / auto_augment._LEVEL_DENOM
     return (level,)
@@ -98,6 +126,10 @@ def _level_to_arg(level, _hparams, max):
 
 
 def rand_augment_transform(magnitude=5, num_layers=3):
+    """Return a RandAugment transform compatible with torchvision's Compose.
+
+    The signature is identical to the original so all call-sites are unchanged.
+    """
     # These are tuned for magnitude=5, which means that effective magnitudes are half of these values.
     hparams = {
         'rotate_deg': 30,
diff --git a/strhub/data/dataset.py b/strhub/data/dataset.py
index 0b73c39b2..9a61773ad 100644
--- a/strhub/data/dataset.py
+++ b/strhub/data/dataset.py
@@ -29,7 +29,20 @@
 log = logging.getLogger(__name__)
 
 
+def _is_ascii_charset(charset: str) -> bool:
+    """Return True if every character in *charset* is a plain ASCII character."""
+    return all(ord(c) < 128 for c in charset)
+
+
 def build_tree_dataset(root: Union[PurePath, str], *args, **kwargs):
+    """Recursively collect every LMDB store found under *root*.
+
+    Supports two layouts transparently:
+    - Flat:  root/data.mdb  (single LMDB at the root level)
+    - Tree:  root/**/data.mdb  (one or more LMDBs in sub-directories)
+
+    Both layouts may be mixed freely inside the same root.
+    """
     try:
         kwargs.pop('root')  # prevent 'root' from being passed via kwargs
     except KeyError:
@@ -44,6 +57,13 @@ def build_tree_dataset(root: Union[PurePath, str], *args, **kwargs):
         dataset = LmdbDataset(ds_root, *args, **kwargs)
         log.info(f'\tlmdb:\t{ds_name}\tnum samples: {len(dataset)}')
         datasets.append(dataset)
+    # Also handle the case where the LMDB lives directly at *root*
+    # (i.e. root/data.mdb exists but would be missed by the recursive glob
+    # when root itself is the LMDB directory).
+    if not datasets and (root / 'data.mdb').is_file():
+        dataset = LmdbDataset(str(root), *args, **kwargs)
+        log.info(f'\tlmdb:\t.\tnum samples: {len(dataset)}')
+        datasets.append(dataset)
     return ConcatDataset(datasets)
 
 
@@ -94,6 +114,12 @@ def env(self):
 
     def _preprocess_labels(self, charset, remove_whitespace, normalize_unicode, max_label_len, min_image_dim):
         charset_adapter = CharsetAdapter(charset)
+        # Determine whether the charset is pure ASCII.
+        # For ASCII charsets (all existing English datasets) we keep the original
+        # NFKD → ASCII stripping behaviour.
+        # For non-ASCII charsets (e.g. Devanagari / Marathi) we only apply NFC
+        # normalisation so that composed Unicode characters are preserved.
+        ascii_charset = _is_ascii_charset(charset)
         with self._create_env() as env, env.begin() as txn:
             num_samples = int(txn.get('num-samples'.encode()))
             if self.unlabelled:
@@ -105,9 +131,16 @@ def _preprocess_labels(self, charset, remove_whitespace, normalize_unicode, max_
                 # Normally, whitespace is removed from the labels.
                 if remove_whitespace:
                     label = ''.join(label.split())
-                # Normalize unicode composites (if any) and convert to compatible ASCII characters
+                # Unicode normalisation:
+                #   ASCII charset  → NFKD + strip non-ASCII (original behaviour,
+                #                    keeps English datasets unchanged).
+                #   Non-ASCII charset → NFC only, so composed characters such as
+                #                    Devanagari matras are preserved intact.
                 if normalize_unicode:
-                    label = unicodedata.normalize('NFKD', label).encode('ascii', 'ignore').decode()
+                    if ascii_charset:
+                        label = unicodedata.normalize('NFKD', label).encode('ascii', 'ignore').decode()
+                    else:
+                        label = unicodedata.normalize('NFC', label)
                 # Filter by length before removing unsupported characters. The original label might be too long.
                 if len(label) > max_label_len:
                     continue
diff --git a/strhub/models/abinet/system.py b/strhub/models/abinet/system.py
index f56e9d1df..9426bfed7 100644
--- a/strhub/models/abinet/system.py
+++ b/strhub/models/abinet/system.py
@@ -23,7 +23,6 @@
 from torch.optim import AdamW
 from torch.optim.lr_scheduler import OneCycleLR
 
-from pytorch_lightning.utilities.types import STEP_OUTPUT
 from timm.optim.optim_factory import param_groups_weight_decay
 
 from strhub.models.base import CrossEntropySystem
@@ -177,7 +176,7 @@ def _prepare_inputs_and_targets(self, labels):
         lengths = torch.as_tensor(list(map(len, labels)), device=self.device) + 1  # +1 for eos
         return inputs, lengths, targets
 
-    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
+    def training_step(self, batch, batch_idx) -> Tensor:
         images, labels = batch
         inputs, lengths, targets = self._prepare_inputs_and_targets(labels)
         if self.lm_only:
diff --git a/strhub/models/base.py b/strhub/models/base.py
index 353cf1a22..b3bdde5b3 100644
--- a/strhub/models/base.py
+++ b/strhub/models/base.py
@@ -16,7 +16,7 @@
 import math
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Optional
+from typing import Any, Optional
 
 from nltk import edit_distance
 
@@ -27,7 +27,6 @@
 from torch.optim.lr_scheduler import OneCycleLR
 
 import pytorch_lightning as pl
-from pytorch_lightning.utilities.types import STEP_OUTPUT
 from timm.optim import create_optimizer_v2
 
 from strhub.data.utils import BaseTokenizer, CharsetAdapter, CTCTokenizer, Tokenizer
@@ -44,9 +43,78 @@ class BatchResult:
     loss_numel: int
 
 
+# Type alias for a list of per-step output dicts collected across a validation epoch.
 EPOCH_OUTPUT = list[dict[str, BatchResult]]
 
 
+# Logical groupings of inner-model submodules for freeze / LR control.
+# Each key is the user-facing config name; the value is the attribute name on
+# the inner nn.Module (self.model for PARSeq-like systems).
+_COMPONENT_MAP: dict[str, str] = {
+    'encoder':    'encoder',
+    'decoder':    'decoder',
+    'head':       'head',
+    'text_embed': 'text_embed',
+}
+
+# Components treated as the "backbone" (shared visual/sequential features).
+_BACKBONE_COMPONENTS = ('encoder', 'decoder')
+# Components treated as the "head" (task-specific classifier + embeddings).
+_HEAD_COMPONENTS      = ('head', 'text_embed')
+
+
+def _apply_freeze(inner_model: torch.nn.Module, freeze_cfg: dict[str, bool]) -> None:
+    """Set ``requires_grad`` on submodules of *inner_model* according to *freeze_cfg*.
+
+    Parameters
+    ----------
+    inner_model:
+        The bare ``nn.Module`` (e.g. ``system.model`` for PARSeq).
+    freeze_cfg:
+        Dict mapping component name → True (freeze) / False (keep trainable).
+        Missing keys default to False (trainable).
+    """
+    for cfg_key, module_attr in _COMPONENT_MAP.items():
+        if not freeze_cfg.get(cfg_key, False):
+            continue
+        module = getattr(inner_model, module_attr, None)
+        if module is None:
+            continue
+        for param in module.parameters():
+            param.requires_grad_(False)
+
+
+def _print_finetune_summary(
+    freeze_cfg: dict[str, bool],
+    backbone_lr: float | None,
+    head_lr: float | None,
+    default_lr: float,
+) -> None:
+    """Print a one-time fine-tuning configuration table."""
+    sep = '=' * 48
+    print(sep)
+    print('Fine-tuning Configuration'.center(48))
+    print(sep)
+
+    labels = {
+        'encoder':    'Encoder    ',
+        'decoder':    'Decoder    ',
+        'head':       'Head       ',
+        'text_embed': 'Text Embed ',
+    }
+    for cfg_key, label in labels.items():
+        state = 'Frozen' if freeze_cfg.get(cfg_key, False) else 'Trainable'
+        print(f'  {label} : {state}')
+
+    if backbone_lr is not None and head_lr is not None:
+        print(f'  Backbone LR : {backbone_lr}')
+        print(f'  Head LR     : {head_lr}')
+    else:
+        print(f'  Learning Rate : {default_lr}')
+
+    print(sep)
+
+
 class BaseSystem(pl.LightningModule, ABC):
 
     def __init__(
@@ -57,6 +125,10 @@ def __init__(
         lr: float,
         warmup_pct: float,
         weight_decay: float,
+        # Optional staged fine-tuning controls.  Both default to None so that
+        # existing configs that do not specify them continue to work unchanged.
+        backbone_lr: Optional[float] = None,
+        head_lr: Optional[float] = None,
     ) -> None:
         super().__init__()
         self.tokenizer = tokenizer
@@ -65,6 +137,8 @@ def __init__(
         self.lr = lr
         self.warmup_pct = warmup_pct
         self.weight_decay = weight_decay
+        self.backbone_lr = backbone_lr
+        self.head_lr = head_lr
         self.outputs: EPOCH_OUTPUT = []
 
     @abstractmethod
@@ -95,21 +169,75 @@ def forward_logits_loss(self, images: Tensor, labels: list[str]) -> tuple[Tensor
         """
         raise NotImplementedError
 
-    def configure_optimizers(self):
+    def _lr_scale(self) -> float:
+        """Compute the linear LR scale factor (DDP + grad accumulation)."""
         agb = self.trainer.accumulate_grad_batches
-        # Linear scaling so that the effective learning rate is constant regardless of the number of GPUs used with DDP.
-        lr_scale = agb * math.sqrt(self.trainer.num_devices) * self.batch_size / 256.0
-        lr = lr_scale * self.lr
-        optim = create_optimizer_v2(self, 'adamw', lr, self.weight_decay)
-        sched = OneCycleLR(
-            optim, lr, self.trainer.estimated_stepping_batches, pct_start=self.warmup_pct, cycle_momentum=False
-        )
+        return agb * math.sqrt(self.trainer.num_devices) * self.batch_size / 256.0
+
+    def configure_optimizers(self):
+        scale = self._lr_scale()
+        lr = scale * self.lr
+
+        use_diff_lr = (self.backbone_lr is not None) and (self.head_lr is not None)
+
+        if use_diff_lr:
+            # Build explicit parameter groups so each component gets its own LR.
+            # Frozen parameters (requires_grad=False) are automatically excluded
+            # by create_optimizer_v2 / AdamW since they have no gradient.
+            inner = getattr(self, 'model', None)
+            if inner is not None:
+                backbone_params = [
+                    p for comp in _BACKBONE_COMPONENTS
+                    for p in getattr(inner, comp, torch.nn.Module()).parameters()
+                    if p.requires_grad
+                ]
+                head_params = [
+                    p for comp in _HEAD_COMPONENTS
+                    for p in getattr(inner, comp, torch.nn.Module()).parameters()
+                    if p.requires_grad
+                ]
+                # Any parameters not covered by the two groups (e.g. pos_queries,
+                # dropout) fall back to the global lr.
+                grouped_ids = {id(p) for p in backbone_params + head_params}
+                other_params = [
+                    p for p in self.parameters()
+                    if p.requires_grad and id(p) not in grouped_ids
+                ]
+                param_groups = []
+                if backbone_params:
+                    param_groups.append({'params': backbone_params,
+                                         'lr': scale * self.backbone_lr})
+                if head_params:
+                    param_groups.append({'params': head_params,
+                                         'lr': scale * self.head_lr})
+                if other_params:
+                    param_groups.append({'params': other_params, 'lr': lr})
+            else:
+                # Fallback: no inner model — treat all params uniformly.
+                param_groups = [{'params': list(self.parameters()), 'lr': lr}]
+
+            # Use the maximum LR for the scheduler's max_lr so that no group
+            # exceeds its configured ceiling during the warm-up phase.
+            max_lr = [g['lr'] for g in param_groups]
+            optim = torch.optim.AdamW(param_groups, weight_decay=self.weight_decay)
+            sched = OneCycleLR(
+                optim, max_lr, self.trainer.estimated_stepping_batches,
+                pct_start=self.warmup_pct, cycle_momentum=False,
+            )
+        else:
+            # Original single-LR path — unchanged from the upstream repository.
+            optim = create_optimizer_v2(self, 'adamw', lr, self.weight_decay)
+            sched = OneCycleLR(
+                optim, lr, self.trainer.estimated_stepping_batches,
+                pct_start=self.warmup_pct, cycle_momentum=False,
+            )
+
         return {'optimizer': optim, 'lr_scheduler': {'scheduler': sched, 'interval': 'step'}}
 
     def optimizer_zero_grad(self, epoch: int, batch_idx: int, optimizer: Optimizer) -> None:
         optimizer.zero_grad(set_to_none=True)
 
-    def _eval_step(self, batch, validation: bool) -> Optional[STEP_OUTPUT]:
+    def _eval_step(self, batch, validation: bool) -> Optional[dict[str, Any]]:
         images, labels = batch
 
         correct = 0
@@ -163,7 +291,7 @@ def _aggregate_results(outputs: EPOCH_OUTPUT) -> tuple[float, float, float]:
         loss = total_loss / total_loss_numel
         return acc, ned, loss
 
-    def validation_step(self, batch, batch_idx) -> Optional[STEP_OUTPUT]:
+    def validation_step(self, batch, batch_idx) -> Optional[dict[str, Any]]:
         result = self._eval_step(batch, True)
         self.outputs.append(result)
         return result
@@ -176,17 +304,19 @@ def on_validation_epoch_end(self) -> None:
         self.log('val_loss', loss, sync_dist=True)
         self.log('hp_metric', acc, sync_dist=True)
 
-    def test_step(self, batch, batch_idx) -> Optional[STEP_OUTPUT]:
+    def test_step(self, batch, batch_idx) -> Optional[dict[str, Any]]:
         return self._eval_step(batch, False)
 
 
 class CrossEntropySystem(BaseSystem):
 
     def __init__(
-        self, charset_train: str, charset_test: str, batch_size: int, lr: float, warmup_pct: float, weight_decay: float
+        self, charset_train: str, charset_test: str, batch_size: int, lr: float, warmup_pct: float, weight_decay: float,
+        backbone_lr: Optional[float] = None, head_lr: Optional[float] = None,
     ) -> None:
         tokenizer = Tokenizer(charset_train)
-        super().__init__(tokenizer, charset_test, batch_size, lr, warmup_pct, weight_decay)
+        super().__init__(tokenizer, charset_test, batch_size, lr, warmup_pct, weight_decay,
+                         backbone_lr=backbone_lr, head_lr=head_lr)
         self.bos_id = tokenizer.bos_id
         self.eos_id = tokenizer.eos_id
         self.pad_id = tokenizer.pad_id
@@ -204,10 +334,12 @@ def forward_logits_loss(self, images: Tensor, labels: list[str]) -> tuple[Tensor
 class CTCSystem(BaseSystem):
 
     def __init__(
-        self, charset_train: str, charset_test: str, batch_size: int, lr: float, warmup_pct: float, weight_decay: float
+        self, charset_train: str, charset_test: str, batch_size: int, lr: float, warmup_pct: float, weight_decay: float,
+        backbone_lr: Optional[float] = None, head_lr: Optional[float] = None,
     ) -> None:
         tokenizer = CTCTokenizer(charset_train)
-        super().__init__(tokenizer, charset_test, batch_size, lr, warmup_pct, weight_decay)
+        super().__init__(tokenizer, charset_test, batch_size, lr, warmup_pct, weight_decay,
+                         backbone_lr=backbone_lr, head_lr=head_lr)
         self.blank_id = tokenizer.blank_id
 
     def forward_logits_loss(self, images: Tensor, labels: list[str]) -> tuple[Tensor, Tensor, int]:
diff --git a/strhub/models/crnn/system.py b/strhub/models/crnn/system.py
index a69dfdd13..e135abab5 100644
--- a/strhub/models/crnn/system.py
+++ b/strhub/models/crnn/system.py
@@ -17,8 +17,6 @@
 
 from torch import Tensor
 
-from pytorch_lightning.utilities.types import STEP_OUTPUT
-
 from strhub.models.base import CTCSystem
 from strhub.models.utils import init_weights
 
@@ -49,7 +47,7 @@ def __init__(
     def forward(self, images: Tensor, max_length: Optional[int] = None) -> Tensor:
         return self.model.forward(images)
 
-    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
+    def training_step(self, batch, batch_idx) -> Tensor:
         images, labels = batch
         loss = self.forward_logits_loss(images, labels)[1]
         self.log('loss', loss)
diff --git a/strhub/models/parseq/system.py b/strhub/models/parseq/system.py
index be8581e2b..c726daabe 100644
--- a/strhub/models/parseq/system.py
+++ b/strhub/models/parseq/system.py
@@ -23,9 +23,7 @@
 import torch.nn.functional as F
 from torch import Tensor
 
-from pytorch_lightning.utilities.types import STEP_OUTPUT
-
-from strhub.models.base import CrossEntropySystem
+from strhub.models.base import CrossEntropySystem, _apply_freeze, _print_finetune_summary
 
 from .model import PARSeq as Model
 
@@ -56,9 +54,14 @@ def __init__(
         decode_ar: bool,
         refine_iters: int,
         dropout: float,
+        # Staged fine-tuning options — both optional; None preserves original behaviour.
+        freeze: Optional[dict] = None,
+        backbone_lr: Optional[float] = None,
+        head_lr: Optional[float] = None,
         **kwargs: Any,
     ) -> None:
-        super().__init__(charset_train, charset_test, batch_size, lr, warmup_pct, weight_decay)
+        super().__init__(charset_train, charset_test, batch_size, lr, warmup_pct, weight_decay,
+                         backbone_lr=backbone_lr, head_lr=head_lr)
         self.save_hyperparameters()
 
         self.model = Model(
@@ -78,12 +81,24 @@ def __init__(
             dropout,
         )
 
+        # Apply layer freezing if requested.  A missing or empty freeze dict is a
+        # no-op so that existing configs continue to work without any changes.
+        self._freeze_cfg: dict[str, bool] = dict(freeze) if freeze else {}
+        if self._freeze_cfg:
+            _apply_freeze(self.model, self._freeze_cfg)
+
         # Perm/attn mask stuff
         self.rng = np.random.default_rng()
         self.max_gen_perms = perm_num // 2 if perm_mirrored else perm_num
         self.perm_forward = perm_forward
         self.perm_mirrored = perm_mirrored
 
+    def on_train_start(self) -> None:
+        """Print the fine-tuning summary once before the first training step."""
+        # Only print from rank 0 to avoid duplicate output in DDP.
+        if self.global_rank == 0:
+            _print_finetune_summary(self._freeze_cfg, self.backbone_lr, self.head_lr, self.lr)
+
     def forward(self, images: Tensor, max_length: Optional[int] = None) -> Tensor:
         return self.model.forward(self.tokenizer, images, max_length)
 
@@ -166,7 +181,7 @@ def generate_attn_masks(self, perm):
         query_mask = mask[1:, :-1]
         return content_mask, query_mask
 
-    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
+    def training_step(self, batch, batch_idx) -> Tensor:
         images, labels = batch
         tgt = self.tokenizer.encode(labels, self._device)
 
diff --git a/strhub/models/trba/system.py b/strhub/models/trba/system.py
index eabc5bacf..b73e5b4ef 100644
--- a/strhub/models/trba/system.py
+++ b/strhub/models/trba/system.py
@@ -19,8 +19,6 @@
 import torch
 import torch.nn.functional as F
 from torch import Tensor
-
-from pytorch_lightning.utilities.types import STEP_OUTPUT
 from timm.models.helpers import named_apply
 
 from strhub.models.base import CrossEntropySystem, CTCSystem
@@ -70,7 +68,7 @@ def forward(self, images: Tensor, max_length: Optional[int] = None) -> Tensor:
         text = images.new_full([1], self.bos_id, dtype=torch.long)
         return self.model.forward(images, max_length, text)
 
-    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
+    def training_step(self, batch, batch_idx) -> Tensor:
         images, labels = batch
         encoded = self.tokenizer.encode(labels, self.device)
         inputs = encoded[:, :-1]  # remove <eos>
@@ -118,7 +116,7 @@ def forward(self, images: Tensor, max_length: Optional[int] = None) -> Tensor:
         # max_label_length is unused in CTC prediction
         return self.model.forward(images, None)
 
-    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
+    def training_step(self, batch, batch_idx) -> Tensor:
         images, labels = batch
         loss = self.forward_logits_loss(images, labels)[1]
         self.log('loss', loss)
diff --git a/strhub/models/utils.py b/strhub/models/utils.py
index 4e641bbc4..f56936e4f 100644
--- a/strhub/models/utils.py
+++ b/strhub/models/utils.py
@@ -70,6 +70,223 @@ def get_pretrained_weights(experiment):
     return torch.hub.load_state_dict_from_url(url=url, map_location='cpu', check_hash=True)
 
 
+def _extract_state_dict(checkpoint: dict) -> dict:
+    """Extract a plain ``state_dict`` from any common checkpoint format.
+
+    Handles three layouts seen in the wild:
+
+    * **Plain state dict** – the file *is* the state dict (keys are parameter
+      names mapping directly to tensors).  Used by the official PARSeq weights
+      downloaded via ``get_pretrained_weights()``.
+    * **PyTorch Lightning checkpoint** – a dict that contains a ``"state_dict"``
+      key alongside trainer metadata (``epoch``, ``global_step``, …).
+    * **Generic model checkpoint** – a dict that contains a ``"model"`` key
+      (common in many custom training scripts).
+
+    Parameters
+    ----------
+    checkpoint:
+        The raw object returned by ``torch.hub.load_state_dict_from_url`` or
+        ``torch.load``.
+
+    Returns
+    -------
+    A flat dict mapping parameter names → tensors, ready for key comparison.
+    """
+    if 'state_dict' in checkpoint:
+        return checkpoint['state_dict']
+    if 'model' in checkpoint:
+        return checkpoint['model']
+    # Assume it is already a plain state dict.
+    return checkpoint
+
+
+def _strip_prefix(state_dict: dict, prefixes: tuple[str, ...] = ('model.', 'module.')) -> dict:
+    """Remove well-known wrapper prefixes from checkpoint keys.
+
+    Checkpoint keys sometimes carry a leading ``model.`` (from a
+    ``LightningModule`` that wraps an inner ``nn.Module``) or ``module.``
+    (from ``DataParallel`` / ``DistributedDataParallel``).  Stripping these
+    prefixes normalises the key space so that the checkpoint can be matched
+    against a bare ``nn.Module`` state dict without any manual renaming.
+
+    Only one prefix is stripped per key (the first matching one).  If a key
+    does not start with any of the given prefixes it is returned unchanged.
+
+    Parameters
+    ----------
+    state_dict:
+        Raw checkpoint state dict, potentially containing prefixed keys.
+    prefixes:
+        Tuple of prefix strings to try, in order.  Defaults to the two most
+        common ones: ``'model.'`` and ``'module.'``.
+
+    Returns
+    -------
+    A new dict with prefixes removed where applicable.
+    """
+    stripped: dict = {}
+    for key, value in state_dict.items():
+        new_key = key
+        for prefix in prefixes:
+            if key.startswith(prefix):
+                new_key = key[len(prefix):]
+                break
+        stripped[new_key] = value
+    return stripped
+
+
+def safe_load_pretrained(model: nn.Module, experiment: str) -> dict[str, list[str]]:
+    """Load pretrained weights into *model*, skipping any incompatible tensors.
+
+    Motivation — multilingual OCR changes the classifier dimensions
+    --------------------------------------------------------------
+    PARSeq (and the other models in this hub) follow a consistent architecture:
+
+    ``Image → Encoder (ViT backbone) → Decoder (attention) → Head (linear)``
+
+    The **encoder** and **decoder** learn general visual and sequential
+    features that are *script-agnostic* — they transfer well across languages
+    because they capture stroke patterns, spatial relationships, and sequence
+    context that are useful regardless of whether the target script is Latin,
+    Devanagari, or any other writing system.
+
+    The **classifier head** (``head.weight`` / ``head.bias`` in PARSeq; the
+    final linear projection in ABINet, CRNN, ViTSTR, TRBA) is the only
+    layer tied to the vocabulary size.  Its weight matrix has shape
+    ``(len(charset_train) + num_special_tokens, embed_dim)``.  When the target
+    charset differs from the pretrained charset — e.g. switching from 94
+    English characters to ~80 Devanagari characters — this dimension changes
+    and the tensor **cannot** be reused.
+
+    Why incompatible classifier layers are intentionally skipped
+    ------------------------------------------------------------
+    Forcing a size-mismatched tensor into the model would either raise a
+    ``RuntimeError`` (strict loading) or silently corrupt activations
+    (truncation/padding).  The correct approach for transfer learning is to:
+
+    1. Load the encoder/decoder weights from the source model (same shape).
+    2. Re-initialise the classifier head randomly (new vocab size).
+    3. Fine-tune end-to-end on the target language dataset.
+
+    This is standard practice in cross-lingual OCR and NLP transfer learning.
+
+    Why ``strict=False`` is not used
+    ---------------------------------
+    ``nn.Module.load_state_dict(strict=False)`` silently ignores *all* missing
+    and unexpected keys.  A typo in a layer name, an accidental architecture
+    divergence, or a subtle checkpoint format change would go undetected.
+
+    Instead, this function performs **explicit key-by-key shape comparison**:
+
+    1. Extract the state dict from the checkpoint (supports plain, Lightning,
+       and ``{"model": ...}`` formats).
+    2. Normalise common key prefixes (``model.``, ``module.``).
+    3. Compare each checkpoint key against the current model by shape.
+    4. Build a filtered copy of the current state dict, replacing only
+       shape-compatible tensors.
+    5. Call ``load_state_dict(strict=True)`` on the filtered dict — PyTorch
+       still validates that every key in the model is accounted for.
+
+    Parameters
+    ----------
+    model:
+        The ``nn.Module`` to initialise.  For PARSeq this is ``system.model``
+        (the inner ``nn.Module``); for other systems it is the
+        ``LightningModule`` itself.  See ``train.py`` for the call-site
+        convention.
+    experiment:
+        Pretrained model identifier (e.g. ``'parseq'``, ``'parseq-tiny'``).
+        Passed directly to :func:`get_pretrained_weights`.
+
+    Returns
+    -------
+    dict with four keys:
+
+    * ``"loaded"``     – keys copied successfully from the checkpoint.
+    * ``"skipped"``    – checkpoint keys skipped due to shape mismatch or
+                          absence from the model (typically the classifier head
+                          when fine-tuning on a different charset).
+    * ``"missing"``    – model keys absent from the checkpoint (new layers not
+                          present in the pretrained weights).
+    * ``"unexpected"`` – checkpoint keys that had no corresponding model key
+                          after prefix normalisation.
+    """
+    raw_checkpoint = get_pretrained_weights(experiment)
+    # Step 1 — extract a plain state dict from whatever format was loaded.
+    raw_state = _extract_state_dict(raw_checkpoint)
+    # Step 2 — strip common wrapper prefixes so keys match bare nn.Module keys.
+    ckpt_state = _strip_prefix(raw_state)
+
+    current_state = model.state_dict()
+    ckpt_keys = set(ckpt_state.keys())
+    model_keys = set(current_state.keys())
+
+    loaded: list[str] = []
+    skipped: list[str] = []      # shape mismatch (present in both, but incompatible)
+    missing: list[str] = []      # in model, absent from checkpoint
+    unexpected: list[str] = []   # in checkpoint, absent from model
+
+    # Step 3 — build a filtered state dict starting from current model weights.
+    # Keys not overwritten will retain their randomly-initialised values.
+    updated_state = {k: v.clone() for k, v in current_state.items()}
+
+    for key, ckpt_tensor in ckpt_state.items():
+        if key not in current_state:
+            # Present in checkpoint but not in the current model.
+            unexpected.append(key)
+            continue
+        if ckpt_tensor.shape != current_state[key].shape:
+            # Shape mismatch — almost always the classifier head when the
+            # target charset differs from the pretrained charset.
+            skipped.append(key)
+            continue
+        updated_state[key] = ckpt_tensor
+        loaded.append(key)
+
+    # Collect model keys absent from the (prefix-normalised) checkpoint.
+    for key in model_keys:
+        if key not in ckpt_keys:
+            missing.append(key)
+
+    # Step 4 — strict=True load against the pre-filtered dict.
+    # Every model key is present (unchanged or replaced), so PyTorch's
+    # validation passes while we retain explicit control over what was loaded.
+    model.load_state_dict(updated_state, strict=True)
+
+    # ── Console summary ───────────────────────────────────────────────────────
+    sep = '=' * 48
+    print(sep)
+    print('Safe Pretrained Loading'.center(48))
+    print(sep)
+    print(f'  Checkpoint tensors : {len(ckpt_state)}')
+    print(f'  Model tensors      : {len(current_state)}')
+    print(f'  Loaded             : {len(loaded)}')
+    print(f'  Skipped            : {len(skipped)}')
+    print(f'  Missing            : {len(missing)}')
+    print(f'  Unexpected         : {len(unexpected)}')
+    if skipped:
+        print('  Skipped layers:')
+        for k in skipped:
+            print(f'    {k}')
+    if missing:
+        print('  Missing layers:')
+        for k in missing:
+            print(f'    {k}')
+    if unexpected:
+        print('  Unexpected layers:')
+        for k in unexpected:
+            print(f'    {k}')
+    print(sep)
+
+    return {
+        'loaded': loaded,
+        'skipped': skipped,
+        'missing': missing,
+        'unexpected': unexpected,
+    }
+
+
 def create_model(experiment: str, pretrained: bool = False, **kwargs):
     try:
         config = _get_config(experiment, **kwargs)
@@ -79,7 +296,10 @@ def create_model(experiment: str, pretrained: bool = False, **kwargs):
     model = ModelClass(**config)
     if pretrained:
         m = model.model if 'parseq' in experiment else model
-        m.load_state_dict(get_pretrained_weights(experiment))
+        # Use safe_load_pretrained so that models built with a different charset
+        # (e.g. for multilingual fine-tuning) can still reuse the encoder and
+        # decoder weights even when the classifier head size differs.
+        safe_load_pretrained(m, experiment)
     return model
 
 
diff --git a/strhub/models/vitstr/system.py b/strhub/models/vitstr/system.py
index 37b762e1a..1ce91b7e1 100644
--- a/strhub/models/vitstr/system.py
+++ b/strhub/models/vitstr/system.py
@@ -18,8 +18,6 @@
 import torch
 from torch import Tensor
 
-from pytorch_lightning.utilities.types import STEP_OUTPUT
-
 from strhub.models.base import CrossEntropySystem
 from strhub.models.utils import init_weights
 
@@ -72,7 +70,7 @@ def forward(self, images: Tensor, max_length: Optional[int] = None) -> Tensor:
         logits = logits[:, 1:]
         return logits
 
-    def training_step(self, batch, batch_idx) -> STEP_OUTPUT:
+    def training_step(self, batch, batch_idx) -> Tensor:
         images, labels = batch
         loss = self.forward_logits_loss(images, labels)[1]
         self.log('loss', loss)
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/test_safe_pretrained.py b/tests/test_safe_pretrained.py
new file mode 100644
index 000000000..a9a4ce996
--- /dev/null
+++ b/tests/test_safe_pretrained.py
@@ -0,0 +1,621 @@
+"""Tests for safe_load_pretrained() and its private helpers.
+
+All tests are fully offline — no checkpoint is downloaded from the internet.
+``get_pretrained_weights`` is patched with ``unittest.mock.patch`` so that
+every test constructs its own synthetic checkpoint dict and feeds it directly
+into the functions under test.
+
+Test layout
+-----------
+TestExtractStateDict
+    Tests for _extract_state_dict(): plain dict, Lightning format, model-key format.
+
+TestStripPrefix
+    Tests for _strip_prefix(): model. prefix, module. prefix, no prefix,
+    mixed keys, custom prefixes.
+
+TestSafeLoadPretrained
+    Integration tests for safe_load_pretrained() via a small synthetic model
+    (MinimalModel) that mirrors the encoder + head pattern of PARSeq:
+
+    test_same_architecture_loads_all_layers
+        Same shapes → every key is loaded, nothing is skipped.
+
+    test_different_head_skips_classifier
+        Head with different vocab size → encoder/decoder load, head is skipped.
+
+    test_encoder_values_are_transferred
+        Verifies that encoder weight *values* in the model actually match the
+        checkpoint after loading (not just that no error was raised).
+
+    test_head_values_unchanged_when_skipped
+        Verifies that skipped head tensors retain their original init values.
+
+    test_model_prefix_stripped
+        Checkpoint keys with leading ``model.`` are matched correctly.
+
+    test_module_prefix_stripped
+        Checkpoint keys with leading ``module.`` are matched correctly.
+
+    test_lightning_checkpoint_format
+        Checkpoint wrapped as ``{"state_dict": {...}, "epoch": 5}`` is handled.
+
+    test_plain_state_dict_format
+        A plain state dict (no wrapper keys) is handled identically.
+
+    test_missing_keys_reported
+        Keys present in the model but absent from the checkpoint appear in
+        the ``"missing"`` list and do not cause a RuntimeError.
+
+    test_unexpected_keys_reported
+        Keys in the checkpoint that do not exist in the model appear in
+        the ``"unexpected"`` list.
+
+    test_return_statistics_correct
+        Asserts that all four returned lists contain exactly the right keys.
+
+    test_no_runtime_error_on_shape_mismatch
+        Confirms that a size mismatch does NOT raise RuntimeError.
+"""
+from __future__ import annotations
+
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.nn as nn
+
+# Functions under test
+from strhub.models.utils import (
+    _extract_state_dict,
+    _strip_prefix,
+    safe_load_pretrained,
+)
+
+
+# ---------------------------------------------------------------------------
+# Synthetic models that mirror the PARSeq encoder → head pattern
+# ---------------------------------------------------------------------------
+
+EMBED_DIM = 8   # tiny embedding dimension — keeps tests fast
+
+
+class MinimalModel(nn.Module):
+    """Minimal model with an encoder block and a classifier head.
+
+    Structure mirrors PARSeq:
+      encoder.fc   – weight (EMBED_DIM, EMBED_DIM), bias (EMBED_DIM,)
+      decoder.fc   – weight (EMBED_DIM, EMBED_DIM), bias (EMBED_DIM,)
+      head         – weight (vocab, EMBED_DIM),      bias (vocab,)
+
+    ``vocab`` is the only dimension that changes between charsets.
+    """
+
+    def __init__(self, vocab: int = 10) -> None:
+        super().__init__()
+        self.encoder = nn.Linear(EMBED_DIM, EMBED_DIM)
+        self.decoder = nn.Linear(EMBED_DIM, EMBED_DIM)
+        self.head = nn.Linear(EMBED_DIM, vocab, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:  # pragma: no cover
+        return self.head(self.decoder(self.encoder(x)))
+
+
+def _state_dict_for(vocab: int) -> dict[str, torch.Tensor]:
+    """Return a freshly initialised state dict for ``MinimalModel(vocab)``."""
+    return MinimalModel(vocab).state_dict()
+
+
+def _make_checkpoint(state_dict: dict, fmt: str = 'plain') -> dict:
+    """Wrap *state_dict* in the requested checkpoint format.
+
+    Parameters
+    ----------
+    fmt : ``'plain'`` | ``'lightning'`` | ``'model_key'``
+    """
+    if fmt == 'plain':
+        return state_dict
+    if fmt == 'lightning':
+        return {'state_dict': state_dict, 'epoch': 5, 'global_step': 1000}
+    if fmt == 'model_key':
+        return {'model': state_dict, 'optimizer': {}}
+    raise ValueError(f'Unknown fmt: {fmt!r}')
+
+
+def _add_prefix(state_dict: dict, prefix: str) -> dict:
+    """Return a copy of *state_dict* with every key prefixed by *prefix*."""
+    return {prefix + k: v for k, v in state_dict.items()}
+
+
+# ---------------------------------------------------------------------------
+# Helpers: _extract_state_dict
+# ---------------------------------------------------------------------------
+
+class TestExtractStateDict:
+    """Unit tests for _extract_state_dict()."""
+
+    def test_plain_dict_returned_unchanged(self):
+        sd = _state_dict_for(vocab=10)
+        result = _extract_state_dict(sd)
+        assert result is sd
+
+    def test_lightning_format_extracts_state_dict(self):
+        sd = _state_dict_for(vocab=10)
+        ckpt = {'state_dict': sd, 'epoch': 3, 'global_step': 500}
+        result = _extract_state_dict(ckpt)
+        assert result is sd
+
+    def test_model_key_format_extracts_model(self):
+        sd = _state_dict_for(vocab=10)
+        ckpt = {'model': sd, 'optimizer': {'param_groups': []}}
+        result = _extract_state_dict(ckpt)
+        assert result is sd
+
+    def test_state_dict_takes_priority_over_model_key(self):
+        """'state_dict' is checked before 'model'."""
+        sd_real = _state_dict_for(vocab=10)
+        sd_other = _state_dict_for(vocab=5)
+        ckpt = {'state_dict': sd_real, 'model': sd_other}
+        result = _extract_state_dict(ckpt)
+        assert result is sd_real
+
+
+# ---------------------------------------------------------------------------
+# Helpers: _strip_prefix
+# ---------------------------------------------------------------------------
+
+class TestStripPrefix:
+    """Unit tests for _strip_prefix()."""
+
+    def test_model_prefix_stripped(self):
+        sd = {'model.encoder.weight': torch.zeros(4, 4),
+              'model.head.weight': torch.zeros(10, 4)}
+        result = _strip_prefix(sd)
+        assert set(result.keys()) == {'encoder.weight', 'head.weight'}
+
+    def test_module_prefix_stripped(self):
+        sd = {'module.encoder.weight': torch.zeros(4, 4),
+              'module.head.bias': torch.zeros(10)}
+        result = _strip_prefix(sd)
+        assert set(result.keys()) == {'encoder.weight', 'head.bias'}
+
+    def test_no_prefix_unchanged(self):
+        sd = {'encoder.weight': torch.zeros(4, 4),
+              'head.weight': torch.zeros(10, 4)}
+        result = _strip_prefix(sd)
+        assert set(result.keys()) == {'encoder.weight', 'head.weight'}
+
+    def test_mixed_prefixes_handled(self):
+        """Some keys have a prefix, others do not."""
+        sd = {'model.encoder.weight': torch.zeros(4, 4),
+              'head.weight': torch.zeros(10, 4)}
+        result = _strip_prefix(sd)
+        assert set(result.keys()) == {'encoder.weight', 'head.weight'}
+
+    def test_only_first_prefix_stripped(self):
+        """Double-prefixed key: only the outermost prefix is removed."""
+        sd = {'model.model.encoder.weight': torch.zeros(4, 4)}
+        result = _strip_prefix(sd)
+        # 'model.' stripped once → 'model.encoder.weight'
+        assert set(result.keys()) == {'model.encoder.weight'}
+
+    def test_custom_prefixes(self):
+        sd = {'backbone.layer.weight': torch.zeros(4, 4)}
+        result = _strip_prefix(sd, prefixes=('backbone.',))
+        assert set(result.keys()) == {'layer.weight'}
+
+    def test_values_preserved(self):
+        t = torch.tensor([1.0, 2.0, 3.0])
+        sd = {'model.vec': t}
+        result = _strip_prefix(sd)
+        assert torch.equal(result['vec'], t)
+
+
+# ---------------------------------------------------------------------------
+# Integration: safe_load_pretrained
+# ---------------------------------------------------------------------------
+
+MOCK_TARGET = 'strhub.models.utils.get_pretrained_weights'
+
+
+class TestSafeLoadPretrained:
+    """Integration tests for safe_load_pretrained().
+
+    Every test patches ``get_pretrained_weights`` so no network call is made.
+    The synthetic checkpoint is built from a MinimalModel state dict with
+    controlled shapes so that each scenario is deterministic.
+    """
+
+    # ------------------------------------------------------------------ #
+    # 1. Same architecture — all layers should load                        #
+    # ------------------------------------------------------------------ #
+
+    def test_same_architecture_loads_all_layers(self):
+        """When checkpoint and model have identical architectures, every
+        tensor is loaded and nothing is skipped."""
+        vocab = 10
+        ckpt_sd = _state_dict_for(vocab)
+        model = MinimalModel(vocab)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert len(result['loaded']) == len(ckpt_sd)
+        assert result['skipped'] == []
+        assert result['missing'] == []
+        assert result['unexpected'] == []
+
+    def test_same_architecture_all_layers_count(self):
+        """The total number of loaded keys equals the number of tensors
+        in a MinimalModel state dict (6: enc.w, enc.b, dec.w, dec.b,
+        head.w, head.b)."""
+        vocab = 10
+        model = MinimalModel(vocab)
+        ckpt_sd = _state_dict_for(vocab)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert len(result['loaded']) == 6
+
+    # ------------------------------------------------------------------ #
+    # 2. Different classifier dimensions — encoder/decoder load, head skipped
+    # ------------------------------------------------------------------ #
+
+    def test_different_head_skips_classifier_weight(self):
+        """Checkpoint has vocab=10; model has vocab=80 (Marathi).
+        head.weight should appear in skipped."""
+        ckpt_sd = _state_dict_for(vocab=10)
+        model = MinimalModel(vocab=80)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert 'head.weight' in result['skipped']
+
+    def test_different_head_skips_classifier_bias(self):
+        ckpt_sd = _state_dict_for(vocab=10)
+        model = MinimalModel(vocab=80)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert 'head.bias' in result['skipped']
+
+    def test_different_head_loads_encoder(self):
+        """encoder.weight and encoder.bias must be in 'loaded'."""
+        ckpt_sd = _state_dict_for(vocab=10)
+        model = MinimalModel(vocab=80)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert 'encoder.weight' in result['loaded']
+        assert 'encoder.bias' in result['loaded']
+
+    def test_different_head_loads_decoder(self):
+        """decoder.weight and decoder.bias must be in 'loaded'."""
+        ckpt_sd = _state_dict_for(vocab=10)
+        model = MinimalModel(vocab=80)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert 'decoder.weight' in result['loaded']
+        assert 'decoder.bias' in result['loaded']
+
+    def test_skipped_count_is_two_for_head_mismatch(self):
+        """Only head.weight and head.bias are skipped — exactly 2 tensors."""
+        ckpt_sd = _state_dict_for(vocab=10)
+        model = MinimalModel(vocab=80)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert len(result['skipped']) == 2
+        assert len(result['loaded']) == 4   # enc.w, enc.b, dec.w, dec.b
+
+
+    # ------------------------------------------------------------------ #
+    # 3. Tensor values are actually transferred / preserved                #
+    # ------------------------------------------------------------------ #
+
+    def test_encoder_values_are_transferred(self):
+        """After loading, the model's encoder.weight must equal the
+        checkpoint's encoder.weight (not just 'no error')."""
+        vocab = 10
+        ckpt_sd = _state_dict_for(vocab)
+        # Give the checkpoint encoder a distinct, recognisable value.
+        sentinel = torch.full((EMBED_DIM, EMBED_DIM), fill_value=42.0)
+        ckpt_sd['encoder.weight'] = sentinel
+
+        model = MinimalModel(vocab)
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            safe_load_pretrained(model, 'parseq')
+
+        assert torch.equal(model.encoder.weight.data, sentinel)
+
+    def test_head_values_unchanged_when_skipped(self):
+        """When the head is skipped, the model retains its original
+        (randomly-initialised) head weights, not the checkpoint values."""
+        ckpt_sd = _state_dict_for(vocab=10)
+        model = MinimalModel(vocab=80)
+
+        # Record the model's initial head weight before loading.
+        original_head_weight = model.head.weight.data.clone()
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            safe_load_pretrained(model, 'parseq')
+
+        assert torch.equal(model.head.weight.data, original_head_weight)
+
+    def test_decoder_values_are_transferred(self):
+        """decoder.weight in the model matches the checkpoint value."""
+        vocab = 10
+        ckpt_sd = _state_dict_for(vocab)
+        sentinel = torch.full((EMBED_DIM, EMBED_DIM), fill_value=7.0)
+        ckpt_sd['decoder.weight'] = sentinel
+
+        model = MinimalModel(vocab)
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            safe_load_pretrained(model, 'parseq')
+
+        assert torch.equal(model.decoder.weight.data, sentinel)
+
+    # ------------------------------------------------------------------ #
+    # 4. No RuntimeError on shape mismatch                                 #
+    # ------------------------------------------------------------------ #
+
+    def test_no_runtime_error_on_head_shape_mismatch(self):
+        """A size-mismatched head must NOT raise RuntimeError."""
+        ckpt_sd = _state_dict_for(vocab=10)
+        model = MinimalModel(vocab=80)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            # Should complete without exception.
+            safe_load_pretrained(model, 'parseq')
+
+    # ------------------------------------------------------------------ #
+    # 5. Prefix handling                                                   #
+    # ------------------------------------------------------------------ #
+
+    def test_model_prefix_stripped_and_loaded(self):
+        """Checkpoint keys prefixed with 'model.' are normalised and loaded."""
+        vocab = 10
+        bare_sd = _state_dict_for(vocab)
+        prefixed_sd = _add_prefix(bare_sd, 'model.')
+        model = MinimalModel(vocab)
+
+        with patch(MOCK_TARGET, return_value=prefixed_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert len(result['loaded']) == len(bare_sd)
+        assert result['skipped'] == []
+        assert result['unexpected'] == []
+
+    def test_module_prefix_stripped_and_loaded(self):
+        """Checkpoint keys prefixed with 'module.' are normalised and loaded."""
+        vocab = 10
+        bare_sd = _state_dict_for(vocab)
+        prefixed_sd = _add_prefix(bare_sd, 'module.')
+        model = MinimalModel(vocab)
+
+        with patch(MOCK_TARGET, return_value=prefixed_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert len(result['loaded']) == len(bare_sd)
+        assert result['skipped'] == []
+        assert result['unexpected'] == []
+
+    def test_model_prefix_with_head_mismatch(self):
+        """model. prefix + different vocab: encoder loads, head skipped."""
+        bare_sd = _state_dict_for(vocab=10)
+        prefixed_sd = _add_prefix(bare_sd, 'model.')
+        model = MinimalModel(vocab=80)
+
+        with patch(MOCK_TARGET, return_value=prefixed_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert 'encoder.weight' in result['loaded']
+        assert 'decoder.weight' in result['loaded']
+        assert 'head.weight' in result['skipped']
+        assert 'head.bias' in result['skipped']
+
+
+    # ------------------------------------------------------------------ #
+    # 6. Checkpoint format handling                                        #
+    # ------------------------------------------------------------------ #
+
+    def test_plain_state_dict_format(self):
+        """A checkpoint that is already a plain state dict is handled."""
+        vocab = 10
+        ckpt = _make_checkpoint(_state_dict_for(vocab), fmt='plain')
+        model = MinimalModel(vocab)
+
+        with patch(MOCK_TARGET, return_value=ckpt):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert len(result['loaded']) == 6
+        assert result['skipped'] == []
+
+    def test_lightning_checkpoint_format(self):
+        """A PyTorch Lightning checkpoint (state_dict key + metadata) is
+        unwrapped automatically."""
+        vocab = 10
+        ckpt = _make_checkpoint(_state_dict_for(vocab), fmt='lightning')
+        model = MinimalModel(vocab)
+
+        with patch(MOCK_TARGET, return_value=ckpt):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert len(result['loaded']) == 6
+        assert result['skipped'] == []
+
+    def test_lightning_checkpoint_with_head_mismatch(self):
+        """Lightning-wrapped checkpoint + different charset: head skipped."""
+        ckpt = _make_checkpoint(_state_dict_for(vocab=10), fmt='lightning')
+        model = MinimalModel(vocab=80)
+
+        with patch(MOCK_TARGET, return_value=ckpt):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert 'head.weight' in result['skipped']
+        assert 'encoder.weight' in result['loaded']
+
+    def test_model_key_checkpoint_format(self):
+        """A checkpoint dict with a 'model' key is unwrapped automatically."""
+        vocab = 10
+        ckpt = _make_checkpoint(_state_dict_for(vocab), fmt='model_key')
+        model = MinimalModel(vocab)
+
+        with patch(MOCK_TARGET, return_value=ckpt):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert len(result['loaded']) == 6
+        assert result['skipped'] == []
+
+    # ------------------------------------------------------------------ #
+    # 7. Missing and unexpected key reporting                              #
+    # ------------------------------------------------------------------ #
+
+    def test_missing_keys_reported(self):
+        """Keys in the model but absent from the checkpoint appear in
+        'missing', and no RuntimeError is raised."""
+        vocab = 10
+        # Checkpoint is missing decoder keys entirely.
+        bare_sd = _state_dict_for(vocab)
+        reduced_sd = {k: v for k, v in bare_sd.items()
+                      if not k.startswith('decoder')}
+        model = MinimalModel(vocab)
+
+        with patch(MOCK_TARGET, return_value=reduced_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert 'decoder.weight' in result['missing']
+        assert 'decoder.bias' in result['missing']
+        # The present layers still load.
+        assert 'encoder.weight' in result['loaded']
+
+    def test_unexpected_keys_reported(self):
+        """Keys in the checkpoint that do not exist in the model appear in
+        'unexpected', and no error is raised."""
+        vocab = 10
+        bare_sd = _state_dict_for(vocab)
+        # Add a key that the model doesn't have.
+        bare_sd['ghost.weight'] = torch.zeros(4, 4)
+        model = MinimalModel(vocab)
+
+        with patch(MOCK_TARGET, return_value=bare_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert 'ghost.weight' in result['unexpected']
+        # All real layers still load.
+        assert len(result['loaded']) == 6
+
+    # ------------------------------------------------------------------ #
+    # 8. Return statistics correctness                                     #
+    # ------------------------------------------------------------------ #
+
+    def test_return_statistics_correct_same_arch(self):
+        """Full correctness check: same architecture, plain checkpoint."""
+        vocab = 10
+        ckpt_sd = _state_dict_for(vocab)
+        model = MinimalModel(vocab)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        expected_keys = {'encoder.weight', 'encoder.bias',
+                         'decoder.weight', 'decoder.bias',
+                         'head.weight', 'head.bias'}
+        assert set(result['loaded']) == expected_keys
+        assert result['skipped'] == []
+        assert result['missing'] == []
+        assert result['unexpected'] == []
+
+    def test_return_statistics_correct_head_mismatch(self):
+        """Full correctness check: different classifier dimension."""
+        ckpt_sd = _state_dict_for(vocab=10)
+        model = MinimalModel(vocab=80)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert set(result['loaded']) == {
+            'encoder.weight', 'encoder.bias',
+            'decoder.weight', 'decoder.bias',
+        }
+        assert set(result['skipped']) == {'head.weight', 'head.bias'}
+        assert result['missing'] == []
+        assert result['unexpected'] == []
+
+    def test_return_dict_has_all_four_keys(self):
+        """The returned dict always contains exactly the four expected keys."""
+        model = MinimalModel(vocab=10)
+        ckpt_sd = _state_dict_for(vocab=10)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        assert set(result.keys()) == {'loaded', 'skipped', 'missing', 'unexpected'}
+
+    def test_all_return_values_are_lists(self):
+        """Every value in the returned dict is a list (not a set or tuple)."""
+        model = MinimalModel(vocab=10)
+        ckpt_sd = _state_dict_for(vocab=10)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            result = safe_load_pretrained(model, 'parseq')
+
+        for key, value in result.items():
+            assert isinstance(value, list), f'result[{key!r}] should be a list'
+
+
+    # ------------------------------------------------------------------ #
+    # 9. Idempotency                                                       #
+    # ------------------------------------------------------------------ #
+
+    def test_calling_twice_is_idempotent(self):
+        """Calling safe_load_pretrained twice on the same model with the same
+        checkpoint should produce identical results both times."""
+        vocab = 10
+        ckpt_sd = _state_dict_for(vocab)
+        model = MinimalModel(vocab)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            result1 = safe_load_pretrained(model, 'parseq')
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            result2 = safe_load_pretrained(model, 'parseq')
+
+        assert result1 == result2
+
+    # ------------------------------------------------------------------ #
+    # 10. Model remains in evaluation-ready state after loading            #
+    # ------------------------------------------------------------------ #
+
+    def test_model_can_forward_after_loading(self):
+        """After safe_load_pretrained the model should accept a forward pass
+        without raising any exception."""
+        vocab = 10
+        ckpt_sd = _state_dict_for(vocab)
+        model = MinimalModel(vocab)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            safe_load_pretrained(model, 'parseq')
+
+        x = torch.zeros(2, EMBED_DIM)
+        out = model(x)
+        assert out.shape == (2, vocab)
+
+    def test_model_can_forward_after_loading_with_head_mismatch(self):
+        """After partial loading (head skipped), the model still runs inference
+        with the new vocab dimension."""
+        ckpt_sd = _state_dict_for(vocab=10)
+        model = MinimalModel(vocab=80)
+
+        with patch(MOCK_TARGET, return_value=ckpt_sd):
+            safe_load_pretrained(model, 'parseq')
+
+        x = torch.zeros(3, EMBED_DIM)
+        out = model(x)
+        assert out.shape == (3, 80)
diff --git a/tools/evaluate_marathi.py b/tools/evaluate_marathi.py
new file mode 100644
index 000000000..7ec8fd4b0
--- /dev/null
+++ b/tools/evaluate_marathi.py
@@ -0,0 +1,231 @@
+#!/usr/bin/env python3
+"""Evaluate a PARSeq checkpoint on a Marathi (or any) LMDB validation dataset.
+
+Usage
+-----
+python tools/evaluate_marathi.py \\
+    --checkpoint outputs/parseq/best.ckpt \\
+    --data_root   data \\
+    --lmdb_path   data/val \\
+    --batch_size  64 \\
+    --num_workers 4 \\
+    --output      evaluation_predictions.csv
+
+The script:
+  1. Loads the model from *checkpoint* using the existing load_from_checkpoint().
+  2. Opens the LMDB at *lmdb_path* directly (no intermediate DataModule needed).
+  3. Runs inference on every sample with torch.inference_mode().
+  4. Decodes predictions via the model's own tokenizer (no duplication).
+  5. Compares predictions to ground-truth with NFC normalisation (not NFKD) so
+     Devanagari matras and composed characters are handled correctly.
+  6. Prints a metrics table: Exact Match Accuracy, Character Accuracy, CER, NED.
+  7. Writes evaluation_predictions.csv with columns:
+       index, ground_truth, prediction, confidence
+
+Metrics definitions
+-------------------
+Exact Match Accuracy : fraction of samples where prediction == ground_truth
+Character Accuracy   : fraction of individual characters that are correct
+                       (1 - CER) × 100, capped at 0
+CER                  : Character Error Rate = edit_distance / len(ground_truth)
+                       averaged across all samples
+NED                  : Normalised Edit Distance per ICDAR 2019 definition
+                       = 1 - mean(edit_distance / max(len(pred), len(gt)))
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import sys
+import unicodedata
+from pathlib import Path
+
+from nltk import edit_distance
+from tqdm import tqdm
+
+import torch
+from torch.utils.data import DataLoader
+
+from strhub.data.dataset import LmdbDataset
+from strhub.data.module import SceneTextDataModule
+from strhub.models.utils import load_from_checkpoint
+
+
+# ---------------------------------------------------------------------------
+# Unicode helpers
+# ---------------------------------------------------------------------------
+
+def _nfc(s: str) -> str:
+    """NFC-normalise a string.  Used for Marathi (Devanagari) comparisons so
+    that composed characters such as matras are preserved.  NFKD is explicitly
+    NOT used here because it strips Devanagari characters to ASCII."""
+    return unicodedata.normalize('NFC', s)
+
+
+# ---------------------------------------------------------------------------
+# Metrics
+# ---------------------------------------------------------------------------
+
+def compute_metrics(
+    ground_truths: list[str],
+    predictions: list[str],
+) -> dict[str, float]:
+    """Compute evaluation metrics over a list of (gt, pred) pairs.
+
+    All strings are NFC-normalised before comparison.
+
+    Returns
+    -------
+    dict with keys: exact_match, char_accuracy, cer, ned
+    """
+    assert len(ground_truths) == len(predictions)
+
+    n = len(ground_truths)
+    if n == 0:
+        return {'exact_match': 0.0, 'char_accuracy': 0.0, 'cer': 0.0, 'ned': 0.0}
+
+    exact_correct = 0
+    total_char_errors = 0
+    total_gt_chars = 0
+    total_ned = 0.0
+
+    for gt_raw, pred_raw in zip(ground_truths, predictions):
+        gt = _nfc(gt_raw)
+        pred = _nfc(pred_raw)
+
+        # Exact match
+        if pred == gt:
+            exact_correct += 1
+
+        # Edit distance for CER and NED
+        dist = edit_distance(pred, gt)
+        total_char_errors += dist
+        total_gt_chars += max(len(gt), 1)  # avoid division by zero on empty gt
+
+        # NED: ICDAR 2019 definition
+        denom = max(len(pred), len(gt))
+        total_ned += dist / denom if denom > 0 else 0.0
+
+    exact_match = 100.0 * exact_correct / n
+    cer = 100.0 * total_char_errors / total_gt_chars
+    char_accuracy = max(0.0, 100.0 - cer)
+    ned = 100.0 * (1.0 - total_ned / n)
+
+    return {
+        'exact_match': exact_match,
+        'char_accuracy': char_accuracy,
+        'cer': cer,
+        'ned': ned,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+@torch.inference_mode()
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description='Evaluate a PARSeq checkpoint on a Marathi LMDB dataset.',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument('--checkpoint',   required=True,
+                        help="Path to .ckpt file, or 'pretrained=<id>'")
+    parser.add_argument('--data_root',    default='data',
+                        help='Root data directory (used when lmdb_path is relative)')
+    parser.add_argument('--lmdb_path',    default=None,
+                        help='Path to the LMDB directory to evaluate. '
+                             'Defaults to <data_root>/val')
+    parser.add_argument('--batch_size',   type=int, default=64)
+    parser.add_argument('--num_workers',  type=int, default=4)
+    parser.add_argument('--device',       default='cuda')
+    parser.add_argument('--output',       default='evaluation_predictions.csv',
+                        help='Path for the CSV output file')
+    args = parser.parse_args()
+
+    # ── Resolve LMDB path ────────────────────────────────────────────────────
+    if args.lmdb_path is None:
+        lmdb_path = str(Path(args.data_root) / 'val')
+    else:
+        lmdb_path = args.lmdb_path
+
+    device = torch.device(args.device if torch.cuda.is_available()
+                          or args.device == 'cpu' else 'cpu')
+
+    # ── Load model ───────────────────────────────────────────────────────────
+    print(f'Loading checkpoint: {args.checkpoint}')
+    model = load_from_checkpoint(args.checkpoint).eval().to(device)
+    hp = model.hparams
+
+    # ── Build dataset / dataloader ────────────────────────────────────────────
+    transform = SceneTextDataModule.get_transform(hp.img_size)
+    dataset = LmdbDataset(
+        root=lmdb_path,
+        charset=hp.charset_test,
+        max_label_len=hp.max_label_length,
+        remove_whitespace=True,
+        normalize_unicode=True,
+        transform=transform,
+    )
+    loader = DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=(device.type == 'cuda'),
+        shuffle=False,
+    )
+    print(f'Dataset: {lmdb_path}  ({len(dataset)} samples)')
+
+    # ── Inference ─────────────────────────────────────────────────────────────
+    all_gts: list[str] = []
+    all_preds: list[str] = []
+    all_confs: list[float] = []
+
+    for imgs, labels in tqdm(loader, desc='Evaluating', unit='batch'):
+        imgs = imgs.to(device)
+        logits = model(imgs)                     # (N, L, C)
+        probs = logits.softmax(-1)
+        preds, prob_seqs = model.tokenizer.decode(probs)
+
+        for pred, prob_seq, gt in zip(preds, prob_seqs, labels):
+            # Apply the model's charset adapter so evaluation mirrors training.
+            pred = model.charset_adapter(pred)
+            confidence = prob_seq.prod().item()
+
+            all_gts.append(gt)
+            all_preds.append(pred)
+            all_confs.append(confidence)
+
+    # ── Metrics ───────────────────────────────────────────────────────────────
+    metrics = compute_metrics(all_gts, all_preds)
+
+    sep = '=' * 48
+    print(f'\n{sep}')
+    print('Evaluation Results'.center(48))
+    print(sep)
+    print(f'  Samples          : {len(all_gts)}')
+    print(f'  Exact Match Acc  : {metrics["exact_match"]:.2f} %')
+    print(f'  Character Acc    : {metrics["char_accuracy"]:.2f} %')
+    print(f'  CER              : {metrics["cer"]:.2f} %')
+    print(f'  NED              : {metrics["ned"]:.2f} %')
+    print(sep)
+
+    # ── CSV export ────────────────────────────────────────────────────────────
+    output_path = Path(args.output)
+    with output_path.open('w', newline='', encoding='utf-8') as f:
+        writer = csv.DictWriter(f, fieldnames=['index', 'ground_truth',
+                                               'prediction', 'confidence'])
+        writer.writeheader()
+        for idx, (gt, pred, conf) in enumerate(zip(all_gts, all_preds, all_confs)):
+            writer.writerow({
+                'index':        idx,
+                'ground_truth': gt,
+                'prediction':   pred,
+                'confidence':   f'{conf:.6f}',
+            })
+
+    print(f'\nPredictions saved to: {output_path.resolve()}')
+
+
+if __name__ == '__main__':
+    main()
diff --git a/tools/infer_marathi.py b/tools/infer_marathi.py
new file mode 100644
index 000000000..8918fb6b8
--- /dev/null
+++ b/tools/infer_marathi.py
@@ -0,0 +1,197 @@
+#!/usr/bin/env python3
+"""Run inference with a PARSeq checkpoint on one image or a folder of images.
+
+Single image
+------------
+python tools/infer_marathi.py \\
+    --checkpoint best.ckpt \\
+    --image      image.jpg
+
+Folder of images
+----------------
+python tools/infer_marathi.py \\
+    --checkpoint best.ckpt \\
+    --folder     images/
+
+Both modes print per-image results to stdout and write predictions.csv.
+
+Output CSV columns
+------------------
+filename  |  prediction  |  confidence
+
+Notes
+-----
+- Images are loaded as RGB and resized to the model's expected input size.
+- Predictions are NFC-normalised before printing and saving.
+- Supported image extensions: .jpg .jpeg .png .bmp .tiff .webp
+- When --folder is given, images are sorted alphabetically for reproducibility.
+- Confidence is the product of per-token probabilities (same as training).
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import unicodedata
+from pathlib import Path
+
+from PIL import Image
+from tqdm import tqdm
+
+import torch
+
+from strhub.data.module import SceneTextDataModule
+from strhub.models.utils import load_from_checkpoint
+
+# Image extensions accepted for folder inference.
+_IMAGE_EXTS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.webp'}
+
+
+# ---------------------------------------------------------------------------
+# Unicode helpers
+# ---------------------------------------------------------------------------
+
+def _nfc(s: str) -> str:
+    """NFC-normalise *s*.  Used for Marathi (Devanagari) strings so that
+    composed characters such as matras and vowel signs are preserved.
+    NFKD is explicitly NOT applied here."""
+    return unicodedata.normalize('NFC', s)
+
+
+# ---------------------------------------------------------------------------
+# Core inference helpers
+# ---------------------------------------------------------------------------
+
+def _collect_images(folder: Path) -> list[Path]:
+    """Return all image files under *folder*, sorted alphabetically."""
+    images = sorted(
+        p for p in folder.iterdir()
+        if p.is_file() and p.suffix.lower() in _IMAGE_EXTS
+    )
+    return images
+
+
+def _infer_batch(
+    model: torch.nn.Module,
+    image_paths: list[Path],
+    transform,
+    device: torch.device,
+) -> list[tuple[str, str, float]]:
+    """Run inference on a list of image paths.
+
+    Returns a list of (filename, prediction, confidence) tuples.
+    """
+    results: list[tuple[str, str, float]] = []
+
+    for path in tqdm(image_paths, desc='Inferring', unit='img'):
+        try:
+            img = Image.open(path).convert('RGB')
+        except Exception as exc:  # noqa: BLE001
+            print(f'  [WARN] Could not open {path}: {exc}')
+            results.append((path.name, '', 0.0))
+            continue
+
+        tensor = transform(img).unsqueeze(0).to(device)
+        logits = model(tensor)                     # (1, L, C)
+        probs = logits.softmax(-1)
+        preds, prob_seqs = model.tokenizer.decode(probs)
+
+        pred = model.charset_adapter(preds[0])
+        pred = _nfc(pred)
+        confidence = prob_seqs[0].prod().item()
+
+        results.append((path.name, pred, confidence))
+
+    return results
+
+
+def _print_and_save(
+    results: list[tuple[str, str, float]],
+    output_csv: Path,
+) -> None:
+    """Print a summary table and write the CSV file."""
+    print('\n' + '─' * 60)
+    print(f'  {"Filename":<30}  {"Prediction":<20}  Confidence')
+    print('─' * 60)
+    for filename, pred, conf in results:
+        print(f'  {filename:<30}  {pred:<20}  {conf:.4f}')
+    print('─' * 60)
+
+    with output_csv.open('w', newline='', encoding='utf-8') as f:
+        writer = csv.DictWriter(f, fieldnames=['filename', 'prediction', 'confidence'])
+        writer.writeheader()
+        for filename, pred, conf in results:
+            writer.writerow({
+                'filename':   filename,
+                'prediction': pred,
+                'confidence': f'{conf:.6f}',
+            })
+    print(f'\nPredictions saved to: {output_csv.resolve()}')
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+@torch.inference_mode()
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description='Run PARSeq inference on a single image or a folder.',
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument('--checkpoint', required=True,
+                        help="Path to .ckpt file, or 'pretrained=<id>'")
+    parser.add_argument('--device',     default='cuda')
+    parser.add_argument('--output',     default='predictions.csv',
+                        help='Output CSV file path')
+
+    # Input source — exactly one of --image or --folder must be provided.
+    src = parser.add_mutually_exclusive_group(required=True)
+    src.add_argument('--image',  metavar='FILE',
+                     help='Single image file to process')
+    src.add_argument('--folder', metavar='DIR',
+                     help='Directory of images to process')
+
+    args = parser.parse_args()
+
+    device = torch.device(args.device if torch.cuda.is_available()
+                          or args.device == 'cpu' else 'cpu')
+
+    # ── Load model ───────────────────────────────────────────────────────────
+    print(f'Loading checkpoint: {args.checkpoint}')
+    model = load_from_checkpoint(args.checkpoint).eval().to(device)
+    transform = SceneTextDataModule.get_transform(model.hparams.img_size)
+
+    # ── Collect input images ─────────────────────────────────────────────────
+    if args.image:
+        image_path = Path(args.image)
+        if not image_path.is_file():
+            print(f'Error: file not found: {image_path}', flush=True)
+            raise SystemExit(1)
+        image_paths = [image_path]
+    else:
+        folder = Path(args.folder)
+        if not folder.is_dir():
+            print(f'Error: directory not found: {folder}', flush=True)
+            raise SystemExit(1)
+        image_paths = _collect_images(folder)
+        if not image_paths:
+            print(f'No images found in {folder} '
+                  f'(supported: {", ".join(sorted(_IMAGE_EXTS))})')
+            raise SystemExit(0)
+        print(f'Found {len(image_paths)} image(s) in {folder}')
+
+    # ── Inference ─────────────────────────────────────────────────────────────
+    results = _infer_batch(model, image_paths, transform, device)
+
+    # ── Single-image convenience output ──────────────────────────────────────
+    if args.image and results:
+        _, pred, conf = results[0]
+        print(f'\nPrediction : {pred}')
+        print(f'Confidence : {conf:.4f}')
+
+    # ── Print table + save CSV ────────────────────────────────────────────────
+    _print_and_save(results, Path(args.output))
+
+
+if __name__ == '__main__':
+    main()
diff --git a/train.py b/train.py
index e4c3aff5c..c73aa48e2 100755
--- a/train.py
+++ b/train.py
@@ -26,11 +26,25 @@
 from pytorch_lightning.callbacks import ModelCheckpoint, StochasticWeightAveraging
 from pytorch_lightning.loggers import TensorBoardLogger
 from pytorch_lightning.strategies import DDPStrategy
-from pytorch_lightning.utilities.model_summary import summarize
+from pytorch_lightning.utilities.model_summary import ModelSummary
 
 from strhub.data.module import SceneTextDataModule
 from strhub.models.base import BaseSystem
-from strhub.models.utils import get_pretrained_weights
+from strhub.models.utils import safe_load_pretrained
+
+
+def _get_autocast_dtype(device_type: str = 'cuda'):
+    """Return the autocast dtype for the given device.
+
+    Wraps both the current and deprecated PyTorch APIs so the code runs on
+    PyTorch 2.0 through 2.x without warnings.
+    """
+    try:
+        # PyTorch >= 2.1
+        return torch.get_autocast_dtype(device_type)
+    except AttributeError:
+        # PyTorch < 2.1 fallback
+        return torch.get_autocast_gpu_dtype()  # type: ignore[attr-defined]
 
 
 # Copied from OneCycleLR
@@ -60,8 +74,9 @@ def main(config: DictConfig):
         gpu = config.trainer.get('accelerator') == 'gpu'
         devices = config.trainer.get('devices', 0)
         if gpu:
-            # Use mixed-precision training
-            config.trainer.precision = 'bf16-mixed' if torch.get_autocast_gpu_dtype() is torch.bfloat16 else '16-mixed'
+            # Use mixed-precision training.
+            # 'bf16-mixed' / '16-mixed' are the PL 2.x precision strings.
+            config.trainer.precision = 'bf16-mixed' if _get_autocast_dtype('cuda') is torch.bfloat16 else '16-mixed'
         if gpu and devices > 1:
             # Use DDP with optimizations
             trainer_strategy = DDPStrategy(find_unused_parameters=False, gradient_as_bucket_view=True)
@@ -75,11 +90,15 @@ def main(config: DictConfig):
         assert config.model.perm_num % 2 == 0, 'perm_num should be even if perm_mirrored = True'
 
     model: BaseSystem = hydra.utils.instantiate(config.model)
-    # If specified, use pretrained weights to initialize the model
+    # If specified, use pretrained weights to initialise the model.
+    # safe_load_pretrained() skips classifier layers whose shape does not match
+    # the current model (e.g. when fine-tuning on a different charset such as
+    # Marathi) so that the encoder/decoder weights are always reused even when
+    # the output head must be re-initialised from scratch.
     if config.pretrained is not None:
         m = model.model if config.model._target_.endswith('PARSeq') else model
-        m.load_state_dict(get_pretrained_weights(config.pretrained))
-    print(summarize(model, max_depth=2))
+        safe_load_pretrained(m, config.pretrained)
+    print(ModelSummary(model, max_depth=2))
 
     datamodule: SceneTextDataModule = hydra.utils.instantiate(config.data)
 
diff --git a/tune.py b/tune.py
index a5020322d..3de895df4 100755
--- a/tune.py
+++ b/tune.py
@@ -137,9 +137,10 @@ def main(config: DictConfig):
         assert config.model.perm_num % 2 == 0, 'perm_num should be even if perm_mirrored = True'
     # Modify config
     with open_dict(config):
-        # Use mixed-precision training
-        if config.trainer.get('gpus', 0):
-            config.trainer.precision = 16
+        # Use mixed-precision training.
+        # PL 2.x uses the string form '16-mixed'; the old integer form (16) is deprecated.
+        if config.trainer.get('accelerator') == 'gpu':
+            config.trainer.precision = '16-mixed'
         # Resolve absolute path to data.root_dir
         config.data.root_dir = hydra.utils.to_absolute_path(config.data.root_dir)
 
@@ -187,7 +188,7 @@ def main(config: DictConfig):
                 name=out_dir.name,
                 stop=MetricTracker('NED', max_t),
                 progress_reporter=reporter,
-                local_dir=str(out_dir.parent.absolute()),
+                storage_path=str(out_dir.parent.absolute()),
             ),
         )
     else: