Enable GPU auto-selection (CUDA if available) with DEVICE override; load TorchScript model on selected device. Reference: https://huggingface.co/spaces/LogicGoInfotechSpaces/Hair_stable_new/tree/main

src/core.py

@@ -64,9 +64,10 @@ ENERGY_MASK_CONST = 100000.0              # large energy value for protective ma
 MASK_THRESHOLD = 10                       # minimum pixel intensity for binary mask
 USE_FORWARD_ENERGY = True                 # if True, use forward energy algorithm
 
-device = torch.device("cpu")
+device_str = os.environ.get("DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device(device_str)
 model_path = "./assets/big-lama.pt"
-model = torch.jit.load(model_path, map_location="cpu")
+model = torch.jit.load(model_path, map_location=device)
 model = model.to(device)
 model.eval()