diffusers/pr_13881/zh/_app/immutable/nodes/58.CcFf2eVu.js

import"../chunks/DsnmJJEf.js";import{i as E,h as x,C as z,H as i,a as s,b as D,E as Y,s as q}from"../chunks/DdZvggmf.js";import{p as H,o as L,s as l,f as c,a as o,b as O,c as t,n as A}from"../chunks/BbekZcyp.js";import{D as P}from"../chunks/Bje6zuSL.js";import{H as p}from"../chunks/BcnRgdDK.js";const K='{"title":"加载调度器与模型","local":"加载调度器与模型","sections":[{"title":"加载调度器","local":"加载调度器","sections":[],"depth":2},{"title":"调度器对比","local":"调度器对比","sections":[{"title":"Flax调度器","local":"flax调度器","sections":[],"depth":3}],"depth":2},{"title":"模型加载","local":"模型加载","sections":[],"depth":2}],"depth":1}';var $=t('<meta name="hf:doc:metadata"/>'),ll=t("<p><code>LMSDiscreteScheduler</code>通常能生成比默认调度器更高质量的图像。</p> <!>",1),sl=t("<p><code>EulerDiscreteScheduler</code>仅需30步即可生成高质量图像。</p> <!>",1),el=t("<p><code>EulerAncestralDiscreteScheduler</code>同样可在30步内生成高质量图像。</p> <!>",1),al=t("<p><code>DPMSolverMultistepScheduler</code>在速度与质量间取得平衡，仅需20步即可生成优质图像。</p> <!>",1),nl=t("<!> <!> <!> <!>",1),ol=t('<p></p> <!> <!> <!> <p>Diffusion管道是由可互换的调度器(schedulers)和模型(models)组成的集合，可通过混合搭配来定制特定用例的流程。调度器封装了整个去噪过程（如去噪步数和寻找去噪样本的算法），其本身不包含可训练参数，因此内存占用极低。模型则主要负责从含噪输入到较纯净样本的前向传播过程。</p> <p>本指南将展示如何加载调度器和模型来自定义流程。我们将全程使用<a href="https://hf.co/stable-diffusion-v1-5/stable-diffusion-v1-5" rel="nofollow">stable-diffusion-v1-5/stable-diffusion-v1-5</a>检查点，首先加载基础管道：</p> <!> <p>通过<code>pipeline.scheduler</code>属性可查看当前管道使用的调度器：</p> <!> <!> <p>调度器通过配置文件定义，同一配置文件可被多种调度器共享。使用<code>SchedulerMixin.from_pretrained()</code>方法加载时，需指定<code>subfolder</code>参数以定位配置文件在仓库中的正确子目录。</p> <p>例如加载<code>DDIMScheduler</code>：</p> <!> <p>然后将新调度器传入管道：</p> <!> <!> <p>不同调度器各有优劣，难以定量评估哪个最适合您的流程。通常需要在去噪速度与质量之间权衡。我们建议尝试多种调度器以找到最佳方案。通过<code>pipeline.scheduler.compatibles</code>属性可查看兼容当前管道的所有调度器。</p> <p>下面我们使用相同提示词和随机种子，对比<code>LMSDiscreteScheduler</code>、<code>EulerDiscreteScheduler</code>、<code>EulerAncestralDiscreteScheduler</code>和<code>DPMSolverMultistepScheduler</code>的表现：</p> <!> <p>使用<code>from_config()</code>方法加载不同调度器的配置来切换管道调度器：</p> <!> <div class="flex gap-4"><div><img class="rounded-xl" src="https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/diffusers_docs/astronaut_lms.png"/> <figcaption class="mt-2 text-center text-sm text-gray-500">LMSDiscreteScheduler</figcaption></div> <div><img class="rounded-xl" src="https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/diffusers_docs/astronaut_euler_discrete.png"/> <figcaption class="mt-2 text-center text-sm text-gray-500">EulerDiscreteScheduler</figcaption></div></div> <div class="flex gap-4"><div><img class="rounded-xl" src="https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/diffusers_docs/astronaut_euler_ancestral.png"/> <figcaption class="mt-2 text-center text-sm text-gray-500">EulerAncestralDiscreteScheduler</figcaption></div> <div><img class="rounded-xl" src="https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/diffusers_docs/astronaut_dpm.png"/> <figcaption class="mt-2 text-center text-sm text-gray-500">DPMSolverMultistepScheduler</figcaption></div></div> <p>多数生成图像质量相近，实际选择需根据具体场景测试多种调度器进行比较。</p> <!> <p>对比Flax调度器时，需额外将调度器状态加载到模型参数中。例如将<code>FlaxStableDiffusionPipeline</code>的默认调度器切换为超高效的<code>FlaxDPMSolverMultistepScheduler</code>：</p> <blockquote><p>[!警告] <code>FlaxLMSDiscreteScheduler</code>和<code>FlaxDDPMScheduler</code>目前暂不兼容<code>FlaxStableDiffusionPipeline</code>。</p></blockquote> <!> <p>利用Flax对TPU的兼容性实现并行图像生成。需为每个设备复制模型参数，并分配输入数据：</p> <!> <!> <p>通过<code>ModelMixin.from_pretrained()</code>方法加载模型，该方法会下载并缓存模型权重和配置的最新版本。若本地缓存已存在最新文件，则直接复用缓存而非重复下载。</p> <p>通过<code>subfolder</code>参数可从子目录加载模型。例如<a href="https://hf.co/stable-diffusion-v1-5/stable-diffusion-v1-5" rel="nofollow">stable-diffusion-v1-5/stable-diffusion-v1-5</a>的模型权重存储在<a href="https://hf.co/stable-diffusion-v1-5/stable-diffusion-v1-5/tree/main/unet" rel="nofollow">unet</a>子目录中：</p> <!> <p>也可直接从<a href="https://huggingface.co/google/ddpm-cifar10-32/tree/main" rel="nofollow">仓库</a>加载：</p> <!> <p>加载和保存模型变体时，需在<code>ModelMixin.from_pretrained()</code>和<code>ModelMixin.save_pretrained()</code>中指定<code>variant</code>参数：</p> <!> <p>使用<code>from_pretrained()</code>的<code>torch_dtype</code>参数指定模型加载精度：</p> <!> <p>也可使用<a href="https://docs.pytorch.org/docs/stable/generated/torch.Tensor.to.html" rel="nofollow">torch.Tensor.to</a>方法即时转换精度，但会转换所有权重（不同于<code>torch_dtype</code>参数会保留<code>_keep_in_fp32_modules</code>中的层）。这对某些必须保持fp32精度的层尤为重要（参见<a href="https://github.com/huggingface/diffusers/blob/f864a9a352fa4a220d860bfdd1782e3e5af96382/src/diffusers/models/transformers/transformer_wan.py#L374" rel="nofollow">示例</a>）。</p> <!> <p></p>',1);function dl(_,I){H(I,!1),L(()=>{new URLSearchParams(window.location.search).get("fw")}),E();var M=ol();x("alcvz4",r=>{var d=$();q(d,"content",K),o(r,d)});var h=l(c(M),2);z(h,{containerStyle:"float: right; margin-left: 10px; display: inline-flex; position: relative; z-index: 10;"});var y=l(h,2);P(y,{containerStyle:"float: right; margin-left: 10px; display: inline-flex; position: relative; z-index: 10;",options:[{label:"Mixed",value:"https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers_doc/zh/schedulers.ipynb"},{label:"PyTorch",value:"https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers_doc/zh/pytorch/schedulers.ipynb"},{label:"TensorFlow",value:"https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers_doc/zh/tensorflow/schedulers.ipynb"},{label:"Mixed",value:"https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/diffusers_doc/zh/schedulers.ipynb"},{label:"PyTorch",value:"https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/diffusers_doc/zh/pytorch/schedulers.ipynb"},{label:"TensorFlow",value:"https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/diffusers_doc/zh/tensorflow/schedulers.ipynb"}]});var m=l(y,2);i(m,{title:"加载调度器与模型",local:"加载调度器与模型",headingTag:"h1"});var U=l(m,6);s(U,{code:"aW1wb3J0JTIwdG9yY2glMEFmcm9tJTIwZGlmZnVzZXJzJTIwaW1wb3J0JTIwRGlmZnVzaW9uUGlwZWxpbmUlMEElMEFwaXBlbGluZSUyMCUzRCUyMERpZmZ1c2lvblBpcGVsaW5lLmZyb21fcHJldHJhaW5lZCglMEElMjAlMjAlMjAlMjAlMjJzdGFibGUtZGlmZnVzaW9uLXYxLTUlMkZzdGFibGUtZGlmZnVzaW9uLXYxLTUlMjIlMkMlMjB0b3JjaF9kdHlwZSUzRHRvcmNoLmZsb2F0MTYlMkMlMjB1c2Vfc2FmZXRlbnNvcnMlM0RUcnVlJTBBKS50byglMjJjdWRhJTIyKQ==",highlighted:`<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">&quot;stable-diffusion-v1-5/stable-diffusion-v1-5&quot;</span>, torch_dtype=torch.float16, use_safetensors=<span class="hljs-literal">True</span>
).to(<span class="hljs-string">&quot;cuda&quot;</span>)`,lang:"python",wrap:!1});var f=l(U,4);s(f,{code:"cGlwZWxpbmUuc2NoZWR1bGVyJTBBUE5ETVNjaGVkdWxlciUyMCU3QiUwQSUyMCUyMCUyMl9jbGFzc19uYW1lJTIyJTNBJTIwJTIyUE5ETVNjaGVkdWxlciUyMiUyQyUwQSUyMCUyMCUyMl9kaWZmdXNlcnNfdmVyc2lvbiUyMiUzQSUyMCUyMjAuMjEuNCUyMiUyQyUwQSUyMCUyMCUyMmJldGFfZW5kJTIyJTNBJTIwMC4wMTIlMkMlMEElMjAlMjAlMjJiZXRhX3NjaGVkdWxlJTIyJTNBJTIwJTIyc2NhbGVkX2xpbmVhciUyMiUyQyUwQSUyMCUyMCUyMmJldGFfc3RhcnQlMjIlM0ElMjAwLjAwMDg1JTJDJTBBJTIwJTIwJTIyY2xpcF9zYW1wbGUlMjIlM0ElMjBmYWxzZSUyQyUwQSUyMCUyMCUyMm51bV90cmFpbl90aW1lc3RlcHMlMjIlM0ElMjAxMDAwJTJDJTBBJTIwJTIwJTIyc2V0X2FscGhhX3RvX29uZSUyMiUzQSUyMGZhbHNlJTJDJTBBJTIwJTIwJTIyc2tpcF9wcmtfc3RlcHMlMjIlM0ElMjB0cnVlJTJDJTBBJTIwJTIwJTIyc3RlcHNfb2Zmc2V0JTIyJTNBJTIwMSUyQyUwQSUyMCUyMCUyMnRpbWVzdGVwX3NwYWNpbmclMjIlM0ElMjAlMjJsZWFkaW5nJTIyJTJDJTBBJTIwJTIwJTIydHJhaW5lZF9iZXRhcyUyMiUzQSUyMG51bGwlMEElN0Q=",highlighted:`pipeline.scheduler
PNDMScheduler {
  <span class="hljs-string">&quot;_class_name&quot;</span>: <span class="hljs-string">&quot;PNDMScheduler&quot;</span>,
  <span class="hljs-string">&quot;_diffusers_version&quot;</span>: <span class="hljs-string">&quot;0.21.4&quot;</span>,
  <span class="hljs-string">&quot;beta_end&quot;</span>: <span class="hljs-number">0.012</span>,
  <span class="hljs-string">&quot;beta_schedule&quot;</span>: <span class="hljs-string">&quot;scaled_linear&quot;</span>,
  <span class="hljs-string">&quot;beta_start&quot;</span>: <span class="hljs-number">0.00085</span>,
  <span class="hljs-string">&quot;clip_sample&quot;</span>: false,
  <span class="hljs-string">&quot;num_train_timesteps&quot;</span>: <span class="hljs-number">1000</span>,
  <span class="hljs-string">&quot;set_alpha_to_one&quot;</span>: false,
  <span class="hljs-string">&quot;skip_prk_steps&quot;</span>: true,
  <span class="hljs-string">&quot;steps_offset&quot;</span>: <span class="hljs-number">1</span>,
  <span class="hljs-string">&quot;timestep_spacing&quot;</span>: <span class="hljs-string">&quot;leading&quot;</span>,
  <span class="hljs-string">&quot;trained_betas&quot;</span>: null
}`,lang:"python",wrap:!1});var b=l(f,2);i(b,{title:"加载调度器",local:"加载调度器",headingTag:"h2"});var J=l(b,6);s(J,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMERESU1TY2hlZHVsZXIlMkMlMjBEaWZmdXNpb25QaXBlbGluZSUwQSUwQWRkaW0lMjAlM0QlMjBERElNU2NoZWR1bGVyLmZyb21fcHJldHJhaW5lZCglMjJzdGFibGUtZGlmZnVzaW9uLXYxLTUlMkZzdGFibGUtZGlmZnVzaW9uLXYxLTUlMjIlMkMlMjBzdWJmb2xkZXIlM0QlMjJzY2hlZHVsZXIlMjIp",highlighted:`<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> DDIMScheduler, DiffusionPipeline

ddim = DDIMScheduler.from_pretrained(<span class="hljs-string">&quot;stable-diffusion-v1-5/stable-diffusion-v1-5&quot;</span>, subfolder=<span class="hljs-string">&quot;scheduler&quot;</span>)`,lang:"python",wrap:!1});var Z=l(J,4);s(Z,{code:"cGlwZWxpbmUlMjAlM0QlMjBEaWZmdXNpb25QaXBlbGluZS5mcm9tX3ByZXRyYWluZWQoJTBBJTIwJTIwJTIwJTIwJTIyc3RhYmxlLWRpZmZ1c2lvbi12MS01JTJGc3RhYmxlLWRpZmZ1c2lvbi12MS01JTIyJTJDJTIwc2NoZWR1bGVyJTNEZGRpbSUyQyUyMHRvcmNoX2R0eXBlJTNEdG9yY2guZmxvYXQxNiUyQyUyMHVzZV9zYWZldGVuc29ycyUzRFRydWUlMEEpLnRvKCUyMmN1ZGElMjIp",highlighted:`pipeline = DiffusionPipeline.from_pretrained(
    <span class="hljs-string">&quot;stable-diffusion-v1-5/stable-diffusion-v1-5&quot;</span>, scheduler=ddim, torch_dtype=torch.float16, use_safetensors=<span class="hljs-literal">True</span>
).to(<span class="hljs-string">&quot;cuda&quot;</span>)`,lang:"python",wrap:!1});var g=l(Z,2);i(g,{title:"调度器对比",local:"调度器对比",headingTag:"h2"});var T=l(g,6);s(T,{code:"aW1wb3J0JTIwdG9yY2glMEFmcm9tJTIwZGlmZnVzZXJzJTIwaW1wb3J0JTIwRGlmZnVzaW9uUGlwZWxpbmUlMEElMEFwaXBlbGluZSUyMCUzRCUyMERpZmZ1c2lvblBpcGVsaW5lLmZyb21fcHJldHJhaW5lZCglMEElMjAlMjAlMjAlMjAlMjJzdGFibGUtZGlmZnVzaW9uLXYxLTUlMkZzdGFibGUtZGlmZnVzaW9uLXYxLTUlMjIlMkMlMjB0b3JjaF9kdHlwZSUzRHRvcmNoLmZsb2F0MTYlMkMlMjB1c2Vfc2FmZXRlbnNvcnMlM0RUcnVlJTBBKS50byglMjJjdWRhJTIyKSUwQSUwQXByb21wdCUyMCUzRCUyMCUyMkElMjBwaG90b2dyYXBoJTIwb2YlMjBhbiUyMGFzdHJvbmF1dCUyMHJpZGluZyUyMGElMjBob3JzZSUyMG9uJTIwTWFycyUyQyUyMGhpZ2glMjByZXNvbHV0aW9uJTJDJTIwaGlnaCUyMGRlZmluaXRpb24uJTIyJTBBZ2VuZXJhdG9yJTIwJTNEJTIwdG9yY2guR2VuZXJhdG9yKGRldmljZSUzRCUyMmN1ZGElMjIpLm1hbnVhbF9zZWVkKDgp",highlighted:`<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">&quot;stable-diffusion-v1-5/stable-diffusion-v1-5&quot;</span>, torch_dtype=torch.float16, use_safetensors=<span class="hljs-literal">True</span>
).to(<span class="hljs-string">&quot;cuda&quot;</span>)

prompt = <span class="hljs-string">&quot;A photograph of an astronaut riding a horse on Mars, high resolution, high definition.&quot;</span>
generator = torch.Generator(device=<span class="hljs-string">&quot;cuda&quot;</span>).manual_seed(<span class="hljs-number">8</span>)`,lang:"python",wrap:!1});var j=l(T,4);D(j,{id:"schedulers",options:["LMSDiscreteScheduler","EulerDiscreteScheduler","EulerAncestralDiscreteScheduler","DPMSolverMultistepScheduler"],children:(r,d)=>{var C=nl(),F=c(C);p(F,{id:"schedulers",option:"LMSDiscreteScheduler",children:(a,u)=>{var e=ll(),n=l(c(e),2);s(n,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMExNU0Rpc2NyZXRlU2NoZWR1bGVyJTBBJTBBcGlwZWxpbmUuc2NoZWR1bGVyJTIwJTNEJTIwTE1TRGlzY3JldGVTY2hlZHVsZXIuZnJvbV9jb25maWcocGlwZWxpbmUuc2NoZWR1bGVyLmNvbmZpZyklMEFpbWFnZSUyMCUzRCUyMHBpcGVsaW5lKHByb21wdCUyQyUyMGdlbmVyYXRvciUzRGdlbmVyYXRvcikuaW1hZ2VzJTVCMCU1RCUwQWltYWdl",highlighted:`<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> LMSDiscreteScheduler

pipeline.scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config)
image = pipeline(prompt, generator=generator).images[<span class="hljs-number">0</span>]
image`,lang:"python",wrap:!1}),o(a,e)},$$slots:{default:!0}});var N=l(F,2);p(N,{id:"schedulers",option:"EulerDiscreteScheduler",children:(a,u)=>{var e=sl(),n=l(c(e),2);s(n,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMEV1bGVyRGlzY3JldGVTY2hlZHVsZXIlMEElMEFwaXBlbGluZS5zY2hlZHVsZXIlMjAlM0QlMjBFdWxlckRpc2NyZXRlU2NoZWR1bGVyLmZyb21fY29uZmlnKHBpcGVsaW5lLnNjaGVkdWxlci5jb25maWcpJTBBaW1hZ2UlMjAlM0QlMjBwaXBlbGluZShwcm9tcHQlMkMlMjBnZW5lcmF0b3IlM0RnZW5lcmF0b3IpLmltYWdlcyU1QjAlNUQlMEFpbWFnZQ==",highlighted:`<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> EulerDiscreteScheduler

pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
image = pipeline(prompt, generator=generator).images[<span class="hljs-number">0</span>]
image`,lang:"python",wrap:!1}),o(a,e)},$$slots:{default:!0}});var k=l(N,2);p(k,{id:"schedulers",option:"EulerAncestralDiscreteScheduler",children:(a,u)=>{var e=el(),n=l(c(e),2);s(n,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMEV1bGVyQW5jZXN0cmFsRGlzY3JldGVTY2hlZHVsZXIlMEElMEFwaXBlbGluZS5zY2hlZHVsZXIlMjAlM0QlMjBFdWxlckFuY2VzdHJhbERpc2NyZXRlU2NoZWR1bGVyLmZyb21fY29uZmlnKHBpcGVsaW5lLnNjaGVkdWxlci5jb25maWcpJTBBaW1hZ2UlMjAlM0QlMjBwaXBlbGluZShwcm9tcHQlMkMlMjBnZW5lcmF0b3IlM0RnZW5lcmF0b3IpLmltYWdlcyU1QjAlNUQlMEFpbWFnZQ==",highlighted:`<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> EulerAncestralDiscreteScheduler

pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
image = pipeline(prompt, generator=generator).images[<span class="hljs-number">0</span>]
image`,lang:"python",wrap:!1}),o(a,e)},$$slots:{default:!0}});var Q=l(k,2);p(Q,{id:"schedulers",option:"DPMSolverMultistepScheduler",children:(a,u)=>{var e=al(),n=l(c(e),2);s(n,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMERQTVNvbHZlck11bHRpc3RlcFNjaGVkdWxlciUwQSUwQXBpcGVsaW5lLnNjaGVkdWxlciUyMCUzRCUyMERQTVNvbHZlck11bHRpc3RlcFNjaGVkdWxlci5mcm9tX2NvbmZpZyhwaXBlbGluZS5zY2hlZHVsZXIuY29uZmlnKSUwQWltYWdlJTIwJTNEJTIwcGlwZWxpbmUocHJvbXB0JTJDJTIwZ2VuZXJhdG9yJTNEZ2VuZXJhdG9yKS5pbWFnZXMlNUIwJTVEJTBBaW1hZ2U=",highlighted:`<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> DPMSolverMultistepScheduler

pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
image = pipeline(prompt, generator=generator).images[<span class="hljs-number">0</span>]
image`,lang:"python",wrap:!1}),o(a,e)},$$slots:{default:!0}}),o(r,C)},$$slots:{default:!0}});var w=l(j,8);i(w,{title:"Flax调度器",local:"flax调度器",headingTag:"h3"});var V=l(w,6);s(V,{code:"aW1wb3J0JTIwamF4JTBBaW1wb3J0JTIwbnVtcHklMjBhcyUyMG5wJTBBZnJvbSUyMGZsYXguamF4X3V0aWxzJTIwaW1wb3J0JTIwcmVwbGljYXRlJTBBZnJvbSUyMGZsYXgudHJhaW5pbmcuY29tbW9uX3V0aWxzJTIwaW1wb3J0JTIwc2hhcmQlMEFmcm9tJTIwZGlmZnVzZXJzJTIwaW1wb3J0JTIwRmxheFN0YWJsZURpZmZ1c2lvblBpcGVsaW5lJTJDJTIwRmxheERQTVNvbHZlck11bHRpc3RlcFNjaGVkdWxlciUwQSUwQXNjaGVkdWxlciUyQyUyMHNjaGVkdWxlcl9zdGF0ZSUyMCUzRCUyMEZsYXhEUE1Tb2x2ZXJNdWx0aXN0ZXBTY2hlZHVsZXIuZnJvbV9wcmV0cmFpbmVkKCUwQSUyMCUyMCUyMCUyMCUyMnN0YWJsZS1kaWZmdXNpb24tdjEtNSUyRnN0YWJsZS1kaWZmdXNpb24tdjEtNSUyMiUyQyUwQSUyMCUyMCUyMCUyMHN1YmZvbGRlciUzRCUyMnNjaGVkdWxlciUyMiUwQSklMEFwaXBlbGluZSUyQyUyMHBhcmFtcyUyMCUzRCUyMEZsYXhTdGFibGVEaWZmdXNpb25QaXBlbGluZS5mcm9tX3ByZXRyYWluZWQoJTBBJTIwJTIwJTIwJTIwJTIyc3RhYmxlLWRpZmZ1c2lvbi12MS01JTJGc3RhYmxlLWRpZmZ1c2lvbi12MS01JTIyJTJDJTBBJTIwJTIwJTIwJTIwc2NoZWR1bGVyJTNEc2NoZWR1bGVyJTJDJTBBJTIwJTIwJTIwJTIwdmFyaWFudCUzRCUyMmJmMTYlMjIlMkMlMEElMjAlMjAlMjAlMjBkdHlwZSUzRGpheC5udW1weS5iZmxvYXQxNiUyQyUwQSklMEFwYXJhbXMlNUIlMjJzY2hlZHVsZXIlMjIlNUQlMjAlM0QlMjBzY2hlZHVsZXJfc3RhdGU=",highlighted:`<span class="hljs-keyword">import</span> jax
<span class="hljs-keyword">import</span> numpy <span class="hljs-keyword">as</span> np
<span class="hljs-keyword">from</span> flax.jax_utils <span class="hljs-keyword">import</span> replicate
<span class="hljs-keyword">from</span> flax.training.common_utils <span class="hljs-keyword">import</span> shard
<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> FlaxStableDiffusionPipeline, FlaxDPMSolverMultistepScheduler

scheduler, scheduler_state = FlaxDPMSolverMultistepScheduler.from_pretrained(
    <span class="hljs-string">&quot;stable-diffusion-v1-5/stable-diffusion-v1-5&quot;</span>,
    subfolder=<span class="hljs-string">&quot;scheduler&quot;</span>
)
pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
    <span class="hljs-string">&quot;stable-diffusion-v1-5/stable-diffusion-v1-5&quot;</span>,
    scheduler=scheduler,
    variant=<span class="hljs-string">&quot;bf16&quot;</span>,
    dtype=jax.numpy.bfloat16,
)
params[<span class="hljs-string">&quot;scheduler&quot;</span>] = scheduler_state`,lang:"python",wrap:!1});var G=l(V,4);s(G,{code:"JTIzJTIwJUU2JUFGJThGJUU0JUI4JUFBJUU1JUI5JUI2JUU4JUExJThDJUU4JUFFJUJFJUU1JUE0JTg3JUU3JTk0JTlGJUU2JTg4JTkwMSVFNSVCQyVBMCVFNSU5QiVCRSVFNSU4MyU4RiVFRiVCQyU4OFRQVXYyLTglMkZUUFV2My04JUU2JTk0JUFGJUU2JThDJTgxOCVFOCVBRSVCRSVFNSVBNCU4NyVFNSVCOSVCNiVFOCVBMSU4QyVFRiVCQyU4OSUwQXByb21wdCUyMCUzRCUyMCUyMiVFNCVCOCU4MCVFNSVCQyVBMCVFNSVBRSU4NyVFOCU4OCVBQSVFNSU5MSU5OCVFNSU5QyVBOCVFNyU4MSVBQiVFNiU5OCU5RiVFNCVCOCU4QSVFOSVBQSU5MSVFOSVBOSVBQyVFNyU5QSU4NCVFOSVBQiU5OCVFNiVCOCU4NSVFNyU4NSVBNyVFNyU4OSU4NyVFRiVCQyU4QyVFOSVBQiU5OCVFNSU4OCU4NiVFOCVCRSVBOCVFNyU4RSU4NyVFRiVCQyU4QyVFOSVBQiU5OCVFNyU5NCVCQiVFOCVCNCVBOCVFMyU4MCU4MiUyMiUwQW51bV9zYW1wbGVzJTIwJTNEJTIwamF4LmRldmljZV9jb3VudCgpJTBBcHJvbXB0X2lkcyUyMCUzRCUyMHBpcGVsaW5lLnByZXBhcmVfaW5wdXRzKCU1QnByb21wdCU1RCUyMColMjBudW1fc2FtcGxlcyklMEElMEFwcm5nX3NlZWQlMjAlM0QlMjBqYXgucmFuZG9tLlBSTkdLZXkoMCklMEFudW1faW5mZXJlbmNlX3N0ZXBzJTIwJTNEJTIwMjUlMEElMEElMjMlMjAlRTUlODglODYlRTklODUlOEQlRTglQkUlOTMlRTUlODUlQTUlRTUlOTIlOEMlRTklOUElOEYlRTYlOUMlQkElRTclQTclOEQlRTUlQUQlOTAlMEFwYXJhbXMlMjAlM0QlMjByZXBsaWNhdGUocGFyYW1zKSUwQXBybmdfc2VlZCUyMCUzRCUyMGpheC5yYW5kb20uc3BsaXQocHJuZ19zZWVkJTJDJTIwamF4LmRldmljZV9jb3VudCgpKSUwQXByb21wdF9pZHMlMjAlM0QlMjBzaGFyZChwcm9tcHRfaWRzKSUwQSUwQWltYWdlcyUyMCUzRCUyMHBpcGVsaW5lKHByb21wdF9pZHMlMkMlMjBwYXJhbXMlMkMlMjBwcm5nX3NlZWQlMkMlMjBudW1faW5mZXJlbmNlX3N0ZXBzJTJDJTIwaml0JTNEVHJ1ZSkuaW1hZ2VzJTBBaW1hZ2VzJTIwJTNEJTIwcGlwZWxpbmUubnVtcHlfdG9fcGlsKG5wLmFzYXJyYXkoaW1hZ2VzLnJlc2hhcGUoKG51bV9zYW1wbGVzJTJDKSUyMCUyQiUyMGltYWdlcy5zaGFwZSU1Qi0zJTNBJTVEKSkp",highlighted:`<span class="hljs-comment"># 每个并行设备生成1张图像（TPUv2-8/TPUv3-8支持8设备并行）</span>
prompt = <span class="hljs-string">&quot;一张宇航员在火星上骑马的高清照片，高分辨率，高画质。&quot;</span>
num_samples = jax.device_count()
prompt_ids = pipeline.prepare_inputs([prompt] * num_samples)

prng_seed = jax.random.PRNGKey(<span class="hljs-number">0</span>)
num_inference_steps = <span class="hljs-number">25</span>

<span class="hljs-comment"># 分配输入和随机种子</span>
params = replicate(params)
prng_seed = jax.random.split(prng_seed, jax.device_count())
prompt_ids = shard(prompt_ids)

images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=<span class="hljs-literal">True</span>).images
images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-<span class="hljs-number">3</span>:])))`,lang:"python",wrap:!1});var v=l(G,2);i(v,{title:"模型加载",local:"模型加载",headingTag:"h2"});var S=l(v,6);s(S,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMFVOZXQyRENvbmRpdGlvbk1vZGVsJTBBJTBBdW5ldCUyMCUzRCUyMFVOZXQyRENvbmRpdGlvbk1vZGVsLmZyb21fcHJldHJhaW5lZCglMjJzdGFibGUtZGlmZnVzaW9uLXYxLTUlMkZzdGFibGUtZGlmZnVzaW9uLXYxLTUlMjIlMkMlMjBzdWJmb2xkZXIlM0QlMjJ1bmV0JTIyJTJDJTIwdXNlX3NhZmV0ZW5zb3JzJTNEVHJ1ZSk=",highlighted:`<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> UNet2DConditionModel

unet = UNet2DConditionModel.from_pretrained(<span class="hljs-string">&quot;stable-diffusion-v1-5/stable-diffusion-v1-5&quot;</span>, subfolder=<span class="hljs-string">&quot;unet&quot;</span>, use_safetensors=<span class="hljs-literal">True</span>)`,lang:"python",wrap:!1});var W=l(S,4);s(W,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMFVOZXQyRE1vZGVsJTBBJTBBdW5ldCUyMCUzRCUyMFVOZXQyRE1vZGVsLmZyb21fcHJldHJhaW5lZCglMjJnb29nbGUlMkZkZHBtLWNpZmFyMTAtMzIlMjIlMkMlMjB1c2Vfc2FmZXRlbnNvcnMlM0RUcnVlKQ==",highlighted:`<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> UNet2DModel

unet = UNet2DModel.from_pretrained(<span class="hljs-string">&quot;google/ddpm-cifar10-32&quot;</span>, use_safetensors=<span class="hljs-literal">True</span>)`,lang:"python",wrap:!1});var R=l(W,4);s(R,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMFVOZXQyRENvbmRpdGlvbk1vZGVsJTBBJTBBdW5ldCUyMCUzRCUyMFVOZXQyRENvbmRpdGlvbk1vZGVsLmZyb21fcHJldHJhaW5lZCglMEElMjAlMjAlMjAlMjAlMjJzdGFibGUtZGlmZnVzaW9uLXYxLTUlMkZzdGFibGUtZGlmZnVzaW9uLXYxLTUlMjIlMkMlMjBzdWJmb2xkZXIlM0QlMjJ1bmV0JTIyJTJDJTIwdmFyaWFudCUzRCUyMm5vbl9lbWElMjIlMkMlMjB1c2Vfc2FmZXRlbnNvcnMlM0RUcnVlJTBBKSUwQXVuZXQuc2F2ZV9wcmV0cmFpbmVkKCUyMi4lMkZsb2NhbC11bmV0JTIyJTJDJTIwdmFyaWFudCUzRCUyMm5vbl9lbWElMjIp",highlighted:`<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> UNet2DConditionModel

unet = UNet2DConditionModel.from_pretrained(
    <span class="hljs-string">&quot;stable-diffusion-v1-5/stable-diffusion-v1-5&quot;</span>, subfolder=<span class="hljs-string">&quot;unet&quot;</span>, variant=<span class="hljs-string">&quot;non_ema&quot;</span>, use_safetensors=<span class="hljs-literal">True</span>
)
unet.save_pretrained(<span class="hljs-string">&quot;./local-unet&quot;</span>, variant=<span class="hljs-string">&quot;non_ema&quot;</span>)`,lang:"python",wrap:!1});var B=l(R,4);s(B,{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMEF1dG9Nb2RlbCUwQSUwQXVuZXQlMjAlM0QlMjBBdXRvTW9kZWwuZnJvbV9wcmV0cmFpbmVkKCUwQSUyMCUyMCUyMCUyMCUyMnN0YWJpbGl0eWFpJTJGc3RhYmxlLWRpZmZ1c2lvbi14bC1iYXNlLTEuMCUyMiUyQyUyMHN1YmZvbGRlciUzRCUyMnVuZXQlMjIlMkMlMjB0b3JjaF9kdHlwZSUzRHRvcmNoLmZsb2F0MTYlMEEp",highlighted:`<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> AutoModel

unet = AutoModel.from_pretrained(
    <span class="hljs-string">&quot;stabilityai/stable-diffusion-xl-base-1.0&quot;</span>, subfolder=<span class="hljs-string">&quot;unet&quot;</span>, torch_dtype=torch.float16
)`,lang:"python",wrap:!1});var X=l(B,4);Y(X,{source:"https://github.com/huggingface/diffusers/blob/main/docs/source/zh/using-diffusers/schedulers.md"}),A(2),o(_,M),O()}export{dl as component};