bug fix

gradio_app/custom_models/mvimg_prediction.py

@@ -11,10 +11,11 @@ from scripts.utils import session, simple_preprocess
 
 training_config = "gradio_app/custom_models/image2mvimage.yaml"
 checkpoint_path = "ckpt/img2mvimg/unet_state_dict.pth"
-trainer, pipeline = load_pipeline(training_config, checkpoint_path)
-pipeline.enable_model_cpu_offload()
 
 def predict(img_list: List[Image.Image], guidance_scale=2., **kwargs):
+    trainer, pipeline = load_pipeline(training_config, checkpoint_path)
+    # pipeline.enable_model_cpu_offload()
+    
     if isinstance(img_list, Image.Image):
         img_list = [img_list]
     img_list = [rgba_to_rgb(i) if i.mode == 'RGBA' else i for i in img_list]

gradio_app/custom_models/normal_prediction.py

@@ -7,10 +7,11 @@ from scripts.all_typing import *
 
 training_config = "gradio_app/custom_models/image2normal.yaml"
 checkpoint_path = "ckpt/image2normal/unet_state_dict.pth"
-trainer, pipeline = load_pipeline(training_config, checkpoint_path)
-pipeline.enable_model_cpu_offload()
 
 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()
+    
     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]
     images = trainer.pipeline_forward(