From f3c294d2da1520ea7374cd844149fd8f3a46bc23 Mon Sep 17 00:00:00 2001
From: sdbds <865105819@qq.com>
Date: Mon, 11 Nov 2024 09:13:23 +0800
Subject: [PATCH] update

---
 library/train_util.py | 16 ++++++++++------
 requirements-uv.txt   |  2 +-
 requirements.txt      |  2 +-
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/library/train_util.py b/library/train_util.py
index 2ba0bc73f..82cbbc20d 100644
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -1420,12 +1420,16 @@ def get_image_size(self, image_path):
         image_size = imagesize.get(image_path)
         if image_size[0] <= 0:
             # imagesize doesn't work for some images, so use cv2
-            img = cv2.imread(image_path)
-            if img is not None:
-                image_size = (img.shape[1], img.shape[0])
-            else:
-                logger.warning(f"failed to get image size: {image_path}")
-                image_size = (0, 0)
+            # img = cv2.imread(image_path)
+            # if img is not None:
+            #     image_size = (img.shape[1], img.shape[0])
+            
+            with Image.open(image_path) as img:
+                if img is not None:
+                    image_size = img.size
+                else:
+                    logger.warning(f"failed to get image size: {image_path}")
+                    image_size = (0, 0)
         return image_size
 
     def load_image_with_face_info(self, subset: BaseSubset, image_path: str, alpha_mask=False):
diff --git a/requirements-uv.txt b/requirements-uv.txt
index fe072812a..c92c6ec4b 100644
--- a/requirements-uv.txt
+++ b/requirements-uv.txt
@@ -95,7 +95,7 @@ h5py==3.11.0
     # via
     #   keras
     #   tensorflow-intel
-heavyball==0.9.0
+heavyball==0.10.2
     # via -r ./requirements.txt
 huggingface-hub==0.24.5
     # via
diff --git a/requirements.txt b/requirements.txt
index 398106cb7..64e5ed326 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -11,7 +11,7 @@ prodigyopt
 lion-pytorch
 scipy
 dadaptation
-heavyball>=0.9.0
+heavyball>=0.10.2
 wandb==0.18.0
 pillow-avif-plugin
 schedulefree==1.3.0