Update app.py

app.py

@@ -1,7 +1,8 @@
 import torch
 from diffusers import UniPCMultistepScheduler
 from diffusers import WanPipeline, AutoencoderKLWan  # Use Wan-specific VAE
-from diffusers.hooks import apply_first_block_cache, FirstBlockCacheConfig
+# from diffusers.hooks import apply_first_block_cache, FirstBlockCacheConfig
+from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
 from diffusers.models import UNetSpatioTemporalConditionModel
 from transformers import T5EncoderModel, T5Tokenizer
 
@@ -24,7 +25,8 @@ def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_ste
         pipe.unload_lora_weights()
         pipe.load_lora_weights(lora_id.strip())
     pipe.to("cuda")
-    apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold=0.2))
+    # apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold=0.2))
+    apply_cache_on_pipe(pipe)
     try:
         output = pipe(
             prompt=prompt,