FLUX: nonblocking and scheduler cpu indexing

[changjonathanc]

What does this PR do?

• add nonblocking and avoids creating tensor on cpu then immediately move to device
• use cpu indexing in scheduler to avoid cuda sync/graph break

Who can review?

@a-r-r-o-w

Tests

main branch: there are 2 visible cuda sync

• first one happens here (it triggers cuda sync and has to wait for all preceding cuda ops to finish):

  diffusers/src/diffusers/pipelines/flux/pipeline_flux.py

  Line 363 in 99f6082

  text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)

• the second one happens inside of the first scheduler step:
  

  diffusers/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py

  Lines 288 to 289 in 99f6082

  	if self.step_index is None:
  	self._init_step_index(timestep)

this branch: fixing both causes results in a nice profile view

profile is done by running the following code in a loop

def bench_fn():
    image = pipe(prompt, num_inference_steps=4, width=1024, height=1024, output_type="pt")
    image = image.cpu()

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[yiyixuxu]

are you testing with flux pipeline? you can just add a set_begin_index from the pipeline to avoid that search step

diffusers/src/diffusers/pipelines/deepfloyd_if/pipeline_if.py

Line 677 in acd6d2c

if hasattr(self.scheduler, "set_begin_index"):

[changjonathanc]

Thanks for the pointer @yiyixuxu ! I updated the PR to use set_begin_index

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[a-r-r-o-w]

Hi, I apologize for the delay here. I have been profiling our pipelines and finding places for improvements continuosly to make Diffusers more production ready out-of-the-box. I will eventually club all PRs similar to this in the future and make sure that we don't have any synchronizations from start to end of pipeline execution

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.