Skip to content

GitLab

Unverified Commit 0e0db625 authored by Patrick von Platen's avatar Patrick von Platen Committed by GitHub
Browse files

Fix safety checker seq offload (#4998) 

* fix safety checker

* fix safety checker

* fix safety checker
parent 1f948109
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_pix2pix_zero.py
View file @ 0e0db625
...@@ -315,6 +315,7 @@ class StableDiffusionPix2PixZeroPipeline(DiffusionPipeline): ...@@ -315,6 +315,7 @@ class StableDiffusionPix2PixZeroPipeline(DiffusionPipeline):
"caption_processor", "caption_processor",
"inverse_scheduler", "inverse_scheduler",
] ]
_exclude_from_cpu_offload = ["safety_checker"]
def __init__( def __init__(
self, self,
......
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_sag.py
View file @ 0e0db625
...@@ -119,6 +119,7 @@ class StableDiffusionSAGPipeline(DiffusionPipeline, TextualInversionLoaderMixin) ...@@ -119,6 +119,7 @@ class StableDiffusionSAGPipeline(DiffusionPipeline, TextualInversionLoaderMixin)
""" """
model_cpu_offload_seq = "text_encoder->unet->vae" model_cpu_offload_seq = "text_encoder->unet->vae"
_optional_components = ["safety_checker", "feature_extractor"] _optional_components = ["safety_checker", "feature_extractor"]
_exclude_from_cpu_offload = ["safety_checker"]
def __init__( def __init__(
self, self,
......
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py
View file @ 0e0db625
...@@ -92,6 +92,7 @@ class StableDiffusionUpscalePipeline(DiffusionPipeline, TextualInversionLoaderMi ...@@ -92,6 +92,7 @@ class StableDiffusionUpscalePipeline(DiffusionPipeline, TextualInversionLoaderMi
""" """
model_cpu_offload_seq = "text_encoder->unet->vae" model_cpu_offload_seq = "text_encoder->unet->vae"
_optional_components = ["watermarker", "safety_checker", "feature_extractor"] _optional_components = ["watermarker", "safety_checker", "feature_extractor"]
_exclude_from_cpu_offload = ["safety_checker"]
def __init__( def __init__(
self, self,
......
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