JPEG XL support

modules/api/helpers.py

@@ -75,6 +75,13 @@ def save_image(image, fn, ext):
             image = image.point(lambda p: p * 0.0038910505836576).convert("RGB")
         exif_bytes = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") } })
         image.save(fn, format=image_format, quality=shared.opts.jpeg_quality, lossless=shared.opts.webp_lossless, exif=exif_bytes)
+    elif image_format == 'JXL':
+        if image.mode == 'I;16':
+            image = image.point(lambda p: p * 0.0038910505836576).convert("RGB")
+        elif image.mode not in {"RGB", "RGBA"}:
+            image = image.convert("RGBA")
+        exif_bytes = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(parameters or "", encoding="unicode") } })
+        image.save(fn, format=image_format, quality=shared.opts.jpeg_quality, lossless=shared.opts.webp_lossless, exif=exif_bytes)
     else:
         # shared.log.warning(f'Unrecognized image format: {extension} attempting save as {image_format}')
         image.save(fn, format=image_format, quality=shared.opts.jpeg_quality)

modules/generation_parameters_copypaste.py

@@ -67,6 +67,8 @@ def image_from_url_text(filedata):
         filedata = filedata[len("data:image/webp;base64,"):]
     if filedata.startswith("data:image/jpeg;base64,"):
         filedata = filedata[len("data:image/jpeg;base64,"):]
+    if filedata.startswith("data:image/jxl;base64,"):
+        filedata = filedata[len("data:image/jxl;base64,"):]
     filedata = base64.decodebytes(filedata.encode('utf-8'))
     image = Image.open(io.BytesIO(filedata))
     images.read_info_from_image(image)

modules/gr_tempdir.py

@@ -93,7 +93,7 @@ def cleanup_tmpdr():
     for root, _dirs, files in os.walk(temp_dir, topdown=False):
         for name in files:
             _, extension = os.path.splitext(name)
-            if extension != ".png" and extension != ".jpg" and extension != ".webp":
+            if extension not in {".png", ".jpg", ".webp", ".jxl"}:
                 continue
             filename = os.path.join(root, name)
             os.remove(filename)

modules/images.py

@@ -70,6 +70,14 @@ def atomically_save_image():
             save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'lossless': shared.opts.webp_lossless }
             if shared.opts.image_metadata:
                 save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
+        elif image_format == 'JXL':
+            if image.mode == 'I;16':
+                image = image.point(lambda p: p * 0.0038910505836576).convert("RGB")
+            elif image.mode not in {"RGB", "RGBA"}:
+                image = image.convert("RGBA")
+            save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'lossless': shared.opts.webp_lossless }
+            if shared.opts.image_metadata:
+                save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
         else:
             save_args = { 'quality': shared.opts.jpeg_quality }
         try:

modules/loader.py

@@ -72,6 +72,9 @@
 logging.getLogger("diffusers.loaders.single_file").setLevel(logging.ERROR)
 timer.startup.record("diffusers")
 
+import pillow_jxl # pylint: disable=W0611,C0411
+from PIL import Image # pylint: disable=W0611,C0411
+timer.startup.record("pillow")
 
 # patch different progress bars
 import tqdm as tqdm_lib # pylint: disable=C0411

modules/shared.py

@@ -679,7 +679,7 @@ def get_default_modes():
 options_templates.update(options_section(('saving-images', "Image Options"), {
     "keep_incomplete": OptionInfo(True, "Keep incomplete images"),
     "samples_save": OptionInfo(True, "Save all generated images"),
-    "samples_format": OptionInfo('jpg', 'File format', gr.Dropdown, {"choices": ["jpg", "png", "webp", "tiff", "jp2"]}),
+    "samples_format": OptionInfo('jpg', 'File format', gr.Dropdown, {"choices": ["jpg", "png", "webp", "tiff", "jp2", "jxl"]}),
     "jpeg_quality": OptionInfo(90, "Image quality", gr.Slider, {"minimum": 1, "maximum": 100, "step": 1}),
     "img_max_size_mp": OptionInfo(1000, "Maximum image size (MP)", gr.Slider, {"minimum": 100, "maximum": 2000, "step": 1}),
     "webp_lossless": OptionInfo(False, "WebP lossless compression"),
@@ -694,7 +694,7 @@ def get_default_modes():
     "save_log_fn": OptionInfo("", "Append image info JSON file", component_args=hide_dirs),
     "image_sep_grid": OptionInfo("<h2>Grid Options</h2>", "", gr.HTML),
     "grid_save": OptionInfo(True, "Save all generated image grids"),
-    "grid_format": OptionInfo('jpg', 'File format', gr.Dropdown, {"choices": ["jpg", "png", "webp", "tiff", "jp2"]}),
+    "grid_format": OptionInfo('jpg', 'File format', gr.Dropdown, {"choices": ["jpg", "png", "webp", "tiff", "jp2", "jxl"]}),
     "n_rows": OptionInfo(-1, "Row count", gr.Slider, {"minimum": -1, "maximum": 16, "step": 1}),
     "grid_background": OptionInfo("#000000", "Grid background color", gr.ColorPicker, {}),
     "font": OptionInfo("", "Font file"),

modules/ui.py

@@ -16,6 +16,7 @@
 mimetypes.init()
 mimetypes.add_type('application/javascript', '.js')
 mimetypes.add_type('image/webp', '.webp')
+mimetypes.add_type('image/jxl', '.jxl')
 log = shared.log
 opts = shared.opts
 cmd_opts = shared.cmd_opts

modules/ui_extra_networks.py

@@ -331,7 +331,7 @@ def find_preview_file(self, path):
             return 'html/card-no-preview.png'
         if os.path.join('models', 'Reference') in path:
             return path
-        exts = ["jpg", "jpeg", "png", "webp", "tiff", "jp2"]
+        exts = ["jpg", "jpeg", "png", "webp", "tiff", "jp2", "jxl"]
         reference_path = os.path.abspath(os.path.join('models', 'Reference'))
         files = list(files_cache.list_files(reference_path, ext_filter=exts, recursive=False))
         if shared.opts.diffusers_dir in path:
@@ -360,7 +360,7 @@ def update_all_previews(self, items):
         t0 = time.time()
         reference_path = os.path.abspath(os.path.join('models', 'Reference'))
         possible_paths = list(set([os.path.dirname(item['filename']) for item in items] + [reference_path]))
-        exts = ["jpg", "jpeg", "png", "webp", "tiff", "jp2"]
+        exts = ["jpg", "jpeg", "png", "webp", "tiff", "jp2", "jxl"]
         all_previews = list(files_cache.list_files(*possible_paths, ext_filter=exts, recursive=False))
         all_previews_fn = [os.path.basename(x) for x in all_previews]
         for item in items:
@@ -680,7 +680,7 @@ def fn_save_img(image):
         return image
 
     def fn_delete_img(_image):
-        preview_extensions = ["jpg", "jpeg", "png", "webp", "tiff", "jp2"]
+        preview_extensions = ["jpg", "jpeg", "png", "webp", "tiff", "jp2", "jxl"]
         fn = os.path.splitext(ui.last_item.filename)[0]
         for file in [f'{fn}{mid}{ext}' for ext in preview_extensions for mid in ['.thumb.', '.preview.', '.']]:
             if os.path.exists(file):

requirements.txt

@@ -55,6 +55,7 @@ tokenizers==0.20.3
 transformers==4.46.2
 urllib3==1.26.19
 Pillow==10.4.0
+pillow-jxl-plugin==1.3.0
 timm==0.9.16
 pydantic==1.10.15
 pyparsing==3.1.4