#!/usr/bin/env python3
from diffusers import DiffusionPipeline
from safetensors.torch import load_file
import torch
from pathlib import Path
from huggingface_hub import HfApi, hf_hub_download
import os
import hf_image_uploader as hiu
import time

api = HfApi()


pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
pipe.to("cuda")

# pipe.load_lora_weights("stabilityai/stable-diffusion-xl-base-1.0", weight_name="sd_xl_offset_example-lora_1.0.safetensors", low_cpu_mem_usage=True)
# file = hf_hub_download("TheLastBen/Papercut_SDXL", filename="papercut.safetensors")
file = hf_hub_download("hf-internal-testing/sdxl-0.9-daiton-lora", filename="daiton-xl-lora-test.safetensors")

start_time = time.time()
pipe.load_lora_weights("Pclanglais/TintinIA")
pipe.fuse_lora()
pipe.unfuse_lora()
pipe.unload_lora_weights()

pipe.load_lora_weights("ProomptEngineer/pe-balloon-diffusion-style")
pipe.fuse_lora()
pipe.unload_lora_weights()
pipe.unfuse_lora()

pipe.load_lora_weights("ostris/crayon_style_lora_sdxl")
pipe.fuse_lora()
pipe.unload_lora_weights()
pipe.unfuse_lora()
print(time.time() - start_time)

prompt = "masterpiece, best quality, mountain"

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

hiu.upload(images[0], "patrickvonplaten/images")