updating paths

handler.py

@@ -47,35 +47,36 @@ 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(
             pretrained_base_model_path,
             subfolder="unet"
         ).to(device, dtype=self.weight_dtype)
 
-        inference_config_path = self.config.inference_config
+        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')
+        denoising_unet_path = os.path.join(base_dir, 'pretrained_weights', 'denoising_unet.pth')
+        reference_unet_path = os.path.join(base_dir, 'pretrained_weights', 'reference_unet.pth')
+        pose_guider_path = os.path.join(base_dir, 'pretrained_weights', 'pose_guider.pth')
+        image_encoder_path = os.path.join(base_dir, 'pretrained_weights', 'image_encoder')
+
         infer_config = OmegaConf.load(inference_config_path)
         denoising_unet = UNet3DConditionModel.from_pretrained_2d(
-            self.config.pretrained_base_model_path,
-            self.config.motion_module_path,
+            pretrained_base_model_path,
+            motion_module_path,
             subfolder="unet",
             unet_additional_kwargs=infer_config.unet_additional_kwargs,
         ).to(device, dtype=self.weight_dtype)
 
         pose_guider = PoseGuider(320, block_out_channels=(16, 32, 96, 256)).to(device, dtype=self.weight_dtype)
-        image_enc = CLIPVisionModelWithProjection.from_pretrained(self.config.image_encoder_path).to(device, dtype=self.weight_dtype)
+        image_enc = CLIPVisionModelWithProjection.from_pretrained(image_encoder_path).to(device, dtype=self.weight_dtype)
         sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)
         scheduler = DDIMScheduler(**sched_kwargs)
 
-        denoising_unet.load_state_dict(torch.load(self.config.denoising_unet_path, map_location="cpu"), strict=False)
-        reference_unet.load_state_dict(torch.load(self.config.reference_unet_path, map_location="cpu"))
-        pose_guider.load_state_dict(torch.load(self.config.pose_guider_path, map_location="cpu"))
+        denoising_unet.load_state_dict(torch.load(denoising_unet_path, map_location="cpu"), strict=False)
+        reference_unet.load_state_dict(torch.load(reference_unet_path, map_location="cpu"))
+        pose_guider.load_state_dict(torch.load(pose_guider_path, map_location="cpu"))
 
         self.pipeline = Pose2VideoPipeline(
             vae=vae,