Unverified Commit 1e7f9654 authored by Will Berman's avatar Will Berman Committed by GitHub
Browse files

convert ckpt script docstring fixes (#2293)



* convert ckpt script docstring fixes

* Update src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py
Co-authored-by: default avatarPedro Cuenca <pedro@huggingface.co>

* Update src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py
Co-authored-by: default avatarPedro Cuenca <pedro@huggingface.co>

---------
Co-authored-by: default avatarPedro Cuenca <pedro@huggingface.co>
parent beb59abf
...@@ -17,6 +17,7 @@ ...@@ -17,6 +17,7 @@
import os import os
import re import re
import tempfile import tempfile
from typing import Optional
import requests import requests
import torch import torch
...@@ -787,8 +788,8 @@ def load_pipeline_from_original_stable_diffusion_ckpt( ...@@ -787,8 +788,8 @@ def load_pipeline_from_original_stable_diffusion_ckpt(
model_type: str = None, model_type: str = None,
extract_ema: bool = False, extract_ema: bool = False,
scheduler_type: str = "pndm", scheduler_type: str = "pndm",
num_in_channels: int = None, num_in_channels: Optional[int] = None,
upcast_attention: bool = None, upcast_attention: Optional[bool] = None,
device: str = None, device: str = None,
from_safetensors: bool = False, from_safetensors: bool = False,
) -> StableDiffusionPipeline: ) -> StableDiffusionPipeline:
...@@ -800,28 +801,36 @@ def load_pipeline_from_original_stable_diffusion_ckpt( ...@@ -800,28 +801,36 @@ def load_pipeline_from_original_stable_diffusion_ckpt(
global step count, which will likely fail for models that have undergone further fine-tuning. Therefore, it is global step count, which will likely fail for models that have undergone further fine-tuning. Therefore, it is
recommended that you override the default values and/or supply an `original_config_file` wherever possible. recommended that you override the default values and/or supply an `original_config_file` wherever possible.
:param checkpoint_path: Path to `.ckpt` file. :param original_config_file: Path to `.yaml` config file Args:
corresponding to the original architecture. If `None`, will be checkpoint_path (`str`): Path to `.ckpt` file.
automatically inferred by looking for a key that only exists in SD2.0 models. original_config_file (`str`):
:param image_size: The image size that the model was trained on. Use 512 for Stable Diffusion v1.X and Stable Path to `.yaml` config file corresponding to the original architecture. If `None`, will be automatically
Diffusion v2 inferred by looking for a key that only exists in SD2.0 models.
image_size (`int`, *optional*, defaults to 512):
The image size that the model was trained on. Use 512 for Stable Diffusion v1.X and Stable Diffusion v2
Base. Use 768 for Stable Diffusion v2. Base. Use 768 for Stable Diffusion v2.
:param prediction_type: The prediction type that the model was trained on. Use `'epsilon'` for Stable Diffusion prediction_type (`str`, *optional*):
v1.X and Stable The prediction type that the model was trained on. Use `'epsilon'` for Stable Diffusion v1.X and Stable
Diffusion v2 Base. Use `'v_prediction'` for Stable Diffusion v2. Diffusion v2 Base. Use `'v_prediction'` for Stable Diffusion v2.
:param num_in_channels: The number of input channels. If `None` number of input channels will be automatically num_in_channels (`int`, *optional*, defaults to None):
inferred. :param scheduler_type: Type of scheduler to use. Should be one of `["pndm", "lms", "heun", "euler", The number of input channels. If `None`, it will be automatically inferred.
"euler-ancestral", "dpm", "ddim"]`. :param model_type: The pipeline type. `None` to automatically infer, or one of scheduler_type (`str`, *optional*, defaults to 'pndm'):
`["FrozenOpenCLIPEmbedder", "FrozenCLIPEmbedder", "PaintByExample"]`. :param extract_ema: Only relevant for Type of scheduler to use. Should be one of `["pndm", "lms", "heun", "euler", "euler-ancestral", "dpm",
checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights "ddim"]`.
or not. Defaults to `False`. Pass `True` to extract the EMA weights. EMA weights usually yield higher model_type (`str`, *optional*, defaults to `None`):
quality images for inference. Non-EMA weights are usually better to continue fine-tuning. The pipeline type. `None` to automatically infer, or one of `["FrozenOpenCLIPEmbedder",
:param upcast_attention: Whether the attention computation should always be upcasted. This is necessary when "FrozenCLIPEmbedder", "PaintByExample"]`.
running extract_ema (`bool`, *optional*, defaults to `False`): Only relevant for
stable diffusion 2.1. checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights or not. Defaults to
:param device: The device to use. Pass `None` to determine automatically. :param from_safetensors: If `False`. Pass `True` to extract the EMA weights. EMA weights usually yield higher quality images for
`checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch. :return: A inference. Non-EMA weights are usually better to continue fine-tuning.
StableDiffusionPipeline object representing the passed-in `.ckpt`/`.safetensors` file. upcast_attention (`bool`, *optional*, defaults to `None`):
Whether the attention computation should always be upcasted. This is necessary when running stable
diffusion 2.1.
device (`str`, *optional*, defaults to `None`):
The device to use. Pass `None` to determine automatically. :param from_safetensors: If `checkpoint_path` is
in `safetensors` format, load checkpoint with safetensors instead of PyTorch. :return: A
StableDiffusionPipeline object representing the passed-in `.ckpt`/`.safetensors` file.
""" """
if prediction_type == "v-prediction": if prediction_type == "v-prediction":
prediction_type = "v_prediction" prediction_type = "v_prediction"
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment