cymic/Waifu_Diffusion_Webui

1import contextlib2 3import torch4 5from modules import errors6 7# has_mps is only available in nightly pytorch (for now), `getattr` for compatibility8has_mps = getattr(torch, 'has_mps', False)9 10cpu = torch.device("cpu")11 12 13def get_optimal_device():14    if torch.cuda.is_available():15        return torch.device("cuda")16 17    if has_mps:18        return torch.device("mps")19 20    return cpu21 22 23def torch_gc():24    if torch.cuda.is_available():25        torch.cuda.empty_cache()26        torch.cuda.ipc_collect()27 28 29def enable_tf32():30    if torch.cuda.is_available():31        torch.backends.cuda.matmul.allow_tf32 = True32        torch.backends.cudnn.allow_tf32 = True33 34 35errors.run(enable_tf32, "Enabling TF32")36 37device = device_gfpgan = device_bsrgan = device_esrgan = device_scunet = device_codeformer = get_optimal_device()38dtype = torch.float1639 40def randn(seed, shape):41    # Pytorch currently doesn't handle setting randomness correctly when the metal backend is used.42    if device.type == 'mps':43        generator = torch.Generator(device=cpu)44        generator.manual_seed(seed)45        noise = torch.randn(shape, generator=generator, device=cpu).to(device)46        return noise47 48    torch.manual_seed(seed)49    return torch.randn(shape, device=device)50 51 52def randn_without_seed(shape):53    # Pytorch currently doesn't handle setting randomness correctly when the metal backend is used.54    if device.type == 'mps':55        generator = torch.Generator(device=cpu)56        noise = torch.randn(shape, generator=generator, device=cpu).to(device)57        return noise58 59    return torch.randn(shape, device=device)60 61 62def autocast():63    from modules import shared64 65    if dtype == torch.float32 or shared.cmd_opts.precision == "full":66        return contextlib.nullcontext()67 68    return torch.autocast("cuda")69