From afff30fc0a4d11be4823ccce78d281a4e504c914 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 6 Mar 2023 10:50:50 -0500 Subject: [PATCH] Add --cpu to use the cpu for inference. --- comfy/model_management.py | 16 +++++++++++++++- comfy/samplers.py | 2 +- comfy/sd.py | 16 +++++++++++----- main.py | 1 + nodes.py | 22 +++++++--------------- 5 files changed, 35 insertions(+), 22 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 32159b82..4b061c32 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -31,6 +31,8 @@ try: except: pass +if "--cpu" in sys.argv: + vram_state = CPU if "--lowvram" in sys.argv: set_vram_to = LOW_VRAM if "--novram" in sys.argv: @@ -118,6 +120,8 @@ def load_model_gpu(model): def load_controlnet_gpu(models): global current_gpu_controlnets global vram_state + if vram_state == CPU: + return if vram_state == LOW_VRAM or vram_state == NO_VRAM: #don't load controlnets like this if low vram because they will be loaded right before running and unloaded right after @@ -144,10 +148,20 @@ def unload_if_low_vram(model): return model.cpu() return model +def get_torch_device(): + if vram_state == CPU: + return torch.device("cpu") + else: + return torch.cuda.current_device() + +def get_autocast_device(dev): + if hasattr(dev, 'type'): + return dev.type + return "cuda" def get_free_memory(dev=None, torch_free_too=False): if dev is None: - dev = torch.cuda.current_device() + dev = get_torch_device() if hasattr(dev, 'type') and dev.type == 'cpu': mem_free_total = psutil.virtual_memory().available diff --git a/comfy/samplers.py b/comfy/samplers.py index 3562f89d..569c32f4 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -438,7 +438,7 @@ class KSampler: else: max_denoise = True - with precision_scope(self.device): + with precision_scope(model_management.get_autocast_device(self.device)): if self.sampler == "uni_pc": samples = uni_pc.sample_unipc(self.model_wrap, noise, latent_image, sigmas, sampling_function=sampling_function, max_denoise=max_denoise, extra_args=extra_args, noise_mask=denoise_mask) elif self.sampler == "uni_pc_bh2": diff --git a/comfy/sd.py b/comfy/sd.py index eb4ea793..67a207cb 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -299,7 +299,7 @@ class CLIP: return cond class VAE: - def __init__(self, ckpt_path=None, scale_factor=0.18215, device="cuda", config=None): + def __init__(self, ckpt_path=None, scale_factor=0.18215, device=None, config=None): if config is None: #default SD1.x/SD2.x VAE parameters ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0} @@ -308,6 +308,8 @@ class VAE: self.first_stage_model = AutoencoderKL(**(config['params']), ckpt_path=ckpt_path) self.first_stage_model = self.first_stage_model.eval() self.scale_factor = scale_factor + if device is None: + device = model_management.get_torch_device() self.device = device def decode(self, samples): @@ -381,11 +383,13 @@ def resize_image_to(tensor, target_latent_tensor, batched_number): return torch.cat([tensor] * batched_number, dim=0) class ControlNet: - def __init__(self, control_model, device="cuda"): + def __init__(self, control_model, device=None): self.control_model = control_model self.cond_hint_original = None self.cond_hint = None self.strength = 1.0 + if device is None: + device = model_management.get_torch_device() self.device = device self.previous_controlnet = None @@ -406,7 +410,7 @@ class ControlNet: else: precision_scope = contextlib.nullcontext - with precision_scope(self.device): + with precision_scope(model_management.get_autocast_device(self.device)): self.control_model = model_management.load_if_low_vram(self.control_model) control = self.control_model(x=x_noisy, hint=self.cond_hint, timesteps=t, context=cond_txt) self.control_model = model_management.unload_if_low_vram(self.control_model) @@ -481,7 +485,7 @@ def load_controlnet(ckpt_path, model=None): context_dim = controlnet_data[key].shape[1] use_fp16 = False - if controlnet_data[key].dtype == torch.float16: + if model_management.should_use_fp16() and controlnet_data[key].dtype == torch.float16: use_fp16 = True control_model = cldm.ControlNet(image_size=32, @@ -527,10 +531,12 @@ def load_controlnet(ckpt_path, model=None): return control class T2IAdapter: - def __init__(self, t2i_model, channels_in, device="cuda"): + def __init__(self, t2i_model, channels_in, device=None): self.t2i_model = t2i_model self.channels_in = channels_in self.strength = 1.0 + if device is None: + device = model_management.get_torch_device() self.device = device self.previous_controlnet = None self.control_input = None diff --git a/main.py b/main.py index 43dff955..ca8674b5 100644 --- a/main.py +++ b/main.py @@ -24,6 +24,7 @@ if __name__ == "__main__": print("\t--lowvram\t\t\tSplit the unet in parts to use less vram.") print("\t--novram\t\t\tWhen lowvram isn't enough.") print() + print("\t--cpu\t\t\tTo use the CPU for everything (slow).") exit() if '--dont-upcast-attention' in sys.argv: diff --git a/nodes.py b/nodes.py index 84510a05..e5800d0d 100644 --- a/nodes.py +++ b/nodes.py @@ -628,9 +628,10 @@ class SetLatentNoiseMask: return (s,) -def common_ksampler(device, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False): +def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False): latent_image = latent["samples"] noise_mask = None + device = model_management.get_torch_device() if disable_noise: noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") @@ -646,12 +647,9 @@ def common_ksampler(device, model, seed, steps, cfg, sampler_name, scheduler, po noise_mask = noise_mask.to(device) real_model = None - if device != "cpu": - model_management.load_model_gpu(model) - real_model = model.model - else: - #TODO: cpu support - real_model = model.patch_model() + model_management.load_model_gpu(model) + real_model = model.model + noise = noise.to(device) latent_image = latent_image.to(device) @@ -697,9 +695,6 @@ def common_ksampler(device, model, seed, steps, cfg, sampler_name, scheduler, po return (out, ) class KSampler: - def __init__(self, device="cuda"): - self.device = device - @classmethod def INPUT_TYPES(s): return {"required": @@ -721,12 +716,9 @@ class KSampler: CATEGORY = "sampling" def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0): - return common_ksampler(self.device, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise) + return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise) class KSamplerAdvanced: - def __init__(self, device="cuda"): - self.device = device - @classmethod def INPUT_TYPES(s): return {"required": @@ -757,7 +749,7 @@ class KSamplerAdvanced: disable_noise = False if add_noise == "disable": disable_noise = True - return common_ksampler(self.device, model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise) + return common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise) class SaveImage: def __init__(self):