Switch to Gradio.

.gitignore

@@ -179,3 +179,5 @@ configs/
 data/
 notebooks/
 output/
+flagged/
+*.pth

ARTICLE.md

@@ -0,0 +1,20 @@
+##### Citation 📝
+
+If our work is useful for your research, please consider citing:
+
+```bibtex
+@inproceedings{xie2024citydreamer,
+  title     = {City{D}reamer: Compositional Generative Model of Unbounded 3{D} Cities},
+  author    = {Xie, Haozhe and 
+               Chen, Zhaoxi and 
+               Hong, Fangzhou and 
+               Liu, Ziwei},
+  booktitle = {CVPR},
+  year      = {2024}
+}
+```
+
+##### License 📋
+
+This project is licensed under [S-Lab License 1.0](https://huggingface.co/hzxie/city-dreamer/blob/main/LICENSE).
+Redistribution and use for non-commercial purposes should follow this license.

README.md

@@ -1,10 +1,10 @@
 ---
 title: CityDreamer
 emoji: 🏙️
-colorFrom: pink
-colorTo: indigo
-sdk: streamlit
-sdk_version: 1.31.1
+colorFrom: blue
+colorTo: green
+sdk: gradio
+sdk_version: 3.41.0
 app_file: app.py
 pinned: false
 license: other

app.py

@@ -4,54 +4,102 @@
 # @Author: Haozhe Xie
 # @Date:   2024-03-02 16:30:00
 # @Last Modified by: Haozhe Xie
-# @Last Modified at: 2024-03-02 22:33:17
+# @Last Modified at: 2024-03-03 10:25:43
 # @Email:  root@haozhexie.com
 
-import streamlit as st
+import logging
+import os
+import torch
+import gradio as gr
+import subprocess
+import urllib.request
+import ssl
+import sys
 
-from PIL import Image
+# Fix: ssl.SSLCertVerificationError: [SSL: CERTIFICATE_VERIFY_FAILED] certificate verify failed
+ssl._create_default_https_context = ssl._create_unverified_context
 
-LOGGER = st.logger.get_logger(__name__)
+sys.path.append(os.path.join(os.path.dirname(__file__), "citydreamer"))
+# Import CityDreamer modules
+# import citydreamer.model
+# import citydreamer.inference
 
 
 def setup_runtime_env():
-    pass
+    subprocess.call(["pip", "freeze"])
+    ext_dir = os.path.join(os.path.dirname(__file__), "citydreamer", "extensions")
+    for e in os.listdir(ext_dir):
+        if not os.path.isdir(e):
+            continue
+        subprocess.call(["pip", "install", "."], workdir=os.path.join(ext_dir, e))
 
 
 def get_models():
-    return None, None
+    if not os.path.exists("CityDreamer-Fgnd.pth"):
+        urllib.request.urlretrieve(
+            "https://huggingface.co/hzxie/city-dreamer/resolve/main/CityDreamer-Fgnd.pth",
+            "CityDreamer-Fgnd.pth",
+        )
+    if not os.path.exists("CityDreamer-Bgnd.pth"):
+        urllib.request.urlretrieve(
+            "https://huggingface.co/hzxie/city-dreamer/resolve/main/CityDreamer-Bgnd.pth",
+            "CityDreamer-Bgnd.pth",
+        )
+
+    bgm_ckpt = torch.load("CityDreamer-Bgnd.pth")
+    fgm_ckpt = torch.load("CityDreamer-Fgnd.pth")
+    bgm = citydreamer.model.GanCraftGenerator(bgm_ckpt["cfg"])
+    fgm = citydreamer.model.GanCraftGenerator(fgm_ckpt["cfg"])
+    if torch.cuda.is_available():
+        fgm = torch.nn.DataParallel(fgm).cuda().eval()
+        bgm = torch.nn.DataParallel(bgm).cuda().eval()
+
+    return bgm, fgm
 
 
 def get_generated_city(radius, altitude, azimuth):
-    pass
+    print(radius, altitude, azimuth)
 
 
-def main(fgm, bgm):
-    st.set_page_config(
-        page_title="CityDreamer Demo",
-        page_icon="🏙️",
-    )
-    # Main
-    st.write("# CityDreamer Minimal Demo 🏙️")
+def main(debug):
+    title = "CityDreamer Demo 🏙️"
     with open("README.md", "r") as f:
         markdown = f.read()
-        st.markdown(markdown[markdown.rfind("---") :])
-    imgbox = st.empty()
+        desc = markdown[markdown.rfind("---") + 3:]
+    with open("ARTICLE.md", "r") as f:
+        arti = f.read()
+    with open("assets/style.css") as f:
+        css = f.read()
 
-    # Sidebar
-    st.sidebar.header("CityDreamer Settings")
-    radius = st.sidebar.slider("Camera Radius (m)", 128, 512, 320, 5)
-    altitude = st.sidebar.slider("Camera Altitude (m)", 256, 512, 384, 5)
-    azimuth = st.sidebar.slider("Camera Azimuth (°)", 0, 360, 180, 5)
-    if st.sidebar.button("Generate", type="primary"):
-        img = get_generated_city(radius, altitude, azimuth)
-        imgbox.image(img, caption="CityDreamer Generation")
+    app = gr.Interface(
+        get_generated_city,
+        [
+            gr.Slider(
+                128, 512, value=320, step=5, label="Camera Radius (m)"
+            ),
+            gr.Slider(
+                256, 512, value=384, step=5, label="Camera Altitude (m)"
+            ),
+            gr.Slider(0, 360, value=180, step=5, label="Camera Azimuth (°)"),
+        ],
+        [gr.Image(type="numpy", label="Generated City")],
+        title=title,
+        description=desc,
+        article=arti,
+        allow_flagging="never",
+        css=css,
+    )
+    app.queue(api_open=False)
+    app.launch(debug=debug)
 
 
 if __name__ == "__main__":
-    LOGGER.info("Setting up runtime environment...")
-    setup_runtime_env()
-    fgm, bgm = get_models()
-
-    LOGGER.info("Starting the main application...")
-    main(fgm, bgm)
+    logging.basicConfig(
+        format="[%(levelname)s] %(asctime)s %(message)s", level=logging.INFO
+    )
+    logging.info("Compile CUDA extensions...")
+    # setup_runtime_env()
+    logging.info("Downloading pretrained models...")
+    # fgm, bgm = get_models()
+    logging.info("Starting the main application...")
+    main(os.getenv("DEBUG") == "1")

assets/style.css

@@ -0,0 +1,7 @@
+h1 {
+    display: block;
+}
+
+p img {
+    display: inline-block;
+}

requirements.txt

@@ -1,2 +1,7 @@
-requests
-streamlit
+--extra-index-url https://download.pytorch.org/whl/cu113
+torch==1.12.0
+torchvision
+
+numpy
+opencv-python
+gradio