File size: 2,111 Bytes
a5b2a9d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
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")
# Initialize the argument parser
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.')
# Parse the arguments
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]
# Extract the filename and extension from args.op_image
filename, extension = os.path.splitext(os.path.basename(args.op_image))
# Construct the full output path with the specified output directory
output_path = os.path.join(args.output_dir, filename + extension)
print(type(image))
# Save the image using PIL
image.save(output_path) # Assuming image is from PIL
print("saved in output directory!") |