Check if gpu is available

app.py

@@ -7,27 +7,35 @@ from diffusers import StableDiffusionPipeline
 # get hf user access token as an environment variable
 TOKEN_KEY = os.getenv('AUTH_TOKEN')
 
+# choose GPU else fallback to CPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+device_name = torch.cuda.get_device_name(0)
+
 # setup pipeline
 pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, use_auth_token=TOKEN_KEY)
-pipe = pipe.to('cuda')
+pipe = pipe.to(device)
 
 # define gradio function
 def generate(prompt:str, seed:int, guidance:float):
-    generator = torch.Generator("cuda").manual_seed(int(seed))
-    with autocast("cuda"):
+    generator = torch.Generator(device).manual_seed(int(seed))
+    with autocast(device):
         image = pipe(prompt=prompt, generator=generator, guidance_scale=guidance, steps=50).images[0]
     return image
 
-# create the gradio UI
-demo = gr.Interface(
-    fn=generate,
-    inputs=[gr.Textbox(placeholder="castle on a mountain"), gr.Number(value=123456), gr.Slider(0,10)],
-    outputs="image",
-    allow_flagging="never",
-)
+if device == "cuda":
+    print(device_name + " available.")
+    # create the gradio UI
+    demo = gr.Interface(
+        fn=generate,
+        inputs=[gr.Textbox(placeholder="castle on a mountain"), gr.Number(value=123456), gr.Slider(0,10)],
+        outputs="image",
+        allow_flagging="never",
+    )
 
-# allow queueing or incoming requests, max=3 
-demo.queue(concurrency_count=3)
+    # allow queueing or incoming requests, max=3 
+    demo.queue(concurrency_count=3)
 
-# launch demo
-demo.launch()
+    # launch demo
+    demo.launch()
+else:
+    print("GPU unavailable.")