Update app.py

app.py

@@ -54,6 +54,20 @@ MODEL_ID = os.environ.get("MODEL_PATH", "diffusers-internal-dev/world-engine-mod
 pipe = ModularPipeline.from_pretrained(MODEL_ID, trust_remote_code=True, revision="aot-compatible")
 pipe.load_components(["transformer", "vae"], trust_remote_code=True, revision="aot-compatible", torch_dtype=torch.bfloat16)
 pipe.load_components(["text_encoder", "tokenizer"], trust_remote_code=True, torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+pipe.blocks.sub_blocks['before_denoise'].sub_blocks['setup_kv_cache']._setup_kv_cache(pipe.transformer, pipe.device, torch.bfloat16)
+aoti_load_(
+    pipe.transformer,
+    "diffusers-internal-dev/world-engine-aot",
+    "transformer-fp8.pt2",
+    "transformer-fp8-constants.pt"
+)
+aoti_load_(
+    pipe.vae.decoder,
+    "diffusers-internal-dev/world-engine-aot",
+    "decoder.pt2",
+    "decoder-constants.pt"
+)
 
 SEED_FRAME_URLS = [
     "https://gist.github.com/user-attachments/assets/5d91c49a-2ae9-418f-99c0-e93ae387e1de",
@@ -147,21 +161,7 @@ def create_gpu_game_loop(command_queue: Queue, initial_seed_image=None, initial_
         Generator that keeps GPU allocated and processes commands.
         Yields (frame, frame_count) tuples.
         """
-        pipe.to("cuda")
-        pipe.blocks.sub_blocks['before_denoise'].sub_blocks['setup_kv_cache']._setup_kv_cache(pipe.transformer, pipe.device, torch.bfloat16)
-        aoti_load_(
-            pipe.transformer,
-            "diffusers-internal-dev/world-engine-aot",
-            "transformer-fp8.pt2",
-            "transformer-fp8-constants.pt"
-        )
-        aoti_load_(
-            pipe.vae.decoder,
-            "diffusers-internal-dev/world-engine-aot",
-            "decoder.pt2",
-            "decoder-constants.pt"
-        )
-
+        
         n_frames = pipe.transformer.config.n_frames
         print(f"Model loaded! (n_frames={n_frames})")