Update app.py

app.py

@@ -32,7 +32,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # =========================
 WEIGHTS_REPO = "epfl-vita/flux2-klein-1step-rdm"
 AE_REPO, AE_FILE = "black-forest-labs/FLUX.2-dev", "ae.safetensors"
-RES = 64                     # WARNING: 512px is borderline on 16GB CPU.
+RES = 125                     # WARNING: 512px is borderline on 16GB CPU.
                               # Change to 256 if OOM occurs.
 BASE_CTX_LEN = 48
 TOKEN = os.environ.get("HF_TOKEN")
@@ -54,7 +54,7 @@ class Qwen3_0_6B_Embedder:
     def __init__(self, device="cpu", torch_dtype=torch.bfloat16):
         self.device = device
         self.torch_dtype = torch_dtype
-        self.max_length = 256
+        self.max_length = 125
 
         self.tokenizer = AutoTokenizer.from_pretrained(
             "Qwen/Qwen3-0.6B",
@@ -147,7 +147,7 @@ def _encode(prompt: str) -> torch.Tensor:
     
     real_len = len(tok(text)["input_ids"])
     # Adaptive context length (interface compatibility)
-    emb.max_length = min(64, max(BASE_CTX_LEN, int(math.ceil((real_len + 4) / 8) * 8)))
+    emb.max_length = min(125, max(BASE_CTX_LEN, int(math.ceil((real_len + 4) / 8) * 8)))
     
     with torch.no_grad():
         return emb([prompt])  # (1, 1, 7680)
@@ -182,7 +182,7 @@ def generate(prompt: str, seed: int, randomize: bool):
         )
         
         with torch.no_grad(), torch.autocast("cpu", dtype=torch.bfloat16):
-            lat = _state["model"].sample_images_with_grad(noise, ctx, {"num_steps": 2})
+            lat = _state["model"].sample_images_with_grad(noise, ctx, {"num_steps": 3})
             pix = _state["vae"].detokenize(lat)
         
         arr = (pix[0].float().clamp(0, 1) * 255).round().byte()
@@ -220,7 +220,7 @@ HEADER = """
     <span style="display:inline-block; background: rgba(236,72,153,.10); border: 1px solid rgba(236,72,153,.35);
       border-radius: 999px; padding: 3px 12px; margin: 2px;">PickScore-COCO <b>22.76</b></span>
     <span style="display:inline-block; background: rgba(245,158,11,.10); border: 1px solid rgba(245,158,11,.35);
-      border-radius: 999px; padding: 3px 12px; margin: 2px;">⚠️ 64px Borderline</span>
+      border-radius: 999px; padding: 3px 12px; margin: 2px;">⚠️ 125px Borderline</span>
     <a href="https://huggingface.co/epfl-vita/flux2-klein-1step-rdm" target="_blank"
       style="display:inline-block; border: 1px solid rgba(120,120,120,.4); border-radius: 999px;
       padding: 3px 12px; margin: 2px; text-decoration: none;">🤗 model card</a>