SkalskiP commited on
Commit
7e2035e
·
1 Parent(s): 3b99a8a

make sure the space code runs with cuda

Browse files
Files changed (2) hide show
  1. app.py +2 -1
  2. utils/models.py +1 -1
app.py CHANGED
@@ -52,7 +52,8 @@ OCR_WITH_REGION_EXAMPLES = [
52
  ["microsoft/Florence-2-large-ft", OCR_WITH_REGION_TASK_NAME, "https://media.roboflow.com/inference/license_plate_1.jpg"]
53
  ]
54
 
55
- DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
56
  MODELS, PROCESSORS = load_models(DEVICE)
57
 
58
 
 
52
  ["microsoft/Florence-2-large-ft", OCR_WITH_REGION_TASK_NAME, "https://media.roboflow.com/inference/license_plate_1.jpg"]
53
  ]
54
 
55
+ # DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
56
+ DEVICE = "cuda"
57
  MODELS, PROCESSORS = load_models(DEVICE)
58
 
59
 
utils/models.py CHANGED
@@ -20,7 +20,7 @@ def load_models(device: torch.device) -> Tuple[Dict[str, Any], Dict[str, Any]]:
20
  processors = {}
21
  for checkpoint in CHECKPOINTS:
22
  models[checkpoint] = AutoModelForCausalLM.from_pretrained(
23
- checkpoint, trust_remote_code=True).to(device)
24
  processors[checkpoint] = AutoProcessor.from_pretrained(
25
  checkpoint, trust_remote_code=True)
26
  return models, processors
 
20
  processors = {}
21
  for checkpoint in CHECKPOINTS:
22
  models[checkpoint] = AutoModelForCausalLM.from_pretrained(
23
+ checkpoint, trust_remote_code=True).to(device).eval()
24
  processors[checkpoint] = AutoProcessor.from_pretrained(
25
  checkpoint, trust_remote_code=True)
26
  return models, processors