diffusers/pr_14313/en/_app/immutable/nodes/314.DdgVPvwV.js

import"../chunks/DsnmJJEf.js";import{i as X,h as I,C as R,H as t,a,E as x,s as G}from"../chunks/CmJXCtRL.js";import{p as V,o as B,s as e,f as k,a as U,b as C,d as T,n as v}from"../chunks/DK803DsY.js";const N='{"title":"Wuerstchen","local":"wuerstchen","sections":[{"title":"Script parameters","local":"script-parameters","sections":[],"depth":2},{"title":"Training script","local":"training-script","sections":[],"depth":2},{"title":"Launch the script","local":"launch-the-script","sections":[],"depth":2}],"depth":1}';var F=T('<meta name="hf:doc:metadata"/>'),E=T('<p></p> <!> <!> <p>The <a href="https://hf.co/papers/2306.00637" rel="nofollow">Wuerstchen</a> model drastically reduces computational costs by compressing the latent space by 42x, without compromising image quality and accelerating inference. During training, Wuerstchen uses two models (VQGAN + autoencoder) to compress the latents, and then a third model (text-conditioned latent diffusion model) is conditioned on this highly compressed space to generate an image.</p> <p>To fit the prior model into GPU memory and to speedup training, try enabling <code>gradient_accumulation_steps</code>, <code>gradient_checkpointing</code>, and <code>mixed_precision</code> respectively.</p> <p>This guide explores the <a href="https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/train_text_to_image_prior.py" rel="nofollow">train_text_to_image_prior.py</a> script to help you become more familiar with it, and how you can adapt it for your own use-case.</p> <p>Before running the script, make sure you install the library from source:</p> <!> <p>Then navigate to the example folder containing the training script and install the required dependencies for the script you’re using:</p> <!> <blockquote class="tip"><p>🤗 Accelerate is a library for helping you train on multiple GPUs/TPUs or with mixed-precision. It’ll automatically configure your training setup based on your hardware and environment. Take a look at the 🤗 Accelerate <a href="https://huggingface.co/docs/accelerate/quicktour" rel="nofollow">Quick tour</a> to learn more.</p></blockquote> <p>Initialize an 🤗 Accelerate environment:</p> <!> <p>To setup a default 🤗 Accelerate environment without choosing any configurations:</p> <!> <p>Or if your environment doesn’t support an interactive shell, like a notebook, you can use:</p> <!> <p>Lastly, if you want to train a model on your own dataset, take a look at the <a href="create_dataset">Create a dataset for training</a> guide to learn how to create a dataset that works with the training script.</p> <blockquote class="tip"><p>The following sections highlight parts of the training scripts that are important for understanding how to modify it, but it doesn’t cover every aspect of the <a href="https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/train_text_to_image_prior.py" rel="nofollow">script</a> in detail. If you’re interested in learning more, feel free to read through the scripts and let us know if you have any questions or concerns.</p></blockquote> <!> <p>The training scripts provides many parameters to help you customize your training run. All of the parameters and their descriptions are found in the <a href="https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/wuerstchen/text_to_image/train_text_to_image_prior.py#L192" rel="nofollow"><code>parse_args()</code></a> function. It provides default values for each parameter, such as the training batch size and learning rate, but you can also set your own values in the training command if you’d like.</p> <p>For example, to speedup training with mixed precision using the fp16 format, add the <code>--mixed_precision</code> parameter to the training command:</p> <!> <p>Most of the parameters are identical to the parameters in the <a href="text2image#script-parameters">Text-to-image</a> training guide, so let’s dive right into the Wuerstchen training script!</p> <!> <p>The training script is also similar to the <a href="text2image#training-script">Text-to-image</a> training guide, but it’s been modified to support Wuerstchen. This guide focuses on the code that is unique to the Wuerstchen training script.</p> <p>The <a href="https://github.com/huggingface/diffusers/blob/6e68c71503682c8693cb5b06a4da4911dfd655ee/examples/wuerstchen/text_to_image/train_text_to_image_prior.py#L441" rel="nofollow"><code>main()</code></a> function starts by initializing the image encoder - an <a href="https://github.com/huggingface/diffusers/blob/main/examples/wuerstchen/text_to_image/modeling_efficient_net_encoder.py" rel="nofollow">EfficientNet</a> - in addition to the usual scheduler and tokenizer.</p> <!> <p>You’ll also load the <code>WuerstchenPrior</code> model for optimization.</p> <!> <p>Next, you’ll apply some <a href="https://github.com/huggingface/diffusers/blob/65ef7a0c5c594b4f84092e328fbdd73183613b30/examples/wuerstchen/text_to_image/train_text_to_image_prior.py#L656" rel="nofollow">transforms</a> to the images and <a href="https://github.com/huggingface/diffusers/blob/65ef7a0c5c594b4f84092e328fbdd73183613b30/examples/wuerstchen/text_to_image/train_text_to_image_prior.py#L637" rel="nofollow">tokenize</a> the captions:</p> <!> <p>Finally, the <a href="https://github.com/huggingface/diffusers/blob/65ef7a0c5c594b4f84092e328fbdd73183613b30/examples/wuerstchen/text_to_image/train_text_to_image_prior.py#L656" rel="nofollow">training loop</a> handles compressing the images to latent space with the <code>EfficientNetEncoder</code>, adding noise to the latents, and predicting the noise residual with the <code>WuerstchenPrior</code> model.</p> <!> <p>If you want to learn more about how the training loop works, check out the <a href="../using-diffusers/write_own_pipeline">Understanding pipelines, models and schedulers</a> tutorial which breaks down the basic pattern of the denoising process.</p> <!> <p>Once you’ve made all your changes or you’re okay with the default configuration, you’re ready to launch the training script! 🚀</p> <p>Set the <code>DATASET_NAME</code> environment variable to the dataset name from the Hub. This guide uses the <a href="https://huggingface.co/datasets/lambdalabs/naruto-blip-captions" rel="nofollow">Naruto BLIP captions</a> dataset, but you can create and train on your own datasets as well (see the <a href="create_dataset">Create a dataset for training</a> guide).</p> <blockquote class="tip"><p>To monitor training progress with Weights & Biases, add the <code>--report_to=wandb</code> parameter to the training command. You’ll also need to add the <code>--validation_prompt</code> to the training command to keep track of results. This can be really useful for debugging the model and viewing intermediate results.</p></blockquote> <!> <p>Once training is complete, you can use your newly trained model for inference!</p> <!> <!> <p></p>',1);function S(j,Z){V(Z,!1),B(()=>{new URLSearchParams(window.location.search).get("fw")}),X();var l=E();I("q4kcug",_=>{var J=F();G(J,"content",N),U(_,J)});var s=e(k(l),2);R(s,{containerStyle:"float: right; margin-left: 10px; display: inline-flex; position: relative; z-index: 10;"});var n=e(s,2);t(n,{title:"Wuerstchen",local:"wuerstchen",headingTag:"h1"});var i=e(n,10);a(i,{code:"Z2l0JTIwY2xvbmUlMjBodHRwcyUzQSUyRiUyRmdpdGh1Yi5jb20lMkZodWdnaW5nZmFjZSUyRmRpZmZ1c2VycyUwQWNkJTIwZGlmZnVzZXJzJTBBcGlwJTIwaW5zdGFsbCUyMC4=",highlighted:`git <span class="hljs-built_in">clone</span> https://github.com/huggingface/diffusers
<span class="hljs-built_in">cd</span> diffusers
pip install .`,lang:"bash",wrap:!1});var o=e(i,4);a(o,{code:"Y2QlMjBleGFtcGxlcyUyRnd1ZXJzdGNoZW4lMkZ0ZXh0X3RvX2ltYWdlJTBBcGlwJTIwaW5zdGFsbCUyMC1yJTIwcmVxdWlyZW1lbnRzLnR4dA==",highlighted:`<span class="hljs-built_in">cd</span> examples/wuerstchen/text_to_image
pip install -r requirements.txt`,lang:"bash",wrap:!1});var r=e(o,6);a(r,{code:"YWNjZWxlcmF0ZSUyMGNvbmZpZw==",highlighted:"accelerate config",lang:"bash",wrap:!1});var c=e(r,4);a(c,{code:"YWNjZWxlcmF0ZSUyMGNvbmZpZyUyMGRlZmF1bHQ=",highlighted:"accelerate config default",lang:"bash",wrap:!1});var p=e(c,4);a(p,{code:"ZnJvbSUyMGFjY2VsZXJhdGUudXRpbHMlMjBpbXBvcnQlMjB3cml0ZV9iYXNpY19jb25maWclMEElMEF3cml0ZV9iYXNpY19jb25maWcoKQ==",highlighted:`<span class="hljs-keyword">from</span> accelerate.utils <span class="hljs-keyword">import</span> write_basic_config

write_basic_config()`,lang:"py",wrap:!1});var d=e(p,6);t(d,{title:"Script parameters",local:"script-parameters",headingTag:"h2"});var h=e(d,6);a(h,{code:"YWNjZWxlcmF0ZSUyMGxhdW5jaCUyMHRyYWluX3RleHRfdG9faW1hZ2VfcHJpb3IucHklMjAlNUMlMEElMjAlMjAtLW1peGVkX3ByZWNpc2lvbiUzRCUyMmZwMTYlMjI=",highlighted:`accelerate launch train_text_to_image_prior.py \\
  --mixed_precision=<span class="hljs-string">&quot;fp16&quot;</span>`,lang:"bash",wrap:!1});var m=e(h,4);t(m,{title:"Training script",local:"training-script",headingTag:"h2"});var u=e(m,6);a(u,{code:"d2l0aCUyMENvbnRleHRNYW5hZ2VycyhkZWVwc3BlZWRfemVyb19pbml0X2Rpc2FibGVkX2NvbnRleHRfbWFuYWdlcigpKSUzQSUwQSUyMCUyMCUyMCUyMHByZXRyYWluZWRfY2hlY2twb2ludF9maWxlJTIwJTNEJTIwaGZfaHViX2Rvd25sb2FkKCUyMmRvbWUyNzIlMkZ3dWVyc3RjaGVuJTIyJTJDJTIwZmlsZW5hbWUlM0QlMjJtb2RlbF92Ml9zdGFnZV9iLnB0JTIyKSUwQSUyMCUyMCUyMCUyMHN0YXRlX2RpY3QlMjAlM0QlMjB0b3JjaC5sb2FkKHByZXRyYWluZWRfY2hlY2twb2ludF9maWxlJTJDJTIwbWFwX2xvY2F0aW9uJTNEJTIyY3B1JTIyKSUwQSUyMCUyMCUyMCUyMGltYWdlX2VuY29kZXIlMjAlM0QlMjBFZmZpY2llbnROZXRFbmNvZGVyKCklMEElMjAlMjAlMjAlMjBpbWFnZV9lbmNvZGVyLmxvYWRfc3RhdGVfZGljdChzdGF0ZV9kaWN0JTVCJTIyZWZmbmV0X3N0YXRlX2RpY3QlMjIlNUQpJTBBJTIwJTIwJTIwJTIwaW1hZ2VfZW5jb2Rlci5ldmFsKCk=",highlighted:`<span class="hljs-keyword">with</span> ContextManagers(deepspeed_zero_init_disabled_context_manager()):
    pretrained_checkpoint_file = hf_hub_download(<span class="hljs-string">&quot;dome272/wuerstchen&quot;</span>, filename=<span class="hljs-string">&quot;model_v2_stage_b.pt&quot;</span>)
    state_dict = torch.load(pretrained_checkpoint_file, map_location=<span class="hljs-string">&quot;cpu&quot;</span>)
    image_encoder = EfficientNetEncoder()
    image_encoder.load_state_dict(state_dict[<span class="hljs-string">&quot;effnet_state_dict&quot;</span>])
    image_encoder.<span class="hljs-built_in">eval</span>()`,lang:"py",wrap:!1});var y=e(u,4);a(y,{code:"cHJpb3IlMjAlM0QlMjBXdWVyc3RjaGVuUHJpb3IuZnJvbV9wcmV0cmFpbmVkKGFyZ3MucHJldHJhaW5lZF9wcmlvcl9tb2RlbF9uYW1lX29yX3BhdGglMkMlMjBzdWJmb2xkZXIlM0QlMjJwcmlvciUyMiklMEElMEFvcHRpbWl6ZXIlMjAlM0QlMjBvcHRpbWl6ZXJfY2xzKCUwQSUyMCUyMCUyMCUyMHByaW9yLnBhcmFtZXRlcnMoKSUyQyUwQSUyMCUyMCUyMCUyMGxyJTNEYXJncy5sZWFybmluZ19yYXRlJTJDJTBBJTIwJTIwJTIwJTIwYmV0YXMlM0QoYXJncy5hZGFtX2JldGExJTJDJTIwYXJncy5hZGFtX2JldGEyKSUyQyUwQSUyMCUyMCUyMCUyMHdlaWdodF9kZWNheSUzRGFyZ3MuYWRhbV93ZWlnaHRfZGVjYXklMkMlMEElMjAlMjAlMjAlMjBlcHMlM0RhcmdzLmFkYW1fZXBzaWxvbiUyQyUwQSk=",highlighted:`prior = WuerstchenPrior.from_pretrained(args.pretrained_prior_model_name_or_path, subfolder=<span class="hljs-string">&quot;prior&quot;</span>)

optimizer = optimizer_cls(
    prior.parameters(),
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)`,lang:"py",wrap:!1});var M=e(y,4);a(M,{code:"ZGVmJTIwcHJlcHJvY2Vzc190cmFpbihleGFtcGxlcyklM0ElMEElMjAlMjAlMjAlMjBpbWFnZXMlMjAlM0QlMjAlNUJpbWFnZS5jb252ZXJ0KCUyMlJHQiUyMiklMjBmb3IlMjBpbWFnZSUyMGluJTIwZXhhbXBsZXMlNUJpbWFnZV9jb2x1bW4lNUQlNUQlMEElMjAlMjAlMjAlMjBleGFtcGxlcyU1QiUyMmVmZm5ldF9waXhlbF92YWx1ZXMlMjIlNUQlMjAlM0QlMjAlNUJlZmZuZXRfdHJhbnNmb3JtcyhpbWFnZSklMjBmb3IlMjBpbWFnZSUyMGluJTIwaW1hZ2VzJTVEJTBBJTIwJTIwJTIwJTIwZXhhbXBsZXMlNUIlMjJ0ZXh0X2lucHV0X2lkcyUyMiU1RCUyQyUyMGV4YW1wbGVzJTVCJTIydGV4dF9tYXNrJTIyJTVEJTIwJTNEJTIwdG9rZW5pemVfY2FwdGlvbnMoZXhhbXBsZXMpJTBBJTIwJTIwJTIwJTIwcmV0dXJuJTIwZXhhbXBsZXM=",highlighted:`<span class="hljs-keyword">def</span> <span class="hljs-title function_">preprocess_train</span>(<span class="hljs-params">examples</span>):
    images = [image.convert(<span class="hljs-string">&quot;RGB&quot;</span>) <span class="hljs-keyword">for</span> image <span class="hljs-keyword">in</span> examples[image_column]]
    examples[<span class="hljs-string">&quot;effnet_pixel_values&quot;</span>] = [effnet_transforms(image) <span class="hljs-keyword">for</span> image <span class="hljs-keyword">in</span> images]
    examples[<span class="hljs-string">&quot;text_input_ids&quot;</span>], examples[<span class="hljs-string">&quot;text_mask&quot;</span>] = tokenize_captions(examples)
    <span class="hljs-keyword">return</span> examples`,lang:"py",wrap:!1});var g=e(M,4);a(g,{code:"cHJlZF9ub2lzZSUyMCUzRCUyMHByaW9yKG5vaXN5X2xhdGVudHMlMkMlMjB0aW1lc3RlcHMlMkMlMjBwcm9tcHRfZW1iZWRzKQ==",highlighted:"pred_noise = prior(noisy_latents, timesteps, prompt_embeds)",lang:"py",wrap:!1});var b=e(g,4);t(b,{title:"Launch the script",local:"launch-the-script",headingTag:"h2"});var f=e(b,8);a(f,{code:"ZXhwb3J0JTIwREFUQVNFVF9OQU1FJTNEJTIybGFtYmRhbGFicyUyRm5hcnV0by1ibGlwLWNhcHRpb25zJTIyJTBBJTBBYWNjZWxlcmF0ZSUyMGxhdW5jaCUyMCUyMHRyYWluX3RleHRfdG9faW1hZ2VfcHJpb3IucHklMjAlNUMlMEElMjAlMjAtLW1peGVkX3ByZWNpc2lvbiUzRCUyMmZwMTYlMjIlMjAlNUMlMEElMjAlMjAtLWRhdGFzZXRfbmFtZSUzRCUyNERBVEFTRVRfTkFNRSUyMCU1QyUwQSUyMCUyMC0tcmVzb2x1dGlvbiUzRDc2OCUyMCU1QyUwQSUyMCUyMC0tdHJhaW5fYmF0Y2hfc2l6ZSUzRDQlMjAlNUMlMEElMjAlMjAtLWdyYWRpZW50X2FjY3VtdWxhdGlvbl9zdGVwcyUzRDQlMjAlNUMlMEElMjAlMjAtLWdyYWRpZW50X2NoZWNrcG9pbnRpbmclMjAlNUMlMEElMjAlMjAtLWRhdGFsb2FkZXJfbnVtX3dvcmtlcnMlM0Q0JTIwJTVDJTBBJTIwJTIwLS1tYXhfdHJhaW5fc3RlcHMlM0QxNTAwMCUyMCU1QyUwQSUyMCUyMC0tbGVhcm5pbmdfcmF0ZSUzRDFlLTA1JTIwJTVDJTBBJTIwJTIwLS1tYXhfZ3JhZF9ub3JtJTNEMSUyMCU1QyUwQSUyMCUyMC0tY2hlY2twb2ludHNfdG90YWxfbGltaXQlM0QzJTIwJTVDJTBBJTIwJTIwLS1scl9zY2hlZHVsZXIlM0QlMjJjb25zdGFudCUyMiUyMCU1QyUwQSUyMCUyMC0tbHJfd2FybXVwX3N0ZXBzJTNEMCUyMCU1QyUwQSUyMCUyMC0tdmFsaWRhdGlvbl9wcm9tcHRzJTNEJTIyQSUyMHJvYm90JTIwbmFydXRvJTJDJTIwNGslMjBwaG90byUyMiUyMCU1QyUwQSUyMCUyMC0tcmVwb3J0X3RvJTNEJTIyd2FuZGIlMjIlMjAlNUMlMEElMjAlMjAtLXB1c2hfdG9faHViJTIwJTVDJTBBJTIwJTIwLS1vdXRwdXRfZGlyJTNEJTIyd3VlcnN0Y2hlbi1wcmlvci1uYXJ1dG8tbW9kZWwlMjI=",highlighted:`<span class="hljs-built_in">export</span> DATASET_NAME=<span class="hljs-string">&quot;lambdalabs/naruto-blip-captions&quot;</span>

accelerate launch  train_text_to_image_prior.py \\
  --mixed_precision=<span class="hljs-string">&quot;fp16&quot;</span> \\
  --dataset_name=<span class="hljs-variable">$DATASET_NAME</span> \\
  --resolution=768 \\
  --train_batch_size=4 \\
  --gradient_accumulation_steps=4 \\
  --gradient_checkpointing \\
  --dataloader_num_workers=4 \\
  --max_train_steps=15000 \\
  --learning_rate=1e-05 \\
  --max_grad_norm=1 \\
  --checkpoints_total_limit=3 \\
  --lr_scheduler=<span class="hljs-string">&quot;constant&quot;</span> \\
  --lr_warmup_steps=0 \\
  --validation_prompts=<span class="hljs-string">&quot;A robot naruto, 4k photo&quot;</span> \\
  --report_to=<span class="hljs-string">&quot;wandb&quot;</span> \\
  --push_to_hub \\
  --output_dir=<span class="hljs-string">&quot;wuerstchen-prior-naruto-model&quot;</span>`,lang:"bash",wrap:!1});var w=e(f,4);a(w,{code:"aW1wb3J0JTIwdG9yY2glMEFmcm9tJTIwZGlmZnVzZXJzJTIwaW1wb3J0JTIwQXV0b1BpcGVsaW5lRm9yVGV4dDJJbWFnZSUwQWZyb20lMjBkaWZmdXNlcnMucGlwZWxpbmVzLnd1ZXJzdGNoZW4lMjBpbXBvcnQlMjBERUZBVUxUX1NUQUdFX0NfVElNRVNURVBTJTBBJTBBcGlwZWxpbmUlMjAlM0QlMjBBdXRvUGlwZWxpbmVGb3JUZXh0MkltYWdlLmZyb21fcHJldHJhaW5lZCglMjJwYXRoJTJGdG8lMkZzYXZlZCUyRm1vZGVsJTIyJTJDJTIwZHR5cGUlM0R0b3JjaC5mbG9hdDE2KS50byglMjJjdWRhJTIyKSUwQSUwQWNhcHRpb24lMjAlM0QlMjAlMjJBJTIwY3V0ZSUyMGJpcmQlMjBuYXJ1dG8lMjBob2xkaW5nJTIwYSUyMHNoaWVsZCUyMiUwQWltYWdlcyUyMCUzRCUyMHBpcGVsaW5lKCUwQSUyMCUyMCUyMCUyMGNhcHRpb24lMkMlMEElMjAlMjAlMjAlMjB3aWR0aCUzRDEwMjQlMkMlMEElMjAlMjAlMjAlMjBoZWlnaHQlM0QxNTM2JTJDJTBBJTIwJTIwJTIwJTIwcHJpb3JfdGltZXN0ZXBzJTNEREVGQVVMVF9TVEFHRV9DX1RJTUVTVEVQUyUyQyUwQSUyMCUyMCUyMCUyMHByaW9yX2d1aWRhbmNlX3NjYWxlJTNENC4wJTJDJTBBJTIwJTIwJTIwJTIwbnVtX2ltYWdlc19wZXJfcHJvbXB0JTNEMiUyQyUwQSkuaW1hZ2Vz",highlighted:`<span class="hljs-keyword">import</span> torch
<span class="hljs-keyword">from</span> diffusers <span class="hljs-keyword">import</span> AutoPipelineForText2Image
<span class="hljs-keyword">from</span> diffusers.pipelines.wuerstchen <span class="hljs-keyword">import</span> DEFAULT_STAGE_C_TIMESTEPS

pipeline = AutoPipelineForText2Image.from_pretrained(<span class="hljs-string">&quot;path/to/saved/model&quot;</span>, dtype=torch.float16).to(<span class="hljs-string">&quot;cuda&quot;</span>)

caption = <span class="hljs-string">&quot;A cute bird naruto holding a shield&quot;</span>
images = pipeline(
    caption,
    width=<span class="hljs-number">1024</span>,
    height=<span class="hljs-number">1536</span>,
    prior_timesteps=DEFAULT_STAGE_C_TIMESTEPS,
    prior_guidance_scale=<span class="hljs-number">4.0</span>,
    num_images_per_prompt=<span class="hljs-number">2</span>,
).images`,lang:"py",wrap:!1});var W=e(w,2);x(W,{source:"https://github.com/huggingface/diffusers/blob/main/docs/source/en/training/wuerstchen.md"}),v(2),U(j,l),C()}export{S as component};