Views
No views yet




pip install accelerate transformers safetensors diffusers1import torch
2import numpy as np
3from PIL import Image
4
5from transformers import DPTFeatureExtractor, DPTForDepthEstimation
6from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
7from diffusers.utils import load_image
8
9
10depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to("cuda")
11feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-hybrid-midas")
12controlnet = ControlNetModel.from_pretrained(
13 "diffusers/controlnet-depth-sdxl-1.0-small",
14 variant="fp16",
15 use_safetensors=True,
16 torch_dtype=torch.float16,
17).to("cuda")
18vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to("cuda")
19pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
20 "stabilityai/stable-diffusion-xl-base-1.0",
21 controlnet=controlnet,
22 vae=vae,
23 variant="fp16",
24 use_safetensors=True,
25 torch_dtype=torch.float16,
26).to("cuda")
27pipe.enable_model_cpu_offload()
28
29def get_depth_map(image):
30 image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
31 with torch.no_grad(), torch.autocast("cuda"):
32 depth_map = depth_estimator(image).predicted_depth
33
34 depth_map = torch.nn.functional.interpolate(
35 depth_map.unsqueeze(1),
36 size=(1024, 1024),
37 mode="bicubic",
38 align_corners=False,
39 )
40 depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
41 depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
42 depth_map = (depth_map - depth_min) / (depth_max - depth_min)
43 image = torch.cat([depth_map] * 3, dim=1)
44
45 image = image.permute(0, 2, 3, 1).cpu().numpy()[0]
46 image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
47 return image
48
49
50prompt = "stormtrooper lecture, photorealistic"
51image = load_image("https://huggingface.co/lllyasviel/sd-controlnet-depth/resolve/main/images/stormtrooper.png")
52controlnet_conditioning_scale = 0.5 # recommended for good generalization
53
54depth_image = get_depth_map(image)
55
56images = pipe(
57 prompt, image=depth_image, num_inference_steps=30, controlnet_conditioning_scale=controlnet_conditioning_scale,
58).images
59images[0]
60
61images[0].save(f"stormtrooper_grid.png")
StableDiffusionXLControlNetPipeline.controlnet_conditioning_scale and guidance_scale arguments for potentially better
image generation quality.