diff --git a/src/Umbraco.Web.UI.Client/src/backoffice/shared/collection/views/collection-view-media-grid.element.ts b/src/Umbraco.Web.UI.Client/src/backoffice/shared/collection/views/collection-view-media-grid.element.ts
index c1132038f8..6d2325ff23 100644
--- a/src/Umbraco.Web.UI.Client/src/backoffice/shared/collection/views/collection-view-media-grid.element.ts
+++ b/src/Umbraco.Web.UI.Client/src/backoffice/shared/collection/views/collection-view-media-grid.element.ts
@@ -3,8 +3,8 @@ import { css, html } from 'lit';
 import { customElement, state } from 'lit/decorators.js';
 import { repeat } from 'lit/directives/repeat.js';
 import { UmbCollectionContext, UMB_COLLECTION_CONTEXT_TOKEN } from '../collection.context';
-import type { MediaTreeItem } from '../../../media/media/media.tree.store';
 import { UmbLitElement } from '@umbraco-cms/element';
+import { EntityTreeItem } from '@umbraco-cms/backend-api';
 
 @customElement('umb-collection-view-media-grid')
 export class UmbCollectionViewsMediaGridElement extends UmbLitElement {
@@ -65,12 +65,12 @@ export class UmbCollectionViewsMediaGridElement extends UmbLitElement {
 	];
 
 	@state()
-	private _mediaItems?: Array<MediaTreeItem>;
+	private _mediaItems?: Array<EntityTreeItem>;
 
 	@state()
 	private _selection: Array<string> = [];
 
-	private _collectionContext?: UmbCollectionContext<MediaTreeItem>;
+	private _collectionContext?: UmbCollectionContext<EntityTreeItem>;
 
 	constructor() {
 		super();
@@ -115,31 +115,31 @@ export class UmbCollectionViewsMediaGridElement extends UmbLitElement {
 		});
 	}
 
-	private _handleOpenItem(mediaItem: MediaTreeItem) {
+	private _handleOpenItem(mediaItem: EntityTreeItem) {
 		//TODO: Fix when we have dynamic routing
 		history.pushState(null, '', 'section/media/media/edit/' + mediaItem.key);
 	}
 
-	private _handleSelect(mediaItem: MediaTreeItem) {
+	private _handleSelect(mediaItem: EntityTreeItem) {
 		if (mediaItem.key) {
 			this._collectionContext?.select(mediaItem.key);
 		}
 	}
 
-	private _handleDeselect(mediaItem: MediaTreeItem) {
+	private _handleDeselect(mediaItem: EntityTreeItem) {
 		if (mediaItem.key) {
 			this._collectionContext?.deselect(mediaItem.key);
 		}
 	}
 
-	private _isSelected(mediaItem: MediaTreeItem) {
+	private _isSelected(mediaItem: EntityTreeItem) {
 		if (mediaItem.key) {
 			return this._selection.includes(mediaItem.key);
 		}
 		return false;
 	}
 
-	private _renderMediaItem(item: MediaTreeItem) {
+	private _renderMediaItem(item: EntityTreeItem) {
 		const name = item.name || '';
 		//TODO: fix the file extension when media items have a file extension.
 		return html`<uui-card-media
diff --git a/src/Umbraco.Web.UI.Client/src/backoffice/shared/components/workspace/workspace-content/views/collection/workspace-view-collection.element.ts b/src/Umbraco.Web.UI.Client/src/backoffice/shared/components/workspace/workspace-content/views/collection/workspace-view-collection.element.ts
index 0be6058a3c..a085014066 100644
--- a/src/Umbraco.Web.UI.Client/src/backoffice/shared/components/workspace/workspace-content/views/collection/workspace-view-collection.element.ts
+++ b/src/Umbraco.Web.UI.Client/src/backoffice/shared/components/workspace/workspace-content/views/collection/workspace-view-collection.element.ts
@@ -2,7 +2,6 @@ import { css, html } from 'lit';
 import { UUITextStyles } from '@umbraco-ui/uui-css/lib';
 import { customElement } from 'lit/decorators.js';
 import { ifDefined } from 'lit-html/directives/if-defined.js';
-import { UmbMediaTreeStore } from '../../../../../../media/media/media.tree.store';
 import {
 	UmbCollectionContext,
 	UMB_COLLECTION_CONTEXT_TOKEN,
@@ -31,7 +30,8 @@ export class UmbWorkspaceViewCollectionElement extends UmbLitElement {
 
 	private _workspaceContext?: UmbWorkspaceEntityContextInterface;
 
-	private _collectionContext?: UmbCollectionContext<FolderTreeItem, UmbMediaTreeStore>;
+	// TODO: add type for the collection context.
+	private _collectionContext?: UmbCollectionContext<FolderTreeItem, any>;
 
 	constructor() {
 		super();