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

Fix height and width argument names

Browse Source
This commit is contained in:
filipstrand 2024-08-18 21:49:41 +02:00
parent 826b58dcb1
commit c50a1fb3d4
2 changed files with 9 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/flux_1_schnell/flux.py
View File

@ -61,10 +61,10 @@ class Flux1Schnell:
return ImageUtil.to_image(decoded)
@staticmethod
def _unpack_latents(latents: mx.array, width: int, height: int) -> mx.array:
latents = mx.reshape(latents, (1, height//16, width//16, 16, 2, 2))
def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array:
latents = mx.reshape(latents, (1, width // 16, height // 16, 16, 2, 2))
latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5))
latents = mx.reshape(latents, (1, 16, height//16 *2, width//16 * 2))
latents = mx.reshape(latents, (1, 16, width // 16 * 2, height // 16 * 2))
return latents
def encode(self, path: str) -> mx.array:

12
src/flux_1_schnell/models/transformer/transformer.py
View File

@ -67,14 +67,14 @@ class Transformer(nn.Module):
return noise
@staticmethod
def _prepare_latent_image_ids(width: int, height: int) -> mx.array:
latent_height = height // 16
def _prepare_latent_image_ids(height: int, width: int) -> mx.array:
latent_width = width // 16
latent_image_ids = mx.zeros((latent_height, latent_width, 3))
latent_image_ids = latent_image_ids.at[:, :, 1].add(mx.arange(0, latent_height)[:, None])
latent_image_ids = latent_image_ids.at[:, :, 2].add(mx.arange(0, latent_width)[None, :])
latent_height = height // 16
latent_image_ids = mx.zeros((latent_width, latent_height, 3))
latent_image_ids = latent_image_ids.at[:, :, 1].add(mx.arange(0, latent_width)[:, None])
latent_image_ids = latent_image_ids.at[:, :, 2].add(mx.arange(0, latent_height)[None, :])
latent_image_ids = mx.repeat(latent_image_ids[None, :], 1, axis=0)
latent_image_ids = mx.reshape(latent_image_ids, (1, latent_height * latent_width, 3))
latent_image_ids = mx.reshape(latent_image_ids, (1, latent_width * latent_height, 3))
return latent_image_ids
@staticmethod