app

app.py

@@ -14,7 +14,7 @@ import tempfile
 
 # tempfile.tempdir = "/home/dachuang/gradio/tmp/"
 
-# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
 
 title_markdown = ("""
 欢迎来到甲骨文文字演变模拟器

models/ddm.py

@@ -124,7 +124,10 @@ class DenoisingDiffusion(object):
         self.num_timesteps = betas.shape[0]
 
     def load_ddm_ckpt(self, load_path, ema=False):
-        checkpoint = utils.logging.load_checkpoint(load_path, None)
+        if self.device == torch.device('cpu'):
+            checkpoint = utils.logging.load_checkpoint(load_path, self.device)
+        else:
+            checkpoint = utils.logging.load_checkpoint(load_path, None)
         self.start_epoch = checkpoint['epoch']
         self.step = checkpoint['step']
         self.model.load_state_dict(checkpoint['state_dict'], strict=True)