From e60cd79f0ca024e0485d600b8cf9e8c3e2635cf7 Mon Sep 17 00:00:00 2001
From: Fabio <f.xuzzo@gmail.com>
Date: Mon, 19 Aug 2024 20:41:11 +0200
Subject: [PATCH] make config static and freeze after initialisation

---
 src/flux_1_schnell/config/config.py           | 40 +++++++++++--------
 src/flux_1_schnell/flux.py                    | 10 +++--
 .../models/transformer/transformer.py         |  5 ++-
 3 files changed, 33 insertions(+), 22 deletions(-)

diff --git a/src/flux_1_schnell/config/config.py b/src/flux_1_schnell/config/config.py
index 3524b9c..13b63d8 100644
--- a/src/flux_1_schnell/config/config.py
+++ b/src/flux_1_schnell/config/config.py
@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 import mlx.core as mx
 import numpy as np
 import logging
@@ -5,37 +6,44 @@ import logging
 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
-    num_train_steps = 1000
 
     def __init__(
             self,
+            num_train_steps: int = 1000,
             num_inference_steps: int = 4,
             width: int = 1024,
             height: int = 1024,
             guidance: float = 4.0,
     ):
+        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)
-        base_sigmas = Config.base_sigmas(num_inference_steps)
         self.num_inference_steps = num_inference_steps
         self.guidance = guidance
-        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
+    def __post_init__(self, **data):
+        super().__init__(**data)
+        self.__config__.frozen = True
+    
     
-    def shift_sigmas(self):
-        y1 = 0.5
-        x1 = 256
-        m = (1.15 - y1) / (4096 - x1)
-        b = y1 - m * x1
-        mu = m * self.width * self.height / 256   + b
-        self.sigmas = mx.exp(mu) / (mx.exp(mu) + (1 / self.sigmas - 1))
-        self.sigmas[-1] = 0
diff --git a/src/flux_1_schnell/flux.py b/src/flux_1_schnell/flux.py
index f5eb5dd..a034c68 100644
--- a/src/flux_1_schnell/flux.py
+++ b/src/flux_1_schnell/flux.py
@@ -3,7 +3,7 @@ import mlx.core as mx
 from PIL import Image
 from tqdm import tqdm
 
-from flux_1_schnell.config.config import Config
+from flux_1_schnell.config.config import Config, get_sigmas, shift_sigmas
 from flux_1_schnell.latent_creator.latent_creator import LatentCreator
 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
@@ -31,8 +31,9 @@ class Flux1:
         self.clip_text_encoder = CLIPEncoder(weights.clip_encoder)
 
     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:
-            config.shift_sigmas()
+            sigmas = shift_sigmas(sigmas)
         latents = LatentCreator.create(config.height, config.width, seed)
 
         t5_tokens = self.t5_tokenizer.tokenize(prompt)
@@ -46,10 +47,11 @@ class Flux1:
                 prompt_embeds=prompt_embeds,
                 pooled_prompt_embeds=pooled_prompt_embeds,
                 hidden_states=latents,
-                config=config
+                config=config,
+                sigmas=sigmas
             )
 
-            dt = config.sigmas[t + 1] - config.sigmas[t]
+            dt = sigmas[t + 1] - sigmas[t]
             latents += noise * dt
 
             mx.eval(latents)
diff --git a/src/flux_1_schnell/models/transformer/transformer.py b/src/flux_1_schnell/models/transformer/transformer.py
index 6638501..11daa28 100644
--- a/src/flux_1_schnell/models/transformer/transformer.py
+++ b/src/flux_1_schnell/models/transformer/transformer.py
@@ -32,9 +32,10 @@ class Transformer(nn.Module):
             prompt_embeds: mx.array,
             pooled_prompt_embeds: mx.array,
             hidden_states: mx.array,
-            config: Config
+            config: Config,
+            sigmas: mx.array,
     ) -> mx.array:
-        time_step = config.sigmas[t] * config.num_train_steps
+        time_step = 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)