diff --git a/docling/datamodel/layout_model_specs.py b/docling/datamodel/layout_model_specs.py index 491d82df..b91fa7fe 100644 --- a/docling/datamodel/layout_model_specs.py +++ b/docling/datamodel/layout_model_specs.py @@ -29,42 +29,42 @@ class LayoutModelConfig(BaseModel): # HuggingFace Layout Models # Default Docling Layout Model -docling_layout_v2 = LayoutModelConfig( +DOCLING_LAYOUT_V2 = LayoutModelConfig( name="docling_layout_v2", repo_id="ds4sd/docling-layout-old", revision="main", model_path="", ) -docling_layout_heron = LayoutModelConfig( +DOCLING_LAYOUT_HERON = LayoutModelConfig( name="docling_layout_heron", repo_id="ds4sd/docling-layout-heron", revision="main", model_path="", ) -docling_layout_heron_101 = LayoutModelConfig( +DOCLING_LAYOUT_HERON_101 = LayoutModelConfig( name="docling_layout_heron_101", repo_id="ds4sd/docling-layout-heron-101", revision="main", model_path="", ) -docling_layout_egret_medium = LayoutModelConfig( +DOCLING_LAYOUT_EGRET_MEDIUM = LayoutModelConfig( name="docling_layout_egret_medium", repo_id="ds4sd/docling-layout-egret-medium", revision="main", model_path="", ) -docling_layout_egret_large = LayoutModelConfig( +DOCLING_LAYOUT_EGRET_LARGE = LayoutModelConfig( name="docling_layout_egret_large", repo_id="ds4sd/docling-layout-egret-large", revision="main", model_path="", ) -docling_layout_egret_xlarge = LayoutModelConfig( +DOCLING_LAYOUT_EGRET_XLARGE = LayoutModelConfig( name="docling_layout_egret_xlarge", repo_id="ds4sd/docling-layout-egret-xlarge", revision="main", diff --git a/docling/datamodel/pipeline_options.py b/docling/datamodel/pipeline_options.py index fec3db76..6ab8fb04 100644 --- a/docling/datamodel/pipeline_options.py +++ b/docling/datamodel/pipeline_options.py @@ -17,13 +17,13 @@ from docling.datamodel import asr_model_specs # Import the following for backwards compatibility from docling.datamodel.accelerator_options import AcceleratorDevice, AcceleratorOptions from docling.datamodel.layout_model_specs import ( + DOCLING_LAYOUT_EGRET_LARGE, + DOCLING_LAYOUT_EGRET_MEDIUM, + DOCLING_LAYOUT_EGRET_XLARGE, + DOCLING_LAYOUT_HERON, + DOCLING_LAYOUT_HERON_101, + DOCLING_LAYOUT_V2, LayoutModelConfig, - docling_layout_egret_large, - docling_layout_egret_medium, - docling_layout_egret_xlarge, - docling_layout_heron, - docling_layout_heron_101, - docling_layout_v2, ) from docling.datamodel.pipeline_options_asr_model import ( InlineAsrOptions, @@ -279,7 +279,7 @@ class LayoutOptions(BaseModel): """Options for layout processing.""" create_orphan_clusters: bool = True # Whether to create clusters for orphaned cells - model_spec: LayoutModelConfig = docling_layout_v2 + model_spec: LayoutModelConfig = DOCLING_LAYOUT_V2 class AsrPipelineOptions(PipelineOptions): diff --git a/docling/models/layout_model.py b/docling/models/layout_model.py index fbe04313..05a86f31 100644 --- a/docling/models/layout_model.py +++ b/docling/models/layout_model.py @@ -12,7 +12,7 @@ from PIL import Image from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.base_models import BoundingBox, Cluster, LayoutPrediction, Page from docling.datamodel.document import ConversionResult -from docling.datamodel.layout_model_specs import LayoutModelConfig, docling_layout_v2 +from docling.datamodel.layout_model_specs import DOCLING_LAYOUT_V2, LayoutModelConfig from docling.datamodel.pipeline_options import LayoutOptions from docling.datamodel.settings import settings from docling.models.base_model import BasePageModel @@ -91,7 +91,7 @@ class LayoutModel(BasePageModel): local_dir: Optional[Path] = None, force: bool = False, progress: bool = False, - layout_model_config: LayoutModelConfig = docling_layout_v2, + layout_model_config: LayoutModelConfig = DOCLING_LAYOUT_V2, ) -> Path: return download_hf_model( repo_id=layout_model_config.repo_id, diff --git a/docling/utils/model_downloader.py b/docling/utils/model_downloader.py index b93efc83..a2994fb7 100644 --- a/docling/utils/model_downloader.py +++ b/docling/utils/model_downloader.py @@ -2,7 +2,7 @@ import logging from pathlib import Path from typing import Optional -from docling.datamodel.layout_model_specs import docling_layout_v2 +from docling.datamodel.layout_model_specs import DOCLING_LAYOUT_V2 from docling.datamodel.pipeline_options import ( granite_picture_description, smolvlm_picture_description, @@ -47,7 +47,7 @@ def download_models( if with_layout: _log.info("Downloading layout model...") LayoutModel.download_models( - local_dir=output_dir / docling_layout_v2.model_repo_folder, + local_dir=output_dir / DOCLING_LAYOUT_V2.model_repo_folder, force=force, progress=progress, )