diff --git a/Source/Editor/Viewport/PrefabWindowViewport.cs b/Source/Editor/Viewport/PrefabWindowViewport.cs
index acac4767f..256b41290 100644
--- a/Source/Editor/Viewport/PrefabWindowViewport.cs
+++ b/Source/Editor/Viewport/PrefabWindowViewport.cs
@@ -5,10 +5,12 @@ using System.Collections.Generic;
 using System.Linq;
 using FlaxEditor.Content;
 using FlaxEditor.Gizmo;
+using FlaxEditor.GUI.ContextMenu;
 using FlaxEditor.SceneGraph;
 using FlaxEditor.Scripting;
 using FlaxEditor.Viewport.Cameras;
 using FlaxEditor.Viewport.Previews;
+using FlaxEditor.Viewport.Widgets;
 using FlaxEditor.Windows.Assets;
 using FlaxEngine;
 using FlaxEngine.GUI;
@@ -44,7 +46,7 @@ namespace FlaxEditor.Viewport
         private sealed class PrefabUIEditorRoot : UIEditorRoot
         {
             private readonly PrefabWindowViewport _viewport;
-            private bool UI => _viewport._hasUILinkedCached;
+            private bool UI => _viewport.ShowUI;
 
             public PrefabUIEditorRoot(PrefabWindowViewport viewport)
             : base(true)
@@ -67,8 +69,78 @@ namespace FlaxEditor.Viewport
         private PrefabSpritesRenderer _spritesRenderer;
         private IntPtr _tempDebugDrawContext;
 
-        private bool _hasUILinkedCached;
         private PrefabUIEditorRoot _uiRoot;
+        private bool _showUI = false;
+        
+        private ContextMenuButton _uiModeButton;
+
+        /// <summary>
+        /// Event fired when the UI Mode is toggled.
+        /// </summary>
+        public event Action<bool> UIModeToggled;
+
+        /// <summary>
+        /// set the initial UI mod
+        /// </summary>
+        /// <param name="value">the initial ShowUI value</param>
+        public void SetInitialUIMode(bool value)
+        {
+            ShowUI = value;
+            _uiModeButton.Checked = value;
+        }
+
+        /// <summary>
+        /// Whether to show the UI mode or not.
+        /// </summary>
+        public bool ShowUI
+        {
+            get => _showUI;
+            set
+            {
+                _showUI = value;
+                if (_showUI)
+                {
+                    // UI widget
+                    Gizmos.Active = null;
+                    ViewportCamera = new UIEditorCamera { UIEditor = _uiRoot };
+
+                    // Hide 3D visuals
+                    ShowEditorPrimitives = false;
+                    ShowDefaultSceneActors = false;
+                    ShowDebugDraw = false;
+
+                    // Show whole UI on startup
+                    var canvas = (CanvasRootControl)_uiParentLink.Children.FirstOrDefault(x => x is CanvasRootControl);
+                    if (canvas != null)
+                        ViewportCamera.ShowActor(canvas.Canvas);
+                    else if (Instance is UIControl)
+                        ViewportCamera.ShowActor(Instance);
+                    _uiRoot.Visible = true;
+                }
+                else
+                {
+                    // Generic prefab
+                    Gizmos.Active = TransformGizmo;
+                    ViewportCamera = new FPSCamera();
+                    _uiRoot.Visible = false;
+                }
+
+                // Update default components usage
+                bool defaultFeatures = !_showUI;
+                _disableInputUpdate = _showUI;
+                _spritesRenderer.Enabled = defaultFeatures;
+                SelectionOutline.Enabled = defaultFeatures;
+                _showDefaultSceneButton.Visible = defaultFeatures;
+                _cameraWidget.Visible = defaultFeatures;
+                _cameraButton.Visible = defaultFeatures;
+                _orthographicModeButton.Visible = defaultFeatures;
+                Task.Enabled = defaultFeatures;
+                UseAutomaticTaskManagement = defaultFeatures;
+                ShowDefaultSceneActors = defaultFeatures;
+                TintColor = defaultFeatures ? Color.White : Color.Transparent;
+                UIModeToggled?.Invoke(_showUI);
+            }
+        }
 
         /// <summary>
         /// Drag and drop handlers
@@ -138,6 +210,11 @@ namespace FlaxEditor.Viewport
             _uiRoot.IndexInParent = 0; // Move viewport down below other widgets in the viewport
             _uiParentLink = _uiRoot.UIRoot;
 
+            // UI mode buton
+            _uiModeButton = ViewWidgetShowMenu.AddButton("UI Mode", (button) => ShowUI = button.Checked);
+            _uiModeButton.AutoCheck = true;
+            _uiModeButton.VisibleChanged += control => (control as ContextMenuButton).Checked = ShowUI;
+
             EditorGizmoViewport.AddGizmoViewportWidgets(this, TransformGizmo);
 
             // Setup input actions
@@ -146,58 +223,8 @@ namespace FlaxEditor.Viewport
             SetUpdate(ref _update, OnUpdate);
         }
 
-        /// <summary>
-        /// Updates the viewport's gizmos, especially to toggle between 3D and UI editing modes.
-        /// </summary>
-        internal void UpdateGizmoMode()
-        {
-            // Skip if gizmo mode was unmodified
-            if (_hasUILinked == _hasUILinkedCached)
-                return;
-            _hasUILinkedCached = _hasUILinked;
-
-            if (_hasUILinked)
-            {
-                // UI widget
-                Gizmos.Active = null;
-                ViewportCamera = new UIEditorCamera { UIEditor = _uiRoot };
-
-                // Hide 3D visuals
-                ShowEditorPrimitives = false;
-                ShowDefaultSceneActors = false;
-                ShowDebugDraw = false;
-
-                // Show whole UI on startup
-                var canvas = (CanvasRootControl)_uiParentLink.Children.FirstOrDefault(x => x is CanvasRootControl);
-                if (canvas != null)
-                    ViewportCamera.ShowActor(canvas.Canvas);
-                else if (Instance is UIControl)
-                    ViewportCamera.ShowActor(Instance);
-            }
-            else
-            {
-                // Generic prefab
-                Gizmos.Active = TransformGizmo;
-                ViewportCamera = new FPSCamera();
-            }
-
-            // Update default components usage
-            bool defaultFeatures = !_hasUILinked;
-            _disableInputUpdate = _hasUILinked;
-            _spritesRenderer.Enabled = defaultFeatures;
-            SelectionOutline.Enabled = defaultFeatures;
-            _showDefaultSceneButton.Visible = defaultFeatures;
-            _cameraWidget.Visible = defaultFeatures;
-            _cameraButton.Visible = defaultFeatures;
-            _orthographicModeButton.Visible = defaultFeatures;
-            Task.Enabled = defaultFeatures;
-            UseAutomaticTaskManagement = defaultFeatures;
-            TintColor = defaultFeatures ? Color.White : Color.Transparent;
-        }
-
         private void OnUpdate(float deltaTime)
         {
-            UpdateGizmoMode();
             for (int i = 0; i < Gizmos.Count; i++)
             {
                 Gizmos[i].Update(deltaTime);
diff --git a/Source/Editor/Windows/Assets/PrefabWindow.cs b/Source/Editor/Windows/Assets/PrefabWindow.cs
index 7ee8ecff0..5bf3236ee 100644
--- a/Source/Editor/Windows/Assets/PrefabWindow.cs
+++ b/Source/Editor/Windows/Assets/PrefabWindow.cs
@@ -355,13 +355,22 @@ namespace FlaxEditor.Windows.Assets
         private void OnPrefabOpened()
         {
             _viewport.Prefab = _asset;
-            _viewport.UpdateGizmoMode();
+            if (Editor.ProjectCache.TryGetCustomData($"UIMode:{_asset.ID}", out bool value))
+                _viewport.SetInitialUIMode(value);
+            else
+                _viewport.SetInitialUIMode(_viewport._hasUILinked);
+            _viewport.UIModeToggled += OnUIModeToggled;
             Graph.MainActor = _viewport.Instance;
             Selection.Clear();
             Select(Graph.Main);
             Graph.Root.TreeNode.Expand(true);
         }
 
+        private void OnUIModeToggled(bool value)
+        {
+            Editor.ProjectCache.SetCustomData($"UIMode:{_asset.ID}", value);
+        }
+
         /// <inheritdoc />
         public override void Save()
         {