work with tuple shape (#393)

stable_diffusion/README.md

@@ -48,8 +48,6 @@ latent_generator = sd.generate_latents("A photo of an astronaut riding a horse o
 # Here we are evaluating each diffusion step but we could also evaluate
 # once at the end.
 for x_t in latent_generator:
-    mx.simplify(x_t) # remove possible redundant computation eg reuse
-                     # scalars etc
     mx.eval(x_t)
 
 # Now x_t is the last latent from the reverse process aka x_0. We can

stable_diffusion/requirements.txt

@@ -1,4 +1,4 @@
-mlx
+mlx>=0.1
 safetensors
 huggingface-hub
 regex

stable_diffusion/stable_diffusion/__init__.py

@@ -16,21 +16,6 @@ from .model_io import (
 from .sampler import SimpleEulerSampler
 
 
-def _repeat(x, n, axis):
-    # Make the expanded shape
-    s = x.shape
-    s.insert(axis + 1, n)
-
-    # Expand
-    x = mx.broadcast_to(mx.expand_dims(x, axis + 1), s)
-
-    # Make the flattened shape
-    s.pop(axis + 1)
-    s[axis] *= n
-
-    return x.reshape(s)
-
-
 class StableDiffusion:
     def __init__(self, model: str = _DEFAULT_MODEL, float16: bool = False):
         self.dtype = mx.float16 if float16 else mx.float32
@@ -62,7 +47,7 @@ class StableDiffusion:
 
         # Repeat the conditioning for each of the generated images
         if n_images > 1:
-            conditioning = _repeat(conditioning, n_images, axis=0)
+            conditioning = mx.repeat(conditioning, n_images, axis=0)
 
         return conditioning