ComfyBootlegOffload.py

# Force model to always use specified device
# Place in `ComfyUI\custom_nodes` to use
#  City96 [Apache2]
#
import types
import torch
import comfy.model_management

class OverrideDevice:
    @classmethod
    def INPUT_TYPES(s):
        devices = ["cpu",]
        for k in range(0, torch.cuda.device_count()):
            devices.append(f"cuda:{k}")

        return {
            "required": {
                "device": (devices, {"default":"cpu"}),
            }
        }

    FUNCTION = "patch"
    CATEGORY = "other"

    def override(self, model, model_attr, device):
        # set model/patcher attributes
        model.device = device
        patcher = getattr(model, "patcher", model) #.clone()
        for name in ["device", "load_device", "offload_device", "current_device", "output_device"]:
            setattr(patcher, name, device)

        # move model to device
        py_model = getattr(model, model_attr)
        py_model.to = types.MethodType(torch.nn.Module.to, py_model)
        py_model.to(device)

        # remove ability to move model
        def to(*args, **kwargs):
            pass
        py_model.to = types.MethodType(to, py_model)
        return (model,)

    def patch(self, *args, **kwargs):
        raise NotImplementedError

class OverrideCLIPDevice(OverrideDevice):
    @classmethod
    def INPUT_TYPES(s):
        k = super().INPUT_TYPES()
        k["required"]["clip"] = ("CLIP",)
        return k

    RETURN_TYPES = ("CLIP",)
    TITLE = "Force/Set CLIP Device"

    def patch(self, clip, device):
        return self.override(clip, "cond_stage_model", torch.device(device))

class OverrideVAEDevice(OverrideDevice):
    @classmethod
    def INPUT_TYPES(s):
        k = super().INPUT_TYPES()
        k["required"]["vae"] = ("VAE",)
        return k

    RETURN_TYPES = ("VAE",)
    TITLE = "Force/Set VAE Device"

    def patch(self, vae, device):
        return self.override(vae, "first_stage_model", torch.device(device))


NODE_CLASS_MAPPINGS = {
    "OverrideCLIPDevice": OverrideCLIPDevice,
    "OverrideVAEDevice": OverrideVAEDevice,
}
NODE_DISPLAY_NAME_MAPPINGS = {k:v.TITLE for k,v in NODE_CLASS_MAPPINGS.items()}

Yes, with this extension.. https://github.com/neuratech-ai/ComfyUI-MultiGPU Here is a picture of all the nodes.. It's awesome. [image: image.png]

…

On Mon, Aug 12, 2024 at 6:57 PM kobechenyang ***@***.***> wrote: ***@***.**** commented on this gist. ------------------------------ Is it possible to offload loading controlnet model to a second gpu? — Reply to this email directly, view it on GitHub <https://gist.github.com/city96/30743dfdfe129b331b5676a79c3a8a39#gistcomment-5152194> or unsubscribe <https://github.com/notifications/unsubscribe-auth/AEVP2RFOCTF6X5GQKOQJUFLZRE4V3BFKMF2HI4TJMJ2XIZLTSKBKK5TBNR2WLJDUOJ2WLJDOMFWWLO3UNBZGKYLEL5YGC4TUNFRWS4DBNZ2F6YLDORUXM2LUPGBKK5TBNR2WLJDHNFZXJJDOMFWWLK3UNBZGKYLEL52HS4DFVRZXKYTKMVRXIX3UPFYGLK2HNFZXIQ3PNVWWK3TUUZ2G64DJMNZZDAVEOR4XAZNEM5UXG5FFOZQWY5LFVEYTGMJYGA2DAMRVU52HE2LHM5SXFJTDOJSWC5DF> . You are receiving this email because you commented on the thread. Triage notifications on the go with GitHub Mobile for iOS <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675> or Android <https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub> .

Could you add storage to the mix as well to offload to? As i want to offload the clip device to storage, and keep it there, so it doesn't delete and rewrite the entire cache for every generation/text change, And if possible use some form of directIO/directstorage also being multithreaded to read sequentially from the storage, 500MB's from sata, or in my case, 6.5GB's from my gen 4 nvme. That way it can read from the written cache as fast as the nvme is capable of.