#!/usr/bin/env python3
from huggingface_hub import HfApi
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, StableDiffusionInpaintPipeline
from pathlib import Path
import os
import sys
import requests
from PIL import Image
from io import BytesIO

org_name = "lykon-models"
org_name = "lykon-absolute-realism"

file_path = sys.argv[1]
file_name = file_path.split("/")[-1].split(".safetensors")[0].replace("_", "-").replace(".", "-").lower()

if file_name.endswith("-"):
    file_name = file_name[:-1]

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

kwargs = {}
if "xl" in file_name or "alpha2" in file_name:
    pipe = StableDiffusionXLPipeline.from_single_file(file_path)
elif "inpaint" in file_name:
    pipe = StableDiffusionInpaintPipeline.from_single_file(file_path)
    img_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 = download_image(img_url).resize((512, 512))
    mask_image = download_image(mask_url).resize((512, 512))
    kwargs = {"image": init_image, "mask_image": mask_image}
else:
    pipe = StableDiffusionPipeline.from_single_file(file_path)

print("Test...")
prompt = "A lion in galaxies, spirals, nebulae, stars, smoke, iridescent, intricate detail, octane render, 8k"
pipe.to("cuda")
images = pipe(prompt=prompt, num_inference_steps=25, **kwargs).images
pipe.to("cpu")

api = HfApi()
for i, image in enumerate(images):
    image_file_name = f"bb_1_{i}"
    path = os.path.join(Path.home(), "images", f"{image_file_name}.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(f"https://huggingface.co/datasets/patrickvonplaten/images/blob/main/{image_file_name}.png")

print("Upload...")
model_id = os.path.join(org_name, file_name)
pipe.push_to_hub(model_id, private=True)