add offload

gradio_app/custom_models/mvimg_prediction.py

@@ -14,7 +14,7 @@ checkpoint_path = "ckpt/img2mvimg/unet_state_dict.pth"
 
 def predict(img_list: List[Image.Image], guidance_scale=2., **kwargs):
     trainer, pipeline = load_pipeline(training_config, checkpoint_path)
-    # pipeline.enable_model_cpu_offload()
+    pipeline.enable_model_cpu_offload()
     
     if isinstance(img_list, Image.Image):
         img_list = [img_list]

gradio_app/custom_models/normal_prediction.py

@@ -10,7 +10,7 @@ checkpoint_path = "ckpt/image2normal/unet_state_dict.pth"
 
 def predict_normals(image: List[Image.Image], guidance_scale=2., do_rotate=True, num_inference_steps=30, **kwargs):
     trainer, pipeline = load_pipeline(training_config, checkpoint_path)
-    # pipeline.enable_model_cpu_offload()
+    pipeline.enable_model_cpu_offload()
     
     img_list = image if isinstance(image, list) else [image]
     img_list = [rgba_to_rgb(i) if i.mode == 'RGBA' else i for i in img_list]

scripts/sd_model_zoo.py

@@ -119,9 +119,8 @@ def load_common_sd15_pipe(base_model=DEFAULT_BASE_MODEL, device="balanced", cont
     if enable_sequential_cpu_offload:
         pipe.enable_sequential_cpu_offload()
     else:
-        # pipe = pipe.to("cuda")
         pass
-        # pipe.enable_model_cpu_offload()
+        pipe.enable_model_cpu_offload()
     if vae_slicing:
         pipe.enable_vae_slicing()