diff --git a/automatic.sh b/automatic.sh index 9ae467a08..9e8d0341e 100755 --- a/automatic.sh +++ b/automatic.sh @@ -1,4 +1,4 @@ #/bin/env bash export PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.9,max_split_size_mb:512 -python launch.py --api --xformers --disable-console-progressbars +python launch.py --api --xformers --disable-console-progressbars --load_gpu diff --git a/modules/sd_models.py b/modules/sd_models.py index 76a89e88c..87ef1e652 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -170,13 +170,13 @@ def get_state_dict_from_checkpoint(pl_sd): def read_state_dict(checkpoint_file, print_global_state=False, map_location=None): _, extension = os.path.splitext(checkpoint_file) + device = map_location or shared.weight_load_location + if device is None: + device = devices.get_cuda_device_string() if torch.cuda.is_available() else "cpu" if extension.lower() == ".safetensors": - device = map_location or shared.weight_load_location - if device is None: - device = devices.get_cuda_device_string() if torch.cuda.is_available() else "cpu" pl_sd = safetensors.torch.load_file(checkpoint_file, device=device) else: - pl_sd = torch.load(checkpoint_file, map_location=map_location or shared.weight_load_location) + pl_sd = torch.load(checkpoint_file, map_location=device) if print_global_state and "global_step" in pl_sd: print(f"Global Step: {pl_sd['global_step']}") diff --git a/modules/shared.py b/modules/shared.py index a6712dae9..c00c10efa 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -89,6 +89,7 @@ parser.add_argument("--api-auth", type=str, help='Set authentication for API lik parser.add_argument("--api-log", action='store_true', help="use api-log=True to enable logging of all API requests") parser.add_argument("--nowebui", action='store_true', help="use api=True to launch the API instead of the webui") parser.add_argument("--ui-debug-mode", action='store_true', help="Don't load model to quickly launch UI") +parser.add_argument("--load_gpu", action='store_true', help="Force loading model in GPU memory") parser.add_argument("--device-id", type=str, help="Select the default CUDA device to use (export CUDA_VISIBLE_DEVICES=0,1,etc might be needed before)", default=None) parser.add_argument("--administrator", action='store_true', help="Administrator rights", default=False) parser.add_argument("--cors-allow-origins", type=str, help="Allowed CORS origin(s) in the form of a comma-separated list (no spaces)", default=None) @@ -131,7 +132,12 @@ devices.device, devices.device_interrogate, devices.device_gfpgan, devices.devic (devices.cpu if any(y in cmd_opts.use_cpu for y in [x, 'all']) else devices.get_optimal_device() for x in ['sd', 'interrogate', 'gfpgan', 'esrgan', 'codeformer']) device = devices.device -weight_load_location = None if cmd_opts.lowram else "cpu" +if cmd_opts.lowram: + weight_load_location = None +elif cmd_opts.load_gpu: + weight_load_location = device +else: + weight_load_location = "cpu" batch_cond_uncond = cmd_opts.always_batch_cond_uncond or not (cmd_opts.lowvram or cmd_opts.medvram) parallel_processing_allowed = not cmd_opts.lowvram and not cmd_opts.medvram