Update free model hooks (huggingface#5680)

update free model hooks

Author: DN6

src/diffusers/pipelines/deepfloyd_if/pipeline_if_inpainting_superresolution.py

@@ -1109,8 +1109,6 @@ def __call__(
             nsfw_detected = None
             watermark_detected = None
 
-            if hasattr(self, "unet_offload_hook") and self.unet_offload_hook is not None:
-                self.unet_offload_hook.offload()
         else:
             # 10. Post-processing
             image = (image / 2 + 0.5).clamp(0, 1)
@@ -1119,9 +1117,7 @@ def __call__(
             # 11. Run safety checker
             image, nsfw_detected, watermark_detected = self.run_safety_checker(image, device, prompt_embeds.dtype)
 
-        # Offload last model to CPU
-        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
-            self.final_offload_hook.offload()
+        self.maybe_free_model_hooks()
 
         if not return_dict:
             return (image, nsfw_detected, watermark_detected)

src/diffusers/pipelines/kandinsky/pipeline_kandinsky.py

@@ -388,6 +388,8 @@ def __call__(
         # post-processing
         image = self.movq.decode(latents, force_not_quantize=True)["sample"]
 
+        self.maybe_free_model_hooks()
+
         if output_type not in ["pt", "np", "pil"]:
             raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}")
 

src/diffusers/pipelines/kandinsky/pipeline_kandinsky_combined.py

@@ -321,6 +321,9 @@ def __call__(
             callback_steps=callback_steps,
             return_dict=return_dict,
         )
+
+        self.maybe_free_model_hooks()
+
         return outputs
 
 
@@ -558,6 +561,9 @@ def __call__(
             callback_steps=callback_steps,
             return_dict=return_dict,
         )
+
+        self.maybe_free_model_hooks()
+
         return outputs
 
 
@@ -593,7 +599,7 @@ class KandinskyInpaintCombinedPipeline(DiffusionPipeline):
     """
 
     _load_connected_pipes = True
-    model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->prior_prior->" "text_encoder->unet->movq"
+    model_cpu_offload_seq = "prior_text_encoder->prior_image_encoder->prior_prior->text_encoder->unet->movq"
 
     def __init__(
         self,
@@ -802,4 +808,7 @@ def __call__(
             callback_steps=callback_steps,
             return_dict=return_dict,
         )
+
+        self.maybe_free_model_hooks()
+
         return outputs

src/diffusers/pipelines/kandinsky/pipeline_kandinsky_img2img.py

@@ -481,6 +481,8 @@ def __call__(
         # 7. post-processing
         image = self.movq.decode(latents, force_not_quantize=True)["sample"]
 
+        self.maybe_free_model_hooks()
+
         if output_type not in ["pt", "np", "pil"]:
             raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}")
 

src/diffusers/pipelines/kandinsky/pipeline_kandinsky_inpaint.py

@@ -616,6 +616,8 @@ def __call__(
         # post-processing
         image = self.movq.decode(latents, force_not_quantize=True)["sample"]
 
+        self.maybe_free_model_hooks()
+
         if output_type not in ["pt", "np", "pil"]:
             raise ValueError(f"Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}")
 

src/diffusers/pipelines/kandinsky/pipeline_kandinsky_prior.py

@@ -527,7 +527,7 @@ def __call__(
         if negative_prompt is None:
             zero_embeds = self.get_zero_embed(latents.shape[0], device=latents.device)
 
-            self.maybe_free_model_hooks
+            self.maybe_free_model_hooks()
         else:
             image_embeddings, zero_embeds = image_embeddings.chunk(2)
 

src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py

@@ -326,6 +326,8 @@ def __call__(
             callback_on_step_end=callback_on_step_end,
             callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
         )
+        self.maybe_free_model_hooks()
+
         return outputs
 
 
@@ -572,6 +574,8 @@ def __call__(
             callback_on_step_end=callback_on_step_end,
             callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
         )
+
+        self.maybe_free_model_hooks()
         return outputs
 
 
@@ -842,4 +846,6 @@ def __call__(
             callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
             **kwargs,
         )
+        self.maybe_free_model_hooks()
+
         return outputs

src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py

@@ -531,14 +531,10 @@ def __call__(
         # if negative prompt has been defined, we retrieve split the image embedding into two
         if negative_prompt is None:
             zero_embeds = self.get_zero_embed(latents.shape[0], device=latents.device)
-
-            if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
-                self.final_offload_hook.offload()
         else:
             image_embeddings, zero_embeds = image_embeddings.chunk(2)
 
-            if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
-                self.prior_hook.offload()
+        self.maybe_free_model_hooks()
 
         if output_type not in ["pt", "np"]:
             raise ValueError(f"Only the output types `pt` and `np` are supported not output_type={output_type}")

src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior_emb2emb.py

@@ -545,12 +545,10 @@ def __call__(
         # if negative prompt has been defined, we retrieve split the image embedding into two
         if negative_prompt is None:
             zero_embeds = self.get_zero_embed(latents.shape[0], device=latents.device)
-            if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
-                self.final_offload_hook.offload()
         else:
             image_embeddings, zero_embeds = image_embeddings.chunk(2)
-            if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
-                self.prior_hook.offload()
+
+        self.maybe_free_model_hooks()
 
         if output_type not in ["pt", "np"]:
             raise ValueError(f"Only the output types `pt` and `np` are supported not output_type={output_type}")

src/diffusers/pipelines/stable_diffusion/pipeline_cycle_diffusion.py

@@ -918,6 +918,7 @@ def __call__(
             do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
 
         image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
+        self.maybe_free_model_hooks()
 
         if not return_dict:
             return (image, has_nsfw_concept)