File size: 434 Bytes
e18985e
815ab0f
e18985e
76b138f
 
815ab0f
76b138f
815ab0f
 
e2aa61c
76b138f
 
 
 
9fca8be
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
from diffusers import StableDiffusionPipeline
import torch

def execute_nvidia(prompt, filename):
    DEVICE = 'cuda'
    model_id = "runwayml/stable-diffusion-v1-5"
    pipe = StableDiffusionPipeline.from_pretrained(
        model_id,
        use_auth_token=True,
        torch_dtype=torch.float16
    ).to(DEVICE)
    pipe.enable_attention_slicing()
    image = pipe(prompt).images[0]
    image.save("pictures/"+ filename +".png")