| """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) |
|
|
