getsep/instantid-dependencies

1import torch2import numpy as np3from PIL import Image4from controlnet_aux import OpenposeDetector5from src.dependencies.instantid.model_util import get_torch_device6import cv27 8 9from transformers import DPTImageProcessor, DPTForDepthEstimation10 11device = get_torch_device()12depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to(device)13feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas")14openpose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")15 16def get_depth_map(image):17    image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")18    with torch.no_grad(), torch.autocast("cuda"):19        depth_map = depth_estimator(image).predicted_depth20 21    depth_map = torch.nn.functional.interpolate(22        depth_map.unsqueeze(1),23        size=(1024, 1024),24        mode="bicubic",25        align_corners=False,26    )27    depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)28    depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)29    depth_map = (depth_map - depth_min) / (depth_max - depth_min)30    image = torch.cat([depth_map] * 3, dim=1)31 32    image = image.permute(0, 2, 3, 1).cpu().numpy()[0]33    image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))34    return image35 36def get_canny_image(image, t1=100, t2=200):37    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)38    edges = cv2.Canny(image, t1, t2)39    return Image.fromarray(edges, "L")