From 5e89252ab693aa9d06121b80dc00f359f18d6a09 Mon Sep 17 00:00:00 2001
From: Max Tretikov <max@tretikov.com>
Date: Thu, 13 Jun 2024 23:40:44 -0600
Subject: [PATCH] Fix cv2 lint errors

---
 comfy/images.py   | 5 +++--
 comfy/open_exr.py | 8 +++++---
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/comfy/images.py b/comfy/images.py
index b674c8867..85f6529b6 100644
--- a/comfy/images.py
+++ b/comfy/images.py
@@ -1,5 +1,6 @@
 import os.path
 from contextlib import contextmanager
+from typing import Iterator
 
 import cv2
 from PIL import Image
@@ -8,11 +9,11 @@ from . import node_helpers
 
 
 def _open_exr(exr_path) -> Image.Image:
-    return Image.fromarray(cv2.imread(exr_path, cv2.IMREAD_COLOR))
+    return Image.fromarray(cv2.imread(exr_path, cv2.IMREAD_COLOR)) # pylint: disable=no-member
 
 
 @contextmanager
-def open_image(file_path: str) -> Image.Image:
+def open_image(file_path: str) -> Iterator[Image.Image]:
     _, ext = os.path.splitext(file_path)
     if ext == ".exr":
         yield _open_exr(file_path)
diff --git a/comfy/open_exr.py b/comfy/open_exr.py
index acd76f1bb..9785800ea 100644
--- a/comfy/open_exr.py
+++ b/comfy/open_exr.py
@@ -36,6 +36,8 @@ from torch import Tensor
 from .component_model.images_types import RgbMaskTuple
 
 
+read_exr = lambda fp: cv.imread(fp, cv.IMREAD_UNCHANGED).astype(np.float32) # pylint: disable=no-member
+
 def mut_srgb_to_linear(np_array) -> None:
     less = np_array <= 0.0404482362771082
     np_array[less] = np_array[less] / 12.92
@@ -49,7 +51,7 @@ def mut_linear_to_srgb(np_array) -> None:
 
 
 def load_exr(file_path: str, srgb: bool) -> RgbMaskTuple:
-    image = cv.imread(file_path, cv.IMREAD_UNCHANGED).astype(np.float32)
+    image = read_exr(file_path)
     rgb = np.flip(image[:, :, :3], 2).copy()
     if srgb:
         mut_linear_to_srgb(rgb)
@@ -64,7 +66,7 @@ def load_exr(file_path: str, srgb: bool) -> RgbMaskTuple:
 
 
 def load_exr_latent(file_path: str) -> Tuple[Tensor]:
-    image = cv.imread(file_path, cv.IMREAD_UNCHANGED).astype(np.float32)
+    image = read_exr(file_path)
     image = image[:, :, np.array([2, 1, 0, 3])]
     image = torch.unsqueeze(torch.from_numpy(image), 0)
     image = torch.movedim(image, -1, 1)
@@ -83,4 +85,4 @@ def save_exr(images: Tensor, filepaths_batched: Sequence[str], colorspace="linea
         bgr[:, :, :, 3] = np.clip(1 - linear[:, :, :, 3], 0, 1)  # invert alpha
 
     for i in range(len(linear.shape[0])):
-        cv.imwrite(filepaths_batched[i], bgr[i])
+        cv.imwrite(filepaths_batched[i], bgr[i]) # pylint: disable=no-member