import os import sys import glob import gc import torch import requests from PIL import Image from safetensors.torch import load_file import torch.nn.functional as F # Set allocator hint before large torch usage os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") # Import model components from sinq.patch_model import AutoSINQHFModel from transformers import AutoTokenizer # Configuration BASE_MODEL_DIR = "Ovis2.5-9B-sinq-4bit-experimental" #change to your model path TARGET_IMAGE_SIZE = (1024, 1024) #Change if needed DEFAULT_MAX_NEW_TOKENS = 512 #increase for longer responses DEFAULT_THINKING = False #set True to enable "thinking" mode (not recommended for grounding tasks) # Add model repo to path if BASE_MODEL_DIR not in sys.path: sys.path.insert(0, BASE_MODEL_DIR) # Import model-specific modules try: import configuration_ovis2_5 import modeling_ovis2_5 except ImportError as e: print(f"CRITICAL: Could not import model files from {BASE_MODEL_DIR}") print(f"Please ensure 'configuration_ovis2_5.py' and 'modeling_ovis2_5.py' exist in that directory.") sys.exit(1) class PrecomputedVisualTokenizer(torch.nn.Module): """Tokenizer that returns precomputed visual tokens.""" def __init__(self, visual_tokens): super().__init__() if visual_tokens is not None: self.register_buffer("_visual_tokens", visual_tokens, persistent=False) else: self.register_buffer("_visual_tokens", torch.zeros(1, 1, dtype=torch.bfloat16), persistent=False) def forward(self, pixel_values, grid_thws): return self._visual_tokens def load_model_components(): """Load all model components and return the configured model.""" # Load tokenizer print(f"Loading tokenizer from {BASE_MODEL_DIR}...") tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_DIR, trust_remote_code=True) # Load quantized LLM quantized_llm_dir = f"{BASE_MODEL_DIR}/quantized_llm" print(f"Loading quantized LLM from {quantized_llm_dir} (on cuda:0)...") qmodel_llm = AutoSINQHFModel.from_quantized( quantized_llm_dir, device="cuda:0", compute_dtype=torch.bfloat16, ) # Create full model print("Instantiating full Ovis2.5 architecture...") config = configuration_ovis2_5.Ovis2_5_Config.from_pretrained(BASE_MODEL_DIR) config.name_or_path = BASE_MODEL_DIR model = modeling_ovis2_5.Ovis2_5(config) model.eval() # Replace LLM with quantized version model.llm = qmodel_llm # Load non-LLM weights load_non_llm_weights(model) # Attach tokenizer model.text_tokenizer = tokenizer print("Model loaded successfully (LLM on cuda:0, rest on CPU).") return model def load_non_llm_weights(model): """Load weights for non-LLM components.""" # Load base model weights (excluding LLM) safetensors_files = sorted(glob.glob(os.path.join(BASE_MODEL_DIR, "*.safetensors"))) state_dict = {} for shard_path in safetensors_files: shard_dict = load_file(shard_path, device="cpu") shard_dict = {k: v for k, v in shard_dict.items() if not k.startswith("llm.")} state_dict.update(shard_dict) # Load vision encoder weights vision_shards = sorted(glob.glob(os.path.join(BASE_MODEL_DIR, "*.safetensors"))) if vision_shards: vision_state_dict = {} for shard_path in vision_shards: shard_dict = load_file(shard_path, device="cpu") vision_state_dict.update(shard_dict) state_dict.update(vision_state_dict) # Load state dict model.load_state_dict(state_dict, strict=False) def process_image(model, img_url): """Download, resize, and process the image to get visual tokens.""" print("Downloading and resizing input image...") img = Image.open(requests.get(img_url, stream=True).raw).convert("RGB") img = img.resize(TARGET_IMAGE_SIZE, resample=Image.LANCZOS) # Create message with image messages = [{ "role": "user", "content": [ {"type": "image", "image": img}, {"type": "text", "text": "Calculate the sum of the numbers in the middle box in figure (c)."}, ], }] # Preprocess inputs print("Preprocessing inputs...") input_ids, pixel_values, grid_thws = model.preprocess_inputs( messages=messages, add_generation_prompt=True, enable_thinking=DEFAULT_THINKING ) # Determine visual device visual_device = "cuda:1" if torch.cuda.is_available() and torch.cuda.device_count() > 1 else "cuda:0" # use cpu if no second gpu and not enough vram print("if you get OOM, and have only one GPU, change visual_device to 'cpu'") print(f"Using {visual_device} for visual processing") # Precompute visual tokens precomputed_visual_tokens = None if pixel_values is not None: print("Precomputing visual tokens...") original_device = next(model.visual_tokenizer.parameters()).device model.visual_tokenizer = model.visual_tokenizer.to(visual_device) with torch.no_grad(), torch.inference_mode(): vt = model.visual_tokenizer( pixel_values.to(visual_device), grid_thws.to(visual_device) if grid_thws is not None else None ) precomputed_visual_tokens = vt.cpu().to(dtype=torch.bfloat16) model.visual_tokenizer = model.visual_tokenizer.to(original_device) del pixel_values, grid_thws, vt gc.collect() torch.cuda.empty_cache() # Replace visual tokenizer with precomputed tokens model.visual_tokenizer = PrecomputedVisualTokenizer(precomputed_visual_tokens) return input_ids def prepare_for_generation(model, input_ids): """Prepare model and inputs for generation.""" device = "cuda:0" # Move non-LLM components to device print(f"Moving non-LLM components to {device}...") for name, module in model.named_children(): if name != 'llm': module.to(device, dtype=torch.bfloat16) # Move top-level buffers if hasattr(model, 'indicator_token_indices'): model.indicator_token_indices = model.indicator_token_indices.to(device) if hasattr(model, 'visual_token_placeholder'): model.visual_token_placeholder = model.visual_token_placeholder.to(device) # Move input_ids to device input_ids = input_ids.to(device) if input_ids is not None else None # Create placeholder tensors pixel_values = torch.zeros(1, 3, 8, 8, device=device, dtype=torch.bfloat16) grid_thws = torch.zeros(1, device=device, dtype=torch.int32) # Clean up memory gc.collect() torch.cuda.empty_cache() return input_ids, pixel_values, grid_thws def generate_response(model, input_ids, pixel_values, grid_thws): """Generate response from the model.""" print("\nGenerating response...") with torch.no_grad(), torch.inference_mode(): outputs = model.generate( inputs=input_ids, pixel_values=pixel_values, grid_thws=grid_thws, enable_thinking=DEFAULT_THINKING, enable_thinking_budget=False, max_new_tokens=DEFAULT_MAX_NEW_TOKENS, thinking_budget=2048, ) # Decode and print response print("\n--- Raw Generated Token IDs ---") print(outputs[0]) response = model.text_tokenizer.decode(outputs[0], skip_special_tokens=True) print("\n--- Response ---") print(response) def main(): """Main execution function.""" # Load model model = load_model_components() # Process image and get input_ids img_url = "https://cdn-uploads.huggingface.co/production/uploads/658a8a837959448ef5500ce5/TIlymOb86R6_Mez3bpmcB.png" input_ids = process_image(model, img_url) # Prepare for generation input_ids, pixel_values, grid_thws = prepare_for_generation(model, input_ids) # Generate response generate_response(model, input_ids, pixel_values, grid_thws) if __name__ == "__main__": main()