Update README.md

README.md

@@ -40,9 +40,8 @@ import torch
 from PIL import Image
 from transformers import VisionEncoderDecoderModel
 from transformers.models.nougat import NougatTokenizerFast
-
 from nougat_latex import NougatLaTexProcessor
-from nougat_latex.image_processing_nougat import NougatImageProcessor
+
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 # init model
@@ -51,15 +50,14 @@ model = VisionEncoderDecoderModel.from_pretrained("Norm/nougat-latex-base").to(d
 # init processor
 tokenizer = NougatTokenizerFast.from_pretrained("Norm/nougat-latex-base")
 
-image_processor = NougatImageProcessor.from_pretrained("Norm/nougat-latex-base")
-latex_processor = NougatLaTexProcessor(image_processor=image_processor)
+latex_processor = NougatLaTexProcessor.from_pretrained("Norm/nougat-latex-base")
 
 # run test
 image = Image.open("path/to/latex/image.png")
 if not image.mode == "RGB":
     image = image.convert('RGB')
 
-pixel_values = latex_processor(image)
+pixel_values = latex_processor(image, return_tensors="pt").pixel_values
 
 decoder_input_ids = tokenizer(tokenizer.bos_token, add_special_tokens=False,
                               return_tensors="pt").input_ids