from_pt

handler.py

@@ -49,11 +49,10 @@ class EndpointHandler():
 
         pretrained_base_model_path_unet = os.path.join(base_dir, 'pretrained_weights', 'stable-diffusion-v1-5', 'unet')
 
-        print("model path is " + pretrained_base_model_path_unet)
-            reference_unet = UNet2DConditionModel.from_pretrained(
-                self.config.pretrained_base_model_path,
-                subfolder="unet",
-            ).to(dtype=self.weight_dtype, device="cuda")
+        reference_unet = UNet2DConditionModel.from_pretrained(
+            pretrained_base_model_path_unet,
+            from_pt=True
+        ).to(device, dtype=self.weight_dtype)
 
         inference_config_path = os.path.join(base_dir, 'configs', 'inference', 'inference_v2.yaml')
         motion_module_path = os.path.join(base_dir, 'pretrained_weights', 'motion_module.pth')