83a7bf3519
* Use magic trailing commas instead of disabling formatting to keep args on separate lines * Scope `fmt: off`s better where that's not possible
58 lines
2.1 KiB
Python
58 lines
2.1 KiB
Python
import mlx.core as mx
|
|
|
|
from mflux.config.runtime_config import RuntimeConfig
|
|
from mflux.models.vae.vae import VAE
|
|
from mflux.post_processing.array_util import ArrayUtil
|
|
from mflux.post_processing.image_util import ImageUtil
|
|
|
|
|
|
class MaskUtil:
|
|
@staticmethod
|
|
def create_masked_latents(
|
|
vae: VAE,
|
|
config: RuntimeConfig,
|
|
latents: mx.array,
|
|
img_path: str,
|
|
mask_path: str | None,
|
|
) -> mx.array:
|
|
if not img_path or not mask_path:
|
|
# Return empty latents if no image or mask is provided
|
|
return mx.zeros((1, 0, 0))
|
|
|
|
# 1. Get the reference image
|
|
scaled_image = ImageUtil.scale_to_dimensions(
|
|
image=ImageUtil.load_image(img_path).convert("RGB"),
|
|
target_width=config.width,
|
|
target_height=config.height,
|
|
)
|
|
image = ImageUtil.to_array(scaled_image)
|
|
|
|
# 2. Get the mask
|
|
scaled = ImageUtil.scale_to_dimensions(
|
|
image=ImageUtil.load_image(mask_path).convert("RGB"),
|
|
target_width=config.width,
|
|
target_height=config.height,
|
|
)
|
|
the_mask = ImageUtil.to_array(scaled, is_mask=True)
|
|
|
|
# 3. Create and pack the masked image
|
|
masked_image = image * (1 - the_mask)
|
|
masked_image = vae.encode(masked_image)
|
|
masked_image = ArrayUtil.pack_latents(latents=masked_image, height=config.height, width=config.width)
|
|
|
|
# 4. Resize mask and pack latents
|
|
mask = MaskUtil._reshape_mask(the_mask=the_mask, height=config.height, width=config.width)
|
|
mask = ArrayUtil.pack_latents(latents=mask, height=config.height, width=config.width, num_channels_latents=64)
|
|
|
|
# 5. Concat the masked_image and the mask
|
|
masked_image_latents = mx.concatenate([masked_image, mask], axis=-1)
|
|
return masked_image_latents
|
|
|
|
@staticmethod
|
|
def _reshape_mask(the_mask: mx.array, height: int, width: int) -> mx.array:
|
|
mask = the_mask[:, 0, :, :]
|
|
mask = mx.reshape(mask, (1, height // 8, 8, width // 8, 8))
|
|
mask = mx.transpose(mask, (0, 2, 4, 1, 3))
|
|
mask = mx.reshape(mask, (1, 64, height // 8, width // 8))
|
|
return mask
|