(
model_name: str,
load_pretrain: bool,
patch_size=2,
in_channel=16,
base_repo=None,
strict: bool = True,
model_kwargs=None,
)
| 6 | from .wan.modules import WanModelT2M, WanModelTM2M |
| 7 | |
| 8 | def get_transformer3d( |
| 9 | model_name: str, |
| 10 | load_pretrain: bool, |
| 11 | patch_size=2, |
| 12 | in_channel=16, |
| 13 | base_repo=None, |
| 14 | strict: bool = True, |
| 15 | model_kwargs=None, |
| 16 | ): |
| 17 | if model_kwargs is not None: |
| 18 | dense_interval = model_kwargs.get('dense_interval', 1) |
| 19 | rope_mode = model_kwargs.get('rope_mode', 'naive') |
| 20 | force_no_sincos_embed = model_kwargs.get('force_no_sincos_embed', True) |
| 21 | load_path = model_kwargs.get('load_path', None) |
| 22 | |
| 23 | if 'wanvideotm2m' in model_name.lower(): |
| 24 | import json |
| 25 | with open(f'{base_repo}/config.json') as f: |
| 26 | config = json.load(f) |
| 27 | # remove 'class_name' and '_diffusers_version' keys |
| 28 | config.pop('_class_name', None) |
| 29 | config.pop('_diffusers_version', None) |
| 30 | config['in_dim'] = in_channel |
| 31 | config['out_dim'] = in_channel |
| 32 | model = WanModelTM2M(**config, **model_kwargs) |
| 33 | if load_pretrain or load_path is not None: |
| 34 | print(f'Loading pretrain from {base_repo}') |
| 35 | state_dict_base = load_safetensors(base_repo) |
| 36 | if load_path is not None: |
| 37 | print(f'Loading pretrain from {load_path}') |
| 38 | state_dict = torch.load(load_path, map_location='cpu') |
| 39 | else: |
| 40 | state_dict = deepcopy(state_dict_base) |
| 41 | import torch.distributed as dist |
| 42 | if dist.is_initialized(): |
| 43 | if dist.get_rank() == 0: |
| 44 | model.load_state_dict(state_dict, strict=strict) |
| 45 | dist.barrier() |
| 46 | else: |
| 47 | model.load_state_dict(state_dict, strict=strict) |
| 48 | return model |
| 49 | |
| 50 | elif 'wanvideot2m' in model_name.lower(): |
| 51 | import json |
| 52 | with open(f'{base_repo}/config.json') as f: |
| 53 | config = json.load(f) |
| 54 | # remove 'class_name' and '_diffusers_version' keys |
| 55 | config.pop('_class_name', None) |
| 56 | config.pop('_diffusers_version', None) |
| 57 | config['in_dim'] = in_channel |
| 58 | config['out_dim'] = in_channel |
| 59 | model = WanModelT2M(**config, **model_kwargs) |
| 60 | if load_pretrain: |
| 61 | if load_path is not None: |
| 62 | print(f'Loading pretrain from {load_path}') |
| 63 | state_dict = torch.load(load_path, map_location='cpu') |
| 64 | else: |
| 65 | print(f'Loading pretrain from {base_repo}') |
no test coverage detected