DeepMount00 commited on
Commit
9a58a0a
·
verified ·
1 Parent(s): 9b818aa

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -2
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=4096,
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