huggingface / huggingface/diffusers

Lower VRAM usage in CPU offload for Flux ControlNet Pipeline

Open
#10,790 15 comments 0 reactions 0 assignees View on GitHub
stale
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

**Is your feature request related to a problem? Please describe.**

I have a 24GB VRAM GPU. When running a diffusion model like Flux1, I can barely fit the model in memory during inference with batch size 1. Enabling CPU offload does not help because the offload does not occur between the controlnet forward pass and the transformer foward pass (which makes sense perfromance-wise).

I would be great to enable offloading between controlnet call and transformer denoising steps (or any other auxiliary model that does not currently get offloaded in the middle of the denoising process) to further reduce VRAM requirements.

**Describe the solution you'd like.**

What I would suggest is having a "slow" offload mode where the models do get offloaded to CPU, even if it is really slow.

```
def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Union[torch.device, str] = "cuda", enable_slow_mode: bool = False)
...
```

For instance, in the [image-to-image pipeline](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py#L927) on line 927:

```python
controlnet_block_samples, controlnet_single_block_samples = self.controlnet(
....
)

if self._enable_slow_cpu_offload:
self.maybe_free_model_hooks()

...

noise_pred = self.transformer(
...
)[0]

if self._enable_slow_cpu_offload:
self.maybe_free_model_hooks()
```

**Describe alternatives you've considered.**

I am not sure there are alternatives if the usage of these models is to be allowed at the desired fp precision (in my case bfloat16).

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.