fromdiffusersimportFluxTransformer2DModel, FluxPipelinefromdiffusers.utils.torch_utilsimportrandn_tensorimporttorch.utils.benchmarkasbenchmarkfromcontextlibimportnullcontextimportargparseimporttorchtorch.fx.experimental._config.use_duck_shape=FalseHEIGHT_WIDTH= [(1024, 1024), (1536, 768), (2048, 2048)]
defbenchmark_fn(f, *args, **kwargs):
t0=benchmark.Timer(
stmt="f(*args, **kwargs)",
globals={"args": args, "kwargs": kwargs, "f": f},
num_threads=1,
)
returnf"{(t0.blocked_autorange().mean):.3f}"defprepare_latents(
batch_size=1,
num_channels_latents=16,
height=1024,
width=1024,
dtype=torch.bfloat16,
device="cuda",
):
vae_scale_factor=8height=2* (int(height) // (vae_scale_factor*2))
width=2* (int(width) // (vae_scale_factor*2))
shape= (batch_size, num_channels_latents, height, width)
latents=randn_tensor(shape, device=device, dtype=dtype)
latents=FluxPipeline._pack_latents(latents, batch_size, num_channels_latents, height, width)
latent_image_ids=FluxPipeline._prepare_latent_image_ids(
batch_size, height//2, width//2, device, dtype
)
returnlatents, latent_image_idsdefget_conditional_inputs(batch_size, dtype=torch.bfloat16, device="cuda"):
prompt_embeds=torch.randn(batch_size, 512, 4096, dtype=dtype, device=device)
pooled_prompt_embeds=torch.randn(batch_size, 768, dtype=dtype, device=device)
text_ids=torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)
returnprompt_embeds, pooled_prompt_embeds, text_idsdefload_transformer(do_compile=False):
transformer=FluxTransformer2DModel.from_pretrained(
"black-forest-labs/FLUX.1-dev", subfolder="transformer", torch_dtype=torch.bfloat16
).to("cuda")
ifdo_compile:
transformer=torch.compile(transformer, fullgraph=True, dynamic=True)
returntransformerdefrun_inference(transformer, **kwargs):
_=transformer(**kwargs)
@torch.no_grad()defmain(transformer, batch_size, height, width):
latents, latent_image_ids=prepare_latents(batch_size=batch_size, height=height, width=width)
prompt_embeds, pooled_prompt_embeds, text_ids=get_conditional_inputs(batch_size=batch_size)
timestep=torch.full([1], 1.0, device="cuda", dtype=torch.float32)
timestep=timestep.expand(latents.shape[0]).to(latents.dtype)
timestep=timestep/1000guidance=torch.full([1], 4.5, device="cuda", dtype=torch.float32)
guidance=guidance.expand(latents.shape[0])
input_dict= {
"hidden_states": latents,
"timestep": timestep,
"guidance": guidance,
"pooled_projections": pooled_prompt_embeds,
"encoder_hidden_states": prompt_embeds,
"txt_ids": text_ids,
"img_ids": latent_image_ids
}
run_inference(transformer, **input_dict)
# time = benchmark_fn(run_inference, transformer, **input_dict)# print(f"{height}x{width}: {time} secs")if__name__=="__main__":
parser=argparse.ArgumentParser()
parser.add_argument("--batch_size", default=1, type=int)
parser.add_argument("--compile", action="store_true")
args=parser.parse_args()
transformer=load_transformer(args.compile)
context=torch._dynamo.config.patch(error_on_recompile=True) ifargs.compileelsenullcontext()
withcontext:
forheight, widthinHEIGHT_WIDTH:
main(transformer=transformer, batch_size=args.batch_size, height=height, width=width)
Similar to #11297, I was investigating potential recompilations for Flux on resolution changes.
Code
It currently fails when run with
python check_flux_recompilation.py --compile:Trace
My env:
@StrongerXi, @anijain2305 would you have any pointers?