Update app.py

app.py

@@ -6,6 +6,7 @@ import torch
 import time
 from diffusers import DiffusionPipeline, AutoencoderTiny
 from custom_pipeline import FluxWithCFGPipeline
+from diffusers.hooks import apply_first_block_cache, FirstBlockCacheConfig
 
 # --- Torch Optimizations ---
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -30,6 +31,8 @@ pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtyp
 
 pipe.to(device)
 
+apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold=0.4))
+
 # --- Inference Function ---
 @spaces.GPU
 def generate_image(prompt: str, seed: int = 42, width: int = DEFAULT_WIDTH, height: int = DEFAULT_HEIGHT, randomize_seed: bool = False, num_inference_steps: int = DEFAULT_INFERENCE_STEPS, is_enhance: bool = False):