.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "tutorials/_rendered_examples/dynamo/torch_compile_stable_diffusion.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_tutorials__rendered_examples_dynamo_torch_compile_stable_diffusion.py>`
        to download the full example code

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_tutorials__rendered_examples_dynamo_torch_compile_stable_diffusion.py:


.. _torch_compile_stable_diffusion:

Torch Compile Stable Diffusion
======================================================

This interactive script is intended as a sample of the Torch-TensorRT workflow with `torch.compile` on a Stable Diffusion model. A sample output is featured below:

.. image:: /tutorials/images/majestic_castle.png
   :width: 512px
   :height: 512px
   :scale: 50 %
   :align: right

.. GENERATED FROM PYTHON SOURCE LINES 17-19

Imports and Model Definition
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

.. GENERATED FROM PYTHON SOURCE LINES 19-47

.. code-block:: python


    import torch
    from diffusers import DiffusionPipeline

    import torch_tensorrt

    model_id = "CompVis/stable-diffusion-v1-4"
    device = "cuda:0"

    # Instantiate Stable Diffusion Pipeline with FP16 weights
    pipe = DiffusionPipeline.from_pretrained(
        model_id, revision="fp16", torch_dtype=torch.float16
    )
    pipe = pipe.to(device)

    backend = "torch_tensorrt"

    # Optimize the UNet portion with Torch-TensorRT
    pipe.unet = torch.compile(
        pipe.unet,
        backend=backend,
        options={
            "truncate_long_and_double": True,
            "precision": torch.float16,
        },
        dynamic=False,
    )


.. GENERATED FROM PYTHON SOURCE LINES 48-50

Inference
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

.. GENERATED FROM PYTHON SOURCE LINES 50-56

.. code-block:: python


    prompt = "a majestic castle in the clouds"
    image = pipe(prompt).images[0]

    image.save("images/majestic_castle.png")
    image.show()


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** ( 0 minutes  0.000 seconds)


.. _sphx_glr_download_tutorials__rendered_examples_dynamo_torch_compile_stable_diffusion.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example




    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: torch_compile_stable_diffusion.py <torch_compile_stable_diffusion.py>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: torch_compile_stable_diffusion.ipynb <torch_compile_stable_diffusion.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_