ComfyUI-IDM-VTON 降低显存的方法 (供参考)

在下面函数(ComfyUI-IDM-VTON\src\nodes\pipeline_loader.py)， def load_pipeline(self, weight_dtype):

修改两点：修改1> 把所有的 .to(DEVICE) ，全部注释掉，所有的。

修改2> 函数结尾处修改前: pipe.unet_encoder = unet_encoder pipe = pipe.to(DEVICE) pipe.weight_dtype = weight_dtype 修改为:

在显卡12G测试，完全无压力。查看显存占用大概6G多点，估计在8G下也能跑。

Jun 15 '24 03:06 peizhiluo007

def load_pipeline(self, weight_dtype):
    if weight_dtype == "float32":
        weight_dtype = torch.float32
    elif weight_dtype == "float16":
        weight_dtype = torch.float16
    elif weight_dtype == "bfloat16":
        weight_dtype = torch.bfloat16
    noise_scheduler = DDPMScheduler.from_pretrained(
        WEIGHTS_PATH, 
        subfolder="scheduler"
    )
    vae = AutoencoderKL.from_pretrained(
        WEIGHTS_PATH,
        subfolder="vae",
        torch_dtype=weight_dtype
    ).requires_grad_(False).eval()#.to(DEVICE)
    unet = UNet2DConditionModel.from_pretrained(
        WEIGHTS_PATH,
        subfolder="unet",
        torch_dtype=weight_dtype
    ).requires_grad_(False).eval()#.to(DEVICE)
    image_encoder = CLIPVisionModelWithProjection.from_pretrained(
        WEIGHTS_PATH,
        subfolder="image_encoder",
        torch_dtype=weight_dtype
    ).requires_grad_(False).eval()#.to(DEVICE)
    unet_encoder = UNet2DConditionModel_ref.from_pretrained(
        WEIGHTS_PATH,
        subfolder="unet_encoder",
        torch_dtype=weight_dtype
    ).requires_grad_(False).eval()#.to(DEVICE)
    text_encoder_one = CLIPTextModel.from_pretrained(
        WEIGHTS_PATH,
        subfolder="text_encoder",
        torch_dtype=weight_dtype
    ).requires_grad_(False).eval()#.to(DEVICE)
    text_encoder_two = CLIPTextModelWithProjection.from_pretrained(
        WEIGHTS_PATH,
        subfolder="text_encoder_2",
        torch_dtype=weight_dtype
    ).requires_grad_(False).eval()#.to(DEVICE)
    tokenizer_one = AutoTokenizer.from_pretrained(
        WEIGHTS_PATH,
        subfolder="tokenizer",
        revision=None,
        use_fast=False,
    )
    tokenizer_two = AutoTokenizer.from_pretrained(
        WEIGHTS_PATH,
        subfolder="tokenizer_2",
        revision=None,
        use_fast=False,
    )
    pipe = TryonPipeline.from_pretrained(
        WEIGHTS_PATH,
        unet=unet,
        vae=vae,
        feature_extractor=CLIPImageProcessor(),
        text_encoder=text_encoder_one,
        text_encoder_2=text_encoder_two,
        tokenizer=tokenizer_one,
        tokenizer_2=tokenizer_two,
        scheduler=noise_scheduler,
        image_encoder=image_encoder,
        torch_dtype=weight_dtype,
    )
    pipe.weight_dtype = weight_dtype
    pipe.unet_encoder = unet_encoder
    pipe.enable_sequential_cpu_offload()
    pipe.unet_encoder.to(DEVICE)
    #pipe.to(DEVICE)
    #
    return (pipe, )

Jun 15 '24 03:06 peizhiluo007

Wow that's awesome! Thanks! Could you open a PR with these changes?