Skip to content

[Stable Diffusion] Allow users to disable Safety checker if loading model from checkpoint #2768

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
patrickvonplaten merged 7 commits into from
Mar 28, 2023
Merged

[Stable Diffusion] Allow users to disable Safety checker if loading model from checkpoint #2768

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
01d4fef
Allow user to disable SafetyChecker and enable dtypes if loading mode…
Stax124 Mar 21, 2023
5d41b53
Fix Import sorting (Ruff error)
Stax124 Mar 21, 2023
280cb2a
Get rid of the dtype convert method as it was implemented all along
Stax124 Mar 23, 2023
fcf1963
Fix the docstring
Stax124 Mar 23, 2023
75fd2d6
Fix ruff formatting
Stax124 Mar 23, 2023
9c2c364
Merge branch 'main' into main
Stax124 Mar 25, 2023
69ebe82
Merge branch 'main' into main
patrickvonplaten Mar 28, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -989,6 +989,7 @@ def download_from_original_stable_diffusion_ckpt(
stable_unclip_prior: Optional[str] = None,
clip_stats_path: Optional[str] = None,
controlnet: Optional[bool] = None,
load_safety_checker: bool = True,
) -> StableDiffusionPipeline:
"""
Load a Stable Diffusion pipeline object from a CompVis-style `.ckpt`/`.safetensors` file and (ideally) a `.yaml`
Expand Down Expand Up @@ -1028,6 +1029,8 @@ def download_from_original_stable_diffusion_ckpt(
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.
load_safety_checker (`bool`, *optional*, defaults to `True`):
Whether to load the safety checker or not. Defaults to `True`.
"""
if prediction_type == "v-prediction":
prediction_type = "v_prediction"
Expand Down Expand Up @@ -1270,8 +1273,13 @@ def download_from_original_stable_diffusion_ckpt(
elif model_type == "FrozenCLIPEmbedder":
text_model = convert_ldm_clip_checkpoint(checkpoint)
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
feature_extractor = AutoFeatureExtractor.from_pretrained("CompVis/stable-diffusion-safety-checker")

if load_safety_checker:
safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
feature_extractor = AutoFeatureExtractor.from_pretrained("CompVis/stable-diffusion-safety-checker")
else:
safety_checker = None
feature_extractor = None

if controlnet:
pipe = StableDiffusionControlNetPipeline(
Expand Down