import torch
from diffusers import StableDiffusionXLPipeline, StableDiffusionPipeline
from typing import cast
from diffusers import DPMSolverMultistepScheduler

sdxl_model = cast(StableDiffusionXLPipeline, StableDiffusionXLPipeline.from_pretrained(
    'stabilityai/stable-diffusion-xl-base-1.0',
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16",
)).to('cuda')
sd_model = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
    revision="fp16",
).to('cuda')

common_config = {'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': 'scaled_linear'}
dpmpp_2m_sde = DPMSolverMultistepScheduler(**common_config, algorithm_type="sde-dpmsolver++")
sdxl_model.scheduler = dpmpp_2m_sde
sd_model.scheduler = dpmpp_2m_sde

sdxl_model.watermark = None
generator = torch.Generator(device='cuda')
generator.manual_seed(12345)

params = {
    "prompt": ['a cat'],
    "num_inference_steps": 50,
    "height": 1024,
    "width": 1024,
    "guidance_scale": 7,
}

sdxl_res = sdxl_model(**params, denoising_end=1.0, generator=generator)
sdxl_img = sdxl_res.images[0]

generator.manual_seed(12345)
sd_res = sd_model(**params, generator=generator)
sd_img = sd_res.images[0]

display(sdxl_img)
display(sd_img)