Skip to content

Commit

Permalink
JPEG XL support
Browse files Browse the repository at this point in the history
  • Loading branch information
Disty0 committed Jan 11, 2025
1 parent 389bb78 commit 705556d
Show file tree
Hide file tree
Showing 9 changed files with 31 additions and 6 deletions.
1 change: 1 addition & 0 deletions installer.py
Original file line number Diff line number Diff line change
Expand Up @@ -1050,6 +1050,7 @@ def install_optional():
install('basicsr')
install('gfpgan')
install('clean-fid')
install('pillow-jxl-plugin==1.3.1', ignore=True)
install('optimum-quanto=0.2.6', ignore=True)
install('bitsandbytes==0.45.0', ignore=True)
install('pynvml', ignore=True)
Expand Down
7 changes: 7 additions & 0 deletions modules/api/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
2 changes: 2 additions & 0 deletions modules/generation_parameters_copypaste.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion modules/gr_tempdir.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,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)
8 changes: 8 additions & 0 deletions modules/images.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,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:
Expand Down
6 changes: 6 additions & 0 deletions modules/loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,12 @@
logging.getLogger("diffusers.loaders.single_file").setLevel(logging.ERROR)
timer.startup.record("diffusers")

try:
import pillow_jxl # pylint: disable=W0611,C0411
except:
pass
from PIL import Image # pylint: disable=W0611,C0411
timer.startup.record("pillow")

# patch different progress bars
import tqdm as tqdm_lib # pylint: disable=C0411
Expand Down
4 changes: 2 additions & 2 deletions modules/shared.py
Original file line number Diff line number Diff line change
Expand Up @@ -701,7 +701,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"),
Expand All @@ -716,7 +716,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"),
Expand Down
1 change: 1 addition & 0 deletions modules/ui.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions modules/ui_extra_networks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -685,7 +685,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):
Expand Down

0 comments on commit 705556d

Please sign in to comment.