update app.py

app.py

@@ -52,7 +52,8 @@ def add_text(rgb, caption):
     return rgb
 
 config = "3DTopia/configs/default.yaml"
-local_ckpt = "3DTopia/checkpoints/3dtopia_diffusion_state_dict.ckpt"
+# local_ckpt = "3DTopia/checkpoints/3dtopia_diffusion_state_dict.ckpt"
+local_ckpt = "/data/3DTopia_all/3DTopia_code/checkpoints/model.safetensors"
 if os.path.exists(local_ckpt):
     ckpt = local_ckpt
 else:
@@ -62,56 +63,65 @@ configs = OmegaConf.load(config)
 os.makedirs("tmp", exist_ok=True)
 print("download finish")
 
-if ckpt.endswith(".ckpt"):
-    model = get_obj_from_str(configs.model["target"]).load_from_checkpoint(ckpt, map_location='cpu', strict=False, **configs.model.params)
-elif ckpt.endswith(".safetensors"):
-    print("download finish")
-    model = get_obj_from_str(configs.model["target"])(**configs.model.params)
-    print("download finish")
-    model_ckpt = load_file(ckpt)
-    print("download finish")
-    model.load_state_dict(model_ckpt)
+import sys
+import traceback
+
+try:
+    if ckpt.endswith(".ckpt"):
+        model = get_obj_from_str(configs.model["target"]).load_from_checkpoint(ckpt, map_location='cpu', strict=False, **configs.model.params)
+    elif ckpt.endswith(".safetensors"):
+        print("download finish")
+        model = get_obj_from_str(configs.model["target"])(**configs.model.params)
+        print("download finish")
+        model_ckpt = load_file(ckpt)
+        print("download finish")
+        model.load_state_dict(model_ckpt)
+        print("download finish")
+    else:
+        raise NotImplementedError
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    model = model.to(device)
     print("download finish")
-else:
-    raise NotImplementedError
-device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-model = model.to(device)
-print("download finish")
-sampler = DDIMSampler(model)
-
-img_size = configs.model.params.unet_config.params.image_size
-channels = configs.model.params.unet_config.params.in_channels
-shape = [channels, img_size, img_size * 3]
-
-pose_folder = '3DTopia/assets/sample_data/pose'
-poses_fname = sorted([os.path.join(pose_folder, f) for f in os.listdir(pose_folder)])
-batch_rays_list = []
-H = 128
-ratio = 512 // H
-for p in poses_fname:
-    c2w = np.loadtxt(p).reshape(4, 4)
-    c2w[:3, 3] *= 2.2
-    c2w = np.array([
-        [1, 0, 0, 0],
-        [0, 0, -1, 0],
-        [0, 1, 0, 0],
-        [0, 0, 0, 1]
-    ]) @ c2w
-
-    k = np.array([
-        [560 / ratio, 0, H * 0.5],
-        [0, 560 / ratio, H * 0.5],
-        [0, 0, 1]
-    ])
-
-    rays_o, rays_d = get_rays(H, H, torch.Tensor(k), torch.Tensor(c2w[:3, :4]))
-    coords = torch.stack(torch.meshgrid(torch.linspace(0, H-1, H), torch.linspace(0, H-1, H), indexing='ij'), -1)
-    coords = torch.reshape(coords, [-1,2]).long()
-    rays_o = rays_o[coords[:, 0], coords[:, 1]]
-    rays_d = rays_d[coords[:, 0], coords[:, 1]]
-    batch_rays = torch.stack([rays_o, rays_d], 0)
-    batch_rays_list.append(batch_rays)
-batch_rays_list = torch.stack(batch_rays_list, 0)
+    sampler = DDIMSampler(model)
+
+    img_size = configs.model.params.unet_config.params.image_size
+    channels = configs.model.params.unet_config.params.in_channels
+    shape = [channels, img_size, img_size * 3]
+
+    pose_folder = '3DTopia/assets/sample_data/pose'
+    poses_fname = sorted([os.path.join(pose_folder, f) for f in os.listdir(pose_folder)])
+    batch_rays_list = []
+    H = 128
+    ratio = 512 // H
+    for p in poses_fname:
+        c2w = np.loadtxt(p).reshape(4, 4)
+        c2w[:3, 3] *= 2.2
+        c2w = np.array([
+            [1, 0, 0, 0],
+            [0, 0, -1, 0],
+            [0, 1, 0, 0],
+            [0, 0, 0, 1]
+        ]) @ c2w
+
+        k = np.array([
+            [560 / ratio, 0, H * 0.5],
+            [0, 560 / ratio, H * 0.5],
+            [0, 0, 1]
+        ])
+
+        rays_o, rays_d = get_rays(H, H, torch.Tensor(k), torch.Tensor(c2w[:3, :4]))
+        coords = torch.stack(torch.meshgrid(torch.linspace(0, H-1, H), torch.linspace(0, H-1, H), indexing='ij'), -1)
+        coords = torch.reshape(coords, [-1,2]).long()
+        rays_o = rays_o[coords[:, 0], coords[:, 1]]
+        rays_d = rays_d[coords[:, 0], coords[:, 1]]
+        batch_rays = torch.stack([rays_o, rays_d], 0)
+        batch_rays_list.append(batch_rays)
+    batch_rays_list = torch.stack(batch_rays_list, 0)
+except Exception as e:
+    print(e)
+    print(traceback.format_exc())
+    print(sys.exc_info()[2])
+
 
 print("download finish")
 def marching_cube(b, text, global_info):