3 files changed, 41 insertions, 7 deletions

diff --git a/source/blender/editors/space_image/image_undo.c b/source/blender/editors/space_image/image_undo.c
index 7cfdea3495d..577c4e24b11 100644
--- a/source/blender/editors/space_image/image_undo.c
+++ b/source/blender/editors/space_image/image_undo.c
@@ -457,6 +457,31 @@ static void ubuf_from_image_all_tiles(UndoImageBuf *ubuf, const ImBuf *ibuf)
   IMB_freeImBuf(tmpibuf);
 }
 
+/** Ensure we can copy the ubuf into the ibuf. */
+static void ubuf_ensure_compat_ibuf(const UndoImageBuf *ubuf, ImBuf *ibuf)
+{
+  /* We could have both float and rect buffers,
+   * in this case free the float buffer if it's unused. */
+  if ((ibuf->rect_float != NULL) && (ubuf->image_state.use_float == false)) {
+    imb_freerectfloatImBuf(ibuf);
+  }
+
+  if (ibuf->x == ubuf->image_dims[0] && ibuf->y == ubuf->image_dims[1] &&
+      (ubuf->image_state.use_float ? (void *)ibuf->rect_float : (void *)ibuf->rect)) {
+    return;
+  }
+
+  imb_freerectImbuf_all(ibuf);
+  IMB_rect_size_set(ibuf, ubuf->image_dims);
+
+  if (ubuf->image_state.use_float) {
+    imb_addrectfloatImBuf(ibuf);
+  }
+  else {
+    imb_addrectImBuf(ibuf);
+  }
+}
+
 static void ubuf_free(UndoImageBuf *ubuf)
 {
   UndoImageBuf *ubuf_post = ubuf->post;
@@ -510,7 +535,8 @@ static void uhandle_restore_list(ListBase *undo_handles, bool use_init)
     bool changed = false;
     for (UndoImageBuf *ubuf_iter = uh->buffers.first; ubuf_iter; ubuf_iter = ubuf_iter->next) {
       UndoImageBuf *ubuf = use_init ? ubuf_iter : ubuf_iter->post;
-      IMB_rect_size_set(ibuf, ubuf->image_dims);
+      ubuf_ensure_compat_ibuf(ubuf, ibuf);
+
       int i = 0;
       for (uint y_tile = 0; y_tile < ubuf->tiles_dims[1]; y_tile += 1) {
         uint y = y_tile << ED_IMAGE_UNDO_TILE_BITS;
diff --git a/source/blender/imbuf/IMB_imbuf.h b/source/blender/imbuf/IMB_imbuf.h
index 9e9c47194e1..173c8135f96 100644
--- a/source/blender/imbuf/IMB_imbuf.h
+++ b/source/blender/imbuf/IMB_imbuf.h
@@ -700,6 +700,8 @@ void imb_freemipmapImBuf(struct ImBuf *ibuf);
 bool imb_addtilesImBuf(struct ImBuf *ibuf);
 void imb_freetilesImBuf(struct ImBuf *ibuf);
 
+void imb_freerectImbuf_all(struct ImBuf *ibuf);
+
 /* threaded processors */
 void IMB_processor_apply_threaded(
     int buffer_lines,
diff --git a/source/blender/imbuf/intern/allocimbuf.c b/source/blender/imbuf/intern/allocimbuf.c
index 75db3fd3c73..381de9de610 100644
--- a/source/blender/imbuf/intern/allocimbuf.c
+++ b/source/blender/imbuf/intern/allocimbuf.c
@@ -197,6 +197,17 @@ void IMB_freezbuffloatImBuf(ImBuf *ibuf)
   ibuf->mall &= ~IB_zbuffloat;
 }
 
+/** Free all pixel data (assosiated with image size). */
+void imb_freerectImbuf_all(ImBuf *ibuf)
+{
+  imb_freerectImBuf(ibuf);
+  imb_freerectfloatImBuf(ibuf);
+  imb_freetilesImBuf(ibuf);
+  IMB_freezbufImBuf(ibuf);
+  IMB_freezbuffloatImBuf(ibuf);
+  freeencodedbufferImBuf(ibuf);
+}
+
 void IMB_freeImBuf(ImBuf *ibuf)
 {
   if (ibuf) {
@@ -212,12 +223,7 @@ void IMB_freeImBuf(ImBuf *ibuf)
     BLI_spin_unlock(&refcounter_spin);
 
     if (needs_free) {
-      imb_freerectImBuf(ibuf);
-      imb_freerectfloatImBuf(ibuf);
-      imb_freetilesImBuf(ibuf);
-      IMB_freezbufImBuf(ibuf);
-      IMB_freezbuffloatImBuf(ibuf);
-      freeencodedbufferImBuf(ibuf);
+      imb_freerectImbuf_all(ibuf);
       IMB_metadata_free(ibuf->metadata);
       colormanage_cache_free(ibuf);