ResearcherXman commited on
Commit
5853f77
1 Parent(s): 2192aaf
Files changed (1) hide show
  1. controlnet_util.py +1 -2
controlnet_util.py CHANGED
@@ -2,13 +2,12 @@ import torch
2
  import numpy as np
3
  from PIL import Image
4
  from controlnet_aux import OpenposeDetector
5
- from model_util import get_torch_device
6
  import cv2
7
 
8
 
9
  from transformers import DPTImageProcessor, DPTForDepthEstimation
10
 
11
- device = get_torch_device()
12
  depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to(device)
13
  feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas")
14
  openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
 
2
  import numpy as np
3
  from PIL import Image
4
  from controlnet_aux import OpenposeDetector
 
5
  import cv2
6
 
7
 
8
  from transformers import DPTImageProcessor, DPTForDepthEstimation
9
 
10
+ device = "cuda" if torch.cuda.is_available() else "cpu"
11
  depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to(device)
12
  feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas")
13
  openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")