Skip to content

Added flux demo #3418

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.

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

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Added flux demo #3418

wants to merge 3 commits into from

Conversation

cehongwang
Copy link
Collaborator

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • This change requires a documentation update

Checklist:

  • My code follows the style guidelines of this project (You can use the linters)
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas and hacks
  • I have made corresponding changes to the documentation
  • I have added tests to verify my fix or my feature
  • New and existing unit tests pass locally with my changes
  • I have added the relevant labels to my PR in so that relevant reviewers are notified

@github-actions github-actions bot added component: conversion Issues re: Conversion stage component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Feb 27, 2025
@cehongwang cehongwang marked this pull request as draft February 27, 2025 00:39
@github-actions github-actions bot requested a review from peri044 February 27, 2025 00:39
@narendasan
Copy link
Collaborator

Can the app display the inference time, might be nice to have some stats rendered live as you generate

@cehongwang cehongwang force-pushed the flux-demo branch 4 times, most recently from 48a7c94 to 5a528f1 Compare March 18, 2025 04:44
@github-actions github-actions bot added the component: tests Issues re: Tests label Mar 18, 2025
@cehongwang cehongwang force-pushed the flux-demo branch 6 times, most recently from 361fb76 to 0aeea36 Compare March 25, 2025 09:26
@cehongwang cehongwang marked this pull request as ready for review March 26, 2025 07:53
@cehongwang cehongwang force-pushed the flux-demo branch 4 times, most recently from 9964674 to cfbc9ea Compare March 26, 2025 07:59
@cehongwang cehongwang force-pushed the flux-demo branch 4 times, most recently from c9cca30 to 27dee53 Compare April 24, 2025 06:03
@cehongwang cehongwang force-pushed the flux-demo branch 5 times, most recently from a791b42 to 3dcf128 Compare May 9, 2025 07:00
@cehongwang cehongwang force-pushed the flux-demo branch 3 times, most recently from 41139e9 to f536ac6 Compare June 2, 2025 20:38
from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
clean_up_graph_after_modifications,
)

Copy link
Collaborator

Choose a reason for hiding this comment

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

use them from examples

Copy link
Collaborator

Choose a reason for hiding this comment

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

I think we should avoid copying the whole model scripts for measuring perf. Try using the sys.path approach and importing the model and just a perf loop. something like

import sys
import os
sys.path.append(torchtrt_root + "examples/dynamo/apps")
from flux_demo import *
model = <insert FLUX model (fp16 or fp8) > 
results = measure_flux_perf(.... )

@cehongwang cehongwang force-pushed the flux-demo branch 5 times, most recently from a031a02 to 9acbee6 Compare June 12, 2025 22:51
@github-actions github-actions bot removed the component: lowering Issues re: The lowering / preprocessing passes label Jun 12, 2025
Copy link

@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.

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/quantize.py	2025-06-12 23:00:13.565607+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/impl/quantize.py	2025-06-12 23:00:38.626299+00:00
@@ -73,11 +73,10 @@
            max_bound = 127
        elif num_bits == 8 and exponent_bits == 4:
            dtype = trt.DataType.FP8
            max_bound = 448

-
        axis = None
        # int8 weight quantization is per-channel quantization(it can have one or multiple amax values)
        if dtype == trt.DataType.INT8 and amax.numel() > 1:
            # if the amax has more than one element, calculate the axis, otherwise axis value will be ignored
            amax_init_shape = amax.shape
--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py	2025-06-12 23:00:13.567607+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py	2025-06-12 23:00:39.100847+00:00
@@ -98,16 +98,17 @@
class _TorchTensorRTConstantFolder(ConstantFolder):  # type: ignore[misc]
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)

    def is_impure(self, node: torch.fx.node.Node) -> bool:
-        # Set of known quantization ops to be excluded from constant folding. 
+        # Set of known quantization ops to be excluded from constant folding.
        # Currently, we exclude all quantization ops coming from modelopt library.
        quantization_ops = {}
        try:
-            # modelopt import ensures torch.ops.tensorrt.quantize_op.default is registered 
+            # modelopt import ensures torch.ops.tensorrt.quantize_op.default is registered
            import modelopt.torch.quantization as mtq
+
            assert torch.ops.tensorrt.quantize_op.default
            quantization_ops.add(torch.ops.tensorrt.quantize_op.default)
        except Exception as e:
            pass
        if quantization_ops and node.target in quantization_ops:

Comment on lines +47 to +50
if args.debug:
pipe.transformer = FluxTransformer2DModel(
num_layers=1, num_single_layers=1, guidance_embeds=True
).to(torch.float16)
Copy link
Collaborator

Choose a reason for hiding this comment

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

remove this

@github-actions github-actions bot added the component: lowering Issues re: The lowering / preprocessing passes label Jun 13, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
cla signed component: api [Python] Issues re: Python API component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: lowering Issues re: The lowering / preprocessing passes component: runtime component: tests Issues re: Tests
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants