add missing warmup

flumina.py

@@ -17,6 +17,20 @@ from typing import Optional, Set, Tuple
 from flux_pipeline import FluxPipeline
 from util import load_config, ModelVersion
 
+_ASPECT_RATIOS = [
+    (1, 1),
+    (21, 9),
+    (16, 9),
+    (3, 2),
+    (5, 4),
+    (4, 5),
+    (2, 3),
+    (9, 16),
+    (9, 21),
+    (4, 3),
+    (3, 4),
+]
+
 # Util
 def _aspect_ratio_to_width_height(aspect_ratio: str) -> Tuple[int, int]:
     """
@@ -35,22 +49,10 @@ def _aspect_ratio_to_width_height(aspect_ratio: str) -> Tuple[int, int]:
             f"Invalid aspect ratio: {aspect_ratio}. Aspect ratio must be in w:h format, e.g. 16:9"
         )
 
-    valid_aspect_ratios = [
-        (1, 1),
-        (21, 9),
-        (16, 9),
-        (3, 2),
-        (5, 4),
-        (4, 5),
-        (2, 3),
-        (9, 16),
-        (9, 21),
-        (4, 3),
-        (3, 4),
-    ]
-    if (w, h) not in valid_aspect_ratios:
+
+    if (w, h) not in _ASPECT_RATIOS:
         raise ValueError(
-            f"Invalid aspect ratio: {aspect_ratio}. Aspect ratio must be one of {valid_aspect_ratios}"
+            f"Invalid aspect ratio: {aspect_ratio}. Aspect ratio must be one of {_ASPECT_RATIOS}"
         )
 
     # We consider megapixel not 10^6 pixels but 2^20 (1024x1024) pixels
@@ -203,8 +205,24 @@ class FluminaModule(FluminaModule):
         # Initialize LoRA adapters
         self.lora_adapters: Set[str] = set()
         self.active_lora_adapter: Optional[str] = None
+
+        self._warm_up()
+
         self._test_return_sync_response = False
 
+    def _warm_up(self):
+        for f, s in _ASPECT_RATIOS:
+            print(f"Warm-up for aspect ratio {f}:{s}")
+            width, height = _aspect_ratio_to_width_height(f"{f}:{s}")
+            self.pipeline.generate(
+                prompt="a quick brown fox",
+                height=height,
+                width=width,
+                guidance=3.5,
+                num_steps=1,
+                seed=0,
+            )
+
     def _error_response(self, code: int, message: str) -> Response:
         response_json = ErrorResponse(
             error=Error(message=message),