|
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!") |