Merge pull request #10 from Xuzzo/feature/add_flux1dev

Include FLUX.1-Dev

README.md

@@ -20,7 +20,7 @@ like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) f
 ### Models
 
 - [x] FLUX.1-Scnhell
-- [ ] FLUX.1-Dev
+- [x] FLUX.1-Dev
 
 ### Installation
 1. Clone the repo:
@@ -37,19 +37,27 @@ like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) f
    ```
 ### Generating an image
 
-Run the provided [main.py](main.py) by specifying a prompt and some optional arguments like so:
+Run the provided [main.py](main.py) by specifying a prompt and some optional arguments like so using the default `Schnell` model:
 
 ```
 python main.py --prompt "Luxury food photograph" --steps 2 --seed 2
 ```
 
+or use the slower, but more powerful `Dev` model and run it with more time steps:
+
+```
+python main.py --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
+```
+
+⚠️ *If the specific model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell or Dev model respectively).* ⚠️
+
 #### Full list of Command-Line Arguments 
 
 - **`--prompt`** (required, `str`): Text description of the image to generate.
 
 - **`--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 (`"schnell"` or `"dev"`).
 
 - **`--seed`** (optional, `int`, default: `0`): Seed for random number generation. Default is time-based.
 
@@ -67,10 +75,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,
@@ -85,8 +93,6 @@ image = flux.generate_image(
 ImageUtil.save_image(image, "image.png")
 ```
 
-If the model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell model).
-
 
 ### Image generation speed (updated)
 
@@ -169,5 +175,4 @@ Luxury food photograph of an italian Linguine pasta alle vongole dish with lots
 
 ### TODO
 
-- FLUX Dev implementation
 - LoRA adapters

main.py

@@ -6,7 +6,7 @@ import time
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '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
 
 
@@ -14,17 +14,18 @@ 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('--seed', type=int, default=0, help='Entropy Seed (Default is time-based random-seed)')
+    parser.add_argument('--model', type=str, default="schnell", help='The model to use ("schnell" or "dev"). 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)')
     parser.add_argument('--steps', type=int, default=4, help='Inference Steps')
+    parser.add_argument('--guidance', type=float, default=3.5, help='Guidance Scale (Default is 3.5)')
 
     args = parser.parse_args()
 
-    seed = args.seed or int(time.time())
+    seed = int(time.time()) if args.seed is None else args.seed
 
-    flux = Flux1Schnell(args.model)
+    flux = Flux1.from_alias(args.model)
 
     image = flux.generate_image(
         seed=seed,
@@ -33,6 +34,7 @@ def main():
             num_inference_steps=args.steps,
             height=args.height,
             width=args.width,
+            guidance=args.guidance,
         )
     )
 

src/flux_1_schnell/config/config.py

@@ -1,31 +1,23 @@
-import mlx.core as mx
-import numpy as np
 import logging
 
+import mlx.core as mx
+
 log = logging.getLogger(__name__)
 
 
 class Config:
-    precision: mx.Dtype = mx.float16
-    num_train_steps = 1000
+    precision: mx.Dtype = mx.bfloat16
 
     def __init__(
             self,
             num_inference_steps: int = 4,
             width: int = 1024,
             height: int = 1024,
+            guidance: float = 4.0,
     ):
         if width % 16 != 0 or height % 16 != 0:
             log.warning("Width and height should be multiples of 16. Rounding down.")
         self.width = 16 * (height // 16)
         self.height = 16 * (width // 16)
-        base_sigmas = Config.base_sigmas(num_inference_steps)
         self.num_inference_steps = num_inference_steps
-        self.time_steps = base_sigmas * self.num_train_steps
-        self.sigmas = mx.concatenate([base_sigmas, mx.zeros(1)])
-
-    @staticmethod
-    def base_sigmas(num_inference_steps):
-        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
-        sigmas = mx.array(sigmas).astype(mx.float32)
-        return sigmas
+        self.guidance = guidance

src/flux_1_schnell/config/model_config.py

@@ -0,0 +1,36 @@
+from enum import Enum
+
+
+class ModelConfig(Enum):
+    FLUX1_DEV = ("black-forest-labs/FLUX.1-dev", "dev", 1000, 512)
+    FLUX1_SCHNELL = ("black-forest-labs/FLUX.1-schnell", "schnell", 1000, 256)
+
+    def __init__(
+            self,
+            model_name: str,
+            alias: str,
+            num_train_steps: int,
+            max_sequence_length: int,
+    ):
+        self.alias = alias
+        self.model_name = model_name
+        self.num_train_steps = num_train_steps
+        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/config/runtime_config.py

@@ -0,0 +1,62 @@
+import mlx.core as mx
+import numpy as np
+
+from flux_1_schnell.config.config import Config
+from flux_1_schnell.config.model_config import ModelConfig
+
+
+class RuntimeConfig:
+
+    def __init__(self, config: Config, model_config: ModelConfig):
+        self.config = config
+        self.model_config = model_config
+        self.sigmas = self._create_sigmas(config, model_config)
+
+    @property
+    def height(self):
+        return self.config.height
+
+    @property
+    def width(self):
+        return self.config.height
+
+    @property
+    def guidance(self):
+        return self.config.guidance
+
+    @property
+    def num_inference_steps(self):
+        return self.config.num_inference_steps
+
+    @property
+    def precision(self):
+        return self.config.precision
+
+    @property
+    def num_train_steps(self):
+        return self.model_config.num_train_steps
+
+    @staticmethod
+    def _create_sigmas(config, model):
+        sigmas = RuntimeConfig._create_sigmas_values(config.num_inference_steps)
+        if model == ModelConfig.FLUX1_DEV:
+            sigmas = RuntimeConfig._shift_sigmas(sigmas, config.width, config.height)
+        return sigmas
+
+    @staticmethod
+    def _create_sigmas_values(num_inference_steps: int) -> mx.array:
+        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) -> mx.array:
+        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/flux.py

@@ -4,59 +4,77 @@ from PIL import Image
 from tqdm import tqdm
 
 from flux_1_schnell.config.config import Config
+from flux_1_schnell.config.runtime_config import RuntimeConfig
 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
 from flux_1_schnell.models.vae.vae import VAE
 from flux_1_schnell.post_processing.image_util import ImageUtil
-from flux_1_schnell.scheduler.scheduler import FlowMatchEulerDiscreteNoiseScheduler
 from flux_1_schnell.tokenizer.clip_tokenizer import TokenizerCLIP
 from flux_1_schnell.tokenizer.t5_tokenizer import TokenizerT5
 from flux_1_schnell.tokenizer.tokenizer_handler import TokenizerHandler
 from flux_1_schnell.weights.weight_handler import WeightHandler
 
 
-class Flux1Schnell:
+class Flux1:
 
     def __init__(self, repo_id: str):
-        tokenizers = TokenizerHandler.load_from_disk_or_huggingface(repo_id)
-        self.t5_tokenizer = TokenizerT5(tokenizers.t5)
+        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:
+        # Create a new runtime config based on the model type and input parameters
+        config = RuntimeConfig(config, 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
+                config=config,
             )
 
-            latents = FlowMatchEulerDiscreteNoiseScheduler.denoise(
-                t=t,
-                noise=noise,
-                latent=latents,
-                config=config
-            )
+            # Take one denoise step
+            dt = config.sigmas[t + 1] - config.sigmas[t]
+            latents += noise * dt
 
+            # To enable progress tracking
             mx.eval(latents)
 
-        latents = Flux1Schnell._unpack_latents(latents, config.height, config.width)
+        # 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/text_encoder/t5_encoder/t5_self_attention.py

@@ -19,7 +19,8 @@ class T5SelfAttention(nn.Module):
         key_states = T5SelfAttention.shape(self.k(hidden_states))
         value_states = T5SelfAttention.shape(self.v(hidden_states))
         scores = mx.matmul(query_states, mx.transpose(key_states, (0, 1, 3, 2)))
-        position_bias = self._compute_bias()
+        seq_length = hidden_states.shape[1]
+        position_bias = self._compute_bias(seq_length=seq_length)
         scores += position_bias
         attn_weights = nn.softmax(scores, axis=-1)
         attn_output = T5SelfAttention.un_shape(mx.matmul(attn_weights, value_states))
@@ -34,9 +35,9 @@ class T5SelfAttention(nn.Module):
     def un_shape(states):
         return mx.reshape(mx.transpose(states, (0, 2, 1, 3)), (1, -1, 4096))
 
-    def _compute_bias(self):
-        context_position = mx.arange(start=0, stop=256, step=1)[:, None]
-        memory_position = mx.arange(start=0, stop=256, step=1)[None, :]
+    def _compute_bias(self, seq_length):
+        context_position = mx.arange(start=0, stop=seq_length, step=1)[:, None]
+        memory_position = mx.arange(start=0, stop=seq_length, step=1)[None, :]
         relative_position = memory_position - context_position
         relative_position_bucket = T5SelfAttention._relative_position_bucket(relative_position)
         values = self.relative_attention_bias(relative_position_bucket)

src/flux_1_schnell/models/transformer/guidance_embedder.py

@@ -0,0 +1,16 @@
+from mlx import nn
+import mlx.core as mx
+
+
+class GuidanceEmbedder(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.linear_1 = nn.Linear(256, 3072)
+        self.linear_2 = nn.Linear(3072, 3072)
+
+    def forward(self, sample: mx.array) -> mx.array:
+        sample = self.linear_1(sample)
+        sample = nn.silu(sample)
+        sample = self.linear_2(sample)
+        return sample

src/flux_1_schnell/models/transformer/time_text_embed.py

@@ -5,18 +5,24 @@ import mlx.core as mx
 from flux_1_schnell.config.config import Config
 from flux_1_schnell.models.transformer.text_embedder import TextEmbedder
 from flux_1_schnell.models.transformer.timestep_embedder import TimestepEmbedder
+from flux_1_schnell.models.transformer.guidance_embedder import GuidanceEmbedder
 
 
 class TimeTextEmbed(nn.Module):
 
-    def __init__(self):
+    def __init__(self, with_guidance_embed: bool = False):
         super().__init__()
         self.text_embedder = TextEmbedder()
+        self.with_guidance_embed = with_guidance_embed
+        if self.with_guidance_embed:
+            self.guidance_embedder = GuidanceEmbedder()
         self.timestep_embedder = TimestepEmbedder()
 
-    def forward(self, time_step: mx.array, pooled_projection: mx.array) -> mx.array:
-        time_steps_proj = TimeTextEmbed._time_proj(time_step)
+    def forward(self, time_step: mx.array, pooled_projection: mx.array, guidance: mx.array) -> mx.array:
+        time_steps_proj = self._time_proj(time_step)
         time_steps_emb = self.timestep_embedder.forward(time_steps_proj)
+        if self.with_guidance_embed:
+            time_steps_emb += self.guidance_embedder.forward(self._time_proj(guidance))
         pooled_projections = self.text_embedder.forward(pooled_projection)
         conditioning = time_steps_emb + pooled_projections
         return conditioning.astype(Config.precision)

src/flux_1_schnell/models/transformer/transformer.py

@@ -2,6 +2,7 @@ import mlx.core as mx
 from mlx import nn
 
 from flux_1_schnell.config.config import Config
+from flux_1_schnell.config.runtime_config import RuntimeConfig
 from flux_1_schnell.models.transformer.ada_layer_norm_continous import AdaLayerNormContinuous
 from flux_1_schnell.models.transformer.embed_nd import EmbedND
 from flux_1_schnell.models.transformer.joint_transformer_block import JointTransformerBlock
@@ -15,7 +16,8 @@ class Transformer(nn.Module):
         super().__init__()
         self.pos_embed = EmbedND()
         self.x_embedder = nn.Linear(64, 3072)
-        self.time_text_embed = TimeTextEmbed()
+        with_guidance_embed = "guidance_embedder" in weights["time_text_embed"].keys()
+        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)]
@@ -31,14 +33,15 @@ class Transformer(nn.Module):
             prompt_embeds: mx.array,
             pooled_prompt_embeds: mx.array,
             hidden_states: mx.array,
-            config: Config
+            config: RuntimeConfig,
     ) -> mx.array:
-        time_step = config.time_steps[t]
-        time_step = mx.broadcast_to(time_step, (1,))
+        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)
-        text_embeddings = self.time_text_embed.forward(time_step, pooled_prompt_embeds)
+        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()
+        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)
@@ -78,5 +81,5 @@ class Transformer(nn.Module):
         return latent_image_ids
 
     @staticmethod
-    def _prepare_text_ids() -> mx.array:
-        return mx.zeros((1, 256, 3))
+    def _prepare_text_ids(seq_len: mx.array) -> mx.array:
+        return mx.zeros((1, seq_len, 3))

src/flux_1_schnell/scheduler/scheduler.py

@@ -1,20 +0,0 @@
-import mlx.core as mx
-
-from flux_1_schnell.config.config import Config
-
-
-class FlowMatchEulerDiscreteNoiseScheduler:
-
-    @staticmethod
-    def denoise(
-            t: int,
-            noise: mx.array,
-            latent: mx.array,
-            config: Config,
-    ) -> mx.array:
-        sigma = config.sigmas[t]
-        denoised = latent - noise * sigma
-        derivative = (latent - denoised) / sigma
-        dt = config.sigmas[t + 1] - sigma
-        prev_sample = latent + derivative * dt
-        return prev_sample

src/flux_1_schnell/tokenizer/t5_tokenizer.py

@@ -3,16 +3,16 @@ from transformers import T5Tokenizer
 
 
 class TokenizerT5:
-    MAX_TOKEN_LENGTH = 256
 
-    def __init__(self, tokenizer: T5Tokenizer):
+    def __init__(self, tokenizer: T5Tokenizer, max_length: int = 256):
         self.tokenizer = tokenizer
+        self.max_length = max_length
 
     def tokenize(self, prompt: str) -> mx.array:
         return self.tokenizer(
             [prompt],
             padding="max_length",
-            max_length=TokenizerT5.MAX_TOKEN_LENGTH,
+            max_length=self.max_length,
             truncation=True,
             return_length=False,
             return_overflowing_tokens=False,

src/flux_1_schnell/tokenizer/tokenizer_handler.py

@@ -9,7 +9,7 @@ from flux_1_schnell.tokenizer.t5_tokenizer import TokenizerT5
 
 class TokenizerHandler:
 
-    def __init__(self, repo_id: str):
+    def __init__(self, repo_id: str, max_t5_length: int = 256):
         root_path = TokenizerHandler._download_or_get_cached_tokenizers(repo_id)
 
         self.clip = transformers.CLIPTokenizer.from_pretrained(
@@ -20,12 +20,12 @@ class TokenizerHandler:
         self.t5 = transformers.T5Tokenizer.from_pretrained(
             pretrained_model_name_or_path=root_path / "tokenizer_2",
             local_files_only=True,
-            max_length=TokenizerT5.MAX_TOKEN_LENGTH
+            max_length=max_t5_length
         )
 
     @staticmethod
-    def load_from_disk_or_huggingface(repo_id: str) -> "TokenizerHandler":
-        return TokenizerHandler(repo_id)
+    def load_from_disk_or_huggingface(repo_id: str, max_t5_length: int = 256) -> "TokenizerHandler":
+        return TokenizerHandler(repo_id, max_t5_length)
 
     @staticmethod
     def _download_or_get_cached_tokenizers(repo_id: str) -> Path: