_checker": pipe.safety_checker,
...     "watermarker": pipe.watermarker,
... }
>>> super_res_2_pipe = DiffusionPipeline.from_pretrained(
...     "stabilityai/stable-diffusion-x4-upscaler", **safety_modules, torch_dtype=torch.float16
... )
>>> super_res_2_pipe.enable_model_cpu_offload()

>>> image = super_res_2_pipe(
...     prompt=prompt,
...     image=image,
... ).images
>>> image[0].save("./if_stage_II.png")

enable_model_cpu_offload

<
sourc