ARCQUB commited on
Commit
1fa4f4e
·
verified ·
1 Parent(s): 97fc69c

Update models/gpt4o.py

Browse files
Files changed (1) hide show
  1. models/gpt4o.py +1 -1
models/gpt4o.py CHANGED
@@ -58,7 +58,7 @@ def extract_all_text_pix2struct(image: Image.Image):
58
  processor = Pix2StructProcessor.from_pretrained("google/pix2struct-textcaps-base")
59
  pix2struct_model = Pix2StructForConditionalGeneration.from_pretrained(
60
  "google/pix2struct-textcaps-base"
61
- ).to("cuda" if torch.cuda.is_available() else "cpu")
62
 
63
  inputs = processor(images=image, return_tensors="pt").to(pix2struct_model.device)
64
  predictions = pix2struct_model.generate(**inputs, max_new_tokens=512)
 
58
  processor = Pix2StructProcessor.from_pretrained("google/pix2struct-textcaps-base")
59
  pix2struct_model = Pix2StructForConditionalGeneration.from_pretrained(
60
  "google/pix2struct-textcaps-base"
61
+ ).to("cpu")
62
 
63
  inputs = processor(images=image, return_tensors="pt").to(pix2struct_model.device)
64
  predictions = pix2struct_model.generate(**inputs, max_new_tokens=512)