# # !pip install opencv-python transformers accelerate
# from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
# from diffusers.utils import load_image
# import numpy as np
# import torch

# import cv2
# from PIL import Image

# # download an image
# image = load_image(
#     "/eph/nvme0/azureml/cr/j/8569d5e3aa08485780b67a53d671e109/exe/wd/1_2M_Dataset/2.jpg"
# )
# image = np.array(image)

# # get canny image
# image = cv2.Canny(image, 100, 200)
# image = image[:, :, None]
# image = np.concatenate([image, image, image], axis=2)
# canny_image = Image.fromarray(image)

# # load control net and stable diffusion v1-5
# controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
# pipe = StableDiffusionControlNetPipeline.from_pretrained(
#     "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
# )

# # speed up diffusion process with faster scheduler and memory optimization
# pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
# # remove following line if xformers is not installed
# pipe.enable_xformers_memory_efficient_attention()

# pipe.enable_model_cpu_offload()

# # generate image
# generator = torch.manual_seed(0)
# image = pipe(
#     "Make the cup in front of Eiffle Tower", num_inference_steps=20, generator=generator, image=canny_image
# ).images[0].save("result.png")

import PIL
import requests
import torch
from io import BytesIO
from diffusers.utils import load_image

from diffusers import StableDiffusionInstructPix2PixPipeline


# def download_image(url):
#     response = requests.get(url)
#     return PIL.Image.open(BytesIO(response.content)).convert("RGB")


img_url = "/eph/nvme0/azureml/cr/j/8569d5e3aa08485780b67a53d671e109/exe/wd/1_2M_Dataset/2.jpg"

image = load_image(img_url).resize((512, 512))

pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
    "timbrooks/instruct-pix2pix", torch_dtype=torch.float16
)
pipe = pipe.to("cuda")

prompt = "make the cup in front of Eiffle Tower"
image = pipe(prompt=prompt, image=image).images[0].save("result.png")



accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks cup" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=400 \
  --push_to_hub