import torch from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler, DPMSolverMultistepScheduler, \ OnnxStableDiffusionPipeline import pipeline_openvino_stable_diffusion def get_sd_21(): model_id = "stabilityai/stable-diffusion-2-1-base" scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler") if torch.cuda.is_available(): pipe = StableDiffusionPipeline.from_pretrained( model_id, scheduler=scheduler, # safety_checker=None, revision="fp16", torch_dtype=torch.float16) pipe = pipe.to('cuda') else: pipe = StableDiffusionPipeline.from_pretrained( model_id, scheduler=scheduler, # safety_checker=None, revision="fp16", torch_dtype=torch.float16) return pipe def get_sd_small(): model_id = 'OFA-Sys/small-stable-diffusion-v0' scheduler = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler") onnx_pipe = OnnxStableDiffusionPipeline.from_pretrained( "OFA-Sys/small-stable-diffusion-v0", scheduler=scheduler, revision="onnx", provider="CPUExecutionProvider", ) pipe = pipeline_openvino_stable_diffusion.OpenVINOStableDiffusionPipeline.from_onnx_pipeline(onnx_pipe) return pipe