WanImageToVideoPipeline - swap out a limited number of blocks

[spezialspezial]

I can fit WanImageToVideoPipeline on a 24GB card but it does scrape the ceiling and is a bit too close for comfort to OOMing at some random system event.

The kijai/ComfyUI-WanVideoWrapper has a nice option to swap a limited and user defined number of blocks out of VRAM. Can a similar thing be done right now with sequential_cpu_offload? If not I would like to request something along these lines to shave off 2-4 GB.

I'm open to other ideas for a small VRAM reduction.

[a-r-r-o-w]

Check apply_group_offloading related docs. It'll be going into current release soon and we can work on improving discoverability/docs. It has the minimal VRAM requirements without much overhead to generation time if you're on a modern cuda gpu. Currently, the RAM requirements are high but it's a WIP to improve it (happy to accept any improvement PRs 🤗)

Some numbers: #10847 (comment)

If you combine it with #10623 and precompute text embeddings and do tiled VAE decode, you can run in under ~7-10 GB.

[spezialspezial]

Superbe, I'll check it out

[rolux]

@a-r-r-o-w: Thanks for the hints!

I did some tests on Google Colab, and the following code will saturate the System RAM (83.5 GB) and then hang.

pipe = WanImageToVideoPipeline.from_pretrained(
    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
)
pipe.transformer.enable_group_offload(
    onload_device=torch.device("cuda"),
    offload_device=torch.device("cpu"),
    offload_type="leaf_level",
    use_stream=True
)

Meanwhile, pipe.vae.enable_tiling() throws an error:

AttributeError: 'AutoencoderKLWan' object has no attribute 'enable_tiling'

[a-r-r-o-w]

Group offloading with streams currently has a significant limitation in that it pins weight tensors on the CPU -- makes it require a lot more RAM than other methods. Could you try without streams, and also with offload_type="block_level"? It should probably lower it based on my finding here We're working on optimizing it/figuring out how to improve it.

Tiling support hasn't been added to Wan yet I believe (cc @yiyixuxu).

[rolux]

@a-r-r-o-w: Thanks for the info!

The following code currently throws, see below.

code

!pip install git+https://github.com/huggingface/diffusers.git
!pip install ftfy


import os
import torch
from diffusers import (
    WanImageToVideoPipeline,
    WanTransformer3DModel
)
from diffusers.hooks import apply_group_offloading
from diffusers.utils import export_to_video
from PIL import Image


onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
model_id = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
pipe = WanImageToVideoPipeline.from_pretrained(model_id)

pipe.transformer.enable_group_offload(
    onload_device=onload_device,
    offload_device=offload_device,
    offload_type="block_level",
    num_blocks_per_group=2
)
apply_group_offloading(
    pipe.text_encoder,
    onload_device=onload_device,
    offload_type="block_level",
    num_blocks_per_group=2
)
apply_group_offloading(
    pipe.vae,
    onload_device=onload_device,
    offload_type="block_level",
    num_blocks_per_group=2
)


def render(
    filename,
    image,
    prompt,
    seed=0,
    width=832,
    height=480,
    num_frames=81,
    num_inference_steps=30,
    guidance_scale=5.0,
    fps=16
):
    video = pipe(
        image=image,
        prompt=prompt,
        generator=torch.Generator(device=pipe.device).manual_seed(seed),
        width=width,
        height=height,
        num_frames=num_frames,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale
    ).frames[0]
    os.makedirs(os.path.dirname(filename), exist_ok=True)
    export_to_video(video, filename, fps=fps)


render(
    filename="/content/test.mp4",
    image=Image.open("/content/test.png"),
    prompt="a woman in a yellow coat is dancing in the desert",
    seed=42
)

RuntimeError                              Traceback (most recent call last)
[<ipython-input-3-41f8c8c49ba8>](https://localhost:8080/#) in <cell line: 0>()
     72 
     73 
---> 74 render(
     75     filename="/content/test.mp4",
     76     image=Image.open("/content/test.png"),

16 frames
[<ipython-input-3-41f8c8c49ba8>](https://localhost:8080/#) in render(filename, image, prompt, seed, width, height, num_frames, num_inference_steps, guidance_scale, fps)
     58     fps=16
     59 ):
---> 60     video = pipe(
     61         image=image,
     62         prompt=prompt,

[/usr/local/lib/python3.11/dist-packages/torch/utils/_contextlib.py](https://localhost:8080/#) in decorate_context(*args, **kwargs)
    114     def decorate_context(*args, **kwargs):
    115         with ctx_factory():
--> 116             return func(*args, **kwargs)
    117 
    118     return decorate_context

[/usr/local/lib/python3.11/dist-packages/diffusers/pipelines/wan/pipeline_wan_i2v.py](https://localhost:8080/#) in __call__(self, image, prompt, negative_prompt, height, width, num_frames, num_inference_steps, guidance_scale, num_videos_per_prompt, generator, latents, prompt_embeds, negative_prompt_embeds, output_type, return_dict, attention_kwargs, callback_on_step_end, callback_on_step_end_tensor_inputs, max_sequence_length)
    585             negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype)
    586 
--> 587         image_embeds = self.encode_image(image)
    588         image_embeds = image_embeds.repeat(batch_size, 1, 1)
    589         image_embeds = image_embeds.to(transformer_dtype)

[/usr/local/lib/python3.11/dist-packages/diffusers/pipelines/wan/pipeline_wan_i2v.py](https://localhost:8080/#) in encode_image(self, image)
    222     def encode_image(self, image: PipelineImageInput):
    223         image = self.image_processor(images=image, return_tensors="pt").to(self.device)
--> 224         image_embeds = self.image_encoder(**image, output_hidden_states=True)
    225         return image_embeds.hidden_states[-2]
    226 

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1734             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1735         else:
-> 1736             return self._call_impl(*args, **kwargs)
   1737 
   1738     # torchrec tests the code consistency with the following code

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1745                 or _global_backward_pre_hooks or _global_backward_hooks
   1746                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1747             return forward_call(*args, **kwargs)
   1748 
   1749         result = None

[/usr/local/lib/python3.11/dist-packages/transformers/models/clip/modeling_clip.py](https://localhost:8080/#) in forward(self, pixel_values, output_attentions, output_hidden_states, interpolate_pos_encoding, return_dict)
   1552         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
   1553 
-> 1554         vision_outputs = self.vision_model(
   1555             pixel_values=pixel_values,
   1556             output_attentions=output_attentions,

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1734             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1735         else:
-> 1736             return self._call_impl(*args, **kwargs)
   1737 
   1738     # torchrec tests the code consistency with the following code

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1745                 or _global_backward_pre_hooks or _global_backward_hooks
   1746                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1747             return forward_call(*args, **kwargs)
   1748 
   1749         result = None

[/usr/local/lib/python3.11/dist-packages/transformers/models/clip/modeling_clip.py](https://localhost:8080/#) in forward(self, pixel_values, output_attentions, output_hidden_states, return_dict, interpolate_pos_encoding)
   1091             raise ValueError("You have to specify pixel_values")
   1092 
-> 1093         hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
   1094         hidden_states = self.pre_layrnorm(hidden_states)
   1095 

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1734             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1735         else:
-> 1736             return self._call_impl(*args, **kwargs)
   1737 
   1738     # torchrec tests the code consistency with the following code

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1745                 or _global_backward_pre_hooks or _global_backward_hooks
   1746                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1747             return forward_call(*args, **kwargs)
   1748 
   1749         result = None

[/usr/local/lib/python3.11/dist-packages/transformers/models/clip/modeling_clip.py](https://localhost:8080/#) in forward(self, pixel_values, interpolate_pos_encoding)
    246             )
    247         target_dtype = self.patch_embedding.weight.dtype
--> 248         patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))  # shape = [*, width, grid, grid]
    249         patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
    250 

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1734             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1735         else:
-> 1736             return self._call_impl(*args, **kwargs)
   1737 
   1738     # torchrec tests the code consistency with the following code

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1745                 or _global_backward_pre_hooks or _global_backward_hooks
   1746                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1747             return forward_call(*args, **kwargs)
   1748 
   1749         result = None

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/conv.py](https://localhost:8080/#) in forward(self, input)
    552 
    553     def forward(self, input: Tensor) -> Tensor:
--> 554         return self._conv_forward(input, self.weight, self.bias)
    555 
    556 

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/conv.py](https://localhost:8080/#) in _conv_forward(self, input, weight, bias)
    547                 self.groups,
    548             )
--> 549         return F.conv2d(
    550             input, weight, bias, self.stride, self.padding, self.dilation, self.groups
    551         )

RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

[rolux]

If I comment out the two instances of apply_group_offloading (t5 + vae), it throws here:

RuntimeError                              Traceback (most recent call last)
[<ipython-input-1-edbc9a66ffc4>](https://localhost:8080/#) in <cell line: 0>()
     73 
     74 
---> 75 render(
     76     filename="/content/test.mp4",
     77     image=Image.open("/content/test.png"),

14 frames
[<ipython-input-1-edbc9a66ffc4>](https://localhost:8080/#) in render(filename, image, prompt, seed, width, height, num_frames, num_inference_steps, guidance_scale, fps)
     59     fps=16
     60 ):
---> 61     video = pipe(
     62         image=image,
     63         prompt=prompt,

[/usr/local/lib/python3.11/dist-packages/torch/utils/_contextlib.py](https://localhost:8080/#) in decorate_context(*args, **kwargs)
    114     def decorate_context(*args, **kwargs):
    115         with ctx_factory():
--> 116             return func(*args, **kwargs)
    117 
    118     return decorate_context

[/usr/local/lib/python3.11/dist-packages/diffusers/pipelines/wan/pipeline_wan_i2v.py](https://localhost:8080/#) in __call__(self, image, prompt, negative_prompt, height, width, num_frames, num_inference_steps, guidance_scale, num_videos_per_prompt, generator, latents, prompt_embeds, negative_prompt_embeds, output_type, return_dict, attention_kwargs, callback_on_step_end, callback_on_step_end_tensor_inputs, max_sequence_length)
    568 
    569         # 3. Encode input prompt
--> 570         prompt_embeds, negative_prompt_embeds = self.encode_prompt(
    571             prompt=prompt,
    572             negative_prompt=negative_prompt,

[/usr/local/lib/python3.11/dist-packages/diffusers/pipelines/wan/pipeline_wan_i2v.py](https://localhost:8080/#) in encode_prompt(self, prompt, negative_prompt, do_classifier_free_guidance, num_videos_per_prompt, prompt_embeds, negative_prompt_embeds, max_sequence_length, device, dtype)
    273 
    274         if prompt_embeds is None:
--> 275             prompt_embeds = self._get_t5_prompt_embeds(
    276                 prompt=prompt,
    277                 num_videos_per_prompt=num_videos_per_prompt,

[/usr/local/lib/python3.11/dist-packages/diffusers/pipelines/wan/pipeline_wan_i2v.py](https://localhost:8080/#) in _get_t5_prompt_embeds(self, prompt, num_videos_per_prompt, max_sequence_length, device, dtype)
    206         seq_lens = mask.gt(0).sum(dim=1).long()
    207 
--> 208         prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state
    209         prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
    210         prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1734             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1735         else:
-> 1736             return self._call_impl(*args, **kwargs)
   1737 
   1738     # torchrec tests the code consistency with the following code

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1745                 or _global_backward_pre_hooks or _global_backward_hooks
   1746                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1747             return forward_call(*args, **kwargs)
   1748 
   1749         result = None

[/usr/local/lib/python3.11/dist-packages/transformers/models/umt5/modeling_umt5.py](https://localhost:8080/#) in forward(self, input_ids, attention_mask, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict)
   1607         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
   1608 
-> 1609         encoder_outputs = self.encoder(
   1610             input_ids=input_ids,
   1611             attention_mask=attention_mask,

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1734             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1735         else:
-> 1736             return self._call_impl(*args, **kwargs)
   1737 
   1738     # torchrec tests the code consistency with the following code

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1745                 or _global_backward_pre_hooks or _global_backward_hooks
   1746                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1747             return forward_call(*args, **kwargs)
   1748 
   1749         result = None

[/usr/local/lib/python3.11/dist-packages/transformers/models/umt5/modeling_umt5.py](https://localhost:8080/#) in forward(self, input_ids, attention_mask, encoder_hidden_states, encoder_attention_mask, inputs_embeds, head_mask, cross_attn_head_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
    686             if self.embed_tokens is None:
    687                 raise ValueError("You have to initialize the model with valid token embeddings")
--> 688             inputs_embeds = self.embed_tokens(input_ids)
    689 
    690         batch_size, seq_length = input_shape

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1734             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1735         else:
-> 1736             return self._call_impl(*args, **kwargs)
   1737 
   1738     # torchrec tests the code consistency with the following code

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1745                 or _global_backward_pre_hooks or _global_backward_hooks
   1746                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1747             return forward_call(*args, **kwargs)
   1748 
   1749         result = None

[/usr/local/lib/python3.11/dist-packages/torch/nn/modules/sparse.py](https://localhost:8080/#) in forward(self, input)
    188 
    189     def forward(self, input: Tensor) -> Tensor:
--> 190         return F.embedding(
    191             input,
    192             self.weight,

[/usr/local/lib/python3.11/dist-packages/torch/nn/functional.py](https://localhost:8080/#) in embedding(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)
   2549         # remove once script supports set_grad_enabled
   2550         _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)
-> 2551     return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
   2552 
   2553 

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

[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.

[a-r-r-o-w]

@rolux There have been some updates to fix the behaviour of block_level offloading since you last tried it out. I haven't tested the code snippet here again, but I think it should be resolved now. If you still face an error, LMK and I'll take a look