#!/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 = sys.argv[1]
path = "gsdf/Counterfeit-V2.5"
# path = "stabilityai/stable-diffusion-2-1"

api = HfApi()
start_time = time.time()
pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    pipe.scheduler.config, use_karras_sigmas=True
)
pipe = pipe.to("cuda")

pipe.load_lora_weights(".", weight_name="light_and_shadow.safetensors")

prompt = "masterpiece, best quality, 1girl, at dusk"
negative_prompt = ("(low quality, worst quality:1.4), (bad anatomy), (inaccurate limb:1.2), "
                   "bad composition, inaccurate eyes, extra digit, fewer digits, (extra arms:1.2), large breasts")

pipe.enable_xformers_memory_efficient_attention()
images = pipe(prompt=prompt, 
    negative_prompt=negative_prompt, 
    width=512, 
    height=768, 
    num_inference_steps=15, 
    num_images_per_prompt=4,
    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")