Update app.py

app.py

@@ -62,14 +62,14 @@ logger.info("Model downloaded to: %s", model_path)
 logger.info('Loading ControlNet model.')
 cache_dir = "./model_cache"  # Customize your cache directory
 controlnet = FluxControlNetModel.from_pretrained(
-    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.float16,cache_dir=cache_dir,low_cpu_mem_usage=True,device_map="auto"
-)
+    "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.float16,cache_dir=cache_dir,low_cpu_mem_usage=True
+).to(device)
 logger.info("ControlNet model loaded successfully.")
 
 logger.info('Loading pipeline.')
 pipe = FluxControlNetPipeline.from_pretrained(
-    model_path, controlnet=controlnet, torch_dtype=torch.float16,cache_dir=cache_dir,low_cpu_mem_usage=True,device_map="auto"
-)
+    model_path, controlnet=controlnet, torch_dtype=torch.float16,cache_dir=cache_dir,low_cpu_mem_usage=True
+).to(device)
 logger.info("Pipeline loaded successfully.")
 
 MAX_SEED = 1000000