diff --git a/src/seaborn_image/_core.py b/src/seaborn_image/_core.py
index 2029f2c3..35329672 100644
--- a/src/seaborn_image/_core.py
+++ b/src/seaborn_image/_core.py
@@ -69,7 +69,7 @@ def _setup_figure(self):
         if self.ax is None:
             f, ax = plt.subplots()
         else:
-            f = plt.gcf()
+            f = self.ax.get_figure() or plt.gcf()
             ax = self.ax
 
         return f, ax