nyrt/docker_python

1print("Start")2import torch3from diffusers import StableDiffusionPipeline4 5model_id = "CompVis/stable-diffusion-v1-4"6device = "cuda"7 8print("Init Pipeline")9pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)10print("Add device to pipeline")11pipe = pipe.to(device)12 13prompt = "a photo of an astronaut riding a horse on mars"14print("Generating image")15image = pipe(prompt).images[0]16print("Done")