Fix using only generated_text[0]

app.py

@@ -7,7 +7,7 @@ from transformers import pipeline
 from ultralytics import YOLO
 from PIL import Image
 
-def process(path, progress = gr.Progress()):
+def process(path, progress = gr.Progress(), device = 'cpu'):
     progress(0, desc="Starting")
     LINE_MODEL_PATH = "Kansallisarkisto/multicentury-textline-detection"
     OCR_MODEL_PATH = "microsoft/trocr-large-handwritten"
@@ -16,6 +16,7 @@ def process(path, progress = gr.Progress()):
     # Load the model and processor
     processor = TrOCRProcessor.from_pretrained(OCR_MODEL_PATH)
     model = VisionEncoderDecoderModel.from_pretrained(OCR_MODEL_PATH)
+    model.to(device)
     
     # Open an image of handwritten text
     image = Image.open(path).convert("RGB")
@@ -44,9 +45,12 @@ def process(path, progress = gr.Progress()):
 
     #Predict and decode the entire batch
     progress(0, desc="Recognizing..")
+    batch = torch.cat(batch).to(device)
+    print("batch.shape", batch.shape)
     generated_ids = model.generate(torch.cat(batch))
     progress(0, desc="Decoding (token -> str)")
-    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
+
     print(generated_text)
     full_text = " ".join(generated_text)
     print(full_text)