diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
index f393880..43d510c 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -24,7 +24,6 @@
 import com.android.customization.picker.color.shared.model.ColorOptionModel
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.systemui.monet.Style
-import com.android.wallpaper.config.BaseFlags
 import com.android.wallpaper.picker.customization.data.repository.WallpaperColorsRepository
 import com.android.wallpaper.picker.customization.shared.model.WallpaperColorsModel
 import javax.inject.Inject
@@ -47,8 +46,6 @@
     private val colorManager: ColorCustomizationManager,
 ) : ColorPickerRepository {
 
-    private val isNewPickerUi = BaseFlags.get().isNewPickerUi()
-
     private val homeWallpaperColors: StateFlow<WallpaperColorsModel?> =
         wallpaperColorsRepository.homeWallpaperColors
     private val lockWallpaperColors: StateFlow<WallpaperColorsModel?> =
@@ -59,7 +56,7 @@
     private val _isApplyingSystemColor = MutableStateFlow(false)
     override val isApplyingSystemColor = _isApplyingSystemColor.asStateFlow()
 
-    private val generatedColorOptions: Flow<Map<ColorType, List<ColorOptionImpl>>> =
+    override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
         combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
                 homeColors to lockColors
             }
@@ -88,15 +85,16 @@
                     colorManager.fetchOptions(
                         object : CustomizationManager.OptionsFetchedListener<ColorOption?> {
                             override fun onOptionsLoaded(options: MutableList<ColorOption?>?) {
-                                val wallpaperColorOptions: MutableList<ColorOptionImpl> =
+                                val wallpaperColorOptions: MutableList<ColorOptionModel> =
                                     mutableListOf()
-                                val presetColorOptions: MutableList<ColorOptionImpl> =
+                                val presetColorOptions: MutableList<ColorOptionModel> =
                                     mutableListOf()
                                 options?.forEach { option ->
                                     when ((option as ColorOptionImpl).type) {
                                         ColorType.WALLPAPER_COLOR ->
-                                            wallpaperColorOptions.add(option)
-                                        ColorType.PRESET_COLOR -> presetColorOptions.add(option)
+                                            wallpaperColorOptions.add(option.toModel())
+                                        ColorType.PRESET_COLOR ->
+                                            presetColorOptions.add(option.toModel())
                                     }
                                 }
                                 continuation.resumeWith(
@@ -123,83 +121,6 @@
                 }
             }
 
-    override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
-        if (isNewPickerUi) {
-            // Convert to ColorOptionModel. When the selected color option changes, update each
-            // ColorOptionModel's isSelected by calling toModel again.
-            combine(generatedColorOptions, selectedColorOption) { generatedColorOptions, _ ->
-                generatedColorOptions
-                    .map { entry ->
-                        entry.key to entry.value.map { colorOption -> colorOption.toModel() }
-                    }
-                    .toMap()
-            }
-        } else {
-            combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
-                    homeColors to lockColors
-                }
-                .map { (homeColors, lockColors) ->
-                    suspendCancellableCoroutine { continuation ->
-                        if (
-                            homeColors is WallpaperColorsModel.Loading ||
-                                lockColors is WallpaperColorsModel.Loading
-                        ) {
-                            continuation.resumeWith(
-                                Result.success(
-                                    mapOf(
-                                        ColorType.WALLPAPER_COLOR to listOf(),
-                                        ColorType.PRESET_COLOR to listOf(),
-                                    )
-                                )
-                            )
-                            return@suspendCancellableCoroutine
-                        }
-                        val homeColorsLoaded = homeColors as WallpaperColorsModel.Loaded
-                        val lockColorsLoaded = lockColors as WallpaperColorsModel.Loaded
-                        colorManager.setWallpaperColors(
-                            homeColorsLoaded.colors,
-                            lockColorsLoaded.colors,
-                        )
-                        colorManager.fetchOptions(
-                            object : CustomizationManager.OptionsFetchedListener<ColorOption?> {
-                                override fun onOptionsLoaded(options: MutableList<ColorOption?>?) {
-                                    val wallpaperColorOptions: MutableList<ColorOptionModel> =
-                                        mutableListOf()
-                                    val presetColorOptions: MutableList<ColorOptionModel> =
-                                        mutableListOf()
-                                    options?.forEach { option ->
-                                        when ((option as ColorOptionImpl).type) {
-                                            ColorType.WALLPAPER_COLOR ->
-                                                wallpaperColorOptions.add(option.toModel())
-                                            ColorType.PRESET_COLOR ->
-                                                presetColorOptions.add(option.toModel())
-                                        }
-                                    }
-                                    continuation.resumeWith(
-                                        Result.success(
-                                            mapOf(
-                                                ColorType.WALLPAPER_COLOR to wallpaperColorOptions,
-                                                ColorType.PRESET_COLOR to presetColorOptions,
-                                            )
-                                        )
-                                    )
-                                }
-
-                                override fun onError(throwable: Throwable?) {
-                                    Log.e(TAG, "Error loading theme bundles", throwable)
-                                    continuation.resumeWith(
-                                        Result.failure(
-                                            throwable ?: Throwable("Error loading theme bundles")
-                                        )
-                                    )
-                                }
-                            },
-                            /* reload= */ false,
-                        )
-                    }
-                }
-        }
-
     override suspend fun select(colorOptionModel: ColorOptionModel) {
         _isApplyingSystemColor.value = true
         suspendCancellableCoroutine { continuation ->
diff --git a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt b/tests/common/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
similarity index 100%
rename from src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
rename to tests/common/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt