Add llms subdir + update README (#145)

* add llms subdir + update README

* nits

* use same pre-commit as mlx

* update readmes a bit

* format

stable_diffusion/stable_diffusion/__init__.py

@@ -6,12 +6,12 @@ from typing import Tuple
 import mlx.core as mx
 
 from .model_io import (
-    load_unet,
-    load_text_encoder,
+    _DEFAULT_MODEL,
     load_autoencoder,
     load_diffusion_config,
+    load_text_encoder,
     load_tokenizer,
-    _DEFAULT_MODEL,
+    load_unet,
 )
 from .sampler import SimpleEulerSampler
 

stable_diffusion/stable_diffusion/model_io.py

@@ -3,20 +3,18 @@
 import json
 from functools import partial
 
+import mlx.core as mx
 import numpy as np
 from huggingface_hub import hf_hub_download
+from mlx.utils import tree_unflatten
 from safetensors import safe_open as safetensor_open
 
-import mlx.core as mx
-from mlx.utils import tree_unflatten
-
 from .clip import CLIPTextModel
-from .config import UNetConfig, CLIPTextModelConfig, AutoencoderConfig, DiffusionConfig
+from .config import AutoencoderConfig, CLIPTextModelConfig, DiffusionConfig, UNetConfig
 from .tokenizer import Tokenizer
 from .unet import UNetModel
 from .vae import Autoencoder
 
-
 _DEFAULT_MODEL = "stabilityai/stable-diffusion-2-1-base"
 _MODELS = {
     # See https://huggingface.co/stabilityai/stable-diffusion-2-1-base for the model details and license

stable_diffusion/stable_diffusion/sampler.py

@@ -1,9 +1,9 @@
 # Copyright © 2023 Apple Inc.
 
-from .config import DiffusionConfig
-
 import mlx.core as mx
 
+from .config import DiffusionConfig
+
 
 def _linspace(a, b, num):
     x = mx.arange(0, num) / (num - 1)

stable_diffusion/txt2image.py

@@ -1,14 +1,13 @@
 # Copyright © 2023 Apple Inc.
 
 import argparse
+
+import mlx.core as mx
 from PIL import Image
 from tqdm import tqdm
 
-import mlx.core as mx
-
 from stable_diffusion import StableDiffusion
 
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description="Generate images from a textual prompt using stable diffusion"