Fix cpu offloading

README.md

@@ -33,7 +33,7 @@ pip install accelerate transformers safetensors diffusers
 
 And then setup the zoe-depth model
 
-```
+```python
 import torch
 import matplotlib
 import matplotlib.cm
@@ -106,8 +106,8 @@ controlnet = ControlNetModel.from_pretrained(
   "diffusers/controlnet-zoe-depth-sdxl-1.0",
   use_safetensors=True,
   torch_dtype=torch.float16,
-).to("cuda")
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to("cuda")
+)
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
   "stabilityai/stable-diffusion-xl-base-1.0",
   controlnet=controlnet,
@@ -115,7 +115,7 @@ pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
   variant="fp16",
   use_safetensors=True,
   torch_dtype=torch.float16,
-).to("cuda")
+)
 pipe.enable_model_cpu_offload()