Update app.py

app.py

@@ -1,6 +1,7 @@
 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.models import UNetSpatioTemporalConditionModel
 from transformers import T5EncoderModel, T5Tokenizer
 
@@ -23,6 +24,7 @@ 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))
     try:
         output = pipe(
             prompt=prompt,