mirror of
https://github.com/jantic/DeOldify.git
synced 2026-08-30 18:02:24 +08:00
Adding graceful and informative gpu out of memory error handling.
This commit is contained in:
+11
-4
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user