From 7653d21cf5d1bc392d28e2786488465a308117e9 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Tue, 28 Nov 2023 07:37:39 +1100
Subject: [PATCH] feat(backend): rename realesrgan class & upscale method

---
 invokeai/app/invocations/upscale.py                  | 6 +++---
 invokeai/backend/image_util/realesrgan/realesrgan.py | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/invokeai/app/invocations/upscale.py b/invokeai/app/invocations/upscale.py
index d2f53a5b9f..9397ea3655 100644
--- a/invokeai/app/invocations/upscale.py
+++ b/invokeai/app/invocations/upscale.py
@@ -11,7 +11,7 @@ from pydantic import ConfigDict
 
 from invokeai.app.invocations.primitives import ImageField, ImageOutput
 from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
-from invokeai.backend.image_util.realesrgan.realesrgan import RealESRGANer
+from invokeai.backend.image_util.realesrgan.realesrgan import RealESRGAN
 from invokeai.backend.util.devices import choose_torch_device
 
 from .baseinvocation import BaseInvocation, InputField, InvocationContext, WithMetadata, WithWorkflow, invocation
@@ -92,7 +92,7 @@ class ESRGANInvocation(BaseInvocation, WithWorkflow, WithMetadata):
 
         esrgan_model_path = Path(f"core/upscaling/realesrgan/{self.model_name}")
 
-        upsampler = RealESRGANer(
+        upscaler = RealESRGAN(
             scale=netscale,
             model_path=models_path / esrgan_model_path,
             model=rrdbnet_model,
@@ -107,7 +107,7 @@ class ESRGANInvocation(BaseInvocation, WithWorkflow, WithMetadata):
         # We can pass an `outscale` value here, but it just resizes the image by that factor after
         # upscaling, so it's kinda pointless for our purposes. If you want something other than 4x
         # upscaling, you'll need to add a resize node after this one.
-        upscaled_image = upsampler.enhance(cv2_image)
+        upscaled_image = upscaler.upscale(cv2_image)
 
         # back to PIL
         pil_image = Image.fromarray(cv2.cvtColor(upscaled_image, cv2.COLOR_BGR2RGB)).convert("RGBA")
diff --git a/invokeai/backend/image_util/realesrgan/realesrgan.py b/invokeai/backend/image_util/realesrgan/realesrgan.py
index 042b1b2f4e..4d41dabc1e 100644
--- a/invokeai/backend/image_util/realesrgan/realesrgan.py
+++ b/invokeai/backend/image_util/realesrgan/realesrgan.py
@@ -32,7 +32,7 @@ class ImageMode(str, Enum):
     RGBA = "RGBA"
 
 
-class RealESRGANer:
+class RealESRGAN:
     """A helper class for upsampling images with RealESRGAN.
 
     Args:
@@ -202,7 +202,7 @@ class RealESRGANer:
         return self.output
 
     @torch.no_grad()
-    def enhance(self, img: MatLike, esrgan_alpha_upscale: bool = True) -> npt.NDArray[Any]:
+    def upscale(self, img: MatLike, esrgan_alpha_upscale: bool = True) -> npt.NDArray[Any]:
         np_img = img.astype(np.float32)
         alpha: Optional[np.ndarray] = None
         if np.max(np_img) > 256: