Duplicate from manpk-ai/OpenCUA-32B-Inference

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,411 @@
+---
+base_model:
+- Qwen/Qwen2.5-VL-32B-Instruct
+datasets:
+- xlangai/AgentNet
+- xlangai/aguvis-stage1
+- xlangai/aguvis-stage2
+- osunlp/UGround-V1-Data
+language:
+- en
+license: mit
+metrics:
+- code_eval
+- accuracy
+pipeline_tag: image-text-to-text
+tags:
+- VLM
+- Computer-Use-Agent
+- OS-Agent
+- GUI
+- Grounding
+library_name: transformers
+---
+
+<h1 style="
+  font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',Helvetica,Arial,sans-serif;
+  font-size:48px;
+  font-weight:700;
+  line-height:1.25;
+  text-align:center;
+  margin:0 0 24px;">
+  OpenCUA: Open Foundations for Computer-Use Agents
+</h1>
+
+<div style="
+  display:flex;
+  justify-content:center;
+  gap:12px;
+  flex-wrap:wrap;
+  margin-bottom:28px;">
+  
+  <a href="https://opencua.xlang.ai/" style="
+     display:inline-block;
+     padding:8px 24px;
+     background:#2b2b2b;
+     color:#ffffff;
+     border-radius:36px;
+     text-decoration:none;
+     font-weight:600;
+     font-size:16px;">
+    🌐 Website
+  </a>
+
+  <a href="https://arxiv.org/abs/2508.09123" style="
+     display:inline-block;
+     padding:8px 24px;
+     background:#2b2b2b;
+     color:#ffffff;
+     border-radius:36px;
+     text-decoration:none;
+     font-weight:600;
+     font-size:16px;">
+    📝 Paper
+  </a>
+
+  <a href="https://github.com/xlang-ai/OpenCUA" style="
+     display:inline-block;
+     padding:8px 24px;
+     background:#2b2b2b;
+     color:#ffffff;
+     border-radius:36px;
+     text-decoration:none;
+     font-weight:600;
+     font-size:16px;">
+    💻 Code
+  </a>
+</div>
+
+<div style="max-width:900px;margin:0 auto;">
+
+#  Introduction
+<div style="
+  max-width: 880px;              /* 可按需调节整体宽度 */
+  margin: 0 auto;               /* 居中容器 */
+  text-align: justify;          /* 关键：两端对齐 */
+  text-justify: inter-word;     /* 优化英文对齐效果 */
+  line-height: 1.6;">
+  
+OpenCUA models (OpenCUA-7B and OpenCUA-32B) are end-to-end computer-use foundation models than can produce executable actions in the computer environments.  They are based on the weights of Qwen2.5-VL-7B-Instruction and Qwen2.5-VL-32B-Instruction. 
+They demonstrate superior performance across CUA benchmarks. In particular, <b>OpenCUA-32B</b> achieves an average success rate of **34.8%** on [OSWorld-Verified](https://os-world.github.io/), 
+establishing a new state-of-the-art (SOTA) among open-source models and surpassing OpenAI CUA (GPT-4o). Both models also have strong grounding performance, OpenCUA-32B achieves 59.6% on [OSWorld-G](https://osworld-grounding.github.io/) and 55.3% on [Screenspot-Pro](https://arxiv.org/abs/2504.07981).
+</div>
+
+### Key Features
+
+- **Superior Computer-Use Capablity**: Able to execute multi-step computer-use actions with effective planning and reasoning
+- **Multi-OS Support**: Trained on demonstrations across Ubuntu, Windows, and macOS
+- **Visual Grounding**: Strong GUI element recognition and spatial reasoning capabilities
+- **Multi-Image Context**: Processes up to 3 screenshot history for better context understanding
+- **Reflective Reasoning**: Enhanced with reflective long Chain-of-Thought that identifies errors and provides corrective reasoning
+
+
+# Performance
+
+### Online Agent Evaluation
+OpenCUA models achieves strong performance on **[OSWorld-Verified](https://os-world.github.io/)**. 
+OPENCUA-32B achieves the best performance among all open-source models with an average success rate of 34.8%, outperforming prior baselines by large margins. 
+It also closes the gap to proprietary Claude models.
+<div align="center">
+
+| **Model**                        | **15 Steps** | **50 Steps** | **100 Steps** |
+|-------------------------------|:--------:|:--------:|:---------:|
+| **Proprietary**               |          |          |           |
+| OpenAI CUA                    | 26.0     | 31.3     | 31.4      |
+| Seed 1.5-VL                   | 27.9     | —        | 34.1      |
+| Claude 3.7 Sonnet             | 27.1     | 35.8     | 35.9      |
+| Claude 4 Sonnet               | 31.2     | 43.9     | 41.5      |
+| **Open-Source**               |          |          |           |
+| Qwen 2.5-VL-32B-Instruct       | 3.0      | —        | 3.9       |
+| Qwen 2.5-VL-72B-Instruct       | 4.4      | —        | 5.0       |
+| Kimi-VL-A3B                   | 9.7      | —        | 10.3      |
+| UI-TARS-72B-DPO               | 24.0     | 25.8     | 27.1      |
+| UI-TARS-1.5-7B                | 24.5     | 27.3     | 27.4      |
+| OpenCUA-7B *(Ours)*           | 24.3     | 27.9     | 26.6      |
+| **OpenCUA-32B *(Ours)***      | **29.7** | **34.1** | **34.8**  |
+</div>
+
+*OpenCUA scores are the mean of 3 independent runs.*
+
+### GUI Grounding Performance
+<div align="center">
+
+| **Model** | **OSWorld-G** | **ScreenSpot-V2** | **ScreenSpot-Pro** |
+|-------|-----------|---------------|----------------|
+| Qwen2.5-VL-7B | 31.4 | 88.8 | 27.6 |  
+| Qwen2.5-VL-32B | 46.5 | 87.0 | 39.4 |
+| UI-TARS-72B | 57.1 | 90.3 | 38.1 |
+| **OpenCUA-A3B** | 48.6 | 91.4 | 28.5 |
+| **OpenCUA-Qwen2-7B** | 45.7 | 88.5 | 23.7 |
+| **OpenCUA-7B** | 55.3 | 92.3 | 50.0 |
+| **OpenCUA-32B** | **59.6** | **93.4** | **55.3** |
+</div>
+
+
+### AgentNetBench (Offline Evaluation)
+<div align="center">
+
+| **Model** | **Coordinate Actions** | **Content Actions** | **Function Actions** | **Average** |
+|-------|-------------------|-----------------|------------------|---------|
+| Qwen2.5-VL-7B | 50.7 | 40.8 | 3.1 | 48.0 |
+| Qwen2.5-VL-32B | 66.6 | 47.2 | 41.5 | 64.8 |
+| Qwen2.5-VL-72B | 67.2 | 52.6 | 50.5 | 67.0 |
+| OpenAI CUA          | 71.7 | 57.3 | **80.0** | 73.1 |
+| **OpenCUA-7B**  | 79.0 | 62.0 | 44.3 | 75.2 |
+| **OpenCUA-32B** | **81.9** | 66.1 | 55.7 | **79.1** |
+</div>
+
+#  🚀 Quick Start
+<div style="border-left: 6px solid #f28c28; background: #fff8e6; padding: 12px 16px; margin: 16px 0;">
+  <strong>⚠️ Important for Qwen-based Models (OpenCUA-7B, OpenCUA-32B):</strong>
+  
+  To align with our training infrastructure, we have modified the model in two places:
+  <ul style="margin-top: 8px;">
+    <li>1. Multimodal Rotary Position Embedding (M-RoPE) has been replaced with 1D RoPE</strong>.</li>
+    <li>2. Using the same Tokenizer and ChatTemplate as Kimi-VL.</li>
+    <li>Do not use the default transformers and vllm classes to load the model. Tokenizer and Chat Template should be aligned if training the models.</li>
+  </ul>
+</div>
+
+
+## Installation & Download
+
+First, install the required transformers dependencies:
+
+```bash
+conda create -n opencua python=3.10
+conda activate opencua
+pip install -r requirement.txt
+```
+
+Download the model weight from huggingface:
+```bash
+from huggingface_hub import snapshot_download
+snapshot_download(
+    repo_id="xlangai/OpenCUA-32B",
+    local_dir="OpenCUA-32B",                
+    local_dir_use_symlinks=False  
+)
+```
+
+## 🎯 GUI Grounding 
+
+The following code demonstrates how to use OpenCUA models for GUI grounding tasks:
+
+```python
+import base64
+import torch
+from transformers import AutoTokenizer, AutoModel, AutoImageProcessor
+from PIL import Image
+import json
+
+def encode_image(image_path: str) -> str:
+    """Encode image to base64 string for model input."""
+    with open(image_path, "rb") as f:
+        return base64.b64encode(f.read()).decode()
+
+def load_opencua_model(model_path: str):
+    """Load OpenCUA model, tokenizer, and image processor."""
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    model = AutoModel.from_pretrained(
+        model_path, 
+        torch_dtype="auto", 
+        device_map="auto", 
+        trust_remote_code=True
+    )
+    image_processor = AutoImageProcessor.from_pretrained(model_path, trust_remote_code=True)
+    
+    return model, tokenizer, image_processor
+
+def create_grounding_messages(image_path: str, instruction: str):
+    """Create chat messages for GUI grounding task."""
+    system_prompt = (
+        "You are a GUI agent. You are given a task and a screenshot of the screen. "
+        "You need to perform a series of pyautogui actions to complete the task."
+    )
+    
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": f"data:image/png;base64,{encode_image(image_path)}"},
+                {"type": "text", "text": instruction},
+            ],
+        },
+    ]
+    return messages
+
+def run_inference(model, tokenizer, image_processor, messages, image_path):
+    """Run inference on the model."""
+    # Prepare text input
+    input_ids = tokenizer.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True
+    )
+    input_ids = torch.tensor([input_ids]).to(model.device)
+    
+    # Prepare image input  
+    image = Image.open(image_path).convert('RGB')
+    image_info = image_processor.preprocess(images=[image])
+    pixel_values = torch.tensor(image_info['pixel_values']).to(
+        dtype=torch.bfloat16, device=model.device
+    )
+    grid_thws = torch.tensor(image_info['image_grid_thw'])
+    
+    # Generate response
+    with torch.no_grad():
+        generated_ids = model.generate(
+            input_ids,
+            pixel_values=pixel_values,
+            grid_thws=grid_thws,
+            max_new_tokens=512,
+            temperature=0
+        )
+    
+    # Decode output
+    prompt_len = input_ids.shape[1]
+    generated_ids = generated_ids[:, prompt_len:]
+    output_text = tokenizer.batch_decode(
+        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )[0]
+    
+    return output_text
+
+# Example usage
+model_path = "xlangai/OpenCUA-32B"  # or other model variants
+image_path = "screenshot.png"
+instruction = "Click on the submit button"
+
+# Load model
+model, tokenizer, image_processor = load_opencua_model(model_path)
+
+# Create messages and run inference
+messages = create_grounding_messages(image_path, instruction)
+result = run_inference(model, tokenizer, image_processor, messages, image_path)
+
+print("Model output:", result)
+```
+
+<div style="border-left: 6px solid #9ca3af; background: #f5f5f5; padding: 12px 16px; margin: 16px 0;">
+  <em>Expected result: ```python
+pyautogui.click(x=1432, y=344)
+```</em>
+</div>
+
+## 🖥️ Computer Use Agent
+**[OpenCUAAgent](https://github.com/xlang-ai/OSWorld/blob/main/mm_agents/opencua_agent.py)** is developed in the [OSWorld](https://github.com/xlang-ai/OSWorld) environment based on OpenCUA models. It iteratively perceives the environment via screenshots, produces reflective long CoT as inner monologue, and predicts the next action to be executed. OpenCUAAgent uses 3 images in total and L2 CoT format in default.
+
+Command for running OpenCUA-7B and OpenCUA-32B in OSWorld:
+```
+    python run_multienv_opencua.py \
+        --headless \
+        --observation_type screenshot \
+        --model OpenCUA-32B \
+        --result_dir ./results --test_all_meta_path evaluation_examples/test_all_no_gdrive.json \
+        --max_steps 100 \
+        --num_envs 30  \
+        --coordinate_type qwen25
+```
+<div style="border-left: 6px solid #9ca3af; background: #f5f5f5; padding: 12px 16px; margin: 16px 0;">
+  <em>Currently we only supports huggingface inference. We are implementing the vLLM supports of OpenCUA models. Please stay tuned.</em>
+</div>
+
+## Important Notes on Coordinate Systems
+<div style="border-left: 6px solid #9ca3af; background: #f5f5f5; padding: 12px 16px; margin: 16px 0;">
+  <ul style="margin: 0;">
+    <li><strong><code>xlangai/OpenCUA-A3B</code></strong> – Relative coordinates <em>(not supported in this code)</em></li>
+    <li><strong><code>xlangai/OpenCUA-Qwen2-7B</code></strong> – Relative coordinates</li>
+    <li><strong><code>xlangai/OpenCUA-7B</code></strong> – Absolute coordinates</li>
+    <li><strong><code>xlangai/OpenCUA-32B</code></strong> – Absolute coordinates</li>
+  </ul>
+</div>
+
+**OpenCUA models use different coordinate systems depending on the base model:**
+
+- **OpenCUA-Qwen2-7B**: Outputs **relative coordinates** (0.0 to 1.0 range)
+  ```python
+  # Example output: pyautogui.click(x=0.5, y=0.3)
+  # x=0.5 means 50% from left edge, y=0.3 means 30% from top edge
+  
+  # Convert to absolute coordinates:
+  def qwen2_relative_to_absolute(rel_x, rel_y, original_width, original_height):
+      abs_x = int(rel_x * original_width)
+      abs_y = int(rel_y * original_height)
+      return abs_x, abs_y
+  ```
+
+- **OpenCUA-7B and OpenCUA-32B** (Qwen2.5-based): Output **absolute coordinates** after smart resize
+  ```python
+  # Example output: pyautogui.click(x=960, y=324)  
+  # These are coordinates on the smart-resized image, not the original image
+  
+  # Convert to original image coordinates:
+  # Please refer to the smart_resize function in: https://github.com/huggingface/transformers/blob/67ddc82fbc7e52c6f42a395b4a6d278c55b77a39/src/transformers/models/qwen2_vl/image_processing_qwen2_vl.py#L55
+  def qwen25_smart_resize_to_absolute(model_x, model_y, original_width, original_height):
+      # First, calculate the smart-resized dimensions
+      resized_height, resized_width = smart_resize(original_height, original_width, factor = 28, min_pixels = 3136, max_pixels = 12845056)
+      
+      # Convert model output to relative coordinates on original image
+      rel_x = model_x / resized_width
+      rel_y = model_y / resized_height
+      
+      # Then convert to absolute coordinates on original image
+      abs_x = int(rel_x * original_width)
+      abs_y = int(rel_y * original_height)
+      return abs_x, abs_y
+  ```
+
+<div style="border-left: 6px solid #9ca3af; background: #f5f5f5; padding: 12px 16px; margin: 16px 0;">
+  <strong>Understanding Smart Resize for Qwen2.5-based Models:</strong>
+  <p style="margin: 8px 0 0;">
+    The Qwen2.5-VL models use a “smart resize” preprocessing that maintains aspect ratio while fitting within pixel constraints.
+    For coordinate conversion, you need the smart resize function from the
+    <a href="https://github.com/QwenLM/Qwen2.5-VL/blob/d2240f11656bfe404b9ba56db4e51cd09f522ff1/qwen-vl-utils/src/qwen_vl_utils/vision_process.py#L60">
+      official Qwen2.5-VL implementation</a>.
+  </p>
+</div>
+
+
+# TODO
+## vLLM Support
+We are actively working with the vLLM team to add support for OpenCUA models.
+
+**Workaround:** For now, please use the standard transformers library as shown in the examples above. We will update this section once vLLM support becomes available.
+
+## Training Code
+OpenCUA models are developed based on the training infrastructure of Kimi Team. We are developting the training pipeline based on the open-source infrastructure as well.
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file in the root folder for details.
+
+## Research Use and Disclaimer
+
+OpenCUA models are intended for **research and educational purposes only**. 
+
+### Prohibited Uses
+- The model may **not** be used for any purpose or activity that violates applicable laws or regulations in any jurisdiction
+- Use for illegal, unethical, or harmful activities is strictly prohibited
+
+### Disclaimer
+- The authors, contributors, and copyright holders are **not responsible** for any illegal, unethical, or harmful use of the Software, nor for any direct or indirect damages resulting from such use
+- Use of the "OpenCUA" name, logo, or trademarks does **not** imply any endorsement or affiliation unless separate written permission is obtained
+- Users are solely responsible for ensuring their use complies with applicable laws and regulations
+
+## Citation
+
+If you use OpenCUA models in your research, please cite our work:
+
+```bibtex
+@misc{wang2025opencuaopenfoundationscomputeruse,
+      title={OpenCUA: Open Foundations for Computer-Use Agents}, 
+      author={Xinyuan Wang and Bowen Wang and Dunjie Lu and Junlin Yang and Tianbao Xie and Junli Wang and Jiaqi Deng and Xiaole Guo and Yiheng Xu and Chen Henry Wu and Zhennan Shen and Zhuokai Li and Ryan Li and Xiaochuan Li and Junda Chen and Boyuan Zheng and Peihang Li and Fangyu Lei and Ruisheng Cao and Yeqiao Fu and Dongchan Shin and Martin Shin and Jiarui Hu and Yuyan Wang and Jixuan Chen and Yuxiao Ye and Danyang Zhang and Dikang Du and Hao Hu and Huarong Chen and Zaida Zhou and Haotian Yao and Ziwei Chen and Qizheng Gu and Yipu Wang and Heng Wang and Diyi Yang and Victor Zhong and Flood Sung and Y. Charles and Zhilin Yang and Tao Yu},
+      year={2025},
+      eprint={2508.09123},
+      archivePrefix={arXiv},
+      primaryClass={cs.AI},
+      url={https://arxiv.org/abs/2508.09123}, 
+}
+```
+
+</div>

config.json

@@ -0,0 +1,69 @@
+{
+  "architectures": [
+    "OpenCUAForConditionalGeneration"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_opencua.OpenCUAConfig",
+    "AutoModel": "modeling_opencua.OpenCUAForConditionalGeneration",
+    "AutoModelForCausalLM": "modeling_opencua.OpenCUAForConditionalGeneration"
+  },
+  "ignore_index": -100,
+  "media_placeholder_token_id": 151664,
+  "model_type": "opencua",
+  "pad_token_id": 0,
+  "text_config": {
+    "bos_token_id": 151643,
+    "eos_token_id": 151644,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 5120,
+    "initializer_range": 0.02,
+    "intermediate_size": 27648,
+    "k_proj_bias": true,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "qwen2",
+    "num_attention_heads": 40,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 64,
+    "num_key_value_heads": 8,
+    "pad_token_id": 152063,
+    "pretraining_sequence_length": 131072,
+    "q_proj_bias": true,
+    "rms_norm_eps": 1e-05,
+    "rope_theta": 1000000.0,
+    "tie_word_embeddings": false,
+    "torch_dtype": "bfloat16",
+    "use_bfloat16": false,
+    "use_cache": true,
+    "v_proj_bias": true,
+    "vocab_size": 152064
+  },
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.48.3",
+  "vision_config": {
+    "depth": 32,
+    "fullatt_block_indexes": [
+      7,
+      15,
+      23,
+      31
+    ],
+    "hidden_act": "silu",
+    "hidden_size": 1280,
+    "num_heads": 16,
+    "in_chans": 3,
+    "intermediate_size": 3456,
+    
+    "patch_size": 14,
+    "spatial_merge_size": 2,
+    "spatial_patch_size": 14,
+    "temporal_patch_size": 2,
+    "out_hidden_size": 5120,
+    "tokens_per_second": 2,
+    "window_size": 112
+  },
+  "vocab_size": 152064
+}

configuration_opencua.py

@@ -0,0 +1,37 @@
+from transformers.configuration_utils import PretrainedConfig
+from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLVisionConfig
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+
+
+class OpenCUAConfig(PretrainedConfig):
+    """OpenCUA-2.5-32B model configuration.
+
+    Args:
+        vision_config: Configuration for the vision model.Qwen2_5_VLVisionConfig
+        text_config: Configuration for the text model. Qwen2Config
+        pad_token_id: The token ID to use for padding.
+    """
+
+    model_type = "opencua"
+
+    def __init__(
+        self,
+        vision_config: dict | Qwen2_5_VLVisionConfig | None = None,
+        text_config: dict | Qwen2Config | None = None,
+        ignore_index: int = -100,
+        media_placeholder_token_id: int = 151664,
+        pad_token_id: int = 0,
+        **kwargs
+    ):
+        if isinstance(vision_config, dict):
+            vision_config = Qwen2_5_VLVisionConfig(**vision_config)
+        self.vision_config = vision_config
+
+        if isinstance(text_config, dict):
+            text_config = Qwen2Config(**text_config)
+        self.text_config = text_config
+
+        self.ignore_index = ignore_index
+        self.media_placeholder_token_id = media_placeholder_token_id
+
+        super().__init__(pad_token_id=pad_token_id, **kwargs)

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "max_length": 32768,
+  "eos_token_id": 151644
+}

handler.py

@@ -0,0 +1,171 @@
+from typing import Dict, List, Any
+from transformers import (
+    AutoTokenizer,
+    AutoModel,
+    AutoImageProcessor,
+)
+import torch
+from PIL import Image
+import base64
+import io
+
+# get dtype and device
+dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        print(f"Initializing model on device: {device}")
+        print(f"Using dtype: {dtype}")
+        
+        # load the model - using AutoModel like in local inference
+        self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
+        self.image_processor = AutoImageProcessor.from_pretrained(path, trust_remote_code=True)
+        
+        # Load model with explicit device mapping
+        if device == "cuda":
+            self.model = AutoModel.from_pretrained(
+                path, 
+                torch_dtype=dtype, 
+                trust_remote_code=True,
+                device_map="auto"  # Automatically map to available GPUs
+            )
+        else:
+            self.model = AutoModel.from_pretrained(
+                path, 
+                torch_dtype=dtype, 
+                trust_remote_code=True
+            )
+            self.model = self.model.to(device)
+        
+        print(f"Model loaded successfully on device: {self.model.device}")
+        print(f"Model dtype: {next(self.model.parameters()).dtype}")
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+       data args:
+            inputs (:obj: `str` or `list`): messages in chat format or text input
+            parameters (:obj: `dict`): generation parameters
+      Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+        print("Call inside handler")
+        # get inputs
+        inputs = data.pop("inputs", data)
+        parameters = data.pop("parameters", {})
+        print("parameters", parameters)
+        
+        # Remove parameters that might cause issues
+        parameters.pop("details", None)
+        parameters.pop("stop", None)
+        parameters.pop("return_full_text", None)
+        if "do_sample" in parameters:
+            parameters["do_sample"] = True
+        
+        # Set default generation parameters
+        max_new_tokens = parameters.pop("max_new_tokens", 512)
+        temperature = parameters.pop("temperature", 0)
+        
+        try:
+            # Handle different input formats
+            if isinstance(inputs, str):
+                # If it's a string, treat it as a simple text prompt
+                input_ids = self.tokenizer.encode(inputs, return_tensors="pt").to(self.model.device)
+                generated_ids = self.model.generate(
+                    input_ids,
+                    max_new_tokens=max_new_tokens,
+                    temperature=temperature,
+                    **parameters
+                )
+                prompt_len = input_ids.shape[1]
+                generated_ids = generated_ids[:, prompt_len:]
+                output_text = self.tokenizer.batch_decode(
+                    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+                )
+                return [{"generated_text": output_text[0]}]
+            
+            elif isinstance(inputs, list):
+                # Handle chat format with images
+                messages = inputs
+                
+                # Apply chat template
+                input_ids = self.tokenizer.apply_chat_template(
+                    messages, tokenize=True, add_generation_prompt=True
+                )
+                input_text = self.tokenizer.decode(input_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)
+                print(input_text)
+
+                input_ids = torch.tensor([input_ids]).to(self.model.device)
+                
+                # Process ALL images if present
+                pixel_values_list = []
+                grid_thws_list = []
+                
+                # Look for images in the messages
+                for message in messages:
+                    if isinstance(message.get("content"), list):
+                        for content_item in message["content"]:
+                            if content_item.get("type") == "image_url":
+                                image_data = content_item.get("image_url").get("url", "")
+                                if image_data.startswith("data:image"):
+                                    # Decode base64 image
+                                    image_data = image_data.split(",")[1]
+                                    image_bytes = base64.b64decode(image_data)
+                                    image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
+                                    
+                                    # Process each image individually
+                                    info = self.image_processor.preprocess(images=[image])
+                                    pixel_values = torch.tensor(info['pixel_values']).to(dtype=dtype, device=self.model.device)
+                                    grid_thws = torch.tensor(info['image_grid_thw']).to(self.model.device)
+                                    
+                                    pixel_values_list.append(pixel_values)
+                                    grid_thws_list.append(grid_thws)
+                
+                # Generate response
+                if pixel_values_list and grid_thws_list:
+                    # Multi-modal generation with images
+                    # Concatenate all pixel_values and grid_thws for batch processing
+                    all_pixel_values = torch.cat(pixel_values_list, dim=0)
+                    all_grid_thws = torch.cat(grid_thws_list, dim=0)
+                    
+                    print(f"Processing {len(pixel_values_list)} images")
+                    print(f"pixel_values shape: {all_pixel_values.shape}")
+                    print(f"grid_thws shape: {all_grid_thws.shape}")
+                    print("grid_thws", all_grid_thws)
+                    
+                    # Ensure all tensors are on the same device as the model
+                    all_pixel_values = all_pixel_values.to(self.model.device)
+                    all_grid_thws = all_grid_thws.to(self.model.device)
+                    
+                    with torch.no_grad():
+                        generated_ids = self.model.generate(
+                            input_ids, 
+                            pixel_values=all_pixel_values, 
+                            grid_thws=all_grid_thws,
+                            max_new_tokens=max_new_tokens,
+                            temperature=temperature,
+                            **parameters
+                        )
+                else:
+                    # Text-only generation
+                    generated_ids = self.model.generate(
+                        input_ids,
+                        max_new_tokens=max_new_tokens,
+                        temperature=temperature,
+                        **parameters
+                    )
+                
+                prompt_len = input_ids.shape[1]
+                generated_ids = generated_ids[:, prompt_len:]
+                output_text = self.tokenizer.batch_decode(
+                    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+                )
+                print("##Model Response##", output_text)
+                return [{"generated_text": output_text[0]}]
+            
+            else:
+                raise ValueError(f"Unsupported input type: {type(inputs)}")
+                
+        except Exception as e:
+            print(f"Error during inference: {str(e)}")
+            return [{"error": str(e)}]

model-1-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8aca1939cb672391efc7d1776450497a468c7cb81c4b1e3ed1eea53af67b616b
+size 3910076520

model-10-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:314e7614356d3c95e09e544e9fd47562d85951ab89dbee7ce6a287e063ecf293
+size 975212312

model-11-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad1cf3ef5479dd3b80dd9481ddf61d286c0aff0df324e00eb780d0441c25ca6e
+size 975212320

model-12-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:101a00c1f24a1384d2fc4392fc4e8a158f5af10178386bbd1d68f7b9c8f55f28
+size 975212320

model-13-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d29d0f6d70549ea8e0fc5931ac5af1f8e3fed24341bfd0609968f0cd8af0883
+size 975212320

model-14-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3829ef7b272767f8b7dbfaedfa6c646c970481d7adac8589bcc4eab83e205f1e
+size 975212320

model-15-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54289f90134d13eb68235d633421b91f8523b7ab1d3bdcedb56abbc842e1abc4
+size 975212320

model-16-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b58ab69762e4a80f82297f29e0181cd272da54c51824dfb0cf837bc84f66fbb
+size 975212320

model-17-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fc03b542261f4832f9cb76227d3a6ed57820b515b7621e5980e4781768bd6308
+size 975212320

model-18-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7aa625aef03f270e160526c9d767f9223b49161715a60000970ba61da8002af6
+size 975212320

model-19-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10f0e43ec131e47545673ccf8f82fb280b5cfe9eb8b176de09af3412b83ebd3c
+size 975212320

model-2-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32e4ea7efa92c57f60465442d016cf63f3d91e21c667ce97180d7b461aeec2ee
+size 975212312

model-20-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f55d2a8fb141e8f2680859419041dbf9ed9bf3476fe1b04ad63a1c80b515d2d1
+size 975212320

model-21-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a44ba3081602764e1e60f2097fd3f643c6bedde6f61d34101e40c19ecdf9a845
+size 975212320

model-22-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90147767925bbb504d648eac8677128aa77e201d8c005365f16e90042946d7fb
+size 975212320

model-23-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e23cae09d39bc395bf12220c44ee49de74c2e2adb4bb74c42a31300404f6e97b
+size 975212320

model-24-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db13aedf91d13cb8817816a746fa64e34d171e90ad8e1e2dda6c2cc311d3cb5d
+size 975212320

model-25-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:238eb5f6b8452295aef90060bb2e6f72edebfd316f63eb0b4ded48aab7285c87
+size 975212320

model-26-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:232379d9dafa70056ab2a6031226ddce4d0d2b5e8a49b6e60ef7e9722e84b087
+size 975212320

model-27-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cf53d1a0e56fb2d110c32cbffa1e29bfa2ebab91ef5e338c236e8a1b4197725
+size 975212320

model-28-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f77ad190724c9426fdaa7e4394f69cde049e4890425543d1e4fd8ff2ca949bb0
+size 975212320

model-29-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0754b041c3d5cfc3dea0161698e788139329f71866e0de83d2836fe9ac67c671
+size 975212320

model-3-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a90489efdefc82a92eab3393dda904cf81ed0c2f72e76af3f2bf26c7f53da370
+size 975212312

model-30-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fdd7bbde153529db81a39a92a6702ce5bee72aedf4c90e1db484d66966f375bb
+size 975212320

model-31-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e51a5b5edfd8ccb34a2ddde73bf787991b7ca7fb17bf71aa9eaa2059a320552
+size 975212320

model-32-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:96e43de41ad3132d829f9639840f0da60ac3388e891379c57358e07718230c8d
+size 975212320

model-33-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13ba5a591d8e06a48b1e808ffdb5a42df90139b15e8df42233614c92ccc3f6b5
+size 975212320

model-34-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08a2fe4e193b08d56c60d4511e9647cca566d0042317959c816e0698a7e31724
+size 975212320

model-35-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:504b5985bc286473c0b8fb7708873492f0849fc127bc568aad2cb6a722860de2
+size 975212320

model-36-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c8fa50a75d040f40481eedc2fc539ad1077f8233b0041be408ae3504b03b67fd
+size 975212320

model-37-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:12dc0628c16279a89233334470722c31f3932a83968563f4e8d42c68677ba5d4
+size 975212320

model-38-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d942860b48b66ced151d242caaf4439e7cb62061ea84dd41395b43b08e5aa536
+size 975212320

model-39-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8cf8c139d14df728793a85e6c4e947f13816e3da2f12df02382b8299580081bb
+size 975212320

model-4-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9acdd3d8b5218fbf943de6cf948480aa19abe0776a46c044d644aab920ddfdac
+size 975212312

model-40-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:510cd3e52d48e1175966d691ebc0739012b31fcf766a31d1d41c50ce011a7c60
+size 975212320

model-41-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f2467e759d9ca1d233b04f9b751f012732a4db2ba58c91dccf00df81b9fe0a63
+size 975212320

model-42-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:767ee823997ecda01d8841c6f8277bbc33adfe07639c954b5b29c04f1a7d1942
+size 975212320

model-43-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:497e9690ec2b244e6cd5d6ed8e905dbc25308b46fd750fec2da5bca4b2ae0d6f
+size 975212320

model-44-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f5eb867d8ded2a16569f2d681a26f336acce5198bef1310c56a7a231104a978
+size 975212320

model-45-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb08c3f5843186b2757693bcdeffdf6bf3517edf70f803c764f1b9ff12d367ce
+size 975212320

model-46-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe59b664845ef521a7980cf77dd2e03457686a3900bfbddc16a0a4685e3abc06
+size 975212320

model-47-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32d200705d0f038dcdb56a0420541546cfb379f75a25673ccd53c4fc86abb25f
+size 975212320

model-48-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d028ef374ea67dbd9237ffa7fed95567e3152b22e887a237f7130612c5e03957
+size 975212320

model-49-of-64.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:795aa39683c39d7f37046ab3630e863d29837f0007d6013086bdfd92d48a7240
+size 975212320