from diffusers import DiffusionPipeline
from diffusers import (
    DDPMScheduler,
    DDIMScheduler,
    PNDMScheduler,
    LMSDiscreteScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    DPMSolverMultistepScheduler,
)

import torch, time, datetime, os
import hashlib


repo_id = "stabilityai/stable-diffusion-xl-base-1.0"
scheduler = DPMSolverMultistepScheduler.from_pretrained(repo_id, subfolder="scheduler")

pipe = DiffusionPipeline.from_pretrained(repo_id, scheduler=scheduler, use_safetensors=True, torch_dtype=torch.float16, variant="fp16")
# pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
# 保存模型
# pipe.save_pretrained("models/stable-diffusion-xl-base-1.0")

pipe.to("cuda")

# if using torch < 2.0
# pipe.enable_xformers_memory_efficient_attention()

prompt = "chinese modern girl，Real photo of a girl, super high definition, super beautiful, lots of details, super good lighting, beautiful, cute, and lovely，anime"


# images = pipe(prompt=prompt).images[0]
# 保存图片
# images.save("images/stable-diffusion-xl-base-1.0.png")
today = datetime.date.today()
m = hashlib.md5()
m.update(prompt.encode("utf-8"))
prompt_hash = m.hexdigest()


save_path = f"images/{today}/{int(time.time())}_{prompt_hash[:8]}"

os.makedirs(save_path, exist_ok=True)

# 指定采样器


# 制定随机数并生成多张图片
num_inference_steps = 50
while True:
    torch.cuda.empty_cache()
    current = int(time.time())
    generator = torch.Generator("cuda").manual_seed(current)
    images = pipe(prompt=prompt, generator=generator, num_images_per_prompt=2, num_inference_steps=num_inference_steps, width=1024, height=1024).images
    for index, img in enumerate(images):
        img.save(f"{save_path}/{current}_{index}_step_{num_inference_steps}.jpg")

import os, time
# os.environ['CUDA_VISIBLE_DEVICES'] = '1'



from diffusers import DiffusionPipeline
import torch

# load both base & refiner
base = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
base.to("cuda")
refiner = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    text_encoder_2=base.text_encoder_2,
    vae=base.vae,
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16",
)
refiner.to("cuda")

# Define how many steps and what % of steps to be run on each experts (80/20) here
n_steps = 40
high_noise_frac = 0.8

prompt = "panda couples wearing hoodie with flowers, 4k, --chaos 60"

# monkey
for _ in range(5):
    seed = int(time.time()*1000)
    generator = torch.Generator("cuda").manual_seed(seed)
    # run both experts
    image = base(
        prompt=prompt,
        num_inference_steps=n_steps,
        denoising_end=high_noise_frac,
        generator=generator,
        output_type="latent",
        width=768,height=512,
        num_images_per_prompt=2
    ).images

    refiner_images = refiner(
        prompt=prompt,
        num_inference_steps=n_steps,
        denoising_start=high_noise_frac,
        generator=generator,
        image=image,
        width=768,height=512,
        num_images_per_prompt=2
    ).images

    for index, img in enumerate(refiner_images):
        img.save(f"data/img_{seed//1000}_{index}.jpg")
    torch.cuda.empty_cache()