from diffusers import DiffusionPipeline
import torch

device = "cuda"
generator = torch.Generator(device).manual_seed(2500)
prompt = "An astronaut riding a green horse"

base = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
base.to(device)

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(device)

generator = torch.Generator(device).manual_seed(2500)
prompt = "An astronaut riding a green horse"

n_steps = 40
high_noise_frac = 0.8

# run both experts
image = base(
    prompt=prompt,
    num_inference_steps=n_steps,
    denoising_end=high_noise_frac,
    output_type="latent",
    generator=generator,
    
).images

image = refiner(
    prompt=prompt,
    num_inference_steps=n_steps,
    denoising_start=high_noise_frac,
    image=image,
    generator=generator,
).images[0]
