ai-art

 1import os
 2import io
 3import pickle
 4
 5import requests
 6
 7import torch
 8from torch.serialization import (
 9    _get_restore_location,
10    _maybe_decode_ascii,
11    _open_file_like,
12    _open_zipfile_reader,
13)
14
15from tqdm import tqdm
16
17
18def safe_load(
19    f,
20    map_location=None,
21    pickle_module=pickle,
22    pickle_file="data.pkl",
23    **pickle_load_args,
24):
25    with _open_file_like(f, "rb") as opened_file:
26        with _open_zipfile_reader(opened_file) as zip_file:
27            restore_location = _get_restore_location(map_location)
28
29            loaded_storages = {}
30
31            def load_tensor(data_type, size, key, location):
32                name = f"data/{key}"
33                dtype = data_type(0).dtype
34
35                storage = zip_file.get_storage_from_record(name, size, dtype).storage()
36                loaded_storages[key] = restore_location(storage, location)
37
38            def persistent_load(saved_id):
39                assert isinstance(saved_id, tuple)
40                typename = _maybe_decode_ascii(saved_id[0])
41                data = saved_id[1:]
42
43                assert (
44                    typename == "storage"
45                ), f"Unknown typename for persistent_load, expected 'storage' but got '{typename}'"
46                data_type, key, location, size = data
47                if key not in loaded_storages:
48                    load_tensor(data_type, size, key, _maybe_decode_ascii(location))
49                storage = loaded_storages[key]
50                return storage
51
52            load_module_mapping = {"torch.tensor": "torch._tensor"}
53
54            class UnpicklerWrapper(pickle_module.Unpickler):
55                def find_class(self, mod_name, name):
56                    try:
57                        mod_name = load_module_mapping.get(mod_name, mod_name)
58                        return super().find_class(mod_name, name)
59                    except Exception:
60                        pass
61
62            # Load the data (which may in turn use `persistent_load` to load tensors)
63            data_file = io.BytesIO(zip_file.get_record(pickle_file))
64
65            unpickler = UnpicklerWrapper(data_file, **pickle_load_args)
66            unpickler.persistent_load = persistent_load
67            result = unpickler.load()
68
69            torch._utils._validate_loaded_sparse_tensors()
70
71            return result
72
73
74def download(url, local_path, chunk_size=1024):
75    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
76    with requests.get(url, stream=True) as r:
77        total_size = int(r.headers.get("content-length", 0))
78        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
79            with open(local_path, "wb") as f:
80                for data in r.iter_content(chunk_size=chunk_size):
81                    if data:
82                        f.write(data)
83                        pbar.update(chunk_size)