From 4446d4b58d06c26a39e4bcae6c1fa96284a856e7 Mon Sep 17 00:00:00 2001 From: JM Date: Sat, 27 Sep 2025 16:33:22 +0300 Subject: [PATCH 1/6] better non srgb handling --- library/train_util.py | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/library/train_util.py b/library/train_util.py index b9d08f25..a24a9beb 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -66,7 +66,7 @@ from library import custom_train_functions from library.original_unet import UNet2DConditionModel from huggingface_hub import hf_hub_download import numpy as np -from PIL import Image +from PIL import Image, ImageCms import imagesize import cv2 import safetensors.torch @@ -2490,10 +2490,34 @@ def load_arbitrary_dataset(args, tokenizer) -> MinimalDataset: def load_image(image_path, alpha=False): try: with Image.open(image_path) as image: + if getattr(image, "is_animated", False): + raise Exception("Image is animated") + + # Convert image to sRGB + icc = image.info.get('icc_profile', '') + if icc: + src_profile = ImageCms.ImageCmsProfile( BytesIO(icc) ) + srgb_profile = ImageCms.createProfile("sRGB") + try: + ImageCms.profileToProfile(image, src_profile, srgb_profile, inPlace=True) + image.info["icc_profile"] = ImageCms.ImageCmsProfile(srgb_profile).tobytes() + except Exception as e: + raise Exception( f"Could not convert to sRGB: {src_profile.profile.model} {src_profile.profile.profile_description}\n{e}" ) + if alpha: if not image.mode == "RGBA": image = image.convert("RGBA") else: + if image.mode == "P": + # Palette images with alpha are easier to handle as RGBA. + image = image.convert('RGBA') + + if "A" in image.getbands(): + # Replace transparency with white background. + bg = Image.new("RGBA", image.size, (255, 255, 255, 255) ) + bg.paste(image, mask=alpha) + image = bg.convert('RGB') + if not image.mode == "RGB": image = image.convert("RGB") img = np.array(image, np.uint8) From f73221038e77c3704f1e73393f6743fb2d7711bc Mon Sep 17 00:00:00 2001 From: JM Date: Sat, 27 Sep 2025 16:45:21 +0300 Subject: [PATCH 2/6] show filenames in errors --- library/train_util.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/library/train_util.py b/library/train_util.py index a24a9beb..91057fa8 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -2491,7 +2491,7 @@ def load_image(image_path, alpha=False): try: with Image.open(image_path) as image: if getattr(image, "is_animated", False): - raise Exception("Image is animated") + raise Exception( f"{image_path} is animated" ) # Convert image to sRGB icc = image.info.get('icc_profile', '') @@ -2502,7 +2502,7 @@ def load_image(image_path, alpha=False): ImageCms.profileToProfile(image, src_profile, srgb_profile, inPlace=True) image.info["icc_profile"] = ImageCms.ImageCmsProfile(srgb_profile).tobytes() except Exception as e: - raise Exception( f"Could not convert to sRGB: {src_profile.profile.model} {src_profile.profile.profile_description}\n{e}" ) + raise Exception( f"Could not convert {image_path} to sRGB: {src_profile.profile.model} {src_profile.profile.profile_description}\n{e}" ) if alpha: if not image.mode == "RGBA": From 8fb676cb6909ced4eac3689e7d0a388b0d0c0c84 Mon Sep 17 00:00:00 2001 From: JM Date: Sun, 28 Sep 2025 13:59:03 +0300 Subject: [PATCH 3/6] fix reused variable --- library/train_util.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/library/train_util.py b/library/train_util.py index 91057fa8..deb6dbc4 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -2514,8 +2514,9 @@ def load_image(image_path, alpha=False): if "A" in image.getbands(): # Replace transparency with white background. + alpha_layer = image.convert('RGBA').split()[-1] bg = Image.new("RGBA", image.size, (255, 255, 255, 255) ) - bg.paste(image, mask=alpha) + bg.paste( image, mask=alpha_layer ) image = bg.convert('RGB') if not image.mode == "RGB": From 1d4c578ac8aaab23cfc5e487e4bbef6cb0e95614 Mon Sep 17 00:00:00 2001 From: JM Date: Sun, 28 Sep 2025 14:04:11 +0300 Subject: [PATCH 4/6] move color profile creation inside try --- library/train_util.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/library/train_util.py b/library/train_util.py index deb6dbc4..0d294e7f 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -2496,9 +2496,9 @@ def load_image(image_path, alpha=False): # Convert image to sRGB icc = image.info.get('icc_profile', '') if icc: - src_profile = ImageCms.ImageCmsProfile( BytesIO(icc) ) - srgb_profile = ImageCms.createProfile("sRGB") try: + src_profile = ImageCms.ImageCmsProfile( BytesIO(icc) ) + srgb_profile = ImageCms.createProfile("sRGB") ImageCms.profileToProfile(image, src_profile, srgb_profile, inPlace=True) image.info["icc_profile"] = ImageCms.ImageCmsProfile(srgb_profile).tobytes() except Exception as e: From b9ddf0ce37733e2168c81e1451077aa50962c021 Mon Sep 17 00:00:00 2001 From: JM Date: Sun, 28 Sep 2025 15:10:09 +0300 Subject: [PATCH 5/6] works without ImageCms --- library/train_util.py | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/library/train_util.py b/library/train_util.py index 0d294e7f..c1b02b8c 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -66,7 +66,12 @@ from library import custom_train_functions from library.original_unet import UNet2DConditionModel from huggingface_hub import hf_hub_download import numpy as np -from PIL import Image, ImageCms +import sys +from PIL import Image +try: + from PIL import ImageCms +except: + print( "ImageCms not available. Images will not be converted to sRGB. Colours may be handled incorrectly." ) import imagesize import cv2 import safetensors.torch @@ -2494,15 +2499,16 @@ def load_image(image_path, alpha=False): raise Exception( f"{image_path} is animated" ) # Convert image to sRGB - icc = image.info.get('icc_profile', '') - if icc: - try: - src_profile = ImageCms.ImageCmsProfile( BytesIO(icc) ) - srgb_profile = ImageCms.createProfile("sRGB") - ImageCms.profileToProfile(image, src_profile, srgb_profile, inPlace=True) - image.info["icc_profile"] = ImageCms.ImageCmsProfile(srgb_profile).tobytes() - except Exception as e: - raise Exception( f"Could not convert {image_path} to sRGB: {src_profile.profile.model} {src_profile.profile.profile_description}\n{e}" ) + if "PIL.ImageCms" in sys.modules: + icc = image.info.get('icc_profile', '') + if icc: + try: + src_profile = ImageCms.ImageCmsProfile( BytesIO(icc) ) + srgb_profile = ImageCms.createProfile("sRGB") + ImageCms.profileToProfile(image, src_profile, srgb_profile, inPlace=True) + image.info["icc_profile"] = ImageCms.ImageCmsProfile(srgb_profile).tobytes() + except Exception as e: + raise Exception( f"Could not convert {image_path} to sRGB: {src_profile.profile.model} {src_profile.profile.profile_description}\n{e}" ) if alpha: if not image.mode == "RGBA": From c19c90d66e6c39bb160c01124318cb1f84a45c77 Mon Sep 17 00:00:00 2001 From: JM Date: Sun, 5 Oct 2025 15:33:41 +0300 Subject: [PATCH 6/6] replace errors with warnings --- library/train_util.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/library/train_util.py b/library/train_util.py index c1b02b8c..33af29f2 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -2496,7 +2496,7 @@ def load_image(image_path, alpha=False): try: with Image.open(image_path) as image: if getattr(image, "is_animated", False): - raise Exception( f"{image_path} is animated" ) + logger.warning( f"{image_path} is animated" ) # Convert image to sRGB if "PIL.ImageCms" in sys.modules: @@ -2508,7 +2508,7 @@ def load_image(image_path, alpha=False): ImageCms.profileToProfile(image, src_profile, srgb_profile, inPlace=True) image.info["icc_profile"] = ImageCms.ImageCmsProfile(srgb_profile).tobytes() except Exception as e: - raise Exception( f"Could not convert {image_path} to sRGB: {src_profile.profile.model} {src_profile.profile.profile_description}\n{e}" ) + logger.warning( f"Could not convert {image_path} to sRGB: {src_profile.profile.model} {src_profile.profile.profile_description}\n{e}" ) if alpha: if not image.mode == "RGBA":