| from diffusers import StableDiffusionPipeline | |
| def execute_nvidia(prompt, filename): | |
| DEVICE = 'cuda' | |
| pipe = StableDiffusionPipeline.from_pretrained( | |
| "CompVis/stable-diffusion-v1-4", | |
| use_auth_token=True | |
| ).to(DEVICE) | |
| # prompt = "Elon Musk riding elephant on the moon." | |
| pipe.enable_attention_slicing() | |
| # _ = pipe(prompt, num_inference_steps=1) | |
| image = pipe(prompt).images[0] | |
| image.save("pictures/"+ filename +".png") | |