From d599394f6086b7291535e510e67b30c18fef4c07 Mon Sep 17 00:00:00 2001
From: ddPn08 <pyn.g.1125@gmail.com>
Date: Sat, 8 Jul 2023 15:47:56 +0900
Subject: [PATCH] support avif

---
 library/train_util.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/library/train_util.py b/library/train_util.py
index 62cd145e..a3017694 100644
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -89,6 +89,12 @@ STEP_DIFFUSERS_DIR_NAME = "{}-step{:08d}"
 
 IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".PNG", ".JPG", ".JPEG", ".WEBP", ".BMP"]
 
+try:
+    import pillow_avif
+    IMAGE_EXTENSIONS.extend([".avif", ".AVIF"])
+except:
+    pass
+
 
 class ImageInfo:
     def __init__(self, image_key: str, num_repeats: int, caption: str, is_reg: bool, absolute_path: str) -> None: