diff --git a/README.md b/README.md index d32d41f..bed2eb2 100644 --- a/README.md +++ b/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, diff --git a/main.py b/main.py index 2e859ea..ff903b0 100644 --- a/main.py +++ b/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, diff --git a/src/flux_1_schnell/config/config.py b/src/flux_1_schnell/config/config.py index 13b63d8..0ea5122 100644 --- a/src/flux_1_schnell/config/config.py +++ b/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 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..0f70952 --- /dev/null +++ b/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") diff --git a/src/flux_1_schnell/flux.py b/src/flux_1_schnell/flux.py index e4da8b2..cbea48f 100644 --- a/src/flux_1_schnell/flux.py +++ b/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) diff --git a/src/flux_1_schnell/models/transformer/guidance_embedder.py b/src/flux_1_schnell/models/transformer/guidance_embedder.py index 5ed915d..ca21128 100644 --- a/src/flux_1_schnell/models/transformer/guidance_embedder.py +++ b/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 \ No newline at end of file + return sample diff --git a/src/flux_1_schnell/models/transformer/transformer.py b/src/flux_1_schnell/models/transformer/transformer.py index 11daa28..22d1f28 100644 --- a/src/flux_1_schnell/models/transformer/transformer.py +++ b/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))