File size: 3,401 Bytes
c28ffff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """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()
|