from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler
import torch

# this will substitute the default PNDM scheduler for K-LMS  
lms = LMSDiscreteScheduler(
    beta_start=0.00085, 
    beta_end=0.012, 
    beta_schedule="scaled_linear"
)

guidance_scale=8.5
seed=777
steps=50

cartoon_model_path = "Norod78/sd2-simpsons-blip"
cartoon_pipe = StableDiffusionPipeline.from_pretrained(cartoon_model_path, scheduler=lms, torch_dtype=torch.float16)
cartoon_pipe.to("cuda")

def generate(prompt, file_prefix ,samples):
    torch.manual_seed(seed)
    prompt += ", Very detailed, clean, high quality, sharp image"
    cartoon_images = cartoon_pipe([prompt] * samples, num_inference_steps=steps, guidance_scale=guidance_scale)["images"]
    for idx, image in enumerate(cartoon_images):
        image.save(f"{file_prefix}-{idx}-{seed}-sd2-simpsons-blip.jpg")

generate("An oil painting of Snoop Dogg as a simpsons character", "01_SnoopDog", 4)
generate("Gal Gadot, cartoon", "02_GalGadot", 4)
generate("A cartoony Simpsons town", "03_SimpsonsTown", 4)
generate("Pikachu with the Simpsons, Eric Wallis", "04_PikachuSimpsons", 4)