Add support for JPEG-XL (file extension: jxl) (#1413)

* Add support for JPEG-XL (file extension: jxl)

Author: StableLlama

INSTALL.md

@@ -41,6 +41,8 @@ poetry install
 poetry install -C install/rocm
 ```
 
+**Note:** When you want to use JPEG XL images, you have to install the optional `jxl` dependency by running `poetry install --with jxl`.
+
 #### NVIDIA Hopper / Blackwell follow-up steps
 
 Optionally, Hopper (or newer) equipment can make use of FlashAttention3 for improved inference and training performance when making use of `torch.compile`

helpers/image_manipulation/load.py

@@ -5,6 +5,10 @@
 
 from io import BytesIO
 from typing import Union, IO, Any
+try:
+    import pillow_jxl
+except ModuleNotFoundError:
+    pass
 from PIL import Image, PngImagePlugin
 
 logger = logging.getLogger(__name__)

helpers/image_manipulation/training_sample.py

@@ -1,3 +1,7 @@
+try:
+    import pillow_jxl
+except ModuleNotFoundError:
+    pass
 from PIL import Image
 from PIL.ImageOps import exif_transpose
 from helpers.multiaspect.image import MultiaspectImage, resize_helpers

helpers/models/sdxl/pipeline.py

@@ -13,6 +13,10 @@
 # limitations under the License.
 
 import inspect
+try:
+    import pillow_jxl
+except ModuleNotFoundError:
+    pass
 import PIL
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 from diffusers.callbacks import PipelineCallback, MultiPipelineCallbacks

helpers/training/__init__.py

@@ -22,22 +22,18 @@
         # Hopper! Or blackwell+.
         quantised_precision_levels.append("fp8-torchao")
 
+try:
+    import pillow_jxl
+except ModuleNotFoundError:
+    pass
+from PIL import Image
+
+supported_extensions = Image.registered_extensions()
 image_file_extensions = set(
-    [
-        "jpg",
-        "jpeg",
-        "png",
-        "webp",
-        "bmp",
-        "tiff",
-        "tif",
-        "mp4",
-        "avi",
-        "gif",
-        "mov",
-        "webm",
-    ]
+    ext.lower().lstrip(".") for ext, img_format in supported_extensions.items()
+    if img_format in Image.OPEN
 )
+
 video_file_extensions = set(["mp4", "avi", "gif", "mov", "webm"])
 
 lycoris_defaults = {

helpers/training/validation.py

@@ -9,6 +9,10 @@
 from helpers.training.wrappers import unwrap_model
 from helpers.models.common import VideoModelFoundation, ImageModelFoundation
 from helpers.models.common import ModelFoundation
+try:
+    import pillow_jxl
+except ModuleNotFoundError:
+    pass
 from PIL import Image
 from helpers.training.state_tracker import StateTracker
 from helpers.models.common import PredictionTypes, PipelineTypes

inference_comparison.py

@@ -5,6 +5,10 @@
 from diffusers import AutoPipelineForText2Image
 from torch import manual_seed, float16
 import os
+try:
+    import pillow_jxl
+except ModuleNotFoundError:
+    pass
 from PIL import Image, ImageDraw, ImageFont
 from helpers.prompts import prompts
 

install/apple/pyproject.toml

@@ -51,6 +51,8 @@ prodigy-plus-schedule-free = "^1.9.0"
 imageio-ffmpeg = "^0.6.0"
 imageio = {extras = ["pyav"], version = "^2.37.0"}
 
+[tool.poetry.group.jxl.dependencies]
+pillow-jxl-plugin = "^1.3.1"
 
 [build-system]
 requires = ["poetry-core"]

install/rocm/pyproject.toml

@@ -50,6 +50,8 @@ huggingface-hub = "^0.29.1"
 imageio-ffmpeg = "^0.6.0"
 imageio = {extras = ["pyav"], version = "^2.37.0"}
 
+[tool.poetry.group.jxl.dependencies]
+pillow-jxl-plugin = "^1.3.1"
 
 [build-system]
 requires = ["poetry-core"]
@@ -58,4 +60,4 @@ build-backend = "poetry.core.masonry.api"
 [[tool.poetry.source]]
 #priority = "explicit"
 name = "pytorch-rocm"
-url = "https://download.pytorch.org/whl/test/rocm6.3"
+url = "https://download.pytorch.org/whl/test/rocm6.3"

poetry.lock

@@ -58,7 +58,8 @@ huggingface-hub = "^0.30.2"
 imageio-ffmpeg = "^0.6.0"
 imageio = {extras = ["pyav"], version = "^2.37.0"}
 
-
+[tool.poetry.group.jxl.dependencies]
+pillow-jxl-plugin = "^1.3.1"
 
 [build-system]
 requires = ["poetry-core", "setuptools", "wheel", "torch"]

pyproject.toml

@@ -1,6 +1,10 @@
 import unittest
 import pandas as pd
 from unittest.mock import patch, Mock, MagicMock
+try:
+    import pillow_jxl
+except ModuleNotFoundError:
+    pass
 from PIL import Image
 from pathlib import Path
 from helpers.multiaspect.dataset import MultiAspectDataset

tests/test_dataset.py

@@ -4,6 +4,10 @@
 logger.setLevel(os.environ.get("SIMPLETUNER_LOG_LEVEL", logging.INFO))
 from unittest.mock import patch
 from unittest.mock import Mock, MagicMock
+try:
+    import pillow_jxl
+except ModuleNotFoundError:
+    pass
 from PIL import Image
 from io import BytesIO
 from helpers.multiaspect.image import MultiaspectImage

tests/test_image.py

@@ -1,5 +1,9 @@
 import unittest, os, logging
 from math import ceil
+try:
+    import pillow_jxl
+except ModuleNotFoundError:
+    pass
 from PIL import Image
 from unittest import skip
 from unittest.mock import Mock, MagicMock, patch