diff --git a/src/slic3r/GUI/Gizmos/GLGizmoFlatten.cpp b/src/slic3r/GUI/Gizmos/GLGizmoFlatten.cpp
index 78f5da58b..cb996a104 100644
--- a/src/slic3r/GUI/Gizmos/GLGizmoFlatten.cpp
+++ b/src/slic3r/GUI/Gizmos/GLGizmoFlatten.cpp
@@ -1,6 +1,7 @@
 // Include GLGizmoBase.hpp before I18N.hpp as it includes some libigl code, which overrides our localization "L" macro.
 #include "GLGizmoFlatten.hpp"
 #include "slic3r/GUI/GLCanvas3D.hpp"
+#include "slic3r/GUI/GUI_App.hpp"
 
 #include <numeric>
 
@@ -23,6 +24,22 @@ bool GLGizmoFlatten::on_init()
     return true;
 }
 
+void GLGizmoFlatten::on_set_state()
+{
+    // m_model_object pointer can be invalid (for instance because of undo/redo action),
+    // we should recover it from the object id
+    m_model_object = nullptr;
+    for (const auto mo : *wxGetApp().model_objects()) {
+        if (mo->id() == m_model_object_id) {
+            m_model_object = mo;
+            break;
+        }
+    }
+
+    if (m_state == On && is_plane_update_necessary())
+        update_planes();
+}
+
 std::string GLGizmoFlatten::on_get_name() const
 {
     return (_(L("Place on face")) + " [F]").ToUTF8().data();
@@ -120,6 +137,7 @@ void GLGizmoFlatten::set_flattening_data(const ModelObject* model_object)
         m_planes_valid = false;
     }
     m_model_object = model_object;
+    m_model_object_id = model_object ? model_object->id() : 0;
 }
 
 void GLGizmoFlatten::update_planes()
diff --git a/src/slic3r/GUI/Gizmos/GLGizmoFlatten.hpp b/src/slic3r/GUI/Gizmos/GLGizmoFlatten.hpp
index c8bd056bc..c69d64134 100644
--- a/src/slic3r/GUI/Gizmos/GLGizmoFlatten.hpp
+++ b/src/slic3r/GUI/Gizmos/GLGizmoFlatten.hpp
@@ -31,6 +31,7 @@ private:
     bool m_planes_valid = false;
     mutable Vec3d m_starting_center;
     const ModelObject* m_model_object = nullptr;
+    ObjectID m_model_object_id = 0;
     std::vector<const Transform3d*> instances_matrices;
 
     void update_planes();
@@ -49,11 +50,7 @@ protected:
     virtual void on_start_dragging();
     virtual void on_render() const;
     virtual void on_render_for_picking() const;
-    virtual void on_set_state()
-    {
-        if (m_state == On && is_plane_update_necessary())
-            update_planes();
-    }
+    virtual void on_set_state() override;
 };
 
 } // namespace GUI