update space

app.py

@@ -5,7 +5,7 @@ from functools import partial
 from typing import Optional
 from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
 from shap_e.diffusion.sample import sample_latents
-from shap_e.models.download import load_model, load_config
+from shap_e.models.download import load_model, load_config, load_checkpoint
 from shap_e.util.notebooks import create_pan_cameras, decode_latent_mesh
 import trimesh
 import torch.nn as nn
@@ -275,10 +275,25 @@ def main():
     """
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     print("device:", device)
-    latent_model = load_model('text300M', device=device)
+    # latent_model = load_model('text300M', device=device)
+
+    latent_model = model_from_config(load_config('text300M'), device=device)
+    # print(model_name, kwargs)
+    # print(model)
+    latent_model.load_state_dict(load_checkpoint('text300M', device='cpu'))
+    latent_model.eval()
     print("loaded latent model")
-    xm = load_model('transmitter', device=device)
+    latent_model.to(device)
+    # xm = load_model('transmitter', device=device)
+
+    xm = model_from_config(load_config('transmitter'), device=device)
+    # print(model_name, kwargs)
+    # print(model)
+    xm.load_state_dict(load_checkpoint('transmitter', device='cpu'))
+    xm.eval()
     print("loaded transmitter")
+    xm.to(device)
+
     diffusion = diffusion_from_config(load_config('diffusion'))
     freeze_params(xm.parameters())
     models = dict()