diff --git a/README.md b/README.md index 2bdb71c..315e99a 100644 --- a/README.md +++ b/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 diff --git a/main.py b/main.py index c55537a..0372c54 100644 --- a/main.py +++ b/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, ) ) diff --git a/src/flux_1_schnell/config/config.py b/src/flux_1_schnell/config/config.py index 51e9599..ab8a2d5 100644 --- a/src/flux_1_schnell/config/config.py +++ b/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 diff --git a/src/flux_1_schnell/config/model_config.py b/src/flux_1_schnell/config/model_config.py new file mode 100644 index 0000000..d236cc2 --- /dev/null +++ b/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") diff --git a/src/flux_1_schnell/config/runtime_config.py b/src/flux_1_schnell/config/runtime_config.py new file mode 100644 index 0000000..e4b3937 --- /dev/null +++ b/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 diff --git a/src/flux_1_schnell/flux.py b/src/flux_1_schnell/flux.py index a9feeb0..4b38c2c 100644 --- a/src/flux_1_schnell/flux.py +++ b/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) diff --git a/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_self_attention.py b/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_self_attention.py index dc2794d..a628d85 100644 --- a/src/flux_1_schnell/models/text_encoder/t5_encoder/t5_self_attention.py +++ b/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) diff --git a/src/flux_1_schnell/models/transformer/guidance_embedder.py b/src/flux_1_schnell/models/transformer/guidance_embedder.py new file mode 100644 index 0000000..ca21128 --- /dev/null +++ b/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 diff --git a/src/flux_1_schnell/models/transformer/time_text_embed.py b/src/flux_1_schnell/models/transformer/time_text_embed.py index a825740..596d271 100644 --- a/src/flux_1_schnell/models/transformer/time_text_embed.py +++ b/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) diff --git a/src/flux_1_schnell/models/transformer/transformer.py b/src/flux_1_schnell/models/transformer/transformer.py index 1d132c9..15b6119 100644 --- a/src/flux_1_schnell/models/transformer/transformer.py +++ b/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)) diff --git a/src/flux_1_schnell/scheduler/__init__.py b/src/flux_1_schnell/scheduler/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/flux_1_schnell/scheduler/scheduler.py b/src/flux_1_schnell/scheduler/scheduler.py deleted file mode 100644 index 3fb3e7b..0000000 --- a/src/flux_1_schnell/scheduler/scheduler.py +++ /dev/null @@ -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 diff --git a/src/flux_1_schnell/tokenizer/t5_tokenizer.py b/src/flux_1_schnell/tokenizer/t5_tokenizer.py index e227505..2ce659f 100644 --- a/src/flux_1_schnell/tokenizer/t5_tokenizer.py +++ b/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, diff --git a/src/flux_1_schnell/tokenizer/tokenizer_handler.py b/src/flux_1_schnell/tokenizer/tokenizer_handler.py index 8c754bb..0269e53 100644 --- a/src/flux_1_schnell/tokenizer/tokenizer_handler.py +++ b/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: