fix

handler.py

@@ -47,8 +47,15 @@ class EndpointHandler():
 
         vae = AutoencoderKL.from_pretrained(config_path).to(device, dtype=self.weight_dtype)
 
+
+        pretrained_base_model_path = os.path.join(base_dir, 'pretrained_weights', 'stable-diffusion-v1-5')
+
+        # Ensure the path exists
+        if not os.path.exists(pretrained_base_model_path):
+            raise FileNotFoundError(f"The folder was not found at: {pretrained_base_model_path}")
+
         reference_unet = UNet2DConditionModel.from_pretrained(
-            self.config.pretrained_base_model_path,
+            pretrained_base_model_path,
             subfolder="unet"
         ).to(device, dtype=self.weight_dtype)