Update app.py

app.py

@@ -32,13 +32,13 @@ app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True,
 # ThreadPoolExecutor for managing image processing threads
 executor = ThreadPoolExecutor()
 
-# Determine the device (GPU or CPU)
-# if torch.cuda.is_available():
-#     device = "cuda"
-#     logger.info("CUDA is available. Using GPU.")
-# else:
-device = "cpu"
-    # logger.info("CUDA is not available. Using CPU.")
+#Determine the device (GPU or CPU)
+if torch.cuda.is_available():
+    device = "cuda"
+    logger.info("CUDA is available. Using GPU.")
+else:
+    device = "cpu"
+    logger.info("CUDA is not available. Using CPU.")
 
 # Load model from Huggingface Hub
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
@@ -60,14 +60,15 @@ logger.info("Model downloaded to: %s", model_path)
 
 # Load pipeline
 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.bfloat16
+    "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.bfloat16
+    model_path, controlnet=controlnet, torch_dtype=torch.float16,cache_dir=cache_dir, low_cpu_mem_usage=True
 ).to(device)
 logger.info("Pipeline loaded successfully.")