#!/usr/bin/env python3
from huggingface_hub import HfApi
import torch

# https://github.com/pix2pixzero/pix2pix-zero/blob/main/src/edit_synthetic.py
import requests

from diffusers import DDIMScheduler, StableDiffusionPix2PixZeroPipeline

api = HfApi()


def download(embedding_url, local_filepath):
    r = requests.get(embedding_url)
    with open(local_filepath, "wb") as f:
        f.write(r.content)


model_ckpt = "CompVis/stable-diffusion-v1-4"
pipeline = StableDiffusionPix2PixZeroPipeline.from_pretrained(
    model_ckpt, conditions_input_image=False, torch_dtype=torch.float16
)
pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
pipeline.to("cuda")


prompt = "a high resolution painting of a cat in the style of van gough"
source_embedding_url = "https://github.com/pix2pixzero/pix2pix-zero/raw/main/assets/embeddings_sd_1.4/cat.pt"
target_embedding_url = "https://github.com/pix2pixzero/pix2pix-zero/raw/main/assets/embeddings_sd_1.4/dog.pt"

for url in [source_embedding_url, target_embedding_url]:
    download(url, url.split("/")[-1])

source_embeds = torch.load(source_embedding_url.split("/")[-1])
target_embeds = torch.load(target_embedding_url.split("/")[-1])

image = pipeline(
    prompt,
    source_embeds=source_embeds,
    target_embeds=target_embeds,
    num_inference_steps=50,
    cross_attention_guidance_amount=0.15,
).images[0]

path = "/home/patrick_huggingface_co/images/aa.png"
image.save(path)

api.upload_file(
    path_or_fileobj=path,
    path_in_repo=path.split("/")[-1],
    repo_id="patrickvonplaten/images",
    repo_type="dataset",
)
print("https://huggingface.co/datasets/patrickvonplaten/images/blob/main/aa.png")