device integration

app.py

@@ -4,9 +4,10 @@ import torch
 
 model_id = "stabilityai/stable-diffusion-2"
 scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
+device = troch.device("cuda" if torch.cuda.is_avaliable() else "cpu")
 
 pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, revision="fp16")
-pipe = pipe.to("cuda")
+pipe = pipe.to(device)
 
 def txt2img(prompt):
   image = pipe(prompt, height=768, width=768, guidance_scale = 10).images[0]