Update app.py

app.py

@@ -167,6 +167,7 @@ def enhance_image(
             pipe.enable_vae_tiling()
             pipe.enable_vae_slicing()
             if device == "cuda":
+                pipe.reset_device_map()
                 pipe.enable_model_cpu_offload()
 
             if USE_ESRGAN:
@@ -191,7 +192,7 @@ def enhance_image(
                 "black-forest-labs/FLUX.1-dev",
                 torch_dtype=dtype,
                 low_cpu_mem_usage=True,
-                device_map="balanced"
+                device_map=None
             )
             pipe.enable_vae_tiling()
             pipe.enable_vae_slicing()