Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -4,15 +4,19 @@ import torch
|
|
4 |
import re
|
5 |
from PIL import Image
|
6 |
import spaces # Add spaces import for Hugging Face Spaces
|
|
|
|
|
|
|
7 |
|
8 |
# Model information
|
9 |
MODEL_ID = "DeepMount00/SmolVLM-Base-ocr_base"
|
10 |
OCR_INSTRUCTION = "Sei un assistente esperto di OCR, converti il testo in formato MD."
|
11 |
|
12 |
# Load processor and model
|
13 |
-
processor = AutoProcessor.from_pretrained(MODEL_ID)
|
14 |
model = AutoModelForVision2Seq.from_pretrained(
|
15 |
MODEL_ID,
|
|
|
16 |
torch_dtype=torch.bfloat16,
|
17 |
).to("cuda") # Ensure model loads on CUDA for Spaces
|
18 |
|
@@ -52,7 +56,7 @@ def process_image(image, progress=gr.Progress()):
|
|
52 |
with torch.no_grad():
|
53 |
generated_ids = model.generate(
|
54 |
**inputs,
|
55 |
-
max_new_tokens=
|
56 |
temperature=0.1
|
57 |
)
|
58 |
|
|
|
4 |
import re
|
5 |
from PIL import Image
|
6 |
import spaces # Add spaces import for Hugging Face Spaces
|
7 |
+
import os
|
8 |
+
|
9 |
+
hf_token = os.environ.get("API_KEY")
|
10 |
|
11 |
# Model information
|
12 |
MODEL_ID = "DeepMount00/SmolVLM-Base-ocr_base"
|
13 |
OCR_INSTRUCTION = "Sei un assistente esperto di OCR, converti il testo in formato MD."
|
14 |
|
15 |
# Load processor and model
|
16 |
+
processor = AutoProcessor.from_pretrained(MODEL_ID, token=hf_token)
|
17 |
model = AutoModelForVision2Seq.from_pretrained(
|
18 |
MODEL_ID,
|
19 |
+
token=hf_token,
|
20 |
torch_dtype=torch.bfloat16,
|
21 |
).to("cuda") # Ensure model loads on CUDA for Spaces
|
22 |
|
|
|
56 |
with torch.no_grad():
|
57 |
generated_ids = model.generate(
|
58 |
**inputs,
|
59 |
+
max_new_tokens=8192,
|
60 |
temperature=0.1
|
61 |
)
|
62 |
|