update v0.0.5.post1

2024-04-29 16:02:23 +08:00
parent f502f8063d
commit 6be496d5ac
5 changed files with 11 additions and 3 deletions
@@ -255,8 +255,11 @@ class DiffusionInference():
            module['device'] = 'offline'
            print('delete module')
        else:
-            module['model'] = module['model'].to('cpu')
-            module['device'] = 'cpu'
+            if module['model'] is not None:
+                module['model'] = module['model'].to('cpu')
+                module['device'] = 'cpu'
+            else:
+                module['device'] = 'offline'
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
        return module