Browse Source

Adding graceful and informative gpu out of memory error handling.

Jason Antic 5 years ago
parent
commit
696001361a
2 changed files with 15 additions and 5 deletions
  1. 11 4
      deoldify/filters.py
  2. 4 1
      deoldify/visualize.py

+ 11 - 4
deoldify/filters.py

@@ -47,10 +47,17 @@ class BaseFilter(IFilter):
         x = x.to(self.device)
         x.div_(255)
         x, y = self.norm((x, x), do_x=True)
-
-        result = self.learn.pred_batch(
-            ds_type=DatasetType.Valid, batch=(x[None], y[None]), reconstruct=True
-        )
+        
+        try:
+            result = self.learn.pred_batch(
+                ds_type=DatasetType.Valid, batch=(x[None], y[None]), reconstruct=True
+            )
+        except RuntimeError as rerr:
+            if 'memory' not in str(rerr):
+                raise rerr
+            print('Warning: render_factor was set too high, and out of memory error resulted. Returning original image.')
+            return model_image
+            
         out = result[0]
         out = self.denorm(out.px, do_x=False)
         out = image2np(out * 255).astype(np.uint8)

+ 4 - 1
deoldify/visualize.py

@@ -112,7 +112,10 @@ class ModelImageVisualizer:
         else:
             self._plot_solo(figsize, render_factor, display_render_factor, result)
 
-        return self._save_result_image(path, result)
+        orig.close()
+        result_path = self._save_result_image(path, result)
+        result.close()
+        return result_path
 
     def _plot_comparison(
         self,