Spaces:
Sleeping
Sleeping
| """Preserve codec buffer values while giving CPU storage views their own storage.""" | |
| def instantiate_cpu_model(model_type, config): | |
| import amphion_utils | |
| original_load = amphion_utils.load_model | |
| def load_codec(module, *args, **kwargs): | |
| for name, buffer in list(module.named_buffers(remove_duplicate=False)): | |
| if buffer.device.type == "cpu" and ( | |
| buffer.storage_offset() != 0 or | |
| buffer.untyped_storage().nbytes() != buffer.numel() * buffer.element_size() | |
| ): | |
| owner, _, key = name.rpartition(".") | |
| setattr(module.get_submodule(owner) if owner else module, key, buffer.clone()) | |
| return original_load(module, *args, **kwargs) | |
| # The CLI and build check construct one model at a time. Always restore the | |
| # source package's loader, including when loading raises an exception. | |
| amphion_utils.load_model = load_codec | |
| try: | |
| return model_type(config) | |
| finally: | |
| amphion_utils.load_model = original_load | |