diffusers/v0.22.0/en/_app/immutable/nodes/178.ecea4dba.js

import{s as Ot,o as el,n as Kt}from"../chunks/scheduler.182ea377.js";import{S as tl,i as ll,g as i,s as a,r as m,A as sl,h as o,f as l,c as n,j as qt,u as c,x as p,k as Pt,y as al,a as s,v as u,d,t as h,w as M}from"../chunks/index.abf12888.js";import{T as Dt}from"../chunks/Tip.230e2334.js";import{C as f}from"../chunks/CodeBlock.57fe6e13.js";import{D as nl}from"../chunks/DocNotebookDropdown.d9060979.js";import{H as Te}from"../chunks/Heading.16916d63.js";function il(fe){let r,w='Learn more details in the <a href="#how-does-parallelization-work">How does parallelization work?</a> section.';return{c(){r=i("p"),r.innerHTML=w},l(y){r=o(y,"P",{"data-svelte-h":!0}),p(r)!=="svelte-5pbnhk"&&(r.innerHTML=w)},m(y,T){s(y,r,T)},p:Kt,d(y){y&&l(r)}}}function ol(fe){let r,w="You need to ensure all your inputs have the same shape in subsequent calls, other JAX will need to recompile the code which is slower.";return{c(){r=i("p"),r.textContent=w},l(y){r=o(y,"P",{"data-svelte-h":!0}),p(r)!=="svelte-10f6ur"&&(r.textContent=w)},m(y,T){s(y,r,T)},p:Kt,d(y){y&&l(r)}}}function pl(fe){let r,w,y,T,j,be,U,Je,g,wt="🤗 Diffusers supports Flax for super fast inference on Google TPUs, such as those available in Colab, Kaggle or Google Cloud Platform. This guide shows you how to run inference with Stable Diffusion using JAX/Flax.",je,$,Tt="Before you begin, make sure you have the necessary libraries installed:",Ue,v,ge,B,bt="You should also make sure you’re using a TPU backend. While JAX does not run exclusively on TPUs, you’ll get the best performance on a TPU because each server has 8 TPU accelerators working in parallel.",$e,_,Jt="If you are running this guide in Colab, select <em>Runtime</em> in the menu above, select the option <em>Change runtime type</em>, and then select <em>TPU</em> under the <em>Hardware accelerator</em> setting. Import JAX and quickly check whether you’re using a TPU:",ve,C,Be,k,jt="Great, now you can import the rest of the dependencies you’ll need:",_e,I,Ce,x,ke,G,Ut="Flax is a functional framework, so models are stateless and parameters are stored outside of them. Loading a pretrained Flax pipeline returns <em>both</em> the pipeline and the model weights (or parameters). In this guide, you’ll use <code>bfloat16</code>, a more efficient half-float type that is supported by TPUs (you can also use <code>float32</code> for full precision if you want).",Ie,W,xe,X,Ge,Z,gt="TPUs usually have 8 devices working in parallel, so let’s use the same prompt for each device. This means you can perform inference on 8 devices at once, with each device generating one image. As a result, you’ll get 8 images in the same amount of time it takes for one chip to generate a single image!",We,b,Xe,H,$t="After replicating the prompt, get the tokenized text ids by calling the <code>prepare_inputs</code> function on the pipeline. The length of the tokenized text is set to 77 tokens as required by the configuration of the underlying CLIP text model.",Ze,F,He,Y,vt='Model parameters and inputs have to be replicated across the 8 parallel devices. The parameters dictionary is replicated with <a href="https://flax.readthedocs.io/en/latest/api_reference/flax.jax_utils.html#flax.jax_utils.replicate" rel="nofollow"><code>flax.jax_utils.replicate</code></a> which traverses the dictionary and changes the shape of the weights so they are repeated 8 times. Arrays are replicated using <code>shard</code>.',Fe,E,Ye,Q,Bt="This shape means each one of the 8 devices receives as an input a <code>jnp</code> array with shape <code>(1, 77)</code>, where <code>1</code> is the batch size per device. On TPUs with sufficient memory, you could have a batch size larger than <code>1</code> if you want to generate multiple images (per chip) at once.",Ee,R,_t="Next, create a random number generator to pass to the generation function. This is standard procedure in Flax, which is very serious and opinionated about random numbers. All functions that deal with random numbers are expected to receive a generator to ensure reproducibility, even when you’re training across multiple distributed devices.",Qe,z,Ct="The helper function below uses a seed to initialize a random number generator. As long as you use the same seed, you’ll get the exact same results. Feel free to use different seeds when exploring results later in the guide.",Re,L,ze,N,kt="The helper function, or <code>rng</code>, is split 8 times so each device receives a different generator and generates a different image.",Le,V,Ne,S,It="To take advantage of JAX’s optimized speed on a TPU, pass <code>jit=True</code> to the pipeline to compile the JAX code into an efficient representation and to ensure the model runs in parallel across the 8 devices.",Ve,J,Se,A,xt="The first inference run takes more time because it needs to compile the code, but subsequent calls (even with different inputs) are much faster. For example, it took more than a minute to compile on a TPU v2-8, but then it takes about <strong>7s</strong> on a future inference run!",Ae,q,qe,P,Gt='The returned array has shape <code>(8, 1, 512, 512, 3)</code> which should be reshaped to remove the second dimension and get 8 images of <code>512 × 512 × 3</code>. Then you can use the <a href="/docs/diffusers/v0.22.0/en/api/utilities#diffusers.utils.numpy_to_pil">numpy_to_pil()</a> function to convert the arrays into images.',Pe,D,De,K,Wt='<img src="https://huggingface.co/datasets/YiYiXu/test-doc-assets/resolve/main/stable_diffusion_jax_how_to_cell_38_output_0.jpeg" alt="img"/>',Ke,O,Oe,ee,Xt="You don’t necessarily have to use the same prompt on all devices. For example, to generate 8 different prompts:",et,te,tt,le,Zt='<img src="https://huggingface.co/datasets/YiYiXu/test-doc-assets/resolve/main/stable_diffusion_jax_how_to_cell_43_output_0.jpeg" alt="img"/>',lt,se,st,ae,Ht="The Flax pipeline in 🤗 Diffusers automatically compiles the model and runs it in parallel on all available devices. Let’s take a closer look at how that process works.",at,ne,Ft='JAX parallelization can be done in multiple ways. The easiest one revolves around using the <a href="https://jax.readthedocs.io/en/latest/_autosummary/jax.pmap.html" rel="nofollow"><code>jax.pmap</code></a> function to achieve single-program multiple-data (SPMD) parallelization. It means running several copies of the same code, each on different data inputs. More sophisticated approaches are possible, and you can go over to the JAX <a href="https://jax.readthedocs.io/en/latest/index.html" rel="nofollow">documentation</a> to explore this topic in more detail if you are interested!',nt,ie,Yt="<code>jax.pmap</code> does two things:",it,oe,Et="<li>Compiles (or ”<code>jit</code>s”) the code which is similar to <code>jax.jit()</code>. This does not happen when you call <code>pmap</code>, and only the first time the <code>pmap</code>ped function is called.</li> <li>Ensures the compiled code runs in parallel on all available devices.</li>",ot,pe,Qt="To demonstrate, call <code>pmap</code> on the pipeline’s <code>_generate</code> method (this is a private method that generates images and may be renamed or removed in future releases of 🤗 Diffusers):",pt,re,rt,me,Rt="After calling <code>pmap</code>, the prepared function <code>p_generate</code> will:",mt,ce,zt="<li>Make a copy of the underlying function, <code>pipeline._generate</code>, on each device.</li> <li>Send each device a different portion of the input arguments (this is why its necessary to call the <em>shard</em> function). In this case, <code>prompt_ids</code> has shape <code>(8, 1, 77, 768)</code> so the array is split into 8 and each copy of <code>_generate</code> receives an input with shape <code>(1, 77, 768)</code>.</li>",ct,ue,Lt="The most important thing to pay attention to here is the batch size (1 in this example), and the input dimensions that make sense for your code. You don’t have to change anything else to make the code work in parallel.",ut,de,Nt="The first time you call the pipeline takes more time, but the calls afterward are much faster. The <code>block_until_ready</code> function is used to correctly measure inference time because JAX uses asynchronous dispatch and returns control to the Python loop as soon as it can. You don’t need to use that in your code; blocking occurs automatically when you want to use the result of a computation that has not yet been materialized.",dt,he,ht,Me,Vt="Check your image dimensions to see if they’re correct:",Mt,ye,yt,we,ft;return j=new Te({props:{title:"JAX/Flax",local:"jaxflax",headingTag:"h1"}}),U=new nl({props:{classNames:"absolute z-10 right-0 top-0",options:[{label:"Mixed",value:"https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers_doc/en/stable_diffusion_jax_how_to.ipynb"},{label:"PyTorch",value:"https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers_doc/en/pytorch/stable_diffusion_jax_how_to.ipynb"},{label:"TensorFlow",value:"https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers_doc/en/tensorflow/stable_diffusion_jax_how_to.ipynb"},{label:"Mixed",value:"https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/diffusers_doc/en/stable_diffusion_jax_how_to.ipynb"},{label:"PyTorch",value:"https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/diffusers_doc/en/pytorch/stable_diffusion_jax_how_to.ipynb"},{label:"TensorFlow",value:"https://studiolab.sagemaker.aws/import/github/huggingface/notebooks/blob/main/diffusers_doc/en/tensorflow/stable_diffusion_jax_how_to.ipynb"}]}}),v=new f({props:{code:"JTIzJTIwdW5jb21tZW50JTIwdG8lMjBpbnN0YWxsJTIwdGhlJTIwbmVjZXNzYXJ5JTIwbGlicmFyaWVzJTIwaW4lMjBDb2xhYiUwQSUyMyFwaXAlMjBpbnN0YWxsJTIwLXElMjBqYXglM0QlM0QwLjMuMjUlMjBqYXhsaWIlM0QlM0QwLjMuMjUlMjBmbGF4JTIwdHJhbnNmb3JtZXJzJTIwZnRmeSUwQSUyMyFwaXAlMjBpbnN0YWxsJTIwLXElMjBkaWZmdXNlcnM=",highlighted:`<span class="hljs-comment"># uncomment to install the necessary libraries in Colab</span>
<span class="hljs-comment">#!pip install -q jax==0.3.25 jaxlib==0.3.25 flax transformers ftfy</span>
<span class="hljs-comment">#!pip install -q diffusers</span>`,wrap:!1}}),C=new f({props:{code:"aW1wb3J0JTIwamF4JTBBaW1wb3J0JTIwamF4LnRvb2xzLmNvbGFiX3RwdSUwQWpheC50b29scy5jb2xhYl90cHUuc2V0dXBfdHB1KCklMEElMEFudW1fZGV2aWNlcyUyMCUzRCUyMGpheC5kZXZpY2VfY291bnQoKSUwQWRldmljZV90eXBlJTIwJTNEJTIwamF4LmRldmljZXMoKSU1QjAlNUQuZGV2aWNlX2tpbmQlMEElMEFwcmludChmJTIyRm91bmQlMjAlN0JudW1fZGV2aWNlcyU3RCUyMEpBWCUyMGRldmljZXMlMjBvZiUyMHR5cGUlMjAlN0JkZXZpY2VfdHlwZSU3RC4lMjIpJTBBYXNzZXJ0JTIwKCUwQSUyMCUyMCUyMCUyMCUyMlRQVSUyMiUyMGluJTIwZGV2aWNlX3R5cGUlMkMlMjAlMEElMjAlMjAlMjAlMjAlMjJBdmFpbGFibGUlMjBkZXZpY2UlMjBpcyUyMG5vdCUyMGElMjBUUFUlMkMlMjBwbGVhc2UlMjBzZWxlY3QlMjBUUFUlMjBmcm9tJTIwRWRpdCUyMCUzRSUyME5vdGVib29rJTIwc2V0dGluZ3MlMjAlM0UlMjBIYXJkd2FyZSUyMGFjY2VsZXJhdG9yJTIyJTBBKSUwQSUyMkZvdW5kJTIwOCUyMEpBWCUyMGRldmljZXMlMjBvZiUyMHR5cGUlMjBDbG91ZCUyMFRQVS4lMjI=",highlighted:`<span class="hljs-keyword">import</span> jax
<span class="hljs-keyword">import</span> jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()

num_devices = jax.device_count()
device_type = jax.devices()[<span class="hljs-number">0</span>].device_kind

<span class="hljs-built_in">print</span>(<span class="hljs-string">f&quot;Found <span class="hljs-subst">{num_devices}</span> JAX devices of type <span class="hljs-subst">{device_type}</span>.&quot;</span>)
<span class="hljs-keyword">assert</span> (
    <span class="hljs-string">&quot;TPU&quot;</span> <span class="hljs-keyword">in</span> device_type, 
    <span class="hljs-string">&quot;Available device is not a TPU, please select TPU from Edit &gt; Notebook settings &gt; Hardware accelerator&quot;</span>
)
<span class="hljs-string">&quot;Found 8 JAX devices of type Cloud TPU.&quot;</span>`,wrap:!1}}),I=new f({props:{code:"aW1wb3J0JTIwbnVtcHklMjBhcyUyMG5wJTBBaW1wb3J0JTIwamF4Lm51bXB5JTIwYXMlMjBqbnAlMEElMEFmcm9tJTIwcGF0aGxpYiUyMGltcG9ydCUyMFBhdGglMEFmcm9tJTIwamF4JTIwaW1wb3J0JTIwcG1hcCUwQWZyb20lMjBmbGF4LmpheF91dGlscyUyMGltcG9ydCUyMHJlcGxpY2F0ZSUwQWZyb20lMjBmbGF4LnRyYWluaW5nLmNvbW1vbl91dGlscyUyMGltcG9ydCUyMHNoYXJkJTBBZnJvbSUyMFBJTCUyMGltcG9ydCUyMEltYWdlJTBBJTBBZnJvbSUyMGh1Z2dpbmdmYWNlX2h1YiUyMGltcG9ydCUyMG5vdGVib29rX2xvZ2luJTBBZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMEZsYXhTdGFibGVEaWZmdXNpb25QaXBlbGluZQ==",highlighted:`<span class="hljs-keyword">import</span> numpy <span class="hljs-keyword">as</span> np
<span class="hljs-keyword">import</span> jax.numpy <span class="hljs-keyword">as</span> jnp

<span class="hljs-keyword">from</span> pathlib <span class="hljs-keyword">import</span> Path
<span class="hljs-keyword">from</span> jax <span class="hljs-keyword">import</span> pmap
<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> PIL <span class="hljs-keyword">import</span> Image

<span class="hljs-keyword">from</span> huggingface_hub <span class="hljs-keyword">import</span> notebook_login
<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> FlaxStableDiffusionPipeline`,wrap:!1}}),x=new Te({props:{title:"Load a model",local:"load-a-model",headingTag:"h2"}}),W=new f({props:{code:"ZHR5cGUlMjAlM0QlMjBqbnAuYmZsb2F0MTYlMEFwaXBlbGluZSUyQyUyMHBhcmFtcyUyMCUzRCUyMEZsYXhTdGFibGVEaWZmdXNpb25QaXBlbGluZS5mcm9tX3ByZXRyYWluZWQoJTBBJTIwJTIwJTIwJTIwJTIyQ29tcFZpcyUyRnN0YWJsZS1kaWZmdXNpb24tdjEtNCUyMiUyQyUwQSUyMCUyMCUyMCUyMHJldmlzaW9uJTNEJTIyYmYxNiUyMiUyQyUwQSUyMCUyMCUyMCUyMGR0eXBlJTNEZHR5cGUlMkMlMEEp",highlighted:`dtype = jnp.bfloat16
pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
    <span class="hljs-string">&quot;CompVis/stable-diffusion-v1-4&quot;</span>,
    revision=<span class="hljs-string">&quot;bf16&quot;</span>,
    dtype=dtype,
)`,wrap:!1}}),X=new Te({props:{title:"Inference",local:"inference",headingTag:"h2"}}),b=new Dt({props:{$$slots:{default:[il]},$$scope:{ctx:fe}}}),F=new f({props:{code:"cHJvbXB0JTIwJTNEJTIwJTIyQSUyMGNpbmVtYXRpYyUyMGZpbG0lMjBzdGlsbCUyMG9mJTIwTW9yZ2FuJTIwRnJlZW1hbiUyMHN0YXJyaW5nJTIwYXMlMjBKaW1pJTIwSGVuZHJpeCUyQyUyMHBvcnRyYWl0JTJDJTIwNDBtbSUyMGxlbnMlMkMlMjBzaGFsbG93JTIwZGVwdGglMjBvZiUyMGZpZWxkJTJDJTIwY2xvc2UlMjB1cCUyQyUyMHNwbGl0JTIwbGlnaHRpbmclMkMlMjBjaW5lbWF0aWMlMjIlMEFwcm9tcHQlMjAlM0QlMjAlNUJwcm9tcHQlNUQlMjAqJTIwamF4LmRldmljZV9jb3VudCgpJTBBcHJvbXB0X2lkcyUyMCUzRCUyMHBpcGVsaW5lLnByZXBhcmVfaW5wdXRzKHByb21wdCklMEFwcm9tcHRfaWRzLnNoYXBlJTBBJTIyKDglMkMlMjA3NyklMjI=",highlighted:`prompt = <span class="hljs-string">&quot;A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of field, close up, split lighting, cinematic&quot;</span>
prompt = [prompt] * jax.device_count()
prompt_ids = pipeline.prepare_inputs(prompt)
prompt_ids.shape
<span class="hljs-string">&quot;(8, 77)&quot;</span>`,wrap:!1}}),E=new f({props:{code:"JTIzJTIwcGFyYW1ldGVycyUwQXBfcGFyYW1zJTIwJTNEJTIwcmVwbGljYXRlKHBhcmFtcyklMEElMEElMjMlMjBhcnJheXMlMEFwcm9tcHRfaWRzJTIwJTNEJTIwc2hhcmQocHJvbXB0X2lkcyklMEFwcm9tcHRfaWRzLnNoYXBlJTBBJTIyKDglMkMlMjAxJTJDJTIwNzcpJTIy",highlighted:`<span class="hljs-comment"># parameters</span>
p_params = replicate(params)

<span class="hljs-comment"># arrays</span>
prompt_ids = shard(prompt_ids)
prompt_ids.shape
<span class="hljs-string">&quot;(8, 1, 77)&quot;</span>`,wrap:!1}}),L=new f({props:{code:"ZGVmJTIwY3JlYXRlX2tleShzZWVkJTNEMCklM0ElMEElMjAlMjAlMjAlMjByZXR1cm4lMjBqYXgucmFuZG9tLlBSTkdLZXkoc2VlZCk=",highlighted:`<span class="hljs-keyword">def</span> <span class="hljs-title function_">create_key</span>(<span class="hljs-params">seed=<span class="hljs-number">0</span></span>):
    <span class="hljs-keyword">return</span> jax.random.PRNGKey(seed)`,wrap:!1}}),V=new f({props:{code:"cm5nJTIwJTNEJTIwY3JlYXRlX2tleSgwKSUwQXJuZyUyMCUzRCUyMGpheC5yYW5kb20uc3BsaXQocm5nJTJDJTIwamF4LmRldmljZV9jb3VudCgpKQ==",highlighted:`rng = create_key(<span class="hljs-number">0</span>)
rng = jax.random.split(rng, jax.device_count())`,wrap:!1}}),J=new Dt({props:{warning:!0,$$slots:{default:[ol]},$$scope:{ctx:fe}}}),q=new f({props:{code:"JTI1JTI1dGltZSUwQWltYWdlcyUyMCUzRCUyMHBpcGVsaW5lKHByb21wdF9pZHMlMkMlMjBwX3BhcmFtcyUyQyUyMHJuZyUyQyUyMGppdCUzRFRydWUpJTVCMCU1RCUwQSUwQSUyMkNQVSUyMHRpbWVzJTNBJTIwdXNlciUyMDU2LjIlMjBzJTJDJTIwc3lzJTNBJTIwNDIuNSUyMHMlMkMlMjB0b3RhbCUzQSUyMDFtaW4lMjAzOHMlMjIlMEElMjJXYWxsJTIwdGltZSUzQSUyMDFtaW4lMjAyOXMlMjI=",highlighted:`%%time
images = pipeline(prompt_ids, p_params, rng, jit=<span class="hljs-literal">True</span>)[<span class="hljs-number">0</span>]

<span class="hljs-string">&quot;CPU times: user 56.2 s, sys: 42.5 s, total: 1min 38s&quot;</span>
<span class="hljs-string">&quot;Wall time: 1min 29s&quot;</span>`,wrap:!1}}),D=new f({props:{code:"ZnJvbSUyMGRpZmZ1c2VycyUyMGltcG9ydCUyMG1ha2VfaW1hZ2VfZ3JpZCUwQSUwQWltYWdlcyUyMCUzRCUyMGltYWdlcy5yZXNoYXBlKChpbWFnZXMuc2hhcGUlNUIwJTVEJTIwKiUyMGltYWdlcy5zaGFwZSU1QjElNUQlMkMpJTIwJTJCJTIwaW1hZ2VzLnNoYXBlJTVCLTMlM0ElNUQpJTBBaW1hZ2VzJTIwJTNEJTIwcGlwZWxpbmUubnVtcHlfdG9fcGlsKGltYWdlcyklMEFtYWtlX2ltYWdlX2dyaWQoaW1hZ2VzJTJDJTIwMiUyQyUyMDQp",highlighted:`<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> make_image_grid

images = images.reshape((images.shape[<span class="hljs-number">0</span>] * images.shape[<span class="hljs-number">1</span>],) + images.shape[-<span class="hljs-number">3</span>:])
images = pipeline.numpy_to_pil(images)
make_image_grid(images, <span class="hljs-number">2</span>, <span class="hljs-number">4</span>)`,wrap:!1}}),O=new Te({props:{title:"Using different prompts",local:"using-different-prompts",headingTag:"h2"}}),te=new f({props:{code:"cHJvbXB0cyUyMCUzRCUyMCU1QiUwQSUyMCUyMCUyMCUyMCUyMkxhYnJhZG9yJTIwaW4lMjB0aGUlMjBzdHlsZSUyMG9mJTIwSG9rdXNhaSUyMiUyQyUwQSUyMCUyMCUyMCUyMCUyMlBhaW50aW5nJTIwb2YlMjBhJTIwc3F1aXJyZWwlMjBza2F0aW5nJTIwaW4lMjBOZXclMjBZb3JrJTIyJTJDJTBBJTIwJTIwJTIwJTIwJTIySEFMLTkwMDAlMjBpbiUyMHRoZSUyMHN0eWxlJTIwb2YlMjBWYW4lMjBHb2doJTIyJTJDJTBBJTIwJTIwJTIwJTIwJTIyVGltZXMlMjBTcXVhcmUlMjB1bmRlciUyMHdhdGVyJTJDJTIwd2l0aCUyMGZpc2glMjBhbmQlMjBhJTIwZG9scGhpbiUyMHN3aW1taW5nJTIwYXJvdW5kJTIyJTJDJTBBJTIwJTIwJTIwJTIwJTIyQW5jaWVudCUyMFJvbWFuJTIwZnJlc2NvJTIwc2hvd2luZyUyMGElMjBtYW4lMjB3b3JraW5nJTIwb24lMjBoaXMlMjBsYXB0b3AlMjIlMkMlMEElMjAlMjAlMjAlMjAlMjJDbG9zZS11cCUyMHBob3RvZ3JhcGglMjBvZiUyMHlvdW5nJTIwYmxhY2slMjB3b21hbiUyMGFnYWluc3QlMjB1cmJhbiUyMGJhY2tncm91bmQlMkMlMjBoaWdoJTIwcXVhbGl0eSUyQyUyMGJva2VoJTIyJTJDJTBBJTIwJTIwJTIwJTIwJTIyQXJtY2hhaXIlMjBpbiUyMHRoZSUyMHNoYXBlJTIwb2YlMjBhbiUyMGF2b2NhZG8lMjIlMkMlMEElMjAlMjAlMjAlMjAlMjJDbG93biUyMGFzdHJvbmF1dCUyMGluJTIwc3BhY2UlMkMlMjB3aXRoJTIwRWFydGglMjBpbiUyMHRoZSUyMGJhY2tncm91bmQlMjIlMkMlMEElNUQlMEElMEFwcm9tcHRfaWRzJTIwJTNEJTIwcGlwZWxpbmUucHJlcGFyZV9pbnB1dHMocHJvbXB0cyklMEFwcm9tcHRfaWRzJTIwJTNEJTIwc2hhcmQocHJvbXB0X2lkcyklMEElMEFpbWFnZXMlMjAlM0QlMjBwaXBlbGluZShwcm9tcHRfaWRzJTJDJTIwcF9wYXJhbXMlMkMlMjBybmclMkMlMjBqaXQlM0RUcnVlKS5pbWFnZXMlMEFpbWFnZXMlMjAlM0QlMjBpbWFnZXMucmVzaGFwZSgoaW1hZ2VzLnNoYXBlJTVCMCU1RCUyMColMjBpbWFnZXMuc2hhcGUlNUIxJTVEJTJDKSUyMCUyQiUyMGltYWdlcy5zaGFwZSU1Qi0zJTNBJTVEKSUwQWltYWdlcyUyMCUzRCUyMHBpcGVsaW5lLm51bXB5X3RvX3BpbChpbWFnZXMpJTBBJTBBbWFrZV9pbWFnZV9ncmlkKGltYWdlcyUyQyUyMDIlMkMlMjA0KQ==",highlighted:`prompts = [
    <span class="hljs-string">&quot;Labrador in the style of Hokusai&quot;</span>,
    <span class="hljs-string">&quot;Painting of a squirrel skating in New York&quot;</span>,
    <span class="hljs-string">&quot;HAL-9000 in the style of Van Gogh&quot;</span>,
    <span class="hljs-string">&quot;Times Square under water, with fish and a dolphin swimming around&quot;</span>,
    <span class="hljs-string">&quot;Ancient Roman fresco showing a man working on his laptop&quot;</span>,
    <span class="hljs-string">&quot;Close-up photograph of young black woman against urban background, high quality, bokeh&quot;</span>,
    <span class="hljs-string">&quot;Armchair in the shape of an avocado&quot;</span>,
    <span class="hljs-string">&quot;Clown astronaut in space, with Earth in the background&quot;</span>,
]

prompt_ids = pipeline.prepare_inputs(prompts)
prompt_ids = shard(prompt_ids)

images = pipeline(prompt_ids, p_params, rng, jit=<span class="hljs-literal">True</span>).images
images = images.reshape((images.shape[<span class="hljs-number">0</span>] * images.shape[<span class="hljs-number">1</span>],) + images.shape[-<span class="hljs-number">3</span>:])
images = pipeline.numpy_to_pil(images)

make_image_grid(images, <span class="hljs-number">2</span>, <span class="hljs-number">4</span>)`,wrap:!1}}),se=new Te({props:{title:"How does parallelization work?",local:"how-does-parallelization-work",headingTag:"h2"}}),re=new f({props:{code:"cF9nZW5lcmF0ZSUyMCUzRCUyMHBtYXAocGlwZWxpbmUuX2dlbmVyYXRlKQ==",highlighted:"p_generate = pmap(pipeline._generate)",wrap:!1}}),he=new f({props:{code:"JTI1JTI1dGltZSUwQWltYWdlcyUyMCUzRCUyMHBfZ2VuZXJhdGUocHJvbXB0X2lkcyUyQyUyMHBfcGFyYW1zJTJDJTIwcm5nKSUwQWltYWdlcyUyMCUzRCUyMGltYWdlcy5ibG9ja191bnRpbF9yZWFkeSgpJTBBJTIyQ1BVJTIwdGltZXMlM0ElMjB1c2VyJTIwMW1pbiUyMDE1cyUyQyUyMHN5cyUzQSUyMDE4LjIlMjBzJTJDJTIwdG90YWwlM0ElMjAxbWluJTIwMzRzJTIyJTBBJTIyV2FsbCUyMHRpbWUlM0ElMjAxbWluJTIwMTVzJTIy",highlighted:`%%time
images = p_generate(prompt_ids, p_params, rng)
images = images.block_until_ready()
<span class="hljs-string">&quot;CPU times: user 1min 15s, sys: 18.2 s, total: 1min 34s&quot;</span>
<span class="hljs-string">&quot;Wall time: 1min 15s&quot;</span>`,wrap:!1}}),ye=new f({props:{code:"aW1hZ2VzLnNoYXBlJTBBJTIyKDglMkMlMjAxJTJDJTIwNTEyJTJDJTIwNTEyJTJDJTIwMyklMjI=",highlighted:`images.shape
<span class="hljs-string">&quot;(8, 1, 512, 512, 3)&quot;</span>`,wrap:!1}}),{c(){r=i("meta"),w=a(),y=i("p"),T=a(),m(j.$$.fragment),be=a(),m(U.$$.fragment),Je=a(),g=i("p"),g.textContent=wt,je=a(),$=i("p"),$.textContent=Tt,Ue=a(),m(v.$$.fragment),ge=a(),B=i("p"),B.textContent=bt,$e=a(),_=i("p"),_.innerHTML=Jt,ve=a(),m(C.$$.fragment),Be=a(),k=i("p"),k.textContent=jt,_e=a(),m(I.$$.fragment),Ce=a(),m(x.$$.fragment),ke=a(),G=i("p"),G.innerHTML=Ut,Ie=a(),m(W.$$.fragment),xe=a(),m(X.$$.fragment),Ge=a(),Z=i("p"),Z.textContent=gt,We=a(),m(b.$$.fragment),Xe=a(),H=i("p"),H.innerHTML=$t,Ze=a(),m(F.$$.fragment),He=a(),Y=i("p"),Y.innerHTML=vt,Fe=a(),m(E.$$.fragment),Ye=a(),Q=i("p"),Q.innerHTML=Bt,Ee=a(),R=i("p"),R.textContent=_t,Qe=a(),z=i("p"),z.textContent=Ct,Re=a(),m(L.$$.fragment),ze=a(),N=i("p"),N.innerHTML=kt,Le=a(),m(V.$$.fragment),Ne=a(),S=i("p"),S.innerHTML=It,Ve=a(),m(J.$$.fragment),Se=a(),A=i("p"),A.innerHTML=xt,Ae=a(),m(q.$$.fragment),qe=a(),P=i("p"),P.innerHTML=Gt,Pe=a(),m(D.$$.fragment),De=a(),K=i("p"),K.innerHTML=Wt,Ke=a(),m(O.$$.fragment),Oe=a(),ee=i("p"),ee.textContent=Xt,et=a(),m(te.$$.fragment),tt=a(),le=i("p"),le.innerHTML=Zt,lt=a(),m(se.$$.fragment),st=a(),ae=i("p"),ae.textContent=Ht,at=a(),ne=i("p"),ne.innerHTML=Ft,nt=a(),ie=i("p"),ie.innerHTML=Yt,it=a(),oe=i("ol"),oe.innerHTML=Et,ot=a(),pe=i("p"),pe.innerHTML=Qt,pt=a(),m(re.$$.fragment),rt=a(),me=i("p"),me.innerHTML=Rt,mt=a(),ce=i("ol"),ce.innerHTML=zt,ct=a(),ue=i("p"),ue.textContent=Lt,ut=a(),de=i("p"),de.innerHTML=Nt,dt=a(),m(he.$$.fragment),ht=a(),Me=i("p"),Me.textContent=Vt,Mt=a(),m(ye.$$.fragment),yt=a(),we=i("p"),this.h()},l(e){const t=sl("svelte-u9bgzb",document.head);r=o(t,"META",{name:!0,content:!0}),t.forEach(l),w=n(e),y=o(e,"P",{}),qt(y).forEach(l),T=n(e),c(j.$$.fragment,e),be=n(e),c(U.$$.fragment,e),Je=n(e),g=o(e,"P",{"data-svelte-h":!0}),p(g)!=="svelte-12svr29"&&(g.textContent=wt),je=n(e),$=o(e,"P",{"data-svelte-h":!0}),p($)!=="svelte-vd07ss"&&($.textContent=Tt),Ue=n(e),c(v.$$.fragment,e),ge=n(e),B=o(e,"P",{"data-svelte-h":!0}),p(B)!=="svelte-15famve"&&(B.textContent=bt),$e=n(e),_=o(e,"P",{"data-svelte-h":!0}),p(_)!=="svelte-1jp1nmd"&&(_.innerHTML=Jt),ve=n(e),c(C.$$.fragment,e),Be=n(e),k=o(e,"P",{"data-svelte-h":!0}),p(k)!=="svelte-5ufrkr"&&(k.textContent=jt),_e=n(e),c(I.$$.fragment,e),Ce=n(e),c(x.$$.fragment,e),ke=n(e),G=o(e,"P",{"data-svelte-h":!0}),p(G)!=="svelte-349jsc"&&(G.innerHTML=Ut),Ie=n(e),c(W.$$.fragment,e),xe=n(e),c(X.$$.fragment,e),Ge=n(e),Z=o(e,"P",{"data-svelte-h":!0}),p(Z)!=="svelte-19grb2g"&&(Z.textContent=gt),We=n(e),c(b.$$.fragment,e),Xe=n(e),H=o(e,"P",{"data-svelte-h":!0}),p(H)!=="svelte-ih3iqi"&&(H.innerHTML=$t),Ze=n(e),c(F.$$.fragment,e),He=n(e),Y=o(e,"P",{"data-svelte-h":!0}),p(Y)!=="svelte-b3wqu4"&&(Y.innerHTML=vt),Fe=n(e),c(E.$$.fragment,e),Ye=n(e),Q=o(e,"P",{"data-svelte-h":!0}),p(Q)!=="svelte-18q6dmb"&&(Q.innerHTML=Bt),Ee=n(e),R=o(e,"P",{"data-svelte-h":!0}),p(R)!=="svelte-j1l5pg"&&(R.textContent=_t),Qe=n(e),z=o(e,"P",{"data-svelte-h":!0}),p(z)!=="svelte-1kfts0e"&&(z.textContent=Ct),Re=n(e),c(L.$$.fragment,e),ze=n(e),N=o(e,"P",{"data-svelte-h":!0}),p(N)!=="svelte-bsl0kg"&&(N.innerHTML=kt),Le=n(e),c(V.$$.fragment,e),Ne=n(e),S=o(e,"P",{"data-svelte-h":!0}),p(S)!=="svelte-1sp1ijt"&&(S.innerHTML=It),Ve=n(e),c(J.$$.fragment,e),Se=n(e),A=o(e,"P",{"data-svelte-h":!0}),p(A)!=="svelte-1xv75ai"&&(A.innerHTML=xt),Ae=n(e),c(q.$$.fragment,e),qe=n(e),P=o(e,"P",{"data-svelte-h":!0}),p(P)!=="svelte-1qaylo"&&(P.innerHTML=Gt),Pe=n(e),c(D.$$.fragment,e),De=n(e),K=o(e,"P",{"data-svelte-h":!0}),p(K)!=="svelte-e1s2k"&&(K.innerHTML=Wt),Ke=n(e),c(O.$$.fragment,e),Oe=n(e),ee=o(e,"P",{"data-svelte-h":!0}),p(ee)!=="svelte-9e9hdw"&&(ee.textContent=Xt),et=n(e),c(te.$$.fragment,e),tt=n(e),le=o(e,"P",{"data-svelte-h":!0}),p(le)!=="svelte-1nea4f2"&&(le.innerHTML=Zt),lt=n(e),c(se.$$.fragment,e),st=n(e),ae=o(e,"P",{"data-svelte-h":!0}),p(ae)!=="svelte-1217ir9"&&(ae.textContent=Ht),at=n(e),ne=o(e,"P",{"data-svelte-h":!0}),p(ne)!=="svelte-1toapy0"&&(ne.innerHTML=Ft),nt=n(e),ie=o(e,"P",{"data-svelte-h":!0}),p(ie)!=="svelte-1pwzekg"&&(ie.innerHTML=Yt),it=n(e),oe=o(e,"OL",{"data-svelte-h":!0}),p(oe)!=="svelte-12i8xv7"&&(oe.innerHTML=Et),ot=n(e),pe=o(e,"P",{"data-svelte-h":!0}),p(pe)!=="svelte-1gtj740"&&(pe.innerHTML=Qt),pt=n(e),c(re.$$.fragment,e),rt=n(e),me=o(e,"P",{"data-svelte-h":!0}),p(me)!=="svelte-1774hdu"&&(me.innerHTML=Rt),mt=n(e),ce=o(e,"OL",{"data-svelte-h":!0}),p(ce)!=="svelte-ngcyce"&&(ce.innerHTML=zt),ct=n(e),ue=o(e,"P",{"data-svelte-h":!0}),p(ue)!=="svelte-w1m4t9"&&(ue.textContent=Lt),ut=n(e),de=o(e,"P",{"data-svelte-h":!0}),p(de)!=="svelte-1ue3hjw"&&(de.innerHTML=Nt),dt=n(e),c(he.$$.fragment,e),ht=n(e),Me=o(e,"P",{"data-svelte-h":!0}),p(Me)!=="svelte-735q12"&&(Me.textContent=Vt),Mt=n(e),c(ye.$$.fragment,e),yt=n(e),we=o(e,"P",{}),qt(we).forEach(l),this.h()},h(){Pt(r,"name","hf:doc:metadata"),Pt(r,"content",rl)},m(e,t){al(document.head,r),s(e,w,t),s(e,y,t),s(e,T,t),u(j,e,t),s(e,be,t),u(U,e,t),s(e,Je,t),s(e,g,t),s(e,je,t),s(e,$,t),s(e,Ue,t),u(v,e,t),s(e,ge,t),s(e,B,t),s(e,$e,t),s(e,_,t),s(e,ve,t),u(C,e,t),s(e,Be,t),s(e,k,t),s(e,_e,t),u(I,e,t),s(e,Ce,t),u(x,e,t),s(e,ke,t),s(e,G,t),s(e,Ie,t),u(W,e,t),s(e,xe,t),u(X,e,t),s(e,Ge,t),s(e,Z,t),s(e,We,t),u(b,e,t),s(e,Xe,t),s(e,H,t),s(e,Ze,t),u(F,e,t),s(e,He,t),s(e,Y,t),s(e,Fe,t),u(E,e,t),s(e,Ye,t),s(e,Q,t),s(e,Ee,t),s(e,R,t),s(e,Qe,t),s(e,z,t),s(e,Re,t),u(L,e,t),s(e,ze,t),s(e,N,t),s(e,Le,t),u(V,e,t),s(e,Ne,t),s(e,S,t),s(e,Ve,t),u(J,e,t),s(e,Se,t),s(e,A,t),s(e,Ae,t),u(q,e,t),s(e,qe,t),s(e,P,t),s(e,Pe,t),u(D,e,t),s(e,De,t),s(e,K,t),s(e,Ke,t),u(O,e,t),s(e,Oe,t),s(e,ee,t),s(e,et,t),u(te,e,t),s(e,tt,t),s(e,le,t),s(e,lt,t),u(se,e,t),s(e,st,t),s(e,ae,t),s(e,at,t),s(e,ne,t),s(e,nt,t),s(e,ie,t),s(e,it,t),s(e,oe,t),s(e,ot,t),s(e,pe,t),s(e,pt,t),u(re,e,t),s(e,rt,t),s(e,me,t),s(e,mt,t),s(e,ce,t),s(e,ct,t),s(e,ue,t),s(e,ut,t),s(e,de,t),s(e,dt,t),u(he,e,t),s(e,ht,t),s(e,Me,t),s(e,Mt,t),u(ye,e,t),s(e,yt,t),s(e,we,t),ft=!0},p(e,[t]){const St={};t&2&&(St.$$scope={dirty:t,ctx:e}),b.$set(St);const At={};t&2&&(At.$$scope={dirty:t,ctx:e}),J.$set(At)},i(e){ft||(d(j.$$.fragment,e),d(U.$$.fragment,e),d(v.$$.fragment,e),d(C.$$.fragment,e),d(I.$$.fragment,e),d(x.$$.fragment,e),d(W.$$.fragment,e),d(X.$$.fragment,e),d(b.$$.fragment,e),d(F.$$.fragment,e),d(E.$$.fragment,e),d(L.$$.fragment,e),d(V.$$.fragment,e),d(J.$$.fragment,e),d(q.$$.fragment,e),d(D.$$.fragment,e),d(O.$$.fragment,e),d(te.$$.fragment,e),d(se.$$.fragment,e),d(re.$$.fragment,e),d(he.$$.fragment,e),d(ye.$$.fragment,e),ft=!0)},o(e){h(j.$$.fragment,e),h(U.$$.fragment,e),h(v.$$.fragment,e),h(C.$$.fragment,e),h(I.$$.fragment,e),h(x.$$.fragment,e),h(W.$$.fragment,e),h(X.$$.fragment,e),h(b.$$.fragment,e),h(F.$$.fragment,e),h(E.$$.fragment,e),h(L.$$.fragment,e),h(V.$$.fragment,e),h(J.$$.fragment,e),h(q.$$.fragment,e),h(D.$$.fragment,e),h(O.$$.fragment,e),h(te.$$.fragment,e),h(se.$$.fragment,e),h(re.$$.fragment,e),h(he.$$.fragment,e),h(ye.$$.fragment,e),ft=!1},d(e){e&&(l(w),l(y),l(T),l(be),l(Je),l(g),l(je),l($),l(Ue),l(ge),l(B),l($e),l(_),l(ve),l(Be),l(k),l(_e),l(Ce),l(ke),l(G),l(Ie),l(xe),l(Ge),l(Z),l(We),l(Xe),l(H),l(Ze),l(He),l(Y),l(Fe),l(Ye),l(Q),l(Ee),l(R),l(Qe),l(z),l(Re),l(ze),l(N),l(Le),l(Ne),l(S),l(Ve),l(Se),l(A),l(Ae),l(qe),l(P),l(Pe),l(De),l(K),l(Ke),l(Oe),l(ee),l(et),l(tt),l(le),l(lt),l(st),l(ae),l(at),l(ne),l(nt),l(ie),l(it),l(oe),l(ot),l(pe),l(pt),l(rt),l(me),l(mt),l(ce),l(ct),l(ue),l(ut),l(de),l(dt),l(ht),l(Me),l(Mt),l(yt),l(we)),l(r),M(j,e),M(U,e),M(v,e),M(C,e),M(I,e),M(x,e),M(W,e),M(X,e),M(b,e),M(F,e),M(E,e),M(L,e),M(V,e),M(J,e),M(q,e),M(D,e),M(O,e),M(te,e),M(se,e),M(re,e),M(he,e),M(ye,e)}}}const rl='{"title":"JAX/Flax","local":"jaxflax","sections":[{"title":"Load a model","local":"load-a-model","sections":[],"depth":2},{"title":"Inference","local":"inference","sections":[],"depth":2},{"title":"Using different prompts","local":"using-different-prompts","sections":[],"depth":2},{"title":"How does parallelization work?","local":"how-does-parallelization-work","sections":[],"depth":2}],"depth":1}';function ml(fe){return el(()=>{new URLSearchParams(window.location.search).get("fw")}),[]}class fl extends tl{constructor(r){super(),ll(this,r,ml,pl,Ot,{})}}export{fl as component};