1import torch 2from diffusers import EulerDiscreteScheduler, StableDiffusionPipeline 3 4 5class StableDiffusion: 6 def __init__(self, seed=42): 7 self.model_id: str = "stabilityai/stable-diffusion-2-1-base" 8 self.generator = torch.manual_seed(seed) 9 self.scheduler = EulerDiscreteScheduler.from_pretrained( 10 self.model_id, subfolder="scheduler" 11 ) 12 13 self.pipe = StableDiffusionPipeline.from_pretrained( 14 self.model_id, scheduler=self.scheduler, torch_dtype=torch.float32 15 ) 16 self.pipe = self.pipe.to("cpu") 17 18 def __call__(self, prompt, height, width, num_time_steps): 19 image = self.pipe( 20 prompt, height, width, num_time_steps, generator=self.generator 21 ).images[0] 22 return image 23