AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00
Code Activity

Update final CPU offloading code for more diffusion pipelines (#5589)

Browse Source 
* Update final model offload for more pipelines

Add test to ensure all pipeline components are returned to CPU after
execution with model offloading

* Add comment to explain early UNet offload in Text-to-Video pipeline

* Style
This commit is contained in:
clarencechen
2023-11-01 13:22:56 -07:00
committed by GitHub
parent d1eb14bc35
commit 151998e1c2
10 changed files with 25 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint_sd_xl.py
View File

@@ -1604,9 +1604,8 @@ class StableDiffusionXLControlNetInpaintPipeline(
image = self.image_processor.postprocess(image, output_type=output_type)
# Offload last model to CPU
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
self.final_offload_hook.offload()
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image,)

5
src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl_img2img.py
View File

@@ -1433,9 +1433,8 @@ class StableDiffusionXLControlNetImg2ImgPipeline(
image = self.image_processor.postprocess(image, output_type=output_type)
# Offload last model to CPU
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
self.final_offload_hook.offload()
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image,)

5
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_gligen.py
View File

@@ -864,9 +864,8 @@ class StableDiffusionGLIGENPipeline(DiffusionPipeline):
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
# Offload last model to CPU
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
self.final_offload_hook.offload()
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image, has_nsfw_concept)

5
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_gligen_text_image.py
View File

@@ -1031,9 +1031,8 @@ class StableDiffusionGLIGENTextImagePipeline(DiffusionPipeline):
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
# Offload last model to CPU
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
self.final_offload_hook.offload()
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image, has_nsfw_concept)

5
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_upscale.py
View File

@@ -820,9 +820,8 @@ class StableDiffusionUpscalePipeline(
if output_type == "pil" and self.watermarker is not None:
image = self.watermarker.apply_watermark(image)
# Offload last model to CPU
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
self.final_offload_hook.offload()
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image, has_nsfw_concept)

5
src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip.py
View File

@@ -942,9 +942,8 @@ class StableUnCLIPPipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraL
image = self.image_processor.postprocess(image, output_type=output_type)
# Offload last model to CPU
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
self.final_offload_hook.offload()
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image,)

5
src/diffusers/pipelines/stable_diffusion/pipeline_stable_unclip_img2img.py
View File

@@ -839,9 +839,8 @@ class StableUnCLIPImg2ImgPipeline(DiffusionPipeline, TextualInversionLoaderMixin
image = self.image_processor.postprocess(image, output_type=output_type)
# Offload last model to CPU
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
self.final_offload_hook.offload()
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image,)

5
src/diffusers/pipelines/t2i_adapter/pipeline_stable_diffusion_xl_adapter.py
View File

@@ -1059,9 +1059,8 @@ class StableDiffusionXLAdapterPipeline(
image = self.image_processor.postprocess(image, output_type=output_type)
# Offload last model to CPU
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
self.final_offload_hook.offload()
# Offload all models
self.maybe_free_model_hooks()
if not return_dict:
return (image,)

1
src/diffusers/pipelines/text_to_video_synthesis/pipeline_text_to_video_synth_img2img.py
View File

@@ -777,6 +777,7 @@ class VideoToVideoSDPipeline(DiffusionPipeline, TextualInversionLoaderMixin, Lor
if output_type == "latent":
return TextToVideoSDPipelineOutput(frames=latents)
# manually for max memory savings
if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
self.unet.to("cpu")

8
tests/pipelines/test_pipelines_common.py
View File

@@ -742,6 +742,14 @@ class PipelineTesterMixin:
max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max()
self.assertLess(max_diff, expected_max_diff, "CPU offloading should not affect the inference results")
self.assertTrue(
all(
v.device == "cpu"
for k, v in pipe.components.values()
if isinstance(v, torch.nn.Module) and k not in pipe._exclude_from_cpu_offload
),
"CPU offloading should leave all pipeline components on the CPU after inference",
)
@unittest.skipIf(
torch_device != "cuda" or not is_xformers_available(),