| from pathlib import Path | |
| import json,torch | |
| from torch import nn | |
| from transformers.configuration_utils import PretrainedConfig | |
| from accelerate import init_empty_weights | |
| from safetensors.torch import load_file | |
| from native.weather.encoder import WeatherEncoder | |
| from native.weather.decoder import WeatherDecoder | |
| from native.weather.projector import IdentityWeatherProjector | |
| def build(config,smoke_geometry=None): | |
| config=dict(config) | |
| if smoke_geometry is not None:config["image_size"]=list(smoke_geometry) | |
| cfg=PretrainedConfig(**config);m=nn.Module();m.model=nn.Module();m.model.modality_router=nn.Module();r=m.model.modality_router | |
| r.encoders=nn.ModuleDict({"weather":WeatherEncoder(cfg)});r.projectors=nn.ModuleDict({"weather":IdentityWeatherProjector(cfg.qwenvl_dim)});r.decoders=nn.ModuleDict({"weather":WeatherDecoder(cfg,r.encoders["weather"])}) | |
| return m | |
| def load_model(path=None,device="cpu",dtype=torch.float32,smoke_geometry=None): | |
| root=Path(path or Path(__file__).parent);cfg=json.loads((root/"weather_config.json").read_text()) | |
| with init_empty_weights():m=build(cfg,smoke_geometry) | |
| m.load_state_dict(load_file(root/"model.safetensors"),strict=True,assign=True) | |
| return m.to(device=device,dtype=dtype).eval() | |