Update routers/donut_inference.py

routers/donut_inference.py

@@ -5,16 +5,12 @@ from transformers import DonutProcessor, VisionEncoderDecoderModel
 from config import settings
 from functools import lru_cache
 import os
-import requests
 
 
 @lru_cache(maxsize=1)
-def load_model(model_url: str):
-    model_name = model_url.replace("https://huggingface.co/", "")
-    print(f"[Model Loader] Loading model: {model_name}")
-
-    processor = DonutProcessor.from_pretrained(model_name)
-    model = VisionEncoderDecoderModel.from_pretrained(model_name)
+def load_model():
+    processor = DonutProcessor.from_pretrained(settings.processor)
+    model = VisionEncoderDecoderModel.from_pretrained(settings.model)
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model.to(device)
@@ -22,30 +18,22 @@ def load_model(model_url: str):
     return processor, model, device
 
 
-def process_document_donut(image, model_url: str):
-    """
-    Process the document using the DONUT model.
-
-    :param image: The input image to process.
-    :param model_url: The model URL to use for inference.
-    :return: A tuple of the result and processing time.
-    """
+def process_document_donut(image):
     worker_pid = os.getpid()
-    print(f"[Inference] Handling request with worker PID: {worker_pid}")
+    print(f"Handling inference request with worker PID: {worker_pid}")
 
     start_time = time.time()
 
-    # Load the model dynamically based on the model_url
-    processor, model, device = load_model(model_url)
+    processor, model, device = load_model()
 
-    # Prepare encoder inputs
+    # prepare encoder inputs
     pixel_values = processor(image, return_tensors="pt").pixel_values
 
-    # Prepare decoder inputs
+    # prepare decoder inputs
     task_prompt = "<s_cord-v2>"
     decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
 
-    # Generate answer
+    # generate answer
     outputs = model.generate(
         pixel_values.to(device),
         decoder_input_ids=decoder_input_ids.to(device),
@@ -59,14 +47,14 @@ def process_document_donut(image, model_url: str):
         return_dict_in_generate=True,
     )
 
-    # Postprocess the result
+    # postprocess
     sequence = processor.batch_decode(outputs.sequences)[0]
     sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
-    sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # Remove first task start token
+    sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first task start token
 
     end_time = time.time()
     processing_time = end_time - start_time
 
-    print(f"[Inference] Done. PID: {worker_pid} | Time taken: {processing_time:.2f} sec")
+    print(f"Inference done, worker PID: {worker_pid}")
 
-    return processor.token2json(sequence), processing_time
+    return processor.token2json(sequence), processing_time