x

gradio_app/custom_models/normal_prediction.py

@@ -7,9 +7,9 @@ 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)
 
-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)
+def predict_normals(image: List[Image.Image], guidance_scale=2., do_rotate=True, num_inference_steps=30, **kwargs):    
     pipeline.enable_model_cpu_offload()
     
     img_list = image if isinstance(image, list) else [image]