Buckets:
| <meta charset="utf-8" /><meta name="hf:doc:metadata" content="{"title":"减少内存使用","local":"reduce-memory-usage","sections":[{"title":"多 GPU","local":"multiple-gpus","sections":[{"title":"分片检查点","local":"sharded-checkpoints","sections":[],"depth":3},{"title":"设备放置","local":"device-placement","sections":[],"depth":3}],"depth":2},{"title":"VAE 切片","local":"vae-slicing","sections":[],"depth":2},{"title":"VAE 分块","local":"vae-tiling","sections":[],"depth":2},{"title":"卸载","local":"offloading","sections":[{"title":"CPU 卸载","local":"cpu-offloading","sections":[],"depth":3},{"title":"模型卸载","local":"model-offloading","sections":[],"depth":3},{"title":"组卸载","local":"group-offloading","sections":[{"title":"CUDA 流","local":"cuda-stream","sections":[],"depth":4},{"title":"卸载到磁盘","local":"offloading-to-disk","sections":[],"depth":4}],"depth":3}],"depth":2},{"title":"逐层类型转换","local":"layerwise-casting","sections":[],"depth":2},{"title":"torch.channels_last","local":"torchchannelslast","sections":[],"depth":2},{"title":"内存高效注意力","local":"memory-efficient-attention","sections":[],"depth":2}],"depth":1}"/> | |
| <link href="/docs/diffusers/main/en/_app/immutable/entry/start.0S0_iEoG.js" rel="modulepreload"> | |
| <link href="/docs/diffusers/main/en/_app/immutable/chunks/I9eooN77.js" rel="modulepreload"> | |
| <link href="/docs/diffusers/main/en/_app/immutable/chunks/Bupkjr9f.js" rel="modulepreload"> | |
| <link href="/docs/diffusers/main/en/_app/immutable/entry/app.CPdOxCbw.js" rel="modulepreload"> | |
| <link href="/docs/diffusers/main/en/_app/immutable/chunks/CyvZ3FS5.js" rel="modulepreload"> | |
| <link href="/docs/diffusers/main/en/_app/immutable/chunks/l0QnnOED.js" rel="modulepreload"> | |
| <link href="/docs/diffusers/main/en/_app/immutable/chunks/DnVIJfC5.js" rel="modulepreload"> | |
| <link href="/docs/diffusers/main/en/_app/immutable/chunks/DsnmJJEf.js" rel="modulepreload"> | |
| <link href="/docs/diffusers/main/en/_app/immutable/nodes/0.DteTvmQZ.js" rel="modulepreload"> | |
| <link href="/docs/diffusers/main/en/_app/immutable/chunks/CoSrdQax.js" rel="modulepreload"> | |
| <link href="/docs/diffusers/main/en/_app/immutable/nodes/339.DJYaCibf.js" rel="modulepreload"> | |
| <link href="/docs/diffusers/main/en/_app/immutable/chunks/DSPDl1gz.js" rel="modulepreload"> | |
| <link href="/docs/diffusers/main/en/_app/immutable/chunks/D5eijdnr.js" rel="modulepreload"> | |
| <!--13d35im--><meta name="hf:doc:metadata" content="{"title":"Reduce memory usage","local":"reduce-memory-usage","sections":[{"title":"Multiple GPUs","local":"multiple-gpus","sections":[{"title":"Sharded checkpoints","local":"sharded-checkpoints","sections":[],"depth":3},{"title":"Device placement","local":"device-placement","sections":[],"depth":3}],"depth":2},{"title":"VAE slicing","local":"vae-slicing","sections":[],"depth":2},{"title":"VAE tiling","local":"vae-tiling","sections":[],"depth":2},{"title":"Offloading","local":"offloading","sections":[{"title":"CPU offloading","local":"cpu-offloading","sections":[],"depth":3},{"title":"Model offloading","local":"model-offloading","sections":[],"depth":3},{"title":"Group offloading","local":"group-offloading","sections":[{"title":"CUDA stream","local":"cuda-stream","sections":[],"depth":4},{"title":"Offloading to disk","local":"offloading-to-disk","sections":[],"depth":4}],"depth":3}],"depth":2},{"title":"Layerwise casting","local":"layerwise-casting","sections":[],"depth":2},{"title":"torch.channels_last","local":"torchchannelslast","sections":[],"depth":2},{"title":"Memory-efficient attention","local":"memory-efficient-attention","sections":[],"depth":2}],"depth":1}"/><!----> | |
| <link href="/docs/diffusers/main/en/_app/immutable/assets/0.tn0RQdqM.css" rel="modulepreload"> <!--[--><!--[0--><!--[--><!--[0--><!--[--><p></p> <div class="items-center shrink-0 min-w-[100px] max-sm:min-w-[50px] justify-end ml-auto flex" style="float: inline-end; margin-inline-start: 10px; display: inline-flex; position: relative; z-index: 10;"><div dir="ltr" class="inline-flex rounded-md max-sm:rounded-sm"><button class="inline-flex items-center gap-1 h-7 max-sm:h-7 px-2 max-sm:px-1.5 text-sm font-medium text-gray-800 border border-r-0 rounded-l-md max-sm:rounded-l-sm border-gray-200 bg-white hover:shadow-inner dark:border-gray-850 dark:bg-gray-950 dark:text-gray-200 dark:hover:bg-gray-800" aria-live="polite"><span class="inline-flex items-center justify-center rounded-md p-0.5 max-sm:p-0 hover:text-gray-800 dark:hover:text-gray-200"><svg class="sm:size-3.5 size-3" 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><!----></span> <span dir="auto">Copy page</span></button> <button class="inline-flex items-center justify-center w-6 max-sm:w-5 h-7 max-sm:h-7 disabled:pointer-events-none text-sm text-gray-500 hover:text-gray-700 dark:hover:text-white rounded-r-md max-sm:rounded-r-sm border border-l transition border-gray-200 bg-white hover:shadow-inner dark:border-gray-850 dark:bg-gray-950 dark:text-gray-200 dark:hover:bg-gray-800" aria-haspopup="menu" aria-expanded="false" aria-label="Open copy menu"><svg class="transition-transform text-gray-400 overflow-visible sm:size-3.5 size-3 rotate-0" width="1em" height="1em" viewBox="0 0 12 7" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M1 1L6 6L11 1" stroke="currentColor"></path></svg><!----></button></div> <!--[-1--><!--]--></div><!----> <!--[0--><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 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><!---->减少内存使用<!----></span></h1><!--]--><!----> <p>像 <a href="../api/pipelines/flux">Flux</a> 和 <a href="../api/pipelines/wan">Wan</a> 这样的现代扩散模型拥有数十亿参数，在硬件上进行推理时会占用大量内存。由于常见的 GPU 往往内存不足，这带来了挑战。为了克服内存限制，您可以使用多块 GPU（如果可用）、将部分流水线组件卸载到 CPU 上，以及采用其他方法。</p> <p>本指南将向你展示如何降低内存使用量。</p> <blockquote class="tip"><p>请注意，这些技术可能需要根据模型进行调整。例如，基于 Transformer 的扩散模型从这些内存优化中获得的收益，可能不如基于 UNet 的模型。</p></blockquote> <!--[1--><h2 class="relative group"><a id="multiple-gpus" 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="#multiple-gpus"><span><svg 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><!---->多 GPU<!----></span></h2><!--]--><!----> <p>如果您拥有多块 GPU，则有几种选项可以高效地将大型模型加载并分布到您的硬件上。这些功能由 <a href="https://huggingface.co/docs/accelerate/index" rel="nofollow">Accelerate</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 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> 已复制</div><!----></button><!----></div> <pre class="language-bash "><!---->pip install -U accelerate<!----></pre></div><!----> <!--[2--><h3 class="relative group"><a id="sharded-checkpoints" 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="#sharded-checkpoints"><span><svg 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><!---->分片检查点<!----></span></h3><!--]--><!----> <p>分片加载大型检查点很有用，因为分片是逐个加载的。这样可以保持内存占用较低，只需满足模型大小和最大分片大小所需的内存即可。我们建议当 fp32 检查点大于 5GB 时进行分片。默认的分片大小为 5GB。</p> <p>使用 <code>max_shard_size</code> 参数在 <a href="/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.save_pretrained">save_pretrained()</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> AutoModel | |
| unet = AutoModel.from_pretrained( | |
| <span class="hljs-string">"stabilityai/stable-diffusion-xl-base-1.0"</span>, subfolder=<span class="hljs-string">"unet"</span> | |
| ) | |
| unet.save_pretrained(<span class="hljs-string">"sdxl-unet-sharded"</span>, max_shard_size=<span class="hljs-string">"5GB"</span>)<!----></pre></div><!----> <p>现在，您可以使用分片检查点（sharded checkpoint）来替代常规检查点，以节省内存。</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> AutoModel, StableDiffusionXLPipeline | |
| unet = AutoModel.from_pretrained( | |
| <span class="hljs-string">"username/sdxl-unet-sharded"</span>, dtype=torch.float16 | |
| ) | |
| pipeline = StableDiffusionXLPipeline.from_pretrained( | |
| <span class="hljs-string">"stabilityai/stable-diffusion-xl-base-1.0"</span>, | |
| unet=unet, | |
| dtype=torch.float16 | |
| ).to(<span class="hljs-string">"cuda"</span>) <span class="hljs-comment"># or "mps", "xpu", "cpu"</span><!----></pre></div><!----> <!--[2--><h3 class="relative group"><a id="device-placement" 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="#device-placement"><span><svg 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><!---->设备放置<!----></span></h3><!--]--><!----> <blockquote class="warning"><p>设备放置是一项实验性功能，API 可能会发生变化。目前仅支持 <code>balanced</code> 策略。我们计划在未来支持更多的映射策略。</p></blockquote> <p><code>device_map</code> 参数用于控制流水线中的模型组件或单个模型中的层如何在不同设备上进行分配。</p> <div class="flex space-x-2 items-center my-1.5 mr-8 h-7 !pl-0 -mx-3 md:mx-0"><!--[--><div class="flex items-center border rounded-lg px-1.5 py-1 leading-none select-none text-smd border-gray-800 bg-black dark:bg-gray-700 text-white">流水线级别</div><div class="flex items-center border rounded-lg px-1.5 py-1 leading-none select-none text-smd text-gray-500 cursor-pointer opacity-90 hover:text-gray-700 dark:hover:text-gray-200 hover:shadow-sm">模型级别</div><!--]--></div> <div class="language-select"><div><p><code>balanced</code> 设备放置策略会将流水线均匀地拆分到所有可用设备上。</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> AutoModel, StableDiffusionXLPipeline | |
| pipeline = StableDiffusionXLPipeline.from_pretrained( | |
| <span class="hljs-string">"stabilityai/stable-diffusion-xl-base-1.0"</span>, | |
| dtype=torch.float16, | |
| device_map=<span class="hljs-string">"balanced"</span> | |
| )<!----></pre></div><!----> <p>您可以通过 <code>hf_device_map</code> 检查 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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-built_in">print</span>(pipeline.hf_device_map) | |
| {<span class="hljs-string">'unet'</span>: <span class="hljs-number">1</span>, <span class="hljs-string">'vae'</span>: <span class="hljs-number">1</span>, <span class="hljs-string">'safety_checker'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'text_encoder'</span>: <span class="hljs-number">0</span>}<!----></pre></div><!----><!----></div><!----> <div hidden=""><p><code>device_map</code> 适用于加载大型模型，例如拥有 12.5B 参数的 Flux 扩散 Transformer。将其设置为 <code>"auto"</code> 可自动将模型优先分配到最快的设备上，然后再分配到较慢的设备。有关更多详细信息，请参阅 <a href="#sharded-checkpoints">分片检查点</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> AutoModel | |
| transformer = AutoModel.from_pretrained( | |
| <span class="hljs-string">"black-forest-labs/FLUX.1-dev"</span>, | |
| subfolder=<span class="hljs-string">"transformer"</span>, | |
| device_map=<span class="hljs-string">"auto"</span>, | |
| dtype=torch.bfloat16 | |
| )<!----></pre></div><!----> <p>您可以通过 <code>hf_device_map</code> 检查模型的设备映射。</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-built_in">print</span>(transformer.hf_device_map)<!----></pre></div><!----><!----></div><!----><!----></div><!----> <p>在设计您自己的 <code>device_map</code> 时，它应该是一个字典，其中包含模型特定的模块名称或层以及设备标识符（GPU 使用整数，CPU 使用 <code>cpu</code>，磁盘使用 <code>disk</code>）。</p> <p>对模型调用 <code>hf_device_map</code> 以查看模型层的分布情况，然后自行设计。</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-built_in">print</span>(transformer.hf_device_map) | |
| {<span class="hljs-string">'pos_embed'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'time_text_embed'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'context_embedder'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'x_embedder'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'transformer_blocks'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.0'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.1'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.2'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.3'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.4'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.5'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.6'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.7'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.8'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.9'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.10'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.11'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.12'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.13'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.14'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.15'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.16'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.17'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.18'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.19'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.20'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.21'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.22'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.23'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.24'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.25'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.26'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.27'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.28'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.29'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.30'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.31'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.32'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.33'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.34'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.35'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.36'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.37'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'norm_out'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'proj_out'</span>: <span class="hljs-string">'cpu'</span>}<!----></pre></div><!----> <p>例如，下面的 <code>device_map</code> 将 <code>single_transformer_blocks.10</code> 到 <code>single_transformer_blocks.20</code> 放置在第二块 GPU（<code>1</code>）上。</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> AutoModel | |
| device_map = { | |
| <span class="hljs-string">'pos_embed'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'time_text_embed'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'context_embedder'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'x_embedder'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'transformer_blocks'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.0'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.1'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.2'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.3'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.4'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.5'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.6'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.7'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.8'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.9'</span>: <span class="hljs-number">0</span>, <span class="hljs-string">'single_transformer_blocks.10'</span>: <span class="hljs-number">1</span>, <span class="hljs-string">'single_transformer_blocks.11'</span>: <span class="hljs-number">1</span>, <span class="hljs-string">'single_transformer_blocks.12'</span>: <span class="hljs-number">1</span>, <span class="hljs-string">'single_transformer_blocks.13'</span>: <span class="hljs-number">1</span>, <span class="hljs-string">'single_transformer_blocks.14'</span>: <span class="hljs-number">1</span>, <span class="hljs-string">'single_transformer_blocks.15'</span>: <span class="hljs-number">1</span>, <span class="hljs-string">'single_transformer_blocks.16'</span>: <span class="hljs-number">1</span>, <span class="hljs-string">'single_transformer_blocks.17'</span>: <span class="hljs-number">1</span>, <span class="hljs-string">'single_transformer_blocks.18'</span>: <span class="hljs-number">1</span>, <span class="hljs-string">'single_transformer_blocks.19'</span>: <span class="hljs-number">1</span>, <span class="hljs-string">'single_transformer_blocks.20'</span>: <span class="hljs-number">1</span>, <span class="hljs-string">'single_transformer_blocks.21'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.22'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.23'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.24'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.25'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.26'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.27'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.28'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.29'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.30'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.31'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.32'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.33'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.34'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.35'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.36'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'single_transformer_blocks.37'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'norm_out'</span>: <span class="hljs-string">'cpu'</span>, <span class="hljs-string">'proj_out'</span>: <span class="hljs-string">'cpu'</span> | |
| } | |
| transformer = AutoModel.from_pretrained( | |
| <span class="hljs-string">"black-forest-labs/FLUX.1-dev"</span>, | |
| subfolder=<span class="hljs-string">"transformer"</span>, | |
| device_map=device_map, | |
| dtype=torch.bfloat16 | |
| )<!----></pre></div><!----> <p>传入一个将最大内存使用量映射到每个设备的字典以强制实施限制。如果某个设备不在 <code>max_memory</code> 中，则会被忽略，且流水线组件不会分发到该设备上。</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> AutoModel, StableDiffusionXLPipeline | |
| max_memory = {<span class="hljs-number">0</span>:<span class="hljs-string">"1GB"</span>, <span class="hljs-number">1</span>:<span class="hljs-string">"1GB"</span>} | |
| pipeline = StableDiffusionXLPipeline.from_pretrained( | |
| <span class="hljs-string">"stabilityai/stable-diffusion-xl-base-1.0"</span>, | |
| dtype=torch.float16, | |
| device_map=<span class="hljs-string">"balanced"</span>, | |
| max_memory=max_memory | |
| )<!----></pre></div><!----> <p>Diffusers 默认使用所有设备的最大内存，但如果模型无法装入 GPU，则需要使用单个 GPU，并通过以下方法将部分计算卸载到 CPU。</p> <ul><li><a href="/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload">enable_model_cpu_offload()</a> 仅适用于单 GPU，但非常大的模型可能无法装入单个 GPU</li> <li><a href="/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_sequential_cpu_offload">enable_sequential_cpu_offload()</a> 可能可以运行，但速度极慢，且仅限于单 GPU</li></ul> <p>使用 <a href="/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.reset_device_map">reset_device_map()</a> 方法重置 <code>device_map</code>。如果您想在已进行设备映射的 pipeline 上使用诸如 <code>.to()</code>、<a href="/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_sequential_cpu_offload">enable_sequential_cpu_offload()</a> 和 <a href="/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload">enable_model_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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!---->pipeline.reset_device_map()<!----></pre></div><!----> <!--[1--><h2 class="relative group"><a id="vae-slicing" 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="#vae-slicing"><span><svg 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><!---->VAE 切片<!----></span></h2><!--]--><!----> <p>VAE 切片通过将一个大批次输入拆分为单个数据批次并分别处理来节省内存。此方法在同时生成多张图像时效果最佳。</p> <p>例如，如果您一次性生成 4 张图像，解码过程会使峰值激活内存增加 4 倍。VAE 切片通过每次仅解码 1 张图像（而非同时解码全部 4 张）来降低这一开销。</p> <p>调用 <code>enable_slicing()</code> 以启用切片式 VAE。在解码多图像批次时，性能会有小幅提升；对于单图像批次，则不会有性能影响。</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> AutoModel, StableDiffusionXLPipeline | |
| pipeline = StableDiffusionXLPipeline.from_pretrained( | |
| <span class="hljs-string">"stabilityai/stable-diffusion-xl-base-1.0"</span>, | |
| dtype=torch.float16, | |
| ).to(<span class="hljs-string">"cuda"</span>) <span class="hljs-comment"># or "mps", "xpu", "cpu"</span> | |
| pipeline.vae.enable_slicing() | |
| pipeline([<span class="hljs-string">"An astronaut riding a horse on Mars"</span>]*<span class="hljs-number">32</span>).images[<span class="hljs-number">0</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>)<!----></pre></div><!----> <blockquote class="warning"><p><a href="/docs/diffusers/v0.43.0/zh/api/models/autoencoder_kl_wan#diffusers.AutoencoderKLWan">AutoencoderKLWan</a> 和 <a href="/docs/diffusers/v0.43.0/zh/api/models/asymmetricautoencoderkl#diffusers.AsymmetricAutoencoderKL">AsymmetricAutoencoderKL</a> 类不支持切片。</p></blockquote> <!--[1--><h2 class="relative group"><a id="vae-tiling" 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="#vae-tiling"><span><svg 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><!---->VAE 分块<!----></span></h2><!--]--><!----> <p>VAE 分块（tiling）通过将图像划分为较小的重叠块，而非一次性处理整张图像，从而节省内存。由于 GPU 每次仅处理一个块，这也有助于降低峰值内存占用。</p> <p>调用 <code>enable_tiling()</code> 以启用 VAE 分块。由于各分块是独立解码的，生成的图像在分块之间可能会出现一些色调差异，但分块之间不应出现明显的接缝。对于低于预设（但可配置）阈值的分辨率，分块功能将被禁用。例如，<a href="/docs/diffusers/v0.43.0/zh/api/pipelines/stable_diffusion/text2img#diffusers.StableDiffusionPipeline">StableDiffusionPipeline</a> 中 VAE 的该阈值为 512x512。</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> AutoPipelineForImage2Image | |
| <span class="hljs-keyword">from</span> diffusers.utils <span class="hljs-keyword">import</span> load_image | |
| pipeline = AutoPipelineForImage2Image.from_pretrained( | |
| <span class="hljs-string">"stabilityai/stable-diffusion-xl-base-1.0"</span>, dtype=torch.float16 | |
| ).to(<span class="hljs-string">"cuda"</span>) <span class="hljs-comment"># or "mps", "xpu", "cpu"</span> | |
| pipeline.vae.enable_tiling() | |
| init_image = load_image(<span class="hljs-string">"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-sdxl-init.png"</span>) | |
| prompt = <span class="hljs-string">"Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"</span> | |
| pipeline(prompt, image=init_image, strength=<span class="hljs-number">0.5</span>).images[<span class="hljs-number">0</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>)<!----></pre></div><!----> <blockquote class="warning"><p><a href="/docs/diffusers/v0.43.0/zh/api/models/autoencoder_kl_wan#diffusers.AutoencoderKLWan">AutoencoderKLWan</a> 和 <a href="/docs/diffusers/v0.43.0/zh/api/models/asymmetricautoencoderkl#diffusers.AsymmetricAutoencoderKL">AsymmetricAutoencoderKL</a> 不支持分块处理。</p></blockquote> <p>若要完全避免在 GPU 上占用 VAE 内存，请使用 <a href="../hybrid_inference/overview">Remote inference</a> 在远程端点上执行编码和解码。</p> <!--[1--><h2 class="relative group"><a id="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="#offloading"><span><svg 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><!---->卸载<!----></span></h2><!--]--><!----> <p>Offloading 策略会将当前未激活的层或模型移至 CPU，以避免增加 GPU 内存占用。这些策略可以与量化和 torch.compile 结合使用，以平衡推理速度和内存使用。</p> <p>有关更多详细信息，请参阅<a href="./speed-memory-optims">编译和卸载量化模型</a>指南。</p> <!--[2--><h3 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 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 卸载<!----></span></h3><!--]--><!----> <p>CPU offloading 会将权重从 GPU 有选择地移至 CPU。当某个组件被需要时，它会被传输到 GPU；当不再需要时，它会被移回 CPU。该方法针对子模块而非整个模型进行操作。通过避免在 GPU 上存储整个模型，此方法可以节省内存。</p> <p>CPU offloading 能大幅降低内存占用，但由于子模块需要在设备之间多次来回传递，其速度<strong>极其缓慢</strong>。因此，在实际应用中往往难以采用。</p> <blockquote class="warning"><p>在调用 <a href="/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_sequential_cpu_offload">enable_sequential_cpu_offload()</a> 之前，请勿将 pipeline 移至 CUDA，否则节省的内存量微乎其微（更多详情请参见此 <a href="https://github.com/huggingface/diffusers/issues/1934" rel="nofollow">issue</a>）。这是一个有状态的操作，会在模型上安装钩子。</p></blockquote> <p>调用 <a href="/docs/diffusers/v0.43.0/zh/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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> DiffusionPipeline | |
| pipeline = DiffusionPipeline.from_pretrained( | |
| <span class="hljs-string">"black-forest-labs/FLUX.1-schnell"</span>, dtype=torch.bfloat16 | |
| ) | |
| pipeline.enable_sequential_cpu_offload() | |
| pipeline( | |
| prompt=<span class="hljs-string">"An astronaut riding a horse on Mars"</span>, | |
| guidance_scale=<span class="hljs-number">0.</span>, | |
| height=<span class="hljs-number">768</span>, | |
| width=<span class="hljs-number">1360</span>, | |
| num_inference_steps=<span class="hljs-number">4</span>, | |
| max_sequence_length=<span class="hljs-number">256</span>, | |
| ).images[<span class="hljs-number">0</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>)<!----></pre></div><!----> <!--[2--><h3 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 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><!---->模型卸载<!----></span></h3><!--]--><!----> <p>模型卸载（Model offloading）会将整个模型移至 GPU，而不是选择性地移动<em>部分</em>层或模型组件。其中一个主要的流水线模型（通常是文本编码器、UNet 和 VAE）会被放置在 GPU 上，而其他组件则保留在 CPU 上。像 UNet 这样需要多次运行的组件会一直留在 GPU 上，直到其完全执行完毕且不再需要为止。这消除了<a href="#cpu-offloading">CPU 卸载</a>带来的通信开销，使模型卸载成为一种更快的替代方案。其代价是内存节省幅度不会那么大。</p> <blockquote class="warning"><p>请注意，如果在安装钩子之后，模型在流水线之外被复用（有关更多详细信息，请参阅 <a href="https://huggingface.co/docs/accelerate/en/package_reference/big_modeling#accelerate.hooks.remove_hook_from_module" rel="nofollow">移除钩子</a>），则需要按照预期的顺序运行整个流水线及模型，才能正确卸载它们。这是一个有状态的操作，会在模型上安装钩子。</p></blockquote> <p>调用 <a href="/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload">enable_model_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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> DiffusionPipeline | |
| pipeline = DiffusionPipeline.from_pretrained( | |
| <span class="hljs-string">"black-forest-labs/FLUX.1-schnell"</span>, dtype=torch.bfloat16 | |
| ) | |
| pipeline.enable_model_cpu_offload() | |
| pipeline( | |
| prompt=<span class="hljs-string">"An astronaut riding a horse on Mars"</span>, | |
| guidance_scale=<span class="hljs-number">0.</span>, | |
| height=<span class="hljs-number">768</span>, | |
| width=<span class="hljs-number">1360</span>, | |
| num_inference_steps=<span class="hljs-number">4</span>, | |
| max_sequence_length=<span class="hljs-number">256</span>, | |
| ).images[<span class="hljs-number">0</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>)<!----></pre></div><!----> <p><a href="/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload">enable_model_cpu_offload()</a> 在单独使用 <a href="/docs/diffusers/v0.43.0/zh/api/pipelines/stable_diffusion/stable_diffusion_xl#diffusers.StableDiffusionXLPipeline.encode_prompt">encode_prompt()</a> 方法生成文本编码器的隐藏状态时也有帮助。</p> <!--[2--><h3 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 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><!---->组卸载<!----></span></h3><!--]--><!----> <p>Group offloading 将内部层的组（<a href="https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html" rel="nofollow">torch.nn.ModuleList</a> 或 <a href="https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html" rel="nofollow">torch.nn.Sequential</a>）移至 CPU。与 <a href="#model-offloading">model offloading</a> 相比，它占用的内存更少；与 <a href="#cpu-offloading">CPU offloading</a> 相比，它速度更快，因为它降低了通信开销。</p> <blockquote class="warning"><p>如果前向实现中包含依赖于权重的输入设备转换，则组卸载可能无法在所有模型上正常工作，因为这可能会与组卸载的设备转换机制发生冲突。</p></blockquote> <p>通过将 <code>offload_type</code> 参数配置为 <code>block_level</code> 或 <code>leaf_level</code> 来启用组卸载。</p> <ul><li><code>block_level</code> 卸载根据 <code>num_blocks_per_group</code> 参数对层进行分组。例如，对于一个具有 40 层的模型，如果 <code>num_blocks_per_group=2</code>，则每次同时加载和卸载 2 层（总共进行 20 次加载/卸载）。这可以大幅降低内存需求。</li> <li><code>leaf_level</code> 在最低层级对单个层进行卸载，其效果等同于 <a href="#cpu-offloading">CPU 卸载</a>。但如果你使用流（streams），可以在不牺牲推理速度的情况下使其运行得更快。</li></ul> <p>组卸载（Group offloading）支持应用于整个流水线或单个模型。将组卸载应用于整个流水线是最简单的选择，而将其选择性地应用于单个模型则能赋予用户更大的灵活性，以便为不同的模型使用不同的卸载技术。</p> <div class="flex space-x-2 items-center my-1.5 mr-8 h-7 !pl-0 -mx-3 md:mx-0"><!--[--><div class="flex items-center border rounded-lg px-1.5 py-1 leading-none select-none text-smd border-gray-800 bg-black dark:bg-gray-700 text-white">pipeline</div><div class="flex items-center border rounded-lg px-1.5 py-1 leading-none select-none text-smd text-gray-500 cursor-pointer opacity-90 hover:text-gray-700 dark:hover:text-gray-200 hover:shadow-sm">model</div><!--]--></div> <div class="language-select"><div><p>在 pipeline 上调用 <a href="/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_group_offload">enable_group_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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><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 | |
| onload_device = torch.device(<span class="hljs-string">"cuda"</span>) | |
| offload_device = torch.device(<span class="hljs-string">"cpu"</span>) | |
| pipeline = CogVideoXPipeline.from_pretrained(<span class="hljs-string">"THUDM/CogVideoX-5b"</span>, dtype=torch.bfloat16) | |
| pipeline.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> | |
| ) | |
| 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 = pipeline(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-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>)<!----></pre></div><!----><!----></div><!----> <div hidden=""><p>在继承自 <a href="/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin">ModelMixin</a> 的标准 Diffusers 模型组件上调用 <a href="/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.enable_group_offload">enable_group_offload()</a>。对于未继承自 <a href="/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin">ModelMixin</a> 的其他模型组件（例如通用的 <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html" rel="nofollow">torch.nn.Module</a>），请改用 <a href="/docs/diffusers/v0.43.0/zh/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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><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 | |
| onload_device = torch.device(<span class="hljs-string">"cuda"</span>) | |
| offload_device = torch.device(<span class="hljs-string">"cpu"</span>) | |
| pipeline = CogVideoXPipeline.from_pretrained(<span class="hljs-string">"THUDM/CogVideoX-5b"</span>, dtype=torch.bfloat16) | |
| <span class="hljs-comment"># Use the enable_group_offload method for Diffusers model implementations</span> | |
| pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type=<span class="hljs-string">"leaf_level"</span>) | |
| pipeline.vae.enable_group_offload(onload_device=onload_device, offload_type=<span class="hljs-string">"leaf_level"</span>) | |
| <span class="hljs-comment"># Use the apply_group_offloading method for other model components</span> | |
| apply_group_offloading(pipeline.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>) | |
| 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 = pipeline(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-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>)<!----></pre></div><!----><!----></div><!----><!----></div><!----> <!--[3--><h4 class="relative group"><a id="cuda-stream" 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="#cuda-stream"><span><svg 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><!---->CUDA 流<!----></span></h4><!--]--><!----> <p><code>use_stream</code> 参数可在支持异步数据传输流的 CUDA 设备上启用，与 <a href="#cpu-offloading">CPU offloading</a> 相比，可降低整体执行时间。它通过层预取来重叠数据传输与计算：在执行当前层的同时，将下一个待执行的层加载到 GPU 上。这可能会显著增加 CPU 内存占用，因此请确保拥有模型大小 2 倍的内存。</p> <p>设置 <code>record_stream=True</code> 可以进一步提速，但会略微增加内存占用。请参阅 <a href="https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html" rel="nofollow">torch.Tensor.record_stream</a> 文档以了解更多信息。</p> <blockquote class="tip"><p>当在启用 tiling 的 VAE 上设置 <code>use_stream=True</code> 时，请在推理前执行一次前向传播（也可以使用虚拟输入），以避免设备不匹配错误。此方法可能不适用于所有实现，因此如果您遇到任何问题，请随时提交 issue。</p></blockquote> <p>如果您在使用 <code>block_level</code> 组卸载且启用了 <code>use_stream</code>，则应将 <code>num_blocks_per_group</code> 参数设置为 <code>1</code>，否则将触发警告。</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!---->pipeline.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>, record_stream=<span class="hljs-literal">True</span>)<!----></pre></div><!----> <p><code>low_cpu_mem_usage</code> 参数可以设置为 <code>True</code>，以在使用流进行组卸载时降低 CPU 内存占用。该设置最适合 <code>leaf_level</code> 卸载以及 CPU 内存成为瓶颈的场景。其通过动态创建固定张量（pinned tensors）而非预先固定来节省内存。不过，这可能会增加整体执行时间。</p> <!--[3--><h4 class="relative group"><a id="offloading-to-disk" 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="#offloading-to-disk"><span><svg 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><!---->卸载到磁盘<!----></span></h4><!--]--><!----> <p>根据模型大小，Group offloading 可能会占用大量系统内存。在内存有限的系统上，可以尝试将 group offloading 到磁盘，将其作为二级内存使用。</p> <p>在 <a href="/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.enable_group_offload">enable_group_offload()</a> 或 <a href="/docs/diffusers/v0.43.0/zh/api/utilities#diffusers.hooks.apply_group_offloading">apply_group_offloading()</a> 中设置 <code>offload_to_disk_path</code> 参数，即可将模型卸载到磁盘。</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!---->pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type=<span class="hljs-string">"leaf_level"</span>, offload_to_disk_path=<span class="hljs-string">"path/to/disk"</span>) | |
| apply_group_offloading(pipeline.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>, offload_to_disk_path=<span class="hljs-string">"path/to/disk"</span>)<!----></pre></div><!----> <p>请参考这<a href="https://github.com/huggingface/diffusers/pull/11682#issue-3129365363" rel="nofollow">两个</a><a href="https://github.com/huggingface/diffusers/pull/11682#issuecomment-2955715126" rel="nofollow">表格</a>，以比较速度和内存之间的权衡。</p> <!--[1--><h2 class="relative group"><a id="layerwise-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="#layerwise-casting"><span><svg 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><!---->逐层类型转换<!----></span></h2><!--]--><!----> <blockquote class="tip"><p>将逐层类型转换与 <a href="#group-offloading">组卸载</a> 结合使用，可进一步节省内存。</p></blockquote> <p>逐层转换（Layerwise casting）将权重存储为更小的数据格式（例如 <code>torch.float8_e4m3fn</code> 和 <code>torch.float8_e5m2</code>）以减少内存占用，并在计算时将权重上转为更高精度，如 <code>torch.float16</code> 或 <code>torch.bfloat16</code>。某些层（与归一化和调制相关的权重）会被跳过，因为以 fp8 存储它们可能会降低生成质量。</p> <blockquote class="warning"><p>如果前向实现中包含对权重的内部类型转换，逐层转换可能无法在所有模型上正常工作。当前逐层转换的实现假设前向传播与权重精度无关，且输入数据类型始终指定为 <code>compute_dtype</code>（参见 <a href="https://github.com/huggingface/transformers/blob/7f5077e53682ca855afc826162b204ebf809f1f9/src/transformers/models/t5/modeling_t5.py#L294-L299" rel="nofollow">此处</a> 了解一个不兼容的实现）。</p> <p>在包含 <a href="https://huggingface.co/docs/peft/index" rel="nofollow">PEFT</a> 层的自定义建模实现中，逐层类型转换也可能失败。虽然存在一些检查机制，但它们并未经过充分测试，也无法保证在所有情况下都能正常工作。</p></blockquote> <p>调用 <a href="/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.enable_layerwise_casting">enable_layerwise_casting()</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><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 | |
| transformer = CogVideoXTransformer3DModel.from_pretrained( | |
| <span class="hljs-string">"THUDM/CogVideoX-5b"</span>, | |
| subfolder=<span class="hljs-string">"transformer"</span>, | |
| dtype=torch.bfloat16 | |
| ) | |
| transformer.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16) | |
| pipeline = CogVideoXPipeline.from_pretrained(<span class="hljs-string">"THUDM/CogVideoX-5b"</span>, | |
| transformer=transformer, | |
| dtype=torch.bfloat16 | |
| ).to(<span class="hljs-string">"cuda"</span>) <span class="hljs-comment"># or "mps", "xpu", "cpu"</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 = pipeline(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-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>)<!----></pre></div><!----> <p>如果需要更多的控制和灵活性，也可以使用 <a href="/docs/diffusers/v0.43.0/zh/api/utilities#diffusers.hooks.apply_layerwise_casting">apply_layerwise_casting()</a> 方法。通过在特定的内部模块上调用该方法，可以将其部分应用于模型层。使用 <code>skip_modules_pattern</code> 或 <code>skip_modules_classes</code> 参数来指定需要避免的模块，例如归一化层和调制层。</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 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> 已复制</div><!----></button><!----></div> <pre class="language-python "><!----><span class="hljs-keyword">import</span> torch | |
| <span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> CogVideoXTransformer3DModel | |
| <span class="hljs-keyword">from</span> diffusers.hooks <span class="hljs-keyword">import</span> apply_layerwise_casting | |
| transformer = CogVideoXTransformer3DModel.from_pretrained( | |
| <span class="hljs-string">"THUDM/CogVideoX-5b"</span>, | |
| subfolder=<span class="hljs-string">"transformer"</span>, | |
| dtype=torch.bfloat16 | |
| ) | |
| <span class="hljs-comment"># skip the normalization layer</span> | |
| apply_layerwise_casting( | |
| transformer, | |
| storage_dtype=torch.float8_e4m3fn, | |
| compute_dtype=torch.bfloat16, | |
| skip_modules_classes=[<span class="hljs-string">"norm"</span>], | |
| non_blocking=<span class="hljs-literal">True</span>, | |
| )<!----></pre></div><!----> <!--[1--><h2 class="relative group"><a id="torchchannelslast" 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="#torchchannelslast"><span><svg 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><!---->torch.channels_last<!----></span></h2><!--]--><!----> <p><a href="https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html" rel="nofollow">torch.channels_last</a> 将张量的存储方式从 <code>(batch size, channels, height, width)</code> 翻转为 <code>(batch size, height, width, channels)</code>。这使得张量的布局与硬件按顺序访问内存中张量的方式保持一致，从而避免了在内存中来回跳跃以访问像素值。</p> <p>并非所有算子目前都支持 channels-last 格式，使用它可能会导致性能下降，但仍值得尝试。</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 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> 已复制</div><!----></button><!----></div> <pre class="language-py "><!----><span class="hljs-built_in">print</span>(pipeline.unet.conv_out.state_dict()[<span class="hljs-string">"weight"</span>].stride()) <span class="hljs-comment"># (2880, 9, 3, 1)</span> | |
| pipeline.unet.to(memory_format=torch.channels_last) <span class="hljs-comment"># in-place operation</span> | |
| <span class="hljs-built_in">print</span>( | |
| pipeline.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><!----></pre></div><!----> <!--[1--><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 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><!---->内存高效注意力<!----></span></h2><!--]--><!----> <p>Diffusers 通过 <a href="/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.set_attention_backend">set_attention_backend()</a> 支持多种内存高效的注意力后端（FlashAttention、xFormers、SageAttention 等）。请参阅 <a href="./attention_backends">Attention backends</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"><svg class="mr-1" 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="M31,16l-7,7l-1.41-1.41L28.17,16l-5.58-5.59L24,9l7,7z"></path><path d="M1,16l7-7l1.41,1.41L3.83,16l5.58,5.59L8,23l-7-7z"></path><path d="M12.419,25.484L17.639,6.552l1.932,0.518L14.351,26.002z"></path></svg><!----> <span><span class="underline">Update</span> on GitHub</span></a><!----> <p></p><!--]--><!----><!--]--><!--]--><!--]--> <!--[-1--><!--]--><!--]--> | |
| <script> | |
| { | |
| __sveltekit_1qsdbwf = { | |
| base: "/docs/diffusers/main/en", | |
| assets: "/docs/diffusers/main/en" | |
| }; | |
| const element = document.currentScript.parentElement; | |
| Promise.all([ | |
| import("/docs/diffusers/main/en/_app/immutable/entry/start.0S0_iEoG.js"), | |
| import("/docs/diffusers/main/en/_app/immutable/entry/app.CPdOxCbw.js") | |
| ]).then(([kit, app]) => { | |
| kit.start(app, element, { | |
| node_ids: [0, 339], | |
| data: [null,null], | |
| form: null, | |
| error: null | |
| }); | |
| }); | |
| } | |
| </script> | |
Xet Storage Details
- Size:
- 99.9 kB
- Xet hash:
- 62cbbff9c5b01657a7102db9bac18dbfa45a7ba46efcb20b4472e395359f1986
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.