@@ -2,12 +2,15 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from pathlib import Path from typing import Any +import filetype import requests from django.core.files import File from messages.models import Message +from ml_model.exceptions import CorruptedFileError, FileExtensionNotSupported from ml_model.services.base import SimpleService from ml_model.tasks import replicate_run from payments.exceptions.insufficient_balance import InsufficientBalance @@ -50,6 +53,17 @@ class Grok_Image(SimpleService): } ) if image := input_message.file: + kind = filetype.guess(image.read(20)) + image.seek(0) + name_ext = Path(image.name).suffix[1:].upper() + if name_ext == 'DNG': + extension = 'DNG' + elif not kind: + raise CorruptedFileError + else: + extension = kind.extension.upper() + if extension not in (extensions := ['DNG', 'JPG', 'JPEG', 'PNG', 'WEBP']): + raise FileExtensionNotSupported(extensions) callback_data.update({'image': image}) start_time = time.time() images = replicate_run('xai/grok-imagine-image', callback_data) @@ -11,11 +11,14 @@ from typing import IO, Any, Dict import deepl import httpx +import rawpy import redis import replicate import requests +from PIL import Image from celery import shared_task from deepl.translator import TextResult +from django.core.files.base import ContentFile from django.core.files.storage import Storage from django.db.models.fields.files import FieldFile from django.utils.translation import gettext as _ @@ -121,18 +124,33 @@ 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': + 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() + elif suffix == '.dng': + image.open('rb') + image.seek(0) + try: + with rawpy.imread(image) as raw: + rgb = raw.postprocess() + buf = BytesIO() + Image.fromarray(rgb).save(buf, format='PNG') + buf.seek(0) + png_name = str(PurePosixPath(image.name).with_suffix('.png')) + saved_name = image.storage.save(png_name, ContentFile(buf.getvalue())) + finally: + image.close() + else: 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: - image.close() - try: image_url = image.storage.url(saved_name) except Exception: @@ -166,7 +184,17 @@ def replicate_run(callback_url: str, payload: dict[str, Any]): and 'is not supported' in error_text ): raise FileExtensionNotSupported(('MP3', 'WAV')) from exc - if any(error in error_text for error in ('E005', 'E006', 'sexual', 'NSFW')): + if any( + error in error_text + for error in ( + 'E005', + 'E006', + 'sexual', + 'NSFW', + 'illegal material', + 'Content flagged', + ) + ): raise RequestBlocked if 'PA' in error_text: raise PredictionInterruptedError @@ -49,6 +49,7 @@ dependencies = [ "pypdf2==3.0.1", "python-dateutil==2.9.0.post0", "python-docx==1.1.2", + "rawpy==0.27.0", "redis>=6.4.0", "replicate==1.0.4", "sentry-sdk[django]==2.39.0", @@ -308,6 +308,7 @@ dependencies = [ { name = "pypdf2" }, { name = "python-dateutil" }, { name = "python-docx" }, + { name = "rawpy" }, { name = "redis" }, { name = "replicate" }, { name = "sentry-sdk", extra = ["django"] }, @@ -374,6 +375,7 @@ requires-dist = [ { name = "pypdf2", specifier = "==3.0.1" }, { name = "python-dateutil", specifier = "==2.9.0.post0" }, { name = "python-docx", specifier = "==1.1.2" }, + { name = "rawpy", specifier = "==0.27.0" }, { name = "redis", specifier = ">=6.4.0" }, { name = "replicate", specifier = "==1.0.4" }, { name = "sentry-sdk", extras = ["django"], specifier = "==2.39.0" }, @@ -3152,6 +3154,29 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" }, ] +[[package]] +name = "rawpy" +version = "0.27.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/86/15/92324724209650167c8c8b0c8c0006c99d07494b9b41f7d6435a37737323/rawpy-0.27.0.tar.gz", hash = "sha256:45251e46c1d891a62919a4ac200a9828f825e3c59f89cea2f1daee0900ff15ec", size = 560718, upload-time = "2026-05-07T08:29:23.974Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/97/46415db86271d977390f607d7a7733d86abb65f2dee29ab81441e092b3d4/rawpy-0.27.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8722fcd00b242e404ab8297d1f7bdd7d0bbe2a30c41a70815cbec74eb4583bed", size = 2051102, upload-time = "2026-05-07T08:28:56.466Z" }, + { url = "https://files.pythonhosted.org/packages/d0/33/206372f73d215c4a379a91081c8445f352b752eab85613445af79cc86855/rawpy-0.27.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a0192f3ce982ea06cdb82196e96dca9e9d252f8715d2b363bcbb79553e2e89e", size = 2945220, upload-time = "2026-05-07T08:28:57.99Z" }, + { url = "https://files.pythonhosted.org/packages/85/fc/6f26df1bc1663366727665f082b00b8cdd70e555e5dad43208dd58fd08bd/rawpy-0.27.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2262088bd4fd768edd7576857a9215cfff5e84d61942336401ee5d9ca8f677e0", size = 2941195, upload-time = "2026-05-07T08:28:59.61Z" }, + { url = "https://files.pythonhosted.org/packages/56/fb/3c754322c080477633644e453bdaebfeb5177249859c37302ad392a93a0e/rawpy-0.27.0-cp312-cp312-win_amd64.whl", hash = "sha256:2965d8c70af1b4ed6177a098871de1cc9204854278b9260d1e79639df40391ae", size = 913608, upload-time = "2026-05-07T08:29:01.593Z" }, + { url = "https://files.pythonhosted.org/packages/f5/02/aec7b4f6befd4752f23a5710bc7687631588006b17c4f9dc3a5198ae12b6/rawpy-0.27.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7758606d5de8497a695511fbace40cb80138329ed0ad7a530d51bdc07eef1a6b", size = 2050987, upload-time = "2026-05-07T08:29:03.374Z" }, + { url = "https://files.pythonhosted.org/packages/f0/29/ff05f8ebd08fa99c185b60b0088c7950016f0ac1e7e56042e5dbe67fd85c/rawpy-0.27.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a1f4a1846247deff8b4c84b5b400541603d48be90bd36fa6d9a1f5fada023cae", size = 2944727, upload-time = "2026-05-07T08:29:05.115Z" }, + { url = "https://files.pythonhosted.org/packages/5c/49/84f15204f80a75ccd6f36b5530878339231faa334fbc47c83a5b45de19cf/rawpy-0.27.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:45b1b08a8d6fdc0139fea4c1389c17cb6f3d01b0581e71eef1337949a486550c", size = 2940136, upload-time = "2026-05-07T08:29:06.803Z" }, + { url = "https://files.pythonhosted.org/packages/b6/18/49c498d363ba3c3935872244b27e01b514d68ea531c325cff173b849ec81/rawpy-0.27.0-cp313-cp313-win_amd64.whl", hash = "sha256:8e42d6218b507584b41ffdb86e1a6d8e2672bfe284a59df1dfdc57b3f8a27ffe", size = 913591, upload-time = "2026-05-07T08:29:08.378Z" }, + { url = "https://files.pythonhosted.org/packages/fc/67/b5aa2517dd4e1c0d0f1644e7c24c78a896df7835aeaf3f20abd31a2455b1/rawpy-0.27.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:cd80c608d741aed9dc014cdac2af3f83796bc73148dbb7b0ec1fa2ba715d7a53", size = 2050437, upload-time = "2026-05-07T08:29:09.918Z" }, + { url = "https://files.pythonhosted.org/packages/8e/23/ec4e4dc5550e96d931ea5d0c3266fc42a9ce18b621572de24f0e683d02c9/rawpy-0.27.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cb018c5d3273d09245bf8d4a7c6da9f2c4191b1ed04f45af4fcf978e2b2372ae", size = 2945830, upload-time = "2026-05-07T08:29:11.613Z" }, + { url = "https://files.pythonhosted.org/packages/b8/48/64e97637398494e6b74f8388c5a17d726219a5ddc9f7c283ff9f77849a05/rawpy-0.27.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a949e9258af1d915b4316b77c8de9d9637d3624e6fe85a5bfb6cb4de15f485d6", size = 2940100, upload-time = "2026-05-07T08:29:13.647Z" }, + { url = "https://files.pythonhosted.org/packages/4d/01/1854b2f789d7a6d94152ba451bdaadae0010d4247099d9c2719b381d4b6b/rawpy-0.27.0-cp314-cp314-win_amd64.whl", hash = "sha256:3cac371c3b6302eb88bd1f33d43c7a92b5e553fdff14258b564f81e173a68c26", size = 940931, upload-time = "2026-05-07T08:29:15.604Z" }, +] + [[package]] name = "redis" version = "6.4.0"