1
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-02 23:14:49 +08:00

Properly handle and reshape masks when used on 3d latents.

This commit is contained in:
comfyanonymous
2024-11-09 07:10:43 -05:00
parent 6ee066a14f
commit 8b90e50979
2 changed files with 23 additions and 6 deletions

View File

@@ -1,14 +1,10 @@
import torch
import comfy.model_management
import comfy.conds
import comfy.utils
def prepare_mask(noise_mask, shape, device):
"""ensures noise mask is of proper dimensions"""
noise_mask = torch.nn.functional.interpolate(noise_mask.reshape((-1, 1, noise_mask.shape[-2], noise_mask.shape[-1])), size=(shape[2], shape[3]), mode="bilinear")
noise_mask = torch.cat([noise_mask] * shape[1], dim=1)
noise_mask = comfy.utils.repeat_to_batch_size(noise_mask, shape[0])
noise_mask = noise_mask.to(device)
return noise_mask
return comfy.utils.reshape_mask(noise_mask, shape).to(device)
def get_models_from_cond(cond, model_type):
models = []