"""North Micro Vision's visual tower and all deepstack mergers."""
import json
from pathlib import Path
from accelerate import init_empty_weights
from safetensors.torch import load_file
from transformers.models.cohere_compass.configuration_cohere_compass import CohereCompassVisionConfig
from transformers.models.cohere_compass.modeling_cohere_compass import CohereCompassVisionModel
from transformers.models.cohere_compass.image_processing_cohere_compass import CohereCompassImageProcessor

def load_model(root):
 root=Path(root);c=json.loads((root/'config.json').read_text())['vision_config'];config=CohereCompassVisionConfig(**c);config._attn_implementation='eager'
 with init_empty_weights():model=CohereCompassVisionModel(config)
 state=load_file(str(root/'model.safetensors'),device='cpu');model.load_state_dict({k.removeprefix('model.visual.'):v for k,v in state.items()},strict=True,assign=True)
 return model.eval()

def load_processor(root):
 config=json.loads((Path(root)/'preprocessor_config.json').read_text());return CohereCompassImageProcessor.from_dict(config)