Akshayram1 commited on
Commit
5f4390d
·
verified ·
1 Parent(s): 69a9801

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +14 -2
app.py CHANGED
@@ -16,8 +16,20 @@ import spaces
16
  model_id = "gv-hf/paligemma2-10b-mix-448"
17
  COLORS = ['#4285f4', '#db4437', '#f4b400', '#0f9d58', '#e48ef1']
18
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
19
- model = PaliGemmaForConditionalGeneration.from_pretrained(model_id).eval().to(device)
20
- processor = PaliGemmaProcessor.from_pretrained(model_id)
 
 
 
 
 
 
 
 
 
 
 
 
21
 
22
  @spaces.GPU
23
  def infer(image: PIL.Image.Image, text: str, max_new_tokens: int) -> str:
 
16
  model_id = "gv-hf/paligemma2-10b-mix-448"
17
  COLORS = ['#4285f4', '#db4437', '#f4b400', '#0f9d58', '#e48ef1']
18
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
19
+ HF_KEY = os.getenv("HF_KEY")
20
+ if not HF_KEY:
21
+ raise ValueError("Please set the HF_KEY environment variable with your Hugging Face API token")
22
+
23
+ model = PaliGemmaForConditionalGeneration.from_pretrained(
24
+ model_id,
25
+ token=HF_KEY,
26
+ trust_remote_code=True
27
+ ).eval().to(device)
28
+
29
+ processor = PaliGemmaProcessor.from_pretrained(
30
+ model_id,
31
+ token=HF_KEY,
32
+ trust_remote_code=True)
33
 
34
  @spaces.GPU
35
  def infer(image: PIL.Image.Image, text: str, max_new_tokens: int) -> str: