Use ImageProcessor

README.md

@@ -33,13 +33,13 @@ import torch
 import numpy as np
 from PIL import Image
 
-from transformers import DPTFeatureExtractor, DPTForDepthEstimation
+from transformers import DPTImageProcessor, DPTForDepthEstimation
 from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
 from diffusers.utils import load_image
 
 
 depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda")
-feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-hybrid-midas")
+processor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas")
 controlnet = ControlNetModel.from_pretrained(
     "diffusers/controlnet-depth-sdxl-1.0",
     variant="fp16",
@@ -58,7 +58,7 @@ pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
 pipe.enable_model_cpu_offload()
 
 def get_depth_map(image):
-    image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
+    image = processor(images=image, return_tensors="pt").pixel_values.to("cuda")
     with torch.no_grad(), torch.autocast("cuda"):
         depth_map = depth_estimator(image).predicted_depth