geetu040 commited on
Commit
3b3903e
·
1 Parent(s): 182701e

use device in example

Browse files
Files changed (1) hide show
  1. README.md +4 -2
README.md CHANGED
@@ -66,13 +66,15 @@ from PIL import Image
66
  import torch
67
  from transformers import DepthProImageProcessorFast, DepthProForDepthEstimation
68
 
 
 
69
  url = 'https://huggingface.co/datasets/mishig/sample_images/resolve/main/tiger.jpg'
70
  image = Image.open(requests.get(url, stream=True).raw)
71
 
72
  image_processor = DepthProImageProcessorFast.from_pretrained("geetu040/DepthPro")
73
- model = DepthProForDepthEstimation.from_pretrained("geetu040/DepthPro")
74
 
75
- inputs = image_processor(images=image, return_tensors="pt")
76
 
77
  with torch.no_grad():
78
  outputs = model(**inputs)
 
66
  import torch
67
  from transformers import DepthProImageProcessorFast, DepthProForDepthEstimation
68
 
69
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
70
+
71
  url = 'https://huggingface.co/datasets/mishig/sample_images/resolve/main/tiger.jpg'
72
  image = Image.open(requests.get(url, stream=True).raw)
73
 
74
  image_processor = DepthProImageProcessorFast.from_pretrained("geetu040/DepthPro")
75
+ model = DepthProForDepthEstimation.from_pretrained("geetu040/DepthPro").to(device)
76
 
77
+ inputs = image_processor(images=image, return_tensors="pt").to(device)
78
 
79
  with torch.no_grad():
80
  outputs = model(**inputs)