Test

app.py

@@ -5,6 +5,7 @@ import json
 import gc
 import torch
 import io
+from PIL import Image
 from transformers import AutoProcessor, AutoModelForImageTextToText
 from qwen_vl_utils import process_vision_info
 import gradio as gr
@@ -12,15 +13,13 @@ import spaces
 
 # Model setup
 MODEL_NAME = "numind/NuExtract-2.0-4B"
-device = "cuda"  # ZeroGPU provides GPU
+device = "cuda"
 
 model = AutoModelForImageTextToText.from_pretrained(
     MODEL_NAME,
     trust_remote_code=True,
     dtype=torch.bfloat16,
-    device_map=None,  # Load on CPU, move to GPU in function
 )
-
 processor = AutoProcessor.from_pretrained(
     MODEL_NAME,
     trust_remote_code=True,
@@ -46,106 +45,106 @@ invoice_schema = {
     ]
 }
 
-def encode_image_to_base64(image_path):
-    with open(image_path, "rb") as img_file:
-        return base64.b64encode(img_file.read()).decode("utf-8")
 
 def encode_image_from_pil(image):
     buffer = io.BytesIO()
     image.save(buffer, format="PNG")
     return base64.b64encode(buffer.getvalue()).decode("utf-8")
 
-def prepare_prompt(image_path):
-    base64_image = encode_image_to_base64(image_path)
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image", "image": f"data:image;base64,{base64_image}"}
-            ]
-        }
-    ]
-    text = processor.tokenizer.apply_chat_template(
-        messages,
-        template=json.dumps(invoice_schema, indent=4),
-        tokenize=False,
-        add_generation_prompt=True
-    )
-    return messages, text
 
 @spaces.GPU
-def process_image(image, schema_str):
-    if image is None:
-        return "No image provided."
-    
+def process_images(files, schema_str):
+    if not files:
+        return "No images provided."
+
     try:
         custom_schema = json.loads(schema_str)
     except json.JSONDecodeError:
-        return "Invalid JSON schema provided."
-    
-    base64_str = encode_image_from_pil(image)
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image", "image": f"data:image;base64,{base64_str}"}
-            ]
-        }
-    ]
-    text = processor.tokenizer.apply_chat_template(
-        messages,
-        template=json.dumps(custom_schema, indent=4),
-        tokenize=False,
-        add_generation_prompt=True
-    )
-    
-    image_inputs = process_vision_info(messages)[0] or []
-    
-    inputs = processor(
-        text=[text],
-        images=image_inputs,
-        padding=True,
-        return_tensors="pt",
-    ).to(device)
-    
-    # Move model to GPU
+        return "Invalid JSON schema."
+
+    results = []
+
     model.to(device)
-    
-    generation_config = {
-        "do_sample": False,
-        "num_beams": 1,
-        "max_new_tokens": 2048,
-    }
-    
-    generated_ids = model.generate(**inputs, **generation_config)
-    
-    generated_ids_trimmed = [
-        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-    ]
-    
-    output_text = processor.batch_decode(
-        generated_ids_trimmed,
-        skip_special_tokens=True,
-        clean_up_tokenization_spaces=False,
-    )[0]
-    
-    return output_text
-
-# Gradio interface
+
+    for file_obj in files:
+        image = Image.open(file_obj.name).convert("RGB")
+        base64_str = encode_image_from_pil(image)
+
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": f"data:image;base64,{base64_str}"}
+                ]
+            }
+        ]
+
+        text = processor.tokenizer.apply_chat_template(
+            messages,
+            template=json.dumps(custom_schema, indent=4),
+            tokenize=False,
+            add_generation_prompt=True
+        )
+
+        image_inputs = process_vision_info(messages)[0] or []
+
+        inputs = processor(
+            text=[text],
+            images=image_inputs,
+            padding=True,
+            return_tensors="pt",
+        ).to(device)
+
+        generated_ids = model.generate(
+            **inputs,
+            do_sample=False,
+            num_beams=1,
+            max_new_tokens=2048,
+        )
+
+        trimmed = [
+            out[len(in_ids):] for in_ids, out in zip(inputs.input_ids, generated_ids)
+        ]
+
+        output_text = processor.batch_decode(
+            trimmed,
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False,
+        )[0]
+
+        results.append({
+            "filename": os.path.basename(file_obj.name),
+            "output": output_text
+        })
+
+    return json.dumps(results, indent=4)
+
+
+# Gradio UI
 iface = gr.Interface(
-    fn=process_image,
+    fn=process_images,
     inputs=[
-        gr.Image(type="pil", label="Upload Invoice Image"),
+        gr.File(
+            label="Upload Invoice Images",
+            type="filepath",
+            file_count="multiple",
+        ),
         gr.Textbox(
             label="Custom Schema (JSON)",
             value=json.dumps(invoice_schema, indent=4),
-            lines=10,
-            placeholder="Enter your custom JSON schema here..."
+            lines=12,
         )
     ],
-    outputs=gr.Textbox(label="Extracted Data (JSON)"),
-    title="Invoice Parser with NuExtract",
-    description="Upload an invoice image and provide a custom JSON schema to extract structured data using AI."
+    outputs=gr.Textbox(
+        label="Extracted JSON Data",
+        lines=40,
+        max_lines=200,
+        autoscroll=True,
+        interactive=True,
+        show_copy_button=True,
+    ),
+    title="Invoice Parser with NuExtract (Multi-Image)",
+    description="Upload one or more invoice images. Each will be processed independently with your custom JSON schema.",
 )
 
-iface.launch()
+iface.launch()