diff --git a/src/flux_1_schnell/config/config.py b/src/flux_1_schnell/config/config.py
index 5c63d07..ca5a9f6 100644
--- a/src/flux_1_schnell/config/config.py
+++ b/src/flux_1_schnell/config/config.py
@@ -34,5 +34,5 @@ class Config:
         x1 = 256
         m = (1.15 - y1) / (4096 - x1)
         b = y1 - m * x1
-        mu = m + b
+        mu = m * self.width * self.height / 256   + b
         self.sigmas = mx.exp(mu) / (mx.exp(mu) + (1 / self.sigmas - 1))
diff --git a/src/flux_1_schnell/flux.py b/src/flux_1_schnell/flux.py
index bbed325..7a02bdc 100644
--- a/src/flux_1_schnell/flux.py
+++ b/src/flux_1_schnell/flux.py
@@ -20,8 +20,7 @@ class Flux1:
 
     def __init__(self, repo_id: str):
         self.is_dev = "FLUX.1-dev" in repo_id
-        # max_t5_length = 512 if self.is_dev else 256
-        max_t5_length = 256
+        max_t5_length = 512 if self.is_dev else 256
         tokenizers = TokenizerHandler.load_from_disk_or_huggingface(repo_id, max_t5_length)
         self.t5_tokenizer = TokenizerT5(tokenizers.t5, max_length=max_t5_length)
         self.clip_tokenizer = TokenizerCLIP(tokenizers.clip)