from diffusers.utils import load_image


def get_image_for_img_to_img(pipeline_to_benchmark):
    if pipeline_to_benchmark == "SD I2I":
        url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
        init_image = load_image(url).convert("RGB")
        init_image = init_image.resize((512, 512))
    elif pipeline_to_benchmark == "SDXL I2I":
        url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl-img2img.png"
        init_image = load_image(url).convert("RGB")
    return init_image


def get_image_for_inpainting(pipeline_to_benchmark):
    if pipeline_to_benchmark == "SD Inpainting":
        image_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
        mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
        init_image = load_image(image_url).convert("RGB").resize((512, 512))
        mask_image = load_image(mask_url).convert("RGB").resize((512, 512))
    elif pipeline_to_benchmark == "SDXL Inpainting":
        image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl-text2img.png"
        mask_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl-inpaint-mask.png"
        init_image = load_image(image_url).convert("RGB")
        mask_image = load_image(mask_url).convert("RGB")
    return init_image, mask_image


def get_image_for_controlnet(pipeline_to_benchmark):
    if pipeline_to_benchmark == "SD ControlNet":
        image = load_image(
            "https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird_canny.png"
        ).convert("RGB")
    elif pipeline_to_benchmark == "SDXL ControlNet":
        image = load_image(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/hf-logo.png"
        ).convert("RGB")

    return image


def get_image_for_adapters(pipeline_to_benchmark):
    if pipeline_to_benchmark == "SD T2I Adapters":
        image = load_image(
            "https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/sd_t2i_canny.png"
        )
    elif pipeline_to_benchmark == "SDXL T2I Adapters":
        image = load_image(
            "https://huggingface.co/Adapter/t2iadapter/resolve/main/figs_SDXLV1.0/cond_canny.png"
        ).convert("RGB")
    return image