change naming

This commit is contained in:
Fabio 2024-08-17 22:50:35 +02:00
parent 65c479b091
commit e76bb3d284
2 changed files with 10 additions and 10 deletions

View File

@ -5,7 +5,7 @@ class LatentCreator:
@staticmethod
def create(height: int, width: int, seed: int) -> (mx.array, mx.array):
r_h = height // 16
r_w = width // 16
latents = mx.random.normal(shape=[1, r_h * r_w, 64], key=mx.random.key(seed))
latent_height = height // 16
latent_width = width // 16
latents = mx.random.normal(shape=[1, latent_height * latent_width, 64], key=mx.random.key(seed))
return latents

View File

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