root commited on
Commit
22af2ed
·
1 Parent(s): 2a94fc4

trying nanother

Browse files
Files changed (1) hide show
  1. handler.py +4 -3
handler.py CHANGED
@@ -50,9 +50,10 @@ class EndpointHandler():
50
  pretrained_base_model_path_unet = os.path.join(base_dir, 'pretrained_weights', 'stable-diffusion-v1-5', 'unet')
51
 
52
  print("model path is " + pretrained_base_model_path_unet)
53
- reference_unet = UNet2DConditionModel.from_pretrained(
54
- pretrained_base_model_path_unet,
55
- ).to(dtype=self.weight_dtype, device="cuda")
 
56
 
57
  inference_config_path = os.path.join(base_dir, 'configs', 'inference', 'inference_v2.yaml')
58
  motion_module_path = os.path.join(base_dir, 'pretrained_weights', 'motion_module.pth')
 
50
  pretrained_base_model_path_unet = os.path.join(base_dir, 'pretrained_weights', 'stable-diffusion-v1-5', 'unet')
51
 
52
  print("model path is " + pretrained_base_model_path_unet)
53
+ reference_unet = UNet2DConditionModel.from_pretrained(
54
+ self.config.pretrained_base_model_path,
55
+ subfolder="unet",
56
+ ).to(dtype=self.weight_dtype, device="cuda")
57
 
58
  inference_config_path = os.path.join(base_dir, 'configs', 'inference', 'inference_v2.yaml')
59
  motion_module_path = os.path.join(base_dir, 'pretrained_weights', 'motion_module.pth')