diff --git a/app/src/main/java/com/maary/shareas/WallpaperViewModel.kt b/app/src/main/java/com/maary/shareas/WallpaperViewModel.kt
index 6835096..e40099a 100644
--- a/app/src/main/java/com/maary/shareas/WallpaperViewModel.kt
+++ b/app/src/main/java/com/maary/shareas/WallpaperViewModel.kt
@@ -136,6 +136,10 @@ class WallpaperViewModel : ViewModel() {
     private val _upscaleProgressState = MutableStateFlow(0)
     val upscaleProgressState = _upscaleProgressState.asStateFlow()
 
+    fun getUpscaleProgress(): Int {
+        return _upscaleProgressState.value
+    }
+
     private val _upscaleToggleState = MutableStateFlow(upscaleToggle)
     val upscaleToggleState = _upscaleToggleState.asStateFlow()
 
diff --git a/app/src/main/java/com/maary/shareas/fragment/editor/UpscaleFragment.kt b/app/src/main/java/com/maary/shareas/fragment/editor/UpscaleFragment.kt
index 7ce523d..20780fd 100644
--- a/app/src/main/java/com/maary/shareas/fragment/editor/UpscaleFragment.kt
+++ b/app/src/main/java/com/maary/shareas/fragment/editor/UpscaleFragment.kt
@@ -5,6 +5,7 @@ import android.os.Bundle
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import android.view.WindowManager
 import androidx.activity.OnBackPressedCallback
 import androidx.fragment.app.Fragment
 import androidx.fragment.app.activityViewModels
@@ -63,8 +64,11 @@ class UpscaleFragment : Fragment() {
                             viewModel.upscale(requireContext(), binding.menuChooseModelTextview.text.toString())
                         }
                         false -> {
+                            activity?.window?.clearFlags(WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON)
                             binding.buttonUpscaleToggle.setIconResource(R.drawable.ic_play)
-                            binding.progressUpscale.setProgressCompat(0, true)
+                            if (viewModel.getUpscaleProgress() != 100) {
+                                binding.progressUpscale.setProgressCompat(0, true)
+                            }
                         }
                     }
                 }
@@ -108,6 +112,7 @@ class UpscaleFragment : Fragment() {
         binding.menuChooseModelTextview.setText(resources.getStringArray(R.array.model_names)[2], false)
 
         binding.buttonUpscaleToggle.setOnClickListener {
+            activity?.window?.addFlags(WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON)
             viewModel.upscaleToggle = !viewModel.upscaleToggle
         }