chore: updates

Author: peri044

examples/dynamo/README.rst

@@ -21,4 +21,4 @@ Model Zoo
 * :ref:`_torch_export_gpt2`: Compiling a GPT2 model using AOT workflow (`ir=dynamo`)
 * :ref:`_torch_export_llama2`: Compiling a Llama2 model using AOT workflow (`ir=dynamo`)
 * :ref:`_torch_export_sam2`: Compiling SAM2 model using AOT workflow (`ir=dynamo`)
-* :ref:`_torch_export_flux_dev`: Compiling FLUX.1-dev model using AOT workflow (`ir=dynamo`)
+* :ref:`_torch_export_sam2`: Compiling FLUX.1-dev model using AOT workflow (`ir=dynamo`)

examples/dynamo/torch_export_flux_dev.py

@@ -7,16 +7,15 @@
 This example illustrates the state of the art model `FLUX.1-dev <https://huggingface.co/black-forest-labs/FLUX.1-dev>`_ optimized using
 Torch-TensorRT.
 
-**FLUX.1 [dev]** is a 12 billion parameter rectified flow transformer capable of generating images from text descriptions. It is an open-weight, guidance-distilled model for non-commercial applications.
-
+**FLUX.1 [dev]** is a 12 billion parameter rectified flow transformer capable of generating images from text descriptions. It is an open-weight, guidance-distilled model for non-commercial applications
 Install the following dependencies before compilation
 
 .. code-block:: python
 
-    pip install sentencepiece=="0.2.0" transformers=="4.48.2" accelerate=="1.3.0" diffusers=="0.32.2"
+    pip install -r requirements.txt
 
-There are different components of the ``FLUX.1-dev`` pipeline such as ``transformer``, ``vae``, ``text_encoder``, ``tokenizer`` and ``scheduler``. In this example,
-we demonstrate optimizing the ``transformer`` component of the model (which typically consumes >95% of the e2e diffusion latency)
+There are different components of the FLUX.1-dev pipeline such as `transformer`, `vae`, `text_encoder`, `tokenizer` and `scheduler`. In this example,
+we demonstrate optimizing the `transformer` component of the model (which typically consumes >95% of the e2el diffusion latency)
 """
 
 # %%
@@ -31,7 +30,7 @@
 # Define the FLUX-1.dev model
 # -----------------------------
 # Load the ``FLUX-1.dev`` pretrained pipeline using ``FluxPipeline`` class.
-# ``FluxPipeline`` includes different components such as ``transformer``, ``vae``, ``text_encoder``, ``tokenizer`` and ``scheduler`` necessary
+# ``FluxPipeline`` includes all the different components such as `transformer`, `vae`, `text_encoder`, `tokenizer` and `scheduler` necessary
 # to generate an image. We load the weights in ``FP16`` precision using ``torch_dtype`` argument
 DEVICE = "cuda:0"
 pipe = FluxPipeline.from_pretrained(
@@ -45,14 +44,14 @@
 
 
 # %%
-# Export the backbone using torch.export
+# Export the backbone using ``torch.export``
 # --------------------------------------------------
-# Define the dummy inputs and their respective dynamic shapes. We export the transformer backbone with dynamic shapes with a ``batch_size=2``
-# due to `0/1 specialization <https://docs.google.com/document/d/16VPOa3d-Liikf48teAOmxLc92rgvJdfosIy-yoT38Io/edit?fbclid=IwAR3HNwmmexcitV0pbZm_x1a4ykdXZ9th_eJWK-3hBtVgKnrkmemz6Pm5jRQ&tab=t.0#heading=h.ez923tomjvyk>`_
+# Define the dummy inputs and their respective dynamic shapes. We export the transformer backbone with dynamic shapes with a batch size of 2
+# due to 0/1 specialization <https://docs.google.com/document/d/16VPOa3d-Liikf48teAOmxLc92rgvJdfosIy-yoT38Io/edit?fbclid=IwAR3HNwmmexcitV0pbZm_x1a4ykdXZ9th_eJWK-3hBtVgKnrkmemz6Pm5jRQ&tab=t.0#heading=h.ez923tomjvyk>`_
 batch_size = 2
 BATCH = torch.export.Dim("batch", min=1, max=2)
 SEQ_LEN = torch.export.Dim("seq_len", min=1, max=512)
-# This particular min, max values for img_id input are recommended by torch dynamo during the export of the model.
+# This particular min, max values are recommended by torch dynamo during the export of the model.
 # To see this recommendation, you can try exporting using min=1, max=4096
 IMG_ID = torch.export.Dim("img_id", min=3586, max=4096)
 dynamic_shapes = {
@@ -93,24 +92,18 @@
 # %%
 # Torch-TensorRT compilation
 # ---------------------------
-# .. note::
-#    The compilation requires a GPU with high memory (> 80GB) since TensorRT is storing the weights in FP32 precision. This is a known issue and will be resolved in the future.
-#
-#
-# We enable ``FP32`` matmul accumulation using ``use_fp32_acc=True`` to ensure accuracy is preserved by introducing cast to ``FP32`` nodes.
-# We also enable explicit typing to ensure TensorRT respects the datatypes set by the user which is a requirement for FP32 matmul accumulation.
-# Since this is a 12 billion parameter model, it takes around 20-30 min to compile on H100 GPU. The model is completely convertible and results in
-# a single TensorRT engine.
+# We enable FP32 matmul accumulation using ``use_fp32_acc=True`` to preserve accuracy with the original Pytorch model.
+# Since this is a 12 billion parameter model, it takes around 20-30 min on H100 GPU
 trt_gm = torch_tensorrt.dynamo.compile(
     ep,
     inputs=dummy_inputs,
     enabled_precisions={torch.float32},
     truncate_double=True,
     min_block_size=1,
+    debug=True,
     use_fp32_acc=True,
     use_explicit_typing=True,
 )
-
 # %%
 # Post Processing
 # ---------------------------
@@ -145,6 +138,4 @@ def generate_image(pipe, prompt, image_name):
 
 # %%
 # The generated image is as shown below
-#
-# .. image:: dog_code.png
-#
+#    .. image:: dog_code.png