diff --git a/__init__.py b/__init__.py index 92494e4..aa46afc 100644 --- a/__init__.py +++ b/__init__.py @@ -9,6 +9,7 @@ import folder_paths node_list = [ "marnodes", + "img2dev", ] NODE_CLASS_MAPPINGS = {} diff --git a/nodes/img2dev.py b/nodes/img2dev.py new file mode 100644 index 0000000..c987b74 --- /dev/null +++ b/nodes/img2dev.py @@ -0,0 +1,46 @@ +import torch +import comfy.sd + +current_device = "cuda:0" + + +def get_torch_device_patched(): + global current_device + if ( + not torch.cuda.is_available() + or comfy.model_management.cpu_state == comfy.model_management.CPUState.CPU + ): + return torch.device("cpu") + + return torch.device(current_device) + + +class ImageToDevice: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "device": ([f"cuda:{i}" for i in range(torch.cuda.device_count())],), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "execute" + CATEGORY = "marduk191/image utils" + + def execute(self, image, device): + image = image.clone().to(device) + torch.cuda.empty_cache() + + return (image,) + +NODE_CLASS_MAPPINGS = { + # Image analysis + "ImageToDevice": ImageToDevice, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + # Image analysis + "ImageToDevice": "marduk191's Image Sender", +} \ No newline at end of file