From 50a8e39b5b544c51bb0752a8fdd8457a69111e65 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Sat, 13 Jan 2024 00:35:26 +0200 Subject: [PATCH] Autodetect device --- nodes.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/nodes.py b/nodes.py index a1857b8..0243115 100644 --- a/nodes.py +++ b/nodes.py @@ -13,6 +13,8 @@ from .model.ccsr_stage1 import ControlLDM from .utils.common import instantiate_from_config, load_state_dict +import comfy.model_management + script_directory = os.path.dirname(os.path.abspath(__file__)) class CCSR_Upscale: @@ -52,13 +54,13 @@ class CCSR_Upscale: config = OmegaConf.load(config_path) model = instantiate_from_config(config) - + device = comfy.model_management.get_torch_device() load_state_dict(model, torch.load(checkpoint_path, map_location="cpu"), strict=True) # reload preprocess model if specified model.freeze() - model.to("cuda") + model.to(device) if (use_fp16): model.half() sampler = SpacedSampler(model, var_type="fixed_small") @@ -75,16 +77,16 @@ class CCSR_Upscale: # Resize the image tensor. resized_image = F.interpolate(image, size=(new_height, new_width), mode='bicubic', align_corners=False) - + # Move the tensor to the GPU. - resized_image = resized_image.to("cuda") + resized_image = resized_image.to(device) strength = 1.0 model.control_scales = [strength] * 13 cond_fn = None height, width = resized_image.size(-2), resized_image.size(-1) shape = (1, 4, height // 8, width // 8) x_T = torch.randn(shape, device=model.device, dtype=torch.float32) - with torch.autocast("cuda", dtype=model.dtype): + with torch.autocast(comfy.model_management.get_autocast_device(device), dtype=model.dtype): if not tiled: samples = sampler.sample_ccsr( steps=steps, t_max=t_max, t_min=t_min, shape=shape, cond_img=resized_image,