| """Loader-path tests run with no torch, model/data or GPU.""" |
| import json |
| from pathlib import Path |
| import sys |
| import tempfile |
| import unittest |
| from unittest.mock import patch |
| _code = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(_code / ('core' if (_code / 'core').is_dir() else 'code_support'))) |
| from model_resources import wan_model_paths, wan_base_root, wan_tokenizer_path, indexed_shards |
|
|
|
|
| class ResourceTests(unittest.TestCase): |
| def setUp(self): |
| self.temporary = tempfile.TemporaryDirectory() |
| self.addCleanup(self.temporary.cleanup) |
| self.root = Path(self.temporary.name) |
|
|
| def add_index(self, relative, names): |
| index = self.root / relative |
| index.parent.mkdir(parents=True, exist_ok=True) |
| for name in names: |
| (index.parent / name).write_bytes(b'path fixture only') |
| index.write_text(json.dumps({'weight_map': {str(i): name for i, name in enumerate(names)}})) |
| return index |
|
|
| def test_original_monolithic_paths_unchanged(self): |
| names = ['diffusion_pytorch_model.safetensors', 'models_t5_umt5-xxl-enc-bf16.pth', 'Wan2.1_VAE.pth'] |
| for name in names: |
| (self.root / name).write_bytes(b'fixture') |
| self.assertEqual(wan_model_paths(self.root, strict=True), [str(self.root / n) for n in names]) |
|
|
| def test_new_export_paths_are_one_list_per_component(self): |
| self.add_index('transformer/diffusion_pytorch_model.safetensors.index.json', ['dit-2.safetensors', 'dit-1.safetensors']) |
| self.add_index('text_encoder/model.safetensors.index.json', ['t5-1.safetensors', 't5-2.safetensors']) |
| (self.root / 'Wan2.1_VAE.pth').write_bytes(b'fixture') |
| paths = wan_model_paths(self.root, strict=True) |
| self.assertEqual(len(paths), 3) |
| self.assertEqual([len(paths[0]), len(paths[1])], [2, 2]) |
| self.assertTrue(paths[0][0].endswith('dit-1.safetensors')) |
| self.assertIsInstance(paths[2], str) |
|
|
| def test_missing_weights_allow_dry_run_but_strict_fails(self): |
| self.assertIsInstance(wan_model_paths(self.root)[0], str) |
| with self.assertRaises(FileNotFoundError): |
| wan_model_paths(self.root, strict=True) |
|
|
| def test_missing_shard_fails_even_without_strict(self): |
| index = self.add_index('transformer/diffusion_pytorch_model.safetensors.index.json', ['a.safetensors']) |
| (index.parent / 'a.safetensors').unlink() |
| with self.assertRaises(FileNotFoundError): |
| wan_model_paths(self.root) |
|
|
| def test_unsafe_index_path_rejected(self): |
| index = self.root / 'index.json' |
| index.write_text('{"weight_map":{"a":"../escape.safetensors"}}') |
| with self.assertRaisesRegex(ValueError, 'unsafe'): |
| indexed_shards(index) |
|
|
| def test_environment_override(self): |
| with patch.dict('os.environ', {'WAN_MODEL_ROOT': str(self.root)}): |
| self.assertEqual(wan_base_root(), str(self.root)) |
|
|
| def test_tokenizer_is_checked_only_in_strict_mode(self): |
| tokenizer = Path(wan_tokenizer_path(self.root)) |
| with self.assertRaises(FileNotFoundError): |
| wan_tokenizer_path(self.root, strict=True) |
| tokenizer.mkdir(parents=True) |
| (tokenizer / 'spiece.model').write_bytes(b'fixture') |
| self.assertEqual(wan_tokenizer_path(self.root, strict=True), str(tokenizer)) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
