diff --git a/src/flux_1_schnell/config/config.py b/src/flux_1_schnell/config/config.py index 0ea5122..ab8a2d5 100644 --- a/src/flux_1_schnell/config/config.py +++ b/src/flux_1_schnell/config/config.py @@ -1,9 +1,6 @@ import logging import mlx.core as mx -import numpy as np - -from flux_1_schnell.config.model_config import ModelConfig log = logging.getLogger(__name__) @@ -13,50 +10,14 @@ class Config: def __init__( self, - num_train_steps: int = 1000, num_inference_steps: int = 4, 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: log.warning("Width and height should be multiples of 16. Rounding down.") self.width = 16 * (height // 16) self.height = 16 * (width // 16) self.num_inference_steps = num_inference_steps self.guidance = guidance - self.sigmas = sigmas - - 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/runtime_config.py b/src/flux_1_schnell/config/runtime_config.py new file mode 100644 index 0000000..3807ec1 --- /dev/null +++ b/src/flux_1_schnell/config/runtime_config.py @@ -0,0 +1,58 @@ +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.num_train_steps = 1000 + 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 + + @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 cbea48f..4b38c2c 100644 --- a/src/flux_1_schnell/flux.py +++ b/src/flux_1_schnell/flux.py @@ -4,6 +4,7 @@ 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 @@ -43,8 +44,8 @@ class 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 config with sigmas based on what model we are running - config = config.copy_with_sigmas(self.model_config) + # 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) diff --git a/src/flux_1_schnell/models/transformer/transformer.py b/src/flux_1_schnell/models/transformer/transformer.py index 22d1f28..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 @@ -32,7 +33,7 @@ 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.sigmas[t] * config.num_train_steps time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision)