How to inference on multi gpu cards？

[kaihekaihe]

def prepare_latents(
self,
ego_prior_video, # ego prior video for latent condition
batch_size: int,
num_channels_latents: int = 16,
height: int = 480,
exo_width: int = 784,
ego_width: int = 448,
num_frames: int = 49,
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.Tensor] = None,
last_images: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
# Calculate latent dimensions for each view type
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
latent_height = height // self.vae_scale_factor_spatial
exo_latent_width = exo_width // self.vae_scale_factor_spatial
ego_latent_width = ego_width // self.vae_scale_factor_spatial
total_latent_width = exo_latent_width + ego_latent_width

    shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, total_latent_width)
    if isinstance(generator, list) and len(generator) != batch_size:
        raise ValueError(
            f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
            f" size of {batch_size}. Make sure the batch size matches the length of the generators."
        )

    # Use shared IC-LoRA latent if available
    if self._shared_latent_exo_randn is not None:
        latents = self._shared_latent_exo_randn.to(device=device, dtype=dtype)
        print(f"Using shared IC-LoRA latent: latents_shape={latents.shape}")
    elif latents is None:
        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
    else:
        latents = latents.to(device=device, dtype=dtype)
    # Process ego prior video and replace ego region in shared latent
    # Add batch dimension if missing (video_processor outputs [C, F, H, W], VAE expects [B, C, F, H, W])
    if ego_prior_video.dim() == 4:
        ego_prior_video = ego_prior_video.unsqueeze(0)  # Add batch dimension
    
    ego_prior_latents = self.encode_video(ego_prior_video.to(device=device, dtype=self.vae.dtype))
    
    # Calculate ego region boundaries
    latent_total_width = self._shared_latent_exo_randn.shape[-1] 
    ego_latent_width = ego_prior_latents.shape[-1]
    exo_latent_width = latent_total_width - ego_latent_width
    
    # Create a copy to avoid modifying the shared variable
    video_condition = self._shared_latent_exo_randn.clone()
    
    # Replace ego region (right part) with ego_prior_latents
    video_condition[:, :, :, :, exo_latent_width:] = ego_prior_latents

    # Use video_condition (already encoded latents) directly as latent_condition
    latent_condition = video_condition

    latent_condition = latent_condition.to(dtype)
    # latent_condition = (latent_condition - latents_mean) * latents_std

    # Create mask for exo and ego views
    # Use num_frames (target frame count) like the parent class, not frame_num (input frame count)
    mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, total_latent_width)
    
    # Set exo_view mask (left part) to 1.0
    mask_lat_size[:, :, :, :, :exo_latent_width] = 1.0
    # Set ego_view mask (right part) to 0.0  
    mask_lat_size[:, :, :, :, exo_latent_width:] = 0.0
    if last_images is None:
        mask_lat_size[:, :, list(range(1, num_frames))] = 0
    else:
        mask_lat_size[:, :, list(range(1, num_frames - 1))] = 0
    first_frame_mask = mask_lat_size[:, :, 0:1]
    first_frame_mask = torch.repeat_interleave(first_frame_mask, dim=2, repeats=self.vae_scale_factor_temporal) # 1 -> 4
    mask_lat_size = torch.concat([first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2) # 48 -> 48 + 4
    mask_lat_size = mask_lat_size.view(batch_size, -1, self.vae_scale_factor_temporal, latent_height, total_latent_width) # 1*52 -> 13*4
    mask_lat_size = mask_lat_size.transpose(1, 2)
    mask_lat_size = mask_lat_size.to(latent_condition.device)

    return latents, torch.concat([mask_lat_size, latent_condition], dim=1), (exo_latent_width, ego_latent_width)，

If this code is run without using --device_map balanced, it does not report an error, but later leads to GPU memory explosion:

bash scripts/multi_infer_itw.sh
Using GPUs: 1
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 14/14 [00:01<00:00, 7.90it/s]
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:01<00:00, 2.84it/s]
Loading pipeline components...: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:02<00:00, 2.85it/s]
loading lora
Traceback (most recent call last):
File "/workspace/huggingface/EgoX/infer.py", line 302, in
main(args)
File "/workspace/huggingface/EgoX/infer.py", line 94, in main
pipe.to("cuda")
File "/workspace/huggingface/EgoX/env/lib/python3.10/site-packages/diffusers/pipelines/pipeline_utils.py", line 541, in to
module.to(device, dtype)
File "/workspace/huggingface/EgoX/env/lib/python3.10/site-packages/diffusers/models/modeling_utils.py", line 1383, in to
return super().to(*args, **kwargs)
File "/workspace/huggingface/EgoX/env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1355, in to
return self._apply(convert)
File "/workspace/huggingface/EgoX/env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 915, in _apply
module._apply(fn)
File "/workspace/huggingface/EgoX/env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 915, in _apply
module._apply(fn)
File "/workspace/huggingface/EgoX/env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 915, in _apply
module._apply(fn)
[Previous line repeated 3 more times]
File "/workspace/huggingface/EgoX/env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 942, in _apply
param_applied = fn(param)
File "/workspace/huggingface/EgoX/env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1341, in convert
return t.to(
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 50.00 MiB. GPU 0 has a total capacity of 39.38 GiB of which 23.38 MiB is free. Process 1593552 has 39.35 GiB memory in use. Of the allocated memory 38.87 GiB is allocated by PyTorch, and 73.49 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
(/workspace/huggingface/EgoX/env) root@devpod2:/workspace/huggingface/EgoX# vi scripts/multi_infer_itw1.sh

However, when starting with --device_map balanced, the following occurs:

bash scripts/multi_infer_itw.sh
Using GPUs: 1
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 14/14 [00:01<00:00, 7.62it/s]
INFO:accelerate.utils.modeling:We will use 90% of the memory on device 0 for storing the model, and 10% for the buffer to avoid OOM. You can set max_memory in to a higher value to use more memory (at your own risk).
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:01<00:00, 2.76it/s]
WARNING:accelerate.big_modeling:Some parameters are on the meta device because they were offloaded to the cpu.███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:01<00:00, 2.68it/s]
Loading pipeline components...: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:03<00:00, 1.99it/s]
loading lora
./example/in_the_wild/videos/joker/exo.mp4
./example/in_the_wild/videos/joker/ego_Prior.mp4
/workspace/huggingface/EgoX/infer.py:138: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad_(True), rather than torch.tensor(sourceTensor).
camera_extrinsic = torch.cat([torch.tensor(camera_extrinsic, dtype=ego_extrinsic.dtype), torch.tensor([[0, 0, 0, 1]], dtype=ego_extrinsic.dtype)], dim=0)
/workspace/huggingface/EgoX/env/lib/python3.10/site-packages/torch/functional.py:554: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at /pytorch/aten/src/ATen/native/TensorShape.cpp:4314.)
return VF.meshgrid(tensors, **kwargs) # type: ignore[attr-defined]
/workspace/huggingface/EgoX/infer.py:211: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.detach().clone() or sourceTensor.detach().clone().requires_grad(True), rather than torch.tensor(sourceTensor).
point_map = torch.tensor(point_map) #! [F, H, W, 3]

0
Setup IC-LoRA latent: exo_shape=torch.Size([1, 16, 49, 56, 98]), ego_shape=torch.Size([1, 16, 49, 56, 56]), total_shape=torch.Size([1, 16, 49, 56, 154])
Decoded image shape: torch.Size([1, 3, 448, 1232]) for CLIP encoding
Using shared IC-LoRA latent: latents_shape=torch.Size([1, 16, 49, 56, 154])
Traceback (most recent call last):
File "/workspace/huggingface/EgoX/infer.py", line 302, in
main(args)
File "/workspace/huggingface/EgoX/infer.py", line 260, in main
video = generate_video(
File "/workspace/huggingface/EgoX/core/inference/wan.py", line 75, in generate_video
video_generate = pipe(
File "/workspace/huggingface/EgoX/env/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/workspace/huggingface/EgoX/core/finetune/models/wan_i2v/sft_trainer.py", line 785, in call
latents, condition, (exo_latent_width, ego_latent_width) = self.prepare_latents(
File "/workspace/huggingface/EgoX/core/finetune/models/wan_i2v/sft_trainer.py", line 191, in prepare_latents
return latents, torch.concat([mask_lat_size, latent_condition], dim=1), (exo_latent_width, ego_latent_width)
RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 13 but got size 49 for tensor number 1 in the list.

[KEH0T0]

Hi, thanks for the report.
Inference requires a single GPU with 80GB VRAM as noted in the README, so a 40GB GPU is unfortunately not sufficient. Also, --device_map balanced is not supported and will cause errors. Please use a single 80GB GPU (e.g., A100 80GB, H100 80GB or H200 140GB) for now.

[kaihekaihe]

Unfortunately,I have only A100 40GB, but I changed code, I can inference on 4cards with device_map balanced:

def prepare_latents(
    self,
    ego_prior_video,  # ego prior video for latent condition
    batch_size: int,
    num_channels_latents: int = 16,
    height: int = 480,
    exo_width: int = 784,
    ego_width: int = 448, 
    num_frames: int = 49,
    dtype: Optional[torch.dtype] = None,
    device: Optional[torch.device] = None,
    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
    latents: Optional[torch.Tensor] = None,
    last_images: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    # Calculate latent dimensions for each view type
    num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
    latent_height = height // self.vae_scale_factor_spatial
    exo_latent_width = exo_width // self.vae_scale_factor_spatial  
    ego_latent_width = ego_width // self.vae_scale_factor_spatial
    total_latent_width = exo_latent_width + ego_latent_width

    shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, total_latent_width)
    if isinstance(generator, list) and len(generator) != batch_size:
        raise ValueError(
            f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
            f" size of {batch_size}. Make sure the batch size matches the length of the generators."
        )

    # Use shared IC-LoRA latent if available
    if self._shared_latent_exo_randn is not None:
        latents = self._shared_latent_exo_randn.to(device=device, dtype=dtype)
        print(f"Using shared IC-LoRA latent: latents_shape={latents.shape}")
    elif latents is None:
        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
    else:
        latents = latents.to(device=device, dtype=dtype)
    # Process ego prior video and replace ego region in shared latent
    # Add batch dimension if missing (video_processor outputs [C, F, H, W], VAE expects [B, C, F, H, W])
    if ego_prior_video.dim() == 4:
        ego_prior_video = ego_prior_video.unsqueeze(0)  # Add batch dimension
    
    ego_prior_latents = self.encode_video(ego_prior_video.to(device=device, dtype=self.vae.dtype))
    target_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
    if ego_prior_latents.shape[2] != target_latent_frames:
        ego_prior_latents = ego_prior_latents[:, :, :: self.vae_scale_factor_temporal, :, :]
        if ego_prior_latents.shape[2] > target_latent_frames:
            ego_prior_latents = ego_prior_latents[:, :, :target_latent_frames, :, :]

    # Calculate ego region boundaries
    latent_total_width = self._shared_latent_exo_randn.shape[-1] 
    ego_latent_width = ego_prior_latents.shape[-1]
    exo_latent_width = latent_total_width - ego_latent_width
    
    # Create a copy to avoid modifying the shared variable
    video_condition = self._shared_latent_exo_randn.clone()
    
    # Replace ego region (right part) with ego_prior_latents
    video_condition[:, :, :, :, exo_latent_width:] = ego_prior_latents

    # Use video_condition (already encoded latents) directly as latent_condition
    latent_condition = video_condition

    latent_condition = latent_condition.to(dtype)
    # latent_condition = (latent_condition - latents_mean) * latents_std

    # Create mask for exo and ego views
    # Use num_frames (target frame count) like the parent class, not frame_num (input frame count)
    mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, total_latent_width)
    
    # Set exo_view mask (left part) to 1.0
    mask_lat_size[:, :, :, :, :exo_latent_width] = 1.0
    # Set ego_view mask (right part) to 0.0  
    mask_lat_size[:, :, :, :, exo_latent_width:] = 0.0
    if last_images is None:
        mask_lat_size[:, :, list(range(1, num_frames))] = 0
    else:
        mask_lat_size[:, :, list(range(1, num_frames - 1))] = 0
    first_frame_mask = mask_lat_size[:, :, 0:1]
    first_frame_mask = torch.repeat_interleave(first_frame_mask, dim=2, repeats=self.vae_scale_factor_temporal) # 1 -> 4
    mask_lat_size = torch.concat([first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2) # 48 -> 48 + 4
    mask_lat_size = mask_lat_size.view(batch_size, -1, self.vae_scale_factor_temporal, latent_height, total_latent_width) # 1*52 -> 13*4
    mask_lat_size = mask_lat_size.transpose(1, 2)
    mask_lat_size = mask_lat_size.to(latent_condition.device)
    return latents, torch.concat([mask_lat_size, latent_condition], dim=1), (exo_latent_width, ego_latent_width)

[kaihekaihe]

Does this have any problems