app.py

import os
import sys
import math
import numpy as np
import tempfile
import torch
import torchvision.transforms as T
from torchvision.transforms.functional import InterpolationMode
from PIL import Image
import gradio as gr
from transformers import AutoModel, AutoTokenizer
import io
import pdf2image
from pptx import Presentation

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

# Configuration
MODEL_NAME = "OpenGVLab/InternVL2_5-8B"
IMAGE_SIZE = 448

# Set up environment variables
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"

# Utility functions for image processing
def build_transform(input_size):
    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
    transform = 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=MEAN, std=STD)
    ])
    return transform

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

    # calculate the existing image aspect ratio
    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])

    # find the closest aspect ratio to the target
    target_aspect_ratio = find_closest_aspect_ratio(
        aspect_ratio, target_ratios, orig_width, orig_height, image_size)

    # calculate the target width and height
    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]

    # resize the image
    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 the image
        split_img = resized_img.crop(box)
        processed_images.append(split_img)
    assert len(processed_images) == blocks
    if use_thumbnail and len(processed_images) != 1:
        thumbnail_img = image.resize((image_size, image_size))
        processed_images.append(thumbnail_img)
    return processed_images

# Load and preprocess image for the model - following the official documentation pattern
def load_image(image_pil, max_num=12):
    # Process the image using dynamic_preprocess
    processed_images = dynamic_preprocess(image_pil, image_size=IMAGE_SIZE, max_num=max_num)
    
    # Convert PIL images to tensor format expected by the model
    transform = build_transform(IMAGE_SIZE)
    pixel_values = [transform(img) for img in processed_images]
    pixel_values = torch.stack(pixel_values)
    
    # Convert to appropriate data type
    if torch.cuda.is_available():
        pixel_values = pixel_values.cuda().to(torch.bfloat16)
    else:
        pixel_values = pixel_values.to(torch.float32)
        
    return pixel_values

# Function to split model across GPUs
def split_model(model_name):
    device_map = {}
    world_size = torch.cuda.device_count()
    if world_size <= 1:
        return "auto"
    
    num_layers = {
        'InternVL2_5-1B': 24, 
        'InternVL2_5-2B': 24, 
        'InternVL2_5-4B': 36, 
        'InternVL2_5-8B': 32,
        'InternVL2_5-26B': 48, 
        'InternVL2_5-38B': 64, 
        'InternVL2_5-78B': 80
    }[model_name]
    
    # Since the first GPU will be used for ViT, treat it as half a GPU.
    num_layers_per_gpu = math.ceil(num_layers / (world_size - 0.5))
    num_layers_per_gpu = [num_layers_per_gpu] * world_size
    num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * 0.5)
    layer_cnt = 0
    for i, num_layer in enumerate(num_layers_per_gpu):
        for j in range(num_layer):
            device_map[f'language_model.model.layers.{layer_cnt}'] = i
            layer_cnt += 1
    device_map['vision_model'] = 0
    device_map['mlp1'] = 0
    device_map['language_model.model.tok_embeddings'] = 0
    device_map['language_model.model.embed_tokens'] = 0
    device_map['language_model.model.rotary_emb'] = 0
    device_map['language_model.output'] = 0
    device_map['language_model.model.norm'] = 0
    device_map['language_model.lm_head'] = 0
    device_map[f'language_model.model.layers.{num_layers - 1}'] = 0

    return device_map

# Get model dtype
def get_model_dtype():
    return torch.bfloat16 if torch.cuda.is_available() else torch.float32

# Model loading function
def load_model():
    print(f"\n=== Loading {MODEL_NAME} ===")
    print(f"CUDA available: {torch.cuda.is_available()}")
    
    model_dtype = get_model_dtype()
    print(f"Using model dtype: {model_dtype}")
    
    if torch.cuda.is_available():
        print(f"GPU count: {torch.cuda.device_count()}")
        for i in range(torch.cuda.device_count()):
            print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
            
        # Memory info
        print(f"Total GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
        print(f"Allocated GPU memory: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
        print(f"Reserved GPU memory: {torch.cuda.memory_reserved() / 1e9:.2f} GB")
    
    # Determine device map
    device_map = "auto"
    if torch.cuda.is_available() and torch.cuda.device_count() > 1:
        model_short_name = MODEL_NAME.split('/')[-1]
        device_map = split_model(model_short_name)
    
    # Load model and tokenizer
    try:
        model = AutoModel.from_pretrained(
            MODEL_NAME,
            torch_dtype=model_dtype,
            low_cpu_mem_usage=True,
            trust_remote_code=True,
            device_map=device_map
        )
        
        tokenizer = AutoTokenizer.from_pretrained(
            MODEL_NAME,
            use_fast=False,
            trust_remote_code=True
        )
        
        print(f"✓ Model and tokenizer loaded successfully!")
        return model, tokenizer
    except Exception as e:
        print(f"❌ Error loading model: {e}")
        import traceback
        traceback.print_exc()
        return None, None

# Extract slides from uploaded PDF or PowerPoint file
def extract_slides(file_obj):
    try:
        file_bytes = file_obj.read()
        file_extension = os.path.splitext(file_obj.name)[1].lower()
        
        # Create temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as temp_file:
            temp_file.write(file_bytes)
            temp_path = temp_file.name
        
        slides = []
        
        if file_extension == '.pdf':
            # Extract images from PDF
            images = pdf2image.convert_from_path(temp_path, dpi=300)
            slides = [(f"Slide {i+1}", img) for i, img in enumerate(images)]
        
        elif file_extension in ['.ppt', '.pptx']:
            # Extract slides from PowerPoint
            prs = Presentation(temp_path)
            for i, slide in enumerate(prs.slides):
                # Create image of slide
                with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as img_file:
                    slide_path = img_file.name
                
                # We need to use pptx-export or other library to render the slide, but for this example
                # we'll create placeholder images for the slides
                img = Image.new('RGB', (1280, 720), color=(255, 255, 255))
                slides.append((f"Slide {i+1}", img))
        
        # Clean up temporary file
        os.unlink(temp_path)
        
        return slides
    
    except Exception as e:
        import traceback
        error_msg = f"Error extracting slides: {str(e)}\n{traceback.format_exc()}"
        print(error_msg)
        return []

# Image analysis function using the chat method from documentation
def analyze_slide(model, tokenizer, image, prompt):
    try:
        # Check if image is valid
        if image is None:
            return "Please upload an image first."
        
        # Process the image following official pattern
        pixel_values = load_image(image)
        
        # Debug info
        print(f"Image processed: tensor shape {pixel_values.shape}, dtype {pixel_values.dtype}")
        
        # Define generation config
        generation_config = {
            "max_new_tokens": 512,
            "do_sample": False
        }
        
        # Use the model.chat method as shown in the official documentation
        question = f"<image>\n{prompt}"
        response, _ = model.chat(
            tokenizer=tokenizer,
            pixel_values=pixel_values,
            question=question,
            generation_config=generation_config,
            history=None,
            return_history=True
        )
        
        return response
    except Exception as e:
        import traceback
        error_msg = f"Error analyzing image: {str(e)}\n{traceback.format_exc()}"
        return error_msg

# Analyze multiple slides from a PDF or PowerPoint
def analyze_multiple_slides(model, tokenizer, file_obj, prompt, num_slides=2):
    try:
        if file_obj is None:
            return "Please upload a PDF or PowerPoint file."
        
        # Extract slides from the file
        slides = extract_slides(file_obj)
        
        if not slides:
            return "No slides were extracted from the file. Please check the file format."
        
        # Limit to the requested number of slides
        slides = slides[:num_slides]
        
        # Analyze each slide
        analyses = []
        for slide_title, slide_image in slides:
            analysis = analyze_slide(model, tokenizer, slide_image, prompt)
            analyses.append((slide_title, analysis))
        
        # Format the results
        result = ""
        for slide_title, analysis in analyses:
            result += f"## {slide_title}\n\n{analysis}\n\n---\n\n"
        
        return result
    
    except Exception as e:
        import traceback
        error_msg = f"Error analyzing slides: {str(e)}\n{traceback.format_exc()}"
        return error_msg

# Main function
def main():
    # Load the model
    model, tokenizer = load_model()
    
    if model is None:
        # Create an error interface if model loading failed
        demo = gr.Interface(
            fn=lambda x: "Model loading failed. Please check the logs for details.",
            inputs=gr.Textbox(),
            outputs=gr.Textbox(),
            title="InternVL2.5 Slide Analyzer - Error",
            description="The model failed to load. Please check the logs for more information."
        )
        return demo
    
    # Create tab for single image analysis
    with gr.Blocks(title="InternVL2.5 Slide Analyzer") as demo:
        gr.Markdown("# InternVL2.5 Slide Analyzer")
        gr.Markdown("Upload an image, PDF, or PowerPoint file and ask the model to analyze it.")
        
        with gr.Tab("Single Image Analysis"):
            # Predefined prompts for analysis
            image_prompts = [
                "Describe this image in detail.",
                "What can you tell me about this image?",
                "Is there any text in this image? If so, can you read it?",
                "What is the main subject of this image?", 
                "What emotions or feelings does this image convey?",
                "Describe the composition and visual elements of this image.",
                "Summarize what you see in this image in one paragraph."
            ]
            
            with gr.Row():
                image_input = gr.Image(type="pil", label="Upload Image")
                image_prompt = gr.Dropdown(
                    choices=image_prompts, 
                    value=image_prompts[0], 
                    label="Select a prompt or write your own below",
                    allow_custom_value=True
                )
            
            image_analyze_btn = gr.Button("Analyze Image")
            image_output = gr.Textbox(label="Analysis Results", lines=15)
            
            # Handle the image analysis action
            image_analyze_btn.click(
                fn=lambda img, prompt: analyze_slide(model, tokenizer, img, prompt),
                inputs=[image_input, image_prompt],
                outputs=image_output
            )
            
            # Add examples
            gr.Examples(
                examples=[
                    ["example_images/example1.jpg", "Describe this image in detail."],
                    ["example_images/example2.jpg", "What can you tell me about this image?"]
                ],
                inputs=[image_input, image_prompt]
            )
        
        with gr.Tab("Multiple Slides Analysis"):
            # Predefined prompts for slides
            slide_prompts = [
                "Analyze this slide and describe its contents.",
                "What is the main message of this slide?",
                "Extract all the text visible in this slide.",
                "What are the key points presented in this slide?",
                "Describe the visual elements and layout of this slide.",
                "Is there any data visualization in this slide? If so, explain it.",
                "How does this slide fit into a typical presentation?"
            ]
            
            with gr.Row():
                file_input = gr.File(label="Upload PDF or PowerPoint")
                slide_prompt = gr.Dropdown(
                    choices=slide_prompts, 
                    value=slide_prompts[0], 
                    label="Select a prompt or write your own below",
                    allow_custom_value=True
                )
            
            num_slides = gr.Slider(
                minimum=1, 
                maximum=10, 
                value=2, 
                step=1, 
                label="Number of Slides to Analyze"
            )
            
            slides_analyze_btn = gr.Button("Analyze Slides")
            slides_output = gr.Markdown(label="Analysis Results")
            
            # Handle the slides analysis action
            slides_analyze_btn.click(
                fn=lambda file, prompt, num: analyze_multiple_slides(model, tokenizer, file, prompt, num),
                inputs=[file_input, slide_prompt, num_slides],
                outputs=slides_output
            )
            
            # Add example
            gr.Examples(
                examples=[
                    ["example_slides/test_slides.pdf", "Extract all the text visible in this slide.", 2]
                ],
                inputs=[file_input, slide_prompt, num_slides]
            )
    
    return demo

# Run the application
if __name__ == "__main__":
    try:
        # Check for GPU
        if not torch.cuda.is_available():
            print("WARNING: CUDA is not available. The model requires a GPU to function properly.")
            
        # Create and launch the interface
        demo = main()
        demo.launch(server_name="0.0.0.0")
    except Exception as e:
        print(f"Error starting the application: {e}")
        import traceback
        traceback.print_exc()