[bugfix] Speed up Dreamverse Modal compile warmup

Author: Davids048

apps/dreamverse/dreamverse/video_generation.py

@@ -280,6 +280,7 @@ def initialize(self, model_config: dict | None = None) -> None:
                 compile=CompileConfig(
                     enabled=enable_compile,
                     text_encoder_enabled=enable_compile,
+                    vae_enabled=enable_compile,
                     backend="inductor",
                     fullgraph=True,
                     mode="max-autotune-no-cudagraphs",

apps/dreamverse/scripts/modal/modal_app.py

@@ -16,11 +16,12 @@
 
 image = modal.Image.from_registry(IMAGE)
 image = image.env({
+    "DREAMVERSE_IMAGE": IMAGE,
     "HF_HOME": "/root/.cache/huggingface",
     "FASTVIDEO_DREAMVERSE_HOME": "/var/lib/dreamverse",
-    "FASTVIDEO_ENABLE_STARTUP_WARMUP": "0",
+    "FASTVIDEO_ENABLE_STARTUP_WARMUP": "1",
     "FASTVIDEO_GPU_COUNT": "1",
-    "ENABLE_TORCH_COMPILE": "0",
+    "ENABLE_TORCH_COMPILE": "1",
     "STREAM_MODE": "av_fmp4",
 })
 

fastvideo/models/encoders/gemma.py

@@ -363,7 +363,32 @@ def named_parameters(self, prefix: str = "", recurse: bool = True):
 
     def prepare_for_compile(self) -> None:
         # Load Gemma outside Dynamo so torch.compile does not trace HF file-system checks.
-        _ = self.gemma_model
+        model = self.gemma_model
+        # Run one tiny eager forward before torch.compile. FastVideo calls the text encoder with
+        # Transformers' output_hidden_states path, which wraps each Gemma layer
+        # forward method and restores it by setting an instance-level forward
+        # attribute. If that attribute appears after Dynamo captures guards, the
+        # next text-encoder call recompiles; doing it here makes the first
+        # compiled call see the stable layer state.
+        token_id = getattr(model.config, "eos_token_id", None)
+        if isinstance(token_id, (list, tuple)):
+            token_id = token_id[0] if token_id else None
+        if token_id is None:
+            token_id = getattr(model.config, "pad_token_id", 0)
+        input_ids = torch.full(
+            (1, 1),
+            int(token_id or 0),
+            dtype=torch.long,
+            device=model.device,
+        )
+        attention_mask = torch.ones_like(input_ids)
+        with torch.no_grad():
+            model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+                return_dict=True,
+            )
 
     @property
     def gemma_model(self) -> Gemma3ForConditionalGeneration: