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