Update app.py

app.py

@@ -26,6 +26,9 @@ from transformers import (
 
 # --- Constants and Model Setup ---
 MAX_INPUT_TOKEN_LENGTH = 4096
+# Note: The following line correctly falls back to CPU if CUDA is not available.
+# The "RuntimeError: CUDA driver initialization failed" is an environment issue,
+# meaning the code is being run where a GPU is expected but not found/configured.
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # --- Prompts for Different Tasks ---
@@ -78,57 +81,64 @@ model_i = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     MODEL_ID_I, trust_remote_code=True, torch_dtype=torch.float16
 ).to(device).eval()
 
-def progress_bar_html(label: str) -> str:
+# --- Utility Functions ---
+def layoutjson2md(layout_data: Any) -> str:
     """
-    Returns an HTML snippet for a thin progress bar with a label.
-    The progress bar is styled as a dark red animated bar.
+    FIXED: Converts the structured JSON from Layout Analysis into formatted Markdown.
+    This version is robust against malformed JSON from the model.
     """
-    return f'''
-<div style="display: flex; align-items: center;">
-    <span style="margin-right: 10px; font-size: 14px;">{label}</span>
-    <div style="width: 110px; height: 5px; background-color: #AFEEEE; border-radius: 2px; overflow: hidden;">
-        <div style="width: 100%; height: 100%; background-color: #00FFFF; animation: loading 1.5s linear infinite;"></div>
-    </div>
-</div>
-<style>
-@keyframes loading {{
-    0% {{ transform: translateX(-100%); }}
-    100% {{ transform: translateX(100%); }}
-}}
-</style>
-    '''
-
-# --- Utility Functions ---
-def layoutjson2md(layout_data: List[Dict]) -> str:
-    """Converts the structured JSON from Layout Analysis into formatted Markdown."""
     markdown_lines = []
+
+    # If the model wraps the list in a dictionary, find and extract the list.
+    if isinstance(layout_data, dict):
+        found_list = None
+        for value in layout_data.values():
+            if isinstance(value, list):
+                found_list = value
+                break
+        if found_list is not None:
+            layout_data = found_list
+        else:
+            return "### Error: Could not find a list of layout items in the JSON object."
+
+    if not isinstance(layout_data, list):
+        return f"### Error: Expected a list of layout items, but received type {type(layout_data).__name__}."
+
     try:
-        # Sort items by reading order (top-to-bottom, left-to-right)
-        sorted_items = sorted(layout_data, key=lambda x: (x.get('bbox', [0,0,0,0])[1], x.get('bbox', [0,0,0,0])[0]))
+        # Filter out any non-dictionary items and sort by reading order.
+        valid_items = [item for item in layout_data if isinstance(item, dict)]
+        sorted_items = sorted(valid_items, key=lambda x: (x.get('bbox', [0, 0, 0, 0])[1], x.get('bbox', [0, 0, 0, 0])[0]))
+
         for item in sorted_items:
-            category = item.get('category', '')
+            category = item.get('category', 'Text') # Default to 'Text' if no category
             text = item.get('text', '')
-            if not text: continue
+            if not text:
+                continue
 
-            if category == 'Title': markdown_lines.append(f"# {text}\n")
-            elif category == 'Section-header': markdown_lines.append(f"## {text}\n")
+            if category == 'Title':
+                markdown_lines.append(f"# {text}\n")
+            elif category == 'Section-header':
+                markdown_lines.append(f"## {text}\n")
             elif category == 'Table':
-                 # Handle structured table JSON
                 if isinstance(text, dict) and 'header' in text and 'rows' in text:
                     header = '| ' + ' | '.join(map(str, text['header'])) + ' |'
                     separator = '| ' + ' | '.join(['---'] * len(text['header'])) + ' |'
                     rows = ['| ' + ' | '.join(map(str, row)) + ' |' for row in text['rows']]
                     markdown_lines.extend([header, separator] + rows)
                     markdown_lines.append("\n")
-                else: # Fallback for simple text
+                else:  # Fallback for simple text or malformed tables
                     markdown_lines.append(f"{text}\n")
             else:
-                 markdown_lines.append(f"{text}\n")
+                markdown_lines.append(f"{text}\n")
+
     except Exception as e:
         print(f"Error converting to markdown: {e}")
-        return "### Error converting JSON to Markdown."
+        traceback.print_exc()
+        return "### Error: An unexpected error occurred while converting JSON to Markdown."
+
     return "\n".join(markdown_lines)
 
+
 # --- Core Application Logic ---
 @spaces.GPU
 def process_document_stream(model_name: str, task_choice: str, image: Image.Image, max_new_tokens: int):
@@ -158,11 +168,10 @@ def process_document_stream(model_name: str, task_choice: str, image: Image.Imag
     inputs = processor(text=[prompt_full], images=[image], return_tensors="pt", padding=True, truncation=True, max_length=MAX_INPUT_TOKEN_LENGTH).to(device)
     streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
     generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
-    
+
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
 
-    
     # 4. Stream raw output to the UI in real-time
     buffer = ""
     for new_text in streamer:
@@ -179,24 +188,32 @@ def process_document_stream(model_name: str, task_choice: str, image: Image.Imag
         try:
             json_match = re.search(r'```json\s*([\s\S]+?)\s*```', buffer)
             if not json_match:
-                raise json.JSONDecodeError("JSON object not found in output.", buffer, 0)
-            
+                # If no JSON block is found, try to parse the whole buffer as a fallback.
+                try:
+                    layout_data = json.loads(buffer)
+                    markdown_content = layoutjson2md(layout_data)
+                    yield buffer, markdown_content, layout_data
+                    return
+                except json.JSONDecodeError:
+                    raise ValueError("JSON object not found in the model's output.")
+
             json_str = json_match.group(1)
             layout_data = json.loads(json_str)
             markdown_content = layoutjson2md(layout_data)
-            
+
             yield buffer, markdown_content, layout_data
         except Exception as e:
-            error_md = f"❌ **Error:** Failed to parse Layout JSON.\n\n**Details:**\n`{str(e)}`"
+            error_md = f"❌ **Error:** Failed to parse Layout JSON.\n\n**Details:**\n`{str(e)}`\n\n**Raw Output:**\n```\n{buffer}\n```"
             error_json = {"error": "ProcessingError", "details": str(e), "raw_output": buffer}
             yield buffer, error_md, error_json
 
+
 # --- Gradio UI Definition ---
 def create_gradio_interface():
     """Builds and returns the Gradio web interface."""
     css = """
     .main-container { max-width: 1400px; margin: 0 auto; }
-    .process-button { border: none !important; color: white !important; font-weight: bold !important; background-color: blue !important;} 
+    .process-button { border: none !important; color: white !important; font-weight: bold !important; background-color: blue !important;}
     .process-button:hover { background-color: darkblue !important; transform: translateY(-2px) !important; box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important; }
     """
     with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
@@ -208,15 +225,15 @@ def create_gradio_interface():
             </p>
         </div>
         """)
-        
+
         with gr.Row():
             # Left Column (Inputs)
             with gr.Column(scale=1):
                 model_choice = gr.Dropdown(
-                    choices=["Camel-Doc-OCR-080125", 
-                             "MonkeyOCR-Recognition", 
+                    choices=["Camel-Doc-OCR-080125",
+                             "MonkeyOCR-Recognition",
                              "olmOCR-7B-0725",
-                             "Nanonets-OCR-s", 
+                             "Nanonets-OCR-s",
                              "Megalodon-OCR-Sync-0713"
                             ],
                     label="Select Model", value="Nanonets-OCR-s"
@@ -228,7 +245,7 @@ def create_gradio_interface():
                 image_input = gr.Image(label="Upload Image", type="pil", sources=['upload'])
                 with gr.Accordion("Advanced Settings", open=False):
                     max_new_tokens = gr.Slider(minimum=512, maximum=8192, value=4096, step=256, label="Max New Tokens")
-                
+
                 process_btn = gr.Button("🚀 Process Document", variant="primary", elem_classes=["process-button"], size="lg")
                 clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
 
@@ -242,33 +259,32 @@ def create_gradio_interface():
                                 examples=["examples/1.png", "examples/2.png", "examples/3.png", "examples/4.png", "examples/5.png"],
                                 inputs=image_input,
                                 label="Examples"
-                    )                    
-                    with gr.Tab("📰 README.md"):  
-                        with gr.Accordion("(Formatted Result)", open=True):
-                            markdown_output = gr.Markdown(label="Formatted Markdown")
-                    
+                    )
+                    with gr.Tab("📰 Formatted Result"):
+                        markdown_output = gr.Markdown(label="Formatted Markdown")
+
                     with gr.Tab("📋 Layout Analysis Results"):
                         json_output = gr.JSON(label="Structured Layout Data (JSON)")
-        
+
         # Event Handlers
         def clear_all_outputs():
             return None, "Raw output will appear here.", "Formatted results will appear here.", None
 
         process_btn.click(
             fn=process_document_stream,
-            inputs=[model_choice, 
-                    task_choice, 
-                    image_input, 
+            inputs=[model_choice,
+                    task_choice,
+                    image_input,
                     max_new_tokens],
-            outputs=[raw_output_stream, 
-                     markdown_output, 
+            outputs=[raw_output_stream,
+                     markdown_output,
                      json_output]
         )
         clear_btn.click(
             clear_all_outputs,
-            outputs=[image_input, 
-                     raw_output_stream, 
-                     markdown_output, 
+            outputs=[image_input,
+                     raw_output_stream,
+                     markdown_output,
                      json_output]
         )
     return demo