Spaces:
Running on Zero
Running on Zero
File size: 929 Bytes
855e1f5 | 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 | 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
|