zerogpu-aoti/FLUX.1-dev-base

1"""2"""3 4from typing import Any5from typing import Callable6from typing import ParamSpec7import spaces8import torch9 10from fa3 import FlashFusedFluxAttnProcessor3_011 12 13P = ParamSpec('P')14 15 16INDUCTOR_CONFIGS = {17    'conv_1x1_as_mm': True,18    'epilogue_fusion': False,19    'coordinate_descent_tuning': True,20    'coordinate_descent_check_all_directions': True,21    'max_autotune': True,22    'triton.cudagraphs': True,23}24 25 26def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):27 28    @spaces.GPU(duration=1500)29    def compile_transformer():30 31        with spaces.aoti_capture(pipeline.transformer) as call:32            pipeline(*args, **kwargs)33 34        exported = torch.export.export(35            mod=pipeline.transformer,36            args=call.args,37            kwargs=call.kwargs,38        )39 40        return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)41 42    pipeline.transformer.fuse_qkv_projections()43    pipeline.transformer.set_attn_processor(FlashFusedFluxAttnProcessor3_0())44    spaces.aoti_apply(compile_transformer(), pipeline.transformer)45