fromdiffusersimportFluxPipeline# Load the pipeline in full-precision and place its model components on CUDA.pipeline=FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell").to("cuda")prompt="A cat playing with a ball of yarn"image=pipe(prompt,num_inference_steps=4).images[0]

fromdiffusersimportFluxPipeline# Load the pipeline in full-precision and place its model components on CUDA.pipeline=FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell",torch_dtype=torch.bfloat16).to("cuda")prompt="A cat playing with a ball of yarn"image=pipe(prompt,num_inference_steps=4).images[0]

fromdiffusersimportFluxPipeline# Load the pipeline in full-precision and place its model components on CUDA.pipeline=FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell").to("cuda")# Compile the compute-intensive portions of the model: denoising transformer / decoder# "max-autotune" mode tunes kernel hyperparameters and applies CUDAGraphspipeline.transformer=torch.compile(pipeline.transformer,mode="max-autotune",fullgraph=True)pipeline.vae.decode=torch.compile(pipeline.vae.decode,mode="max-autotune",fullgraph=True)# warmup for a few iterations; trigger compilationfor_inrange(3):pipeline("dummy prompt to trigger torch compilation",output_type="pil",num_inference_steps=4,    ).images[0]prompt="A cat playing with a ball of yarn"image=pipe(prompt,num_inference_steps=4).images[0]

fromdiffusersimportFluxPipeline# Load the pipeline in full-precision and place its model components on CUDA.pipeline=FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell").to("cuda")# Use channels_last memory formatpipeline.vae=pipeline.vae.to(memory_format=torch.channels_last)# Combine attention projection matrices for (q, k, v)pipeline.transformer.fuse_qkv_projections()pipeline.vae.fuse_qkv_projections()# compilation details omitted (see above)...prompt="A cat playing with a ball of yarn"image=pipe(prompt,num_inference_steps=4).images[0]

fromdiffusersimportFluxPipeline# Load the pipeline in full-precision and place its model components on CUDA.pipeline=FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell").to("cuda")# Use channels_last memory formatpipeline.vae.to(memory_format=torch.channels_last)# compilation details omitted (see above)...prompt="A cat playing with a ball of yarn"image=pipe(prompt,num_inference_steps=4).images[0]

fromdiffusersimportFluxPipeline# Load the pipeline in full-precision and place its model components on CUDA.pipeline=FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell").to("cuda")# Use channels_last memory formatpipeline.vae.to(memory_format=torch.channels_last)# Combine attention projection matrices for (q, k, v)pipeline.transformer.fuse_qkv_projections()pipeline.vae.fuse_qkv_projections()# Use FA3; reference FlashFusedFluxAttnProcessor3_0 impl for detailspipeline.transformer.set_attn_processor(FlashFusedFluxAttnProcessor3_0())# compilation details omitted (see above)...prompt="A cat playing with a ball of yarn"image=pipe(prompt,num_inference_steps=4).images[0]

fromdiffusersimportFluxPipeline# Load the pipeline in full-precision and place its model components on CUDA.pipeline=FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell").to("cuda")# Use channels_last memory formatpipeline.vae.to(memory_format=torch.channels_last)# Combine attention projection matrices for (q, k, v)pipeline.transformer.fuse_qkv_projections()pipeline.vae.fuse_qkv_projections()# Use FA3; reference FlashFusedFluxAttnProcessor3_0 impl for detailspipeline.transformer.set_attn_processor(FlashFusedFluxAttnProcessor3_0())# Apply float8 quantization on weights and activationsfromtorchao.quantizationimportquantize_,float8_dynamic_activation_float8_weightquantize_(pipeline.transformer,float8_dynamic_activation_float8_weight(),)# compilation details omitted (see above)...prompt="A cat playing with a ball of yarn"image=pipe(prompt,num_inference_steps=4).images[0]

fromdiffusersimportFluxPipeline# Load the pipeline in full-precision and place its model components on CUDA.pipeline=FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell").to("cuda")# Use channels_last memory formatpipeline.vae.to(memory_format=torch.channels_last)# Combine attention projection matrices for (q, k, v)pipeline.transformer.fuse_qkv_projections()pipeline.vae.fuse_qkv_projections()# Use FA3; reference FlashFusedFluxAttnProcessor3_0 impl for detailspipeline.transformer.set_attn_processor(FlashFusedFluxAttnProcessor3_0())# Apply float8 quantization on weights and activationsfromtorchao.quantizationimportquantize_,float8_dynamic_activation_float8_weightquantize_(pipeline.transformer,float8_dynamic_activation_float8_weight(),)# Tune Inductor flagsconfig=torch._inductor.configconfig.conv_1x1_as_mm=True# treat 1x1 convolutions as matrix muls# adjust autotuning algorithmconfig.coordinate_descent_tuning=Trueconfig.coordinate_descent_check_all_directions=Trueconfig.epilogue_fusion=False# do not fuse pointwise ops into matmuls# compilation details omitted (see above)...prompt="A cat playing with a ball of yarn"image=pipe(prompt,num_inference_steps=4).images[0]

# Apply torch.export + AOTI. If serialize=True, writes out the exported models within the cache_dir.# Otherwise, attempts to load previously-exported models from the cache_dir.# This function also applies CUDAGraphs on the loaded models.defuse_export_aoti(pipeline,cache_dir,serialize=False):fromtorch._inductor.packageimportload_package# create cache dir if neededpathlib.Path(cache_dir).mkdir(parents=True,exist_ok=True)def_example_tensor(*shape):returntorch.randn(*shape,device="cuda",dtype=torch.bfloat16)# === Transformer export ===# torch.export requires a representative set of example args to be passed intransformer_kwargs= {"hidden_states":_example_tensor(1,4096,64),"timestep":torch.tensor([1.],device="cuda",dtype=torch.bfloat16),"guidance":None,"pooled_projections":_example_tensor(1,768),"encoder_hidden_states":_example_tensor(1,512,4096),"txt_ids":_example_tensor(512,3),"img_ids":_example_tensor(4096,3),"joint_attention_kwargs": {},"return_dict":False,    }# Possibly serialize model outtransformer_package_path=os.path.join(cache_dir,"exported_transformer.pt2")ifserialize:# Apply exportexported_transformer:torch.export.ExportedProgram=torch.export.export(pipeline.transformer,args=(),kwargs=transformer_kwargs        )# Apply AOTIpath=torch._inductor.aoti_compile_and_package(exported_transformer,package_path=transformer_package_path,inductor_configs={"max_autotune":True,"triton.cudagraphs":True},        )loaded_transformer=load_package(transformer_package_path,run_single_threaded=True    )# warmup before cudagraphingwithtorch.no_grad():loaded_transformer(**transformer_kwargs)# Apply CUDAGraphs. CUDAGraphs are utilized in torch.compile with mode="max-autotune", but# they must be manually applied for torch.export + AOTI.loaded_transformer=cudagraph(loaded_transformer)pipeline.transformer.forward=loaded_transformer# warmup after cudagraphingwithtorch.no_grad():pipeline.transformer(**transformer_kwargs)# hack to get around export's limitationspipeline.vae.forward=pipeline.vae.decodevae_decode_kwargs= {"return_dict":False,    }# Possibly serialize model outdecoder_package_path=os.path.join(cache_dir,"exported_decoder.pt2")ifserialize:# Apply exportexported_decoder:torch.export.ExportedProgram=torch.export.export(pipeline.vae,args=(_example_tensor(1,16,128,128),),kwargs=vae_decode_kwargs        )# Apply AOTIpath=torch._inductor.aoti_compile_and_package(exported_decoder,package_path=decoder_package_path,inductor_configs={"max_autotune":True,"triton.cudagraphs":True},        )loaded_decoder=load_package(decoder_package_path,run_single_threaded=True)# warmup before cudagraphingwithtorch.no_grad():loaded_decoder(_example_tensor(1,16,128,128),**vae_decode_kwargs)loaded_decoder=cudagraph(loaded_decoder)pipeline.vae.decode=loaded_decoder# warmup for a few iterationsfor_inrange(3):pipeline("dummy prompt to trigger torch compilation",output_type="pil",num_inference_steps=4,        ).images[0]returnpipeline

# wrapper to automatically handle CUDAGraph record / replay over the given functiondefcudagraph(f):fromtorch.utils._pytreeimporttree_map_only_graphs= {}deff_(*args,**kwargs):key=hash(tuple(tuple(kwargs[a].shape)forainsorted(kwargs.keys())ifisinstance(kwargs[a],torch.Tensor)))ifkeyin_graphs:# use the cached wrapper if one exists. this will perform CUDAGraph replaywrapped,*_=_graphs[key]returnwrapped(*args,**kwargs)# record a new CUDAGraph and cache it for future useg=torch.cuda.CUDAGraph()in_args,in_kwargs=tree_map_only(torch.Tensor,lambdat:t.clone(), (args,kwargs))f(*in_args,**in_kwargs)# stream warmupwithtorch.cuda.graph(g):out_tensors=f(*in_args,**in_kwargs)defwrapped(*args,**kwargs):# note that CUDAGraphs require inputs / outputs to be in fixed memory locations.# inputs must be copied into the fixed input memory locations.            [a.copy_(b)fora,binzip(in_args,args)ifisinstance(a,torch.Tensor)]forkeyinkwargs:ifisinstance(kwargs[key],torch.Tensor):in_kwargs[key].copy_(kwargs[key])g.replay()# clone() outputs on the way out to disconnect them from the fixed output memory# locations. this allows for CUDAGraph reuse without accidentally overwriting memoryreturn [o.clone()foroinout_tensors]# cache function that does CUDAGraph replay_graphs[key]= (wrapped,g,in_args,in_kwargs,out_tensors)returnwrapped(*args,**kwargs)returnf_

fromdiffusersimportFluxPipeline# Load the pipeline in full-precision and place its model components on CUDA.pipeline=FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell").to("cuda")# Use channels_last memory formatpipeline.vae.to(memory_format=torch.channels_last)# Combine attention projection matrices for (q, k, v)pipeline.transformer.fuse_qkv_projections()pipeline.vae.fuse_qkv_projections()# Use FA3; reference FlashFusedFluxAttnProcessor3_0 impl for detailspipeline.transformer.set_attn_processor(FlashFusedFluxAttnProcessor3_0())# Apply float8 quantization on weights and activationsfromtorchao.quantizationimportquantize_,float8_dynamic_activation_float8_weightquantize_(pipeline.transformer,float8_dynamic_activation_float8_weight(),)# Tune Inductor flagsconfig=torch._inductor.configconfig.conv_1x1_as_mm=True# treat 1x1 convolutions as matrix muls# adjust autotuning algorithmconfig.coordinate_descent_tuning=Trueconfig.coordinate_descent_check_all_directions=Trueconfig.epilogue_fusion=False# do not fuse pointwise ops into matmuls# Apply torch.export + AOTI with CUDAGraphspipeline=use_export_aoti(pipeline,cache_dir=args.cache_dir,serialize=False)prompt="A cat playing with a ball of yarn"image=pipe(prompt,num_inference_steps=4).images[0]

# Install: pip install -U cache-ditfromdiffusersimportFluxPipelinefromcache_dit.cache_factoryimportapply_cache_on_pipe,CacheTypepipeline=FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev",torch_dtype=torch.bfloat16,).to("cuda")# cache-dit: DBCache configscache_options= {"cache_type":CacheType.DBCache,"warmup_steps":0,"max_cached_steps":-1,# -1 means no limit"Fn_compute_blocks":1,# Fn, F1, F12, etc."Bn_compute_blocks":0,# Bn, B0, B12, etc."residual_diff_threshold":0.12,# TaylorSeer options"enable_taylorseer":True,"enable_encoder_taylorseer":True,# Taylorseer cache type cache be hidden_states or residual"taylorseer_cache_type":"residual","taylorseer_kwargs": {"n_derivatives":2,    },}apply_cache_on_pipe(pipeline,**cache_options)

apply_cache_on_pipe(pipeline,**cache_options)# The cache-dit relies heavily on dynamic Python operations to maintain the cache_context,# so it is necessary to introduce graph breaks at appropriate positions to be compatible# with torch.compile. Thus, we compile the transformer with `max-autotune-no-cudagraphs`# mode if cache-dit is enabled. Otherwise, we compile with `max-autotune` mode.pipeline.transformer=torch.compile(pipeline.transformer,mode="max-autotune-no-cudagraphs",fullgraph=False, )

Important Notes