"""Lossless UniDDT visual components. Original packed-patch contract, no LLM."""
import hashlib,importlib.util,json,sys,types
from pathlib import Path
import torch
from torch import nn
from accelerate import init_empty_weights
from safetensors.torch import load_file

def load_model(root,component='diffusion_decoder',device='cpu',dtype=None):
    root=Path(root).resolve();key='_uniddt_'+hashlib.sha256(str(root).encode()).hexdigest()[:12]
    if key not in sys.modules:
        pkg=types.ModuleType(key);pkg.__path__=[str(root/'native')];sys.modules[key]=pkg
    import importlib
    name,cls=('dit_decoder','FlattenDiT') if component=='diffusion_decoder' else ('xt_vit_adaln','Qwen3VLVisionModel')
    if component not in ['diffusion_decoder','noisy_encoder']:raise ValueError(component)
    mod=importlib.import_module(key+'.'+name);c=json.loads((root/'config.json').read_text())[component]
    with init_empty_weights(include_buffers=False):model=getattr(mod,cls)(**c)
    state={}
    for p in sorted((root/'weights'/component).glob('*.safetensors')):state.update(load_file(str(p)))
    model.load_state_dict(state,strict=True,assign=True)
    model.to(device=device,**({'dtype':dtype} if dtype is not None else {}))
    return model.eval().requires_grad_(False)