fluxdev/stable-diffusion-webui-forge

1# https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/148552 3import torch4 5from ldm_patched.modules import args_parser6from ldm_patched.modules import model_management7 8 9def stream_context():10    if torch.cuda.is_available():11        return torch.cuda.stream12 13    if model_management.is_intel_xpu():14        return torch.xpu.stream15 16    return None17 18 19def get_current_stream():20    try:21        if torch.cuda.is_available():22            device = torch.device(torch.cuda.current_device())23            stream = torch.cuda.current_stream(device)24            with torch.cuda.stream(stream):25                torch.zeros((1, 1)).to(device, torch.float32)26            stream.synchronize()27            return stream28        if model_management.is_intel_xpu():29            device = torch.device("xpu")30            stream = torch.xpu.current_stream(device)31            with torch.xpu.stream(stream):32                torch.zeros((1, 1)).to(device, torch.float32)33            stream.synchronize()34            return stream35    except:36        return None37 38 39def get_new_stream():40    try:41        if torch.cuda.is_available():42            device = torch.device(torch.cuda.current_device())43            stream = torch.cuda.Stream(device)44            with torch.cuda.stream(stream):45                torch.zeros((1, 1)).to(device, torch.float32)46            stream.synchronize()47            return stream48        if model_management.is_intel_xpu():49            device = torch.device("xpu")50            stream = torch.xpu.Stream(device)51            with torch.xpu.stream(stream):52                torch.zeros((1, 1)).to(device, torch.float32)53            stream.synchronize()54            return stream55    except:56        return None57 58 59current_stream = None60mover_stream = None61using_stream = False62 63if args_parser.args.cuda_stream:64    current_stream = get_current_stream()65    mover_stream = get_new_stream()66    using_stream = current_stream is not None and mover_stream is not None67 68