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")