diff --git a/app/src/main/java/com/anthonyla/paperize/feature/wallpaper/data/repository/SelectedAlbumRepositoryImpl.kt b/app/src/main/java/com/anthonyla/paperize/feature/wallpaper/data/repository/SelectedAlbumRepositoryImpl.kt
index 75ed6a17..e7c92ff6 100644
--- a/app/src/main/java/com/anthonyla/paperize/feature/wallpaper/data/repository/SelectedAlbumRepositoryImpl.kt
+++ b/app/src/main/java/com/anthonyla/paperize/feature/wallpaper/data/repository/SelectedAlbumRepositoryImpl.kt
@@ -4,33 +4,47 @@ import com.anthonyla.paperize.feature.wallpaper.data.data_source.SelectedAlbumDa
 import com.anthonyla.paperize.feature.wallpaper.domain.model.SelectedAlbum
 import com.anthonyla.paperize.feature.wallpaper.domain.model.Wallpaper
 import com.anthonyla.paperize.feature.wallpaper.domain.repository.SelectedAlbumRepository
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.withContext
 
 class SelectedAlbumRepositoryImpl(
-    private val dao: SelectedAlbumDao
+    private val dao: SelectedAlbumDao,
+    private val dispatcher: CoroutineDispatcher = Dispatchers.IO
 ): SelectedAlbumRepository {
     override fun getSelectedAlbum(): Flow<List<SelectedAlbum>> {
         return dao.getSelectedAlbum()
     }
 
     override suspend fun upsertSelectedAlbum(selectedAlbum: SelectedAlbum) {
-        dao.upsertSelectedAlbum(selectedAlbum)
+        withContext(dispatcher) {
+            dao.upsertSelectedAlbum(selectedAlbum)
+        }
     }
 
     override suspend fun deleteAlbum(initialAlbumName: String) {
-        dao.deleteAlbum(initialAlbumName)
+        withContext(dispatcher) {
+            dao.deleteAlbum(initialAlbumName)
+        }
     }
 
     override suspend fun deleteWallpaper(wallpaper: Wallpaper) {
-        dao.deleteWallpaper(wallpaper)
+        withContext(dispatcher) {
+            dao.deleteWallpaper(wallpaper)
+        }
     }
 
     override suspend fun cascadeDeleteAlbum(initialAlbumName: String) {
-        dao.deleteAlbum(initialAlbumName)
-        dao.cascadeDeleteWallpaper(initialAlbumName)
+        withContext(dispatcher) {
+            dao.deleteAlbum(initialAlbumName)
+            dao.cascadeDeleteWallpaper(initialAlbumName)
+        }
     }
 
     override suspend fun deleteAll() {
-        dao.deleteAll()
+        withContext(dispatcher) {
+            dao.deleteAll()
+        }
     }
 }
\ No newline at end of file
diff --git a/app/src/main/java/com/anthonyla/paperize/feature/wallpaper/presentation/wallpaper_screen/WallpaperEvent.kt b/app/src/main/java/com/anthonyla/paperize/feature/wallpaper/presentation/wallpaper_screen/WallpaperEvent.kt
index e9316d46..a79e0adc 100644
--- a/app/src/main/java/com/anthonyla/paperize/feature/wallpaper/presentation/wallpaper_screen/WallpaperEvent.kt
+++ b/app/src/main/java/com/anthonyla/paperize/feature/wallpaper/presentation/wallpaper_screen/WallpaperEvent.kt
@@ -7,5 +7,4 @@ sealed class WallpaperEvent {
     data class AddSelectedAlbum(val album: AlbumWithWallpaperAndFolder, val deleteAlbumName: String? = null): WallpaperEvent()
     data class UpdateSelectedAlbum(val album: SelectedAlbum): WallpaperEvent()
     data class Reset(val album: SelectedAlbum? = null) : WallpaperEvent()
-    data object Refresh : WallpaperEvent()
 }
\ No newline at end of file
diff --git a/app/src/main/java/com/anthonyla/paperize/feature/wallpaper/presentation/wallpaper_screen/WallpaperScreenViewModel.kt b/app/src/main/java/com/anthonyla/paperize/feature/wallpaper/presentation/wallpaper_screen/WallpaperScreenViewModel.kt
index 37f5ea86..9c34a53e 100644
--- a/app/src/main/java/com/anthonyla/paperize/feature/wallpaper/presentation/wallpaper_screen/WallpaperScreenViewModel.kt
+++ b/app/src/main/java/com/anthonyla/paperize/feature/wallpaper/presentation/wallpaper_screen/WallpaperScreenViewModel.kt
@@ -1,52 +1,46 @@
 package com.anthonyla.paperize.feature.wallpaper.presentation.wallpaper_screen
 
-import android.app.Application
-import androidx.lifecycle.AndroidViewModel
+import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.anthonyla.paperize.feature.wallpaper.domain.model.SelectedAlbum
 import com.anthonyla.paperize.feature.wallpaper.domain.model.Wallpaper
 import com.anthonyla.paperize.feature.wallpaper.domain.repository.SelectedAlbumRepository
 import dagger.hilt.android.lifecycle.HiltViewModel
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
-import kotlinx.coroutines.flow.stateIn
-import kotlinx.coroutines.flow.update
+import kotlinx.coroutines.flow.*
 import kotlinx.coroutines.launch
 import javax.inject.Inject
 
 @HiltViewModel
-class WallpaperScreenViewModel @Inject constructor (
-    application: Application,
+class WallpaperScreenViewModel @Inject constructor(
     private val repository: SelectedAlbumRepository,
-) : AndroidViewModel(application) {
-    private var _state = MutableStateFlow(WallpaperState())
-    val state = _state.stateIn(
+) : ViewModel() {
+
+    private val _state = MutableStateFlow(WallpaperState())
+    val state = combine(
+        loadSelectedAlbumFlow(),
+        _state
+    ) { selectedAlbum, currentState ->
+        currentState.copy(
+            selectedAlbum = selectedAlbum,
+            isDataLoaded = true
+        )
+    }.stateIn(
         viewModelScope,
-        SharingStarted.WhileSubscribed(5000), WallpaperState()
+        SharingStarted.WhileSubscribed(5000),
+        WallpaperState()
     )
 
-    init {
-        viewModelScope.launch(Dispatchers.IO) {
-            repository.getSelectedAlbum().collect { selectedAlbum ->
-                _state.update {
-                    it.copy(
-                        selectedAlbum = selectedAlbum,
-                        isDataLoaded = true
-                    )
-                }
-            }
-        }
-    }
+    private fun loadSelectedAlbumFlow(): Flow<List<SelectedAlbum>?> =
+        repository.getSelectedAlbum()
 
     fun onEvent(event: WallpaperEvent) {
         when (event) {
             is WallpaperEvent.AddSelectedAlbum -> {
-                viewModelScope.launch(Dispatchers.IO) {
-                    val deleted = event.deleteAlbumName?.let {
+                viewModelScope.launch {
+                    event.deleteAlbumName?.let {
                         repository.cascadeDeleteAlbum(it)
-                        true
-                    } ?: false
+                    }
+                    
                     val wallpapers = event.album.wallpapers + event.album.folders.flatMap { folder ->
                         folder.wallpapers.map { wallpaper ->
                             Wallpaper(
@@ -56,6 +50,7 @@ class WallpaperScreenViewModel @Inject constructor (
                             )
                         }
                     }
+
                     val newSelectedAlbum = SelectedAlbum(
                         album = event.album.album.copy(
                             lockWallpapersInQueue = wallpapers.map { it.wallpaperUri }.shuffled(),
@@ -64,65 +59,24 @@ class WallpaperScreenViewModel @Inject constructor (
                         wallpapers = wallpapers
                     )
                     repository.upsertSelectedAlbum(newSelectedAlbum)
-                    _state.update {
-                        it.copy(
-                            selectedAlbum = it.selectedAlbum?.let { selectedAlbums ->
-                                selectedAlbums.filterNot { selectedAlbum -> deleted && selectedAlbum.album.initialAlbumName == event.deleteAlbumName } + newSelectedAlbum
-                            }
-                        )
-                    }
                 }
             }
+
             is WallpaperEvent.UpdateSelectedAlbum -> {
-                viewModelScope.launch(Dispatchers.IO) {
+                viewModelScope.launch {
                     repository.upsertSelectedAlbum(event.album)
-                    _state.update {
-                        it.copy(
-                            selectedAlbum = it.selectedAlbum?.map { selectedAlbum ->
-                                if (selectedAlbum.album.initialAlbumName == event.album.album.initialAlbumName) {
-                                    event.album
-                                } else {
-                                    selectedAlbum
-                                }
-                            }
-                        )
-                    }
                 }
             }
-            is WallpaperEvent.Refresh -> {
-                refreshAlbums()
-            }
+
             is WallpaperEvent.Reset -> {
-                viewModelScope.launch(Dispatchers.IO) {
+                viewModelScope.launch {
                     if (event.album == null) {
                         repository.deleteAll()
-                        _state.update {
-                            it.copy(
-                                selectedAlbum = null
-                            )
-                        }
-                    }
-                    else {
+                    } else {
                         repository.cascadeDeleteAlbum(event.album.album.initialAlbumName)
-                        _state.update {
-                            it.copy(
-                                selectedAlbum = it.selectedAlbum?.filter { selectedAlbum -> selectedAlbum.album.initialAlbumName != event.album.album.initialAlbumName }
-                            )
-                        }
                     }
                 }
             }
         }
     }
-
-    // Retrieve album from database into viewModel
-    private fun refreshAlbums() {
-        viewModelScope.launch(Dispatchers.IO) {
-            repository.getSelectedAlbum().collect { selectedAlbum ->
-                _state.update { it.copy(
-                    selectedAlbum = selectedAlbum,
-                ) }
-            }
-        }
-    }
 }