monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #162 from akx/dead-maskutil

Browse Source 
Remove duplicate mask_util file
This commit is contained in:
Filip Strand 2025-03-31 00:03:53 +02:00 committed by GitHub
commit 941d010436
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 0 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

54
src/mflux/post_processing/mask_util.py
View File

@ -1,54 +0,0 @@
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
class MaskUtil:
@staticmethod
def create_masked_latents(
vae: VAE,
config: RuntimeConfig,
latents: mx.array,
img_path: str,
mask_path: str | None
) -> mx.array: # fmt: off
from mflux import ImageUtil
# 1. Get the reference image
scaled_image = ImageUtil.scale_to_dimensions(
image=ImageUtil.load_image(config.image_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):
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