modify on ZeroGPU

app.py

@@ -1,15 +1,20 @@
 import gradio as gr
 import spaces
-import tempfile
 import os
 import torch
 import numpy as np
 from matplotlib.colors import LinearSegmentedColormap
 
-from app_config import CSS, HEADER, FOOTER, DEVICE
-import sample_cond
+from app_config import CSS, HEADER, FOOTER
+from sample_cond import CKPT_PATH, MODEL_CFG, load_model_from_config, sample
 
-model = sample_cond.load_model()
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+def load_model():
+    pl_sd = torch.load(CKPT_PATH, map_location="cpu")
+    model = load_model_from_config(MODEL_CFG.model, pl_sd["state_dict"])
+    return model
 
 
 def create_custom_colormap():
@@ -33,7 +38,7 @@ def colorize_depth(depth, log_scale):
 @spaces.GPU
 @torch.no_grad()
 def generate_lidar(model, cond):
-    img, pcd = sample_cond.sample(model, cond)
+    img, pcd = sample(model, cond)
     return img, pcd
 
 
@@ -46,6 +51,8 @@ def load_camera(image):
     return camera_cond
 
 
+model = load_model().to(DEVICE)
+
 with gr.Blocks(css=CSS) as demo:
     gr.Markdown(HEADER)
 

app_config.py

@@ -14,7 +14,6 @@ CSS = """
     max-height: 70vh;
     }
 """
-DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 HEADER = '''
 # LiDAR Diffusion

sample_cond.py

@@ -6,9 +6,8 @@ from omegaconf import OmegaConf
 from PIL import Image
 
 from lidm.models.diffusion.ddim import DDIMSampler
-from lidm.utils.misc_utils import instantiate_from_config, isimage, ismap
+from lidm.utils.misc_utils import instantiate_from_config
 from lidm.utils.lidar_utils import range2pcd
-from app_config import DEVICE
 
 
 CUSTOM_STEPS = 50
@@ -20,7 +19,7 @@ CFG_PATH = os.path.join(MODEL_PATH, 'config.yaml')
 CKPT_PATH = os.path.join(MODEL_PATH, 'model.ckpt')
 
 # settings
-model_config = OmegaConf.load(CFG_PATH)
+MODEL_CFG = OmegaConf.load(CFG_PATH)
 
 
 def custom_to_pcd(x, config, rgb=None):
@@ -64,20 +63,13 @@ def logs2pil(logs, keys=["sample"]):
     return imgs
 
 
-def load_model_from_config(config, sd, device):
+def load_model_from_config(config, sd):
     model = instantiate_from_config(config)
     model.load_state_dict(sd, strict=False)
-    model.to(device)
     model.eval()
     return model
 
 
-def load_model():
-    pl_sd = torch.load(CKPT_PATH, map_location="cpu")
-    model = load_model_from_config(model_config.model, pl_sd["state_dict"], DEVICE)
-    return model
-
-
 @torch.no_grad()
 def convsample_ddim(model, cond, steps, shape, eta=1.0, verbose=False):
     ddim = DDIMSampler(model)
@@ -103,7 +95,7 @@ def make_convolutional_sample(model, batch, batch_size, custom_steps=None, eta=1
 def sample(model, cond):
     batch = {'camera': cond}
     img = make_convolutional_sample(model, batch, batch_size=1, custom_steps=CUSTOM_STEPS, eta=ETA)  # TODO add arguments for batch_size, custom_steps and eta
-    pcd = custom_to_pcd(img, model_config)[0].astype(np.float32)
+    pcd = custom_to_pcd(img, MODEL_CFG)[0].astype(np.float32)
     img = img.squeeze().detach().cpu().numpy()
     return img, pcd