许可证: openrail++ 基础模型: stabilityai/stable-diffusion-xl-base-1.0 标签:

• stable-diffusion-xl
• stable-diffusion-xl-diffusers
• 文本生成图像
• diffusers
• controlnet 推理: 禁用

SDXL-controlnet: 深度控制

这些是基于stabilityai/stable-diffusion-xl-base-1.0训练的controlnet权重，支持深度条件控制。下方展示了一些示例图像。

提示词: 蜘蛛侠授课，照片级真实感

使用说明

首先确保安装以下库：

pip install accelerate transformers safetensors diffusers

然后即可运行以下代码：

import torch
import numpy as np
from PIL import Image

from transformers import DPTFeatureExtractor, 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")
controlnet = ControlNetModel.from_pretrained(
    "diffusers/controlnet-depth-sdxl-1.0",
    variant="fp16",
    use_safetensors=True,
    torch_dtype=torch.float16,
)
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    controlnet=controlnet,
    vae=vae,
    variant="fp16",
    use_safetensors=True,
    torch_dtype=torch.float16,
)
pipe.enable_model_cpu_offload()

def get_depth_map(image):
    image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
    with torch.no_grad(), torch.autocast("cuda"):
        depth_map = depth_estimator(image).predicted_depth

    depth_map = torch.nn.functional.interpolate(
        depth_map.unsqueeze(1),
        size=(1024, 1024),
        mode="bicubic",
        align_corners=False,
    )
    depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
    depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
    depth_map = (depth_map - depth_min) / (depth_max - depth_min)
    image = torch.cat([depth_map] * 3, dim=1)

    image = image.permute(0, 2, 3, 1).cpu().numpy()[0]
    image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
    return image


prompt = "冲锋队授课，照片级真实感"
image = load_image("https://huggingface.co/lllyasviel/sd-controlnet-depth/resolve/main/images/stormtrooper.png")
controlnet_conditioning_scale = 0.5  # 推荐值以获得良好泛化性

depth_image = get_depth_map(image)

images = pipe(
    prompt, image=depth_image, num_inference_steps=30, controlnet_conditioning_scale=controlnet_conditioning_scale,
).images
images[0]

images[0].save(f"stormtrooper.png")

更多细节请查阅StableDiffusionXLControlNetPipeline官方文档。

训练过程

我们的训练脚本基于官方提供的训练代码构建。

训练数据与算力

模型使用LAION-Aesthetics V2数据集的300万图文对进行训练，在80GB A100 GPU上累计训练700小时。

批处理大小

采用数据并行策略，单GPU批大小为8，总批大小达到256。

超参数设置

恒定学习率设为1e-5。

混合精度

使用fp16混合精度训练