Update routers/donut_inference.py

routers/donut_inference.py

@@ -5,12 +5,21 @@ 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():
-    processor = DonutProcessor.from_pretrained(settings.processor)
-    model = VisionEncoderDecoderModel.from_pretrained(settings.model)
+def load_model(model_url: str):
+    """
+    Load the processor and model dynamically based on the model URL.
+    
+    :param model_url: The URL for the model to use.
+    :return: The processor, model, and device.
+    """
+    # Assuming the model URL follows a pattern like "https://huggingface.co/{model_name}"
+    model_name = model_url.split("/")[-1]  # Extract model name from the URL
+
+    processor = DonutProcessor.from_pretrained(model_name)
+    model = VisionEncoderDecoderModel.from_pretrained(model_name)
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model.to(device)
@@ -18,22 +27,30 @@ def load_model():
     return processor, model, device
 
 
-def process_document_donut(image):
+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.
+    """
     worker_pid = os.getpid()
     print(f"Handling inference request with worker PID: {worker_pid}")
 
     start_time = time.time()
 
-    processor, model, device = load_model()
+    # Load the model dynamically based on the model_url
+    processor, model, device = load_model(model_url)
 
-    # 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),
@@ -47,10 +64,10 @@ def process_document_donut(image):
         return_dict_in_generate=True,
     )
 
-    # postprocess
+    # Postprocess the result
     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