Upload folder using huggingface_hub

app.py

@@ -52,9 +52,9 @@ config = "configs/pipelines/stage_2_base.py"
 config = Config.fromfile(config)
 model = BUILDER.build(config.model).cuda().bfloat16().eval()
 checkpoint_path = "checkpoints/Puffin-Base.pth"
-checkpoint = torch.load(checkpoint_path)
+checkpoint = torch.load(checkpoint_path, map_location='cpu')
 _ = model.load_state_dict(checkpoint, strict=False)
-_ = model.vae.load_state_dict(torch.load('checkpoints/vae.pth'), strict=True)
+_ = model.vae.load_state_dict(torch.load('checkpoints/vae.pth', map_location='cpu'), strict=True)
 
 def fig_to_image(fig):
     buf = io.BytesIO()