Add separate model config and other small updates

README.md

@@ -49,7 +49,7 @@ python main.py --prompt "Luxury food photograph" --steps 2 --seed 2
 
 - **`--output`** (optional, `str`, default: `"image.png"`): Output image filename.
 
-- **`--model`** (optional, `str`, default: `"black-forest-labs/FLUX.1-schnell"`): Model to use for generation.
+- **`--model`** (optional, `str`, default: `"schnell"`): Model to use for generation.
 
 - **`--seed`** (optional, `int`, default: `0`): Seed for random number generation. Default is time-based.
 
@@ -67,10 +67,10 @@ import sys
 sys.path.append("/path/to/mflux/src")
 
 from flux_1_schnell.config.config import Config
-from flux_1_schnell.flux import Flux1Schnell
+from flux_1_schnell.flux import Flux1
 from flux_1_schnell.post_processing.image_util import ImageUtil
 
-flux = Flux1Schnell("black-forest-labs/FLUX.1-schnell")
+flux = Flux1.from_repo("black-forest-labs/FLUX.1-schnell")
 
 image = flux.generate_image(
    seed=3,

main.py

@@ -9,12 +9,12 @@ from flux_1_schnell.config.config import Config
 from flux_1_schnell.flux import Flux1
 from flux_1_schnell.post_processing.image_util import ImageUtil
 
+
 def main():
     parser = argparse.ArgumentParser(description='Generate an image based on a prompt.')
     parser.add_argument('--prompt', type=str, required=True, help='The textual description of the image to generate.')
     parser.add_argument('--output', type=str, default="image.png", help='The filename for the output image. Default is "image.png".')
-    parser.add_argument('--model', type=str, default="black-forest-labs/FLUX.1-schnell", help='The model to use. Default is "black-forest-labs/FLUX.1-schnell".')
-    parser.add_argument('--max_sequence_length', type=int, default=256, help='Max Sequence Length (Default is 256)')
+    parser.add_argument('--model', type=str, default="schnell", help='The model to use. Default is "schnell".')
     parser.add_argument('--seed', type=int, default=None, help='Entropy Seed (Default is time-based random-seed)')
     parser.add_argument('--height', type=int, default=1024, help='Image height (Default is 1024)')
     parser.add_argument('--width', type=int, default=1024, help='Image width (Default is 1024)')
@@ -25,7 +25,7 @@ def main():
 
     seed = int(time.time()) if args.seed is None else args.seed
 
-    flux = Flux1(args.model, max_sequence_length=args.max_sequence_length)
+    flux = Flux1.from_alias(args.model)
 
     image = flux.generate_image(
         seed=seed,

src/flux_1_schnell/config/config.py

@@ -1,28 +1,13 @@
-from dataclasses import dataclass
+import logging
+
 import mlx.core as mx
 import numpy as np
-import logging
+
+from flux_1_schnell.config.model_config import ModelConfig
 
 log = logging.getLogger(__name__)
 
 
-def get_sigmas(num_inference_steps):
-    sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
-    sigmas = mx.array(sigmas).astype(mx.float32)
-    return mx.concatenate([sigmas, mx.zeros(1)])
-
-def shift_sigmas(sigmas, width, height):
-    y1 = 0.5
-    x1 = 256
-    m = (1.15 - y1) / (4096 - x1)
-    b = y1 - m * x1
-    mu = m * width * height / 256   + b
-    shifted_sigmas = mx.exp(mu) / (mx.exp(mu) + (1 / sigmas - 1))
-    shifted_sigmas[-1] = 0
-    return shifted_sigmas
-
-
-@dataclass
 class Config:
     precision: mx.Dtype = mx.bfloat16
 
@@ -33,6 +18,7 @@ class Config:
             width: int = 1024,
             height: int = 1024,
             guidance: float = 4.0,
+            sigmas: mx.array | None = None,
     ):
         self.num_train_steps = num_train_steps
         if width % 16 != 0 or height % 16 != 0:
@@ -41,9 +27,36 @@ class Config:
         self.height = 16 * (width // 16)
         self.num_inference_steps = num_inference_steps
         self.guidance = guidance
+        self.sigmas = sigmas
 
-    def __post_init__(self, **data):
-        super().__init__(**data)
-        self.__config__.frozen = True
-    
-    
+    def copy_with_sigmas(self, model: ModelConfig) -> "Config":
+        sigmas = Config._get_sigmas(self.num_inference_steps)
+        if model == ModelConfig.FLUX1_DEV:
+            sigmas = Config._shift_sigmas(sigmas, self.width, self.height)
+
+        return Config(
+            num_train_steps=self.num_train_steps,
+            num_inference_steps=self.num_inference_steps,
+            width=self.width,
+            height=self.height,
+            guidance=self.guidance,
+            sigmas=sigmas,
+        )
+
+    @staticmethod
+    def _get_sigmas(num_inference_steps):
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+        sigmas = mx.array(sigmas).astype(mx.float32)
+        return mx.concatenate([sigmas, mx.zeros(1)])
+
+    @staticmethod
+    def _shift_sigmas(sigmas: mx.array, width: int, height: int):
+        y1 = 0.5
+        x1 = 256
+        m = (1.15 - y1) / (4096 - x1)
+        b = y1 - m * x1
+        mu = m * width * height / 256 + b
+        mu = mx.array(mu)
+        shifted_sigmas = mx.exp(mu) / (mx.exp(mu) + (1 / sigmas - 1))
+        shifted_sigmas[-1] = 0
+        return shifted_sigmas

src/flux_1_schnell/config/model_config.py

@@ -0,0 +1,29 @@
+from enum import Enum
+
+
+class ModelConfig(Enum):
+    FLUX1_DEV = ("black-forest-labs/FLUX.1-dev", "dev", 512)
+    FLUX1_SCHNELL = ("black-forest-labs/FLUX.1-schnell", "schnell", 256)
+
+    def __init__(self, model_name: str, alias: str, max_sequence_length: int):
+        self.alias = alias
+        self.model_name = model_name
+        self.max_sequence_length = max_sequence_length
+
+    @staticmethod
+    def from_repo(model_name: str) -> "ModelConfig":
+        try:
+            for model in ModelConfig:
+                if model.model_name == model_name:
+                    return model
+        except KeyError:
+            raise ValueError(f"'{model_name}' is not a valid model")
+
+    @staticmethod
+    def from_alias(alias: str) -> "ModelConfig":
+        try:
+            for model in ModelConfig:
+                if model.alias == alias:
+                    return model
+        except KeyError:
+            raise ValueError(f"'{alias}' is not a valid model")

src/flux_1_schnell/flux.py

@@ -3,8 +3,9 @@ import mlx.core as mx
 from PIL import Image
 from tqdm import tqdm
 
-from flux_1_schnell.config.config import Config, get_sigmas, shift_sigmas
+from flux_1_schnell.config.config import Config
 from flux_1_schnell.latent_creator.latent_creator import LatentCreator
+from flux_1_schnell.config.model_config import ModelConfig
 from flux_1_schnell.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
 from flux_1_schnell.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
 from flux_1_schnell.models.transformer.transformer import Transformer
@@ -18,44 +19,60 @@ from flux_1_schnell.weights.weight_handler import WeightHandler
 
 class Flux1:
 
-    def __init__(self, repo_id: str, max_sequence_length: int = 512):
-        self.is_dev = "FLUX.1-dev" in repo_id
-        tokenizers = TokenizerHandler.load_from_disk_or_huggingface(repo_id, max_sequence_length)
-        self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=max_sequence_length)
+    def __init__(self, repo_id: str):
+        self.model_config = ModelConfig.from_repo(repo_id)
+
+        # Initialize the tokenizers
+        tokenizers = TokenizerHandler.load_from_disk_or_huggingface(repo_id, self.model_config.max_sequence_length)
+        self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=self.model_config.max_sequence_length)
         self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)
 
+        # Initialize the models
         weights = WeightHandler.load_from_disk_or_huggingface(repo_id)
         self.vae = VAE(weights.vae)
         self.transformer = Transformer(weights.transformer)
         self.t5_text_encoder = T5Encoder(weights.t5_encoder)
         self.clip_text_encoder = CLIPEncoder(weights.clip_encoder)
 
+    @staticmethod
+    def from_repo(repo_id: str) -> "Flux1":
+        return Flux1(repo_id)
+
+    @staticmethod
+    def from_alias(alias: str) -> "Flux1":
+        return Flux1(ModelConfig.from_alias(alias).model_name)
+
     def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> PIL.Image.Image:
-        sigmas = get_sigmas(config.num_inference_steps)
-        if self.is_dev:
-            sigmas = shift_sigmas(sigmas, config.width, config.height)
+        # Create a new config with sigmas based on what model we are running
+        config = config.copy_with_sigmas(self.model_config)
+
+        # Create the latents
         latents = LatentCreator.create(config.height, config.width, seed)
 
+        # Embedd the prompt
         t5_tokens = self.t5_tokenizer.tokenize(prompt)
         clip_tokens = self.clip_tokenizer.tokenize(prompt)
         prompt_embeds = self.t5_text_encoder.forward(t5_tokens)
         pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens)
 
         for t in tqdm(range(config.num_inference_steps)):
+            # Predict the noise
             noise = self.transformer.predict(
                 t=t,
                 prompt_embeds=prompt_embeds,
                 pooled_prompt_embeds=pooled_prompt_embeds,
                 hidden_states=latents,
                 config=config,
-                sigmas=sigmas
             )
 
-            dt = sigmas[t + 1] - sigmas[t]
+            # Take one denoise step
+            dt = config.sigmas[t + 1] - config.sigmas[t]
             latents += noise * dt
 
+            # To enable progress tracking
             mx.eval(latents)
 
+        # Decode the latent array
         latents = Flux1._unpack_latents(latents, config.height, config.width)
         decoded = self.vae.decode(latents)
         return ImageUtil.to_image(decoded)

src/flux_1_schnell/models/transformer/guidance_embedder.py

@@ -13,4 +13,4 @@ class GuidanceEmbedder(nn.Module):
         sample = self.linear_1(sample)
         sample = nn.silu(sample)
         sample = self.linear_2(sample)
-        return sample
+        return sample

src/flux_1_schnell/models/transformer/transformer.py

@@ -16,7 +16,7 @@ class Transformer(nn.Module):
         self.pos_embed = EmbedND()
         self.x_embedder = nn.Linear(64, 3072)
         with_guidance_embed = "guidance_embedder" in weights["time_text_embed"].keys()
-        self.time_text_embed = TimeTextEmbed(with_guidance_embed = with_guidance_embed)
+        self.time_text_embed = TimeTextEmbed(with_guidance_embed=with_guidance_embed)
         self.context_embedder = nn.Linear(4096, 3072)
         self.transformer_blocks = [JointTransformerBlock(i) for i in range(19)]
         self.single_transformer_blocks = [SingleTransformerBlock(i) for i in range(38)]
@@ -33,15 +33,14 @@ class Transformer(nn.Module):
             pooled_prompt_embeds: mx.array,
             hidden_states: mx.array,
             config: Config,
-            sigmas: mx.array,
     ) -> mx.array:
-        time_step = sigmas[t] * config.num_train_steps
+        time_step = config.sigmas[t] * config.num_train_steps
         time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision)
         hidden_states = self.x_embedder(hidden_states)
         guidance = mx.broadcast_to(config.guidance * config.num_train_steps, (1,)).astype(config.precision)
         text_embeddings = self.time_text_embed.forward(time_step, pooled_prompt_embeds, guidance)
         encoder_hidden_states = self.context_embedder(prompt_embeds)
-        txt_ids = Transformer._prepare_text_ids(seq_len = prompt_embeds.shape[1])
+        txt_ids = Transformer._prepare_text_ids(seq_len=prompt_embeds.shape[1])
         img_ids = Transformer._prepare_latent_image_ids(config.height, config.width)
         ids = mx.concatenate((txt_ids, img_ids), axis=1)
         image_rotary_emb = self.pos_embed.forward(ids)
@@ -81,5 +80,5 @@ class Transformer(nn.Module):
         return latent_image_ids
 
     @staticmethod
-    def _prepare_text_ids(seq_len) -> mx.array:
+    def _prepare_text_ids(seq_len: mx.array) -> mx.array:
         return mx.zeros((1, seq_len, 3))