@@ -11,7 +11,7 @@ import openpyxl from io import BytesIO -from PIL import Image +from PIL import Image, PngImagePlugin from django.core.files.images import get_image_dimensions from django.db.models.fields.files import FieldFile from django.utils.translation import gettext as _ @@ -136,7 +136,8 @@ class ImageFileProcessingService: try: return func(self, *args, **kwargs) finally: - self.image.seek(0) + if not self.image.closed: + self.image.seek(0) return wrapper @@ -164,12 +165,20 @@ class ImageFileProcessingService: except Image.DecompressionBombError: raise ImageTooLargeError(max_pixels) + @__reset_image def get_normalized_image(self, file_bytes: bytes) -> BytesIO: - normalized_image = BytesIO(file_bytes) - with Image.open(normalized_image) as source_image: - img = source_image.convert('RGBA') - normalized_image = BytesIO() - img.save(normalized_image, format='PNG') - img.close() - normalized_image.seek(0) - return normalized_image + old_limit = PngImagePlugin.MAX_TEXT_CHUNK + PngImagePlugin.MAX_TEXT_CHUNK = max(old_limit, 100 * 1024 * 1024) + try: + normalized_image = BytesIO(file_bytes) + with Image.open(normalized_image) as source_image: + img = source_image.convert('RGBA') + img.info.clear() + normalized_image = BytesIO() + img.save(normalized_image, format='PNG') + img.close() + normalized_image.seek(0) + return normalized_image + finally: + PngImagePlugin.MAX_TEXT_CHUNK = old_limit + \ No newline at end of file @@ -14,8 +14,10 @@ import httpx import redis import replicate import requests +from PIL import Image from celery import shared_task from deepl.translator import TextResult +from django.core.files import File from django.core.files.storage import Storage from django.db.models.fields.files import FieldFile from django.utils.translation import gettext as _ @@ -42,6 +44,7 @@ from ml_model.exceptions import ( RequestBlocked, ServiceHighDemandError, ) +from ml_model.services.FileService import ImageFileProcessingService from poller.models import Proxy logger = logging.getLogger(__name__) @@ -121,26 +124,53 @@ def transcript_audio(payload: dict[str, Any]): def _prepare_replicate_image(image: FieldFile) -> tuple[str, tuple[Storage, str] | None]: - if PurePosixPath(image.name).suffix.lower() != '.jfif': - return image.url, None - - # JFIF already contains JPEG data, so copy it without decoding or re-encoding. - image.open('rb') - image.seek(0) - jpeg_name = str(PurePosixPath(image.name).with_suffix('.jpg')) - try: - saved_name = image.storage.save(jpeg_name, image) - finally: + suffix = PurePosixPath(image.name).suffix.lower() + + if suffix == '.jfif': + # JFIF already contains JPEG data, so copy it without decoding or re-encoding. + image.open('rb') + image.seek(0) + jpeg_name = str(PurePosixPath(image.name).with_suffix('.jpg')) + try: + saved_name = image.storage.save(jpeg_name, image) + finally: + image.close() + + try: + image_url = image.storage.url(saved_name) + except Exception: + image.storage.delete(saved_name) + + raise + + return image_url, (image.storage, saved_name) + + if suffix == '.png': + image.open('rb') + image.seek(0) + file_bytes = image.read() image.close() - - try: - image_url = image.storage.url(saved_name) - except Exception: - image.storage.delete(saved_name) - - raise - - return image_url, (image.storage, saved_name) + try: + with Image.open(BytesIO(file_bytes)) as im: + im.verify() + return image.url, None + except ValueError as exc: + if 'MAX_TEXT_CHUNK' not in str(exc): + raise + + clean = ImageFileProcessingService(image).get_normalized_image(file_bytes) + + path = PurePosixPath(image.name) + png_name = str(path.with_stem(f'{path.stem}_clean')) + saved_name = image.storage.save(png_name, File(clean, name=PurePosixPath(png_name).name)) + + try: + return image.storage.url(saved_name), (image.storage, saved_name) + except Exception: + image.storage.delete(saved_name) + raise + + return image.url, None @shared_task