gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/images/gpu/stable-diffusion-xl/download_checkpoints.py (about)

     1  #!/usr/bin/env python3
     2  
     3  # Copyright 2024 The gVisor Authors.
     4  #
     5  # Licensed under the Apache License, Version 2.0 (the "License");
     6  # you may not use this file except in compliance with the License.
     7  # You may obtain a copy of the License at
     8  #
     9  #     http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  
    17  """Download Stable Diffusion XL checkpoints from Hugging Face."""
    18  
    19  import diffusers
    20  import torch
    21  
    22  # Download base model.
    23  base = diffusers.DiffusionPipeline.from_pretrained(
    24      "stabilityai/stable-diffusion-xl-base-1.0",
    25      torch_dtype=torch.float16,
    26      variant="fp16",
    27      use_safetensors=True,
    28  )
    29  
    30  # Download refiner model.
    31  refiner = diffusers.DiffusionPipeline.from_pretrained(
    32      "stabilityai/stable-diffusion-xl-refiner-1.0",
    33      text_encoder_2=base.text_encoder_2,
    34      vae=base.vae,
    35      torch_dtype=torch.float16,
    36      use_safetensors=True,
    37      variant="fp16",
    38  )