Spaces:
Sleeping
Sleeping
jhj0517
commited on
Commit
•
8d01c02
1
Parent(s):
c64775c
Update dtype
Browse files- modules/sam_inference.py +1 -2
modules/sam_inference.py
CHANGED
@@ -40,7 +40,6 @@ class SamInference:
|
|
40 |
self.current_model_type = DEFAULT_MODEL_TYPE
|
41 |
self.model_dir = model_dir
|
42 |
self.output_dir = output_dir
|
43 |
-
self.model_path = os.path.join(self.model_dir, AVAILABLE_MODELS[DEFAULT_MODEL_TYPE][0])
|
44 |
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
45 |
self.dtype = torch.float16 if torch.cuda.is_available() else torch.float32
|
46 |
self.mask_generator = None
|
@@ -272,7 +271,7 @@ class SamInference:
|
|
272 |
)
|
273 |
images = get_frames_from_dir(vid_dir=TEMP_DIR, as_numpy=True)
|
274 |
|
275 |
-
with torch.autocast(device_type=self.device, dtype=
|
276 |
for out_frame_idx, out_obj_ids, out_mask_logits in generator:
|
277 |
mask = (out_mask_logits[0] > 0.0).cpu().numpy()
|
278 |
video_segments[out_frame_idx] = {
|
|
|
40 |
self.current_model_type = DEFAULT_MODEL_TYPE
|
41 |
self.model_dir = model_dir
|
42 |
self.output_dir = output_dir
|
|
|
43 |
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
44 |
self.dtype = torch.float16 if torch.cuda.is_available() else torch.float32
|
45 |
self.mask_generator = None
|
|
|
271 |
)
|
272 |
images = get_frames_from_dir(vid_dir=TEMP_DIR, as_numpy=True)
|
273 |
|
274 |
+
with torch.autocast(device_type=self.device, dtype=self.dtype):
|
275 |
for out_frame_idx, out_obj_ids, out_mask_logits in generator:
|
276 |
mask = (out_mask_logits[0] > 0.0).cpu().numpy()
|
277 |
video_segments[out_frame_idx] = {
|