| from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler |
| import torch |
| import os |
| from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler |
| from diffusers.utils import load_image |
| import argparse |
| from PIL import Image |
| import cv2 |
| import numpy as np |
| import torch |
| import os |
| import shutil |
| from tqdm import tqdm |
| import numpy as np |
|
|
| device = ("cuda") |
|
|
| |
| parser = argparse.ArgumentParser(description='Choose a processor to run.') |
| parser.add_argument('--op_image', type=str, help='path to pose image') |
| parser.add_argument('--dp_image', type=str, help='path to depth image') |
| parser.add_argument('--output_dir', type=str, default='/content/multi', help='The directory to save the output.') |
| |
| args = parser.parse_args() |
|
|
| op_image = load_image(args.op_image) |
| dp_image = load_image(args.dp_image) |
|
|
| controlnet = [ |
| ControlNetModel.from_pretrained("/content/checkpoints/openpose", torch_dtype=torch.float16).to('cuda'), |
| ControlNetModel.from_pretrained("/content/checkpoints/depth", torch_dtype=torch.float16).to('cuda'), |
| ] |
|
|
| pipe = StableDiffusionControlNetPipeline.from_pretrained( |
| "SG161222/Realistic_Vision_V4.0_noVAE", controlnet=controlnet, torch_dtype=torch.float16 |
| ).to('cuda') |
|
|
| pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) |
|
|
| prompt = "a boxer in a boxing ring, best quality" |
| negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality" |
|
|
| images = [op_image, dp_image] |
|
|
| image = pipe( |
| prompt, |
| images, |
| num_inference_steps=20, |
| negative_prompt=negative_prompt, |
| controlnet_conditioning_scale=[1.0, 0.8], |
| ).images[0] |
|
|
| |
| filename, extension = os.path.splitext(os.path.basename(args.op_image)) |
|
|
| |
| output_path = os.path.join(args.output_dir, filename + extension) |
|
|
| print(type(image)) |
| |
| image.save(output_path) |
| print("saved in output directory!") |