gokaygokay commited on
Commit
69958d1
1 Parent(s): d256f3b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -16,7 +16,7 @@ import subprocess
16
  subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
17
 
18
  model_id = 'microsoft/Florence-2-large'
19
- model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).eval()
20
  processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
21
 
22
 
@@ -28,13 +28,14 @@ def fig_to_pil(fig):
28
  fig.savefig(buf, format='png')
29
  buf.seek(0)
30
  return Image.open(buf)
31
-
 
32
  def run_example(task_prompt, image, text_input=None):
33
  if text_input is None:
34
  prompt = task_prompt
35
  else:
36
  prompt = task_prompt + text_input
37
- inputs = processor(text=prompt, images=image, return_tensors="pt")
38
  generated_ids = model.generate(
39
  input_ids=inputs["input_ids"],
40
  pixel_values=inputs["pixel_values"],
 
16
  subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
17
 
18
  model_id = 'microsoft/Florence-2-large'
19
+ model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to("cuda").eval()
20
  processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
21
 
22
 
 
28
  fig.savefig(buf, format='png')
29
  buf.seek(0)
30
  return Image.open(buf)
31
+
32
+ @spaces.GPU
33
  def run_example(task_prompt, image, text_input=None):
34
  if text_input is None:
35
  prompt = task_prompt
36
  else:
37
  prompt = task_prompt + text_input
38
+ inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda")
39
  generated_ids = model.generate(
40
  input_ids=inputs["input_ids"],
41
  pixel_values=inputs["pixel_values"],