Add a separate runtime config

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

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

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)

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)