Update app.py

app.py

@@ -44,15 +44,6 @@ model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to(device).eval()
 
-# Load SmolDocling-256M-preview
-MODEL_ID_X = "ds4sd/SmolDocling-256M-preview"
-processor_x = AutoProcessor.from_pretrained(MODEL_ID_X, trust_remote_code=True)
-model_x = AutoModelForVision2Seq.from_pretrained(
-    MODEL_ID_X,
-    trust_remote_code=True,
-    torch_dtype=torch.float16
-).to(device).eval()
-
 # Load MonkeyOCR
 MODEL_ID_G = "echo840/MonkeyOCR"
 SUBFOLDER = "Recognition"
@@ -68,6 +59,26 @@ model_g = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     torch_dtype=torch.float16
 ).to(device).eval()
 
+# Load NuExtract-2.0-8B
+MODEL_ID_L = "numind/NuExtract-2.0-8B"
+processor_l = AutoProcessor.from_pretrained(MODEL_ID_L, trust_remote_code=True)
+model_l = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    MODEL_ID_L,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+).to(device).eval()
+
+#--------------------------------------------------#
+# Load SmolDocling-256M-preview
+MODEL_ID_X = "ds4sd/SmolDocling-256M-preview"
+processor_x = AutoProcessor.from_pretrained(MODEL_ID_X, trust_remote_code=True)
+model_x = AutoModelForVision2Seq.from_pretrained(
+    MODEL_ID_X,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+).to(device).eval()
+#--------------------------------------------------#
+
 # Preprocessing functions for SmolDocling-256M
 def add_random_padding(image, min_percent=0.1, max_percent=0.10):
     """Add random padding to an image based on its size."""
@@ -132,6 +143,9 @@ def generate_image(model_name: str, text: str, image: Image.Image,
     elif model_name == "SmolDocling-256M-preview":
         processor = processor_x
         model = model_x
+    elif model_name == "NuExtract-2.0-8B":
+        processor = processor_l
+        model = model_l
     else:
         yield "Invalid model selected."
         return
@@ -216,6 +230,9 @@ def generate_video(model_name: str, text: str, video_path: str,
     elif model_name == "SmolDocling-256M-preview":
         processor = processor_x
         model = model_x
+    elif model_name == "NuExtract-2.0-8B":
+        processor = processor_l
+        model = model_l
     else:
         yield "Invalid model selected."
         return
@@ -336,7 +353,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         with gr.Column():
             output = gr.Textbox(label="Output", interactive=False, lines=3, scale=2)
             model_choice = gr.Radio(
-                choices=["Nanonets-OCR-s", "MonkeyOCR-Recognition", "SmolDocling-256M-preview"],
+                choices=["Nanonets-OCR-s", "MonkeyOCR-Recognition", "SmolDocling-256M-preview", "NuExtract-2.0-8B"],
                 label="Select Model",
                 value="Nanonets-OCR-s"
             )