Update app.py

app.py

@@ -51,6 +51,8 @@ css = """
 
 set_seed(666)
 
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 # DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 default_steps = 20
@@ -63,8 +65,10 @@ ckpt_path = hf_hub_download(
 state_dict = torch.load(ckpt_path, map_location="cpu")
 model.load_state_dict(state_dict, strict=False)
 model = model.eval()
+model = model.to(DEVICE)
 
 moge_model = MoGeModel.from_pretrained("Ruicheng/moge-2-vitl-normal").eval()
+moge_model = moge_model.to(DEVICE)
 
 
 def main(share=True):