#!/usr/bin/env python3
from diffusers import StableDiffusionPipeline, KDPM2DiscreteScheduler, StableDiffusionImg2ImgPipeline, HeunDiscreteScheduler, KDPM2AncestralDiscreteScheduler, DDIMScheduler,  DPMSolverMultistepScheduler
import time
import os
from huggingface_hub import HfApi
# from compel import Compel
import torch
import sys
from pathlib import Path
import requests
from PIL import Image
from io import BytesIO

path = "runwayml/stable-diffusion-v1-5"
lora_id = "takuma104/lora-test-text-encoder-lora-target"

api = HfApi()
start_time = time.time()
pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
pipe.load_lora_weights(lora_id)
pipe = pipe.to("cuda")

prompt = "a red sks dog"

images = pipe(prompt=prompt, 
    num_inference_steps=15, 
    cross_attention_kwargs={"scale": 0.5},
    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")