salso commited on
Commit
0e5a0f3
·
verified ·
1 Parent(s): 75a8eb7

Update florence_sam/utils/florence.py

Browse files
Files changed (1) hide show
  1. florence_sam/utils/florence.py +2 -1
florence_sam/utils/florence.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import os
2
  from typing import Union, Any, Tuple, Dict
3
  from unittest.mock import patch
@@ -23,7 +24,7 @@ def fixed_get_imports(filename: Union[str, os.PathLike]) -> list[str]:
23
  #imports.remove("flash_attn")
24
  return imports
25
 
26
-
27
  def load_florence_model(
28
  device: torch.device, checkpoint: str = FLORENCE_CHECKPOINT
29
  ) -> Tuple[Any, Any]:
 
1
+ import spaces
2
  import os
3
  from typing import Union, Any, Tuple, Dict
4
  from unittest.mock import patch
 
24
  #imports.remove("flash_attn")
25
  return imports
26
 
27
+ @spaces.GPU
28
  def load_florence_model(
29
  device: torch.device, checkpoint: str = FLORENCE_CHECKPOINT
30
  ) -> Tuple[Any, Any]: