Wuvin commited on
Commit
0977aa6
1 Parent(s): 8d1277c
gradio_app/custom_models/mvimg_prediction.py CHANGED
@@ -14,7 +14,7 @@ checkpoint_path = "ckpt/img2mvimg/unet_state_dict.pth"
14
 
15
  def predict(img_list: List[Image.Image], guidance_scale=2., **kwargs):
16
  trainer, pipeline = load_pipeline(training_config, checkpoint_path)
17
- # pipeline.enable_model_cpu_offload()
18
 
19
  if isinstance(img_list, Image.Image):
20
  img_list = [img_list]
 
14
 
15
  def predict(img_list: List[Image.Image], guidance_scale=2., **kwargs):
16
  trainer, pipeline = load_pipeline(training_config, checkpoint_path)
17
+ pipeline.enable_model_cpu_offload()
18
 
19
  if isinstance(img_list, Image.Image):
20
  img_list = [img_list]
gradio_app/custom_models/normal_prediction.py CHANGED
@@ -10,7 +10,7 @@ checkpoint_path = "ckpt/image2normal/unet_state_dict.pth"
10
 
11
  def predict_normals(image: List[Image.Image], guidance_scale=2., do_rotate=True, num_inference_steps=30, **kwargs):
12
  trainer, pipeline = load_pipeline(training_config, checkpoint_path)
13
- # pipeline.enable_model_cpu_offload()
14
 
15
  img_list = image if isinstance(image, list) else [image]
16
  img_list = [rgba_to_rgb(i) if i.mode == 'RGBA' else i for i in img_list]
 
10
 
11
  def predict_normals(image: List[Image.Image], guidance_scale=2., do_rotate=True, num_inference_steps=30, **kwargs):
12
  trainer, pipeline = load_pipeline(training_config, checkpoint_path)
13
+ pipeline.enable_model_cpu_offload()
14
 
15
  img_list = image if isinstance(image, list) else [image]
16
  img_list = [rgba_to_rgb(i) if i.mode == 'RGBA' else i for i in img_list]