Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[WIP] Stable Diffusion 3.x and Flux Optimization #22986

Draft
wants to merge 12 commits into
base: main
Choose a base branch
from
Draft

[WIP] Stable Diffusion 3.x and Flux Optimization #22986

wants to merge 12 commits into from
+1,340 −314

Conversation

tianleiwu
Copy link
Contributor

@tianleiwu tianleiwu commented Dec 2, 2024 •
edited
Loading

Description

This work is in progress.

Optimize the ONNX pipeline for Stable Diffusion 3.x and Flux 1.0 models (fp32 or fp16).

  • Update optimize_pipeline script
  • Update benchmkark script
  • Update document about Stable Diffusion 3.x and Flux 1.0 models
  • Add graph optimizations for MMDit model
    • FastGelu fusion
    • RMSNorm fusion
    • MultiHeadAttention fusion
  • Add graph optimizations for Flux transformer models
    • MultiHeadAttention fusion
  • Add graph optimizations for Qwen2VL used in Flux.

Optimize the ONNX pipeline for Stable Diffusion 3.x and Flux 1.0 models

python optimize_pipeline.py -i ./flux1_schnell_onnx/fp32 -o ./flux1_schnell_onnx/fp16 --float16

_optimize_sd_pipeline: Optimize flux1_schnell_onnx/fp32/transformer/model.onnx ...
               apply: Fused LayerNormalization: 115
               apply: Fused SimplifiedLayerNormalization: 152
               apply: Fused FastGelu: 76
               apply: Fused MultiHeadAttention: 57
         prune_graph: Removed 1406 nodes

Motivation and Context

@tianleiwu tianleiwu marked this pull request as draft December 3, 2024 19:19
github-advanced-security[bot]
github-advanced-security bot found potential problems Dec 5, 2024
View reviewed changes
onnxruntime/python/tools/transformers/fusion_fastgelu.py
@@ -358,3 +361,122 @@
self.nodes_to_add.append(fused_node)
self.node_name_to_graph_name[fused_node.name] = self.this_graph_name
return True

def fuse_4(self, tanh_node, input_name_to_nodes: Dict, output_name_to_node: Dict) -> Optional[bool]:

Check notice

Code scanning / CodeQL

Explicit returns mixed with implicit (fall through) returns Note

Mixing implicit and explicit returns may indicate an error as implicit returns always return None.
github-actions[bot]
github-actions bot reviewed Dec 19, 2024
View reviewed changes
Copy link
Contributor

@github-actions github-actions bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can commit the suggested changes from lintrunner.

onnxruntime/python/tools/transformers/fusion_mha_mmdit.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/fusion_mha_mmdit.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/fusion_mha_mmdit.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/fusion_mha_mmdit.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/fusion_mha_mmdit.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py Outdated Show resolved Hide resolved
onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py Outdated Show resolved Hide resolved
github-advanced-security[bot]
github-advanced-security bot found potential problems Dec 20, 2024
View reviewed changes
onnxruntime/python/tools/transformers/onnx_model_mmdit.py
Comment on lines +91 to +96
# if (options is None) or options.enable_skip_layer_norm:
# self.fuse_skip_simplified_layer_norm()
# self.fuse_skip_layer_norm()
# if (options is None) or options.enable_bias_skip_layer_norm:
# # Fuse SkipLayerNormalization and Add Bias before it.
# self.fuse_add_bias_skip_layer_norm()

Check notice

Code scanning / CodeQL

Commented-out code Note

This comment appears to contain commented-out code.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@github-actions github-actions[bot] github-actions[bot] left review comments

@ACinfr ACinfr ACinfr left review comments

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@tianleiwu @ACinfr