Update gradio_app.py

gradio_app.py

@@ -63,6 +63,154 @@ example_files = [
     os.path.join("demo_files/examples", f) for f in os.listdir("demo_files/examples")
 ]
 
+def create_zip_file(glb_file, pc_file, illumination_file):
+    if not all([glb_file, pc_file, illumination_file]):
+        return None
+
+    # Create a temporary zip file
+    temp_dir = tempfile.mkdtemp()
+    zip_path = os.path.join(temp_dir, "spar3d_output.zip")
+
+    with zipfile.ZipFile(zip_path, "w") as zipf:
+        zipf.write(glb_file, "mesh.glb")
+        zipf.write(pc_file, "points.ply")
+        zipf.write(illumination_file, "illumination.hdr")
+
+    generated_files.append(zip_path)
+    return zip_path
+
+def forward_model(
+    batch,
+    system,
+    guidance_scale=3.0,
+    seed=0,
+    device="cuda",
+    remesh_option="none",
+    vertex_count=-1,
+    texture_resolution=1024,
+):
+    batch_size = batch["rgb_cond"].shape[0]
+
+    # prepare the condition for point cloud generation
+    # set seed
+    random.seed(seed)
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    cond_tokens = system.forward_pdiff_cond(batch)
+
+    if "pc_cond" not in batch:
+        sample_iter = system.sampler.sample_batch_progressive(
+            batch_size,
+            cond_tokens,
+            guidance_scale=guidance_scale,
+            device=device,
+        )
+        for x in sample_iter:
+            samples = x["xstart"]
+        batch["pc_cond"] = samples.permute(0, 2, 1).float()
+        batch["pc_cond"] = spar3d_utils.normalize_pc_bbox(batch["pc_cond"])
+
+    # subsample to the 512 points
+    batch["pc_cond"] = batch["pc_cond"][
+        :, torch.randperm(batch["pc_cond"].shape[1])[:512]
+    ]
+
+    # get the point cloud
+    xyz = batch["pc_cond"][0, :, :3].cpu().numpy()
+    color_rgb = (batch["pc_cond"][0, :, 3:6] * 255).cpu().numpy().astype(np.uint8)
+    pc_rgb_trimesh = trimesh.PointCloud(vertices=xyz, colors=color_rgb)
+
+    # forward for the final mesh
+    trimesh_mesh, _glob_dict = model.generate_mesh(
+        batch,
+        texture_resolution,
+        remesh=remesh_option,
+        vertex_count=vertex_count,
+        estimate_illumination=True,
+    )
+    trimesh_mesh = trimesh_mesh[0]
+    illumination = _glob_dict["illumination"]
+
+    return trimesh_mesh, pc_rgb_trimesh, illumination.cpu().detach().numpy()[0]
+
+def process_model_run(
+    fr_res,
+    guidance_scale,
+    random_seed,
+    pc_cond,
+    remesh_option,
+    vertex_count_type,
+    vertex_count,
+    texture_resolution,
+):
+    start = time.time()
+    with torch.no_grad():
+        with (
+            torch.autocast(device_type=device, dtype=torch.bfloat16)
+            if "cuda" in device
+            else nullcontext()
+        ):
+            model_batch = create_batch(fr_res)
+            model_batch = {k: v.to(device) for k, v in model_batch.items()}
+
+            trimesh_mesh, trimesh_pc, illumination_map = forward_model(
+                model_batch,
+                model,
+                guidance_scale=guidance_scale,
+                seed=random_seed,
+                device="cuda",
+                remesh_option=remesh_option.lower(),
+                vertex_count=vertex_count,
+                texture_resolution=texture_resolution,
+            )
+
+    # Create new tmp file
+    temp_dir = tempfile.mkdtemp()
+    tmp_file = os.path.join(temp_dir, "mesh.glb")
+
+    trimesh_mesh.export(tmp_file, file_type="glb", include_normals=True)
+    generated_files.append(tmp_file)
+
+    tmp_file_pc = os.path.join(temp_dir, "points.ply")
+    trimesh_pc.export(tmp_file_pc)
+    generated_files.append(tmp_file_pc)
+
+    tmp_file_illumination = os.path.join(temp_dir, "illumination.hdr")
+    cv2.imwrite(tmp_file_illumination, illumination_map)
+    generated_files.append(tmp_file_illumination)
+
+    print("Generation took:", time.time() - start, "s")
+
+    return tmp_file, tmp_file_pc, tmp_file_illumination, trimesh_pc
+
+def create_batch(input_image: Image) -> dict[str, Any]:
+    img_cond = (
+        torch.from_numpy(
+            np.asarray(input_image.resize((COND_WIDTH, COND_HEIGHT))).astype(np.float32)
+            / 255.0
+        )
+        .float()
+        .clip(0, 1)
+    )
+    mask_cond = img_cond[:, :, -1:]
+    rgb_cond = torch.lerp(
+        torch.tensor(BACKGROUND_COLOR)[None, None, :], img_cond[:, :, :3], mask_cond
+    )
+
+    batch_elem = {
+        "rgb_cond": rgb_cond,
+        "mask_cond": mask_cond,
+        "c2w_cond": c2w_cond.unsqueeze(0),
+        "intrinsic_cond": intrinsic.unsqueeze(0),
+        "intrinsic_normed_cond": intrinsic_normed_cond.unsqueeze(0),
+    }
+    # Add batch dim
+    batched = {k: v.unsqueeze(0) for k, v in batch_elem.items()}
+    return batched
+
+def remove_background(input_image: Image) -> Image:
+    return bg_remover.process(input_image.convert("RGB"))
+
 def auto_process(input_image):
     if input_image is None:
         return None, None, None, None