triphuong57 commited on
Commit
6fd6f53
1 Parent(s): 6dabe1c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -1,11 +1,13 @@
1
  import gradio as gr
2
  from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
3
  from peft import PeftModel
 
4
 
5
-
6
  def greet(image, prompt):
7
- model = PaliGemmaForConditionalGeneration.from_pretrained("/folders")
8
  processor = AutoProcessor.from_pretrained("google/paligemma-3b-mix-224")
 
9
  inputs = processor(prompt, raw_image, return_tensors="pt")
10
  output = model.generate(**inputs, max_new_tokens=20)
11
  return output
 
1
  import gradio as gr
2
  from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
3
  from peft import PeftModel
4
+ import spaces
5
 
6
+ @spaces.GPU
7
  def greet(image, prompt):
8
+ base_model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma-3b-mix-224")
9
  processor = AutoProcessor.from_pretrained("google/paligemma-3b-mix-224")
10
+ model = PeftModel(base_model, "/folders")
11
  inputs = processor(prompt, raw_image, return_tensors="pt")
12
  output = model.generate(**inputs, max_new_tokens=20)
13
  return output