app.py

try:
    import spaces
except ImportError:
    class spaces:
        @staticmethod
        def GPU(fn=None, duration=None):
            if fn is None:
                def decorator(f):
                    return f
                return decorator
            return fn

import os
import shutil
import tempfile
import zipfile
import time
import uuid
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
import gradio as gr

def has_selectable_text(pdf_path):
    """Checks if the PDF has selectable text to skip OCR for digital documents."""
    try:
        import pypdf
        reader = pypdf.PdfReader(pdf_path)
        text_sample = ""
        num_pages = len(reader.pages)
        for i in range(min(10, num_pages)):
            page_text = reader.pages[i].extract_text()
            if page_text:
                text_sample += page_text
                if len(text_sample.strip()) > 50:
                    return True
        return len(text_sample.strip()) > 50
    except Exception as e:
        print(f"Error checking text selectability: {e}")
        return True

def convert_single_doc(file_path, filename, digital_converter, scanned_converter, output_dir):
    """Processes a single document. Executed inside the GPU node's local thread pool."""
    start_time = time.time()
    try:
        print(f"Starting conversion of {filename} in parallel on GPU...")
        is_digital = has_selectable_text(file_path)
        converter = digital_converter if is_digital else scanned_converter
        
        result = converter.convert(file_path)
        markdown_content = result.document.export_to_markdown()
        
        # Save output markdown file
        base_name = os.path.splitext(filename)[0]
        output_path = os.path.join(output_dir, f"{base_name}.md")
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(markdown_content)
        
        elapsed = time.time() - start_time
        try:
            import pypdf
            reader = pypdf.PdfReader(file_path)
            pages = len(reader.pages)
        except Exception:
            pages = "N/A"
            
        return {
            "name": filename,
            "pages": pages,
            "status": "Success",
            "time": f"{elapsed:.1f}s",
            "content": markdown_content
        }
    except Exception as e:
        elapsed = time.time() - start_time
        print(f"Error converting {filename}: {e}")
        return {
            "name": filename,
            "pages": "N/A",
            "status": f"Error: {str(e)}",
            "time": f"{elapsed:.1f}s",
            "content": ""
        }

# Define a batch conversion function decorated with ZeroGPU.
# This processes a list of files on a SINGLE GPU instance concurrently in VRAM
# using a local thread pool, maximizing GPU core occupancy and minimizing GPU-second duration.
@spaces.GPU(duration=480)
def convert_file_batch_on_single_gpu(batch_info: list, output_dir: str) -> list:
    # Local imports to prevent global PyTorch initialization on CPU container
    from docling.datamodel.base_models import InputFormat
    from docling.datamodel.pipeline_options import PdfPipelineOptions
    from docling.document_converter import DocumentConverter, PdfFormatOption
    from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend
    from docling.datamodel.accelerator_options import AcceleratorOptions, AcceleratorDevice

    # 1. Configure Docling pipeline options
    def get_pipeline_options(do_ocr):
        opts = PdfPipelineOptions()
        opts.do_ocr = do_ocr
        opts.do_picture_classification = False
        opts.do_picture_description = False
        opts.do_chart_extraction = False
        opts.do_code_enrichment = False
        opts.do_formula_enrichment = False
        opts.generate_page_images = False
        opts.generate_picture_images = False
        opts.accelerator_options = AcceleratorOptions(
            device=AcceleratorDevice.CUDA,
            num_threads=2
        )
        return opts

    # Pre-instantiate converters to share VRAM weights across the concurrent threads
    # One for digital (no OCR) and one for scanned (OCR enabled)
    digital_converter = DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(
                pipeline_options=get_pipeline_options(do_ocr=False),
                backend=PyPdfiumDocumentBackend
            )
        }
    )
    scanned_converter = DocumentConverter(
        format_options={
            InputFormat.PDF: PdfFormatOption(
                pipeline_options=get_pipeline_options(do_ocr=True),
                backend=PyPdfiumDocumentBackend
            )
        }
    )

    results = []
    
    # Process files concurrently inside the A10G/RTX6000 VRAM using a local thread pool
    # With 48GB VRAM (Blackwell), we can easily handle 32 concurrent document streams sharing weights
    local_concurrency = min(32, len(batch_info))
    print(f"Executing local thread pool inside GPU container with {local_concurrency} workers...")
    
    with ThreadPoolExecutor(max_workers=local_concurrency) as local_executor:
        futures = {
            local_executor.submit(
                convert_single_doc,
                file_path, filename,
                digital_converter, scanned_converter,
                output_dir
            ): filename
            for file_path, filename in batch_info
        }
        for future in as_completed(futures):
            results.append(future.result())
            
    return results

def run_batch_conversion(uploaded_files, max_workers):
    if not uploaded_files:
        return "No files uploaded.", gr.update(visible=False), gr.update(visible=False), []

    session_id = uuid.uuid4().hex
    temp_workspace = os.path.join(tempfile.gettempdir(), f"docling_workspace_{session_id}")
    input_dir = os.path.join(temp_workspace, "inputs")
    output_dir = os.path.join(temp_workspace, "outputs")
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)

    files_to_process = []

    for file_obj in uploaded_files:
        path = file_obj.name
        filename = os.path.basename(path)
        
        if filename.lower().endswith(".zip"):
            print(f"Extracting ZIP archive: {filename}")
            zip_extract_dir = os.path.join(input_dir, f"zip_{uuid.uuid4().hex}")
            os.makedirs(zip_extract_dir, exist_ok=True)
            try:
                with zipfile.ZipFile(path, 'r') as zip_ref:
                    zip_ref.extractall(zip_extract_dir)
                for root, _, files in os.walk(zip_extract_dir):
                    for f in files:
                        ext = os.path.splitext(f)[1].lower()
                        if ext in [".pdf", ".docx", ".pptx", ".html", ".png", ".jpg", ".jpeg"]:
                            full_path = os.path.join(root, f)
                            files_to_process.append((full_path, f))
            except Exception as e:
                print(f"Error extracting ZIP: {e}")
        else:
            dest_path = os.path.join(input_dir, filename)
            shutil.copy(path, dest_path)
            files_to_process.append((dest_path, filename))

    if not files_to_process:
        return "No supported document files found to convert.", gr.update(visible=False), gr.update(visible=False), []

    # Process ALL files in a single GPU call.
    # The GPU function uses ThreadPoolExecutor(max_workers=32) internally,
    # so 32 files run concurrently and as each finishes, the next one starts.
    # This maximizes GPU utilization with a single model load.
    results = []
    print(f"Processing {len(files_to_process)} files on single GPU node (32-way VRAM worker pool)...")
    
    batch_results = convert_file_batch_on_single_gpu(files_to_process, output_dir)
    results.extend(batch_results)

    zip_output_path = os.path.join(temp_workspace, "converted_markdown_files.zip")
    with zipfile.ZipFile(zip_output_path, 'w', zipfile.ZIP_DEFLATED) as zip_out:
        for root, _, files in os.walk(output_dir):
            for f in files:
                full_path = os.path.join(root, f)
                arcname = os.path.relpath(full_path, output_dir)
                zip_out.write(full_path, arcname)

    table_data = [
        [r["name"], r["pages"], r["status"], r["time"]]
        for r in results
    ]

    summary_text = f"Successfully processed {len([r for r in results if r['status'] == 'Success'])} out of {len(results)} files."
    
    preview_dict = {r["name"]: r["content"] for r in results if r["status"] == "Success"}
    dropdown_update = gr.update(choices=list(preview_dict.keys()), value=list(preview_dict.keys())[0] if preview_dict else None, visible=bool(preview_dict))
    preview_box_update = gr.update(visible=bool(preview_dict))

    return (
        summary_text,
        gr.update(value=zip_output_path, visible=True),
        table_data,
        dropdown_update,
        preview_box_update,
        preview_dict
    )

theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
).set(
    body_background_fill="*neutral_950",
    block_background_fill="*neutral_900",
    block_border_color="*neutral_800",
    button_primary_background_fill="linear-gradient(90deg, *primary_600, *secondary_600)",
    button_primary_text_color="white",
    block_title_text_color="*primary_400",
)

css = """
body {
    background-color: #0b0f19;
    font-family: 'Outfit', 'Inter', sans-serif;
}
.gradio-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
}
h1 {
    background: linear-gradient(90deg, #818cf8, #3b82f6);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    font-weight: 800;
}
.glass-panel {
    background: rgba(30, 41, 59, 0.4) !important;
    backdrop-filter: blur(12px);
    border: 1px solid rgba(255, 255, 255, 0.05);
    border-radius: 12px;
}
"""

with gr.Blocks(theme=theme, css=css, title="Docling GPU Batch Converter") as demo:
    preview_state = gr.State({})

    with gr.Column():
        gr.Markdown(
            """
            # 📄 Docling GPU Batch Converter
            ### IBM Docling parallel document-to-markdown parsing powered by Hugging Face ZeroGPU.
            """
        )
        
        with gr.Row():
            with gr.Column(scale=1, elem_classes=["glass-panel"]):
                gr.Markdown("### 📥 Upload Documents")
                file_input = gr.File(
                    file_count="multiple",
                    label="Upload PDFs, DOCX, PPTX or a ZIP archive containing them",
                    file_types=[".pdf", ".docx", ".pptx", ".html", ".png", ".jpg", ".zip"]
                )
                
                with gr.Row():
                    workers_slider = gr.Slider(
                        minimum=1,
                        maximum=4,
                        value=1,
                        step=1,
                        label="Parallel GPU Nodes"
                    )
                
                submit_btn = gr.Button("⚡ Convert Batch", variant="primary")
                download_output = gr.File(label="📦 Download Converted MD (ZIP)", visible=False)
                
            with gr.Column(scale=2, elem_classes=["glass-panel"]):
                gr.Markdown("### 📊 Conversion Progress & Status")
                status_summary = gr.Markdown("Ready to process.")
                progress_table = gr.Dataframe(
                    headers=["File Name", "Pages", "Status", "Time Taken"],
                    datatype=["str", "str", "str", "str"],
                    value=[]
                )
                
        with gr.Column(visible=False) as preview_container:
            gr.Markdown("### 🔍 Document Markdown Preview")
            preview_selector = gr.Dropdown(label="Select document to preview", choices=[], interactive=True)
            preview_markdown = gr.Markdown(label="Converted Output")

    def update_preview_text(selected_file, data_dict):
        if selected_file in data_dict:
            return data_dict[selected_file]
        return "No content available."

    submit_btn.click(
        fn=run_batch_conversion,
        inputs=[file_input, workers_slider],
        outputs=[status_summary, download_output, progress_table, preview_selector, preview_container, preview_state]
    )

    preview_selector.change(
        fn=update_preview_text,
        inputs=[preview_selector, preview_state],
        outputs=preview_markdown
    )

if __name__ == "__main__":
    demo.launch()