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"\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()