cuda

app.py

@@ -10,12 +10,12 @@ from decomp_diffusion.diffusion.respace import SpacedDiffusion
 from decomp_diffusion.gen_image import *
 from download import download_model
 
+import gradio as gr
+
 # fix randomness
 th.manual_seed(0)
 np.random.seed(0)
 
-import gradio as gr
-
 
 def get_pil_im(im, resolution=64):
     im = imresize(im, (resolution, resolution))[:, :, :3]
@@ -112,7 +112,7 @@ model_kwargs.update(dict(
 clevr_model = create_diffusion_model(**model_kwargs)
 clevr_model.eval()
 
-device = 'cuda'
+device = 'cuda' if th.cuda.is_available() else 'cpu'
 clevr_model.to(device)
 
 print(f'loading from {ckpt_path}')