Skip to content

Commit 5424a8b

Browse files
committed
Add environment token support for Hugging Face models
- Use `HF_TOKEN` environment variable as fallback for Hugging Face authentication - Simplify token handling in `HuggingfaceDetectionModel.load_model()` - Ensure consistent token usage across model and processor loading
1 parent e983fff commit 5424a8b

File tree

1 file changed

+5
-3
lines changed

1 file changed

+5
-3
lines changed

sahi/models/huggingface.py

+5-3
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
# Code written by Fatih C Akyon and Devrim Cavusoglu, 2022.
33

44
import logging
5+
import os
56
from typing import Any, Dict, List, Optional, Tuple, Union
67

78
import numpy as np
@@ -69,18 +70,19 @@ def num_categories(self) -> int:
6970
def load_model(self):
7071
from transformers import AutoModelForObjectDetection, AutoProcessor
7172

72-
model = AutoModelForObjectDetection.from_pretrained(self.model_path, token=self._token)
73+
hf_token = os.getenv("HF_TOKEN", self._token)
74+
model = AutoModelForObjectDetection.from_pretrained(self.model_path, token=hf_token)
7375
if self.image_size is not None:
7476
if model.base_model_prefix == "rt_detr_v2":
7577
size = {"height": self.image_size, "width": self.image_size}
7678
else:
7779
size = {"shortest_edge": self.image_size, "longest_edge": None}
7880
# use_fast=True raises error: AttributeError: 'SizeDict' object has no attribute 'keys'
7981
processor = AutoProcessor.from_pretrained(
80-
self.model_path, size=size, do_resize=True, use_fast=False, token=self._token
82+
self.model_path, size=size, do_resize=True, use_fast=False, token=hf_token
8183
)
8284
else:
83-
processor = AutoProcessor.from_pretrained(self.model_path, use_fast=False, token=self._token)
85+
processor = AutoProcessor.from_pretrained(self.model_path, use_fast=False, token=hf_token)
8486
self.set_model(model, processor)
8587

8688
def set_model(self, model: Any, processor: Any = None):

0 commit comments

Comments
 (0)