Update app.py
Browse files
app.py
CHANGED
@@ -12,13 +12,25 @@ if not HF_TOKEN:
|
|
12 |
st.error("❌ Hugging Face API key not found! Set it as 'HF_KEY' in Spaces secrets.")
|
13 |
st.stop()
|
14 |
|
15 |
-
# Load the model and processor
|
16 |
@st.cache_resource
|
17 |
def load_model():
|
18 |
model_id = "google/paligemma2-3b-mix-224"
|
19 |
-
|
20 |
-
|
21 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
22 |
|
23 |
processor, model = load_model()
|
24 |
|
|
|
12 |
st.error("❌ Hugging Face API key not found! Set it as 'HF_KEY' in Spaces secrets.")
|
13 |
st.stop()
|
14 |
|
15 |
+
# Load the model and processor with authentication
|
16 |
@st.cache_resource
|
17 |
def load_model():
|
18 |
model_id = "google/paligemma2-3b-mix-224"
|
19 |
+
|
20 |
+
try:
|
21 |
+
model = PaliGemmaForConditionalGeneration.from_pretrained(
|
22 |
+
model_id, token=HF_TOKEN, torch_dtype=torch.bfloat16, device_map="auto"
|
23 |
+
).eval()
|
24 |
+
|
25 |
+
processor = PaliGemmaProcessor.from_pretrained(
|
26 |
+
model_id, token=HF_TOKEN
|
27 |
+
)
|
28 |
+
|
29 |
+
return processor, model
|
30 |
+
|
31 |
+
except Exception as e:
|
32 |
+
st.error(f"❌ Error loading model: {str(e)}")
|
33 |
+
st.stop()
|
34 |
|
35 |
processor, model = load_model()
|
36 |
|