wehub-resource-sync
162 lines
8.4 KiB
Python
162 lines
8.4 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
# Configure dependencies before any external imports
|
|
from demo_diffusion import deps
|
|
deps.configure("sd")
|
|
|
|
import argparse
|
|
|
|
import controlnet_aux
|
|
import torch
|
|
from cuda.bindings import runtime as cudart
|
|
from PIL import Image
|
|
|
|
from demo_diffusion import dd_argparse
|
|
from demo_diffusion import image as image_module
|
|
from demo_diffusion import pipeline as pipeline_module
|
|
|
|
|
|
def parseArgs():
|
|
parser = argparse.ArgumentParser(description="Options for Stable Diffusion ControlNet Demo", conflict_handler='resolve')
|
|
parser = dd_argparse.add_arguments(parser)
|
|
parser.add_argument('--scheduler', type=str, default="UniPC", choices=["DDIM", "DPM", "EulerA", "LMSD", "PNDM", "UniPC"], help="Scheduler for diffusion process")
|
|
parser.add_argument('--input-image', nargs = '+', type=str, default=[], help="Path to the input image/images already prepared for ControlNet modality. For example: canny edged image for canny ControlNet, not just regular rgb image")
|
|
parser.add_argument('--controlnet-type', nargs='+', type=str, default=["canny"], help="Controlnet type, can be `None`, `str` or `str` list from ['canny', 'depth', 'hed', 'mlsd', 'normal', 'openpose', 'scribble', 'seg']")
|
|
parser.add_argument('--controlnet-scale', nargs='+', type=float, default=[1.0], help="The outputs of the controlnet are multiplied by `controlnet_scale` before they are added to the residual in the original unet, can be `None`, `float` or `float` list")
|
|
return parser.parse_args()
|
|
|
|
if __name__ == "__main__":
|
|
print("[I] Initializing StableDiffusion controlnet demo using TensorRT")
|
|
args = parseArgs()
|
|
|
|
# Controlnet configuration
|
|
if not isinstance(args.controlnet_type, list):
|
|
raise ValueError(f"`--controlnet-type` must be of type `str` or `str` list, but is {type(args.controlnet_type)}")
|
|
|
|
# Controlnet configuration
|
|
if not isinstance(args.controlnet_scale, list):
|
|
raise ValueError(f"`--controlnet-scale`` must be of type `float` or `float` list, but is {type(args.controlnet_scale)}")
|
|
|
|
# Check number of ControlNets to ControlNet scales
|
|
if len(args.controlnet_type) != len(args.controlnet_scale):
|
|
raise ValueError(f"Numbers of ControlNets {len(args.controlnet_type)} should be equal to number of ControlNet scales {len(args.controlnet_scale)}.")
|
|
|
|
# Convert controlnet scales to tensor
|
|
controlnet_scale = torch.FloatTensor(args.controlnet_scale)
|
|
|
|
# Check images
|
|
input_images = []
|
|
if len(args.input_image) > 0:
|
|
for image in args.input_image:
|
|
input_images.append(Image.open(image))
|
|
else:
|
|
for controlnet in args.controlnet_type:
|
|
if controlnet == "canny":
|
|
if args.version == "xl-1.0":
|
|
canny_image = image_module.download_image(
|
|
"https://huggingface.co/diffusers/controlnet-canny-sdxl-1.0/resolve/main/out_bird.png"
|
|
)
|
|
# "out_bird.png" has 5 images combined in a row. We pick the first image which is the input image.
|
|
canny_image = canny_image.crop((0, 0, canny_image.width / 5, canny_image.height))
|
|
elif args.version == "1.5":
|
|
canny_image = image_module.download_image(
|
|
"https://hf.co/datasets/huggingface/documentation-images/resolve/main/diffusers/input_image_vermeer.png"
|
|
)
|
|
canny_image = controlnet_aux.CannyDetector()(canny_image)
|
|
else:
|
|
raise ValueError(
|
|
f"This demo supports ControlNets for v1.4 and SDXL base pipelines only. Version provided: {args.version}"
|
|
)
|
|
input_images.append(canny_image.resize((args.width, args.height)))
|
|
elif controlnet == "normal":
|
|
normal_image = image_module.download_image(
|
|
"https://huggingface.co/lllyasviel/sd-controlnet-normal/resolve/main/images/toy.png"
|
|
)
|
|
normal_image = controlnet_aux.NormalBaeDetector.from_pretrained("lllyasviel/Annotators")(normal_image)
|
|
input_images.append(normal_image.resize((args.width, args.height)))
|
|
elif controlnet == "depth":
|
|
depth_image = image_module.download_image(
|
|
"https://huggingface.co/lllyasviel/sd-controlnet-depth/resolve/main/images/stormtrooper.png"
|
|
)
|
|
depth_image = controlnet_aux.LeresDetector.from_pretrained("lllyasviel/Annotators")(depth_image)
|
|
input_images.append(depth_image.resize((args.width, args.height)))
|
|
elif controlnet == "hed":
|
|
hed_image = image_module.download_image(
|
|
"https://huggingface.co/lllyasviel/sd-controlnet-hed/resolve/main/images/man.png"
|
|
)
|
|
hed_image = controlnet_aux.HEDdetector.from_pretrained("lllyasviel/Annotators")(hed_image)
|
|
input_images.append(hed_image.resize((args.width, args.height)))
|
|
elif controlnet == "mlsd":
|
|
mlsd_image = image_module.download_image(
|
|
"https://huggingface.co/lllyasviel/sd-controlnet-mlsd/resolve/main/images/room.png"
|
|
)
|
|
mlsd_image = controlnet_aux.MLSDdetector.from_pretrained("lllyasviel/Annotators")(mlsd_image)
|
|
input_images.append(mlsd_image.resize((args.width, args.height)))
|
|
elif controlnet == "openpose":
|
|
openpose_image = image_module.download_image(
|
|
"https://huggingface.co/lllyasviel/sd-controlnet-openpose/resolve/main/images/pose.png"
|
|
)
|
|
openpose_image = controlnet_aux.OpenposeDetector.from_pretrained("lllyasviel/Annotators")(openpose_image)
|
|
input_images.append(openpose_image.resize((args.width, args.height)))
|
|
elif controlnet == "scribble":
|
|
scribble_image = image_module.download_image(
|
|
"https://huggingface.co/lllyasviel/sd-controlnet-scribble/resolve/main/images/bag.png"
|
|
)
|
|
scribble_image = controlnet_aux.HEDdetector.from_pretrained("lllyasviel/Annotators")(scribble_image, scribble=True)
|
|
input_images.append(scribble_image.resize((args.width, args.height)))
|
|
elif controlnet == "seg":
|
|
seg_image = image_module.download_image(
|
|
"https://huggingface.co/lllyasviel/sd-controlnet-seg/resolve/main/images/house.png"
|
|
)
|
|
seg_image = controlnet_aux.SamDetector.from_pretrained("ybelkada/segment-anything", subfolder="checkpoints")(seg_image)
|
|
input_images.append(seg_image.resize((args.width, args.height)))
|
|
else:
|
|
raise ValueError(f"You should implement the conditonal image of this controlnet: {controlnet}")
|
|
assert len(input_images) > 0
|
|
|
|
kwargs_init_pipeline, kwargs_load_engine, args_run_demo = dd_argparse.process_pipeline_args(args)
|
|
|
|
# Initialize demo
|
|
demo = pipeline_module.StableDiffusionPipeline(
|
|
pipeline_type=(
|
|
pipeline_module.PIPELINE_TYPE.CONTROLNET
|
|
if args.version != "xl-1.0"
|
|
else pipeline_module.PIPELINE_TYPE.XL_CONTROLNET
|
|
),
|
|
controlnets=args.controlnet_type,
|
|
**kwargs_init_pipeline,
|
|
)
|
|
|
|
# Load TensorRT engines and pytorch modules
|
|
demo.loadEngines(
|
|
args.engine_dir,
|
|
args.framework_model_dir,
|
|
args.onnx_dir,
|
|
**kwargs_load_engine)
|
|
|
|
# Load resources
|
|
_, shared_device_memory = cudart.cudaMalloc(demo.calculateMaxDeviceMemory())
|
|
demo.activateEngines(shared_device_memory)
|
|
demo.loadResources(args.height, args.width, args.batch_size, args.seed)
|
|
|
|
# Run inference
|
|
demo_kwargs = {'input_image': input_images, 'controlnet_scales': controlnet_scale}
|
|
demo.run(*args_run_demo, **demo_kwargs)
|
|
|
|
demo.teardown()
|