Spaces:
Running on Zero
Running on Zero
| from accelerate import Accelerator | |
| from diffusers.utils.torch_utils import is_compiled_module | |
| import torch | |
| global_accelerator = None | |
| def get_accelerator() -> Accelerator: | |
| global global_accelerator | |
| if global_accelerator is None: | |
| global_accelerator = Accelerator() | |
| elif torch.cuda.is_available() and global_accelerator.device.type == 'cpu': | |
| # Re-initialize if Accelerator was first created in a ZeroGPU CPU parent process | |
| # before @spaces.GPU dynamically attached the CUDA device. | |
| global_accelerator = Accelerator() | |
| return global_accelerator | |
| def reset_accelerator(): | |
| global global_accelerator | |
| global_accelerator = None | |
| def unwrap_model(model): | |
| try: | |
| accelerator = get_accelerator() | |
| model = accelerator.unwrap_model(model) | |
| model = model._orig_mod if is_compiled_module(model) else model | |
| except Exception as e: | |
| pass | |
| return model | |