diff --git a/src/com/android/customization/model/color/ColorProvider.kt b/src/com/android/customization/model/color/ColorProvider.kt
index 159d08f..3d2cc7e 100644
--- a/src/com/android/customization/model/color/ColorProvider.kt
+++ b/src/com/android/customization/model/color/ColorProvider.kt
@@ -288,7 +288,7 @@
 
                     if (
                         style == Style.MONOCHROMATIC &&
-                            !InjectorProvider.getInjector().flags.isMonochromaticFlagEnabled()
+                            !InjectorProvider.getInjector().getFlags().isMonochromaticFlagEnabled()
                     ) {
                         continue
                     }
diff --git a/src/com/android/customization/module/CustomizationInjector.java b/src/com/android/customization/module/CustomizationInjector.java
deleted file mode 100644
index 2cc1245..0000000
--- a/src/com/android/customization/module/CustomizationInjector.java
+++ /dev/null
@@ -1,41 +0,0 @@
-/*
- * Copyright (C) 2019 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.customization.module;
-
-import android.content.Context;
-
-import androidx.fragment.app.FragmentActivity;
-
-import com.android.customization.model.theme.OverlayManagerCompat;
-import com.android.customization.model.theme.ThemeBundleProvider;
-import com.android.customization.model.theme.ThemeManager;
-import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor;
-import com.android.wallpaper.module.Injector;
-
-public interface CustomizationInjector extends Injector {
-
-    CustomizationPreferences getCustomizationPreferences(Context context);
-
-    ThemeManager getThemeManager(ThemeBundleProvider provider, FragmentActivity activity,
-            OverlayManagerCompat overlayManagerCompat, ThemesUserEventLogger logger);
-
-
-    /**
-     * Get {@link KeyguardQuickAffordancePickerInteractor}
-     */
-    KeyguardQuickAffordancePickerInteractor getKeyguardQuickAffordancePickerInteractor(
-            Context context);
-}
diff --git a/src/com/android/customization/module/CustomizationInjector.kt b/src/com/android/customization/module/CustomizationInjector.kt
new file mode 100644
index 0000000..3cf8393
--- /dev/null
+++ b/src/com/android/customization/module/CustomizationInjector.kt
@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.customization.module
+
+import android.content.Context
+import androidx.fragment.app.FragmentActivity
+import com.android.customization.model.theme.OverlayManagerCompat
+import com.android.customization.model.theme.ThemeBundleProvider
+import com.android.customization.model.theme.ThemeManager
+import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
+import com.android.wallpaper.module.Injector
+
+interface CustomizationInjector : Injector {
+    fun getCustomizationPreferences(context: Context): CustomizationPreferences
+
+    fun getThemeManager(
+        provider: ThemeBundleProvider,
+        activity: FragmentActivity,
+        overlayManagerCompat: OverlayManagerCompat,
+        logger: ThemesUserEventLogger
+    ): ThemeManager
+
+    fun getKeyguardQuickAffordancePickerInteractor(
+        context: Context
+    ): KeyguardQuickAffordancePickerInteractor
+}
diff --git a/src/com/android/customization/module/ThemePickerInjector.java b/src/com/android/customization/module/ThemePickerInjector.java
deleted file mode 100644
index fb778f1..0000000
--- a/src/com/android/customization/module/ThemePickerInjector.java
+++ /dev/null
@@ -1,236 +0,0 @@
-/*
- * Copyright (C) 2019 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.customization.module;
-
-import static com.android.wallpaper.picker.PreviewFragment.ARG_FULL_SCREEN;
-import static com.android.wallpaper.picker.PreviewFragment.ARG_PREVIEW_MODE;
-import static com.android.wallpaper.picker.PreviewFragment.ARG_TESTING_MODE_ENABLED;
-import static com.android.wallpaper.picker.PreviewFragment.ARG_VIEW_AS_HOME;
-import static com.android.wallpaper.picker.PreviewFragment.ARG_WALLPAPER;
-
-import android.app.Activity;
-import android.content.Context;
-import android.content.Intent;
-import android.net.Uri;
-import android.os.Bundle;
-
-import androidx.fragment.app.Fragment;
-import androidx.fragment.app.FragmentActivity;
-
-import com.android.customization.model.theme.OverlayManagerCompat;
-import com.android.customization.model.theme.ThemeBundleProvider;
-import com.android.customization.model.theme.ThemeManager;
-import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository;
-import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor;
-import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer;
-import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordancePickerViewModel;
-import com.android.systemui.shared.customization.data.content.CustomizationProviderClient;
-import com.android.systemui.shared.customization.data.content.CustomizationProviderClientImpl;
-import com.android.wallpaper.config.BaseFlags;
-import com.android.wallpaper.model.LiveWallpaperInfo;
-import com.android.wallpaper.model.WallpaperInfo;
-import com.android.wallpaper.module.CustomizationSections;
-import com.android.wallpaper.module.FragmentFactory;
-import com.android.wallpaper.module.WallpaperPicker2Injector;
-import com.android.wallpaper.module.WallpaperPreferences;
-import com.android.wallpaper.picker.CustomizationPickerActivity;
-import com.android.wallpaper.picker.ImagePreviewFragment;
-import com.android.wallpaper.picker.LivePreviewFragment;
-import com.android.wallpaper.picker.PreviewFragment;
-import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer;
-
-import java.util.Map;
-
-import kotlinx.coroutines.Dispatchers;
-
-/**
- * A concrete, real implementation of the dependency provider.
- */
-public class ThemePickerInjector extends WallpaperPicker2Injector
-        implements CustomizationInjector {
-    private CustomizationSections mCustomizationSections;
-    private ThemesUserEventLogger mUserEventLogger;
-    private WallpaperPreferences mPrefs;
-    private KeyguardQuickAffordancePickerInteractor mKeyguardQuickAffordancePickerInteractor;
-    private KeyguardQuickAffordancePickerViewModel.Factory
-            mKeyguardQuickAffordancePickerViewModelFactory;
-    private CustomizationProviderClient mCustomizationProviderClient;
-    private FragmentFactory mFragmentFactory;
-    private BaseFlags mFlags;
-    private KeyguardQuickAffordanceSnapshotRestorer mKeyguardQuickAffordanceSnapshotRestorer;
-
-    @Override
-    public CustomizationSections getCustomizationSections(Activity activity) {
-        if (mCustomizationSections == null) {
-            mCustomizationSections = new DefaultCustomizationSections(
-                    mKeyguardQuickAffordancePickerInteractor,
-                    mKeyguardQuickAffordancePickerViewModelFactory);
-        }
-        return mCustomizationSections;
-    }
-
-    @Override
-    public Intent getDeepLinkRedirectIntent(Context context, Uri uri) {
-        Intent intent = new Intent();
-        intent.setClass(context, CustomizationPickerActivity.class);
-        intent.setData(uri);
-        intent.setFlags(Intent.FLAG_ACTIVITY_NEW_TASK | Intent.FLAG_ACTIVITY_CLEAR_TASK);
-        return intent;
-    }
-
-    @Override
-    public String getDownloadableIntentAction() {
-        return null;
-    }
-
-    @Override
-    public Fragment getPreviewFragment(
-            Context context,
-            WallpaperInfo wallpaperInfo,
-            int mode,
-            boolean viewAsHome,
-            boolean viewFullScreen,
-            boolean testingModeEnabled) {
-        Bundle args = new Bundle();
-        args.putParcelable(ARG_WALLPAPER, wallpaperInfo);
-        args.putInt(ARG_PREVIEW_MODE, mode);
-        args.putBoolean(ARG_VIEW_AS_HOME, viewAsHome);
-        args.putBoolean(ARG_FULL_SCREEN, viewFullScreen);
-        args.putBoolean(ARG_TESTING_MODE_ENABLED, testingModeEnabled);
-        PreviewFragment fragment = wallpaperInfo instanceof LiveWallpaperInfo
-                ? new LivePreviewFragment() : new ImagePreviewFragment();
-        fragment.setArguments(args);
-        return fragment;
-    }
-
-    @Override
-    public synchronized ThemesUserEventLogger getUserEventLogger(Context context) {
-        if (mUserEventLogger == null) {
-            mUserEventLogger = new StatsLogUserEventLogger(context);
-        }
-        return mUserEventLogger;
-    }
-
-    @Override
-    public synchronized WallpaperPreferences getPreferences(Context context) {
-        if (mPrefs == null) {
-            mPrefs = new DefaultCustomizationPreferences(context.getApplicationContext());
-        }
-        return mPrefs;
-    }
-
-    //
-    // Functions from {@link CustomizationInjector}
-    //
-    @Override
-    public CustomizationPreferences getCustomizationPreferences(Context context) {
-        return (CustomizationPreferences) getPreferences(context);
-    }
-
-    @Override
-    public ThemeManager getThemeManager(ThemeBundleProvider provider, FragmentActivity activity,
-            OverlayManagerCompat overlayManagerCompat, ThemesUserEventLogger logger) {
-        return new ThemeManager(provider, activity, overlayManagerCompat, logger);
-    }
-
-    @Override
-    public KeyguardQuickAffordancePickerInteractor getKeyguardQuickAffordancePickerInteractor(
-            Context context) {
-        if (mKeyguardQuickAffordancePickerInteractor == null) {
-            final CustomizationProviderClient client =
-                    getKeyguardQuickAffordancePickerProviderClient(context);
-            mKeyguardQuickAffordancePickerInteractor = new KeyguardQuickAffordancePickerInteractor(
-                    new KeyguardQuickAffordancePickerRepository(client, Dispatchers.getIO()),
-                    client,
-                    () -> getKeyguardQuickAffordanceSnapshotRestorer(context));
-        }
-        return mKeyguardQuickAffordancePickerInteractor;
-    }
-
-    /**
-     * Returns a {@link KeyguardQuickAffordancePickerViewModel.Factory}.
-     */
-    public KeyguardQuickAffordancePickerViewModel.Factory
-            getKeyguardQuickAffordancePickerViewModelFactory(Context context) {
-        if (mKeyguardQuickAffordancePickerViewModelFactory == null) {
-            mKeyguardQuickAffordancePickerViewModelFactory =
-                    new KeyguardQuickAffordancePickerViewModel.Factory(
-                            context,
-                            getKeyguardQuickAffordancePickerInteractor(context),
-                            getUndoInteractor(context),
-                            getCurrentWallpaperInfoFactory(context));
-        }
-        return mKeyguardQuickAffordancePickerViewModelFactory;
-    }
-
-    @Override
-    public FragmentFactory getFragmentFactory() {
-        if (mFragmentFactory == null) {
-            mFragmentFactory = new ThemePickerFragmentFactory();
-        }
-        return mFragmentFactory;
-    }
-
-    @Override
-    public BaseFlags getFlags() {
-        if (mFlags == null) {
-            mFlags = new BaseFlags() {};
-        }
-
-        return mFlags;
-    }
-
-    @Override
-    public Map<Integer, SnapshotRestorer> getSnapshotRestorers(Context context) {
-        final Map<Integer, SnapshotRestorer> restorers = super.getSnapshotRestorers(context);
-        restorers.put(
-                KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER,
-                getKeyguardQuickAffordanceSnapshotRestorer(context));
-        return restorers;
-    }
-
-    /** Returns the {@link CustomizationProviderClient}. */
-    protected CustomizationProviderClient getKeyguardQuickAffordancePickerProviderClient(
-            Context context) {
-        if (mCustomizationProviderClient == null) {
-            mCustomizationProviderClient =
-                    new CustomizationProviderClientImpl(context, Dispatchers.getIO());
-        }
-
-        return mCustomizationProviderClient;
-    }
-
-    protected KeyguardQuickAffordanceSnapshotRestorer getKeyguardQuickAffordanceSnapshotRestorer(
-            Context context) {
-        if (mKeyguardQuickAffordanceSnapshotRestorer == null) {
-            mKeyguardQuickAffordanceSnapshotRestorer = new KeyguardQuickAffordanceSnapshotRestorer(
-                    getKeyguardQuickAffordancePickerInteractor(context),
-                    getKeyguardQuickAffordancePickerProviderClient(context));
-        }
-
-        return mKeyguardQuickAffordanceSnapshotRestorer;
-    }
-
-    private static final int KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER =
-            WallpaperPicker2Injector.MIN_SNAPSHOT_RESTORER_KEY;
-
-    /**
-     * When this injector is overridden, this is the minimal value that should be used by restorers
-     * returns in {@link #getSnapshotRestorers(Context)}.
-     */
-    protected static final int MIN_SNAPSHOT_RESTORER_KEY =
-            KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER + 1;
-}
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
new file mode 100644
index 0000000..b3f95d5
--- /dev/null
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -0,0 +1,205 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.customization.module
+
+import android.app.Activity
+import android.content.Context
+import android.content.Intent
+import android.net.Uri
+import android.os.Bundle
+import androidx.fragment.app.Fragment
+import androidx.fragment.app.FragmentActivity
+import com.android.customization.model.theme.OverlayManagerCompat
+import com.android.customization.model.theme.ThemeBundleProvider
+import com.android.customization.model.theme.ThemeManager
+import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
+import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
+import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer
+import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordancePickerViewModel
+import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
+import com.android.systemui.shared.customization.data.content.CustomizationProviderClientImpl
+import com.android.wallpaper.model.LiveWallpaperInfo
+import com.android.wallpaper.model.WallpaperInfo
+import com.android.wallpaper.module.CustomizationSections
+import com.android.wallpaper.module.FragmentFactory
+import com.android.wallpaper.module.UserEventLogger
+import com.android.wallpaper.module.WallpaperPicker2Injector
+import com.android.wallpaper.module.WallpaperPreferences
+import com.android.wallpaper.picker.CustomizationPickerActivity
+import com.android.wallpaper.picker.ImagePreviewFragment
+import com.android.wallpaper.picker.LivePreviewFragment
+import com.android.wallpaper.picker.PreviewFragment
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import kotlinx.coroutines.Dispatchers.IO
+
+open class ThemePickerInjector : WallpaperPicker2Injector(), CustomizationInjector {
+    private var customizationSections: CustomizationSections? = null
+    private var userEventLogger: UserEventLogger? = null
+    private var prefs: WallpaperPreferences? = null
+    private var keyguardQuickAffordancePickerInteractor: KeyguardQuickAffordancePickerInteractor? =
+        null
+    private var keyguardQuickAffordancePickerViewModelFactory:
+        KeyguardQuickAffordancePickerViewModel.Factory? =
+        null
+    private var customizationProviderClient: CustomizationProviderClient? = null
+    private var fragmentFactory: FragmentFactory? = null
+    private var keyguardQuickAffordanceSnapshotRestorer: KeyguardQuickAffordanceSnapshotRestorer? =
+        null
+
+    override fun getCustomizationSections(activity: Activity): CustomizationSections {
+        return customizationSections
+            ?: DefaultCustomizationSections(
+                    getKeyguardQuickAffordancePickerInteractor(activity),
+                    getKeyguardQuickAffordancePickerViewModelFactory(activity)
+                )
+                .also { customizationSections = it }
+    }
+
+    override fun getDeepLinkRedirectIntent(context: Context, uri: Uri): Intent {
+        val intent = Intent()
+        intent.setClass(context, CustomizationPickerActivity::class.java)
+        intent.data = uri
+        intent.flags = Intent.FLAG_ACTIVITY_NEW_TASK or Intent.FLAG_ACTIVITY_CLEAR_TASK
+        return intent
+    }
+
+    override fun getDownloadableIntentAction(): String? {
+        return null
+    }
+
+    override fun getPreviewFragment(
+        context: Context,
+        wallpaperInfo: WallpaperInfo,
+        mode: Int,
+        viewAsHome: Boolean,
+        viewFullScreen: Boolean,
+        testingModeEnabled: Boolean
+    ): Fragment {
+        return if (wallpaperInfo is LiveWallpaperInfo) LivePreviewFragment()
+        else
+            ImagePreviewFragment().apply {
+                arguments =
+                    Bundle().apply {
+                        putParcelable(PreviewFragment.ARG_WALLPAPER, wallpaperInfo)
+                        putInt(PreviewFragment.ARG_PREVIEW_MODE, mode)
+                        putBoolean(PreviewFragment.ARG_VIEW_AS_HOME, viewAsHome)
+                        putBoolean(PreviewFragment.ARG_FULL_SCREEN, viewFullScreen)
+                        putBoolean(PreviewFragment.ARG_TESTING_MODE_ENABLED, testingModeEnabled)
+                    }
+            }
+    }
+
+    @Synchronized
+    override fun getUserEventLogger(context: Context): ThemesUserEventLogger {
+        return if (userEventLogger != null) userEventLogger as ThemesUserEventLogger
+        else StatsLogUserEventLogger(context).also { userEventLogger = it }
+    }
+
+    @Synchronized
+    override fun getPreferences(context: Context): WallpaperPreferences {
+        return prefs
+            ?: DefaultCustomizationPreferences(context.applicationContext).also { prefs = it }
+    }
+
+    override fun getFragmentFactory(): FragmentFactory? {
+        return fragmentFactory ?: ThemePickerFragmentFactory().also { fragmentFactory }
+    }
+
+    override fun getSnapshotRestorers(context: Context): Map<Int, SnapshotRestorer> {
+        return super<WallpaperPicker2Injector>.getSnapshotRestorers(context).toMutableMap().apply {
+            this[KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER] =
+                getKeyguardQuickAffordanceSnapshotRestorer(context)
+        }
+    }
+
+    override fun getCustomizationPreferences(context: Context): CustomizationPreferences {
+        return getPreferences(context) as CustomizationPreferences
+    }
+
+    override fun getThemeManager(
+        provider: ThemeBundleProvider,
+        activity: FragmentActivity,
+        overlayManagerCompat: OverlayManagerCompat,
+        logger: ThemesUserEventLogger
+    ): ThemeManager {
+        return ThemeManager(provider, activity, overlayManagerCompat, logger)
+    }
+
+    override fun getKeyguardQuickAffordancePickerInteractor(
+        context: Context
+    ): KeyguardQuickAffordancePickerInteractor {
+        return keyguardQuickAffordancePickerInteractor
+            ?: getKeyguardQuickAffordancePickerInteractorImpl(context).also {
+                keyguardQuickAffordancePickerInteractor = it
+            }
+    }
+
+    fun getKeyguardQuickAffordancePickerViewModelFactory(
+        context: Context
+    ): KeyguardQuickAffordancePickerViewModel.Factory {
+        return keyguardQuickAffordancePickerViewModelFactory
+            ?: KeyguardQuickAffordancePickerViewModel.Factory(
+                    context,
+                    getKeyguardQuickAffordancePickerInteractor(context),
+                    getUndoInteractor(context),
+                    getCurrentWallpaperInfoFactory(context),
+                )
+                .also { keyguardQuickAffordancePickerViewModelFactory = it }
+    }
+
+    private fun getKeyguardQuickAffordancePickerInteractorImpl(
+        context: Context
+    ): KeyguardQuickAffordancePickerInteractor {
+        val client = getKeyguardQuickAffordancePickerProviderClient(context)
+        return KeyguardQuickAffordancePickerInteractor(
+            KeyguardQuickAffordancePickerRepository(client, IO),
+            client
+        ) { getKeyguardQuickAffordanceSnapshotRestorer(context) }
+    }
+
+    protected fun getKeyguardQuickAffordancePickerProviderClient(
+        context: Context
+    ): CustomizationProviderClient {
+        return customizationProviderClient
+            ?: CustomizationProviderClientImpl(context, IO).also {
+                customizationProviderClient = it
+            }
+    }
+
+    protected fun getKeyguardQuickAffordanceSnapshotRestorer(
+        context: Context
+    ): KeyguardQuickAffordanceSnapshotRestorer {
+        return keyguardQuickAffordanceSnapshotRestorer
+            ?: KeyguardQuickAffordanceSnapshotRestorer(
+                    getKeyguardQuickAffordancePickerInteractor(context),
+                    getKeyguardQuickAffordancePickerProviderClient(context)
+                )
+                .also { keyguardQuickAffordanceSnapshotRestorer = it }
+    }
+
+    companion object {
+        @JvmStatic
+        private val KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER =
+            WallpaperPicker2Injector.MIN_SNAPSHOT_RESTORER_KEY
+
+        /**
+         * When this injector is overridden, this is the minimal value that should be used by
+         * restorers returns in [getSnapshotRestorers].
+         */
+        @JvmStatic
+        protected val MIN_SNAPSHOT_RESTORER_KEY = KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER + 1
+    }
+}