Using `enable_model_cpu_offload` with multiple XL pipelines triggers error Expected all tensors to be on the same device

[alexisrolland]

Hello,

I am having difficulties to use enable_model_cpu_offload in the context of my application where I am loading multiple pipelines in memory (StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline and StableDiffusionXLControlNetPipeline). I read through the documentation here: Model offloading for fast inference and memory savings and I also read through this issue which seemed similar to what I am facing: Diffusers 0.21 breaks model offloading. I couldn't find a complete code snippet that could help...

The error I am facing:

Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)

It is triggered when calling pipeline_controlnet:

import torch
from compel import Compel, ReturnedEmbeddingsType
from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained('stabilityai/stable-diffusion-xl-base-1.0', torch_dtype=torch.float16, variant='fp16', use_safetensors=True)
pipeline.enable_model_cpu_offload()

controlnet = ControlNetModel.from_pretrained('diffusers/controlnet-depth-sdxl-1.0', torch_dtype=torch.float16)
pipeline_controlnet = StableDiffusionXLControlNetPipeline(controlnet=controlnet, **pipeline.components)

compel = Compel(
    tokenizer=[pipeline.tokenizer, pipeline.tokenizer_2],
    text_encoder=[pipeline.text_encoder, pipeline.text_encoder_2],
    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
    requires_pooled=[False, True]
)

prompts = ["a cat playing with a (ball)++++ in the forest"]
negative_prompts = [""]
prompt_embeds, pooled_prompt_embeds = compel(prompts)
negative_prompt_embeds, negative_pooled_prompt_embeds = compel(negative_prompts)
seed = 123456
generators = torch.Generator(device='cuda').manual_seed(seed)

image = pipeline(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    negative_prompt_embeds=negative_prompt_embeds,
    negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
    num_inference_steps=30,
    num_images_per_prompt=1,
    generator=generators
).images[0]

torch.cuda.empty_cache()
torch.cuda.ipc_collect()

image = pipeline_controlnet(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    negative_prompt_embeds=negative_prompt_embeds,
    negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
    image=image,
    controlnet_conditioning_scale=0.5,
    num_inference_steps=30,
    num_images_per_prompt=1,
    generator=generators
).images[0]

torch.cuda.empty_cache()
torch.cuda.ipc_collect()

I have tried to add pipeline_controlnet.enable_model_cpu_offload() like this:

[...]
MODEL_PATH_CONTROLNET = os.getenv('MODEL_PATH_CONTROLNET_DEPTH')
controlnet = ControlNetModel.from_pretrained(MODEL_PATH_CONTROLNET, torch_dtype=torch.float16)
pipeline_controlnet = StableDiffusionXLControlNetPipeline(controlnet=controlnet, **pipeline.components)
pipeline_controlnet.enable_model_cpu_offload() 
[...]

But then calling the StableDiffusionXLPipeline pipeline triggers:

ValueError: Cannot generate a cpu tensor from a generator of type cuda.

@sayakpaul @patrickvonplaten I would really appreciate some guidance how to achieve this!

Thank you <3

[sayakpaul]

Hi the code snippet doesn't look minimal to me. Could you be a little more specific in helping us find out where exactly does it stem from and preferably provide a Colab Notebook?

[alexisrolland]

Hi @sayakpaul and thank you for the prompt answer!

I have simplified the code a bit in my post above. Here is a Colab to reproduce (with the error displayed at the end): https://colab.research.google.com/gist/alexisrolland/64b20e68742cd796b86a3e4ca76ae5af/untitled0.ipynb

Thanks a lot once again for looking into this... I guess I am just misunderstanding how to use enable_model_cpu_offload() when working with multiple pipelines in the same application.

[CesarERamosMedina]

Getting the same error when running inference of the GLIGEN pipeline with enable_model_cpu_offload

[patrickvonplaten]

cc @dhruvrnaik @yiyixuxu here as well

[alexisrolland]

After spending hours testing, here is what I think is happening...

After instantiating pipeline_controlnet = StableDiffusionXLControlNetPipeline(controlnet=controlnet, **pipeline.components), I verified what is the pipeline state. It says it is on CPU and to be executed on GPU. I assume because it inherits the state from the previous call pipeline.enable_model_cpu_offload()...

>>> pipeline_controlnet.device
device(type='cpu')
>>> pipeline_controlnet._execution_device
device(type='cuda', index=0)

I think this is partially correct... since StableDiffusionXLControlNetPipeline is instantiated after enable_model_cpu_offload(), the ControlNet model itself is not set to execute on the GPU. It remains on the CPU and this is what triggers the first error Expected all tensors to be on the same device

[alexisrolland]

I have done some more tests and made some progess. What I did:

• Load ControlNetModel
• Instantiate pipeline StableDiffusionXLControlNetPipeline to load model weights
• Enable model CPU offload on StableDiffusionXLControlNetPipeline
• Instantiate pipeline StableDiffusionXLPipeline from the ControlNet pipeline components
• Call StableDiffusionXLControlNetPipeline -> OK, ControlNet model is offloaded from the GPU
• Call StableDiffusionXLPipeline -> OK, model is offloaded from the GPU

Problem

• Load new/different ControlNetModel
• Instantiate pipeline StableDiffusionXLControlNetPipeline from the previous ControlNet pipeline components
• Call the new StableDiffusionXLControlNetPipeline -> KO, the pipeline runs fine but it seems the new ControlNet model is not offloaded to the CPU after the inference

Question

How to update the instance of a ControlNet pipeline with different ControlNet models without breaking the CPU offloading?

Code to reproduce

import torch
from compel import Compel, ReturnedEmbeddingsType
from diffusers import ControlNetModel
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLControlNetPipeline
from diffusers.utils import load_image

MODEL_PATH = 'stabilityai/stable-diffusion-xl-base-1.0'
MODEL_PATH_CONTROLNET_DEPTH = 'diffusers/controlnet-depth-sdxl-1.0'
MODEL_PATH_CONTROLNET_CANNY = 'diffusers/controlnet-canny-sdxl-1.0'

# Load ControlNet model and instantiate pipeline 
controlnet = ControlNetModel.from_pretrained(MODEL_PATH_CONTROLNET_DEPTH, torch_dtype=torch.float16)
pipeline_controlnet = StableDiffusionXLControlNetPipeline.from_pretrained(MODEL_PATH, controlnet=controlnet, torch_dtype=torch.float16)
pipeline_controlnet.enable_model_cpu_offload()

# Instantiate text to image pipeline from components
pipeline = StableDiffusionXLPipeline(
    scheduler=pipeline_controlnet.components['scheduler'],
    text_encoder=pipeline_controlnet.components['text_encoder'],
    text_encoder_2=pipeline_controlnet.components['text_encoder_2'],
    tokenizer=pipeline_controlnet.components['tokenizer'],
    tokenizer_2=pipeline_controlnet.components['tokenizer_2'],
    unet=pipeline_controlnet.components['unet'],
    vae=pipeline_controlnet.components['vae']
)

prompts = ["hello world"]
negative_prompts = ["ugly"]
seed = 123456
generators = torch.Generator(device='cuda').manual_seed(seed)
image = load_image("https://hf.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/hf-logo.png")

# Convert prompt to embeddings for ControlNet pipeline
compel = Compel(
    tokenizer=[pipeline_controlnet.tokenizer, pipeline_controlnet.tokenizer_2],
    text_encoder=[pipeline_controlnet.text_encoder, pipeline_controlnet.text_encoder_2],
    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
    requires_pooled=[False, True]
)
prompt_embeds, pooled_prompt_embeds = compel(prompts)
negative_prompt_embeds, negative_pooled_prompt_embeds = compel(negative_prompts)

# Call ControlNet pipeline
image = pipeline_controlnet(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    negative_prompt_embeds=negative_prompt_embeds,
    negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
    image=image,
    controlnet_conditioning_scale=0.5,
    num_inference_steps=20,
    num_images_per_prompt=1,
    generator=generators
).images[0]

torch.cuda.empty_cache()

# Convert prompt to embeddings for text to image pipeline
compel = Compel(
    tokenizer=[pipeline.tokenizer, pipeline.tokenizer_2],
    text_encoder=[pipeline.text_encoder, pipeline.text_encoder_2],
    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
    requires_pooled=[False, True]
)
prompt_embeds, pooled_prompt_embeds = compel(prompts)
negative_prompt_embeds, negative_pooled_prompt_embeds = compel(negative_prompts)

# Call text to image pipeline
image = pipeline(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    negative_prompt_embeds=negative_prompt_embeds,
    negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
    width=768,
    height=768,
    num_inference_steps=20,
    num_images_per_prompt=1,
    generator=generators
).images[0]

torch.cuda.empty_cache()

# Load second ControlNet model and instantiate pipeline from components
controlnet2 = ControlNetModel.from_pretrained(MODEL_PATH_CONTROLNET_CANNY, torch_dtype=torch.float16)
pipeline_controlnet.components['controlnet'] = controlnet2
pipeline_controlnet2 = StableDiffusionXLControlNetPipeline(**pipeline_controlnet.components)

# Convert prompt to embeddings for second ControlNet pipeline
compel = Compel(
    tokenizer=[pipeline_controlnet2.tokenizer, pipeline_controlnet2.tokenizer_2],
    text_encoder=[pipeline_controlnet2.text_encoder, pipeline_controlnet2.text_encoder_2],
    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
    requires_pooled=[False, True]
)
prompt_embeds, pooled_prompt_embeds = compel(prompts)
negative_prompt_embeds, negative_pooled_prompt_embeds = compel(negative_prompts)

# Call second ControlNet pipeline
image = pipeline_controlnet2(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    negative_prompt_embeds=negative_prompt_embeds,
    negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
    image=image,
    controlnet_conditioning_scale=0.5,
    num_inference_steps=20,
    num_images_per_prompt=1,
    generator=generators
).images[0]

torch.cuda.empty_cache()

[DN6]

@alexisrolland Since you're instantiating a new pipeline class pipeline_controlnet_2, you would need to call enable_model_cpu_offload on that new pipeline object to ensure offloading is done properly on all components. Could you check to see if that solves your issue.

[alexisrolland]

Hello @DN6 I have tried your suggestion

[...]
pipeline_controlnet2 = StableDiffusionXLControlNetPipeline(**pipeline_controlnet.components)
pipeline_controlnet2.enable_model_cpu_offload()
[..]

When calling the inference it triggers the error "ValueError: Cannot generate a cpu tensor from a generator of type cuda."

I think the reason is because enable_model_cpu_offload() is a stateful operation, so my understanding is that it should not be called more than one time in a single Python session / application. See documentation: Model offloading for fast inference and memory savings

Note: enable_model_cpu_offload() is a stateful operation that installs hooks on the models and state on the pipeline. In order to properly offload models after they are called, it is required that the entire pipeline is run and models are called in the order the pipeline expects them to be. Exercise caution if models are re-used outside the context of the pipeline after hooks have been installed. See accelerate for further docs on removing hooks.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.