chage torch_dtype

app.py

@@ -13,9 +13,17 @@ from transformers import AutoProcessor, AutoModelForVision2Seq
 from transformers.image_utils import load_image
 from pathlib import Path
 import time
+import os
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Optimize for CPU: set float32 and use all CPU threads
+if DEVICE == "cpu":
+    torch.set_num_threads(os.cpu_count() or 1)
+    smoldocling_dtype = torch.float32
+else:
+    smoldocling_dtype = torch.bfloat16
+
 pipeline_options = PdfPipelineOptions(enable_remote_services=True)
 converter = DocumentConverter(
     format_options={
@@ -28,7 +36,7 @@ pipeline = PPStructureV3()
 processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
 model = AutoModelForVision2Seq.from_pretrained(
     "ds4sd/SmolDocling-256M-preview",
-    torch_dtype=torch.bfloat16,
+    torch_dtype=smoldocling_dtype,
     _attn_implementation="flash_attention_2" if DEVICE == "cuda" else "eager",
 ).to(DEVICE)