cua-agent 0.1.35__py3-none-any.whl → 0.1.38__py3-none-any.whl

agent/core/factory.py

@@ -116,6 +116,7 @@ class LoopFactory:
                 base_dir=trajectory_dir,
                 only_n_most_recent_images=only_n_most_recent_images,
                 provider_base_url=provider_base_url,
+                provider=provider,
             )
         else:
             raise ValueError(f"Unsupported loop type: {loop_type}")

agent/core/provider_config.py

@@ -8,6 +8,7 @@ DEFAULT_MODELS = {
     LLMProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
     LLMProvider.OLLAMA: "gemma3:4b-it-q4_K_M",
     LLMProvider.OAICOMPAT: "Qwen2.5-VL-7B-Instruct",
+    LLMProvider.MLXVLM: "mlx-community/UI-TARS-1.5-7B-4bit",
 }
 
 # Map providers to their environment variable names
@@ -16,4 +17,5 @@ ENV_VARS = {
     LLMProvider.ANTHROPIC: "ANTHROPIC_API_KEY",
     LLMProvider.OLLAMA: "none",
     LLMProvider.OAICOMPAT: "none",  # OpenAI-compatible API typically doesn't require an API key
+    LLMProvider.MLXVLM: "none",  # MLX VLM typically doesn't require an API key
 }

agent/core/types.py

@@ -23,6 +23,7 @@ class LLMProvider(StrEnum):
     OPENAI = "openai"
     OLLAMA = "ollama"
     OAICOMPAT = "oaicompat"
+    MLXVLM= "mlxvlm"
 
 
 @dataclass

agent/providers/uitars/clients/mlxvlm.py

@@ -0,0 +1,263 @@
+"""MLX LVM client implementation."""
+
+import io
+import logging
+import base64
+import tempfile
+import os
+import re
+import math
+from typing import Dict, List, Optional, Any, cast, Tuple
+from PIL import Image
+
+from .base import BaseUITarsClient
+import mlx.core as mx
+from mlx_vlm import load, generate
+from mlx_vlm.prompt_utils import apply_chat_template
+from mlx_vlm.utils import load_config
+from transformers.tokenization_utils import PreTrainedTokenizer
+
+logger = logging.getLogger(__name__)
+
+# Constants for smart_resize
+IMAGE_FACTOR = 28
+MIN_PIXELS = 100 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+def round_by_factor(number: float, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+def ceil_by_factor(number: float, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+def floor_by_factor(number: float, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+def smart_resize(
+    height: int, width: int, factor: int = IMAGE_FACTOR, min_pixels: int = MIN_PIXELS, max_pixels: int = MAX_PIXELS
+) -> tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+class MLXVLMUITarsClient(BaseUITarsClient):
+    """MLX LVM client implementation class."""
+
+    def __init__(
+        self, 
+        model: str = "mlx-community/UI-TARS-1.5-7B-4bit"
+    ):
+        """Initialize MLX LVM client.
+
+        Args:
+            model: Model name or path (defaults to mlx-community/UI-TARS-1.5-7B-4bit)
+        """
+        # Load model and processor
+        model_obj, processor = load(
+            model, 
+            processor_kwargs={"min_pixels": MIN_PIXELS, "max_pixels": MAX_PIXELS}
+        )
+        self.config = load_config(model)
+        self.model = model_obj
+        self.processor = processor
+        self.model_name = model
+
+    def _process_coordinates(self, text: str, original_size: Tuple[int, int], model_size: Tuple[int, int]) -> str:
+        """Process coordinates in box tokens based on image resizing using smart_resize approach.
+        
+        Args:
+            text: Text containing box tokens
+            original_size: Original image size (width, height)
+            model_size: Model processed image size (width, height)
+            
+        Returns:
+            Text with processed coordinates
+        """
+        # Find all box tokens
+        box_pattern = r"<\|box_start\|>\((\d+),\s*(\d+)\)<\|box_end\|>"
+        
+        def process_coords(match):
+            model_x, model_y = int(match.group(1)), int(match.group(2))
+            # Scale coordinates from model space to original image space
+            # Both original_size and model_size are in (width, height) format
+            new_x = int(model_x * original_size[0] / model_size[0])  # Width
+            new_y = int(model_y * original_size[1] / model_size[1])  # Height
+            return f"<|box_start|>({new_x},{new_y})<|box_end|>"
+        
+        return re.sub(box_pattern, process_coords, text)
+
+    async def run_interleaved(
+        self, messages: List[Dict[str, Any]], system: str, max_tokens: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """Run interleaved chat completion.
+
+        Args:
+            messages: List of message dicts
+            system: System prompt
+            max_tokens: Optional max tokens override
+
+        Returns:
+            Response dict
+        """
+        # Ensure the system message is included
+        if not any(msg.get("role") == "system" for msg in messages):
+            messages = [{"role": "system", "content": system}] + messages
+        
+        # Create a deep copy of messages to avoid modifying the original
+        processed_messages = messages.copy()
+        
+        # Extract images and process messages
+        images = []
+        original_sizes = {}  # Track original sizes of images for coordinate mapping
+        model_sizes = {}  # Track model processed sizes
+        image_index = 0
+        
+        for msg_idx, msg in enumerate(messages):
+            content = msg.get("content", [])
+            if not isinstance(content, list):
+                continue
+                
+            # Create a copy of the content list to modify
+            processed_content = []
+            
+            for item_idx, item in enumerate(content):
+                if item.get("type") == "image_url":
+                    image_url = item.get("image_url", {}).get("url", "")
+                    pil_image = None
+                    
+                    if image_url.startswith("data:image/"):
+                        # Extract base64 data
+                        base64_data = image_url.split(',')[1]
+                        # Convert base64 to PIL Image
+                        image_data = base64.b64decode(base64_data)
+                        pil_image = Image.open(io.BytesIO(image_data))
+                    else:
+                        # Handle file path or URL
+                        pil_image = Image.open(image_url)
+                    
+                    # Store original image size for coordinate mapping
+                    original_size = pil_image.size
+                    original_sizes[image_index] = original_size
+                    
+                    # Use smart_resize to determine model size
+                    # Note: smart_resize expects (height, width) but PIL gives (width, height)
+                    height, width = original_size[1], original_size[0]
+                    new_height, new_width = smart_resize(height, width)
+                    # Store model size in (width, height) format for consistent coordinate processing
+                    model_sizes[image_index] = (new_width, new_height)
+                    
+                    # Resize the image using the calculated dimensions from smart_resize
+                    resized_image = pil_image.resize((new_width, new_height))
+                    images.append(resized_image)
+                    image_index += 1
+                
+                # Copy items to processed content list
+                processed_content.append(item.copy())
+            
+            # Update the processed message content
+            processed_messages[msg_idx] = msg.copy()
+            processed_messages[msg_idx]["content"] = processed_content
+        
+        logger.info(f"resized {len(images)} from {original_sizes[0]} to {model_sizes[0]}")
+        
+        # Process user text input with box coordinates after image processing
+        # Swap original_size and model_size arguments for inverse transformation
+        for msg_idx, msg in enumerate(processed_messages):
+            if msg.get("role") == "user" and isinstance(msg.get("content"), str):
+                if "<|box_start|>" in msg.get("content") and original_sizes and model_sizes and 0 in original_sizes and 0 in model_sizes:
+                    orig_size = original_sizes[0]
+                    model_size = model_sizes[0]
+                    # Swap arguments to perform inverse transformation for user input
+                    processed_messages[msg_idx]["content"] = self._process_coordinates(msg["content"], model_size, orig_size)
+        
+        try:
+            # Format prompt according to model requirements using the processor directly
+            prompt = self.processor.apply_chat_template(
+                processed_messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
+            tokenizer = cast(PreTrainedTokenizer, self.processor)
+            
+            print("generating response...")
+            
+            # Generate response
+            text_content, usage = generate(
+                self.model, 
+                tokenizer, 
+                str(prompt), 
+                images, 
+                verbose=False,
+                max_tokens=max_tokens
+            )
+            
+            from pprint import pprint
+            print("DEBUG - AGENT GENERATION --------")
+            pprint(text_content)
+            print("DEBUG - AGENT GENERATION --------")
+        except Exception as e:
+            logger.error(f"Error generating response: {str(e)}")
+            return {
+                "choices": [
+                    {
+                        "message": {
+                            "role": "assistant",
+                            "content": f"Error generating response: {str(e)}"
+                        },
+                        "finish_reason": "error"
+                    }
+                ],
+                "model": self.model_name,
+                "error": str(e)
+            }
+        
+        # Process coordinates in the response back to original image space
+        if original_sizes and model_sizes and 0 in original_sizes and 0 in model_sizes:
+            # Get original image size and model size (using the first image)
+            orig_size = original_sizes[0]
+            model_size = model_sizes[0]
+            
+            # Check if output contains box tokens that need processing
+            if "<|box_start|>" in text_content:
+                # Process coordinates from model space back to original image space
+                text_content = self._process_coordinates(text_content, orig_size, model_size)
+        
+        # Format response to match OpenAI format
+        response = {
+            "choices": [
+                {
+                    "message": {
+                        "role": "assistant",
+                        "content": text_content
+                    },
+                    "finish_reason": "stop"
+                }
+            ],
+            "model": self.model_name,
+            "usage": usage
+        }
+        
+        return response

agent/providers/uitars/loop.py

@@ -23,6 +23,7 @@ from .tools.computer import ToolResult
 from .prompts import COMPUTER_USE, SYSTEM_PROMPT, MAC_SPECIFIC_NOTES
 
 from .clients.oaicompat import OAICompatClient
+from .clients.mlxvlm import MLXVLMUITarsClient
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -44,6 +45,7 @@ class UITARSLoop(BaseLoop):
         computer: Computer,
         api_key: str,
         model: str,
+        provider: Optional[LLMProvider] = None,
         provider_base_url: Optional[str] = "http://localhost:8000/v1",
         only_n_most_recent_images: Optional[int] = 2,
         base_dir: Optional[str] = "trajectories",
@@ -64,9 +66,10 @@ class UITARSLoop(BaseLoop):
             max_retries: Maximum number of retries for API calls
             retry_delay: Delay between retries in seconds
             save_trajectory: Whether to save trajectory data
+            provider: The LLM provider to use (defaults to OAICOMPAT if not specified)
         """
         # Set provider before initializing base class
-        self.provider = LLMProvider.OAICOMPAT
+        self.provider = provider or LLMProvider.OAICOMPAT
         self.provider_base_url = provider_base_url
 
         # Initialize message manager with image retention config
@@ -113,7 +116,7 @@ class UITARSLoop(BaseLoop):
             logger.error(f"Error initializing tool manager: {str(e)}")
             logger.warning("Will attempt to initialize tools on first use.")
 
-        # Initialize client for the OAICompat provider
+        # Initialize client for the selected provider
         try:
             await self.initialize_client()
         except Exception as e:
@@ -128,18 +131,28 @@ class UITARSLoop(BaseLoop):
         """Initialize the appropriate client.
 
         Implements abstract method from BaseLoop to set up the specific
-        provider client (OAICompat for UI-TARS).
+        provider client based on the configured provider.
         """
         try:
-            logger.info(f"Initializing OAICompat client for UI-TARS with model {self.model}...")
-
-            self.client = OAICompatClient(
-                api_key=self.api_key or "EMPTY",  # Local endpoints typically don't require an API key
-                model=self.model,
-                provider_base_url=self.provider_base_url,
-            )
-
-            logger.info(f"Initialized OAICompat client with model {self.model}")
+            if self.provider == LLMProvider.MLXVLM:
+                logger.info(f"Initializing MLX VLM client for UI-TARS with model {self.model}...")
+                
+                self.client = MLXVLMUITarsClient(
+                    model=self.model,
+                )
+                
+                logger.info(f"Initialized MLX VLM client with model {self.model}")
+            else:
+                # Default to OAICompat client for other providers
+                logger.info(f"Initializing OAICompat client for UI-TARS with model {self.model}...")
+                
+                self.client = OAICompatClient(
+                    api_key=self.api_key or "EMPTY",  # Local endpoints typically don't require an API key
+                    model=self.model,
+                    provider_base_url=self.provider_base_url,
+                )
+                
+                logger.info(f"Initialized OAICompat client with model {self.model}")
         except Exception as e:
             logger.error(f"Error initializing client: {str(e)}")
             self.client = None

agent/providers/uitars/utils.py

@@ -105,7 +105,7 @@ async def to_agent_response_format(
             }
         ],
         truncation="auto",
-        usage=response["usage"],
+        usage=response.get("usage", {}),
         user=None,
         metadata={},
         response=response

agent/ui/gradio/app.py

@@ -6,7 +6,7 @@ with an advanced UI for model selection and configuration.
 
 Supported Agent Loops and Models:
 - AgentLoop.OPENAI: Uses OpenAI Operator CUA model
-  • computer_use_preview
+  • computer-use-preview
 
 - AgentLoop.ANTHROPIC: Uses Anthropic Computer-Use models
   • claude-3-5-sonnet-20240620
@@ -133,12 +133,12 @@ class GradioChatScreenshotHandler(DefaultCallbackHandler):
 MODEL_MAPPINGS = {
     "openai": {
         # Default to operator CUA model
-        "default": "computer_use_preview",
+        "default": "computer-use-preview",
         # Map standard OpenAI model names to CUA-specific model names
-        "gpt-4-turbo": "computer_use_preview",
-        "gpt-4o": "computer_use_preview",
-        "gpt-4": "computer_use_preview",
-        "gpt-4.5-preview": "computer_use_preview",
+        "gpt-4-turbo": "computer-use-preview",
+        "gpt-4o": "computer-use-preview",
+        "gpt-4": "computer-use-preview",
+        "gpt-4.5-preview": "computer-use-preview",
         "gpt-4o-mini": "gpt-4o-mini",
     },
     "anthropic": {
@@ -164,8 +164,10 @@ MODEL_MAPPINGS = {
         "claude-3-7-sonnet-20250219": "claude-3-7-sonnet-20250219",
     },
     "uitars": {
-        # UI-TARS models default to custom endpoint
-        "default": "ByteDance-Seed/UI-TARS-1.5-7B",
+        # UI-TARS models using MLXVLM provider
+        "default": "mlx-community/UI-TARS-1.5-7B-4bit",
+        "mlx-community/UI-TARS-1.5-7B-4bit": "mlx-community/UI-TARS-1.5-7B-4bit",
+        "mlx-community/UI-TARS-1.5-7B-6bit": "mlx-community/UI-TARS-1.5-7B-6bit"
     },
     "ollama": {
         # For Ollama models, we keep the original name
@@ -215,7 +217,7 @@ def get_provider_and_model(model_name: str, loop_provider: str) -> tuple:
         # Determine provider and clean model name based on the full string from UI
         cleaned_model_name = model_name  # Default to using the name as-is (for custom)
 
-        if model_name == "Custom model...":
+        if model_name == "Custom model (OpenAI compatible API)":
             # Actual model name comes from custom_model_value via model_to_use.
             # Assume OAICOMPAT for custom models unless overridden by URL/key later?
             # get_provider_and_model determines the *initial* provider/model.
@@ -276,8 +278,8 @@ def get_provider_and_model(model_name: str, loop_provider: str) -> tuple:
                     break
             # Note: No fallback needed here as we explicitly check against omni keys
 
-        else:  # Handles unexpected formats or the raw custom name if "Custom model..." selected
-            # Should only happen if user selected "Custom model..."
+        else:  # Handles unexpected formats or the raw custom name if "Custom model (OpenAI compatible API)" selected
+            # Should only happen if user selected "Custom model (OpenAI compatible API)"
             # Or if a model name format isn't caught above
             provider = LLMProvider.OAICOMPAT
             cleaned_model_name = (
@@ -288,8 +290,16 @@ def get_provider_and_model(model_name: str, loop_provider: str) -> tuple:
         model_name_to_use = cleaned_model_name
         # agent_loop remains AgentLoop.OMNI
     elif agent_loop == AgentLoop.UITARS:
-        provider = LLMProvider.OAICOMPAT
-        model_name_to_use = MODEL_MAPPINGS["uitars"]["default"]  # Default 
+        # For UITARS, use MLXVLM provider for the MLX models, OAICOMPAT for custom
+        if model_name == "Custom model (OpenAI compatible API)":
+            provider = LLMProvider.OAICOMPAT
+            model_name_to_use = "tgi"
+        else:
+            provider = LLMProvider.MLXVLM
+            # Get the model name from the mappings or use as-is if not found
+            model_name_to_use = MODEL_MAPPINGS["uitars"].get(
+                model_name, model_name if model_name else MODEL_MAPPINGS["uitars"]["default"]
+            )
     else:
         # Default to OpenAI if unrecognized loop
         provider = LLMProvider.OPENAI
@@ -439,8 +449,12 @@ def create_gradio_ui(
     provider_to_models = {
         "OPENAI": openai_models,
         "ANTHROPIC": anthropic_models,
-        "OMNI": omni_models + ["Custom model..."],  # Add custom model option
-        "UITARS": ["Custom model..."],  # UI-TARS options
+        "OMNI": omni_models + ["Custom model (OpenAI compatible API)", "Custom model (ollama)"],  # Add custom model options
+        "UITARS": [
+            "mlx-community/UI-TARS-1.5-7B-4bit",
+            "mlx-community/UI-TARS-1.5-7B-6bit",
+            "Custom model (OpenAI compatible API)"
+        ],  # UI-TARS options with MLX models
     }
 
     # --- Apply Saved Settings (override defaults if available) ---
@@ -460,9 +474,9 @@ def create_gradio_ui(
             initial_model = anthropic_models[0] if anthropic_models else "No models available"
         else:  # OMNI
             initial_model = omni_models[0] if omni_models else "No models available"
-            if "Custom model..." in available_models_for_loop:
+            if "Custom model (OpenAI compatible API)" in available_models_for_loop:
                 initial_model = (
-                    "Custom model..."  # Default to custom if available and no other default fits
+                    "Custom model (OpenAI compatible API)"  # Default to custom if available and no other default fits
                 )
 
     initial_custom_model = saved_settings.get("custom_model", "Qwen2.5-VL-7B-Instruct")
@@ -485,7 +499,7 @@ def create_gradio_ui(
         
         Args:
             agent_loop_choice: The agent loop type (e.g., UITARS, OPENAI, ANTHROPIC, OMNI)
-            provider: The provider type (e.g., OPENAI, ANTHROPIC, OLLAMA, OAICOMPAT)
+            provider: The provider type (e.g., OPENAI, ANTHROPIC, OLLAMA, OAICOMPAT, MLXVLM)
             model_name: The model name
             tasks: List of tasks to execute
             provider_url: The provider base URL for OAICOMPAT providers
@@ -514,14 +528,58 @@ async def main():
             only_n_most_recent_images={recent_images},
             save_trajectory={save_trajectory},'''
         
-        # Add the model configuration based on provider
-        if provider == LLMProvider.OAICOMPAT:
+        # Add the model configuration based on provider and agent loop
+        if agent_loop_choice == "OPENAI":
+            # For OPENAI loop, always use OPENAI provider with computer-use-preview
+            code += f'''
+            model=LLM(
+                provider=LLMProvider.OPENAI, 
+                name="computer-use-preview"
+            )'''
+        elif agent_loop_choice == "ANTHROPIC":
+            # For ANTHROPIC loop, always use ANTHROPIC provider
             code += f'''
+            model=LLM(
+                provider=LLMProvider.ANTHROPIC, 
+                name="{model_name}"
+            )'''
+        elif agent_loop_choice == "UITARS":
+            # For UITARS, use MLXVLM for mlx-community models, OAICOMPAT for others
+            if provider == LLMProvider.MLXVLM:
+                code += f'''
+            model=LLM(
+                provider=LLMProvider.MLXVLM, 
+                name="{model_name}"
+            )'''
+            else:  # OAICOMPAT
+                code += f'''
+            model=LLM(
+                provider=LLMProvider.OAICOMPAT, 
+                name="{model_name}",
+                provider_base_url="{provider_url}"
+            )'''
+        elif agent_loop_choice == "OMNI":
+            # For OMNI, provider can be OPENAI, ANTHROPIC, OLLAMA, or OAICOMPAT
+            if provider == LLMProvider.OAICOMPAT:
+                code += f'''
             model=LLM(
                 provider=LLMProvider.OAICOMPAT, 
                 name="{model_name}",
                 provider_base_url="{provider_url}"
             )'''
+            else:  # OPENAI, ANTHROPIC, OLLAMA
+                code += f'''
+            model=LLM(
+                provider=LLMProvider.{provider.name}, 
+                name="{model_name}"
+            )'''
+        else:
+            # Default case - just use the provided provider and model
+            code += f'''
+            model=LLM(
+                provider=LLMProvider.{provider.name}, 
+                name="{model_name}"
+            )'''
             
         code += """
         )
@@ -547,6 +605,8 @@ async def main():
         print(f"Executing task: {{task}}")
         async for result in agent.run(task):
             print(result)'''
+
+
         
         # Add the main block
         code += '''
@@ -556,62 +616,6 @@ if __name__ == "__main__":
         
         return code
 
-    # Function to update model choices based on agent loop selection
-    def update_model_choices(loop):
-        models = provider_to_models.get(loop, [])
-        if loop == "OMNI":
-            # For OMNI, include the custom model option
-            if not models:
-                models = ["Custom model..."]
-            elif "Custom model..." not in models:
-                models.append("Custom model...")
-
-            # Show both OpenAI and Anthropic key inputs for OMNI if keys aren't set
-            return [
-                gr.update(choices=models, value=models[0] if models else "Custom model...", interactive=True),
-                gr.update(visible=not has_openai_key),
-                gr.update(visible=not has_anthropic_key)
-            ]
-        elif loop == "OPENAI":
-            # Show only OpenAI key input for OPENAI loop if key isn't set
-            if not models:
-                return [
-                    gr.update(choices=["No models available"], value="No models available", interactive=True),
-                    gr.update(visible=not has_openai_key),
-                    gr.update(visible=False)
-                ]
-            return [
-                gr.update(choices=models, value=models[0] if models else None, interactive=True),
-                gr.update(visible=not has_openai_key),
-                gr.update(visible=False)
-            ]
-        elif loop == "ANTHROPIC":
-            # Show only Anthropic key input for ANTHROPIC loop if key isn't set
-            if not models:
-                return [
-                    gr.update(choices=["No models available"], value="No models available", interactive=True),
-                    gr.update(visible=False),
-                    gr.update(visible=not has_anthropic_key)
-                ]
-            return [
-                gr.update(choices=models, value=models[0] if models else None, interactive=True),
-                gr.update(visible=False),
-                gr.update(visible=not has_anthropic_key)
-            ]
-        else:
-            # For other providers (like UITARS), don't show API key inputs
-            if not models:
-                return [
-                    gr.update(choices=["No models available"], value="No models available", interactive=True),
-                    gr.update(visible=False),
-                    gr.update(visible=False)
-                ]
-            return [
-                gr.update(choices=models, value=models[0] if models else None, interactive=True),
-                gr.update(visible=False),
-                gr.update(visible=False)
-            ]
-
     # Create the Gradio interface with advanced UI
     with gr.Blocks(title="Computer-Use Agent") as demo:
         with gr.Row():
@@ -670,14 +674,52 @@ if __name__ == "__main__":
                         info="Select the agent loop provider",
                     )
 
-                    # Create model selection dropdown with custom value support for OMNI
-                    model_choice = gr.Dropdown(
-                        choices=provider_to_models.get(initial_loop, ["No models available"]),
-                        label="LLM Provider and Model",
-                        value=initial_model,
-                        info="Select model or choose 'Custom model...' to enter a custom name",
-                        interactive=True,
-                    )
+
+                    # Create separate model selection dropdowns for each provider type
+                    # This avoids the Gradio bug with updating choices
+                    with gr.Group() as model_selection_group:
+                        # OpenAI models dropdown
+                        openai_model_choice = gr.Dropdown(
+                            choices=openai_models,
+                            label="OpenAI Model",
+                            value=openai_models[0] if openai_models else "No models available",
+                            info="Select OpenAI model",
+                            interactive=True,
+                            visible=(initial_loop == "OPENAI")
+                        )
+                        
+                        # Anthropic models dropdown
+                        anthropic_model_choice = gr.Dropdown(
+                            choices=anthropic_models,
+                            label="Anthropic Model",
+                            value=anthropic_models[0] if anthropic_models else "No models available",
+                            info="Select Anthropic model",
+                            interactive=True,
+                            visible=(initial_loop == "ANTHROPIC")
+                        )
+                        
+                        # OMNI models dropdown
+                        omni_model_choice = gr.Dropdown(
+                            choices=omni_models + ["Custom model (OpenAI compatible API)", "Custom model (ollama)"],
+                            label="OMNI Model",
+                            value=omni_models[0] if omni_models else "Custom model (OpenAI compatible API)",
+                            info="Select OMNI model or choose a custom model option",
+                            interactive=True,
+                            visible=(initial_loop == "OMNI")
+                        )
+                        
+                        # UITARS models dropdown
+                        uitars_model_choice = gr.Dropdown(
+                            choices=provider_to_models.get("UITARS", ["No models available"]),
+                            label="UITARS Model",
+                            value=provider_to_models.get("UITARS", ["No models available"])[0] if provider_to_models.get("UITARS") else "No models available",
+                            info="Select UITARS model",
+                            interactive=True,
+                            visible=(initial_loop == "UITARS")
+                        )
+                        
+                        # Hidden field to store the selected model (for compatibility with existing code)
+                        model_choice = gr.Textbox(visible=False)
 
                     # Add API key inputs for OpenAI and Anthropic
                     with gr.Group(visible=not has_openai_key and (initial_loop == "OPENAI" or initial_loop == "OMNI")) as openai_key_group:
@@ -699,34 +741,177 @@ if __name__ == "__main__":
                             type="password",
                             info="Required for Anthropic models"
                         )
+                        
+                    # Function to set OpenAI API key environment variable
+                    def set_openai_api_key(key):
+                        if key and key.strip():
+                            os.environ["OPENAI_API_KEY"] = key.strip()
+                            print(f"DEBUG - Set OpenAI API key environment variable")
+                        return key
+                    
+                    # Function to set Anthropic API key environment variable
+                    def set_anthropic_api_key(key):
+                        if key and key.strip():
+                            os.environ["ANTHROPIC_API_KEY"] = key.strip()
+                            print(f"DEBUG - Set Anthropic API key environment variable")
+                        return key
+                    
+                    # Add change event handlers for API key inputs
+                    openai_api_key_input.change(
+                        fn=set_openai_api_key,
+                        inputs=[openai_api_key_input],
+                        outputs=[openai_api_key_input],
+                        queue=False
+                    )
+                    
+                    anthropic_api_key_input.change(
+                        fn=set_anthropic_api_key,
+                        inputs=[anthropic_api_key_input],
+                        outputs=[anthropic_api_key_input],
+                        queue=False
+                    )
 
-                    # Add custom model textbox (only visible when "Custom model..." is selected)
+                    # Combined function to update UI based on selections
+                    def update_ui(loop=None, openai_model=None, anthropic_model=None, omni_model=None, uitars_model=None):
+                        # Default values if not provided
+                        loop = loop or agent_loop.value
+                        
+                        # Determine which model value to use for custom model checks
+                        model_value = None
+                        if loop == "OPENAI" and openai_model:
+                            model_value = openai_model
+                        elif loop == "ANTHROPIC" and anthropic_model:
+                            model_value = anthropic_model
+                        elif loop == "OMNI" and omni_model:
+                            model_value = omni_model
+                        elif loop == "UITARS" and uitars_model:
+                            model_value = uitars_model
+                        
+                        # Show/hide appropriate model dropdown based on loop selection
+                        openai_visible = (loop == "OPENAI")
+                        anthropic_visible = (loop == "ANTHROPIC")
+                        omni_visible = (loop == "OMNI")
+                        uitars_visible = (loop == "UITARS")
+                        
+                        # Show/hide API key inputs based on loop selection
+                        show_openai_key = not has_openai_key and (loop == "OPENAI" or (loop == "OMNI" and model_value and "OpenAI" in model_value and "Custom" not in model_value))
+                        show_anthropic_key = not has_anthropic_key and (loop == "ANTHROPIC" or (loop == "OMNI" and model_value and "Claude" in model_value and "Custom" not in model_value))
+                        
+                        # Determine custom model visibility
+                        is_custom_openai_api = model_value == "Custom model (OpenAI compatible API)"
+                        is_custom_ollama = model_value == "Custom model (ollama)"
+                        is_any_custom = is_custom_openai_api or is_custom_ollama
+                        
+                        # Update the hidden model_choice field based on the visible dropdown
+                        model_choice_value = model_value if model_value else ""
+                        
+                        # Return all UI updates
+                        return [
+                            # Model dropdowns visibility
+                            gr.update(visible=openai_visible),
+                            gr.update(visible=anthropic_visible),
+                            gr.update(visible=omni_visible),
+                            gr.update(visible=uitars_visible),
+                            # API key inputs visibility
+                            gr.update(visible=show_openai_key),
+                            gr.update(visible=show_anthropic_key),
+                            # Custom model fields visibility
+                            gr.update(visible=is_any_custom),  # Custom model name always visible for any custom option
+                            gr.update(visible=is_custom_openai_api),  # Provider base URL only for OpenAI compatible API
+                            gr.update(visible=is_custom_openai_api),   # Provider API key only for OpenAI compatible API
+                            # Update the hidden model_choice field
+                            gr.update(value=model_choice_value)
+                        ]
+                        
+                    # Add custom model textbox (visible for both custom model options)
                     custom_model = gr.Textbox(
                         label="Custom Model Name",
-                        placeholder="Enter custom model name (e.g., Qwen2.5-VL-7B-Instruct)",
+                        placeholder="Enter custom model name (e.g., Qwen2.5-VL-7B-Instruct or llama3)",
                         value=initial_custom_model,
-                        visible=(initial_model == "Custom model..."),
+                        visible=(initial_model == "Custom model (OpenAI compatible API)" or initial_model == "Custom model (ollama)"),
                         interactive=True,
                     )
 
-                    # Add custom provider base URL textbox (only visible when "Custom model..." is selected)
+                    # Add custom provider base URL textbox (only visible for OpenAI compatible API)
                     provider_base_url = gr.Textbox(
                         label="Provider Base URL",
                         placeholder="Enter provider base URL (e.g., http://localhost:1234/v1)",
                         value=initial_provider_base_url,
-                        visible=(initial_model == "Custom model..."),
+                        visible=(initial_model == "Custom model (OpenAI compatible API)"),
                         interactive=True,
                     )
 
-                    # Add custom API key textbox (only visible when "Custom model..." is selected)
+                    # Add custom API key textbox (only visible for OpenAI compatible API)
                     provider_api_key = gr.Textbox(
                         label="Provider API Key",
                         placeholder="Enter provider API key (if required)",
                         value="",
-                        visible=(initial_model == "Custom model..."),
+                        visible=(initial_model == "Custom model (OpenAI compatible API)"),
                         interactive=True,
                         type="password",
                     )
+                    
+                    # Connect agent_loop changes to update all UI elements
+                    agent_loop.change(
+                        fn=update_ui,
+                        inputs=[agent_loop, openai_model_choice, anthropic_model_choice, omni_model_choice, uitars_model_choice],
+                        outputs=[
+                            openai_model_choice, anthropic_model_choice, omni_model_choice, uitars_model_choice, 
+                            openai_key_group, anthropic_key_group,
+                            custom_model, provider_base_url, provider_api_key,
+                            model_choice  # Add model_choice to outputs
+                        ],
+                        queue=False  # Process immediately without queueing
+                    )
+
+                    # Connect each model dropdown to update UI
+                    omni_model_choice.change(
+                        fn=update_ui,
+                        inputs=[agent_loop, openai_model_choice, anthropic_model_choice, omni_model_choice, uitars_model_choice],                        
+                        outputs=[
+                            openai_model_choice, anthropic_model_choice, omni_model_choice, uitars_model_choice, 
+                            openai_key_group, anthropic_key_group,
+                            custom_model, provider_base_url, provider_api_key,
+                            model_choice  # Add model_choice to outputs
+                        ],
+                        queue=False
+                    )
+                    
+                    uitars_model_choice.change(
+                        fn=update_ui,
+                        inputs=[agent_loop, openai_model_choice, anthropic_model_choice, omni_model_choice, uitars_model_choice],             
+                        outputs=[
+                            openai_model_choice, anthropic_model_choice, omni_model_choice, uitars_model_choice, 
+                            openai_key_group, anthropic_key_group,
+                            custom_model, provider_base_url, provider_api_key,
+                            model_choice  # Add model_choice to outputs
+                        ],
+                        queue=False
+                    )
+                    
+                    openai_model_choice.change(
+                        fn=update_ui,
+                        inputs=[agent_loop, openai_model_choice, anthropic_model_choice, omni_model_choice, uitars_model_choice],             
+                        outputs=[
+                            openai_model_choice, anthropic_model_choice, omni_model_choice, uitars_model_choice, 
+                            openai_key_group, anthropic_key_group,
+                            custom_model, provider_base_url, provider_api_key,
+                            model_choice  # Add model_choice to outputs
+                        ],
+                        queue=False
+                    )
+
+                    anthropic_model_choice.change(
+                        fn=update_ui,
+                        inputs=[agent_loop, openai_model_choice, anthropic_model_choice, omni_model_choice, uitars_model_choice],             
+                        outputs=[
+                            openai_model_choice, anthropic_model_choice, omni_model_choice, uitars_model_choice, 
+                            openai_key_group, anthropic_key_group,
+                            custom_model, provider_base_url, provider_api_key,
+                            model_choice  # Add model_choice to outputs
+                        ],
+                        queue=False
+                    )
 
                     save_trajectory = gr.Checkbox(
                         label="Save Trajectory",
@@ -758,6 +943,9 @@ if __name__ == "__main__":
                     placeholder="Ask me to perform tasks in a virtual macOS environment"
                 )
                 clear = gr.Button("Clear")
+                
+                # Add cancel button
+                cancel_button = gr.Button("Cancel", variant="stop")
 
                 # Add examples
                 example_group = gr.Examples(examples=example_messages, inputs=msg)
@@ -768,10 +956,28 @@ if __name__ == "__main__":
                     history.append(gr.ChatMessage(role="user", content=message))
                     return "", history
 
+                # Function to cancel the running agent
+                async def cancel_agent_task(history):
+                    global global_agent
+                    if global_agent and hasattr(global_agent, '_loop'):
+                        print("DEBUG - Cancelling agent task")
+                        # Cancel the agent loop
+                        if hasattr(global_agent._loop, 'cancel') and callable(global_agent._loop.cancel):
+                            await global_agent._loop.cancel()
+                            history.append(gr.ChatMessage(role="assistant", content="Task cancelled by user", metadata={"title": "❌ Cancelled"}))
+                        else:
+                            history.append(gr.ChatMessage(role="assistant", content="Could not cancel task: cancel method not found", metadata={"title": "⚠️ Warning"}))
+                    else:
+                        history.append(gr.ChatMessage(role="assistant", content="No active agent task to cancel", metadata={"title": "ℹ️ Info"}))
+                    return history
+                
                 # Function to process agent response after user input
                 async def process_response(
                     history,
-                    model_choice_value,
+                    openai_model_value,
+                    anthropic_model_value,
+                    omni_model_value,
+                    uitars_model_value,
                     custom_model_value,
                     agent_loop_choice,
                     save_traj,
@@ -788,21 +994,47 @@ if __name__ == "__main__":
                     # Get the last user message
                     last_user_message = history[-1]["content"]
 
+                    # Get the appropriate model value based on the agent loop
+                    if agent_loop_choice == "OPENAI":
+                        model_choice_value = openai_model_value
+                    elif agent_loop_choice == "ANTHROPIC":
+                        model_choice_value = anthropic_model_value
+                    elif agent_loop_choice == "OMNI":
+                        model_choice_value = omni_model_value
+                    elif agent_loop_choice == "UITARS":
+                        model_choice_value = uitars_model_value
+                    else:
+                        model_choice_value = "No models available"
+                    
+                    # Determine if this is a custom model selection and which type
+                    is_custom_openai_api = model_choice_value == "Custom model (OpenAI compatible API)"
+                    is_custom_ollama = model_choice_value == "Custom model (ollama)"
+                    is_custom_model_selected = is_custom_openai_api or is_custom_ollama
+                    
                     # Determine the model name string to analyze: custom or from dropdown
-                    model_string_to_analyze = (
-                        custom_model_value
-                        if model_choice_value == "Custom model..."
-                        else model_choice_value  # Use the full UI string initially
-                    )
-
-                    # Determine if this is a custom model selection
-                    is_custom_model_selected = model_choice_value == "Custom model..."
+                    if is_custom_model_selected:
+                        model_string_to_analyze = custom_model_value
+                    else:
+                        model_string_to_analyze = model_choice_value  # Use the full UI string initially
 
                     try:
-                        # Get the provider, *cleaned* model name, and agent loop type
-                        provider, cleaned_model_name_from_func, agent_loop_type = (
-                            get_provider_and_model(model_string_to_analyze, agent_loop_choice)
-                        )
+                        # Special case for UITARS - use MLXVLM provider
+                        if agent_loop_choice == "UITARS":
+                            provider = LLMProvider.MLXVLM
+                            cleaned_model_name_from_func = model_string_to_analyze
+                            agent_loop_type = AgentLoop.UITARS
+                            print(f"Using MLXVLM provider for UITARS model: {model_string_to_analyze}")
+                        # Special case for Ollama custom model
+                        elif is_custom_ollama and agent_loop_choice == "OMNI":
+                            provider = LLMProvider.OLLAMA
+                            cleaned_model_name_from_func = custom_model_value
+                            agent_loop_type = AgentLoop.OMNI
+                            print(f"Using Ollama provider for custom model: {custom_model_value}")
+                        else:
+                            # Get the provider, *cleaned* model name, and agent loop type
+                            provider, cleaned_model_name_from_func, agent_loop_type = (
+                                get_provider_and_model(model_string_to_analyze, agent_loop_choice)
+                            )
                         
                         print(f"provider={provider} cleaned_model_name_from_func={cleaned_model_name_from_func} agent_loop_type={agent_loop_type} agent_loop_choice={agent_loop_choice}")
 
@@ -814,26 +1046,34 @@ if __name__ == "__main__":
                             else cleaned_model_name_from_func
                         )
 
-                        # Determine if OAICOMPAT should be used (only if custom model explicitly selected)
-                        is_oaicompat = is_custom_model_selected
+                        # Determine if OAICOMPAT should be used (only for OpenAI compatible API custom model)
+                        is_oaicompat = is_custom_openai_api and agent_loop_choice != "UITARS"
 
                         # Get API key based on provider determined by get_provider_and_model
                         if is_oaicompat and custom_api_key:
-                            # Use custom API key if provided for custom model
+                            # Use custom API key if provided for OpenAI compatible API custom model
                             api_key = custom_api_key
                             print(
-                                f"DEBUG - Using custom API key for model: {final_model_name_to_send}"
+                                f"DEBUG - Using custom API key for OpenAI compatible API model: {final_model_name_to_send}"
                             )
+                        elif provider == LLMProvider.OLLAMA:
+                            # No API key needed for Ollama
+                            api_key = ""
+                            print(f"DEBUG - No API key needed for Ollama model: {final_model_name_to_send}")
                         elif provider == LLMProvider.OPENAI:
                             # Use OpenAI key from input if provided, otherwise use environment variable
                             api_key = openai_key_input if openai_key_input else (openai_api_key or os.environ.get("OPENAI_API_KEY", ""))
                             if openai_key_input:
-                                print(f"DEBUG - Using provided OpenAI API key from UI")
+                                # Set the environment variable for the OpenAI API key
+                                os.environ["OPENAI_API_KEY"] = openai_key_input
+                                print(f"DEBUG - Using provided OpenAI API key from UI and set as environment variable")
                         elif provider == LLMProvider.ANTHROPIC:
                             # Use Anthropic key from input if provided, otherwise use environment variable
                             api_key = anthropic_key_input if anthropic_key_input else (anthropic_api_key or os.environ.get("ANTHROPIC_API_KEY", ""))
                             if anthropic_key_input:
-                                print(f"DEBUG - Using provided Anthropic API key from UI")
+                                # Set the environment variable for the Anthropic API key
+                                os.environ["ANTHROPIC_API_KEY"] = anthropic_key_input
+                                print(f"DEBUG - Using provided Anthropic API key from UI and set as environment variable")
                         else:
                             # For Ollama or default OAICOMPAT (without custom key), no key needed/expected
                             api_key = ""
@@ -852,8 +1092,8 @@ if __name__ == "__main__":
 
                         # Create or update the agent
                         create_agent(
-                            # Provider determined by get_provider_and_model unless custom model selected
-                            provider=LLMProvider.OAICOMPAT if is_oaicompat else provider,
+                            # Provider determined by special cases and get_provider_and_model
+                            provider=provider,
                             agent_loop=agent_loop_type,
                             # Pass the FINAL determined model name (cleaned or custom)
                             model_name=final_model_name_to_send,
@@ -966,13 +1206,21 @@ if __name__ == "__main__":
                         # Update with error message
                         history.append(gr.ChatMessage(role="assistant", content=f"Error: {str(e)}"))
                         yield history
-
-                # Connect the components
-                msg.submit(chat_submit, [msg, chatbot_history], [msg, chatbot_history]).then(
-                    process_response,
-                    [
+                        
+                # Connect the submit button to the process_response function
+                submit_event = msg.submit(
+                    fn=chat_submit,
+                    inputs=[msg, chatbot_history],
+                    outputs=[msg, chatbot_history],
+                    queue=False,
+                ).then(
+                    fn=process_response,
+                    inputs=[
                         chatbot_history,
-                        model_choice,
+                        openai_model_choice,
+                        anthropic_model_choice,
+                        omni_model_choice,
+                        uitars_model_choice,
                         custom_model,
                         agent_loop,
                         save_trajectory,
@@ -982,44 +1230,22 @@ if __name__ == "__main__":
                         openai_api_key_input,
                         anthropic_api_key_input,
                     ],
-                    [chatbot_history],
+                    outputs=[chatbot_history],
+                    queue=True,
                 )
 
                 # Clear button functionality
                 clear.click(lambda: None, None, chatbot_history, queue=False)
-
-                # Connect agent_loop changes to model selection
-                agent_loop.change(
-                    fn=update_model_choices,
-                    inputs=[agent_loop],
-                    outputs=[model_choice],
-                    queue=False,  # Process immediately without queueing
-                )
-
-                # Show/hide custom model, provider base URL, and API key textboxes based on dropdown selection
-                def update_custom_model_visibility(model_value):
-                    is_custom = model_value == "Custom model..."
-                    return (
-                        gr.update(visible=is_custom),
-                        gr.update(visible=is_custom),
-                        gr.update(visible=is_custom),
-                    )
-
-                model_choice.change(
-                    fn=update_custom_model_visibility,
-                    inputs=[model_choice],
-                    outputs=[custom_model, provider_base_url, provider_api_key],
-                    queue=False,  # Process immediately without queueing
-                )
                 
-                # Connect agent_loop changes to model selection and API key visibility
-                agent_loop.change(
-                    fn=update_model_choices,
-                    inputs=[agent_loop],
-                    outputs=[model_choice, openai_key_group, anthropic_key_group],
-                    queue=False,  # Process immediately without queueing
+                # Connect cancel button to cancel function
+                cancel_button.click(
+                    cancel_agent_task,
+                    [chatbot_history],
+                    [chatbot_history],
+                    queue=False  # Process immediately without queueing
                 )
 
+
                 # Function to update the code display based on configuration and chat history
                 def update_code_display(agent_loop, model_choice_val, custom_model_val, chat_history, provider_base_url, recent_images_val, save_trajectory_val):
                     # Extract messages from chat history
@@ -1029,9 +1255,72 @@ if __name__ == "__main__":
                             if msg.get("role") == "user":
                                 messages.append(msg.get("content", ""))
                     
-                    # Determine provider and model name based on selection
-                    model_string = custom_model_val if model_choice_val == "Custom model..." else model_choice_val
-                    provider, model_name, _ = get_provider_and_model(model_string, agent_loop)
+                    # Determine if this is a custom model selection and which type
+                    is_custom_openai_api = model_choice_val == "Custom model (OpenAI compatible API)"
+                    is_custom_ollama = model_choice_val == "Custom model (ollama)"
+                    is_custom_model_selected = is_custom_openai_api or is_custom_ollama
+                    
+                    # Determine provider and model name based on agent loop
+                    if agent_loop == "OPENAI":
+                        # For OPENAI loop, always use OPENAI provider with computer-use-preview
+                        provider = LLMProvider.OPENAI
+                        model_name = "computer-use-preview"
+                    elif agent_loop == "ANTHROPIC":
+                        # For ANTHROPIC loop, always use ANTHROPIC provider
+                        provider = LLMProvider.ANTHROPIC
+                        # Extract model name from the UI string
+                        if model_choice_val.startswith("Anthropic: Claude "):
+                            # Extract the model name based on the UI string
+                            model_parts = model_choice_val.replace("Anthropic: Claude ", "").split(" (")
+                            version = model_parts[0]  # e.g., "3.7 Sonnet"
+                            date = model_parts[1].replace(")", "") if len(model_parts) > 1 else ""  # e.g., "20250219"
+                            
+                            # Format as claude-3-7-sonnet-20250219 or claude-3-5-sonnet-20240620
+                            version = version.replace(".", "-").replace(" ", "-").lower()
+                            model_name = f"claude-{version}-{date}"
+                        else:
+                            # Use the model_choice_val directly if it doesn't match the expected format
+                            model_name = model_choice_val
+                    elif agent_loop == "UITARS":
+                        # For UITARS, use MLXVLM for mlx-community models, OAICOMPAT for custom
+                        if model_choice_val == "Custom model (OpenAI compatible API)":
+                            provider = LLMProvider.OAICOMPAT
+                            model_name = custom_model_val
+                        else:
+                            provider = LLMProvider.MLXVLM
+                            model_name = model_choice_val
+                    elif agent_loop == "OMNI":
+                        # For OMNI, provider can be OPENAI, ANTHROPIC, OLLAMA, or OAICOMPAT
+                        if is_custom_openai_api:
+                            provider = LLMProvider.OAICOMPAT
+                            model_name = custom_model_val
+                        elif is_custom_ollama:
+                            provider = LLMProvider.OLLAMA
+                            model_name = custom_model_val
+                        elif model_choice_val.startswith("OMNI: OpenAI "):
+                            provider = LLMProvider.OPENAI
+                            # Extract model name from UI string (e.g., "OMNI: OpenAI GPT-4o" -> "gpt-4o")
+                            model_name = model_choice_val.replace("OMNI: OpenAI ", "").lower().replace(" ", "-")
+                        elif model_choice_val.startswith("OMNI: Claude "):
+                            provider = LLMProvider.ANTHROPIC
+                            # Extract model name from UI string (similar to ANTHROPIC loop case)
+                            model_parts = model_choice_val.replace("OMNI: Claude ", "").split(" (")
+                            version = model_parts[0]  # e.g., "3.7 Sonnet"
+                            date = model_parts[1].replace(")", "") if len(model_parts) > 1 else ""  # e.g., "20250219"
+                            
+                            # Format as claude-3-7-sonnet-20250219 or claude-3-5-sonnet-20240620
+                            version = version.replace(".", "-").replace(" ", "-").lower()
+                            model_name = f"claude-{version}-{date}"
+                        elif model_choice_val.startswith("OMNI: Ollama "):
+                            provider = LLMProvider.OLLAMA
+                            # Extract model name from UI string (e.g., "OMNI: Ollama llama3" -> "llama3")
+                            model_name = model_choice_val.replace("OMNI: Ollama ", "")
+                        else:
+                            # Fallback to get_provider_and_model for any other cases
+                            provider, model_name, _ = get_provider_and_model(model_choice_val, agent_loop)
+                    else:
+                        # Fallback for any other agent loop
+                        provider, model_name, _ = get_provider_and_model(model_choice_val, agent_loop)
                     
                     # Generate and return the code
                     return generate_python_code(

{cua_agent-0.1.35.dist-info → cua_agent-0.1.38.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: cua-agent
-Version: 0.1.35
+Version: 0.1.38
 Summary: CUA (Computer Use) Agent for AI-driven computer interaction
 Author-Email: TryCua <gh@trycua.com>
 Requires-Python: >=3.10
@@ -23,6 +23,7 @@ Requires-Dist: openai<2.0.0,>=1.14.0; extra == "openai"
 Requires-Dist: httpx<0.29.0,>=0.27.0; extra == "openai"
 Provides-Extra: uitars
 Requires-Dist: httpx<0.29.0,>=0.27.0; extra == "uitars"
+Provides-Extra: uitars-mlx
 Provides-Extra: ui
 Requires-Dist: gradio<6.0.0,>=5.23.3; extra == "ui"
 Requires-Dist: python-dotenv<2.0.0,>=1.0.1; extra == "ui"
@@ -104,6 +105,10 @@ pip install "cua-agent[anthropic]" # Anthropic Cua Loop
 pip install "cua-agent[uitars]"    # UI-Tars support
 pip install "cua-agent[omni]" # Cua Loop based on OmniParser (includes Ollama for local models)
 pip install "cua-agent[ui]" # Gradio UI for the agent
+
+# For local UI-TARS with MLX support, you need to manually install mlx-vlm:
+pip install "cua-agent[uitars-mlx]"
+pip install git+https://github.com/ddupont808/mlx-vlm.git@stable/fix/qwen2-position-id # PR: https://github.com/Blaizzy/mlx-vlm/pull/349
 ```
 
 ## Run
@@ -206,7 +211,32 @@ The Gradio UI provides:
 
 ### Using UI-TARS
 
-You can use UI-TARS by first following the [deployment guide](https://github.com/bytedance/UI-TARS/blob/main/README_deploy.md). This will give you a provider URL like this: `https://**************.us-east-1.aws.endpoints.huggingface.cloud/v1` which you can use in the gradio UI.
+The UI-TARS models are available in two forms:
+
+1. **MLX UI-TARS models** (Default): These models run locally using MLXVLM provider
+   - `mlx-community/UI-TARS-1.5-7B-4bit` (default) - 4-bit quantized version
+   - `mlx-community/UI-TARS-1.5-7B-6bit` - 6-bit quantized version for higher quality
+
+   ```python
+   agent = ComputerAgent(
+       computer=macos_computer,
+       loop=AgentLoop.UITARS,
+       model=LLM(provider=LLMProvider.MLXVLM, name="mlx-community/UI-TARS-1.5-7B-4bit")
+   )
+   ```
+
+2. **OpenAI-compatible UI-TARS**: For using the original ByteDance model
+   - If you want to use the original ByteDance UI-TARS model via an OpenAI-compatible API, follow the [deployment guide](https://github.com/bytedance/UI-TARS/blob/main/README_deploy.md)
+   - This will give you a provider URL like `https://**************.us-east-1.aws.endpoints.huggingface.cloud/v1` which you can use in the code or Gradio UI:
+
+   ```python 
+   agent = ComputerAgent(
+       computer=macos_computer,
+       loop=AgentLoop.UITARS,
+       model=LLM(provider=LLMProvider.OAICOMPAT, name="tgi", 
+                provider_base_url="https://**************.us-east-1.aws.endpoints.huggingface.cloud/v1")
+   )
+   ```
 
 ## Agent Loops
 
@@ -216,7 +246,7 @@ The `cua-agent` package provides three agent loops variations, based on differen
 |:-----------|:-----------------|:------------|:-------------|
 | `AgentLoop.OPENAI` | • `computer_use_preview` | Use OpenAI Operator CUA model | Not Required |
 | `AgentLoop.ANTHROPIC` | • `claude-3-5-sonnet-20240620`<br>• `claude-3-7-sonnet-20250219` | Use Anthropic Computer-Use | Not Required |
-| `AgentLoop.UITARS` | • `ByteDance-Seed/UI-TARS-1.5-7B` | Uses ByteDance's UI-TARS 1.5 model | Not Required |
+| `AgentLoop.UITARS` | • `mlx-community/UI-TARS-1.5-7B-4bit` (default)<br>• `mlx-community/UI-TARS-1.5-7B-6bit`<br>• `ByteDance-Seed/UI-TARS-1.5-7B` (via openAI-compatible endpoint) | Uses UI-TARS models with MLXVLM (default) or OAICOMPAT providers | Not Required |
 | `AgentLoop.OMNI` | • `claude-3-5-sonnet-20240620`<br>• `claude-3-7-sonnet-20250219`<br>• `gpt-4.5-preview`<br>• `gpt-4o`<br>• `gpt-4`<br>• `phi4`<br>• `phi4-mini`<br>• `gemma3`<br>• `...`<br>• `Any Ollama or OpenAI-compatible model` | Use OmniParser for element pixel-detection (SoM) and any VLMs for UI Grounding and Reasoning | OmniParser |
 
 ## AgentResponse

{cua_agent-0.1.35.dist-info → cua_agent-0.1.38.dist-info}/RECORD

@@ -4,9 +4,9 @@ agent/core/agent.py,sha256=HUfBe7Uam3TObAmf6KH0GDKuNCNunNmmMcuxS7aZg0Q,8332
 agent/core/base.py,sha256=2sg8B2VqUKImRlkLTNj5lx-Oarlu7_GoMR6MbNzSY9Q,8078
 agent/core/callbacks.py,sha256=FKAxyajJ-ZJ5SxNXoupNcrm0GYBgjOjJEsStqst0EAk,6453
 agent/core/experiment.py,sha256=Ywj6q3JZFDKicfPuQsDl0vSN55HS7-Cnk3u3EcUCKe8,8866
-agent/core/factory.py,sha256=LSamFOjq2WhGGp5EVyLfDMAUrgHH1C_K5PKpdo24rhU,4573
+agent/core/factory.py,sha256=zzlCdibctqhf8Uta-SrvE-G7h59wAw-7SGhHiGvS9GY,4608
 agent/core/messages.py,sha256=-OVMDqcxK5MUHPEkHliK29XFJYMRAc1keFvzrUyrOmM,16231
-agent/core/provider_config.py,sha256=Hr9kDFSXdPeqC6hbid3OTykNF0-XVi0wzZyd44a7kww,627
+agent/core/provider_config.py,sha256=jB3fLsEsf806HQZ8jtzfSq4bCYGYONBeuCOoog_Nv_Y,768
 agent/core/telemetry.py,sha256=HElPd32k_w2SJ6t-Cc3j_2-AKdLbFwh2YlM8QViDgRw,4790
 agent/core/tools.py,sha256=Jes2CFCFqC727WWHbO-sG7V03rBHnQe5X7Oi9ZkuScI,877
 agent/core/tools/__init__.py,sha256=xZen-PqUp2dUaMEHJowXCQm33_5Sxhsx9PSoD0rq6tI,489
@@ -16,7 +16,7 @@ agent/core/tools/collection.py,sha256=NuwTn6dXSyznxWodfmFDQwUlxxaGb4oBPym4AEJABS
 agent/core/tools/computer.py,sha256=lT_aW3huoYpcM8kffuokELupSz_WZG_qkaW1gITRC58,3892
 agent/core/tools/edit.py,sha256=kv4jTKCM0VXrnoNErf7mT-xlr81-7T8v49_VA9y_L4Y,2005
 agent/core/tools/manager.py,sha256=IRsCXjGc076nncQuyIjODoafnHTDhrf9sP5B4q5Pcdo,1742
-agent/core/types.py,sha256=lDMtMFoBRW82X559VJBpbnNAzRo4LL7BbhT5r_QZFmg,2421
+agent/core/types.py,sha256=tkT-PqjgjL0oWVBRFkHAGWVwYx2Byp7PlUWSpvw_-h8,2442
 agent/core/visualization.py,sha256=1DuFF5sSeSf5BRSevBMDxml9-ajl7BQLFm5KBUwMbI8,6573
 agent/providers/__init__.py,sha256=b4tIBAaIB1V7p8V0BWipHVnMhfHH_OuVgP4OWGSHdD8,194
 agent/providers/anthropic/__init__.py,sha256=Mj11IZnVshZ2iHkvg4Z5-jrQIaD1WvzDz2Zk_pMwqIA,149
@@ -68,18 +68,19 @@ agent/providers/openai/types.py,sha256=0mFUxeFy23fJhMwc6lAFVXKngg2fJIXkPS5oV284V
 agent/providers/openai/utils.py,sha256=YeCZWIqOFSeugWoqAS0rhxOKAfL-9uN9nrYSBGBgPdc,3175
 agent/providers/uitars/__init__.py,sha256=sq5OMVJP9E_sok9tIiKJreGkjmNWXPMObjPTClYv1es,38
 agent/providers/uitars/clients/base.py,sha256=5w8Ajmq1JiPyUQJUAq1lSkfpA8_Ts80NQiDxPMTtQrI,948
+agent/providers/uitars/clients/mlxvlm.py,sha256=lMnN6ecMmWHf_l7khJ2iJHHvT7PE4XagUjrWhB0zEhc,10893
 agent/providers/uitars/clients/oaicompat.py,sha256=uYjwrGCVpFi8wj4kcaJ905ABiY6ksJZXaLlM61B2DUA,8907
-agent/providers/uitars/loop.py,sha256=CoZDk4ltz5nsw9yDnFKET5skP1uzibl3QDZOUfJQsKQ,22774
+agent/providers/uitars/loop.py,sha256=4-cgQteixPy03vp7xWezd6jWpuPkBmlLS3tizaOmd0U,23494
 agent/providers/uitars/prompts.py,sha256=_pQNd438mFpZKZT0aMl6Bd0_GgQxuy9y08kQAMPi9UM,2536
 agent/providers/uitars/tools/__init__.py,sha256=0hc3W6u5TvcXYztYKIyve_C2G3XMfwt_y7grmH0ZHC0,29
 agent/providers/uitars/tools/computer.py,sha256=TeIg_aCtMroxWOBJEiYY_YI4krW_C3pYu51tgGsVUYU,11808
 agent/providers/uitars/tools/manager.py,sha256=2dK9STtz6NuZG3i0nH7ZuHJpb7vKJ2mOVbxGsb0t8lQ,1945
-agent/providers/uitars/utils.py,sha256=S6FiZ3P-O4B15P1Gdup2o7SyuIu4nSQbspxcektpwmM,8870
+agent/providers/uitars/utils.py,sha256=493STTEEJcVhVbQgR0e8rNTI1DjkxUx8IgIv3wkJ1SU,8878
 agent/telemetry.py,sha256=pVGxbj0ewnvq4EGj28CydN4a1iOfvZR_XKL3vIOqhOM,390
 agent/ui/__init__.py,sha256=ohhxJLBin6k1hl5sKcmBST8mgh23WXgAXz3pN4f470E,45
 agent/ui/gradio/__init__.py,sha256=ANKZhv1HqsLheWbLVBlyRQ7Q5qGeXuPi5jDs8vu-ZMo,579
-agent/ui/gradio/app.py,sha256=q_nS6JJLlu1Y9xu56YHR26l_ypgaK3zR3v6BfpZT4qc,49396
-cua_agent-0.1.35.dist-info/METADATA,sha256=C7b0g8sHR6-3eWEJFQwuelXf6MKPeD8_Z8Z5aPwoikQ,11335
-cua_agent-0.1.35.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
-cua_agent-0.1.35.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-cua_agent-0.1.35.dist-info/RECORD,,
+agent/ui/gradio/app.py,sha256=M_pqSiN40F1u-8luBdqvTJxQFGzqd0WPsSz8APLbPus,67826
+cua_agent-0.1.38.dist-info/METADATA,sha256=z1UB551Vd-9YISHLzAY_gXSJFZTOgHi7lS-wy0XlgV8,12689
+cua_agent-0.1.38.dist-info/WHEEL,sha256=tSfRZzRHthuv7vxpI4aehrdN9scLjk-dCJkPLzkHxGg,90
+cua_agent-0.1.38.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+cua_agent-0.1.38.dist-info/RECORD,,