diff --git a/src/Avalonia.Diagnostics/ViewModels/ControlDetailsViewModel.cs b/src/Avalonia.Diagnostics/ViewModels/ControlDetailsViewModel.cs
index 34cc32972c..0339d724f7 100644
--- a/src/Avalonia.Diagnostics/ViewModels/ControlDetailsViewModel.cs
+++ b/src/Avalonia.Diagnostics/ViewModels/ControlDetailsViewModel.cs
@@ -3,19 +3,19 @@
 
 using System.Collections.Generic;
 using System.Linq;
-using Avalonia.Controls;
+using Avalonia.VisualTree;
 using ReactiveUI;
 
 namespace Avalonia.Diagnostics.ViewModels
 {
     internal class ControlDetailsViewModel : ReactiveObject
     {
-        public ControlDetailsViewModel(Control control)
+        public ControlDetailsViewModel(IVisual control)
         {
-            if (control != null)
+            if (control is AvaloniaObject avaloniaObject)
             {
-                Properties = AvaloniaPropertyRegistry.Instance.GetRegistered(control)
-                    .Select(x => new PropertyDetails(control, x))
+                Properties = AvaloniaPropertyRegistry.Instance.GetRegistered(avaloniaObject)
+                    .Select(x => new PropertyDetails(avaloniaObject, x))
                     .OrderBy(x => x.IsAttached)
                     .ThenBy(x => x.Name);
             }
diff --git a/src/Avalonia.Diagnostics/ViewModels/TreeNode.cs b/src/Avalonia.Diagnostics/ViewModels/TreeNode.cs
index f186304f7b..0b0b7a94fe 100644
--- a/src/Avalonia.Diagnostics/ViewModels/TreeNode.cs
+++ b/src/Avalonia.Diagnostics/ViewModels/TreeNode.cs
@@ -5,8 +5,8 @@ using System;
 using System.Collections.Specialized;
 using System.Reactive;
 using System.Reactive.Linq;
-using Avalonia.Controls;
 using Avalonia.Styling;
+using Avalonia.VisualTree;
 using ReactiveUI;
 
 namespace Avalonia.Diagnostics.ViewModels
@@ -16,32 +16,35 @@ namespace Avalonia.Diagnostics.ViewModels
         private string _classes;
         private bool _isExpanded;
 
-        public TreeNode(Control control, TreeNode parent)
+        public TreeNode(IVisual visual, TreeNode parent)
         {
-            Control = control;
             Parent = parent;
-            Type = control.GetType().Name;
+            Type = visual.GetType().Name;
+            Visual = visual;
 
-            var classesChanged = Observable.FromEventPattern<
-                    NotifyCollectionChangedEventHandler, 
-                    NotifyCollectionChangedEventArgs>(
-                x => control.Classes.CollectionChanged += x,
-                x => control.Classes.CollectionChanged -= x)
-                .TakeUntil(((IStyleable)control).StyleDetach);
+            if (visual is IStyleable styleable)
+            {
+                var classesChanged = Observable.FromEventPattern<
+                        NotifyCollectionChangedEventHandler,
+                        NotifyCollectionChangedEventArgs>(
+                    x => styleable.Classes.CollectionChanged += x,
+                    x => styleable.Classes.CollectionChanged -= x)
+                    .TakeUntil(((IStyleable)styleable).StyleDetach);
 
-            classesChanged.Select(_ => Unit.Default)
-                .StartWith(Unit.Default)
-                .Subscribe(_ =>
-                {
-                    if (control.Classes.Count > 0)
+                classesChanged.Select(_ => Unit.Default)
+                    .StartWith(Unit.Default)
+                    .Subscribe(_ =>
                     {
-                        Classes = "(" + string.Join(" ", control.Classes) + ")";
-                    }
-                    else
-                    {
-                        Classes = string.Empty;
-                    }
-                });
+                        if (styleable.Classes.Count > 0)
+                        {
+                            Classes = "(" + string.Join(" ", styleable.Classes) + ")";
+                        }
+                        else
+                        {
+                            Classes = string.Empty;
+                        }
+                    });
+            }
         }
 
         public IReadOnlyReactiveList<TreeNode> Children
@@ -56,7 +59,7 @@ namespace Avalonia.Diagnostics.ViewModels
             private set { this.RaiseAndSetIfChanged(ref _classes, value); }
         }
 
-        public Control Control
+        public IVisual Visual
         {
             get;
         }
diff --git a/src/Avalonia.Diagnostics/ViewModels/TreePageViewModel.cs b/src/Avalonia.Diagnostics/ViewModels/TreePageViewModel.cs
index 92be04a056..97bd971e74 100644
--- a/src/Avalonia.Diagnostics/ViewModels/TreePageViewModel.cs
+++ b/src/Avalonia.Diagnostics/ViewModels/TreePageViewModel.cs
@@ -18,7 +18,7 @@ namespace Avalonia.Diagnostics.ViewModels
         {
             Nodes = nodes;
             _details = this.WhenAnyValue(x => x.SelectedNode)
-                .Select(x => x != null ? new ControlDetailsViewModel(x.Control) : null)
+                .Select(x => x != null ? new ControlDetailsViewModel(x.Visual) : null)
                 .ToProperty(this, x => x.Details);
         }
 
@@ -79,7 +79,7 @@ namespace Avalonia.Diagnostics.ViewModels
 
         private TreeNode FindNode(TreeNode node, IControl control)
         {
-            if (node.Control == control)
+            if (node.Visual == control)
             {
                 return node;
             }
diff --git a/src/Avalonia.Diagnostics/ViewModels/VisualTreeNode.cs b/src/Avalonia.Diagnostics/ViewModels/VisualTreeNode.cs
index d155572c14..f1b02a61c9 100644
--- a/src/Avalonia.Diagnostics/ViewModels/VisualTreeNode.cs
+++ b/src/Avalonia.Diagnostics/ViewModels/VisualTreeNode.cs
@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See licence.md file in the project root for full license information.
 
 using Avalonia.Controls;
+using Avalonia.Styling;
 using Avalonia.VisualTree;
 using ReactiveUI;
 
@@ -10,7 +11,7 @@ namespace Avalonia.Diagnostics.ViewModels
     internal class VisualTreeNode : TreeNode
     {
         public VisualTreeNode(IVisual visual, TreeNode parent)
-            : base((Control)visual, parent)
+            : base(visual, parent)
         {
             var host = visual as IVisualTreeHost;
 
@@ -23,9 +24,9 @@ namespace Avalonia.Diagnostics.ViewModels
                 Children = new ReactiveList<VisualTreeNode>(new[] { new VisualTreeNode(host.Root, this) });
             }
 
-            if (Control != null)
+            if ((Visual is IStyleable styleable))
             {
-                IsInTemplate = Control.TemplatedParent != null;
+                IsInTemplate = styleable.TemplatedParent != null;
             }
         }
 
diff --git a/src/Avalonia.Diagnostics/Views/TreePage.xaml.cs b/src/Avalonia.Diagnostics/Views/TreePage.xaml.cs
index 0a79aa90d7..7e4e5a8564 100644
--- a/src/Avalonia.Diagnostics/Views/TreePage.xaml.cs
+++ b/src/Avalonia.Diagnostics/Views/TreePage.xaml.cs
@@ -23,14 +23,14 @@ namespace Avalonia.Diagnostics.Views
         protected void AddAdorner(object sender, PointerEventArgs e)
         {
             var node = (TreeNode)((Control)sender).DataContext;
-            var layer = AdornerLayer.GetAdornerLayer(node.Control);
+            var layer = AdornerLayer.GetAdornerLayer(node.Visual);
 
             if (layer != null)
             {
                 _adorner = new Rectangle
                 {
                     Fill = new SolidColorBrush(0x80a0c5e8),
-                    [AdornerLayer.AdornedElementProperty] = node.Control,
+                    [AdornerLayer.AdornedElementProperty] = node.Visual,
                 };
 
                 layer.Children.Add(_adorner);