change naming
This commit is contained in:
parent
65c479b091
commit
e76bb3d284
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user