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