modal_app.py

"""
Modal deployment for vision-language models.

Two models:
- InternVL3-8B: Video analysis, visual reasoning, 3D understanding
- DeepSeek-VL2-Tiny: Document OCR, text extraction (834 OCRBench vs GPT-4o's 736)

Deploy:
    modal deploy modal_app.py

Run locally:
    modal run modal_app.py --image-path image.png --prompt "Describe this"
    modal run modal_app.py --document invoice.pdf --extract

Endpoints:
    POST /analyze - Image analysis with InternVL3
    POST /extract - Document OCR with DeepSeek-VL2
"""

import modal

app = modal.App("vision-reasoning")

# Container image with all dependencies
image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("libgl1-mesa-glx", "libglib2.0-0", "git", "libheif-dev")
    .pip_install(
        "numpy<2",  # Pin to 1.x for torchvision compatibility
        "torch==2.4.0",
        "torchvision",
        "transformers>=4.45.0",
        "accelerate>=0.25.0",
        "pillow>=10.0.0",
        "pillow-heif",  # HEIC/HEIF support
        "opencv-python-headless>=4.8.0",
        "decord>=0.6.0",
        "pymupdf>=1.23.0",
        "einops",
        "timm",
        "fastapi[standard]",  # Required for web endpoints
        "attrdict",  # Required by DeepSeek-VL2
    )
    # Install DeepSeek-VL2 package from GitHub
    .pip_install("git+https://github.com/deepseek-ai/DeepSeek-VL2.git")
    # Note: Skipping flash-attn (requires CUDA compilation).
    # Model will use PyTorch SDPA instead - slightly slower but works fine.
)


def download_internvl3():
    """Download InternVL3 model weights at image build time."""
    from transformers import AutoModel, AutoTokenizer
    import torch

    model_id = "OpenGVLab/InternVL3-8B"
    print(f"Downloading {model_id}...")

    AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    AutoModel.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
    )
    print("Download complete")


def download_deepseek_vl2():
    """Download DeepSeek-VL2 model weights at image build time."""
    from deepseek_vl2.models import DeepseekVLV2ForCausalLM, DeepseekVLV2Processor
    import torch

    model_id = "deepseek-ai/deepseek-vl2-small"
    print(f"Downloading {model_id}...")

    DeepseekVLV2Processor.from_pretrained(model_id)
    DeepseekVLV2ForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=True,
    )
    print("Download complete")


# Bake model weights into image (avoids download on each cold start)
image_with_internvl = image.run_function(
    download_internvl3,
    gpu="A10G",
    timeout=3600,
)

image_with_deepseek = image.run_function(
    download_deepseek_vl2,
    gpu="A100",  # Small model needs A100 (80GB)
    timeout=3600,
)


@app.cls(
    image=image_with_internvl,
    gpu="A100",
    timeout=600,
    scaledown_window=300,
)
class InternVL3Model:
    """InternVL3-8B for image and video analysis."""

    @modal.enter()
    def load_model(self):
        """Load model when container starts."""
        from transformers import AutoModel, AutoTokenizer
        import torch

        model_id = "OpenGVLab/InternVL3-8B"
        print(f"Loading {model_id}...")

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_id,
            trust_remote_code=True,
        )

        self.model = AutoModel.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        ).eval()

        print("Model loaded")

    def _preprocess_image(self, image, max_num=12):
        """Preprocess image for InternVL3 using dynamic resolution."""
        import torchvision.transforms as T
        from torchvision.transforms.functional import InterpolationMode
        import math
        import torch

        IMAGENET_MEAN = (0.485, 0.456, 0.406)
        IMAGENET_STD = (0.229, 0.224, 0.225)

        def build_transform(input_size):
            return T.Compose([
                T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
                T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
                T.ToTensor(),
                T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
            ])

        def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
            best_ratio_diff = float('inf')
            best_ratio = (1, 1)
            area = width * height
            for ratio in target_ratios:
                target_aspect_ratio = ratio[0] / ratio[1]
                ratio_diff = abs(aspect_ratio - target_aspect_ratio)
                if ratio_diff < best_ratio_diff:
                    best_ratio_diff = ratio_diff
                    best_ratio = ratio
                elif ratio_diff == best_ratio_diff:
                    if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
                        best_ratio = ratio
            return best_ratio

        def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
            orig_width, orig_height = image.size
            aspect_ratio = orig_width / orig_height

            target_ratios = set(
                (i, j) for n in range(min_num, max_num + 1)
                for i in range(1, n + 1) for j in range(1, n + 1)
                if i * j <= max_num and i * j >= min_num
            )
            target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])

            target_aspect_ratio = find_closest_aspect_ratio(
                aspect_ratio, target_ratios, orig_width, orig_height, image_size
            )

            target_width = image_size * target_aspect_ratio[0]
            target_height = image_size * target_aspect_ratio[1]
            blocks = target_aspect_ratio[0] * target_aspect_ratio[1]

            resized_img = image.resize((target_width, target_height))
            processed_images = []
            for i in range(blocks):
                box = (
                    (i % (target_width // image_size)) * image_size,
                    (i // (target_width // image_size)) * image_size,
                    ((i % (target_width // image_size)) + 1) * image_size,
                    ((i // (target_width // image_size)) + 1) * image_size
                )
                split_img = resized_img.crop(box)
                processed_images.append(split_img)

            if use_thumbnail and len(processed_images) != 1:
                thumbnail_img = image.resize((image_size, image_size))
                processed_images.append(thumbnail_img)

            return processed_images

        transform = build_transform(input_size=448)
        images = dynamic_preprocess(image, image_size=448, use_thumbnail=True, max_num=max_num)
        pixel_values = [transform(img) for img in images]
        pixel_values = torch.stack(pixel_values)
        return pixel_values

    @modal.method()
    def chat(self, image_bytes: bytes, prompt: str, max_new_tokens: int = 1024) -> str:
        """Analyze an image with a prompt."""
        from PIL import Image
        import pillow_heif
        import torch
        import io

        # Register HEIF/HEIC support with PIL
        pillow_heif.register_heif_opener()

        img = Image.open(io.BytesIO(image_bytes)).convert("RGB")

        # Preprocess image
        pixel_values = self._preprocess_image(img, max_num=12)
        pixel_values = pixel_values.to(torch.bfloat16).cuda()

        # Build prompt with image placeholder
        num_patches = pixel_values.shape[0]
        image_tokens = '<image>' * num_patches
        question = f'{image_tokens}\n{prompt}'

        # Generate
        generation_config = dict(max_new_tokens=max_new_tokens, do_sample=False)

        response = self.model.chat(
            self.tokenizer,
            pixel_values,
            question,
            generation_config,
        )

        return response

    @modal.method()
    def analyze_frames(self, frame_bytes_list: list[bytes], prompt: str) -> str:
        """Analyze multiple video frames together (A100 has enough memory)."""
        from PIL import Image
        import pillow_heif
        import torch
        import io

        pillow_heif.register_heif_opener()

        # Process each frame
        all_pixel_values = []
        num_patches_list = []

        for frame_bytes in frame_bytes_list:
            img = Image.open(io.BytesIO(frame_bytes)).convert("RGB")
            pixel_values = self._preprocess_image(img, max_num=6)
            all_pixel_values.append(pixel_values)
            num_patches_list.append(pixel_values.shape[0])

        # Combine all frames
        combined_pixels = torch.cat(all_pixel_values, dim=0).to(torch.bfloat16).cuda()

        # Build prompt with frame labels
        prompt_parts = []
        for i, num_patches in enumerate(num_patches_list):
            frame_tokens = '<image>' * num_patches
            prompt_parts.append(f'Frame {i+1}: {frame_tokens}')
        prompt_parts.append(f'\n{prompt}')
        question = '\n'.join(prompt_parts)

        generation_config = dict(max_new_tokens=1024, do_sample=False)

        response = self.model.chat(
            self.tokenizer,
            combined_pixels,
            question,
            generation_config,
        )

        return response


@app.function(image=image_with_internvl, gpu="A100", timeout=300)
@modal.fastapi_endpoint(method="POST")
def analyze(image_url: str, prompt: str = "Describe this image in detail."):
    """
    HTTP endpoint for image analysis.

    POST /analyze
    {
        "image_url": "https://example.com/image.png",
        "prompt": "What's in this image?"
    }
    """
    import urllib.request

    with urllib.request.urlopen(image_url) as response:
        image_bytes = response.read()

    model = InternVL3Model()
    result = model.chat.remote(image_bytes, prompt)

    return {"result": result}


@app.function(image=image_with_internvl, gpu="A100", timeout=1800)
def analyze_video(
    video_bytes: bytes,
    prompt: str = "Describe what happens in this video.",
    sample_interval: float = 2.0,
    max_frames: int = 12,
) -> dict:
    """
    Analyze a video by sampling frames.

    Args:
        video_bytes: Raw video file bytes
        prompt: Question about the video
        sample_interval: Seconds between sampled frames
        max_frames: Maximum frames to analyze (keep low for memory)

    Returns:
        Analysis result
    """
    import tempfile
    import cv2
    from PIL import Image
    import io

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
        f.write(video_bytes)
        video_path = f.name

    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_interval = max(1, int(fps * sample_interval))

    frames = []
    frame_count = 0

    while len(frames) < max_frames:
        ret, frame = cap.read()
        if not ret:
            break

        if frame_count % frame_interval == 0:
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(frame_rgb)

            buf = io.BytesIO()
            img.save(buf, format="JPEG", quality=85)
            frames.append(buf.getvalue())

        frame_count += 1

    cap.release()

    if not frames:
        return {"error": "Could not extract frames from video"}

    model = InternVL3Model()
    result = model.analyze_frames.remote(frames, prompt)

    return {
        "frames_analyzed": len(frames),
        "sample_interval": sample_interval,
        "result": result,
    }


# =============================================================================
# DeepSeek-VL2 for Document OCR
# =============================================================================

@app.cls(
    image=image_with_deepseek,
    gpu="A100",  # Small model needs A100 (80GB)
    timeout=600,
    scaledown_window=300,
)
class DeepSeekVL2Model:
    """DeepSeek-VL2-Small for document OCR and text extraction."""

    @modal.enter()
    def load_model(self):
        """Load model when container starts."""
        from deepseek_vl2.models import DeepseekVLV2ForCausalLM, DeepseekVLV2Processor
        import torch

        model_id = "deepseek-ai/deepseek-vl2-small"
        print(f"Loading {model_id}...")

        self.processor = DeepseekVLV2Processor.from_pretrained(model_id)

        self.model = DeepseekVLV2ForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            low_cpu_mem_usage=True,
        ).eval()

        print("Model loaded")

    def _generate(self, img, prompt: str) -> str:
        """Internal helper to run generation with DeepSeek-VL2."""
        import torch

        conversation = [
            {"role": "<|User|>", "content": f"<image>\n{prompt}", "images": [img]},
            {"role": "<|Assistant|>", "content": ""},
        ]

        inputs = self.processor(
            conversations=conversation,
            images=[img],
            force_batchify=True,
        ).to("cuda", dtype=torch.bfloat16)

        with torch.no_grad():
            input_embeds = self.model.prepare_inputs_embeds(**inputs)
            outputs = self.model.language.generate(
                inputs_embeds=input_embeds,
                attention_mask=inputs.attention_mask,
                pad_token_id=self.processor.tokenizer.eos_token_id,
                bos_token_id=self.processor.tokenizer.bos_token_id,
                eos_token_id=self.processor.tokenizer.eos_token_id,
                max_new_tokens=2048,
                do_sample=False,
                use_cache=True,
                repetition_penalty=1.1,  # Prevent repetition loops
            )

        return self.processor.tokenizer.decode(
            outputs[0].cpu().tolist(),
            skip_special_tokens=True,
        ).strip()

    def _translate_text(self, text: str, target_language: str) -> str:
        """Translate extracted text by creating a dummy image and using vision-language."""
        import torch
        from PIL import Image

        # DeepSeek-VL2 requires an image, so we create a minimal white image
        # and put the text in the prompt for translation
        dummy_img = Image.new("RGB", (64, 64), color="white")

        prompt = f"""Translate the following text to {target_language}.
Translate everything: headers, labels, technical terms, descriptions.
Output ONLY the {target_language} translation, nothing else.

Text:
{text}"""

        conversation = [
            {"role": "<|User|>", "content": f"<image>\n{prompt}", "images": [dummy_img]},
            {"role": "<|Assistant|>", "content": ""},
        ]

        inputs = self.processor(
            conversations=conversation,
            images=[dummy_img],
            force_batchify=True,
        ).to("cuda", dtype=torch.bfloat16)

        with torch.no_grad():
            input_embeds = self.model.prepare_inputs_embeds(**inputs)
            outputs = self.model.language.generate(
                inputs_embeds=input_embeds,
                attention_mask=inputs.attention_mask,
                pad_token_id=self.processor.tokenizer.eos_token_id,
                bos_token_id=self.processor.tokenizer.bos_token_id,
                eos_token_id=self.processor.tokenizer.eos_token_id,
                max_new_tokens=2048,
                do_sample=False,
                use_cache=True,
                repetition_penalty=1.1,  # Prevent repetition loops
            )

        response = self.processor.tokenizer.decode(
            outputs[0].cpu().tolist(),
            skip_special_tokens=True,
        )

        return response.strip()

    @modal.method()
    def extract_text(self, image_bytes: bytes, translate_to: str = None) -> str:
        """Extract all text from a document image, optionally translating."""
        from PIL import Image
        import pillow_heif
        import io

        pillow_heif.register_heif_opener()

        img = Image.open(io.BytesIO(image_bytes)).convert("RGB")

        # Step 1: Extract text (OCR)
        extracted = self._generate(
            img,
            "Extract all text from this document image. Preserve the layout and formatting as much as possible."
        )

        # Step 2: Translate if requested (text-only pass)
        if translate_to:
            return self._translate_text(extracted, translate_to)

        return extracted

    @modal.method()
    def extract_structured(
        self,
        image_bytes: bytes,
        output_format: str = "json",
        translate_to: str = None,
    ) -> str:
        """Extract structured data from a document, optionally translated."""
        from PIL import Image
        import pillow_heif
        import torch
        import io

        pillow_heif.register_heif_opener()

        img = Image.open(io.BytesIO(image_bytes)).convert("RGB")

        # Add translation instruction if requested
        lang_instruction = ""
        if translate_to:
            lang_instruction = f"\nIMPORTANT: All text values in the JSON must be translated to {translate_to}. Translate field names/keys to {translate_to} as well."

        if output_format == "invoice":
            prompt = f"""Extract invoice data from this image.
Return JSON with:
{{
  "vendor_name": "",
  "invoice_number": "",
  "invoice_date": "",
  "due_date": "",
  "subtotal": 0,
  "tax": 0,
  "total": 0,
  "currency": "",
  "line_items": [
    {{"description": "", "quantity": 0, "unit_price": 0, "amount": 0}}
  ]
}}
Return ONLY valid JSON. Use null for missing fields.{lang_instruction}"""
        elif output_format == "receipt":
            prompt = f"""Extract receipt data from this image.
Return JSON with:
{{
  "merchant_name": "",
  "date": "",
  "items": [{{"name": "", "price": 0}}],
  "subtotal": 0,
  "tax": 0,
  "total": 0,
  "payment_method": ""
}}
Return ONLY valid JSON. Use null for missing fields.{lang_instruction}"""
        else:
            prompt = f"""Extract the content from this document image and return it as structured JSON.
Identify and extract:
- Any headers or titles
- Key-value pairs (like form fields)
- Tables (as arrays of objects)
- Lists (as arrays)
- Dates, numbers, and currency values
Return ONLY valid JSON, no explanation.{lang_instruction}"""

        # DeepSeek-VL2 conversation format
        conversation = [
            {"role": "<|User|>", "content": f"<image>\n{prompt}", "images": [img]},
            {"role": "<|Assistant|>", "content": ""},
        ]

        # Prepare inputs
        inputs = self.processor(
            conversations=conversation,
            images=[img],
            force_batchify=True,
        ).to("cuda", dtype=torch.bfloat16)

        # Generate using DeepSeek-VL2's API
        with torch.no_grad():
            input_embeds = self.model.prepare_inputs_embeds(**inputs)
            outputs = self.model.language.generate(
                inputs_embeds=input_embeds,
                attention_mask=inputs.attention_mask,
                pad_token_id=self.processor.tokenizer.eos_token_id,
                bos_token_id=self.processor.tokenizer.bos_token_id,
                eos_token_id=self.processor.tokenizer.eos_token_id,
                max_new_tokens=2048,
                do_sample=False,
                use_cache=True,
                repetition_penalty=1.1,
            )

        response = self.processor.tokenizer.decode(
            outputs[0].cpu().tolist(),
            skip_special_tokens=True,
        )

        return response.strip()


@app.function(image=image_with_deepseek, gpu="A10G", timeout=300)
@modal.fastapi_endpoint(method="POST")
def extract(
    image_url: str,
    output_format: str = "json",
    translate_to: str = None,
):
    """
    HTTP endpoint for document extraction with DeepSeek-VL2.

    POST /extract
    {
        "image_url": "https://example.com/invoice.png",
        "output_format": "json" | "invoice" | "receipt" | "text",
        "translate_to": "English" | "Spanish" | "French" | etc.
    }
    """
    import urllib.request

    with urllib.request.urlopen(image_url) as response:
        image_bytes = response.read()

    model = DeepSeekVL2Model()

    if output_format == "text":
        result = model.extract_text.remote(image_bytes, translate_to)
    else:
        result = model.extract_structured.remote(image_bytes, output_format)

    return {"result": result, "model": "deepseek-vl2-tiny", "format": output_format}


@app.function(image=image_with_deepseek, gpu="A10G", timeout=600)
def extract_document(
    document_bytes: bytes,
    output_format: str = "json",
    translate_to: str = None,
) -> dict:
    """
    Extract text/data from a document (image or PDF).

    Args:
        document_bytes: Raw document bytes
        output_format: "text", "json", "invoice", or "receipt"
        translate_to: Target language for translation (e.g., "English", "Spanish")

    Returns:
        Extraction result
    """
    from PIL import Image
    import pillow_heif
    import fitz  # PyMuPDF
    import io

    pillow_heif.register_heif_opener()

    # Check if it's a PDF
    is_pdf = document_bytes[:4] == b'%PDF'

    if is_pdf:
        # Convert PDF pages to images
        doc = fitz.open(stream=document_bytes, filetype="pdf")
        results = []

        for page_num in range(min(len(doc), 10)):  # Limit to 10 pages
            page = doc[page_num]
            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))  # 2x zoom
            img_bytes = pix.tobytes("png")

            model = DeepSeekVL2Model()
            if output_format == "text":
                text = model.extract_text.remote(img_bytes, translate_to)
            else:
                text = model.extract_structured.remote(img_bytes, output_format, translate_to)

            results.append({"page": page_num + 1, "content": text})

        doc.close()
        return {"pages": results, "total_pages": len(results)}
    else:
        # Single image
        model = DeepSeekVL2Model()
        if output_format == "text":
            result = model.extract_text.remote(document_bytes, translate_to)
        else:
            result = model.extract_structured.remote(document_bytes, output_format, translate_to)

        return {"content": result}


def create_pdf_from_text(text_pages: list, output_path: str):
    """Create a PDF from extracted/translated text pages."""
    import fitz  # PyMuPDF

    doc = fitz.open()

    for page_text in text_pages:
        # Create a new page (Letter size: 612x792 points)
        page = doc.new_page(width=612, height=792)

        # Set margins
        margin = 50
        text_width = 612 - (2 * margin)

        # Insert text with automatic wrapping
        text_rect = fitz.Rect(margin, margin, 612 - margin, 792 - margin)

        # Use a standard font
        page.insert_textbox(
            text_rect,
            page_text,
            fontsize=10,
            fontname="helv",  # Helvetica
            align=0,  # Left align
        )

    doc.save(output_path)
    doc.close()
    return output_path


@app.local_entrypoint()
def main(
    image_path: str = None,
    video_path: str = None,
    document: str = None,
    prompt: str = "Describe what you see in detail.",
    extract: bool = False,
    output_format: str = "json",
    translate_to: str = None,
    output_pdf: str = None,
):
    """
    CLI entry point.

    Examples:
        # Image analysis with InternVL3
        modal run modal_app.py --image-path photo.jpg
        modal run modal_app.py --image-path photo.jpg --prompt "What objects are here?"

        # Video analysis with InternVL3
        modal run modal_app.py --video-path video.mp4 --prompt "What happens?"

        # Document OCR with DeepSeek-VL2
        modal run modal_app.py --document invoice.pdf --extract
        modal run modal_app.py --document receipt.jpg --extract --output-format receipt

        # Document OCR with translation
        modal run modal_app.py --document german_doc.pdf --extract --output-format text --translate-to English

        # Output to PDF
        modal run modal_app.py --document german_doc.pdf --extract --translate-to English --output-pdf translated.pdf
    """
    if document or extract:
        # Use DeepSeek-VL2 for document extraction
        doc_path = document or image_path
        if not doc_path:
            print("Provide --document or --image-path with --extract")
            return

        with open(doc_path, "rb") as f:
            doc_bytes = f.read()

        result = extract_document.remote(doc_bytes, output_format, translate_to)

        # Collect text for PDF output
        text_pages = []

        if "pages" in result:
            print(f"Extracted {result['total_pages']} pages:")
            for page in result["pages"]:
                print(f"\n--- Page {page['page']} ---")
                print(page["content"])
                text_pages.append(page["content"])
        else:
            print(result["content"])
            text_pages.append(result["content"])

        # Save to PDF if requested
        if output_pdf:
            create_pdf_from_text(text_pages, output_pdf)
            print(f"\nSaved to: {output_pdf}")

    elif image_path:
        # Use InternVL3 for image analysis
        with open(image_path, "rb") as f:
            image_bytes = f.read()

        model = InternVL3Model()
        result = model.chat.remote(image_bytes, prompt)
        print(result)

    elif video_path:
        # Use InternVL3 for video analysis
        with open(video_path, "rb") as f:
            video_bytes = f.read()

        result = analyze_video.remote(video_bytes, prompt)
        print(f"Analyzed {result['frames_analyzed']} frames")
        print(result["result"])

    else:
        print("Usage:")
        print("  modal run modal_app.py --image-path photo.jpg")
        print("  modal run modal_app.py --video-path video.mp4")
        print("  modal run modal_app.py --document invoice.pdf --extract")