update

app.py

@@ -52,8 +52,8 @@ torch._dynamo.config.capture_scalar_outputs = True
 # --- Configuration ---
 MODEL_ID = os.environ.get("MODEL_PATH", "diffusers-internal-dev/world-engine-modular")
 pipe = ModularPipeline.from_pretrained(MODEL_ID, trust_remote_code=True, revision="aot-compatible")
-pipe.load_components("transformer", trust_remote_code=True, revision="aot-compatible", torch_dtype=torch.bfloat16)
-pipe.load_components(["vae", "text_encoder", "tokenizer"], trust_remote_code=True, torch_dtype=torch.bfloat16)
+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)
 
 SEED_FRAME_URLS = [
     "https://gist.github.com/user-attachments/assets/5d91c49a-2ae9-418f-99c0-e93ae387e1de",
@@ -113,7 +113,11 @@ def create_gpu_game_loop(command_queue: Queue):
         """
         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.pt2")
+        aoti_load_(pipe.transformer, "diffusers-internal-dev/world-engine-aot", "transformer-inference-patch.pt2")
+        
+        pipe.vae.bake_weight_norm()
+        aoti_load_(pipe.vae.encoder, "diffusers-internal-dev/world-engine-aot", "encoder.pt2")
+        aoti_load_(pipe.vae.decoder, "diffusers-internal-dev/world-engine-aot", "decoder.pt2")
 
         n_frames = pipe.transformer.config.n_frames
         print(f"Model loaded! (n_frames={n_frames})")