Update app.py

app.py

@@ -10,22 +10,27 @@ from typing import Optional, Tuple, Dict, Any, Iterable
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
+# --- Model & Script Download ---
 print("Downloading model snapshot to ensure all scripts are present...")
+# Download the full model repo to ensure postprocessing.py is available locally
 model_dir = snapshot_download(repo_id="nvidia/NVIDIA-Nemotron-Parse-v1.1")
 print(f"Model downloaded to: {model_dir}")
 
+# Add the model directory to sys.path so we can import postprocessing
 sys.path.append(model_dir)
 
 try:
     from postprocessing import extract_classes_bboxes, transform_bbox_to_original, postprocess_text
-    print("Successfully imported postprocessing functions.")
+    print("✅ Successfully imported postprocessing functions.")
 except ImportError as e:
-    print(f" Error importing postprocessing: {e}")
+    print(f"❌ Error importing postprocessing: {e}")
     raise e
 
+# --- Device Setup ---
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
+# --- Theme Definition ---
 colors.steel_blue = colors.Color(
     name="steel_blue",
     c50="#EBF3F8",
@@ -92,6 +97,7 @@ css = """
 #output-title h2 { font-size: 2.1em !important; }
 """
 
+# --- Model Loading ---
 print("Loading Model components...")
 
 processor = AutoProcessor.from_pretrained(model_dir, trust_remote_code=True)
@@ -107,7 +113,7 @@ except Exception as e:
     print(f"Warning: Could not load GenerationConfig: {e}. Using default.")
     generation_config = GenerationConfig(max_new_tokens=4096)
 
-print("Model loaded successfully.")
+print("✅ Model loaded successfully.")
 
 @spaces.GPU
 def process_ocr_task(image):
@@ -124,7 +130,7 @@ def process_ocr_task(image):
     if device.type == 'cuda':
         inputs = {k: v.to(torch.bfloat16) if v.dtype == torch.float32 else v for k, v in inputs.items()}
 
-    print("Running inference...")
+    print("🏃 Running inference...")
     with torch.no_grad():
         outputs = model.generate(
             **inputs, 
@@ -139,6 +145,7 @@ def process_ocr_task(image):
         print(f"Error extracting boxes: {e}")
         return generated_text, image
 
+    # Transform boxes to original image size
     bboxes = [transform_bbox_to_original(bbox, image.width, image.height) for bbox in bboxes]
 
     table_format = 'latex'  
@@ -169,8 +176,15 @@ def process_ocr_task(image):
     final_output_text = ""
     
     for cls, bbox, txt in zip(classes, bboxes, processed_texts):
+        # Normalize coordinates to prevent PIL ValueError (x1 >= x0)
+        x1, y1, x2, y2 = bbox
+        xmin = min(x1, x2)
+        ymin = min(y1, y2)
+        xmax = max(x1, x2)
+        ymax = max(y1, y2)
+        
         color = color_map.get(cls, "red")
-        draw.rectangle([bbox[0], bbox[1], bbox[2], bbox[3]], outline=color, width=3)
+        draw.rectangle([xmin, ymin, xmax, ymax], outline=color, width=3)
         
         if cls == "Table":
             final_output_text += f"\n\n--- [Table] ---\n{txt}\n-----------------\n"
@@ -184,6 +198,7 @@ def process_ocr_task(image):
 
     return final_output_text, result_image
 
+# --- Gradio Interface ---
 with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
     gr.Markdown("# **NVIDIA Nemotron Parse v1.1 [OCR/Parsing]**", elem_id="main-title")
     gr.Markdown("Upload a document image to extract text, tables, and layout structures using NVIDIA's state-of-the-art Parse model.")
@@ -200,9 +215,16 @@ with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
             )
 
         with gr.Column(scale=2):
-            output_text = gr.Textbox(label="Parsed Content (Markdown/LaTeX)", lines=8, show_copy_button=True)
+            output_text = gr.Textbox(label="Parsed Content (Markdown/LaTeX)", lines=20, show_copy_button=True)
             output_image = gr.Image(label="Detected Layout & Bounding Boxes", type="pil")
             
+            with gr.Accordion("Technical Details", open=False):
+                gr.Markdown("""
+                **Model:** [nvidia/NVIDIA-Nemotron-Parse-v1.1](https://huggingface.co/nvidia/NVIDIA-Nemotron-Parse-v1.1)
+                **Architecture:** Llama-3-Vila based.
+                **Capabilities:** High-accuracy OCR, Table extraction (to LaTeX/HTML), Figure detection.
+                """)
+
     submit_btn.click(
         fn=process_ocr_task, 
         inputs=[image_input],