Spaces:
Running on Zero
Running on Zero
| # Copied from https://github.com/lllyasviel/FramePack/tree/main/demo_utils | |
| # Apache-2.0 License | |
| # By lllyasviel | |
| import torch | |
| class DynamicSwapInstaller: | |
| def _install_module(module: torch.nn.Module, **kwargs): | |
| original_class = module.__class__ | |
| module.__dict__['forge_backup_original_class'] = original_class | |
| def hacked_get_attr(self, name: str): | |
| if '_parameters' in self.__dict__: | |
| _parameters = self.__dict__['_parameters'] | |
| if name in _parameters: | |
| p = _parameters[name] | |
| if p is None: | |
| return None | |
| if p.__class__ == torch.nn.Parameter: | |
| return torch.nn.Parameter(p.to(**kwargs), requires_grad=p.requires_grad) | |
| else: | |
| return p.to(**kwargs) | |
| if '_buffers' in self.__dict__: | |
| _buffers = self.__dict__['_buffers'] | |
| if name in _buffers: | |
| return _buffers[name].to(**kwargs) | |
| return super(original_class, self).__getattr__(name) | |
| module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), { | |
| '__getattr__': hacked_get_attr, | |
| }) | |
| return | |
| def _uninstall_module(module: torch.nn.Module): | |
| if 'forge_backup_original_class' in module.__dict__: | |
| module.__class__ = module.__dict__.pop('forge_backup_original_class') | |
| return | |
| def install_model(model: torch.nn.Module, **kwargs): | |
| for m in model.modules(): | |
| DynamicSwapInstaller._install_module(m, **kwargs) | |
| return | |
| def uninstall_model(model: torch.nn.Module): | |
| for m in model.modules(): | |
| DynamicSwapInstaller._uninstall_module(m) | |
| return | |
| def get_cuda_free_memory_gb(device=None): | |
| if device is None: | |
| device = torch.device(f"cuda:{torch.cuda.current_device()}") | |
| memory_stats = torch.cuda.memory_stats(device) | |
| bytes_active = memory_stats['active_bytes.all.current'] | |
| bytes_reserved = memory_stats['reserved_bytes.all.current'] | |
| bytes_free_cuda, _ = torch.cuda.mem_get_info(device) | |
| bytes_inactive_reserved = bytes_reserved - bytes_active | |
| bytes_total_available = bytes_free_cuda + bytes_inactive_reserved | |
| return bytes_total_available / (1024 ** 3) | |