diff --git a/app/ml/siglip.py b/app/ml/siglip.py index 693b28b..4fa6951 100644 --- a/app/ml/siglip.py +++ b/app/ml/siglip.py @@ -3,7 +3,11 @@ from __future__ import annotations import os import numpy as np -from PIL import Image +from PIL import Image, ImageFile + +# Mirror wd14.py: tolerate minutely-truncated source images so embedding +# doesn't fail on the same images WD14 successfully tagged. +ImageFile.LOAD_TRUNCATED_IMAGES = True # Defer torch/transformers imports to lazily-loaded functions to allow # importing MODEL_VERSION in non-ML-worker contexts (e.g., web container diff --git a/app/ml/wd14.py b/app/ml/wd14.py index 640e0cb..9cabea3 100644 --- a/app/ml/wd14.py +++ b/app/ml/wd14.py @@ -6,7 +6,13 @@ from typing import Iterable import numpy as np import onnxruntime as ort -from PIL import Image +from PIL import Image, ImageFile + +# Some images in the library are minutely truncated (e.g. 6 bytes short of +# the JPEG end-of-image marker). The model doesn't care about the last few +# pixels, but PIL's default strict load raises OSError. Tolerate them so a +# single corrupt tail doesn't block tagging the rest of the image. +ImageFile.LOAD_TRUNCATED_IMAGES = True MODEL_VERSION = os.environ.get('WD14_MODEL_VERSION', 'wd-eva02-large-tagger-v3') _MODEL_DIR = os.environ.get('ML_MODEL_DIR', '/models')