Upload handler.py

handler.py

@@ -14,7 +14,8 @@ from torchao.quantization.quant_api import PerRow
 from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, TorchAoConfig
 from transformers import T5EncoderModel, BitsAndBytesConfig
 from optimum.quanto import freeze, qfloat8, quantize
-from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe, parallelize_pipe
+from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
+from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
 from huggingface_inference_toolkit.logging import logger
 
 import subprocess