diff --git a/libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerDataSource.kt b/libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerDataSource.kt
index f66521fb16..ebaba4cdcd 100644
--- a/libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerDataSource.kt
+++ b/libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerDataSource.kt
@@ -18,6 +18,7 @@ import io.element.android.libraries.architecture.AsyncData
 import io.element.android.libraries.matrix.api.media.MatrixMediaLoader
 import io.element.android.libraries.matrix.api.media.MediaFile
 import io.element.android.libraries.matrix.api.timeline.Timeline
+import io.element.android.libraries.mediaviewer.api.MediaViewerEntryPoint.MediaViewerMode
 import io.element.android.libraries.mediaviewer.api.local.LocalMedia
 import io.element.android.libraries.mediaviewer.api.local.LocalMediaFactory
 import io.element.android.libraries.mediaviewer.impl.datasource.MediaGalleryDataSource
@@ -38,7 +39,7 @@ import kotlinx.coroutines.withContext
 import timber.log.Timber
 
 class MediaViewerDataSource(
-    private val galleryMode: MediaGalleryMode,
+    mode: MediaViewerMode,
     private val dispatcher: CoroutineDispatcher,
     private val galleryDataSource: MediaGalleryDataSource,
     private val mediaLoader: MatrixMediaLoader,
@@ -49,6 +50,12 @@ class MediaViewerDataSource(
     // List of media files that are currently being loaded
     private val mediaFiles: MutableList<MediaFile> = mutableListOf()
 
+    private val galleryMode = when (mode) {
+        MediaViewerMode.SingleMedia,
+        MediaViewerMode.TimelineImagesAndVideos -> MediaGalleryMode.Images
+        MediaViewerMode.TimelineFilesAndAudios -> MediaGalleryMode.Files
+    }
+
     // Map of sourceUrl to local media state
     private val localMediaStates: MutableMap<String, MutableState<AsyncData<LocalMedia>>> =
         mutableMapOf()
diff --git a/libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerNode.kt b/libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerNode.kt
index b3f95e2b37..7876641051 100644
--- a/libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerNode.kt
+++ b/libraries/mediaviewer/impl/src/main/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerNode.kt
@@ -26,7 +26,6 @@ import io.element.android.libraries.mediaviewer.api.MediaViewerEntryPoint
 import io.element.android.libraries.mediaviewer.api.local.LocalMediaFactory
 import io.element.android.libraries.mediaviewer.impl.datasource.FocusedTimelineMediaGalleryDataSourceFactory
 import io.element.android.libraries.mediaviewer.impl.datasource.TimelineMediaGalleryDataSource
-import io.element.android.libraries.mediaviewer.impl.gallery.MediaGalleryMode
 import io.element.android.libraries.mediaviewer.impl.model.hasEvent
 import io.element.android.services.toolbox.api.systemclock.SystemClock
 
@@ -84,18 +83,12 @@ class MediaViewerNode @AssistedInject constructor(
         }
     }
 
-    private val galleryMode = when (inputs.mode) {
-        MediaViewerEntryPoint.MediaViewerMode.SingleMedia,
-        MediaViewerEntryPoint.MediaViewerMode.TimelineImagesAndVideos -> MediaGalleryMode.Images
-        MediaViewerEntryPoint.MediaViewerMode.TimelineFilesAndAudios -> MediaGalleryMode.Files
-    }
-
     private val presenter = presenterFactory.create(
         inputs = inputs,
         navigator = this,
         dataSource = MediaViewerDataSource(
+            mode = inputs.mode,
             dispatcher = coroutineDispatchers.computation,
-            galleryMode = galleryMode,
             galleryDataSource = mediaGallerySource,
             mediaLoader = mediaLoader,
             localMediaFactory = localMediaFactory,
diff --git a/libraries/mediaviewer/impl/src/test/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerDataSourceTest.kt b/libraries/mediaviewer/impl/src/test/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerDataSourceTest.kt
index 5a66501c92..9e2320cb61 100644
--- a/libraries/mediaviewer/impl/src/test/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerDataSourceTest.kt
+++ b/libraries/mediaviewer/impl/src/test/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerDataSourceTest.kt
@@ -17,10 +17,10 @@ import io.element.android.libraries.matrix.test.AN_EVENT_ID
 import io.element.android.libraries.matrix.test.AN_EVENT_ID_2
 import io.element.android.libraries.matrix.test.AN_EXCEPTION
 import io.element.android.libraries.matrix.test.media.FakeMatrixMediaLoader
+import io.element.android.libraries.mediaviewer.api.MediaViewerEntryPoint.MediaViewerMode
 import io.element.android.libraries.mediaviewer.api.local.LocalMediaFactory
 import io.element.android.libraries.mediaviewer.impl.datasource.FakeMediaGalleryDataSource
 import io.element.android.libraries.mediaviewer.impl.datasource.MediaGalleryDataSource
-import io.element.android.libraries.mediaviewer.impl.gallery.MediaGalleryMode
 import io.element.android.libraries.mediaviewer.impl.gallery.aGroupedMediaItems
 import io.element.android.libraries.mediaviewer.impl.gallery.ui.aMediaItemDateSeparator
 import io.element.android.libraries.mediaviewer.impl.gallery.ui.aMediaItemFile
@@ -137,7 +137,7 @@ class MediaViewerDataSourceTest {
     fun `test dataFlow with data galleryMode image`() = runTest {
         val galleryDataSource = FakeMediaGalleryDataSource()
         val sut = createMediaViewerDataSource(
-            galleryMode = MediaGalleryMode.Images,
+            mode = MediaViewerMode.TimelineImagesAndVideos,
             galleryDataSource = galleryDataSource,
         )
         sut.dataFlow().test {
@@ -159,7 +159,7 @@ class MediaViewerDataSourceTest {
     fun `test dataFlow with data galleryMode files`() = runTest {
         val galleryDataSource = FakeMediaGalleryDataSource()
         val sut = createMediaViewerDataSource(
-            galleryMode = MediaGalleryMode.Files,
+            mode = MediaViewerMode.TimelineFilesAndAudios,
             galleryDataSource = galleryDataSource,
         )
         sut.dataFlow().test {
@@ -265,12 +265,12 @@ class MediaViewerDataSourceTest {
     }
 
     private fun TestScope.createMediaViewerDataSource(
-        galleryMode: MediaGalleryMode = MediaGalleryMode.Images,
+        mode: MediaViewerMode = MediaViewerMode.TimelineImagesAndVideos,
         galleryDataSource: MediaGalleryDataSource = FakeMediaGalleryDataSource(),
         mediaLoader: MatrixMediaLoader = FakeMatrixMediaLoader(),
         localMediaFactory: LocalMediaFactory = FakeLocalMediaFactory(mockMediaUrl),
     ) = MediaViewerDataSource(
-        galleryMode = galleryMode,
+        mode = mode,
         dispatcher = testCoroutineDispatchers().computation,
         galleryDataSource = galleryDataSource,
         mediaLoader = mediaLoader,
diff --git a/libraries/mediaviewer/impl/src/test/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerPresenterTest.kt b/libraries/mediaviewer/impl/src/test/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerPresenterTest.kt
index e1e5818f2b..9f6b636cc6 100644
--- a/libraries/mediaviewer/impl/src/test/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerPresenterTest.kt
+++ b/libraries/mediaviewer/impl/src/test/kotlin/io/element/android/libraries/mediaviewer/impl/viewer/MediaViewerPresenterTest.kt
@@ -32,7 +32,6 @@ import io.element.android.libraries.mediaviewer.impl.R
 import io.element.android.libraries.mediaviewer.impl.datasource.FakeMediaGalleryDataSource
 import io.element.android.libraries.mediaviewer.impl.datasource.MediaGalleryDataSource
 import io.element.android.libraries.mediaviewer.impl.details.MediaBottomSheetState
-import io.element.android.libraries.mediaviewer.impl.gallery.MediaGalleryMode
 import io.element.android.libraries.mediaviewer.impl.gallery.ui.aMediaItemImage
 import io.element.android.libraries.mediaviewer.impl.gallery.ui.aMediaItemLoadingIndicator
 import io.element.android.libraries.mediaviewer.impl.model.GroupedMediaItems
@@ -782,11 +781,7 @@ class MediaViewerPresenterTest {
             ),
             navigator = mediaViewerNavigator,
             dataSource = MediaViewerDataSource(
-                galleryMode = when (mode) {
-                    MediaViewerEntryPoint.MediaViewerMode.SingleMedia -> MediaGalleryMode.Images
-                    MediaViewerEntryPoint.MediaViewerMode.TimelineImagesAndVideos -> MediaGalleryMode.Images
-                    MediaViewerEntryPoint.MediaViewerMode.TimelineFilesAndAudios -> MediaGalleryMode.Files
-                },
+                mode = mode,
                 dispatcher = testCoroutineDispatchers().computation,
                 galleryDataSource = mediaGalleryDataSource,
                 mediaLoader = matrixMediaLoader,