Spaces:
Running
on
Zero
Running
on
Zero
make sure the space code runs with cuda
Browse files- app.py +2 -1
- 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
|