bug fix
Browse files
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 |
-
|
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 |
-
|
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]
|