64 kB

	<meta charset="utf-8" /><meta name="hf:doc:metadata" content="{"title":"Reduce memory usage","local":"reduce-memory-usage","sections":[{"title":"Sliced VAE","local":"sliced-vae","sections":[],"depth":2},{"title":"Tiled VAE","local":"tiled-vae","sections":[],"depth":2},{"title":"CPU offloading","local":"cpu-offloading","sections":[],"depth":2},{"title":"Model offloading","local":"model-offloading","sections":[],"depth":2},{"title":"Group offloading","local":"group-offloading","sections":[],"depth":2},{"title":"FP8 layerwise weight-casting","local":"fp8-layerwise-weight-casting","sections":[],"depth":2},{"title":"Channels-last memory format","local":"channels-last-memory-format","sections":[],"depth":2},{"title":"Tracing","local":"tracing","sections":[],"depth":2},{"title":"Memory-efficient attention","local":"memory-efficient-attention","sections":[],"depth":2}],"depth":1}">
	<link href="/docs/diffusers/pr_11234/en/_app/immutable/assets/0.e3b0c442.css" rel="modulepreload">
	<link rel="modulepreload" href="/docs/diffusers/pr_11234/en/_app/immutable/entry/start.5a617280.js">
	<link rel="modulepreload" href="/docs/diffusers/pr_11234/en/_app/immutable/chunks/scheduler.8c3d61f6.js">
	<link rel="modulepreload" href="/docs/diffusers/pr_11234/en/_app/immutable/chunks/singletons.0c831ec7.js">
	<link rel="modulepreload" href="/docs/diffusers/pr_11234/en/_app/immutable/chunks/index.0997d446.js">
	<link rel="modulepreload" href="/docs/diffusers/pr_11234/en/_app/immutable/chunks/paths.c606bb8d.js">
	<link rel="modulepreload" href="/docs/diffusers/pr_11234/en/_app/immutable/entry/app.1a1de8b6.js">
	<link rel="modulepreload" href="/docs/diffusers/pr_11234/en/_app/immutable/chunks/index.da70eac4.js">
	<link rel="modulepreload" href="/docs/diffusers/pr_11234/en/_app/immutable/nodes/0.8544c871.js">
	<link rel="modulepreload" href="/docs/diffusers/pr_11234/en/_app/immutable/chunks/each.e59479a4.js">
	<link rel="modulepreload" href="/docs/diffusers/pr_11234/en/_app/immutable/nodes/214.6b9a8d11.js">
	<link rel="modulepreload" href="/docs/diffusers/pr_11234/en/_app/immutable/chunks/Tip.1d9b8c37.js">
	<link rel="modulepreload" href="/docs/diffusers/pr_11234/en/_app/immutable/chunks/CodeBlock.a9c4becf.js">
	<link rel="modulepreload" href="/docs/diffusers/pr_11234/en/_app/immutable/chunks/index.5d4ab994.js"><!-- HEAD_svelte-u9bgzb_START --><meta name="hf:doc:metadata" content="{"title":"Reduce memory usage","local":"reduce-memory-usage","sections":[{"title":"Sliced VAE","local":"sliced-vae","sections":[],"depth":2},{"title":"Tiled VAE","local":"tiled-vae","sections":[],"depth":2},{"title":"CPU offloading","local":"cpu-offloading","sections":[],"depth":2},{"title":"Model offloading","local":"model-offloading","sections":[],"depth":2},{"title":"Group offloading","local":"group-offloading","sections":[],"depth":2},{"title":"FP8 layerwise weight-casting","local":"fp8-layerwise-weight-casting","sections":[],"depth":2},{"title":"Channels-last memory format","local":"channels-last-memory-format","sections":[],"depth":2},{"title":"Tracing","local":"tracing","sections":[],"depth":2},{"title":"Memory-efficient attention","local":"memory-efficient-attention","sections":[],"depth":2}],"depth":1}"><!-- HEAD_svelte-u9bgzb_END --> <p></p> <h1 class="relative group"><a id="reduce-memory-usage" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#reduce-memory-usage"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Reduce memory usage</span></h1> <p data-svelte-h="svelte-18ag1qn">A barrier to using diffusion models is the large amount of memory required. To overcome this challenge, there are several memory-reducing techniques you can use to run even some of the largest models on free-tier or consumer GPUs. Some of these techniques can even be combined to further reduce memory usage.</p> <div class="course-tip bg-gradient-to-br dark:bg-gradient-to-r before:border-green-500 dark:before:border-green-800 from-green-50 dark:from-gray-900 to-white dark:to-gray-950 border border-green-50 text-green-700 dark:text-gray-400"><p data-svelte-h="svelte-15lbm2y">In many cases, optimizing for memory or speed leads to improved performance in the other, so you should try to optimize for both whenever you can. This guide focuses on minimizing memory usage, but you can also learn more about how to <a href="fp16">Speed up inference</a>.</p></div> <p data-svelte-h="svelte-1nfsyhf">The results below are obtained from generating a single 512x512 image from the prompt a photo of an astronaut riding a horse on mars with 50 DDIM steps on a Nvidia Titan RTX, demonstrating the speed-up you can expect as a result of reduced memory consumption.</p> <table data-svelte-h="svelte-qiiy0g"><thead><tr><th></th> <th>latency</th> <th>speed-up</th></tr></thead> <tbody><tr><td>original</td> <td>9.50s</td> <td>x1</td></tr> <tr><td>fp16</td> <td>3.61s</td> <td>x2.63</td></tr> <tr><td>channels last</td> <td>3.30s</td> <td>x2.88</td></tr> <tr><td>traced UNet</td> <td>3.21s</td> <td>x2.96</td></tr> <tr><td>memory-efficient attention</td> <td>2.63s</td> <td>x3.61</td></tr></tbody></table> <h2 class="relative group"><a id="sliced-vae" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#sliced-vae"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Sliced VAE</span></h2> <p data-svelte-h="svelte-dx6t25">Sliced VAE enables decoding large batches of images with limited VRAM or batches with 32 images or more by decoding the batches of latents one image at a time. You’ll likely want to couple this with <a href="/docs/diffusers/pr_11234/en/api/models/overview#diffusers.ModelMixin.enable_xformers_memory_efficient_attention">enable_xformers_memory_efficient_attention()</a> to reduce memory use further if you have xFormers installed.</p> <p data-svelte-h="svelte-1nz82rs">To use sliced VAE, call <a href="/docs/diffusers/pr_11234/en/api/pipelines/latent_consistency_models#diffusers.LatentConsistencyModelPipeline.enable_vae_slicing">enable_vae_slicing()</a> on your pipeline before inference:</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> torch
	<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> StableDiffusionPipeline

	pipe = StableDiffusionPipeline.from_pretrained(
	<span class="hljs-string">"stable-diffusion-v1-5/stable-diffusion-v1-5"</span>,
	torch_dtype=torch.float16,
	use_safetensors=<span class="hljs-literal">True</span>,
	)
	pipe = pipe.to(<span class="hljs-string">"cuda"</span>)

	prompt = <span class="hljs-string">"a photo of an astronaut riding a horse on mars"</span>
	pipe.enable_vae_slicing()
	<span class="hljs-comment">#pipe.enable_xformers_memory_efficient_attention()</span>
	images = pipe([prompt] * <span class="hljs-number">32</span>).images<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-fte8mn">You may see a small performance boost in VAE decoding on multi-image batches, and there should be no performance impact on single-image batches.</p> <h2 class="relative group"><a id="tiled-vae" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#tiled-vae"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Tiled VAE</span></h2> <p data-svelte-h="svelte-2uu404">Tiled VAE processing also enables working with large images on limited VRAM (for example, generating 4k images on 8GB of VRAM) by splitting the image into overlapping tiles, decoding the tiles, and then blending the outputs together to compose the final image. You should also used tiled VAE with <a href="/docs/diffusers/pr_11234/en/api/models/overview#diffusers.ModelMixin.enable_xformers_memory_efficient_attention">enable_xformers_memory_efficient_attention()</a> to reduce memory use further if you have xFormers installed.</p> <p data-svelte-h="svelte-1x966xx">To use tiled VAE processing, call <a href="/docs/diffusers/pr_11234/en/api/pipelines/latent_consistency_models#diffusers.LatentConsistencyModelPipeline.enable_vae_tiling">enable_vae_tiling()</a> on your pipeline before inference:</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> torch
	<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> StableDiffusionPipeline, UniPCMultistepScheduler

	pipe = StableDiffusionPipeline.from_pretrained(
	<span class="hljs-string">"stable-diffusion-v1-5/stable-diffusion-v1-5"</span>,
	torch_dtype=torch.float16,
	use_safetensors=<span class="hljs-literal">True</span>,
	)
	pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
	pipe = pipe.to(<span class="hljs-string">"cuda"</span>)
	prompt = <span class="hljs-string">"a beautiful landscape photograph"</span>
	pipe.enable_vae_tiling()
	<span class="hljs-comment">#pipe.enable_xformers_memory_efficient_attention()</span>

	image = pipe([prompt], width=<span class="hljs-number">3840</span>, height=<span class="hljs-number">2224</span>, num_inference_steps=<span class="hljs-number">20</span>).images[<span class="hljs-number">0</span>]<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-nyzaca">The output image has some tile-to-tile tone variation because the tiles are decoded separately, but you shouldn’t see any sharp and obvious seams between the tiles. Tiling is turned off for images that are 512x512 or smaller.</p> <h2 class="relative group"><a id="cpu-offloading" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#cpu-offloading"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>CPU offloading</span></h2> <p data-svelte-h="svelte-435lli">Offloading the weights to the CPU and only loading them on the GPU when performing the forward pass can also save memory. Often, this technique can reduce memory consumption to less than 3GB.</p> <p data-svelte-h="svelte-ji0gh0">To perform CPU offloading, call <a href="/docs/diffusers/pr_11234/en/api/pipelines/overview#diffusers.DiffusionPipeline.enable_sequential_cpu_offload">enable_sequential_cpu_offload()</a>:</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> torch
	<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> StableDiffusionPipeline

	pipe = StableDiffusionPipeline.from_pretrained(
	<span class="hljs-string">"stable-diffusion-v1-5/stable-diffusion-v1-5"</span>,
	torch_dtype=torch.float16,
	use_safetensors=<span class="hljs-literal">True</span>,
	)

	prompt = <span class="hljs-string">"a photo of an astronaut riding a horse on mars"</span>
	pipe.enable_sequential_cpu_offload()
	image = pipe(prompt).images[<span class="hljs-number">0</span>]<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-153k8wf">CPU offloading works on submodules rather than whole models. This is the best way to minimize memory consumption, but inference is much slower due to the iterative nature of the diffusion process. The UNet component of the pipeline runs several times (as many as <code>num_inference_steps</code>); each time, the different UNet submodules are sequentially onloaded and offloaded as needed, resulting in a large number of memory transfers.</p> <div class="course-tip bg-gradient-to-br dark:bg-gradient-to-r before:border-green-500 dark:before:border-green-800 from-green-50 dark:from-gray-900 to-white dark:to-gray-950 border border-green-50 text-green-700 dark:text-gray-400"><p data-svelte-h="svelte-1yyrc0t">Consider using <a href="#model-offloading">model offloading</a> if you want to optimize for speed because it is much faster. The tradeoff is your memory savings won’t be as large.</p></div> <div class="course-tip course-tip-orange bg-gradient-to-br dark:bg-gradient-to-r before:border-orange-500 dark:before:border-orange-800 from-orange-50 dark:from-gray-900 to-white dark:to-gray-950 border border-orange-50 text-orange-700 dark:text-gray-400"><p data-svelte-h="svelte-16y86kh">When using <a href="/docs/diffusers/pr_11234/en/api/pipelines/overview#diffusers.DiffusionPipeline.enable_sequential_cpu_offload">enable_sequential_cpu_offload()</a>, don’t move the pipeline to CUDA beforehand or else the gain in memory consumption will only be minimal (see this <a href="https://github.com/huggingface/diffusers/issues/1934" rel="nofollow">issue</a> for more information).</p> <p data-svelte-h="svelte-1361y4a"><a href="/docs/diffusers/pr_11234/en/api/pipelines/overview#diffusers.DiffusionPipeline.enable_sequential_cpu_offload">enable_sequential_cpu_offload()</a> is a stateful operation that installs hooks on the models.</p></div> <h2 class="relative group"><a id="model-offloading" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#model-offloading"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Model offloading</span></h2> <div class="course-tip bg-gradient-to-br dark:bg-gradient-to-r before:border-green-500 dark:before:border-green-800 from-green-50 dark:from-gray-900 to-white dark:to-gray-950 border border-green-50 text-green-700 dark:text-gray-400"><p data-svelte-h="svelte-1guxtqy">Model offloading requires 🤗 Accelerate version 0.17.0 or higher.</p></div> <p data-svelte-h="svelte-1w0xgx8"><a href="#cpu-offloading">Sequential CPU offloading</a> preserves a lot of memory but it makes inference slower because submodules are moved to GPU as needed, and they’re immediately returned to the CPU when a new module runs.</p> <p data-svelte-h="svelte-nislc4">Full-model offloading is an alternative that moves whole models to the GPU, instead of handling each model’s constituent <em>submodules</em>. There is a negligible impact on inference time (compared with moving the pipeline to <code>cuda</code>), and it still provides some memory savings.</p> <p data-svelte-h="svelte-11gjm3">During model offloading, only one of the main components of the pipeline (typically the text encoder, UNet and VAE)
	is placed on the GPU while the others wait on the CPU. Components like the UNet that run for multiple iterations stay on the GPU until they’re no longer needed.</p> <p data-svelte-h="svelte-v0rtjh">Enable model offloading by calling <a href="/docs/diffusers/pr_11234/en/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload">enable_model_cpu_offload()</a> on the pipeline:</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> torch
	<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> StableDiffusionPipeline

	pipe = StableDiffusionPipeline.from_pretrained(
	<span class="hljs-string">"stable-diffusion-v1-5/stable-diffusion-v1-5"</span>,
	torch_dtype=torch.float16,
	use_safetensors=<span class="hljs-literal">True</span>,
	)

	prompt = <span class="hljs-string">"a photo of an astronaut riding a horse on mars"</span>
	pipe.enable_model_cpu_offload()
	image = pipe(prompt).images[<span class="hljs-number">0</span>]<!-- HTML_TAG_END --></pre></div> <div class="course-tip course-tip-orange bg-gradient-to-br dark:bg-gradient-to-r before:border-orange-500 dark:before:border-orange-800 from-orange-50 dark:from-gray-900 to-white dark:to-gray-950 border border-orange-50 text-orange-700 dark:text-gray-400"><p data-svelte-h="svelte-bljvwk">In order to properly offload models after they’re called, it is required to run the entire pipeline and models are called in the pipeline’s expected order. Exercise caution if models are reused outside the context of the pipeline after hooks have been installed. See <a href="https://huggingface.co/docs/accelerate/en/package_reference/big_modeling#accelerate.hooks.remove_hook_from_module" rel="nofollow">Removing Hooks</a> for more information.</p> <p data-svelte-h="svelte-vpnyla"><a href="/docs/diffusers/pr_11234/en/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload">enable_model_cpu_offload()</a> is a stateful operation that installs hooks on the models and state on the pipeline.</p></div> <h2 class="relative group"><a id="group-offloading" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#group-offloading"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Group offloading</span></h2> <p data-svelte-h="svelte-n1o14o">Group offloading is the middle ground between sequential and model offloading. It works by offloading groups of internal layers (either <code>torch.nn.ModuleList</code> or <code>torch.nn.Sequential</code>), which uses less memory than model-level offloading. It is also faster than sequential-level offloading because the number of device synchronizations is reduced.</p> <p data-svelte-h="svelte-lupt1d">To enable group offloading, call the <a href="/docs/diffusers/pr_11234/en/api/models/overview#diffusers.ModelMixin.enable_group_offload">enable_group_offload()</a> method on the model if it is a Diffusers model implementation. For any other model implementation, use <a href="/docs/diffusers/pr_11234/en/api/utilities#diffusers.hooks.apply_group_offloading">apply_group_offloading()</a>:</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> torch
	<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> CogVideoXPipeline
	<span class="hljs-keyword">from</span> diffusers.hooks <span class="hljs-keyword">import</span> apply_group_offloading
	<span class="hljs-keyword">from</span> diffusers.utils <span class="hljs-keyword">import</span> export_to_video

	<span class="hljs-comment"># Load the pipeline</span>
	onload_device = torch.device(<span class="hljs-string">"cuda"</span>)
	offload_device = torch.device(<span class="hljs-string">"cpu"</span>)
	pipe = CogVideoXPipeline.from_pretrained(<span class="hljs-string">"THUDM/CogVideoX-5b"</span>, torch_dtype=torch.bfloat16)

	<span class="hljs-comment"># We can utilize the enable_group_offload method for Diffusers model implementations</span>
	pipe.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type=<span class="hljs-string">"leaf_level"</span>, use_stream=<span class="hljs-literal">True</span>)

	<span class="hljs-comment"># Uncomment the following to also allow recording the current streams.</span>
	<span class="hljs-comment"># pipe.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True, record_stream=True)</span>

	<span class="hljs-comment"># For any other model implementations, the apply_group_offloading function can be used</span>
	apply_group_offloading(pipe.text_encoder, onload_device=onload_device, offload_type=<span class="hljs-string">"block_level"</span>, num_blocks_per_group=<span class="hljs-number">2</span>)
	apply_group_offloading(pipe.vae, onload_device=onload_device, offload_type=<span class="hljs-string">"leaf_level"</span>)

	prompt = (
	<span class="hljs-string">"A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. "</span>
	<span class="hljs-string">"The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other "</span>
	<span class="hljs-string">"pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, "</span>
	<span class="hljs-string">"casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. "</span>
	<span class="hljs-string">"The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical "</span>
	<span class="hljs-string">"atmosphere of this unique musical performance."</span>
	)
	video = pipe(prompt=prompt, guidance_scale=<span class="hljs-number">6</span>, num_inference_steps=<span class="hljs-number">50</span>).frames[<span class="hljs-number">0</span>]
	<span class="hljs-comment"># This utilized about 14.79 GB. It can be further reduced by using tiling and using leaf_level offloading throughout the pipeline.</span>
	<span class="hljs-built_in">print</span>(<span class="hljs-string">f"Max memory reserved: <span class="hljs-subst">{torch.cuda.max_memory_allocated() / <span class="hljs-number">1024</span>**<span class="hljs-number">3</span>:<span class="hljs-number">.2</span>f}</span> GB"</span>)
	export_to_video(video, <span class="hljs-string">"output.mp4"</span>, fps=<span class="hljs-number">8</span>)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-fe960w">Group offloading (for CUDA devices with support for asynchronous data transfer streams) overlaps data transfer and computation to reduce the overall execution time compared to sequential offloading. This is enabled using layer prefetching with CUDA streams. The next layer to be executed is loaded onto the accelerator device while the current layer is being executed - this increases the memory requirements slightly. Group offloading also supports leaf-level offloading (equivalent to sequential CPU offloading) but can be made much faster when using streams.</p> <div class="course-tip bg-gradient-to-br dark:bg-gradient-to-r before:border-green-500 dark:before:border-green-800 from-green-50 dark:from-gray-900 to-white dark:to-gray-950 border border-green-50 text-green-700 dark:text-gray-400"><ul data-svelte-h="svelte-15p4zc8"><li>Group offloading may not work with all models out-of-the-box. If the forward implementations of the model contain weight-dependent device-casting of inputs, it may clash with the offloading mechanism’s handling of device-casting.</li> <li>The <code>offload_type</code> parameter can be set to either <code>block_level</code> or <code>leaf_level</code>. <code>block_level</code> offloads groups of <code>torch::nn::ModuleList</code> or <code>torch::nn:Sequential</code> modules based on a configurable attribute <code>num_blocks_per_group</code>. For example, if you set <code>num_blocks_per_group=2</code> on a standard transformer model containing 40 layers, it will onload/offload 2 layers at a time for a total of 20 onload/offloads. This drastically reduces the VRAM requirements. <code>leaf_level</code> offloads individual layers at the lowest level, which is equivalent to sequential offloading. However, unlike sequential offloading, group offloading can be made much faster when using streams, with minimal compromise to end-to-end generation time.</li> <li>The <code>use_stream</code> parameter can be used with CUDA devices to enable prefetching layers for onload. It defaults to <code>False</code>. Layer prefetching allows overlapping computation and data transfer of model weights, which drastically reduces the overall execution time compared to other offloading methods. However, it can increase the CPU RAM usage significantly. Ensure that available CPU RAM that is at least twice the size of the model when setting <code>use_stream=True</code>. You can find more information about CUDA streams <a href="https://pytorch.org/docs/stable/generated/torch.cuda.Stream.html" rel="nofollow">here</a></li> <li>If specifying <code>use_stream=True</code> on VAEs with tiling enabled, make sure to do a dummy forward pass (possibly with dummy inputs) before the actual inference to avoid device-mismatch errors. This may not work on all implementations. Please open an issue if you encounter any problems.</li> <li>The parameter <code>low_cpu_mem_usage</code> can be set to <code>True</code> to reduce CPU memory usage when using streams for group offloading. This is useful when the CPU memory is the bottleneck, but it may counteract the benefits of using streams and increase the overall execution time. The CPU memory savings come from creating pinned-tensors on-the-fly instead of pre-pinning them. This parameter is better suited for using <code>leaf_level</code> offloading.</li> <li>When using <code>use_stream=True</code>, users can additionally specify <code>record_stream=True</code> to get better speedups at the expense of slightly increased memory usage. Refer to the <a href="https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html" rel="nofollow">official PyTorch docs</a> to know more about this.</li></ul> <p data-svelte-h="svelte-1utmqww">For more information about available parameters and an explanation of how group offloading works, refer to <a href="/docs/diffusers/pr_11234/en/api/utilities#diffusers.hooks.apply_group_offloading">apply_group_offloading()</a>.</p></div> <h2 class="relative group"><a id="fp8-layerwise-weight-casting" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#fp8-layerwise-weight-casting"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>FP8 layerwise weight-casting</span></h2> <p data-svelte-h="svelte-lwtukj">PyTorch supports <code>torch.float8_e4m3fn</code> and <code>torch.float8_e5m2</code> as weight storage dtypes, but they can’t be used for computation in many different tensor operations due to unimplemented kernel support. However, you can use these dtypes to store model weights in fp8 precision and upcast them on-the-fly when the layers are used in the forward pass. This is known as layerwise weight-casting.</p> <p data-svelte-h="svelte-r2gdi7">Typically, inference on most models is done with <code>torch.float16</code> or <code>torch.bfloat16</code> weight/computation precision. Layerwise weight-casting cuts down the memory footprint of the model weights by approximately half.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> torch
	<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> CogVideoXPipeline, CogVideoXTransformer3DModel
	<span class="hljs-keyword">from</span> diffusers.utils <span class="hljs-keyword">import</span> export_to_video

	model_id = <span class="hljs-string">"THUDM/CogVideoX-5b"</span>

	<span class="hljs-comment"># Load the model in bfloat16 and enable layerwise casting</span>
	transformer = CogVideoXTransformer3DModel.from_pretrained(model_id, subfolder=<span class="hljs-string">"transformer"</span>, torch_dtype=torch.bfloat16)
	transformer.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)

	<span class="hljs-comment"># Load the pipeline</span>
	pipe = CogVideoXPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch.bfloat16)
	pipe.to(<span class="hljs-string">"cuda"</span>)

	prompt = (
	<span class="hljs-string">"A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. "</span>
	<span class="hljs-string">"The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other "</span>
	<span class="hljs-string">"pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, "</span>
	<span class="hljs-string">"casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. "</span>
	<span class="hljs-string">"The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical "</span>
	<span class="hljs-string">"atmosphere of this unique musical performance."</span>
	)
	video = pipe(prompt=prompt, guidance_scale=<span class="hljs-number">6</span>, num_inference_steps=<span class="hljs-number">50</span>).frames[<span class="hljs-number">0</span>]
	export_to_video(video, <span class="hljs-string">"output.mp4"</span>, fps=<span class="hljs-number">8</span>)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-lgt48u">In the above example, layerwise casting is enabled on the transformer component of the pipeline. By default, certain layers are skipped from the FP8 weight casting because it can lead to significant degradation of generation quality. The normalization and modulation related weight parameters are also skipped by default.</p> <p data-svelte-h="svelte-nk60bm">However, you gain more control and flexibility by directly utilizing the <a href="/docs/diffusers/pr_11234/en/api/utilities#diffusers.hooks.apply_layerwise_casting">apply_layerwise_casting()</a> function instead of <a href="/docs/diffusers/pr_11234/en/api/models/overview#diffusers.ModelMixin.enable_layerwise_casting">enable_layerwise_casting()</a>.</p> <div class="course-tip bg-gradient-to-br dark:bg-gradient-to-r before:border-green-500 dark:before:border-green-800 from-green-50 dark:from-gray-900 to-white dark:to-gray-950 border border-green-50 text-green-700 dark:text-gray-400"><ul data-svelte-h="svelte-1gstadw"><li>Layerwise casting may not work with all models out-of-the-box. Sometimes, the forward implementations of the model might contain internal typecasting of weight values. Such implementations are not supported due to the currently simplistic implementation of layerwise casting, which assumes that the forward pass is independent of the weight precision and that the input dtypes are always in <code>compute_dtype</code>. An example of an incompatible implementation can be found <a href="https://github.com/huggingface/transformers/blob/7f5077e53682ca855afc826162b204ebf809f1f9/src/transformers/models/t5/modeling_t5.py#L294-L299" rel="nofollow">here</a>.</li> <li>Layerwise casting may fail on custom modeling implementations that make use of <a href="https://github.com/huggingface/peft" rel="nofollow">PEFT</a> layers. Some minimal checks to handle this case is implemented but is not extensively tested or guaranteed to work in all cases.</li> <li>It can be also be applied partially to specific layers of a model. Partially applying layerwise casting can either be done manually by calling the <code>apply_layerwise_casting</code> function on specific internal modules, or by specifying the <code>skip_modules_pattern</code> and <code>skip_modules_classes</code> parameters for a root module. These parameters are particularly useful for layers such as normalization and modulation.</li></ul></div> <h2 class="relative group"><a id="channels-last-memory-format" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#channels-last-memory-format"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Channels-last memory format</span></h2> <p data-svelte-h="svelte-sky5la">The channels-last memory format is an alternative way of ordering NCHW tensors in memory to preserve dimension ordering. Channels-last tensors are ordered in such a way that the channels become the densest dimension (storing images pixel-per-pixel). Since not all operators currently support the channels-last format, it may result in worst performance but you should still try and see if it works for your model.</p> <p data-svelte-h="svelte-mhdc10">For example, to set the pipeline’s UNet to use the channels-last format:</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-built_in">print</span>(pipe.unet.conv_out.state_dict()[<span class="hljs-string">"weight"</span>].stride()) <span class="hljs-comment"># (2880, 9, 3, 1)</span>
	pipe.unet.to(memory_format=torch.channels_last) <span class="hljs-comment"># in-place operation</span>
	<span class="hljs-built_in">print</span>(
	pipe.unet.conv_out.state_dict()[<span class="hljs-string">"weight"</span>].stride()
	) <span class="hljs-comment"># (2880, 1, 960, 320) having a stride of 1 for the 2nd dimension proves that it works</span><!-- HTML_TAG_END --></pre></div> <h2 class="relative group"><a id="tracing" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#tracing"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Tracing</span></h2> <p data-svelte-h="svelte-1hj6sqj">Tracing runs an example input tensor through the model and captures the operations that are performed on it as that input makes its way through the model’s layers. The executable or <code>ScriptFunction</code> that is returned is optimized with just-in-time compilation.</p> <p data-svelte-h="svelte-1ca1ko1">To trace a UNet:</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> time
	<span class="hljs-keyword">import</span> torch
	<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> StableDiffusionPipeline
	<span class="hljs-keyword">import</span> functools

	<span class="hljs-comment"># torch disable grad</span>
	torch.set_grad_enabled(<span class="hljs-literal">False</span>)

	<span class="hljs-comment"># set variables</span>
	n_experiments = <span class="hljs-number">2</span>
	unet_runs_per_experiment = <span class="hljs-number">50</span>


	<span class="hljs-comment"># load inputs</span>
	<span class="hljs-keyword">def</span> <span class="hljs-title function_">generate_inputs</span>():
	sample = torch.randn((<span class="hljs-number">2</span>, <span class="hljs-number">4</span>, <span class="hljs-number">64</span>, <span class="hljs-number">64</span>), device=<span class="hljs-string">"cuda"</span>, dtype=torch.float16)
	timestep = torch.rand(<span class="hljs-number">1</span>, device=<span class="hljs-string">"cuda"</span>, dtype=torch.float16) * <span class="hljs-number">999</span>
	encoder_hidden_states = torch.randn((<span class="hljs-number">2</span>, <span class="hljs-number">77</span>, <span class="hljs-number">768</span>), device=<span class="hljs-string">"cuda"</span>, dtype=torch.float16)
	<span class="hljs-keyword">return</span> sample, timestep, encoder_hidden_states


	pipe = StableDiffusionPipeline.from_pretrained(
	<span class="hljs-string">"stable-diffusion-v1-5/stable-diffusion-v1-5"</span>,
	torch_dtype=torch.float16,
	use_safetensors=<span class="hljs-literal">True</span>,
	).to(<span class="hljs-string">"cuda"</span>)
	unet = pipe.unet
	unet.<span class="hljs-built_in">eval</span>()
	unet.to(memory_format=torch.channels_last) <span class="hljs-comment"># use channels_last memory format</span>
	unet.forward = functools.partial(unet.forward, return_dict=<span class="hljs-literal">False</span>) <span class="hljs-comment"># set return_dict=False as default</span>

	<span class="hljs-comment"># warmup</span>
	<span class="hljs-keyword">for</span> _ <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(<span class="hljs-number">3</span>):
	<span class="hljs-keyword">with</span> torch.inference_mode():
	inputs = generate_inputs()
	orig_output = unet(*inputs)

	<span class="hljs-comment"># trace</span>
	<span class="hljs-built_in">print</span>(<span class="hljs-string">"tracing.."</span>)
	unet_traced = torch.jit.trace(unet, inputs)
	unet_traced.<span class="hljs-built_in">eval</span>()
	<span class="hljs-built_in">print</span>(<span class="hljs-string">"done tracing"</span>)


	<span class="hljs-comment"># warmup and optimize graph</span>
	<span class="hljs-keyword">for</span> _ <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(<span class="hljs-number">5</span>):
	<span class="hljs-keyword">with</span> torch.inference_mode():
	inputs = generate_inputs()
	orig_output = unet_traced(*inputs)


	<span class="hljs-comment"># benchmarking</span>
	<span class="hljs-keyword">with</span> torch.inference_mode():
	<span class="hljs-keyword">for</span> _ <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(n_experiments):
	torch.cuda.synchronize()
	start_time = time.time()
	<span class="hljs-keyword">for</span> _ <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(unet_runs_per_experiment):
	orig_output = unet_traced(*inputs)
	torch.cuda.synchronize()
	<span class="hljs-built_in">print</span>(<span class="hljs-string">f"unet traced inference took <span class="hljs-subst">{time.time() - start_time:<span class="hljs-number">.2</span>f}</span> seconds"</span>)
	<span class="hljs-keyword">for</span> _ <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(n_experiments):
	torch.cuda.synchronize()
	start_time = time.time()
	<span class="hljs-keyword">for</span> _ <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(unet_runs_per_experiment):
	orig_output = unet(*inputs)
	torch.cuda.synchronize()
	<span class="hljs-built_in">print</span>(<span class="hljs-string">f"unet inference took <span class="hljs-subst">{time.time() - start_time:<span class="hljs-number">.2</span>f}</span> seconds"</span>)

	<span class="hljs-comment"># save the model</span>
	unet_traced.save(<span class="hljs-string">"unet_traced.pt"</span>)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-mpebhd">Replace the <code>unet</code> attribute of the pipeline with the traced model:</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> StableDiffusionPipeline
	<span class="hljs-keyword">import</span> torch
	<span class="hljs-keyword">from</span> dataclasses <span class="hljs-keyword">import</span> dataclass


	<span class="hljs-meta">@dataclass</span>
	<span class="hljs-keyword">class</span> <span class="hljs-title class_">UNet2DConditionOutput</span>:
	sample: torch.Tensor


	pipe = StableDiffusionPipeline.from_pretrained(
	<span class="hljs-string">"stable-diffusion-v1-5/stable-diffusion-v1-5"</span>,
	torch_dtype=torch.float16,
	use_safetensors=<span class="hljs-literal">True</span>,
	).to(<span class="hljs-string">"cuda"</span>)

	<span class="hljs-comment"># use jitted unet</span>
	unet_traced = torch.jit.load(<span class="hljs-string">"unet_traced.pt"</span>)


	<span class="hljs-comment"># del pipe.unet</span>
	<span class="hljs-keyword">class</span> <span class="hljs-title class_">TracedUNet</span>(torch.nn.Module):
	<span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self</span>):
	<span class="hljs-built_in">super</span>().__init__()
	self.in_channels = pipe.unet.config.in_channels
	self.device = pipe.unet.device

	<span class="hljs-keyword">def</span> <span class="hljs-title function_">forward</span>(<span class="hljs-params">self, latent_model_input, t, encoder_hidden_states</span>):
	sample = unet_traced(latent_model_input, t, encoder_hidden_states)[<span class="hljs-number">0</span>]
	<span class="hljs-keyword">return</span> UNet2DConditionOutput(sample=sample)


	pipe.unet = TracedUNet()

	<span class="hljs-keyword">with</span> torch.inference_mode():
	image = pipe([prompt] * <span class="hljs-number">1</span>, num_inference_steps=<span class="hljs-number">50</span>).images[<span class="hljs-number">0</span>]<!-- HTML_TAG_END --></pre></div> <h2 class="relative group"><a id="memory-efficient-attention" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#memory-efficient-attention"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Memory-efficient attention</span></h2> <p data-svelte-h="svelte-1tu2ji2">Recent work on optimizing bandwidth in the attention block has generated huge speed-ups and reductions in GPU memory usage. The most recent type of memory-efficient attention is <a href="https://arxiv.org/abs/2205.14135" rel="nofollow">Flash Attention</a> (you can check out the original code at <a href="https://github.com/HazyResearch/flash-attention" rel="nofollow">HazyResearch/flash-attention</a>).</p> <div class="course-tip bg-gradient-to-br dark:bg-gradient-to-r before:border-green-500 dark:before:border-green-800 from-green-50 dark:from-gray-900 to-white dark:to-gray-950 border border-green-50 text-green-700 dark:text-gray-400"><p data-svelte-h="svelte-17vr965">If you have PyTorch >= 2.0 installed, you should not expect a speed-up for inference when enabling <code>xformers</code>.</p></div> <p data-svelte-h="svelte-rpfoyf">To use Flash Attention, install the following:</p> <ul data-svelte-h="svelte-1q6gndt"><li>PyTorch > 1.12</li> <li>CUDA available</li> <li><a href="xformers">xFormers</a></li></ul> <p data-svelte-h="svelte-122gdhe">Then call <a href="/docs/diffusers/pr_11234/en/api/models/overview#diffusers.ModelMixin.enable_xformers_memory_efficient_attention">enable_xformers_memory_efficient_attention()</a> on the pipeline:</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class=""><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> DiffusionPipeline
	<span class="hljs-keyword">import</span> torch

	pipe = DiffusionPipeline.from_pretrained(
	<span class="hljs-string">"stable-diffusion-v1-5/stable-diffusion-v1-5"</span>,
	torch_dtype=torch.float16,
	use_safetensors=<span class="hljs-literal">True</span>,
	).to(<span class="hljs-string">"cuda"</span>)

	pipe.enable_xformers_memory_efficient_attention()

	<span class="hljs-keyword">with</span> torch.inference_mode():
	sample = pipe(<span class="hljs-string">"a small cat"</span>)

	<span class="hljs-comment"># optional: You can disable it via</span>
	<span class="hljs-comment"># pipe.disable_xformers_memory_efficient_attention()</span><!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1nf2kxh">The iteration speed when using <code>xformers</code> should match the iteration speed of PyTorch 2.0 as described <a href="torch2.0">here</a>.</p> <a class="!text-gray-400 !no-underline text-sm flex items-center not-prose mt-4" href="https://github.com/huggingface/diffusers/blob/main/docs/source/en/optimization/memory.md" target="_blank"><span data-svelte-h="svelte-1kd6by1"><</span> <span data-svelte-h="svelte-x0xyl0">></span> <span data-svelte-h="svelte-1dajgef"><span class="underline ml-1.5">Update</span> on GitHub</span></a> <p></p>

	<script>
	{
	__sveltekit_225bzv = {
	assets: "/docs/diffusers/pr_11234/en",
	base: "/docs/diffusers/pr_11234/en",
	env: {}
	};

	const element = document.currentScript.parentElement;

	const data = [null,null];

	Promise.all([
	import("/docs/diffusers/pr_11234/en/_app/immutable/entry/start.5a617280.js"),
	import("/docs/diffusers/pr_11234/en/_app/immutable/entry/app.1a1de8b6.js")
	]).then(([kit, app]) => {
	kit.start(app, element, {
	node_ids: [0, 214],
	data,
	form: null,
	error: null
	});
	});
	}
	</script>

Xet Storage Details

Size:: 64 kB
Xet hash:: 42f7220d212b9808f7710d74d1c5d7000b85dec276166fe0d0d9210e94750244

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.