"""Selected active Mobile-O visual branches with strict loading."""
import json
from pathlib import Path
import torch
from torch import nn
from accelerate import init_empty_weights
from safetensors.torch import load_file
from diffusers import AutoencoderDC,SanaTransformer2DModel
from native.mobileclip import MCi,load_model_config
from native.mobile_block import MobileConditioningProjector

class VisionTower(nn.Module):
    def __init__(self):
        super().__init__()
        cfg=load_model_config('mobileclip_l_1024')
        self.vision_tower=MCi(cfg['image_cfg']['model_name'],projection_dim=cfg['embed_dim'])
    def forward(self,pixels):
        out=self.vision_tower(pixels,return_image_embeddings=True)['image_embeddings']
        return out.flatten(2).transpose(1,2)

class VisualComponents(nn.Module):
    def __init__(self,root):
        super().__init__()
        root=Path(root);cfg=json.loads((root/'source_config.json').read_text())
        self.model=nn.Module();m=self.model
        m.vision_tower=VisionTower()
        m.mm_projector=nn.Sequential(nn.Linear(cfg['mm_hidden_size'],cfg['hidden_size']),nn.GELU(),nn.Linear(cfg['hidden_size'],cfg['hidden_size']))
        m.diffusion_connector=MobileConditioningProjector(input_dim=cfg['hidden_size'],hidden_dim=512,output_dim=2304,num_layers=cfg['vlm_num_layers'])
        m.dit=SanaTransformer2DModel.from_config(json.loads((root/'dit_config.json').read_text()))
        m.vae=AutoencoderDC.from_config(json.loads((root/'vae_config.json').read_text()))

    def image_features(self,pixels):
        """Pixels after the official resize/pad preprocessing, rescaled to [0,1]."""
        return self.model.mm_projector(self.model.vision_tower(pixels))

    def condition(self,language_hidden_states):
        """Tuple/list of original VLM hidden states; final four layers are fused."""
        return self.model.diffusion_connector(language_hidden_states)

    def encode_latent(self,rgb):
        return self.model.vae.encode(rgb).latent*self.model.vae.config.scaling_factor

    def decode_latent(self,latent):
        return self.model.vae.decode(latent/self.model.vae.config.scaling_factor).sample

def load_model(root,device='cpu',dtype=torch.float32):
    with init_empty_weights(include_buffers=False):
        m=VisualComponents(root)
    m.load_state_dict(load_file(str(Path(root)/'model.safetensors')),strict=True,assign=True)
    return m.to(device=device,dtype=dtype).eval()