def use_compile()

in utils/pipeline_utils.py [0:0]

• 14 lines of code
• 3 McCabe index (conditional complexity)

def use_compile(pipeline):
    # Compile the compute-intensive portions of the model: denoising transformer / decoder
    pipeline.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 (`num_inference_steps` shouldn't matter)
    for _ in range(3):
        pipeline(
            "dummy prompt to trigger torch compilation",
            output_type="pil",
            num_inference_steps=4,
        ).images[0]

    return pipeline