#!/usr/bin/env python3
from diffusers import AutoPipelineForText2Image
import time
import os
from huggingface_hub import HfApi
import torch
from pathlib import Path

path = "stabilityai/stable-diffusion-xl-base-1.0"

api = HfApi()
start_time = time.time()
pipe = AutoPipelineForText2Image.from_pretrained(path, torch_dtype=torch.float16)
pipe.enable_model_cpu_offload()
lora_model_id = "hf-internal-testing/sdxl-0.9-kamepan-lora"
lora_model_id = "TheLastBen/Papercut_SDXL"
lora_filename = "kame_sdxl_v2-000020-16rank.safetensors"
lora_filename = "papercut.safetensors"
pipe.load_lora_weights(lora_model_id, weight_name=lora_filename)

prompt = "masterpiece, best quality, mountain"
prompt = "papercut sonic"

images = pipe(prompt=prompt, 
    num_inference_steps=20, 
    generator=torch.manual_seed(0)
).images


for i, image in enumerate(images):
    file_name = f"aa_{i}"
    path = os.path.join(Path.home(), "images", f"{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/{file_name}.png")