Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Feature Request: offload memory usage on storage #26

Open
aont opened this issue Jan 6, 2023 · 0 comments
Open

Feature Request: offload memory usage on storage #26

aont opened this issue Jan 6, 2023 · 0 comments

Comments

@aont
Copy link

aont commented Jan 6, 2023

I used Real-ESRGAN on Google Colab and wanted to enlarge the image 16x.
I tried the tile option, but it inevitably crashed in the process.
I examined the code and found that the paste_faces_to_input_image in facexlib's face_restoration_helper.py was allocating a huge amount of memory.
So I thought of offloading memory usage on storage and implemented it using numpy's memmap and numexpr.
This greatly reduced the memory consumption and made it possible to zoom in 16x on the image with Google Colab.

Would you like to incorporate this feature if you like?

--- face_restoration_helper_a.py
+++ face_restoration_helper_b.py
@@ -279,6 +279,7 @@
         self.restored_faces.append(face)
 
     def paste_faces_to_input_image(self, save_path=None, upsample_img=None):
+        import numexpr as ne
         h, w, _ = self.input_img.shape
         h_up, w_up = int(h * self.upscale_factor), int(w * self.upscale_factor)
 
@@ -288,6 +289,12 @@
         else:
             upsample_img = cv2.resize(upsample_img, (w_up, h_up), interpolation=cv2.INTER_LANCZOS4)
 
+        upsample_img_orig = upsample_img
+        upsample_img_memmap_fn = "/content/upsample_img.npy"
+        upsample_img = np.memmap(upsample_img_memmap_fn, shape=upsample_img.shape, dtype="float64", mode="w+")
+        upsample_img[:] = upsample_img_orig
+        del upsample_img_orig
+
         assert len(self.restored_faces) == len(
             self.inverse_affine_matrices), ('length of restored_faces and affine_matrices are different.')
         for restored_face, inverse_affine in zip(self.restored_faces, self.inverse_affine_matrices):
@@ -352,12 +359,14 @@
                 upsample_img = inv_soft_mask * pasted_face + (1 - inv_soft_mask) * upsample_img[:, :, 0:3]
                 upsample_img = np.concatenate((upsample_img, alpha), axis=2)
             else:
-                upsample_img = inv_soft_mask * pasted_face + (1 - inv_soft_mask) * upsample_img
+                ne.evaluate("inv_soft_mask * pasted_face + (1 - inv_soft_mask) * upsample_img", out=upsample_img)
 
         if np.max(upsample_img) > 256:  # 16-bit image
             upsample_img = upsample_img.astype(np.uint16)
         else:
             upsample_img = upsample_img.astype(np.uint8)
+        import os
+        os.unlink(upsample_img_memmap_fn)
         if save_path is not None:
             path = os.path.splitext(save_path)[0]
             save_path = f'{path}.{self.save_ext}'
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant