Fix some bugs.

README.md

@@ -14,3 +14,6 @@ Official demo for **[CityDreamer: Compositional Generative Model of Unbounded 3D
 - 🔥 CityDreamer is a unbounded 3D city generator.
 - 🤗 Try CityDreamer to generate photolistic 3D cities.
 - ⚠️ Due to the limited computational resources at Hugging Face, this demo only generates **A SINGLE IMAGE** based on the New York City layout. If you wish to experience more comprehensive functionality, please visit the demo on [GitHub](https://github.com/hzxie/city-dreamer?tab=readme-ov-file#iterative-demo-%EF%B8%8F).
+
+❕IMPORTANT NOTE: We are working on migrating to ZeroGPU, so the demo might be temporarily unavailable.
+

app.py

@@ -4,13 +4,14 @@
 # @Author: Haozhe Xie
 # @Date:   2024-03-02 16:30:00
 # @Last Modified by: Haozhe Xie
-# @Last Modified at: 2024-09-21 10:36:47
+# @Last Modified at: 2024-09-22 10:31:28
 # @Email:  [email protected]
 
 import gradio as gr
 import logging
 import numpy as np
 import os
+
 import spaces
 import ssl
 import subprocess
@@ -26,25 +27,35 @@ ssl._create_default_https_context = ssl._create_unverified_context
 sys.path.append(os.path.join(os.path.dirname(__file__), "citydreamer"))
 
 
-def setup_runtime_env():
+def _get_output(cmd):
     try:
-        logging.info("CUDA version is %s" % subprocess.check_output(["nvcc", "--version"]))
+        return subprocess.check_output(cmd).decode("utf-8")
     except Exception as ex:
         logging.exception(ex)
 
-    logging.info("GCC version is %s" % subprocess.check_output(["g++", "--version"]))
-    # Instsall Pre-compiled CUDA extensions
+    return None
+
+
+def setup_runtime_env():
+    logging.info("Python Version: %s" % _get_output(["python", "--version"]))
+    logging.info("CUDA Version: %s" % _get_output(["nvcc", "--version"]))
+    logging.info("GCC Version: %s" % _get_output(["gcc", "--version"]))
+
+    # Install Pre-compiled CUDA extensions
     ext_dir = os.path.join(os.path.dirname(__file__), "wheels")
     for e in os.listdir(ext_dir):
-        if not os.path.join(ext_dir, e).endswith(".whl"):
-            subprocess.call(["pip", "install", os.path.join(ext_dir, e)])
-
+        logging.info("Installing Extensions from %s" % e)
+        subprocess.call(
+            ["pip", "install", os.path.join(ext_dir, e)], stderr=subprocess.STDOUT
+        )
     # Compile CUDA extensions
     # ext_dir = os.path.join(os.path.dirname(__file__), "citydreamer", "extensions")
     # for e in os.listdir(ext_dir):
     #     if os.path.isdir(os.path.join(ext_dir, e)):
     #         subprocess.call(["pip", "install", "."], cwd=os.path.join(ext_dir, e))
 
+    logging.info("Installed Python Packages: %s" % _get_output(["pip", "list"]))
+
 
 def get_models(file_name):
     import citydreamer.model
@@ -55,7 +66,8 @@ def get_models(file_name):
             file_name,
         )
 
-    ckpt = torch.load(file_name)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    ckpt = torch.load(file_name, map_location=torch.device(device))
     model = citydreamer.model.GanCraftGenerator(ckpt["cfg"])
     if torch.cuda.is_available():
         model = torch.nn.DataParallel(model).cuda().eval()
@@ -65,8 +77,8 @@ def get_models(file_name):
 
 
 def get_city_layout():
-    hf = np.array(Image.open("assets/NYC-HghtFld.png"))
-    seg = np.array(Image.open("assets/NYC-SegMap.png").convert("P"))
+    hf = np.array(Image.open("assets/NYC-HghtFld.png")).astype(np.int32)
+    seg = np.array(Image.open("assets/NYC-SegMap.png").convert("P")).astype(np.int32)
     return hf, seg
 
 
@@ -74,12 +86,14 @@ def get_city_layout():
 def get_generated_city(
     radius, altitude, azimuth, map_center, progress=gr.Progress(track_tqdm=True)
 ):
+    logging.info("CUDA is available: %s" % torch.cuda.is_available())
+    logging.info("PyTorch is built with CUDA: %s" % torch.version.cuda)
     # The import must be done after CUDA extension compilation
     import citydreamer.inference
 
     return citydreamer.inference.generate_city(
-        get_generated_city.fgm,
-        get_generated_city.bgm,
+        get_generated_city.fgm.to("cuda"),
+        get_generated_city.bgm.to("cuda"),
         get_generated_city.hf.copy(),
         get_generated_city.seg.copy(),
         map_center,

citydreamer/extensions/grid_encoder/__init__.py

@@ -4,7 +4,7 @@
 # @Author: Jiaxiang Tang (@ashawkey)
 # @Date:   2023-04-15 10:39:28
 # @Last Modified by: Haozhe Xie
-# @Last Modified at: 2023-04-15 13:08:46
+# @Last Modified at: 2024-09-22 10:26:30
 # @Email:  [email protected]
 # @Ref: https://github.com/ashawkey/torch-ngp
 
@@ -12,8 +12,6 @@ import math
 import numpy as np
 import torch
 
-import grid_encoder_ext
-
 
 class GridEncoderFunction(torch.autograd.Function):
     @staticmethod
@@ -55,6 +53,9 @@ class GridEncoderFunction(torch.autograd.Function):
                 1, device=inputs.device, dtype=embeddings.dtype
             )  # placeholder... TODO: a better way?
 
+        # Fix: libcudart.so.11.0: cannot open shared object file
+        import grid_encoder_ext
+
         grid_encoder_ext.forward(
             inputs,
             embeddings,
@@ -96,6 +97,9 @@ class GridEncoderFunction(torch.autograd.Function):
         else:
             grad_inputs = torch.zeros(1, device=inputs.device, dtype=embeddings.dtype)
 
+        # Fix: libcudart.so.11.0: cannot open shared object file
+        import grid_encoder_ext
+
         grid_encoder_ext.backward(
             grad,
             inputs,

citydreamer/inference.py

@@ -4,7 +4,7 @@
 # @Author: Haozhe Xie
 # @Date:   2024-03-02 16:30:00
 # @Last Modified by: Haozhe Xie
-# @Last Modified at: 2024-03-03 15:59:00
+# @Last Modified at: 2024-09-22 10:22:05
 # @Email:  [email protected]
 
 import copy
@@ -13,7 +13,6 @@ import logging
 import math
 import numpy as np
 import torch
-import torchvision
 
 import citydreamer.extensions.extrude_tensor
 import citydreamer.extensions.voxlib
@@ -336,6 +335,9 @@ def render_bg(
     assert hf_seg.size(2) == CONSTANTS["LAYOUT_VOL_SIZE"]
     assert hf_seg.size(3) == CONSTANTS["LAYOUT_VOL_SIZE"]
 
+    # Fix: operator torchvision::nms does not exist
+    import torchvision
+
     blurrer = torchvision.transforms.GaussianBlur(kernel_size=3, sigma=(2, 2))
     _voxel_id = copy.deepcopy(voxel_id)
     _voxel_id[voxel_id >= CONSTANTS["BLD_INS_LABEL_MIN"]] = CLASSES["BLD_FACADE"]

requirements.txt

@@ -1,6 +1,5 @@
---extra-index-url https://download.pytorch.org/whl/cu118
-torch==2.2.2
-torchvision
+torch==2.2.0
+torchvision==0.17.0
 
 easydict
 gradio